Fix pad-row top-k masking with custom_routing_function under DP attention (#31838)

test/registered/moe/test_topk_padded_region.py

@@ -190,5 +190,63 @@ class TestPostProcessPaddedMaskingHip(CustomTestCase):
             topk_mod._skip_hip_pad_mask = orig
 
 
+@unittest.skipUnless(torch.cuda.is_available(), "padded-region masking needs a GPU")
+@unittest.skipIf(_IS_HIP, "HIP keeps the pre-mask routing contract (AITER/MORI)")
+class TestSelectExpertsCustomRoutingPadMask(CustomTestCase):
+    """``select_experts`` must accept ``num_token_non_padded`` together with a
+    ``custom_routing_function`` and mask the padded region to -1.
+
+    Bug regression (EP MoE dispatch overflow): DP-attention/SP pad rows
+    carry garbage router input; a custom routing function can emit the same
+    expert id in every top-k slot for them (the masked argmax degenerates on
+    non-finite scores), which overflows an EP dispatch pool's
+    min(top_k, experts_per_rank) distinct-ids sizing bound. The fix routes
+    padded-region masking through the shared post-process; previously this
+    combination was rejected with ``assert num_token_non_padded is None``,
+    so no model with a custom router could mask its pad rows at all.
+    """
+
+    DEVICE = "cuda"
+
+    def test_padded_tail_masked_after_custom_routing(self):
+        from sglang.srt.layers.moe.topk import select_experts
+
+        num_tokens, num_experts, top_k, n_valid = 12, 32, 8, 10
+        hidden = torch.randn((num_tokens, 64), device=self.DEVICE, dtype=torch.bfloat16)
+        router_logits = torch.randn(
+            (num_tokens, num_experts), device=self.DEVICE, dtype=torch.float32
+        )
+        # Pad rows carry non-finite router input, like real DP-pad rows.
+        router_logits[n_valid:] = float("nan")
+
+        def _degenerate_router(hidden_states, gating_output, topk, renormalize):
+            # Mimic the incident: NaN rows collapse to one expert id x top_k.
+            weights = torch.softmax(gating_output.nan_to_num(0.0), dim=-1).topk(
+                topk, dim=-1
+            )
+            ids = weights.indices.to(torch.int32)
+            ids[gating_output.isnan().any(dim=-1)] = 7
+            return weights.values.float(), ids
+
+        out = select_experts(
+            hidden_states=hidden,
+            router_logits=router_logits,
+            topk_config=TopKConfig(
+                top_k=top_k,
+                renormalize=True,
+                custom_routing_function=_degenerate_router,
+            ),
+            layer_id=0,
+            num_token_non_padded=torch.tensor(
+                n_valid, device=self.DEVICE, dtype=torch.int32
+            ),
+        )
+        # Padded tail fully -1 (skipped by every EP dispatch path)...
+        self.assertTrue(torch.all(out.topk_ids[n_valid:] == -1))
+        # ...and real rows untouched (in-range, no -1 leakage).
+        self.assertTrue(torch.all(out.topk_ids[:n_valid] >= 0))
+        self.assertTrue(torch.all(out.topk_ids[:n_valid] < num_experts))
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -45,6 +45,7 @@ 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_tokens_gpu: Optional[torch.Tensor] = None
     global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
     ngram_embedding_info: Optional[object] = None
@@ -1248,6 +1249,80 @@ class TestBuildPrefillRegistry(unittest.TestCase):
         self.assertIs(fb_view.input_embeds, embeds)
 
 
+class TestPrefillNumTokenNonPaddedPostFill(unittest.TestCase):
+    """The prefill registry must re-derive the attn-TP-local pad boundary from
+    the CAPTURE BUCKET, not trust the FB tensor.
+
+    Bug regression: breakable-graph replay pads ``raw`` tokens up to the
+    capture bucket, moving the attn-TP shard boundary to ``bucket/attn_tp``
+    rows — but the FB ``num_token_non_padded`` tensor was localized against
+    the RAW length on the eager prep path. Copying it verbatim made every
+    ``raw < bucket`` replay mask the last ``(bucket - raw)/attn_tp`` shard
+    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
+    post_fill does.
+    """
+
+    def _fill(self, *, attn_tp_rank, attn_tp_size, require_gathered_buffer=True):
+        from unittest import mock
+
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_prefill_registry,
+        )
+
+        reg = build_prefill_registry(
+            device=torch.device("cpu"),
+            max_bs=4,
+            max_num_token=2048,
+            cache_loc_dtype=torch.int64,
+            enable_num_token_non_padded=True,
+            require_gathered_buffer=require_gathered_buffer,
+        )
+        # 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,
+        )
+        with mock.patch(
+            "sglang.srt.model_executor.forward_batch_info.get_parallel",
+            return_value=SimpleNamespace(
+                attn_tp_rank=attn_tp_rank, attn_tp_size=attn_tp_size
+            ),
+        ):
+            reg.fill_from(
+                fb,
+                raw_bs=1,
+                padded_bs=1,
+                raw_num_tokens=1018,
+                padded_num_tokens=1024,
+            )
+        return int(reg.get_slot("num_token_non_padded").buffer.item())
+
+    def test_rank0_uses_bucket_shard_not_raw_localized_value(self):
+        # bucket 1024 / attn_tp 2 -> 512-row shards. Rank 0's shard is fully
+        # real (global rows [0, 512)); the raw-localized FB value (509) would
+        # mask 3 real rows.
+        self.assertEqual(self._fill(attn_tp_rank=0, attn_tp_size=2), 512)
+
+    def test_rank1_masks_exactly_the_true_pads(self):
+        # Rank 1's shard holds global rows [512, 1024): 506 real + 6 bucket
+        # pads. local = clamp(1018 - 512, 0, 512).
+        self.assertEqual(self._fill(attn_tp_rank=1, attn_tp_size=2), 506)
+
+    def test_non_gathered_keeps_plain_fb_copy(self):
+        # Without a gathered buffer there is no attn-TP scatter; the plain FB
+        # copy must be preserved (post_fill no-op), mirroring the decode
+        # registry's contract.
+        self.assertEqual(
+            self._fill(attn_tp_rank=0, attn_tp_size=2, require_gathered_buffer=False),
+            509,
+        )
+
+
 class TestFillOncePolicy(unittest.TestCase):
     """FILL_ONCE initializes the whole buffer at alloc and never resets the
     padded tail per iter (unlike FILL_SENTINEL)."""