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

python/sglang/srt/layers/moe/topk.py

@@ -2152,9 +2152,10 @@ def select_experts(
                 **_fused_topk_kwargs,
             )
     else:
-        assert (
-            num_token_non_padded is None
-        ), "num_token_non_padded is not yet supported in custom_routing_function"
+        # custom_routing_function itself is padding-unaware; its output on
+        # padded rows is garbage. That is fine because _post_process_topk_ids
+        # below masks rows >= num_token_non_padded (-1 on CUDA, 0 + zeroed
+        # weights on HIP) after the logical->physical remap.
         assert not apply_routed_scaling_factor_on_output, "Not implemented"
         topk_weights, topk_ids = custom_routing_function(
             hidden_states=hidden_states,

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -789,6 +789,8 @@ def build_prefill_registry(
     embed_dtype: Optional[torch.dtype] = None,
     enable_mamba_track: bool = False,
     enable_num_token_non_padded: bool = False,
+    require_gathered_buffer: bool = False,
+    enable_prefill_cp: bool = False,
     register_input_embeds: bool = True,
     share_pool: bool = True,
     source: Optional[Any] = None,
@@ -878,12 +880,32 @@ def build_prefill_registry(
         slots.append(GraphSlot("mamba_track_mask", _bs, torch.bool, axis="bs"))
         slots.append(GraphSlot("mamba_track_seqlens", _bs, torch.int32, axis="bs"))
     if enable_num_token_non_padded:
+        from sglang.srt.model_executor.forward_batch_info import (
+            compute_local_num_token_non_padded_cpu,
+        )
+
+        def _prefill_num_token_non_padded_post_fill(buf, fb, ctx):
+            # The FB tensor was attn-TP-localized against the RAW length, but
+            # replay pads rows up to the capture bucket, moving the shard
+            # boundary — copying it verbatim would make the in-graph pad mask
+            # blank real tokens whenever raw < bucket. Recompute the local
+            # count against the padded bucket from the batch's un-adjusted
+            # global count, mirroring the decode registry's post_fill.
+            if require_gathered_buffer and not enable_prefill_cp:
+                buf.fill_(
+                    compute_local_num_token_non_padded_cpu(
+                        global_num_token_non_padded=fb.num_token_non_padded_cpu,
+                        num_tokens_per_dp=ctx.padded_num_tokens,
+                    )
+                )
+
         slots.append(
             GraphSlot(
                 "num_token_non_padded",
                 lambda _bs2, _mt: (1,),
                 torch.int32,
                 axis="none",
+                post_fill=_prefill_num_token_non_padded_post_fill,
             )
         )
 

python/sglang/srt/model_executor/forward_batch_info.py

@@ -218,6 +218,13 @@ class CaptureHiddenMode(IntEnum):
         return self.value < other.value
 
 
+def _attn_tp_local_shard_bounds(num_tokens_per_dp: int) -> Tuple[int, int]:
+    """(tokens_per_rank, rank_offset) of this attn-TP rank's contiguous shard."""
+    parallel = get_parallel()
+    tokens_per_rank = num_tokens_per_dp // parallel.attn_tp_size
+    return tokens_per_rank, tokens_per_rank * parallel.attn_tp_rank
+
+
 def compute_local_num_token_non_padded(
     global_num_token_non_padded: torch.Tensor,
     num_tokens_per_dp: int,
@@ -227,17 +234,29 @@ def compute_local_num_token_non_padded(
     Converts a global count (across all TP ranks) to a local count for this rank.
     The "global" scope is within the current DP rank; DP is handled via num_tokens_per_dp.
     """
-    attn_tp_rank = get_parallel().attn_tp_rank
-    attn_tp_size = get_parallel().attn_tp_size
-    tokens_per_rank = num_tokens_per_dp // attn_tp_size
-
+    tokens_per_rank, rank_offset = _attn_tp_local_shard_bounds(num_tokens_per_dp)
     return torch.clamp(
-        global_num_token_non_padded - tokens_per_rank * attn_tp_rank,
+        global_num_token_non_padded - rank_offset,
         0,
         tokens_per_rank,
     )
 
 
+def compute_local_num_token_non_padded_cpu(
+    global_num_token_non_padded: int,
+    num_tokens_per_dp: int,
+) -> int:
+    """Int-scalar twin of ``compute_local_num_token_non_padded``.
+
+    Replay-time hooks hold the global count as a host int
+    (``num_token_non_padded_cpu``) and write the localized result into a
+    device buffer; keeping the math on ints lets them use ``Tensor.fill_``
+    instead of staging a CPU tensor through a host-to-device copy per replay.
+    """
+    tokens_per_rank, rank_offset = _attn_tp_local_shard_bounds(num_tokens_per_dp)
+    return min(max(global_num_token_non_padded - rank_offset, 0), tokens_per_rank)
+
+
 @dataclass
 class DSV4OutCacheLoc:
     """Per-forward-pass KV cache allocation for DeepSeek-V4 on NPU.

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -47,6 +47,7 @@ import torch
 import tqdm
 
 from sglang.srt.distributed.parallel_state import graph_capture
+from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
     set_dp_buffer_len,
@@ -54,6 +55,7 @@ from sglang.srt.layers.dp_attention import (
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.pooler import EmbeddingPoolerOutput
+from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 from sglang.srt.model_executor.cuda_graph_buffer_registry import (
     CudaGraphBufferRegistry,
     build_prefill_registry,
@@ -232,6 +234,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             embed_dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
             enable_num_token_non_padded=enable_num_token_non_padded(),
+            require_gathered_buffer=require_gathered_buffer(model_runner.server_args),
+            enable_prefill_cp=(
+                is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
+            ),
             source=self.buffers,
         )