MiniMax-H3: sage takes the fused RoPE op, faster video and audio VAE decode (B200 max render 22.6 to 15.4 s, video bit-identical) - #19
Open
danielhanchen wants to merge 7 commits into
Conversation
With --sage-attn, forward_head_major returned nullptr, so every block fell back to the unfused chunk / slice / rope / concat / scale / cast chain (about 0.75 s per 960x544x124 step on B200). ggml_rope_pe_permute already writes [head_dim, tokens, heads, batch], the layout ggml_sage_attn reads: emit F32 Q, F32 K * kv_scale and F16 V * kv_scale and call sage directly. Same arithmetic and rounding points as the unfused sage branch; the rendered video and audio are bit-identical. SD_H3_FAST_SAGE_QKV=0 restores the old chain.
The 960x544x124 decode left the B200 idle for about 2 s between its 105 tile graphs: the host blended each tile, zero-filled and assembled the frames, rebuilt the rotary tables and queried free device memory four times per tile while the GPU waited. - tiling: the blend of a batch into the output runs on a worker thread while the next batch is split and computed; merges stay serial and in tile order with the same arithmetic (SD_TILE_ASYNC_MERGE=0 merges inline). - H3 decode: each temporal chunk is trimmed, cross-faded and copied straight into the trimmed output on a worker thread while the next chunk decodes, instead of concatenating and slicing at the end (SD_H3_VAE_ASYNC_ASSEMBLY=0 restores it). - rotary tables are rebuilt only when the tile shape changes. - DeviceMemoryRequest::reuse_device_query: while a runner repeats identical computes, a capacity check that allocates nothing new (no pending bytes, all parameters resident) reuses the owner's earlier free-memory reading; runner_end() drops it. The H3 decode turns it on for its duration (SD_H3_VAE_REUSE_MEMQUERY=0 disables). Warm decode on B200 6.8-7.0 s -> 5.06 s; frames and audio bit-identical.
…onvs The up/down-sampling filters of every Activation1D went through ggml_conv_1d / ggml_conv_1d_dw: an F16 im2col plus a single-column matmul per call. Route them through CONV_2D_DW (H=1, F32 per-channel kernel) instead, for the MiniMax-H3 audio VAE only. Falls back to the im2col graph when the backend lacks the op. SD_H3_AUDIO_DIRECT_DW=0 restores the previous graph.
…tion The decoder attention chunk-copied q/k/v out of the fused projection, applied the partial RoPE as table mul/add/concat chains and then permuted and cast K/V for flash attention. With one tile per graph, normalise q/k in place on the projection and let ggml_rope_pe_permute (already used by the DiT) write the F32 Q and F16 K/V head-major tensors ggml_ext_attention_prepared reads. Same products, sums and casts as the table path, so the decoded frames are bit-identical. SD_H3_VAE_FUSED_QKV=0 restores the old graph; batched tiles, sage and non-flash decodes keep it.
Carry ggml patch 0005 (ggml_rope_pe_permute_rms) and use it for the decoder's q/k: the per-head norm was 7560 launches of a 64-thread kernel, about 0.45 s per 960x544x124 decode on B200. The fused op normalises each head row with the same block reduction as the standalone norm, so the frames stay bit-identical. Backends without it (HIP, Vulkan, GGML_CUDA_NORM_SMALL_ROWS=0) keep the separate norm; SD_H3_VAE_FUSED_QK_NORM=0 forces that.
Fewer, larger tiles cut the per-tile fixed cost (20x20 decodes 960x544x124 in 3.07 s vs 3.77 s at 16x16 on B200) but move the seams: 36.2 dB PSNR / 0.09 LPIPS vs the 16x16 decode, slightly closer to an untiled decode than 16x16 is. Default stays 16.
- The saved device free-memory reading is now taken only by a check that allocates nothing, so it is never the reading from before this owner's compute buffer was allocated, and any check that does not fit drops the owner's readings so the reclaim / evict retries see fresh ones. - If a worker thread cannot be started, the tile merge and the chunk assembly run inline instead of throwing.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #18 (base branch
perf/h3-gguf-speed-v2). Review that first.Problem
An Nsight Systems profile of a warm MiniMax-H3 render on a B200 (960x544, 124 frames, 4 steps, UD-Q3_K_XL DiT) put this fork's max mode at 21.4 s against 14.6 to 15.2 s for ComfyUI's int8 graph. The gap had three parts:
forward_head_major()returns nullptr, so the fused RoPE / permute op is skipped. About 765 ms per step comes back as f32 RoPE, permute copies, casts and elementwise kernels.cudaMemGetInfocalls per decode;mul_mat_vec_fplus im2col (1.0 s of GPU time, against 37 ms in ComfyUI).Change
Each lever has an environment switch;
=0turns it off.SD_H3_FAST_SAGE_QKVSD_TILE_ASYNC_MERGESD_H3_VAE_ASYNC_ASSEMBLYcudaMemGetInfocalls)SD_H3_VAE_REUSE_MEMQUERYSD_H3_VAE_FUSED_QKVSD_H3_VAE_FUSED_QK_NORMCONV_2D_DWin F32 (LTX unchanged)SD_H3_AUDIO_DIRECT_DWSD_H3_VAE_TILE=Nscripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patchadds the fused norm. Patches 0001 to 0005 apply cleanly to the pinned ggml and reproduce the patched tree; the ggml gitlink is unchanged.Results
Warm renders (renders 1 and 2), seconds:
Accuracy
Video: the output stream is bit-identical before and after on B200, G4 and A100, in max and default modes.
All switches off: video and audio reproduce the base build exactly.
Audio: the direct depthwise path changes the audio on purpose. SNR against the base build: 41.5 dB on B200, 41.7 / 42.7 dB on G4, 35.0 / 41.6 dB on A100 (max / default). Against the same graph run fully in F32 it is closer than before:
Tests
test-backend-opson CUDA (B200):ROPE_PE_PERMUTE(80 cases, including the new norm variants),RMS_NORMandCONV_2D_DWpass.Not covered
GGML_CUDA_QUANT_CUBLAS_MIN_BATCH) was measured on the B200 without sage: 7.20 to 3.12 s per step, 0.4 GB more VRAM, and closer to an F32-dequant reference than MMQ (44.4 vs 37.9 dB PSNR, 0.017 vs 0.044 LPIPS, 23.2 vs 14.7 dB audio SNR). Making it the default is a Studio-side change, proposed separately.