Skip to content

fix: reduce MiniMax-H3 VRAM spikes and redundant token refinemen - #1900

Open
assouan wants to merge 2 commits into
leejet:masterfrom
assouan:model/minimax-h3/token-refiner-cache
Open

fix: reduce MiniMax-H3 VRAM spikes and redundant token refinemen#1900
assouan wants to merge 2 commits into
leejet:masterfrom
assouan:model/minimax-h3/token-refiner-cache

Conversation

@assouan

@assouan assouan commented Aug 22, 2026

Copy link
Copy Markdown

Summary

This fixes two MiniMax-H3 issues that made sampling slower and could cause a VRAM spike with --stream-layers.

Before this PR:

  • The condition projection and token refiner blocks were executed again at every sampling step, even though their input does not change during sampling.
  • The first streaming segment simultaneously loaded the condition projection, both token refiners, and transformer block 0 together. This defeats the one-layer-at-a-time behavior expected from --stream-layers and can use much more VRAM than expected.

After this PR:

  • The refined context is computed once per conditioning context, cached for the current sampling run, and reused for the following steps.
  • The condition projection, token refiner 0, token refiner 1, and transformer block 0 are handled by separate streamed segments. Each set of weights is released before the next one is loaded.
  • The last refiner writes directly to the persistent cache, so there is no extra copy-only segment.
  • Cache entries also track the real source context and active weight adapter, and the cache is cleared when sampling ends.

This removes repeated work from every step while restoring the low-VRAM behavior expected from layer streaming.

Related Issue / Discussion

N/A

Additional Information

Tested with MiniMax-H3 Q4 using:

  • 256x256, 22 frames, 2 sampling steps, CFG 1
  • --stream-layers
  • --max-vram cuda1=0.1

The runtime log confirms this order, with every segment marked as STREAMED:

Load condition projection -> compute -> free
Load token refiner 0      -> compute -> free
Load token refiner 1      -> compute -> free
Load transformer block 0  -> compute -> free

The token refiners ran once for the whole sampling run, while transformer block 0 ran once per step. CPU and CUDA builds passed, and the generated output remained unchanged.

Checklist

Run condition projection and both token-refiner blocks once per
conditioning context, then reuse the refined output across denoising
steps. Keep entries distinct by condition and active weight adapter.

Split each refiner block into its own streaming segment and clear the
sampling-scoped cache when sampling finishes.
Split condition projection and token refiners into sequential streamed segments while folding the persistent copy into the final refiner segment.

Tie cached refined contexts to the source tensor identity, storage, shape, and active weight adapter.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant