[diffusion] MiniMax-H3: run the attention section eager under torch.compile

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):

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