@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 290 lines 11.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """XPU sparse decode for DeepSeek V4 with FP8 KV cache.
4
5 Strategy: dequantize FP8 UE8M0 KV cache pages to BF16 on the fly,
6 then reuse the BF16 sparse MLA attention kernel (xpu_sparse_mla_bf16).
7 This keeps the external KV cache layout identical to CUDA/ROCm.
8 """
9
10 import torch
11
12 from vllm.triton_utils import tl, triton
13 from vllm.v1.attention.ops.xpu_mla_sparse import (
14 triton_bf16_mla_sparse_interface,
15 )
16
17 # FP8 DS MLA cache layout constants
18 TOKEN_FP8_DIM = 448 # NoPE portion in FP8
19 TOKEN_BF16_DIM = 64 # RoPE portion in BF16
20 TOKEN_SCALE_DIM = 8 # UE8M0 scales per token
21 QUANT_BLOCK_SIZE = 64 # Elements per quant block
22 OUTPUT_DIM = 512 # = TOKEN_FP8_DIM + TOKEN_BF16_DIM after dequant
23 TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 # 576 bytes per token
24
25
26 @triton.jit
27 def _dequant_gather_slots_kernel(
28 # Output workspace: [total_slots, OUTPUT_DIM] bf16
29 out_ptr,
30 # FP8 paged cache base pointer (uint8 flat)
31 cache_ptr,
32 # Global slot indices: [total_slots] int32
33 indices_ptr,
34 # Cache geometry
35 cache_block_size: tl.constexpr,
36 token_data_size: tl.constexpr, # 576
37 block_stride: tl.int64, # k_cache.stride(0) — total uint8 per block
38 fp8_dim: tl.constexpr, # 448
39 bf16_dim: tl.constexpr, # 64
40 scale_dim: tl.constexpr, # 8
41 quant_block: tl.constexpr, # 64
42 output_dim: tl.constexpr, # 512
43 n_quant_blocks: tl.constexpr, # 7
44 ):
45 """Dequantize scattered FP8 slots into a flat BF16 workspace.
46
47 Grid: [total_slots] — one program per slot to gather.
48
49 Cache block layout (block_size tokens):
50 [0, block_size*576): Token data, each token 448 FP8 + 128 BF16
51 [block_size*576, block_size*576 + block_size*8): Scales
52 """
53 pid = tl.program_id(0)
54
55 # Load global slot index
56 slot_idx = tl.load(indices_ptr + pid).to(tl.int64)
57
58 # Output pointer for this slot
59 out_row_ptr = out_ptr + pid * output_dim
60
61 # Handle invalid slots (index < 0): write zeros
62 if slot_idx < 0:
63 zero = tl.zeros([quant_block], dtype=tl.bfloat16)
64 for i in tl.static_range(0, 512, 64):
65 offsets = i + tl.arange(0, quant_block)
66 mask = offsets < output_dim
67 tl.store(out_row_ptr + offsets, zero, mask=mask)
68 return
69
70 # Compute block and position within block
71 block_idx = slot_idx // cache_block_size
72 pos_in_block = slot_idx % cache_block_size
73
74 # Block base pointer
75 block_base = cache_ptr + block_idx * block_stride
76
77 # Token data: at offset pos_in_block * token_data_size within block
78 token_data_ptr = block_base + pos_in_block * token_data_size
79
80 # Scale: after all token data, at offset
81 # cache_block_size * token_data_size + pos_in_block * scale_dim
82 scale_region_offset = tl.cast(cache_block_size, tl.int64) * token_data_size
83 token_scale_ptr = block_base + scale_region_offset + pos_in_block * scale_dim
84
85 # ========== Dequantize FP8 portion (448 elements) ==========
86 for qblock_idx in tl.static_range(n_quant_blocks):
87 qblock_start = qblock_idx * quant_block
88 offsets = qblock_start + tl.arange(0, quant_block)
89 mask = offsets < fp8_dim
90
91 # Load FP8 as uint8 and bitcast
92 x_uint8 = tl.load(token_data_ptr + offsets, mask=mask, other=0)
93 x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True)
94 x_float = x_fp8.to(tl.float32)
95
96 # Load UE8M0 scale: scale = 2^(stored_value - 127)
97 encoded_scale = tl.load(token_scale_ptr + qblock_idx)
98 exponent = encoded_scale.to(tl.float32) - 127.0
99 scale = tl.exp2(exponent)
100
101 # Dequantize and store as bf16
102 x_dequant = x_float * scale
103 tl.store(out_row_ptr + offsets, x_dequant.to(tl.bfloat16), mask=mask)
104
105 # ========== Copy BF16 portion (64 elements) directly ==========
106 bf16_src_ptr = (token_data_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))
107 bf16_out_ptr = (out_row_ptr + fp8_dim).to(tl.pointer_type(tl.bfloat16))
108
109 for j in tl.static_range(bf16_dim // 16):
110 chunk_offsets = j * 16 + tl.arange(0, 16)
111 bf16_vals = tl.load(bf16_src_ptr + chunk_offsets)
112 tl.store(bf16_out_ptr + chunk_offsets, bf16_vals)
113
114
115 def dequant_gather_slots(
116 out: torch.Tensor, # [total_slots, 512] bf16, pre-allocated
117 cache: torch.Tensor, # [num_blocks, block_size, head_bytes] uint8
118 indices: torch.Tensor, # [total_slots] int32, global slot IDs
119 cache_block_size: int, # block_size for this cache
120 ) -> None:
121 """Dequantize FP8 UE8M0 pages at scattered slot indices into bf16."""
122 total_slots = indices.shape[0]
123 if total_slots == 0:
124 return
125
126 block_stride = cache.stride(0)
127
128 _dequant_gather_slots_kernel[(total_slots,)](
129 out,
130 cache,
131 indices,
132 cache_block_size=cache_block_size,
133 token_data_size=TOKEN_DATA_SIZE,
134 block_stride=block_stride,
135 fp8_dim=TOKEN_FP8_DIM,
136 bf16_dim=TOKEN_BF16_DIM,
137 scale_dim=TOKEN_SCALE_DIM,
138 quant_block=QUANT_BLOCK_SIZE,
139 output_dim=OUTPUT_DIM,
140 n_quant_blocks=7,
141 )
142
143
144 def xpu_sparse_decode_fp8(
145 q: torch.Tensor, # [num_tokens, num_heads, head_dim]
146 kv_cache: torch.Tensor | None, # [num_blocks, block_size, head_bytes] uint8
147 swa_kv_cache: torch.Tensor, # [num_blocks, swa_block_size, head_bytes] uint8
148 swa_only: bool,
149 topk_indices: torch.Tensor | None, # [num_tokens, 1, topk] global slot IDs
150 topk_lens: torch.Tensor | None,
151 swa_indices: torch.Tensor, # [num_tokens, 1, swa_k] global slot IDs
152 swa_lens: torch.Tensor,
153 attn_sink: torch.Tensor,
154 softmax_scale: float,
155 head_dim: int,
156 nope_head_dim: int,
157 rope_head_dim: int,
158 out: torch.Tensor, # [num_tokens, num_heads, head_dim]
159 ) -> None:
160 """XPU decode: dequant FP8 pages to BF16, then BF16 sparse MLA attention.
161
162 Keeps external FP8 KV cache layout identical to CUDA/ROCm.
163 Performance is slower due to on-the-fly dequant, but correctness is
164 guaranteed by reusing the validated BF16 attention kernel.
165 """
166 num_tokens = q.shape[0]
167 device = q.device
168
169 # Determine max topk and swa widths
170 if not swa_only and topk_indices is not None:
171 topk_idx_2d = (
172 topk_indices.squeeze(1) if topk_indices.dim() == 3 else topk_indices
173 )
174 max_topk = topk_idx_2d.shape[1]
175 else:
176 topk_idx_2d = None
177 max_topk = 0
178
179 swa_idx_2d = swa_indices.squeeze(1) if swa_indices.dim() == 3 else swa_indices
180 max_swa = swa_idx_2d.shape[1]
181
182 K_total = max_topk + max_swa
183
184 # Allocate flat workspace: [num_tokens * K_total, 512] bf16
185 workspace = torch.empty(
186 (num_tokens * K_total, OUTPUT_DIM), dtype=torch.bfloat16, device=device
187 )
188 ws_3d = workspace.view(num_tokens, K_total, OUTPUT_DIM)
189
190 # Dequant+gather topk slots from compressed cache
191 if not swa_only and topk_idx_2d is not None and kv_cache is not None:
192 topk_flat = topk_idx_2d.reshape(-1).to(torch.int32)
193 topk_buf = torch.empty(
194 (num_tokens * max_topk, OUTPUT_DIM), dtype=torch.bfloat16, device=device
195 )
196 compressed_block_size = kv_cache.shape[1]
197 dequant_gather_slots(topk_buf, kv_cache, topk_flat, compressed_block_size)
198 ws_3d[:, :max_topk, :] = topk_buf.view(num_tokens, max_topk, OUTPUT_DIM)
199
200 # Dequant+gather SWA slots
201 swa_flat = swa_idx_2d.reshape(-1).to(torch.int32)
202 swa_buf = torch.empty(
203 (num_tokens * max_swa, OUTPUT_DIM), dtype=torch.bfloat16, device=device
204 )
205 swa_block_size = swa_kv_cache.shape[1]
206 dequant_gather_slots(swa_buf, swa_kv_cache, swa_flat, swa_block_size)
207
208 ws_3d[:, max_topk:, :] = swa_buf.view(num_tokens, max_swa, OUTPUT_DIM)
209
210 # Build combined indices into the flat workspace and combined lengths.
211 # Workspace layout per token t: [topk_0..topk_{max_topk-1}, swa_0..swa_{max_swa-1}]
212 # Flat index for token t, position p = t * K_total + p
213 #
214 # IMPORTANT: The attention kernel uses combined_lens as a position cutoff —
215 # it only reads indices[0..combined_lens-1]. So indices must be PACKED
216 # contiguously: [valid_topk_indices..., valid_swa_indices..., -1 padding...]
217 if not swa_only and topk_lens is not None:
218 combined_lens = (topk_lens + swa_lens).to(torch.int32)
219 else:
220 combined_lens = swa_lens.to(torch.int32)
221
222 max_combined = int(combined_lens.max().item()) if combined_lens.numel() > 0 else 0
223 # Round up to BLOCK_N=16 alignment for kernel efficiency
224 _BLOCK_N = 16
225 max_combined_padded = ((max_combined + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N
226
227 # Build packed index table: [num_tokens, max_combined_padded]
228 # Each token t: [topk_0..topk_{tlen-1}, swa_0..swa_{slen-1}, -1 padding]
229 # Vectorized: for each token, topk indices are t*K_total + 0..tlen-1,
230 # swa indices are t*K_total + max_topk + 0..slen-1
231 combined_indices = torch.full(
232 (num_tokens, max_combined_padded),
233 fill_value=-1,
234 dtype=torch.int32,
235 device=device,
236 )
237
238 token_offsets = (
239 torch.arange(num_tokens, device=device, dtype=torch.int32) * K_total
240 ) # [B]
241
242 if not swa_only and topk_lens is not None:
243 # Pack topk: for each token, write t*K_total + 0..tlen-1 at positions 0..tlen-1
244 max_tlen = int(topk_lens.max().item())
245 topk_range = torch.arange(max_tlen, device=device, dtype=torch.int32).unsqueeze(
246 0
247 )
248 topk_valid = topk_range < topk_lens.unsqueeze(1)
249 topk_ws_indices = token_offsets.unsqueeze(1) + topk_range
250 combined_indices[:, :max_tlen] = torch.where(
251 topk_valid,
252 topk_ws_indices,
253 torch.tensor(-1, dtype=torch.int32, device=device),
254 )
255 # Pack swa after topk: positions tlen..tlen+slen-1
256 # Since tlen varies per token, we need per-token offset
257 swa_range = torch.arange(max_swa, device=device, dtype=torch.int32).unsqueeze(0)
258 swa_valid = swa_range < swa_lens.unsqueeze(1)
259 swa_ws_indices = token_offsets.unsqueeze(1) + max_topk + swa_range
260 # Write at position topk_lens[t] + swa_pos for each token
261 for t_idx in range(num_tokens):
262 tlen = int(topk_lens[t_idx].item())
263 slen = int(swa_lens[t_idx].item())
264 combined_indices[t_idx, tlen : tlen + slen] = swa_ws_indices[t_idx, :slen]
265 else:
266 # SWA-only: pack swa indices at positions 0..slen-1
267 # Use min(max_swa, max_combined_padded) because combined_indices only
268 # has max_combined_padded columns, and all valid entries fit within it.
269 effective_swa = min(max_swa, max_combined_padded)
270 swa_range = torch.arange(
271 effective_swa, device=device, dtype=torch.int32
272 ).unsqueeze(0)
273 swa_valid = swa_range < swa_lens.unsqueeze(1)
274 swa_ws_indices = token_offsets.unsqueeze(1) + swa_range # max_topk=0
275 combined_indices[:, :effective_swa] = torch.where(
276 swa_valid,
277 swa_ws_indices,
278 torch.tensor(-1, dtype=torch.int32, device=device),
279 )
280
281 # Call BF16 sparse MLA kernel
282 out_attn, _, _ = triton_bf16_mla_sparse_interface(
283 q=q,
284 kv=workspace.unsqueeze(1),
285 indices=combined_indices.unsqueeze(1),
286 sm_scale=softmax_scale,
287 d_v=q.shape[-1],
288 block_dpe=0,
289 )
290 out.copy_(out_attn)