[Perf] minimax-h3: Skip attention for the packing's alignment rows in the Ulysses pipeline

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -66,6 +66,9 @@
 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
     AttentionRequirements,
 )
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import (
+    DynamicCudnnSDPAImpl,
+)
 from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn_h3 import (
     vsa_h3_fold_gate,
 )
@@ -735,17 +738,56 @@ def _minimax_h3_pipelined_dense_attention(
         ulysses_pipelined_attention,
     )
 
+    attend = _minimax_h3_group_attend(
+        impl,
+        cu_seqlens=cu_seqlens,
+        cu_seqlens_host=cu_seqlens_host,
+        max_seqlen=max_seqlen,
+    )
+    return ulysses_pipelined_attention(q, k, v, attend, groups, fill=fill)
+
+
+def _minimax_h3_group_attend(
+    impl,
+    *,
+    cu_seqlens: torch.Tensor,
+    cu_seqlens_host: tuple[int, ...] | None,
+    max_seqlen: int,
+):
+    """One head group's attention over the gathered packed sequence.
+
+    Trailing alignment rows (cu_seqlens ``[0, used, total]``) only reach their
+    own rows, so they stay zero; the real rows run cuDNN SDPA varlen's own
+    per-segment call, written straight into the output when the backend can.
+    """
+    used = max_seqlen
+    pad_only_tail = (
+        isinstance(impl, DynamicCudnnSDPAImpl)
+        and cu_seqlens_host is not None
+        and len(cu_seqlens_host) == 3
+        and cu_seqlens_host[:2] == (0, used)
+        and cu_seqlens_host[2] > used
+    )
+
     def attend(q_group, k_group, v_group):
-        return impl.forward_varlen(
-            q_group,
-            k_group,
-            v_group,
-            cu_seqlens=cu_seqlens,
-            max_seqlen=max_seqlen,
-            cu_seqlens_host=cu_seqlens_host,
-        )
+        if not pad_only_tail:
+            return impl.forward_varlen(
+                q_group,
+                k_group,
+                v_group,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+                cu_seqlens_host=cu_seqlens_host,
+            )
+        out = torch.empty_like(q_group)
+        real = [t[:used].unsqueeze(0) for t in (q_group, k_group, v_group)]
+        dst = out[:used].unsqueeze(0)
+        if not impl.forward_into(*real, dst, None):
+            dst.copy_(impl.forward(*real, None))
+        out[used:].zero_()
+        return out
 
-    return ulysses_pipelined_attention(q, k, v, attend, groups, fill=fill)
+    return attend
 
 
 def _minimax_h3_qknorm_rope_pipelined_attention(

python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py

@@ -24,7 +24,10 @@
 from sglang.multimodal_gen.runtime.layers.attention.backends.cube_sparse_attn import (
     CubeSparseAttentionBackend,
 )
-from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPAImpl
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import (
+    DynamicCudnnSDPAImpl,
+    SDPAImpl,
+)
 from sglang.multimodal_gen.runtime.layers.attention.selector import (
     component_attn_backend_context_manager,
 )
@@ -46,6 +49,7 @@
     MiniMaxH3DiTModel,
     _copy_grouped_qkv_tp_shard,
     _diffusers_h3_checkpoint,
+    _minimax_h3_group_attend,
     _modulate_gate,
     _qkv_scale_block_rows,
     _reorder_grouped_qkv_to_qkv,
@@ -415,6 +419,31 @@ def test_cache_dit_out_of_place_gate_preserves_cuda_input():
     torch.testing.assert_close(output, expected, rtol=0, atol=0)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_group_attend_skips_only_the_trailing_alignment_rows():
+    """The pipelined Ulysses group attention gives the real rows the bytes of
+    forward_varlen and leaves the packing's alignment rows zero."""
+    used, total = 16400, 16448  # long enough for FA4 on sm_103
+    impl = DynamicCudnnSDPAImpl(
+        num_heads=2, head_size=128, causal=False, softmax_scale=128**-0.5
+    )
+    torch.manual_seed(0)
+    q, k, v = (
+        torch.randn(total, 2, 128, device="cuda", dtype=torch.bfloat16)
+        for _ in range(3)
+    )
+    bounds = (0, used, total)
+    kwargs = dict(
+        cu_seqlens=torch.tensor(bounds, device="cuda", dtype=torch.int32),
+        cu_seqlens_host=bounds,
+        max_seqlen=used,
+    )
+    ref = impl.forward_varlen(q, k, v, **kwargs)
+    got = _minimax_h3_group_attend(impl, **kwargs)(q, k, v)
+    assert torch.equal(got[:used], ref[:used])
+    assert not got[used:].any()
+
+
 def test_tp_and_ulysses_admission_uses_tp_local_shapes():
     arch = MiniMaxH3DiTArchConfig()
     model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel)