[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"),

python/sglang/srt/models/deepseek_v2.py

@@ -2387,6 +2387,13 @@ class DeepseekV2Model(nn.Module):
     def get_input_embeddings(self) -> torch.Tensor:
         return self.embed_tokens
 
+    def _dsa_forward_uses_topk(self) -> bool:
+        if not self.use_dsa:
+            return False
+        backend = get_attn_backend()
+        backend = getattr(backend, "primary", backend)
+        return not getattr(backend, "use_mha", False)
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -2396,6 +2403,7 @@ class DeepseekV2Model(nn.Module):
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> Union[torch.Tensor, PPProxyTensors]:
         total_num_layers = self.end_layer - self.start_layer
+        dsa_forward_uses_topk = self._dsa_forward_uses_topk()
         if self.pp_group.is_first_rank:
             if input_embeds is None:
                 hidden_states = self.embed_tokens(input_ids)
@@ -2411,6 +2419,7 @@ class DeepseekV2Model(nn.Module):
                 not forward_batch.forward_mode.is_idle()
                 and hidden_states.shape[0] != 0
                 and self.use_dsa
+                and dsa_forward_uses_topk
                 and dsa_layer_skips_topk(self.config, self.start_layer)
                 and topk_indices is None
             ), (
@@ -2515,6 +2524,7 @@ class DeepseekV2Model(nn.Module):
             }
             if (
                 self.use_dsa
+                and dsa_forward_uses_topk
                 and self.end_layer < self.config.num_hidden_layers
                 and dsa_layer_skips_topk(self.config, self.end_layer)
             ):