@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 506 lines 19.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """V4.1 state saving/compression and independently schedulable cache insertion."""
4
5 import torch
6
7 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import _fp32x2_to_fp4x2
8 from vllm.platforms import current_platform
9 from vllm.triton_utils import tl, triton
10
11 if current_platform.is_rocm():
12 from vllm.platforms.rocm import _ON_GFX950
13 else:
14 _ON_GFX950 = False
15
16 # The ring and raw rows a ratio-2 request program pools are not adjacent, so
17 # reading them as one tile means joining two pointers. ROCm's Triton fails to
18 # legalize `tt.join` on pointers, so there the two rows are loaded separately.
19 _JOIN_ROW_PTRS = not current_platform.is_rocm()
20
21
22 def fused_save_compress_norm(
23 kv_score: torch.Tensor,
24 positions: torch.Tensor,
25 state_cache: torch.Tensor | None,
26 slot_mapping: torch.Tensor,
27 query_start_loc: torch.Tensor | None,
28 token_to_req_indices: torch.Tensor | None,
29 rms_norm_weight: torch.Tensor,
30 rms_norm_eps: float,
31 compress_ratio: int,
32 latent_out: torch.Tensor,
33 ) -> None:
34 """Pool each closed group into a normalized BF16 latent; save FP32 states.
35
36 The latent feeds the main-cache insert and the indexer K path, which the
37 attention layer schedules on separate streams.
38
39 Ratio 2 keeps one ring block per request holding the open group's rows:
40 position ``p`` lives in row ``p % capacity`` and ``slot_mapping`` encodes
41 ``block * capacity + p % capacity``. The grid has one program per request
42 followed by one per pair of packed tokens. A request program handles the
43 group that the chunk's first token closes with its predecessor's ring row,
44 then stores the chunk's last ``capacity`` rows to the ring; because the
45 same program does both, ring reads and writes never race. A pair program
46 handles the group that ends inside its pair, reading both rows from the
47 raw input. Ratio 1 has no ring and one program per token; ``slot_mapping``
48 then only marks valid tokens.
49
50 Args:
51 kv_score: FP32 [tokens, 512] for CR1, [tokens, 1024] for CR2.
52 positions: Absolute positions of the packed request tokens.
53 state_cache: Ring FP32 [blocks, capacity, 1024] KV/score states (CR2).
54 slot_mapping: Ring slots (CR2) or main-cache slots (CR1).
55 query_start_loc: [num_reqs + 1] token offsets of each request's chunk.
56 token_to_req_indices: Request indices for the packed token rows.
57 rms_norm_weight: BF16 [512] normalization weight.
58 rms_norm_eps: RMSNorm epsilon.
59 compress_ratio: Group size, either 1 or 2.
60 latent_out: BF16 [tokens, 512], written only at valid group boundaries.
61
62 """
63 assert compress_ratio in (1, 2)
64 assert kv_score.dtype == torch.float32
65 assert kv_score.shape[1] == 512 * compress_ratio and kv_score.stride(1) == 1
66 assert latent_out.shape == (kv_score.shape[0], 512)
67 assert latent_out.is_contiguous() and latent_out.dtype == torch.bfloat16
68 assert positions.is_contiguous() and slot_mapping.is_contiguous()
69 # Rows stay 64-byte aligned, as the kernel's tl.multiple_of hint promises.
70 assert kv_score.stride(0) % 16 == 0
71 if compress_ratio == 2:
72 assert state_cache is not None and query_start_loc is not None
73 assert token_to_req_indices is not None
74 assert query_start_loc.is_contiguous()
75 assert token_to_req_indices.is_contiguous()
76 assert state_cache.dtype == torch.float32
77 assert state_cache.shape[2] == 1024 and state_cache.stride(2) == 1
78 assert state_cache.stride(1) % 16 == 0
79 state_stride, state_row_stride, state_block = (
80 state_cache.stride(0),
81 state_cache.stride(1),
82 state_cache.shape[1],
83 )
84 num_reqs = query_start_loc.numel() - 1
85 else:
86 state_cache = query_start_loc = token_to_req_indices = None
87 state_stride = state_row_stride = state_block = 1
88 num_reqs = 0
89 num_tokens = slot_mapping.numel()
90 assert num_tokens <= min(kv_score.shape[0], positions.numel())
91 if num_tokens == 0:
92 return
93 grid = num_reqs + triton.cdiv(num_tokens, compress_ratio)
94 _fused_save_compress_norm_kernel[(grid,)](
95 kv_score,
96 positions,
97 state_cache,
98 slot_mapping,
99 query_start_loc,
100 token_to_req_indices,
101 rms_norm_weight,
102 latent_out,
103 num_tokens,
104 num_reqs,
105 RAW_STRIDE=kv_score.stride(0),
106 STATE_STRIDE=state_stride,
107 STATE_ROW_STRIDE=state_row_stride,
108 STATE_BLOCK=state_block,
109 COMPRESS_RATIO=compress_ratio,
110 EPS=rms_norm_eps,
111 JOIN_ROW_PTRS=_JOIN_ROW_PTRS,
112 num_warps=4,
113 **({"launch_pdl": False} if current_platform.is_cuda() else {}),
114 )
115
116
117 @triton.jit(do_not_specialize=["num_tokens", "num_reqs"])
118 def _fused_save_compress_norm_kernel(
119 raw,
120 positions,
121 state,
122 state_slots,
123 query_start_loc,
124 req_ids,
125 norm_weight,
126 latent,
127 num_tokens,
128 num_reqs,
129 RAW_STRIDE: tl.constexpr,
130 STATE_STRIDE: tl.constexpr,
131 STATE_ROW_STRIDE: tl.constexpr,
132 STATE_BLOCK: tl.constexpr,
133 COMPRESS_RATIO: tl.constexpr,
134 EPS: tl.constexpr,
135 JOIN_ROW_PTRS: tl.constexpr,
136 ):
137 pid = tl.program_id(0)
138 d = tl.arange(0, 512)
139
140 if COMPRESS_RATIO == 2: # noqa: SIM102 (constexpr branch, then runtime)
141 if pid < num_reqs:
142 # Request program: the group closed by the chunk's first token,
143 # then the ring tail store.
144 start = tl.load(query_start_loc + pid)
145 end = tl.load(query_start_loc + pid + 1)
146 if start >= end:
147 return
148 state_slot = tl.load(state_slots + start)
149 if state_slot < 0:
150 return
151 position = tl.load(positions + start)
152 if (position + 1) % 2 == 0:
153 ring = state + (state_slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE
154 prev = ring + ((position - 1) % STATE_BLOCK) * STATE_ROW_STRIDE
155 current = raw + start.to(tl.int64) * RAW_STRIDE
156 if JOIN_ROW_PTRS:
157 # Joining pointers hides row alignment from Triton; restate it.
158 rows = tl.multiple_of(tl.join(prev, current)[:, None], (16, 16))
159 kv = tl.load(rows + d[None, :])
160 score = tl.load(rows + 512 + d[None, :])
161 # Every lane has read the ring before the tail store below
162 # writes it.
163 tl.debug_barrier()
164 pooled = tl.sum(kv * tl.softmax(score, 0), 0)
165 else:
166 kv_prev = tl.load(prev + d)
167 score_prev = tl.load(prev + 512 + d)
168 kv_current = tl.load(current + d)
169 score_current = tl.load(current + 512 + d)
170 # Every lane has read the ring before the tail store below
171 # writes it.
172 tl.debug_barrier()
173 peak = tl.maximum(score_prev, score_current)
174 weight_prev = tl.exp(score_prev - peak)
175 weight_current = tl.exp(score_current - peak)
176 pooled = (kv_prev * weight_prev + kv_current * weight_current) / (
177 weight_prev + weight_current
178 )
179 _store_latent(pooled, start, norm_weight, latent, EPS)
180 num_rows = tl.minimum(end - start, STATE_BLOCK)
181 for k in tl.range(0, num_rows):
182 token = end - num_rows + k
183 slot = tl.load(state_slots + token)
184 if slot >= 0:
185 row = (
186 state
187 + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE
188 + (slot % STATE_BLOCK) * STATE_ROW_STRIDE
189 )
190 src = raw + token.to(tl.int64) * RAW_STRIDE
191 tl.store(row + d, tl.load(src + d))
192 tl.store(row + 512 + d, tl.load(src + 512 + d))
193 return
194
195 # Group program: one token (ratio 1) or the token of the pair that ends
196 # a group (ratio 2). A chunk's first token belongs to its request program.
197 t = (pid - num_reqs) * COMPRESS_RATIO
198 if COMPRESS_RATIO == 2:
199 t += (tl.load(positions + t) + 1) % 2
200 state_slot = tl.load(state_slots + t, t < num_tokens, other=-1)
201 if state_slot < 0:
202 return
203 if COMPRESS_RATIO == 1:
204 pooled = tl.load(raw + t.to(tl.int64) * RAW_STRIDE + d)
205 else:
206 if t == tl.load(query_start_loc + tl.load(req_ids + t)):
207 return
208 rows = raw + (t.to(tl.int64) - 1 + tl.arange(0, 2)[:, None]) * RAW_STRIDE
209 kv = tl.load(rows + d[None, :])
210 score = tl.load(rows + 512 + d[None, :])
211 pooled = tl.sum(kv * tl.softmax(score, 0), 0)
212 _store_latent(pooled, t, norm_weight, latent, EPS)
213
214
215 @triton.jit
216 def _store_latent(pooled, t, norm_weight, latent, EPS: tl.constexpr):
217 d = tl.arange(0, 512)
218 weight = tl.load(norm_weight + d).to(tl.float32)
219 variance = tl.sum(pooled * pooled, 0) / 512
220 normed = pooled * tl.rsqrt(variance + EPS) * weight
221 tl.store(latent + t.to(tl.int64) * 512 + d, normed.to(tl.bfloat16))
222
223
224 def rope_quant_insert(
225 latent: torch.Tensor,
226 positions: torch.Tensor,
227 cos_sin_cache: torch.Tensor,
228 kv_cache: torch.Tensor,
229 slot_mapping: torch.Tensor,
230 compress_ratio: int,
231 fp8_scale: torch.Tensor | None = None,
232 ) -> None:
233 """Apply GPT-J RoPE and publish a latent to the compressed KV cache.
234
235 The BF16 latent supplies both NoPE quantization and RoPE input. It is read
236 only for valid slots at group boundaries. The cache dtype selects the
237 layout: ``uint8`` is a paged FlashMLA layout, whose record the per-token
238 byte width names -- 584 B for V4 (576 value bytes and eight segregated
239 UE8M0 scale bytes, including one zero padding scale), 528 B for V4.1
240 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 scales of
241 32 dims each), or 288 B for V4.1 NVFP4 (256 bytes of e2m1 pairs then 32
242 e4m3 scales of 16 dims each), which only the compressed cache uses.
243 ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows
244 read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``.
245 """
246 assert compress_ratio in (1, 2)
247 assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16
248 assert latent.is_contiguous()
249 num_tokens = slot_mapping.numel()
250 assert num_tokens <= min(latent.shape[0], positions.numel())
251 if num_tokens == 0:
252 return
253 launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
254 if kv_cache.dtype == torch.uint8:
255 kernel = {
256 584: _rope_quant_insert_kernel,
257 528: _rope_quant_insert_mxfp8_kernel,
258 288: _rope_quant_insert_nvfp4_kernel,
259 }.get(kv_cache.shape[-1])
260 assert kernel is not None, (
261 f"unsupported paged KV record width {kv_cache.shape[-1]}"
262 )
263 kernel[(num_tokens,)](
264 latent,
265 positions,
266 cos_sin_cache,
267 kv_cache,
268 slot_mapping,
269 COS_STRIDE=cos_sin_cache.stride(0),
270 CACHE_STRIDE=kv_cache.stride(0),
271 CACHE_BLOCK=kv_cache.shape[1],
272 COMPRESS_RATIO=compress_ratio,
273 SANITIZE_CACHE_NANS=_ON_GFX950,
274 num_warps=4,
275 **launch_kwargs,
276 )
277 return
278
279 assert kv_cache.dtype in (torch.bfloat16, torch.float8_e4m3fn)
280 assert kv_cache.shape[-1] == 512 and kv_cache.stride(-1) == 1
281 store_fp8 = kv_cache.dtype == torch.float8_e4m3fn
282 if store_fp8:
283 assert fp8_scale is not None and fp8_scale.numel() == 1
284 assert fp8_scale.dtype == torch.float32
285 _rope_plain_insert_kernel[(num_tokens,)](
286 latent,
287 positions,
288 cos_sin_cache,
289 kv_cache,
290 slot_mapping,
291 fp8_scale if store_fp8 else None,
292 COS_STRIDE=cos_sin_cache.stride(0),
293 CACHE_STRIDE=kv_cache.stride(0),
294 ROW_STRIDE=kv_cache.stride(1),
295 CACHE_BLOCK=kv_cache.shape[1],
296 COMPRESS_RATIO=compress_ratio,
297 STORE_FP8=store_fp8,
298 num_warps=4,
299 **launch_kwargs,
300 )
301
302
303 @triton.jit
304 def _rope_quant_insert_kernel(
305 latent,
306 positions,
307 cos_sin,
308 cache,
309 cache_slots,
310 COS_STRIDE: tl.constexpr,
311 CACHE_STRIDE: tl.constexpr,
312 CACHE_BLOCK: tl.constexpr,
313 COMPRESS_RATIO: tl.constexpr,
314 SANITIZE_CACHE_NANS: tl.constexpr,
315 ):
316 t = tl.program_id(0)
317 slot = tl.load(cache_slots + t)
318 if slot < 0:
319 return
320 position = tl.load(positions + t)
321 if (position + 1) % COMPRESS_RATIO != 0:
322 return
323 d = tl.arange(0, 512)
324 normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
325 page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
326 values = page + (slot % CACHE_BLOCK) * 576
327 scales = page + CACHE_BLOCK * 576 + (slot % CACHE_BLOCK) * 8
328
329 quant = tl.reshape(normed, (8, 64))
330 amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4)
331 exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0)))
332 scaled = quant * tl.reshape(tl.exp2(-exponent), (8, 1))
333 fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv)
334 packed = tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,))
335 tl.store(values + d, packed, d < 448)
336 s = tl.arange(0, 8)
337 max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
338 encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded)
339 tl.store(scales + s, encoded.to(tl.uint8), s < 7)
340 tl.store(scales + 7, tl.full((), 0, tl.uint8))
341
342 even, odd = tl.split(tl.reshape(normed, (256, 2)))
343 pair = tl.arange(0, 256) - 224
344 cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
345 c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
346 s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
347 rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
348 if SANITIZE_CACHE_NANS:
349 rotated = tl.where(rotated == rotated, rotated, 0.0)
350 rope_dst = (values + 448).to(tl.pointer_type(tl.bfloat16))
351 tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448)
352
353
354 @triton.jit
355 def _rope_quant_insert_mxfp8_kernel(
356 latent,
357 positions,
358 cos_sin,
359 cache,
360 cache_slots,
361 COS_STRIDE: tl.constexpr,
362 CACHE_STRIDE: tl.constexpr,
363 CACHE_BLOCK: tl.constexpr,
364 COMPRESS_RATIO: tl.constexpr,
365 SANITIZE_CACHE_NANS: tl.constexpr,
366 ):
367 """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims.
368
369 The RoPE dims are quantized here too, so unlike the V4 kernel the rotation
370 has to happen before the scales are picked.
371 """
372 t = tl.program_id(0)
373 slot = tl.load(cache_slots + t)
374 if slot < 0:
375 return
376 position = tl.load(positions + t)
377 if (position + 1) % COMPRESS_RATIO != 0:
378 return
379 d = tl.arange(0, 512)
380 normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
381
382 # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
383 even, odd = tl.split(tl.reshape(normed, (256, 2)))
384 pair = tl.arange(0, 256) - 224
385 cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
386 c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
387 s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
388 rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
389 if SANITIZE_CACHE_NANS:
390 rotated = tl.where(rotated == rotated, rotated, 0.0)
391
392 page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
393 values = page + (slot % CACHE_BLOCK) * 512
394 scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16
395
396 quant = tl.reshape(rotated, (16, 32))
397 amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4)
398 exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0)))
399 scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1))
400 fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv)
401 tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,)))
402
403 max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
404 encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded)
405 tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8))
406
407
408 @triton.jit
409 def _rope_quant_insert_nvfp4_kernel(
410 latent,
411 positions,
412 cos_sin,
413 cache,
414 cache_slots,
415 COS_STRIDE: tl.constexpr,
416 CACHE_STRIDE: tl.constexpr,
417 CACHE_BLOCK: tl.constexpr,
418 COMPRESS_RATIO: tl.constexpr,
419 SANITIZE_CACHE_NANS: tl.constexpr,
420 ):
421 """V4.1 NVFP4 record: RoPE, then e2m1 with one e4m3 scale per 16 dims.
422
423 The scale is ``amax / 6`` (6 is e2m1's largest magnitude) clamped to the
424 e4m3 range, with no per-tensor scale on top.
425 """
426 t = tl.program_id(0)
427 slot = tl.load(cache_slots + t)
428 if slot < 0:
429 return
430 position = tl.load(positions + t)
431 if (position + 1) % COMPRESS_RATIO != 0:
432 return
433 d = tl.arange(0, 512)
434 normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
435
436 # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
437 even, odd = tl.split(tl.reshape(normed, (256, 2)))
438 pair = tl.arange(0, 256) - 224
439 cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
440 c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
441 s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
442 rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
443 if SANITIZE_CACHE_NANS:
444 rotated = tl.where(rotated == rotated, rotated, 0.0)
445
446 tiles = tl.reshape(rotated, (32, 16))
447 amax = tl.max(tl.abs(tiles), 1)
448 # 2**-9 is the smallest normal e4m3 magnitude; 448 the largest.
449 scale = tl.clamp(amax * (1.0 / 6.0), 0.001953125, 448.0).to(tl.float8e4nv)
450 # Round-to-nearest division: Triton's default div.full misplaces values
451 # that land exactly on an e2m1 tie.
452 scaled = tl.math.div_rn(tiles, tl.reshape(scale.to(tl.float32), (32, 1)))
453 lo, hi = tl.split(tl.reshape(tl.reshape(scaled, (512,)), (256, 2)))
454
455 page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
456 tl.store(
457 page + (slot % CACHE_BLOCK) * 256 + tl.arange(0, 256),
458 _fp32x2_to_fp4x2(lo, hi),
459 )
460 scales = page + CACHE_BLOCK * 256 + (slot % CACHE_BLOCK) * 32
461 tl.store(scales + tl.arange(0, 32), scale.to(tl.uint8, bitcast=True))
462
463
464 @triton.jit
465 def _rope_plain_insert_kernel(
466 latent,
467 positions,
468 cos_sin,
469 cache,
470 cache_slots,
471 fp8_scale,
472 COS_STRIDE: tl.constexpr,
473 CACHE_STRIDE: tl.constexpr,
474 ROW_STRIDE: tl.constexpr,
475 CACHE_BLOCK: tl.constexpr,
476 COMPRESS_RATIO: tl.constexpr,
477 STORE_FP8: tl.constexpr,
478 ):
479 t = tl.program_id(0)
480 slot = tl.load(cache_slots + t)
481 if slot < 0:
482 return
483 position = tl.load(positions + t)
484 if (position + 1) % COMPRESS_RATIO != 0:
485 return
486 d = tl.arange(0, 512)
487 normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
488
489 # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there.
490 even, odd = tl.split(tl.reshape(normed, (256, 2)))
491 pair = tl.arange(0, 256) - 224
492 cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
493 c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
494 s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
495 row = tl.interleave(even * c - odd * s, odd * c + even * s).to(tl.bfloat16)
496
497 dst = (
498 cache
499 + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
500 + (slot % CACHE_BLOCK) * ROW_STRIDE
501 )
502 if STORE_FP8:
503 scaled = row.to(tl.float32) * (1.0 / tl.load(fp8_scale))
504 tl.store(dst + d, tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv))
505 else:
506 tl.store(dst + d, row)