From 6370f8b537990be6df88e9074cc5a7c431f36a37 Mon Sep 17 00:00:00 2001 From: Mick Qian Date: Sun, 11 Oct 2026 09:18:42 +0800 Subject: [PATCH] [diffusion] MiniMax-H3: run the attention section eager under torch.compile The attention section gates its fast paths on not is_compiling(): the fused qk-norm + RoPE, the head-group pipelined Ulysses exchange over the copy engine and the IPC all-to-all are untraceable. A compiled DiT therefore took the unfused norm/RoPE and a sequential exchange, and served slower than eager (2xGB300 t2va: 61.5 s compiled vs 57.7 s eager). Move the section from the q/k norm to the attention output into MiniMaxH3Attention._attend. Eager calls it as before; under compile forward calls torch.compiler.disable(_attend), so the section runs the eager fast paths and inductor compiles the projections, MLP, AdaLN and norms around it. Co-Authored-By: Claude Opus 5.5 --- .../runtime/models/dits/minimax_h3.py | 50 ++++++++++++++++--- .../test/unit/test_minimax_h3_dit_contract.py | 32 ++++++++++++ 2 files changed, 74 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 970cdc63ea04..01002aefd6ee 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -1296,6 +1296,43 @@ def forward( ulysses_active=ulysses_active, ring_active=ring_active, ) + # The fused qk-norm/RoPE, pipelined Ulysses exchange and attention kernels + # are untraceable, and the traceable fallback is slower than eager, so a + # compiled DiT runs this section eager and compiles around it. + attend = self._attend_eager if torch.compiler.is_compiling() else self._attend + out = attend( + x, + q, + k, + v, + rope_cache=rope_cache, + cu_seqlens=cu_seqlens, + cu_seqlens_host=cu_seqlens_host, + max_seqlen=max_seqlen, + subblock_sparse_query_block_mask=subblock_sparse_query_block_mask, + ulysses_active=ulysses_active, + ring_active=ring_active, + ) + out, _ = self.out_proj(out) + return out + + def _attend( + self, + x: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + rope_cache: tuple[torch.Tensor, torch.Tensor] | None, + cu_seqlens: torch.Tensor, + cu_seqlens_host: tuple[int, ...] | None, + max_seqlen: int, + subblock_sparse_query_block_mask: torch.Tensor | None, + ulysses_active: bool, + ring_active: bool, + ) -> torch.Tensor: + """q/k/v [T, heads, head_dim] -> attention output [T, heads * head_dim].""" + total = x.shape[0] if rope_cache is None: q, k = _apply_qk_norm( q, @@ -1306,7 +1343,7 @@ def forward( ) else: cos_sin_cache, positions = rope_cache - if self._use_fused_qknorm_rope and not torch.compiler.is_compiling(): + if self._use_fused_qknorm_rope: if ulysses_active and not ring_active: out = _minimax_h3_qknorm_rope_pipelined_attention( self, @@ -1319,10 +1356,7 @@ def forward( max_seqlen=max_seqlen, ) if out is not None: - out, _ = self.out_proj( - out.reshape(total, self.num_heads * self.head_dim) - ) - return out + return out.reshape(total, self.num_heads * self.head_dim) fused_inplace_qknorm_rope( q, k, @@ -1372,9 +1406,9 @@ def forward( ring_active=ring_active, gate_compress=gate_compress, ) - out = out.reshape(total, self.num_heads * self.head_dim) - out, _ = self.out_proj(out) - return out + return out.reshape(total, self.num_heads * self.head_dim) + + _attend_eager = torch.compiler.disable(_attend) class MiniMaxH3MLP(nn.Module): diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py index a6179eda033e..6738dc81cbdf 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py @@ -850,3 +850,35 @@ def test_cuda_ulysses_qkv_pack_is_bit_exact(): actual = pack_qkv_destination_major(q.contiguous(), k.contiguous(), v, world_size) torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("compiling", [False, True]) +def test_compiled_attention_runs_its_fused_section_eager(compiling): + _ensure_single_process_parallel_runtime() + arch = MiniMaxH3DiTArchConfig() + with torch.device("meta"): + attention = MiniMaxH3Attention(arch, None, prefix="blocks.0.attn") + tokens = 4 + out = torch.empty(tokens, attention.local_inner_dim, device="meta") + + with ( + patch.object(torch.compiler, "is_compiling", return_value=compiling), + patch.object(MiniMaxH3Attention, "_attend", return_value=out) as attend, + patch.object( + MiniMaxH3Attention, "_attend_eager", return_value=out + ) as attend_eager, + ): + result = attention( + torch.empty(tokens, arch.hidden_size, device="meta"), + rope_cache=None, + cu_seqlens=torch.empty(2, dtype=torch.int32, device="meta"), + max_seqlen=tokens, + ) + + assert result.shape == (tokens, arch.hidden_size) + assert (attend_eager if compiling else attend).call_count == 1 + assert (attend if compiling else attend_eager).call_count == 0 + + +def test_eager_attention_section_is_hidden_from_dynamo(): + assert MiniMaxH3Attention._attend_eager._torchdynamo_disable