@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 613 lines 20.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 from vllm.model_executor.layers.quantization.utils.quant_utils import (
7 FP8_DTYPE,
8 get_fp8_min_max,
9 )
10 from vllm.platforms import current_platform
11 from vllm.triton_utils import tl, triton
12 from vllm.utils.torch_utils import is_quantized_kv_cache
13 from vllm.v1.kv_cache_interface import KVQuantMode
14
15 FP8_MIN, FP8_MAX = get_fp8_min_max()
16
17 _NATIVE_KV_CACHE_DTYPES = {"auto", "float16", "bfloat16", "float32", "half", "float"}
18
19
20 def _is_supported_kv_cache_dtype(kv_cache_dtype: str) -> bool:
21 if not (
22 kv_cache_dtype in _NATIVE_KV_CACHE_DTYPES
23 or is_quantized_kv_cache(kv_cache_dtype)
24 ):
25 return False
26 if kv_cache_dtype.startswith("fp8"):
27 return current_platform.has_device_capability(89) or current_platform.is_xpu()
28 if kv_cache_dtype == "bfloat16":
29 return current_platform.has_device_capability(80) or current_platform.is_xpu()
30 return True
31
32
33 @triton.jit
34 def reshape_and_cache_kernel_flash(
35 key_ptr, # [num_tokens, num_heads, head_size]
36 value_ptr, # [num_tokens, num_heads, head_size]
37 key_cache_ptr, # [num_blocks, block_size, num_heads, head_size]
38 value_cache_ptr, # [num_blocks, block_size, num_heads, head_size]
39 slot_mapping_ptr, # [num_tokens]
40 k_scale, # float32
41 v_scale, # float32
42 # strides
43 key_stride: tl.int64,
44 value_stride: tl.int64,
45 block_stride: tl.int64,
46 head_stride: tl.int64,
47 dim_stride_k: tl.int64,
48 dim_stride_v: tl.int64,
49 page_stride: tl.int64,
50 num_heads: tl.constexpr,
51 head_size: tl.constexpr,
52 block_size: tl.constexpr,
53 x: tl.constexpr,
54 USE_HEAD_MAJOR_LAYOUT: tl.constexpr,
55 # FP8 flags
56 FP8_KV_CACHE: tl.constexpr,
57 # tune parameters
58 TILE_SIZE: tl.constexpr,
59 ):
60 token_idx = tl.program_id(axis=0)
61 slot_idx = tl.load(slot_mapping_ptr + token_idx).to(tl.int64)
62 if slot_idx < 0:
63 # Padding token that should be ignored.
64 return
65
66 block_idx = slot_idx // block_size
67 block_offset = slot_idx % block_size
68
69 tile_i = tl.program_id(axis=1)
70 tile_offs = tl.arange(0, TILE_SIZE)
71 tile_pos = tile_i * TILE_SIZE + tile_offs
72 src_key_idx = token_idx * key_stride
73 src_value_idx = token_idx * value_stride
74
75 if USE_HEAD_MAJOR_LAYOUT:
76 # Decompose the tile index back into head and dim coordinates.
77 cur_head = tile_pos // head_size
78 cur_dim = tile_pos % head_size
79 # Value addressing (4D): [Block, Head, Dim, Slot]
80 tgt_idx_v = (
81 block_idx * block_stride
82 + cur_head * head_stride
83 + cur_dim * dim_stride_v
84 + block_offset * 1
85 )
86 # Key addressing (5D): [Block, Head, Dim//8, Slot, 8]
87 tgt_idx_k = (
88 block_idx * block_stride
89 + cur_head * head_stride
90 + (cur_dim // x) * dim_stride_k
91 + block_offset * x
92 + (cur_dim % x)
93 )
94 else:
95 cur_head = tile_pos // head_size
96 cur_dim = tile_pos % head_size
97 tgt_idx_k = (
98 block_idx * block_stride
99 + block_offset * page_stride
100 + cur_head * head_stride
101 + cur_dim
102 )
103 tgt_idx_v = tgt_idx_k
104
105 # [TILE_SIZE]
106 key_load = tl.load(
107 key_ptr + src_key_idx + tile_pos, mask=tile_pos < (num_heads * head_size)
108 )
109 if FP8_KV_CACHE:
110 # tl.store will do the correct implicit cast to fp8,
111 # based on the key_cache_ptr.dtype.element_ty
112 key_tile = key_load if key_load.dtype.is_fp8() else key_load / tl.load(k_scale)
113 else:
114 key_tile = key_load
115
116 # [TILE_SIZE]
117 value_load = tl.load(
118 value_ptr + src_value_idx + tile_pos, mask=tile_pos < (num_heads * head_size)
119 )
120 if FP8_KV_CACHE:
121 if value_load.dtype.is_fp8():
122 value_tile = value_load
123 else:
124 # tl.store will do the correct implicit cast to fp8,
125 # based on the value_cache_ptr.dtype.element_ty
126 value_tile = value_load / tl.load(v_scale)
127 else:
128 value_tile = value_load
129
130 tl.store(
131 key_cache_ptr + tgt_idx_k,
132 key_tile,
133 mask=tile_pos < (num_heads * head_size),
134 )
135 tl.store(
136 value_cache_ptr + tgt_idx_v,
137 value_tile,
138 mask=tile_pos < (num_heads * head_size),
139 )
140 return
141
142
143 # ---------------------------------------------------------------------------
144 # Per-token-head dynamic quantization kernel
145 # Grid: (num_tokens, NUM_KV_HEADS)
146 # Each program handles one (token, head) pair:
147 # 1. Loads K (or V) for that single head
148 # 2. Computes absmax across head_size → scale = absmax / QUANT_MAX
149 # 3. Quantizes and stores the data + per-head scale
150 #
151 # Parametrised by QUANT_MAX / QUANT_MIN so the same code path works
152 # for int8 (±127/128), fp8_e4m3 (±448), and other formats.
153 # ---------------------------------------------------------------------------
154 @triton.jit
155 def _reshape_cache_per_token_head(
156 key_ptr, # [num_tokens, num_kv_heads, head_size]
157 value_ptr, # [num_tokens, num_kv_heads, head_size_v]
158 key_cache_ptr, # [num_blocks, block_size, num_kv_heads, head_size]
159 value_cache_ptr, # [num_blocks, block_size, num_kv_heads, head_size_v]
160 k_scale_cache_ptr, # [num_blocks, block_size, num_kv_heads] float32
161 v_scale_cache_ptr, # [num_blocks, block_size, num_kv_heads] float32
162 slot_mapping_ptr, # [num_tokens]
163 stride_key_tok: tl.int64,
164 stride_key_head: tl.int64,
165 stride_val_tok: tl.int64,
166 stride_val_head: tl.int64,
167 stride_kc_blk: tl.int64, # key_cache stride over blocks
168 stride_kc_slot: tl.int64, # key_cache stride over slots
169 stride_kc_head: tl.int64, # key_cache stride over heads
170 stride_vc_blk: tl.int64,
171 stride_vc_slot: tl.int64,
172 stride_vc_head: tl.int64,
173 stride_ks_blk: tl.int64, # k_scale_cache stride[0] (blocks)
174 stride_ks_slot: tl.int64, # k_scale_cache stride[1] (slots)
175 stride_ks_head: tl.int64, # k_scale_cache stride[2] (heads)
176 stride_vs_blk: tl.int64, # v_scale_cache stride[0] (blocks)
177 stride_vs_slot: tl.int64, # v_scale_cache stride[1] (slots)
178 stride_vs_head: tl.int64, # v_scale_cache stride[2] (heads)
179 block_size: tl.constexpr,
180 head_size: tl.constexpr,
181 head_size_v: tl.constexpr,
182 HEAD_SIZE_PADDED: tl.constexpr, # next_power_of_2(max(head_size, head_size_v))
183 QUANT_MAX: tl.constexpr = 127.0,
184 QUANT_MIN: tl.constexpr = -128.0,
185 IS_INT_QUANT: tl.constexpr = False,
186 ):
187 tok = tl.program_id(0)
188 head = tl.program_id(1)
189
190 slot = tl.load(slot_mapping_ptr + tok).to(tl.int64)
191 if slot < 0:
192 return
193
194 blk = slot // block_size
195 slot_in_blk = slot % block_size
196
197 dim_offs = tl.arange(0, HEAD_SIZE_PADDED)
198
199 # ---- Key: load one head → absmax → quantize → store -------------------
200 k_mask = dim_offs < head_size
201 k_h = tl.load(
202 key_ptr + tok * stride_key_tok + head * stride_key_head + dim_offs,
203 mask=k_mask,
204 other=0.0,
205 ).to(tl.float32)
206
207 k_scale = tl.maximum(tl.max(tl.abs(k_h)) / QUANT_MAX, 1e-6)
208 tl.store(
209 k_scale_cache_ptr
210 + blk * stride_ks_blk
211 + slot_in_blk * stride_ks_slot
212 + head * stride_ks_head,
213 k_scale,
214 )
215
216 k_q = k_h * (1.0 / k_scale)
217 if IS_INT_QUANT:
218 # Round half away from zero before the int8 store truncates.
219 k_q = tl.where(k_q >= 0, k_q + 0.5, k_q - 0.5)
220 k_q = tl.clamp(k_q, QUANT_MIN, QUANT_MAX)
221 tl.store(
222 key_cache_ptr
223 + blk * stride_kc_blk
224 + slot_in_blk * stride_kc_slot
225 + head * stride_kc_head
226 + dim_offs,
227 k_q,
228 mask=k_mask,
229 )
230
231 # ---- Value: same per-head approach ------------------------------------
232 v_mask = dim_offs < head_size_v
233 v_h = tl.load(
234 value_ptr + tok * stride_val_tok + head * stride_val_head + dim_offs,
235 mask=v_mask,
236 other=0.0,
237 ).to(tl.float32)
238
239 v_scale = tl.maximum(tl.max(tl.abs(v_h)) / QUANT_MAX, 1e-6)
240 tl.store(
241 v_scale_cache_ptr
242 + blk * stride_vs_blk
243 + slot_in_blk * stride_vs_slot
244 + head * stride_vs_head,
245 v_scale,
246 )
247
248 v_q = v_h * (1.0 / v_scale)
249 if IS_INT_QUANT:
250 # Round half away from zero before the int8 store truncates.
251 v_q = tl.where(v_q >= 0, v_q + 0.5, v_q - 0.5)
252 v_q = tl.clamp(v_q, QUANT_MIN, QUANT_MAX)
253 tl.store(
254 value_cache_ptr
255 + blk * stride_vc_blk
256 + slot_in_blk * stride_vc_slot
257 + head * stride_vc_head
258 + dim_offs,
259 v_q,
260 mask=v_mask,
261 )
262
263
264 # Mapping from cache torch dtype to (QUANT_MAX, QUANT_MIN) for the
265 # per-token-head quantization kernel.
266 _PER_TOKEN_HEAD_QUANT_PARAMS: dict[torch.dtype, tuple[float, float]] = {
267 torch.int8: (127.0, -128.0),
268 FP8_DTYPE: (FP8_MAX, FP8_MIN),
269 }
270
271
272 def triton_reshape_and_cache_flash_per_token_head_quant(
273 key: torch.Tensor, # [num_tokens, num_kv_heads, head_size]
274 value: torch.Tensor, # [num_tokens, num_kv_heads, head_size_v]
275 key_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads, head_size]
276 value_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads, head_size_v]
277 k_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32
278 v_scale_cache: torch.Tensor, # [num_blocks, block_size, num_kv_heads] float32
279 slot_mapping: torch.Tensor, # [num_tokens]
280 kv_quant_mode: KVQuantMode,
281 ):
282 """Quantize key/value per (token, head) and write to paged cache.
283
284 Computes one scale = absmax / QUANT_MAX per (token, head), stores
285 quantized data in key_cache/value_cache, and stores the float32
286 scale in k_scale_cache/v_scale_cache.
287
288 INT4 needs sub-byte packing + a Hadamard rotation, so it is handled by
289 its own kernel; INT8 / FP8 share this kernel, with the quantization
290 range (QUANT_MAX, QUANT_MIN) derived from the cache tensor dtype.
291 """
292 if kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
293 from vllm.v1.attention.ops.int4_per_token_head import (
294 reshape_and_cache_int4,
295 )
296
297 reshape_and_cache_int4(
298 key,
299 value,
300 key_cache,
301 value_cache,
302 slot_mapping,
303 k_scale_cache=k_scale_cache,
304 v_scale_cache=v_scale_cache,
305 )
306 return
307
308 cache_dtype = key_cache.dtype
309 quant_params = _PER_TOKEN_HEAD_QUANT_PARAMS.get(cache_dtype)
310 if quant_params is None:
311 raise ValueError(
312 f"Per-token-head quantization not supported for cache dtype "
313 f"{cache_dtype}. Supported: {list(_PER_TOKEN_HEAD_QUANT_PARAMS)}"
314 )
315 quant_max, quant_min = quant_params
316
317 num_tokens, num_kv_heads, head_size = key.shape
318 head_size_v = value.shape[2]
319 head_size_padded = triton.next_power_of_2(max(head_size, head_size_v))
320
321 block_size = key_cache.shape[1]
322
323 if current_platform.is_rocm() or current_platform.is_xpu():
324 num_warps = 4
325 else:
326 num_warps = min(16, max(1, head_size_padded // 32))
327
328 _reshape_cache_per_token_head[(num_tokens, num_kv_heads)](
329 key_ptr=key,
330 value_ptr=value,
331 key_cache_ptr=key_cache,
332 value_cache_ptr=value_cache,
333 k_scale_cache_ptr=k_scale_cache,
334 v_scale_cache_ptr=v_scale_cache,
335 slot_mapping_ptr=slot_mapping,
336 stride_key_tok=key.stride(0),
337 stride_key_head=key.stride(1),
338 stride_val_tok=value.stride(0),
339 stride_val_head=value.stride(1),
340 stride_kc_blk=key_cache.stride(0),
341 stride_kc_slot=key_cache.stride(1),
342 stride_kc_head=key_cache.stride(2),
343 stride_vc_blk=value_cache.stride(0),
344 stride_vc_slot=value_cache.stride(1),
345 stride_vc_head=value_cache.stride(2),
346 stride_ks_blk=k_scale_cache.stride(0),
347 stride_ks_slot=k_scale_cache.stride(1),
348 stride_ks_head=k_scale_cache.stride(2),
349 stride_vs_blk=v_scale_cache.stride(0),
350 stride_vs_slot=v_scale_cache.stride(1),
351 stride_vs_head=v_scale_cache.stride(2),
352 block_size=block_size,
353 head_size=head_size,
354 head_size_v=head_size_v,
355 HEAD_SIZE_PADDED=head_size_padded,
356 QUANT_MAX=quant_max,
357 QUANT_MIN=quant_min,
358 IS_INT_QUANT=cache_dtype == torch.int8,
359 num_warps=num_warps,
360 )
361
362
363 def triton_reshape_and_cache_flash(
364 key: torch.Tensor, # [num_tokens, num_heads, head_size]
365 value: torch.Tensor, # [num_tokens, num_heads, head_size]
366 # [num_blocks, block_size, num_heads, head_size]
367 key_cache: torch.Tensor,
368 # [num_blocks, block_size, num_heads, head_size]
369 value_cache: torch.Tensor,
370 slot_mapping: torch.Tensor, # [num_tokens]
371 kv_cache_dtype: str, # "auto", "fp8"
372 k_scale: torch.Tensor, # float32
373 v_scale: torch.Tensor, # float32
374 ):
375 num_heads = key.shape[1]
376 head_size = key.shape[2]
377
378 use_head_major_layout = key_cache.ndim == 5
379 if use_head_major_layout:
380 block_size = key_cache.shape[3]
381 x = key_cache.shape[4]
382 head_stride = key_cache.stride(1)
383 dim_stride_k = key_cache.stride(2)
384 dim_stride_v = value_cache.stride(2)
385 else:
386 block_size = key_cache.shape[1]
387 x = 1
388 dim_stride_k = 0
389 dim_stride_v = 0
390 head_stride = key_cache.stride()[2]
391 n = num_heads * head_size
392 key_stride = key.stride()[0]
393 value_stride = value.stride()[0]
394 block_stride = key_cache.stride()[0]
395 page_stride = key_cache.stride()[1]
396
397 assert _is_supported_kv_cache_dtype(kv_cache_dtype), (
398 f"Triton reshape-and-cache cannot store kv_cache_dtype={kv_cache_dtype} "
399 f"on this device: an FP8 KV cache needs native fp8e4nv (SM89+). Use "
400 f"--kv-cache-dtype bfloat16 (or float16 on SM75)."
401 )
402 kv_cache_torch_dtype = (
403 current_platform.fp8_dtype()
404 if is_quantized_kv_cache(kv_cache_dtype)
405 else key_cache.dtype
406 )
407
408 if key_cache.dtype != kv_cache_torch_dtype and is_quantized_kv_cache(
409 kv_cache_dtype
410 ):
411 # to avoid erounous implicit cast in triton kernel (tl.store to uint8)
412 # (e.g. explicit cast to fp8e4m3fnuz is not supported in triton 3.4)
413 key_cache = key_cache.view(kv_cache_torch_dtype)
414 value_cache = value_cache.view(kv_cache_torch_dtype)
415 FP8_KV_CACHE = is_quantized_kv_cache(kv_cache_dtype)
416 # heuristics instead of autotuning
417 TILE_SIZE = min(2048, triton.next_power_of_2(n))
418 if current_platform.is_rocm() or current_platform.is_xpu():
419 num_stages = 4
420 num_warps = 8
421 else: # cuda
422 num_stages = 10
423 num_warps = 16
424 if torch.cuda.get_device_capability(key.device)[0] < 9:
425 TILE_SIZE = min(512, TILE_SIZE)
426
427 # TODO(ngl): maybe replace with static launch grid to avoid overhead if
428 # using cudagraphs
429 grid = lambda meta: (
430 slot_mapping.shape[0],
431 triton.cdiv(n, meta["TILE_SIZE"]),
432 )
433
434 reshape_and_cache_kernel_flash[grid](
435 key_ptr=key,
436 value_ptr=value,
437 key_cache_ptr=key_cache,
438 value_cache_ptr=value_cache,
439 slot_mapping_ptr=slot_mapping,
440 k_scale=k_scale,
441 v_scale=v_scale,
442 # strides
443 key_stride=key_stride,
444 value_stride=value_stride,
445 block_stride=block_stride,
446 head_stride=head_stride,
447 dim_stride_k=dim_stride_k,
448 dim_stride_v=dim_stride_v,
449 page_stride=page_stride,
450 num_heads=num_heads,
451 head_size=head_size,
452 block_size=block_size,
453 x=x,
454 USE_HEAD_MAJOR_LAYOUT=use_head_major_layout,
455 FP8_KV_CACHE=FP8_KV_CACHE,
456 # autotune parameters
457 TILE_SIZE=TILE_SIZE,
458 num_warps=num_warps,
459 num_stages=num_stages,
460 )
461
462
463 @triton.jit
464 def reshape_and_cache_kernel_flash_diffkv(
465 key_ptr, # [num_tokens, num_heads, head_size]
466 value_ptr, # [num_tokens, num_heads, head_size_v]
467 kv_cache_ptr, # [num_blocks, block_size, num_heads, head_size + head_size_v]
468 slot_mapping_ptr, # [num_tokens]
469 k_scale, # float32
470 v_scale, # float32
471 # strides
472 key_stride: tl.int64,
473 value_stride: tl.int64,
474 block_stride: tl.int64,
475 page_stride: tl.int64,
476 head_stride: tl.int64,
477 num_heads: tl.constexpr,
478 head_size_k: tl.constexpr,
479 head_size_v: tl.constexpr,
480 block_size: tl.constexpr,
481 # FP8 flags
482 FP8_KV_CACHE: tl.constexpr,
483 # tune parameters
484 TILE_SIZE: tl.constexpr,
485 ):
486 token_idx = tl.program_id(axis=0)
487 slot_idx = tl.load(slot_mapping_ptr + token_idx).to(tl.int64)
488 if slot_idx < 0:
489 # Padding token that should be ignored.
490 return
491
492 tile_i = tl.program_id(axis=1)
493 tile_offs = tl.arange(0, TILE_SIZE)
494
495 block_idx = slot_idx // block_size
496 block_offset = slot_idx % block_size
497
498 src_key_idx = token_idx * key_stride + tile_i * head_size_k
499 src_value_idx = token_idx * value_stride + tile_i * head_size_v
500
501 tgt_idx = (
502 block_idx * block_stride + block_offset * page_stride + tile_i * head_stride
503 )
504
505 # [TILE_SIZE]
506 key_load = tl.load(key_ptr + src_key_idx + tile_offs, mask=tile_offs < head_size_k)
507 if FP8_KV_CACHE:
508 # tl.store will do the correct implicit cast to fp8,
509 # based on the key_cache_ptr.dtype.element_ty
510 key_tile = key_load if key_load.dtype.is_fp8() else key_load / tl.load(k_scale)
511 else:
512 key_tile = key_load
513
514 # [TILE_SIZE]
515 value_load = tl.load(
516 value_ptr + src_value_idx + tile_offs, mask=tile_offs < head_size_v
517 )
518 if FP8_KV_CACHE:
519 if value_load.dtype.is_fp8():
520 value_tile = value_load
521 else:
522 # tl.store will do the correct implicit cast to fp8,
523 # based on the value_cache_ptr.dtype.element_ty
524 value_tile = value_load / tl.load(v_scale)
525 else:
526 value_tile = value_load
527
528 tl.store(
529 kv_cache_ptr + tgt_idx + tile_offs,
530 key_tile,
531 mask=tile_offs < head_size_k,
532 )
533 tl.store(
534 kv_cache_ptr + tgt_idx + head_size_k + tile_offs,
535 value_tile,
536 mask=tile_offs < head_size_v,
537 )
538 return
539
540
541 def triton_reshape_and_cache_flash_diffkv(
542 key: torch.Tensor, # [num_tokens, num_heads, head_size]
543 value: torch.Tensor, # [num_tokens, num_heads, head_size_v]
544 # Strided [num_blocks, block_size, num_heads, head_size + head_size_v].
545 kv_cache: torch.Tensor,
546 slot_mapping: torch.Tensor, # [num_tokens]
547 kv_cache_dtype: str, # "auto", "fp8"
548 k_scale: torch.Tensor, # float32
549 v_scale: torch.Tensor, # float32
550 ):
551 num_heads = key.shape[1]
552 head_size_k = key.shape[2]
553 head_size_v = value.shape[2]
554 block_size = kv_cache.shape[1]
555
556 k_stride = key.stride()[0]
557 v_stride = value.stride()[0]
558 block_stride = kv_cache.stride()[0]
559 page_stride = kv_cache.stride()[1]
560 head_stride = kv_cache.stride()[2]
561
562 kv_cache_torch_dtype = (
563 current_platform.fp8_dtype()
564 if is_quantized_kv_cache(kv_cache_dtype)
565 else kv_cache.dtype
566 )
567
568 if kv_cache.dtype != kv_cache_torch_dtype and is_quantized_kv_cache(kv_cache_dtype):
569 # to avoid erounous implicit cast in triton kernel (tl.store to uint8)
570 # (e.g. explicit cast to fp8e4m3fnuz is not supported in triton 3.4)
571 kv_cache = kv_cache.view(kv_cache_torch_dtype)
572 FP8_KV_CACHE = is_quantized_kv_cache(kv_cache_dtype)
573 # heuristics instead of autotuning
574 TILE_SIZE = max(head_size_k, head_size_v)
575 TILE_SIZE = triton.next_power_of_2(TILE_SIZE)
576 if current_platform.is_rocm() or current_platform.is_xpu():
577 num_stages = 4
578 num_warps = 8
579 else: # cuda
580 num_stages = 10
581 num_warps = 16
582
583 # TODO(ngl): maybe replace with static launch grid to avoid overhead if
584 # using cudagraphs
585 grid = lambda meta: (
586 slot_mapping.shape[0],
587 num_heads,
588 )
589
590 reshape_and_cache_kernel_flash_diffkv[grid](
591 key_ptr=key,
592 value_ptr=value,
593 kv_cache_ptr=kv_cache,
594 slot_mapping_ptr=slot_mapping,
595 k_scale=k_scale,
596 v_scale=v_scale,
597 # strides
598 key_stride=k_stride,
599 value_stride=v_stride,
600 block_stride=block_stride,
601 page_stride=page_stride,
602 head_stride=head_stride,
603 num_heads=num_heads,
604 head_size_k=head_size_k,
605 head_size_v=head_size_v,
606 block_size=block_size,
607 # FP8 flags
608 FP8_KV_CACHE=FP8_KV_CACHE,
609 # autotune parameters
610 TILE_SIZE=TILE_SIZE,
611 num_warps=num_warps,
612 num_stages=num_stages,
613 )