Repository navigation
Conversation
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
mergify
Bot
added
type: perf
 [Low impact] The change is limited to an opt-in path in MiniMax-H3 video VAE decoding. Fix the rotary layout and gradient guards, and satisfy the comment requirement before merging. Findings
Fix with agent prompt### Issue 1
fastvideo/models/vaes/minimax_h3_video.py:324-327
`_can_fuse_qk_norm_rope` accepts broadcast or strided rotary tables, but the kernel reads them as packed `[B, S, 1, RD]` buffers. Passing `[1, S, 1, RD]` tables with batch size 2 works in the eager path, but the fused path reads past those tables for the second batch. This can produce wrong results or an illegal GPU memory access.
Check both tables' shapes, devices, and contiguity before selecting the kernel; otherwise use the eager path.
### Issue 2
fastvideo/models/vaes/minimax_h3_video.py:386-391
With the flag enabled, the fused branch also runs on CUDA calls that need gradients. Its in-place Triton writes record no backward operation for RMSNorm or RoPE. A decoder forward followed by backward therefore silently computes the wrong gradients for the query/key projections and earlier layers.
The decoder explicitly supports gradient checkpointing. Keep gradient-tracked calls on the eager path, as the existing DiT fusion does.
### Issue 3
fastvideo/models/vaes/minimax_h3_video.py:387-388
This comment narrates which operations the fused call replaces and repeats the kernel documentation. The repository requires removing comments outside its small keep-list, which includes required non-obvious behavior and public API docs-not performance narration. Remove these two lines before merging to satisfy that requirement.
```suggestion
```
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if not key.is_contiguous() or key.shape != query.shape or not query.is_floating_point(): | ||
| return False | ||
| head_dim, rotary_dim = query.shape[-1], rotary_emb[0].shape[-1] | ||
| return 0 < rotary_dim <= head_dim and rotary_dim % 2 == 0 and head_dim & (head_dim - 1) == 0 |
There was a problem hiding this comment.
_can_fuse_qk_norm_rope accepts broadcast or strided rotary tables, but the kernel reads them as packed [B, S, 1, RD] buffers. Passing [1, S, 1, RD] tables with batch size 2 works in the eager path, but the fused path reads past those tables for the second batch. This can produce wrong results or an illegal GPU memory access.
Check both tables' shapes, devices, and contiguity before selecting the kernel; otherwise use the eager path.
Prompt To Fix With AI
This is a comment left during a code review.
Path: fastvideo/models/vaes/minimax_h3_video.py
Line: 324-327
Comment:
**Shared rotary tables break attention**
`_can_fuse_qk_norm_rope` accepts broadcast or strided rotary tables, but the kernel reads them as packed `[B, S, 1, RD]` buffers. Passing `[1, S, 1, RD]` tables with batch size 2 works in the eager path, but the fused path reads past those tables for the second batch. This can produce wrong results or an illegal GPU memory access.
Check both tables' shapes, devices, and contiguity before selecting the kernel; otherwise use the eager path.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.There was a problem hiding this comment.
Fixed in aa2fa42: the gate now requires contiguous [B, S, 1, RD] cos/sin tables on the query's device, with B and S matching the query. Anything else, including broadcast [1, S, 1, RD] tables, takes the eager path. The production rotary embedding still qualifies because its position ids are materialized per batch.
| if envs.FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE.get() and _can_fuse_qk_norm_rope(query, key, rotary_emb): | ||
| # One Triton pass replaces the two fp32 RMSNorm kernels, the two activation-dtype | ||
| # casts and the four cat/neg/mul/add kernels of the eager prologue. | ||
| from fastvideo.models.vaes.minimax_h3_qk_norm_rope import qk_norm_rope_ | ||
| cos, sin = rotary_emb | ||
| cos = cos.to(query.dtype) | ||
| sin = sin.to(query.dtype) | ||
| rotary_dim = cos.shape[-1] | ||
| query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] | ||
| key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] | ||
| query_first, query_second = query_rotary.chunk(2, dim=-1) | ||
| key_first, key_second = key_rotary.chunk(2, dim=-1) | ||
| query_rotated = torch.cat([-query_second, query_first], dim=-1) | ||
| key_rotated = torch.cat([-key_second, key_first], dim=-1) | ||
| query = torch.cat([query_rotary * cos + query_rotated * sin, query_pass], dim=-1) | ||
| key = torch.cat([key_rotary * cos + key_rotated * sin, key_pass], dim=-1) | ||
| qk_norm_rope_(query, key, cos, sin, self.heads, query.shape[-1], cos.shape[-1], float(self.norm_q.eps)) |
There was a problem hiding this comment.
With the flag enabled, the fused branch also runs on CUDA calls that need gradients. Its in-place Triton writes record no backward operation for RMSNorm or RoPE. A decoder forward followed by backward therefore silently computes the wrong gradients for the query/key projections and earlier layers.
The decoder explicitly supports gradient checkpointing. Keep gradient-tracked calls on the eager path, as the existing DiT fusion does.
Prompt To Fix With AI
This is a comment left during a code review.
Path: fastvideo/models/vaes/minimax_h3_video.py
Line: 386-391
Comment:
**Decoder gradients become wrong**
With the flag enabled, the fused branch also runs on CUDA calls that need gradients. Its in-place Triton writes record no backward operation for RMSNorm or RoPE. A decoder forward followed by backward therefore silently computes the wrong gradients for the query/key projections and earlier layers.
The decoder explicitly supports gradient checkpointing. Keep gradient-tracked calls on the eager path, as the existing DiT fusion does.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.There was a problem hiding this comment.
Fixed in aa2fa42: calls that track gradients (torch.is_grad_enabled() with Q or K requiring grad) now use the eager path, matching the DiT fusion gate.
| # One Triton pass replaces the two fp32 RMSNorm kernels, the two activation-dtype | ||
| # casts and the four cat/neg/mul/add kernels of the eager prologue. |
There was a problem hiding this comment.
This comment narrates which operations the fused call replaces and repeats the kernel documentation. The repository requires removing comments outside its small keep-list, which includes required non-obvious behavior and public API docs—not performance narration. Remove these two lines before merging to satisfy that requirement.
| # One Triton pass replaces the two fp32 RMSNorm kernels, the two activation-dtype | |
| # casts and the four cat/neg/mul/add kernels of the eager prologue. |
Rule Used: What: Remove all unwanted code comments in the scoped diff except for a small keep-list (license headers, dependency/protocol-required non-obvious behavior, // prettier-ignore, public API doc comments, and issue/RFC constraint links). Why: Comment... (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: fastvideo/models/vaes/minimax_h3_video.py
Line: 387-388
Comment:
**Added comment violates repository rules**
This comment narrates which operations the fused call replaces and repeats the kernel documentation. The repository requires removing comments outside its small keep-list, which includes required non-obvious behavior and public API docs—not performance narration. Remove these two lines before merging to satisfy that requirement.
```suggestion
```
**Rule Used:** What: Remove all unwanted code comments in the scoped diff except for a small keep-list (license headers, dependency/protocol-required non-obvious behavior, `// prettier-ignore`, public API doc comments, and issue/RFC constraint links). Why: Comment... ([source](https://app.greptile.com/hao-ai-lab/-/custom-context?memory=856788fb-b8e1-47f8-b01b-d83ea205fc6b))
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Summary
FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE=1.How this was found. These changes came out of an agent-driven kernel optimization run with MotionKernel, a tool we're developing that profiles a real generation, finds where the time goes, and only keeps changes that hold up end to end with the quality checks below.
Problem
Every attention layer of the H3 video VAE decoder ran its Q/K prologue as ten kernels:
At 768x1344 that is ~1.9 ms per layer (measured; 36 layers x 21 batched tile decodes at tile
batch 8 = ~1.4 s of a 3.65 s decode), and most of it is kernel-launch and elementwise traffic on
[8, 1797, 32, 64] tensors.
Fix
One Triton kernel (
fastvideo/models/vaes/minimax_h3_qk_norm_rope.py) reads Q and K once,normalizes in fp32, rounds to the activation dtype exactly as eager does, applies the rotate-half
RoPE on the rope prefix with the same per-op rounding, and stores back in place: one launch per
attention layer per chunk instead of ten. Opt-in via
FASTVIDEO_H3_VAE_FUSED_QK_NORM_ROPE=1; the eager prologue is factored into_apply_rotaryandkept as the CPU / unsupported-shape fallback (
_can_fuse_qk_norm_ropechecks CUDA, contiguity,even rope width and a power-of-two head dim).
Evidence
Decode-only, tile batch 8, one loaded VAE, warm, interleaved, GB200 (768x1344, fp16 autocast):
1.875 ms -> 0.090 ms per layer (20.8x); the existing DiT
fused_qknorm_ropekernel takes0.619 ms at the same shape.
Parity (this is an exact-precision fusion, not a bitwise one)
The arithmetic is the eager arithmetic; only the fp32 RMSNorm reduction order differs (Triton's
tree vs
at::nativelayer norm). Measured on the real VAE decode:Every differing pixel is off by exactly one 8-bit code value. Q/K parity on random
[8, 337, 32, 64] fp16 inputs including register/cls suffix tokens: bit-equal 99.998%, rel-L2
3.7e-6, max |d| 1 ULP (bf16: 99.9996%, 9.9e-6). This matches how upstream already treats the DiT
fused_qknorm_rope("Triton's reduction order ... can produce small, expected rounding drift").Tests
Verification on current main (d9e6c3e): the patch applies cleanly; pre-commit (yapf, ruff, codespell, PyMarkdown, mypy) passes on all changed files; the tests listed above were run on GB200.
Independent GB200 verification (after the review fixes)
An agent rechecked this PR on a separate GB200: