[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/attention/unittests/KNOWN_FAILURES.md

@@ -174,9 +174,9 @@ moment production is fixed; no test method invokes them today.
 | Backend | Status | Root cause |
 |---|---|---|
 | Lightning | `[no test]` (`lightning/README.md`) | Backend returns flat `[T, num_heads * head_dim]` at `lightning_backend.py:335`; `RadixAttention` piecewise writes per-head (`radix_attention.py:124-137`). Shape mismatch eager vs piecewise |
-| Mamba2 | `[no test]` (`mamba/README.md`) | `MambaMixer2.forward` projects ALL rows of `hidden_states` before per-layer `num_token_non_padded_cpu` slicing (`mamba.py:467`); trips assert under token-padding |
+| Mamba2 | `[no test]` (`mamba/README.md`) | `MambaMixer2.forward` projects ALL rows of `hidden_states` before per-layer `global_num_token_non_padded_cpu` slicing (`mamba.py:467`); trips assert under token-padding |
 | DSV4 | `[no test]` (`dsv4/README.md`) | `flash_mla.flash_mla_with_kvcache` asserts `indices.shape == (b, s_q, topk)`; metadata sized for live batch, q is static-token-padded |
-| DSA MHA_ONE_SHOT dense fallback | `[no test]` (`dsa/README.md`) | DSA passes K as concatenated `prefix + extend` to `module.attn(save_kv_cache=False)`; `unified_attention_with_output` (`radix_attention.py:170-208`) slices K to `num_token_non_padded_cpu`, dropping the prefix portion — piecewise CG diverges from eager ~50% mismatch (~0.35 max diff) |
+| DSA MHA_ONE_SHOT dense fallback | `[no test]` (`dsa/README.md`) | DSA passes K as concatenated `prefix + extend` to `module.attn(save_kv_cache=False)`; `unified_attention_with_output` (`radix_attention.py:170-208`) slices K to `global_num_token_non_padded_cpu`, dropping the prefix portion — piecewise CG diverges from eager ~50% mismatch (~0.35 max diff) |
 
 ## C.5. Sparse-kernel production bugs
 

test/registered/attention/unittests/dsa/README.md

@@ -73,7 +73,7 @@ hardware/SDK. The variant tests live in `test_dsa.py` as
   `[sum(seq_lens), num_kv_heads, head_dim]`) to `module.attn(q, k, v,
   forward_batch, save_kv_cache=False)`, but `unified_attention_with_output`
   (`radix_attention.py:170-208`, which RadixAttention routes to under
-  piecewise CG) slices K to `forward_batch.num_token_non_padded_cpu` (=
+  piecewise CG) slices K to `forward_batch.global_num_token_non_padded_cpu` (=
   live extend-token count) on the per-token K convention used by
   Triton/FlashInfer/FA. The slice removes the prefix portion, so a
   piecewise CG run diverges from the eager DSA dense fallback by ~50%

test/registered/attention/unittests/dsa/test_dsa.py

@@ -42,7 +42,7 @@ class TestDSAAttentionBackendCorrectness(CustomTestCase):
     # MHA_ONE_SHOT dense fallback passes K as concatenated prefix+extend
     # (length = sum(seq_lens)) to `module.attn`, but
     # `unified_attention_with_output` (`radix_attention.py:170-208`) slices
-    # K to `forward_batch.num_token_non_padded_cpu` (= live extend-token
+    # K to `forward_batch.global_num_token_non_padded_cpu` (= live extend-token
     # count), under the per-token K convention used by Triton/FlashInfer/
     # FA. The K-slice removes the prefix portion, so DSA's dense fallback
     # output diverges by ~50% mismatch under piecewise CG. See

test/registered/attention/unittests/mamba/README.md

@@ -83,7 +83,7 @@ call.
 - **PCG / BCG split-op extend** — `MambaMixer2.forward` asserts
   `num_actual_tokens == projected_states.shape[0]`
   (`mamba.py:467`) at the projection step, BEFORE the
-  `num_token_non_padded_cpu` slicing kicks in at the attention
+  `global_num_token_non_padded_cpu` slicing kicks in at the attention
   dispatch. The shared `_run_split_op_extend_case` pads
   `hidden_states` to a fixed `static_num_tokens` upper bound to
   exercise the per-layer slicing contract, but Mamba2 trips this

test/registered/attention/unittests/mamba/test_mamba2.py

@@ -166,9 +166,9 @@ class TestTritonMamba2BackendCorrectness(CustomTestCase):
     # exactly, with no padding tolerance. The shared split-op runner
     # pads `hidden_states` to a fixed `static_num_tokens` upper bound
     # and then relies on the backend's per-layer slicing contract via
-    # `num_token_non_padded_cpu`. Mamba2 doesn't support this padding
+    # `global_num_token_non_padded_cpu`. Mamba2 doesn't support this padding
     # because its mixer projects BEFORE the attention dispatch sees
-    # `num_token_non_padded_cpu`. Landing this needs either a Mamba2
+    # `global_num_token_non_padded_cpu`. Landing this needs either a Mamba2
     # mixer change to accept padded `hidden_states`, or a split-op
     # runner variant that passes unpadded `hidden_states` while still
     # padding the `forward_batch.input_ids` / `out_cache_loc`.

test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py

@@ -222,7 +222,7 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         num_tokens = 3
         padded_num_tokens = 5
         forward_batch = SimpleNamespace(
-            num_token_non_padded_cpu=num_tokens,
+            global_num_token_non_padded_cpu=num_tokens,
             out_cache_loc=torch.arange(padded_num_tokens, device=self.device),
             positions=torch.arange(padded_num_tokens, device=self.device),
         )

test/registered/kv_canary/test_self_unit_perturb.py

@@ -206,7 +206,7 @@ class TestRealKvPostForwardPerturb(CustomTestCase):
         forward_batch.out_cache_loc = torch.tensor(
             [2], dtype=torch.int32, device=device
         )
-        forward_batch.num_token_non_padded_cpu = 1
+        forward_batch.global_num_token_non_padded_cpu = 1
 
         head_snapshot = group.k_head.clone()
         v_head_snapshot = group.v_head.clone()
@@ -288,7 +288,7 @@ class TestReqToTokenPerturb(CustomTestCase):
         forward_batch.out_cache_loc = torch.tensor(
             [7, 0, 0], dtype=torch.int32, device=device
         )
-        forward_batch.num_token_non_padded_cpu = 1
+        forward_batch.global_num_token_non_padded_cpu = 1
 
         targets = collect_active_slots(
             maybe_inaccurate_forward_batch=forward_batch,

test/registered/kv_canary/test_self_unit_runner_per_forward.py

@@ -91,7 +91,7 @@ class TestLaunchEndpointsPerForward(CanaryManagerTestCase):
         forward_batch.out_cache_loc = torch.tensor(
             [7, 0, 0], dtype=torch.int64, device=self.device
         )
-        forward_batch.num_token_non_padded_cpu = 1
+        forward_batch.global_num_token_non_padded_cpu = 1
 
         kernel_launcher_module.launch_endpoints_per_forward(
             endpoints=(endpoint,),
@@ -144,7 +144,7 @@ class TestLaunchEndpointsPerForward(CanaryManagerTestCase):
         forward_batch.out_cache_loc = torch.tensor(
             [7], dtype=torch.int32, device=self.device
         )
-        forward_batch.num_token_non_padded_cpu = 1
+        forward_batch.global_num_token_non_padded_cpu = 1
 
         kernel_launcher_module.launch_endpoints_per_forward(
             endpoints=(endpoint,),
@@ -182,7 +182,7 @@ class TestLaunchEndpointsPerForward(CanaryManagerTestCase):
         forward_batch.out_cache_loc = torch.tensor(
             [7, 0, 0], dtype=torch.int64, device=self.device
         )
-        forward_batch.num_token_non_padded_cpu = 1
+        forward_batch.global_num_token_non_padded_cpu = 1
 
         kernel_launcher_module.launch_endpoints_per_forward(
             endpoints=(endpoint,),
@@ -218,7 +218,7 @@ class TestLaunchEndpointsPerForward(CanaryManagerTestCase):
         forward_batch.out_cache_loc = torch.tensor(
             [[7, 8]], dtype=torch.int64, device=self.device
         )[:, 0]
-        forward_batch.num_token_non_padded_cpu = 1
+        forward_batch.global_num_token_non_padded_cpu = 1
 
         kernel_launcher_module.launch_endpoints_per_forward(
             endpoints=(endpoint,),

test/registered/spec/dspark/test_dspark_dp_tier.py

@@ -81,9 +81,11 @@ class TestDraftDpSyncMetadata(CustomTestCase):
             [1, 3, 0, 2],
         )
         self.assertEqual(forward_batch.global_num_tokens_cpu, [6, 18, 0, 12])
-        self.assertEqual(forward_batch.num_token_non_padded.item(), 6)
-        self.assertEqual(forward_batch.num_token_non_padded.dtype, torch.int32)
-        self.assertEqual(forward_batch.num_token_non_padded_cpu, 6)
+        # Metadata fill sets only the invariant GLOBAL count; the LOCAL
+        # num_token_non_padded is derived later when the draft forward localizes.
+        self.assertEqual(forward_batch.global_num_token_non_padded.item(), 6)
+        self.assertEqual(forward_batch.global_num_token_non_padded.dtype, torch.int32)
+        self.assertEqual(forward_batch.global_num_token_non_padded_cpu, 6)
         self.assertTrue(forward_batch.can_run_decode_cuda_graph)
 
 

test/registered/unit/batch_overlap/test_tbo_children_dummy_token_mask.py

@@ -6,7 +6,7 @@ by setting ``num_token_non_padded(_cpu)`` to 0 so MoE top-k skips them and
 attention returns early. ``TboForwardBatchPreparer`` runs right after that, and
 it used to (a) recompute the children counts from ``len(batch.input_ids)`` --
 the *padded* token count, which restores the dummy rows as real tokens for MoE
--- and (b) hardcode each child's ``num_token_non_padded_cpu`` to ``None``, so
+-- and (b) hardcode each child's ``global_num_token_non_padded_cpu`` to ``None``, so
 the ``real_num_tokens == 0`` attention skip compared ``None == 0`` and never
 fired. Net effect: DeepEP/pplx MAX_LEN + ``--enable-two-batch-overlap``
 silently lost the whole idle-rank optimization.
@@ -29,7 +29,7 @@ from sglang.test.test_utils import CustomTestCase
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
-def _make_extend_batch(*, padded_num_tokens: int, num_token_non_padded_cpu: int):
+def _make_extend_batch(*, padded_num_tokens: int, global_num_token_non_padded_cpu: int):
     # Only the fields compute_tbo_children_num_token_non_padded reads.
     return SimpleNamespace(
         forward_mode=ForwardMode.EXTEND,
@@ -37,7 +37,7 @@ def _make_extend_batch(*, padded_num_tokens: int, num_token_non_padded_cpu: int)
         tbo_split_seq_index=1,
         extend_seq_lens_cpu=[4, 4],
         input_ids=torch.zeros(padded_num_tokens, dtype=torch.long),
-        num_token_non_padded_cpu=num_token_non_padded_cpu,
+        global_num_token_non_padded_cpu=global_num_token_non_padded_cpu,
     )
 
 
@@ -49,7 +49,7 @@ def _make_decode_capture_batch(*, num_tokens: int):
         tbo_split_seq_index=2,
         extend_seq_lens_cpu=None,
         input_ids=torch.zeros(num_tokens, dtype=torch.long),
-        num_token_non_padded_cpu=None,
+        global_num_token_non_padded_cpu=None,
     )
 
 
@@ -66,30 +66,36 @@ class TestTboChildrenDummyTokenMask(CustomTestCase):
     def test_masked_idle_parent_yields_zero_token_children(self):
         # An idle rank masked to 0 real tokens must split into two 0-token
         # children even though input_ids still holds the padded dummy rows.
-        batch = _make_extend_batch(padded_num_tokens=16, num_token_non_padded_cpu=0)
+        batch = _make_extend_batch(
+            padded_num_tokens=16, global_num_token_non_padded_cpu=0
+        )
         self.assertEqual(self._children_counts(batch), [0, 0])
 
     def test_padding_is_not_counted_as_real_tokens(self):
         # A busy rank padded up to a larger bucket must split on its real token
         # count, not on the padded input_ids length.
-        batch = _make_extend_batch(padded_num_tokens=16, num_token_non_padded_cpu=8)
+        batch = _make_extend_batch(
+            padded_num_tokens=16, global_num_token_non_padded_cpu=8
+        )
         self.assertEqual(self._children_counts(batch), [4, 4])
 
     def test_cpu_pair_matches_device_pair(self):
         # prepare() derives the children's CPU counts separately from the device
         # tensor; the two must not drift apart.
-        batch = _make_extend_batch(padded_num_tokens=16, num_token_non_padded_cpu=5)
+        batch = _make_extend_batch(
+            padded_num_tokens=16, global_num_token_non_padded_cpu=5
+        )
         cpu_pair = TboForwardBatchPreparer._split_num_token_non_padded(
             tbo_split_token_index=TboForwardBatchPreparer._compute_split_token_index(
                 batch
             ),
-            num_token_non_padded=batch.num_token_non_padded_cpu,
+            num_token_non_padded=batch.global_num_token_non_padded_cpu,
         )
         self.assertEqual(list(cpu_pair), self._children_counts(batch))
 
     def test_filter_batch_propagates_cpu_count_to_child(self):
         # Without this the attention 0-token skip (which reads
-        # num_token_non_padded_cpu) compares None == 0 and never fires.
+        # global_num_token_non_padded_cpu) compares None == 0 and never fires.
         bs = 8
         parent = ForwardBatch(
             forward_mode=ForwardMode.TARGET_VERIFY,
@@ -118,7 +124,7 @@ class TestTboChildrenDummyTokenMask(CustomTestCase):
                 out_num_token_non_padded=torch.tensor(0),
                 out_num_token_non_padded_cpu=0,
             )
-        self.assertEqual(child.num_token_non_padded_cpu, 0)
+        self.assertEqual(child.global_num_token_non_padded_cpu, 0)
 
     def test_capture_count_falls_back_to_physical_rows(self):
         batch = _make_decode_capture_batch(num_tokens=8)

test/registered/unit/layers/test_radix_attention.py

@@ -69,7 +69,7 @@ class TestRadixAttentionGraphInterface(CustomTestCase):
         real_num_tokens=2,
     ):
         forward_batch = SimpleNamespace(
-            num_token_non_padded_cpu=real_num_tokens,
+            global_num_token_non_padded_cpu=real_num_tokens,
             out_cache_loc=torch.arange(num_tokens, dtype=torch.int64),
             positions=torch.arange(num_tokens, dtype=torch.int64),
             _attn_output=None,

test/registered/unit/layers/test_radix_linear_attention.py

@@ -79,7 +79,7 @@ class TestRadixLinearAttentionPadding(CustomTestCase):
         original_out_cache_loc = torch.arange(5)
         forward_batch = SimpleNamespace(
             forward_mode=_ExtendMode(),
-            num_token_non_padded_cpu=3,
+            global_num_token_non_padded_cpu=3,
             out_cache_loc=original_out_cache_loc,
         )
 
@@ -119,7 +119,7 @@ class TestRadixLinearAttentionPadding(CustomTestCase):
         original_out_cache_loc = torch.arange(5)
         forward_batch = SimpleNamespace(
             forward_mode=_TargetVerifyMode(),
-            num_token_non_padded_cpu=3,
+            global_num_token_non_padded_cpu=3,
             out_cache_loc=original_out_cache_loc,
         )
 
@@ -158,7 +158,7 @@ class TestRadixLinearAttentionPadding(CustomTestCase):
         original_out_cache_loc = torch.arange(5)
         forward_batch = SimpleNamespace(
             forward_mode=_ExtendMode(),
-            num_token_non_padded_cpu=3,
+            global_num_token_non_padded_cpu=3,
             out_cache_loc=original_out_cache_loc,
         )
 
@@ -189,7 +189,7 @@ class TestRadixLinearAttentionPadding(CustomTestCase):
             with self.subTest(padded_num_tokens=padded_num_tokens):
                 original_out_cache_loc = torch.arange(padded_num_tokens)
                 forward_batch = SimpleNamespace(
-                    num_token_non_padded_cpu=3,
+                    global_num_token_non_padded_cpu=3,
                     out_cache_loc=original_out_cache_loc,
                 )
                 context = SimpleNamespace(

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