[sp] Make attention-TP sequence sharding a per-forward batch property (#37546)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -46,7 +46,8 @@ class _MiniForwardBatch:
     encoder_lens: Optional[torch.Tensor] = None
     mrope_positions: Optional[torch.Tensor] = None
     num_token_non_padded: Optional[torch.Tensor] = None
-    num_token_non_padded_cpu: Optional[int] = None
+    global_num_token_non_padded: Optional[torch.Tensor] = None
+    global_num_token_non_padded_cpu: Optional[int] = None
     global_num_tokens_gpu: Optional[torch.Tensor] = None
     global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
     ngram_embedding_info: Optional[object] = None
@@ -818,8 +819,9 @@ class TestBuildDecodeRegistry(unittest.TestCase):
             global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
             global_num_tokens_for_logprob_gpu=torch.zeros(1, dtype=torch.int32),
         )
-        # Gathered (DP) path: post_fill overwrites the FB copy with the local
-        # count. Pin attn-TP (size=2, rank=0) so the result is deterministic.
+        # Sharded (SP-on) forward: post_fill derives the LOCAL count from the
+        # invariant GLOBAL scalar. Pin attn-TP (size=2, rank=0) so the result
+        # is deterministic.
         with get_parallel().override(attn_tp_size=2, attn_tp_rank=0):
             reg = build_decode_registry(
                 device=torch.device("cpu"),
@@ -829,18 +831,67 @@ class TestBuildDecodeRegistry(unittest.TestCase):
                 cache_loc_dtype=torch.int64,
                 enable_num_token_non_padded=True,
                 require_gathered_buffer=True,
+                attn_tp_sharded_fn=lambda num_tokens: True,
                 source=src,
             )
             fb = _MiniForwardBatch(
-                num_token_non_padded=torch.tensor([100], dtype=torch.int32),
+                global_num_token_non_padded=torch.tensor([100], dtype=torch.int32),
             )
             reg.fill_from(
                 fb, raw_bs=4, padded_bs=4, raw_num_tokens=4, padded_num_tokens=8
             )
         # tokens_per_rank = padded_num_tokens(8) // attn_tp_size(2) = 4;
-        # local = clamp(100 - rank*4, 0, 4) = 4  (NOT the raw FB copy of 100).
+        # local = clamp(global(100) - rank*4, 0, 4) = 4.
         self.assertEqual(int(src.num_token_non_padded.item()), 4)
 
+    def test_num_token_non_padded_bypass_carries_local_count(self):
+        # Regression: the dense SBD draft and TBO sub-batches bypass
+        # ForwardBatch.init_new -- they leave global_num_token_non_padded None and
+        # set the replicated LOCAL count directly. The decode post_fill must carry
+        # that value through verbatim, not derive from the absent global (which
+        # crashed on None - rank_offset).
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_decode_registry,
+        )
+        from sglang.srt.runtime_context import get_parallel
+
+        ntnp = torch.full((1,), 99, dtype=torch.int32)  # poisoned static buffer
+        src = SimpleNamespace(
+            input_ids=torch.zeros(8, dtype=torch.int64),
+            positions=torch.zeros(8, dtype=torch.int64),
+            out_cache_loc=torch.zeros(8, dtype=torch.int64),
+            req_pool_indices=torch.zeros(4, dtype=torch.int64),
+            seq_lens=torch.full((4,), 5, dtype=torch.int64),
+            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int64),
+            mrope_positions=torch.zeros((3, 8), dtype=torch.int64),
+            num_token_non_padded=ntnp,
+            global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
+            global_num_tokens_for_logprob_gpu=torch.zeros(1, dtype=torch.int32),
+        )
+        # attn_tp_sharded_fn=True pins rank 1: were shard math applied it would
+        # clamp to 3; carrying the local count verbatim proves the bypass
+        # short-circuits before any sharding.
+        with get_parallel().override(attn_tp_size=2, attn_tp_rank=1):
+            reg = build_decode_registry(
+                device=torch.device("cpu"),
+                max_bs=4,
+                max_num_token=8,
+                seq_len_fill_value=5,
+                cache_loc_dtype=torch.int64,
+                enable_num_token_non_padded=True,
+                require_gathered_buffer=True,
+                attn_tp_sharded_fn=lambda num_tokens: True,
+                source=src,
+            )
+            fb = _MiniForwardBatch(
+                num_token_non_padded=torch.tensor([7], dtype=torch.int32),
+                global_num_token_non_padded=None,
+            )
+            reg.fill_from(
+                fb, raw_bs=4, padded_bs=4, raw_num_tokens=4, padded_num_tokens=8
+            )
+        self.assertEqual(int(src.num_token_non_padded.item()), 7)
+
     def test_register_global_num_tokens_false_carries_fb_values(self):
         # register_global_num_tokens=False (eager) excludes the computed
         # global_num_tokens_* slots so the batch's DP values are carried, not
@@ -1107,7 +1158,11 @@ class TestBuildPrefillRegistry(unittest.TestCase):
         self.assertTrue(torch.all(ids[3:8] == 0))  # padded tail reset
         self.assertTrue(torch.all(ids[8:] == 7))  # beyond the bucket: untouched
 
-    def test_num_token_non_padded_scalar_copy(self):
+    def test_num_token_non_padded_prefill_buffer_adoption(self):
+        # The prefill num_token_non_padded slot adopts the source's static
+        # buffer (shared storage), and its post_fill writes the LOCAL count
+        # derived from the invariant global host int in place — so the static
+        # buffer exposed by extract_buffer is the same storage.
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
             build_prefill_registry,
         )
@@ -1131,9 +1186,11 @@ class TestBuildPrefillRegistry(unittest.TestCase):
             input_ids=torch.tensor([1, 2, 3], dtype=torch.int64),
             positions=torch.tensor([4, 5, 6], dtype=torch.int64),
             out_cache_loc=torch.tensor([8, 9, 10], dtype=torch.int64),
-            num_token_non_padded=torch.tensor([3], dtype=torch.int32),
+            global_num_token_non_padded_cpu=3,
         )
         reg.fill_from(fb, raw_bs=1, padded_bs=1, raw_num_tokens=3, padded_num_tokens=8)
+        # Not sequence-sharded (default predicate): passthrough of the global
+        # count into the adopted static buffer.
         self.assertTrue(
             torch.equal(
                 reg.get_slot("num_token_non_padded").buffer,
@@ -1379,11 +1436,16 @@ class TestPrefillNumTokenNonPaddedPostFill(unittest.TestCase):
     rows of attn-TP rank 0 — REAL tokens — zeroing their MoE output
     in-graph. The slot's post_fill must instead recompute the local count
     against ``ctx.padded_num_tokens`` from the batch's un-adjusted global
-    count (``num_token_non_padded_cpu``), exactly like the decode registry's
+    count (``global_num_token_non_padded_cpu``), exactly like the decode registry's
     post_fill does.
+
+    Localization is gated solely on the per-forward sharding decision
+    (``attn_tp_sharded_fn``): a sharded bucket re-derives the rank-local
+    count; a replicated one passes the global count through. The cases below
+    drive that predicate directly via ``sharded``.
     """
 
-    def _fill(self, *, attn_tp_rank, attn_tp_size, require_gathered_buffer=True):
+    def _fill(self, *, attn_tp_rank, attn_tp_size, sharded=True, global_count=1018):
         from unittest import mock
 
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
@@ -1396,14 +1458,14 @@ class TestPrefillNumTokenNonPaddedPostFill(unittest.TestCase):
             max_num_token=2048,
             cache_loc_dtype=torch.int64,
             enable_num_token_non_padded=True,
-            require_gathered_buffer=require_gathered_buffer,
+            attn_tp_sharded_fn=lambda num_tokens: sharded,
         )
         # FB tensor carries the RAW-length-localized (stale) value; the CPU
         # field carries the un-adjusted global count.
         fb = _MiniForwardBatch(
             batch_size=1,
             num_token_non_padded=torch.tensor([509], dtype=torch.int32),
-            num_token_non_padded_cpu=1018,
+            global_num_token_non_padded_cpu=global_count,
         )
         with mock.patch(
             "sglang.srt.model_executor.forward_batch_info.get_parallel",
@@ -1431,11 +1493,18 @@ class TestPrefillNumTokenNonPaddedPostFill(unittest.TestCase):
         # pads. local = clamp(1018 - 512, 0, 512).
         self.assertEqual(self._fill(attn_tp_rank=1, attn_tp_size=2), 506)
 
-    def test_non_gathered_uses_raw_token_count(self):
-        # Full prefill graphs need the live raw boundary even without a
-        # gathered buffer so model layers can discard the padded bucket tail.
+    def test_not_sharded_passes_through_global_count(self):
+        # A replicated forward owns every row, so the global count is kept.
         self.assertEqual(
-            self._fill(attn_tp_rank=0, attn_tp_size=2, require_gathered_buffer=False),
+            self._fill(attn_tp_rank=0, attn_tp_size=2, sharded=False),
+            1018,
+        )
+
+    def test_absent_global_count_falls_back_to_raw_tokens(self):
+        # Full prefill graphs still need the live raw boundary when the batch
+        # carries no global count, so layers can discard the bucket tail.
+        self.assertEqual(
+            self._fill(attn_tp_rank=0, attn_tp_size=2, global_count=None),
             1018,
         )
 
@@ -1471,10 +1540,10 @@ class TestFillOncePolicy(unittest.TestCase):
 
 
 class TestComputedSlots(unittest.TestCase):
-    """num_token_non_padded (copy_from_fb + post_fill) and global_num_tokens
+    """num_token_non_padded and global_num_tokens are both computed slots
     (copy_from_fb=False + post_fill fill)."""
 
-    def test_num_token_non_padded_copy_path(self):
+    def test_num_token_non_padded_passthrough_path(self):
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
             build_decode_registry,
         )
@@ -1491,10 +1560,11 @@ class TestComputedSlots(unittest.TestCase):
         self.assertTrue(reg.has_slot("num_token_non_padded"))
         fb = _MiniForwardBatch(
             batch_size=2,
-            num_token_non_padded=torch.tensor([7], dtype=torch.int32),
+            global_num_token_non_padded=torch.tensor([7], dtype=torch.int32),
         )
-        reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=2, padded_num_tokens=2)
-        # Non-gathered: plain FB copy, post_fill is a no-op.
+        reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=7, padded_num_tokens=8)
+        # Not sequence-sharded (default predicate): post_fill passes the global
+        # scalar (7) through to the local buffer unchanged (clamped to bucket 8).
         self.assertTrue(
             torch.equal(
                 reg.get_slot("num_token_non_padded").buffer,

test/registered/unit/model_executor/test_num_token_non_padded_localization.py

@@ -0,0 +1,70 @@
+"""Per-rank-local non-padded token count across attn x MoE parallelism layouts.
+
+``compute_local_num_token_non_padded`` (GPU tensor) and
+``compute_local_num_token_non_padded_cpu`` (host int) convert a dp-group-global
+real-token count into this attention-TP rank's local count. Each rank owns a
+contiguous ``padded_bucket // attn_tp_size`` slice of the padded sequence, so the
+localizer clamps ``real - chunk * attn_tp_rank`` into ``[0, chunk]``: a replicated
+(non-sharded) rank keeps the full count and SP ranks split it. The value is
+identical whether the MoE runs TP or EP -- it is an attention-side quantity both
+backends consume. This table locks the exact per-rank counts and that the GPU
+tensor and host-int twin agree, so a change to the sharding math fails loudly.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+import unittest
+
+import torch
+
+from sglang.srt.model_executor.forward_batch_info import (
+    compute_local_num_token_non_padded,
+    compute_local_num_token_non_padded_cpu,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestNumTokenNonPaddedLayoutTable(CustomTestCase):
+    # (label, attn_tp_size, sharded, padded_bucket, real-per-dp-group,
+    #  expected [per attn-tp rank] per dp group)
+    _LAYOUTS = [
+        # Each dp rank gets a 10-token request, cuda graph pads it to 16.
+        ("TP4", 4, False, 16, [10], [[10, 10, 10, 10]]),
+        ("TP4.SP4", 4, True, 16, [10], [[4, 4, 2, 0]]),
+        ("DP4", 1, False, 16, [10, 10, 10, 10], [[10], [10], [10], [10]]),
+        ("TP2.DP2.SP2", 2, True, 16, [10, 10], [[8, 2], [8, 2]]),
+        # dp0/dp2 get 10 tokens, dp1/dp3 get 20; all padded to 32.
+        ("DP4.EP4", 1, False, 32, [10, 20, 10, 20], [[10], [20], [10], [20]]),
+        ("TP2.DP2.SP2.EP2", 2, True, 32, [10, 20], [[10, 0], [16, 4]]),
+    ]
+
+    def test_layouts_match_expected_per_rank(self):
+        for label, attn_tp, sharded, bucket, dp_reals, expected in self._LAYOUTS:
+            for dp_idx, real in enumerate(dp_reals):
+                for rank in range(attn_tp):
+                    want = expected[dp_idx][rank]
+                    with (
+                        self.subTest(layout=label, dp=dp_idx, rank=rank),
+                        get_parallel().override(
+                            attn_tp_size=attn_tp, attn_tp_rank=rank
+                        ),
+                    ):
+                        got_cpu = compute_local_num_token_non_padded_cpu(
+                            global_num_token_non_padded=real,
+                            num_tokens_per_dp=bucket,
+                            sharded=sharded,
+                        )
+                        got_gpu = compute_local_num_token_non_padded(
+                            global_num_token_non_padded=torch.tensor(real),
+                            num_tokens_per_dp=bucket,
+                            sharded=sharded,
+                        )
+                        self.assertEqual(got_cpu, want)
+                        self.assertEqual(int(got_gpu), want)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -15,6 +15,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardMode,
     PPProxyTensors,
 )
+from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import (
     capture_prefill_graph,
 )
@@ -93,6 +94,19 @@ class _FakeBatchRegistry:
 
 
 class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
+    @patch(
+        "sglang.srt.model_executor.model_runner.require_gathered_buffer",
+        return_value=True,
+    )
+    def test_default_attn_tp_sequence_sharded_uses_runtime_predicate(
+        self, mock_require_gathered_buffer
+    ):
+        runner = ModelRunner.__new__(ModelRunner)
+        runner.server_args = object()
+
+        self.assertTrue(runner.attn_tp_sequence_sharded(num_tokens=4))
+        mock_require_gathered_buffer.assert_called_once_with()
+
     def test_low_free_memory_still_captures_prefill_graph(self):
         eager_runner = object()
         prefill_runner = object()
@@ -198,6 +212,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
         runner.capture_num_tokens = [4]
         runner.buffer_registry = _FakeBatchRegistry()
+        runner.model_runner = SimpleNamespace(attn_tp_sequence_sharded=lambda _: False)
         runner.enable_cp_v2_bcg_capture = False
         runner._is_full_backend = False
         runner.backend = SimpleNamespace()