diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 9387ac88d..5b46fb55b 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -382,6 +382,9 @@ class MQALayer(nn.Module): ) self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32)) + self._attn_sink_local: Optional[torch.Tensor] = ( + self.attn_sink if attn_tp_size == 1 else None + ) self.fuse_wqa_wkv = envs.SGLANG_OPT_FUSE_WQA_WKV.get() if self.fuse_wqa_wkv: self.wqkv_a = ReplicatedLinear( @@ -890,10 +893,23 @@ class MQALayer(nn.Module): tp_slice, q_padded, q_out = slice(None), None, None if self.tp_size > 1: - q_padded = x.new_empty(x.shape[0], self.n_heads, self.head_dim) - rank = self.tp_rank - tp_slice = slice(rank * self.n_local_heads, (rank + 1) * self.n_local_heads) + # FlashMLA's fp8 sparse decode kernel only specializes h_q for {64, 128}. + # Pad the per-rank heads to 64 (not the full n_heads) when they fit, to + # dispatch the cheaper decode::head64 variant; attn_sink is sliced to + # this rank and padded to match. + padded_num_heads = 64 if self.n_local_heads <= 64 else self.n_heads + q_padded = x.new_empty(x.shape[0], padded_num_heads, self.head_dim) + tp_slice = slice(0, self.n_local_heads) q_out = q_padded[:, tp_slice, :] + if self._attn_sink_local is None: + # Build once on the first forward (post weight load); a per-call + # rebuild would replay a fill+copy per layer in the decode graph. + rank = self.tp_rank + sink = self.attn_sink.new_zeros(padded_num_heads) + sink[: self.n_local_heads] = self.attn_sink[ + rank * self.n_local_heads : (rank + 1) * self.n_local_heads + ] + self._attn_sink_local = sink if enable_multi_stream: # Multi-stream path always fuses cache write into the K kernel, @@ -960,7 +976,7 @@ class MQALayer(nn.Module): o, self.attn_mqa.layer_id, self.compress_ratio, - self.attn_sink, + self._attn_sink_local, save_kv_cache, ) else: @@ -971,7 +987,7 @@ class MQALayer(nn.Module): layer=self.attn_mqa, forward_batch=forward_batch, compress_ratio=self.compress_ratio, - attn_sink=self.attn_sink, + attn_sink=self._attn_sink_local, save_kv_cache=save_kv_cache, ) o = o[:, tp_slice, :]