[Bugfix][XPU] Apply attn_sink in the DeepSeek-V4 sparse MLA Triton kernel

tests/kernels/attention/test_xpu_mla_sparse.py

@@ -168,3 +168,77 @@ def test_bf16_triton_sparse_mla_masked_chunks(device_str, dtype):
     # lse/max_logits are large-negative finite rather than the reference's
     # +inf/-inf placeholders, so only the output is compared here.
     assert torch.allclose(out[2], torch.zeros_like(out[2]))
+
+
+@pytest.mark.parametrize("device_str", ["xpu"])
+@pytest.mark.parametrize("num_valid", [1, 4, 16, 128, 512])
+@pytest.mark.skipif(
+    not torch.xpu.is_available(),
+    reason="XPU is required",
+)
+def test_bf16_triton_sparse_mla_attn_sink(device_str, num_valid):
+    """DSv4 attn_sink (a per-head extra logit in the softmax denominator).
+
+    Uses the DSv4 prefill/decode shape (head 512, no rope split,
+    ``block_dpe=0``) and sinks drawn from the range of the released
+    checkpoints. Rows with few valid keys are where ignoring the sink hurts
+    most (start of a sequence / SWA window), so they are covered explicitly.
+    """
+    device = torch.device(device_str)
+    s_q, s_kv, h_q, h_kv, d_qk, d_v, topk = 8, 1024, 64, 1, 512, 512, 512
+    dtype = torch.bfloat16
+    torch.random.manual_seed(1234)
+
+    q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device)
+    kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device)
+    indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device)
+    for t in range(s_q):
+        indices[t, 0, :num_valid] = torch.randperm(s_kv, device=device)[:num_valid].to(
+            torch.int32
+        )
+    # attn_sink values of the released DSv4-Flash checkpoints span ~[-2.2, 2.5].
+    attn_sink = torch.rand(h_q, device=device) * 4.7 - 2.2
+    sm_scale = d_qk**-0.5
+
+    out, _, lse = triton_bf16_mla_sparse_interface(
+        q, kv, indices, sm_scale, d_v, block_dpe=0, attn_sink=attn_sink
+    )
+    ref_out, _, _, ref_lse = reference_mla_sparse_prefill(
+        q, kv, indices, sm_scale, d_v, attn_sink=attn_sink
+    )
+    assert out.isfinite().all()
+    assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2)
+    # lse is reported over the selected keys only (sink excluded).
+    assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3)
+
+    # Without a sink the result must differ markedly for short rows -- this
+    # guards against the sink silently becoming a no-op again.
+    out_nosink, _, _ = triton_bf16_mla_sparse_interface(
+        q, kv, indices, sm_scale, d_v, block_dpe=0
+    )
+    if num_valid <= 16:
+        assert not torch.allclose(out_nosink, ref_out, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not torch.xpu.is_available(),
+    reason="XPU is required",
+)
+def test_bf16_triton_sparse_mla_attn_sink_neg_inf_is_noop():
+    """The default sink value (-inf) must not change the output."""
+    device = torch.device("xpu")
+    s_q, s_kv, h_q, d_qk, d_v, topk = 4, 256, 16, 512, 512, 128
+    torch.random.manual_seed(7)
+    q = torch.randn((s_q, h_q, d_qk), dtype=torch.bfloat16, device=device)
+    kv = torch.randn((s_kv, 1, d_qk), dtype=torch.bfloat16, device=device)
+    indices = torch.stack(
+        [torch.randperm(s_kv, device=device)[:topk] for _ in range(s_q)]
+    ).to(torch.int32)[:, None, :]
+    sink = torch.full((h_q,), float("-inf"), device=device)
+    a, _, _ = triton_bf16_mla_sparse_interface(
+        q, kv, indices, d_qk**-0.5, d_v, block_dpe=0
+    )
+    b, _, _ = triton_bf16_mla_sparse_interface(
+        q, kv, indices, d_qk**-0.5, d_v, block_dpe=0, attn_sink=sink
+    )
+    assert torch.equal(a, b)

vllm/models/deepseek_v4/xpu/xpu_sparse.py

@@ -364,5 +364,6 @@ def _forward_prefill(
                 sm_scale=self.scale,
                 d_v=q.shape[-1],
                 block_dpe=0,
+                attn_sink=self.attn_sink[: q.shape[1]],
             )
             output[query_start:query_end] = out

vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py

@@ -233,5 +233,6 @@ def xpu_sparse_decode_fp8(
         sm_scale=softmax_scale,
         d_v=q.shape[-1],
         block_dpe=0,
+        attn_sink=attn_sink[: q.shape[1]],
     )
     out.copy_(out_attn)

vllm/v1/attention/ops/xpu_mla_sparse.py

@@ -15,6 +15,7 @@ def _bf16_mla_sparse_kernel(
     out_ptr,
     softmax_lse_ptr,
     max_logits_ptr,
+    attn_sink_ptr,
     seq_q,
     seq_kv,
     h_q,
@@ -41,6 +42,7 @@ def _bf16_mla_sparse_kernel(
     BLOCK_DMODEL: tl.constexpr,  # block size for dim_nope
     BLOCK_DPE: tl.constexpr,  # block size for positional embedding
     LOGE2: tl.constexpr,
+    HAS_ATTN_SINK: tl.constexpr,
 ):
     cur_q = tl.program_id(0)
     cur_head_id = tl.program_id(1)
@@ -151,13 +153,24 @@ def _bf16_mla_sparse_kernel(
         e_sum = e_sum * re_scale + tl.sum(p, 1)
         e_max = n_e_max
 
-    # rescaling
-    acc /= e_sum[:, None]
-
     max_logits = e_max * LOGE2
-    # calculate lse
+    # calculate lse (over the selected keys only, sink excluded)
     lse = max_logits + tl.log2(e_sum) * LOGE2
 
+    if HAS_ATTN_SINK:
+        # The sink is an extra logit (natural-log domain, same units as the
+        # scaled QK scores) that joins the softmax denominator but has no value.
+        sink = tl.load(attn_sink_ptr + cur_head, mask=mask_h, other=float("-inf"))
+        # e_max / e_sum live in the log2 domain: log2(e) = 1.4426950408889634
+        sink2 = sink * 1.4426950408889634
+        m_final = tl.maximum(e_max, sink2)
+        sink_scale = tl.exp2(e_max - m_final)
+        acc *= sink_scale[:, None]
+        e_sum = e_sum * sink_scale + tl.exp2(sink2 - m_final)
+
+    # rescaling
+    acc /= e_sum[:, None]
+
     # write output
     offs_o = (
         cur_q * stride_out_token
@@ -184,6 +197,7 @@ def triton_bf16_mla_sparse_interface(
     sm_scale: float,
     d_v: int = 512,
     block_dpe: int = 64,
+    attn_sink: torch.Tensor | None = None,  # [num_heads_q] fp32
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Out : [num_tokens, num_heads_q, d_v]
     max_logits : [num_tokens, num_heads_q]
@@ -198,6 +212,8 @@ def triton_bf16_mla_sparse_interface(
         block_dpe: Size of positional embedding portion of dim_qk.
             Set to 0 when q/kv contain only the nope latent (e.g. DSv4
             prefill where RoPE is not split out).
+        attn_sink: Optional per-head sink logit added to the softmax
+            denominator (DSv4). ``-inf`` entries are a no-op.
 
     """
     num_tokens, num_heads_q, dim_qk = q.shape
@@ -237,6 +253,9 @@ def triton_bf16_mla_sparse_interface(
         (num_tokens, num_heads_q), dtype=torch.float32, device=q.device
     )
 
+    if attn_sink is not None:
+        assert attn_sink.dtype == torch.float32 and attn_sink.shape[0] >= num_heads_q
+
     k = kv
     v = kv[..., :d_v]
 
@@ -248,6 +267,7 @@ def triton_bf16_mla_sparse_interface(
         out_ptr=out,
         softmax_lse_ptr=softmax_lse,
         max_logits_ptr=max_logits,
+        attn_sink_ptr=attn_sink,
         seq_q=num_tokens,
         seq_kv=kv.shape[0],
         h_q=num_heads_q,
@@ -274,6 +294,7 @@ def triton_bf16_mla_sparse_interface(
         BLOCK_DMODEL=BLOCK_DMODEL,
         BLOCK_DPE=BLOCK_DPE,
         LOGE2=LOGE2,
+        HAS_ATTN_SINK=attn_sink is not None,
     )
 
     return out, max_logits, softmax_lse