From 155d23cb008f793d4609d2bd5885cd9882ba38ee Mon Sep 17 00:00:00 2001 From: Yongqi Wang Date: Sun, 4 Oct 2026 18:31:15 +0800 Subject: [PATCH] [Bugfix][XPU] Make DeepSeek V4 FP8 sparse decode graph-capturable (#59159) Signed-off-by: Yongqi Wang Co-authored-by: Kunshang Ji --- .../deepseek_v4/xpu/xpu_sparse_decode_fp8.py | 125 +++++------------- 1 file changed, 36 insertions(+), 89 deletions(-) diff --git a/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py b/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py index 7c0b808de2..4fb8c4059b 100644 --- a/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py +++ b/vllm/models/deepseek_v4/xpu/xpu_sparse_decode_fp8.py @@ -166,8 +166,7 @@ def xpu_sparse_decode_fp8( num_tokens = q.shape[0] device = q.device - # Determine max topk and swa widths - if not swa_only and topk_indices is not None: + if not swa_only and topk_indices is not None and kv_cache is not None: topk_idx_2d = ( topk_indices.squeeze(1) if topk_indices.dim() == 3 else topk_indices ) @@ -176,109 +175,57 @@ def xpu_sparse_decode_fp8( topk_idx_2d = None max_topk = 0 + use_topk = topk_idx_2d is not None swa_idx_2d = swa_indices.squeeze(1) if swa_indices.dim() == 3 else swa_indices max_swa = swa_idx_2d.shape[1] - K_total = max_topk + max_swa - # Allocate flat workspace: [num_tokens * K_total, 512] bf16 + # Topk rows precede SWA rows in the flat workspace. + n_topk_rows = num_tokens * max_topk workspace = torch.empty( (num_tokens * K_total, OUTPUT_DIM), dtype=torch.bfloat16, device=device ) - ws_3d = workspace.view(num_tokens, K_total, OUTPUT_DIM) - # Dequant+gather topk slots from compressed cache - if not swa_only and topk_idx_2d is not None and kv_cache is not None: - topk_flat = topk_idx_2d.reshape(-1).to(torch.int32) - topk_buf = torch.empty( - (num_tokens * max_topk, OUTPUT_DIM), dtype=torch.bfloat16, device=device + if topk_idx_2d is not None and kv_cache is not None and n_topk_rows > 0: + dequant_gather_slots( + workspace[:n_topk_rows], + kv_cache, + topk_idx_2d.reshape(-1).to(torch.int32), + kv_cache.shape[1], ) - compressed_block_size = kv_cache.shape[1] - dequant_gather_slots(topk_buf, kv_cache, topk_flat, compressed_block_size) - ws_3d[:, :max_topk, :] = topk_buf.view(num_tokens, max_topk, OUTPUT_DIM) - # Dequant+gather SWA slots - swa_flat = swa_idx_2d.reshape(-1).to(torch.int32) - swa_buf = torch.empty( - (num_tokens * max_swa, OUTPUT_DIM), dtype=torch.bfloat16, device=device + dequant_gather_slots( + workspace[n_topk_rows:], + swa_kv_cache, + swa_idx_2d.reshape(-1).to(torch.int32), + swa_kv_cache.shape[1], ) - swa_block_size = swa_kv_cache.shape[1] - dequant_gather_slots(swa_buf, swa_kv_cache, swa_flat, swa_block_size) - ws_3d[:, max_topk:, :] = swa_buf.view(num_tokens, max_swa, OUTPUT_DIM) - - # Build combined indices into the flat workspace and combined lengths. - # Workspace layout per token t: [topk_0..topk_{max_topk-1}, swa_0..swa_{max_swa-1}] - # Flat index for token t, position p = t * K_total + p - # - # IMPORTANT: The attention kernel uses combined_lens as a position cutoff — - # it only reads indices[0..combined_lens-1]. So indices must be PACKED - # contiguously: [valid_topk_indices..., valid_swa_indices..., -1 padding...] - if not swa_only and topk_lens is not None: - combined_lens = (topk_lens + swa_lens).to(torch.int32) - else: - combined_lens = swa_lens.to(torch.int32) - - max_combined = int(combined_lens.max().item()) if combined_lens.numel() > 0 else 0 - # Round up to BLOCK_N=16 alignment for kernel efficiency + # Use a shape-derived width so replay accepts changing lens without host sync. _BLOCK_N = 16 - max_combined_padded = ((max_combined + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N + width = max(_BLOCK_N, ((K_total + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N) - # Build packed index table: [num_tokens, max_combined_padded] - # Each token t: [topk_0..topk_{tlen-1}, swa_0..swa_{slen-1}, -1 padding] - # Vectorized: for each token, topk indices are t*K_total + 0..tlen-1, - # swa indices are t*K_total + max_topk + 0..slen-1 - combined_indices = torch.full( - (num_tokens, max_combined_padded), - fill_value=-1, - dtype=torch.int32, - device=device, + pos = torch.arange(width, device=device, dtype=torch.int32).unsqueeze(0) + tok = torch.arange(num_tokens, device=device, dtype=torch.int32).unsqueeze(1) + + slen = swa_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_swa) + if use_topk and topk_lens is not None: + tlen = topk_lens.reshape(-1, 1).to(torch.int32).clamp(min=0, max=max_topk) + else: + tlen = torch.zeros_like(slen) + + in_topk = pos < tlen + in_swa = (pos >= tlen) & (pos < tlen + slen) + + topk_ws = tok * max_topk + pos + swa_ws = n_topk_rows + tok * max_swa + (pos - tlen) + + combined_indices = torch.where( + in_topk, + topk_ws, + torch.where(in_swa, swa_ws, torch.full_like(topk_ws, -1)), ) - token_offsets = ( - torch.arange(num_tokens, device=device, dtype=torch.int32) * K_total - ) # [B] - - if not swa_only and topk_lens is not None: - # Pack topk: for each token, write t*K_total + 0..tlen-1 at positions 0..tlen-1 - max_tlen = int(topk_lens.max().item()) - topk_range = torch.arange(max_tlen, device=device, dtype=torch.int32).unsqueeze( - 0 - ) - topk_valid = topk_range < topk_lens.unsqueeze(1) - topk_ws_indices = token_offsets.unsqueeze(1) + topk_range - combined_indices[:, :max_tlen] = torch.where( - topk_valid, - topk_ws_indices, - torch.tensor(-1, dtype=torch.int32, device=device), - ) - # Pack swa after topk: positions tlen..tlen+slen-1 - # Since tlen varies per token, we need per-token offset - swa_range = torch.arange(max_swa, device=device, dtype=torch.int32).unsqueeze(0) - swa_valid = swa_range < swa_lens.unsqueeze(1) - swa_ws_indices = token_offsets.unsqueeze(1) + max_topk + swa_range - # Write at position topk_lens[t] + swa_pos for each token - for t_idx in range(num_tokens): - tlen = int(topk_lens[t_idx].item()) - slen = int(swa_lens[t_idx].item()) - combined_indices[t_idx, tlen : tlen + slen] = swa_ws_indices[t_idx, :slen] - else: - # SWA-only: pack swa indices at positions 0..slen-1 - # Use min(max_swa, max_combined_padded) because combined_indices only - # has max_combined_padded columns, and all valid entries fit within it. - effective_swa = min(max_swa, max_combined_padded) - swa_range = torch.arange( - effective_swa, device=device, dtype=torch.int32 - ).unsqueeze(0) - swa_valid = swa_range < swa_lens.unsqueeze(1) - swa_ws_indices = token_offsets.unsqueeze(1) + swa_range # max_topk=0 - combined_indices[:, :effective_swa] = torch.where( - swa_valid, - swa_ws_indices, - torch.tensor(-1, dtype=torch.int32, device=device), - ) - - # Call BF16 sparse MLA kernel out_attn, _, _ = triton_bf16_mla_sparse_interface( q=q, kv=workspace.unsqueeze(1),