[PP] Pass DSA topk through PP warmup proxy buffers (#28785)

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

@@ -78,6 +78,7 @@ def _allocate_decode_buffers(
     enable_mamba_track: bool,
     ne_token_table: Optional[torch.Tensor] = None,
     hc_hidden_size: Optional[int] = None,
+    pp_proxy_topk_size: Optional[int] = None,
 ) -> SimpleNamespace:
     """Allocate the FB-shared decode buffers."""
     with torch.device(device):
@@ -115,6 +116,10 @@ def _allocate_decode_buffers(
                 pp_proxy_tensors["residual"] = torch.zeros(
                     (max_bs, hidden_size), dtype=dtype
                 )
+            if pp_proxy_topk_size is not None:
+                pp_proxy_tensors["topk_indices"] = torch.zeros(
+                    (max_num_token, pp_proxy_topk_size), dtype=torch.int32
+                )
         else:
             pp_proxy_tensors = None
 
@@ -429,6 +434,7 @@ class BaseRunner(ABC):
             cache_loc_dtype=torch.int64,
             enable_mamba_track=False,
             hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
+            pp_proxy_topk_size=mr.get_pp_proxy_topk_size(),
         )
 
     def _dummy_run(

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

@@ -190,6 +190,11 @@ class EagerRunner(BaseRunner):
                 pp_proxy_tensors["residual"] = torch.zeros(
                     (rows, hidden_size), dtype=mr.dtype, device=mr.device
                 )
+            pp_proxy_topk_size = mr.get_pp_proxy_topk_size()
+            if pp_proxy_topk_size is not None:
+                pp_proxy_tensors["topk_indices"] = torch.zeros(
+                    (rows, pp_proxy_topk_size), dtype=torch.int32, device=mr.device
+                )
 
         adapter = SimpleNamespace(
             input_ids=_slot("input_ids"),