From 7088f21922edc633dbb38e7543ef2b33252dfcbf Mon Sep 17 00:00:00 2001
From: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com>
Date: Thu, 27 Aug 2026 21:22:07 -0700
Subject: [PATCH] [diffusion] fix: make tail_attn_meta cuda-graph capturable
 (#36658)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
---
 .../runtime/distributed/sp_shard_utils.py     |  6 ++---
 .../multimodal_gen/test/unit/test_sp_shard.py | 23 +++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py b/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py
index bd603d245..826733e6f 100644
--- a/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py
+++ b/python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py
@@ -195,10 +195,10 @@ def tail_attn_meta(
         return None
     seq = shard.sp_size * (shard.local_len + image_seq_len)
     valid = seq - shard.num_pad
-    row = torch.tensor([valid, shard.num_pad], dtype=torch.int32, device=device)
-    seglens = row.repeat(batch_size)
+    row_starts = torch.arange(batch_size, dtype=torch.int32, device=device) * seq
     cu_seqlens = torch.zeros(2 * batch_size + 1, dtype=torch.int32, device=device)
-    cu_seqlens[1:] = torch.cumsum(seglens, dim=0)
+    cu_seqlens[1::2] = row_starts + valid
+    cu_seqlens[2::2] = row_starts + seq
     return {
         "pad_start": valid,
         "pad_end": seq,
diff --git a/python/sglang/multimodal_gen/test/unit/test_sp_shard.py b/python/sglang/multimodal_gen/test/unit/test_sp_shard.py
index b467ffb6a..6edcae21d 100644
--- a/python/sglang/multimodal_gen/test/unit/test_sp_shard.py
+++ b/python/sglang/multimodal_gen/test/unit/test_sp_shard.py
@@ -108,6 +108,29 @@ def test_tail_meta_max_seqlen_covers_pad_segment():
     assert meta["max_seqlen_tail"] == 3
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA graph capture")
+def test_tail_meta_is_cuda_graph_capturable():
+    """The meta is built inside the DiT forward, which breakable-cuda-graph
+    captures; host-staged tensor construction aborts capture with 'Cannot copy
+    between CPU and CUDA tensors'."""
+    device = torch.device("cuda")
+    shard = SpShard(orig_len=15, local_len=8, num_pad=1, sp_size=2, sp_rank=1)
+    eager = tail_attn_meta(shard, 2, device, image_seq_len=100)
+
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        tail_attn_meta(shard, 2, device, image_seq_len=100)
+    torch.cuda.current_stream().wait_stream(side)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = tail_attn_meta(shard, 2, device, image_seq_len=100)
+    graph.replay()
+    torch.cuda.synchronize()
+    assert torch.equal(captured["cu_seqlens_tail"], eager["cu_seqlens_tail"])
+
+
 def test_tail_meta_matches_legacy_gap_formula():
     # The tail layout puts the pad exactly where the legacy per-model gap
     # formula pointed, minus the relocation: end == S (global tail).