| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import Any |
| 6 | |
| 7 | import cutlass |
| 8 | import cutlass.cute as cute |
| 9 | import torch |
| 10 | from cuda.bindings.driver import CUstream |
| 11 | from cutlass import BFloat16, Int32, Uint8, Uint32 |
| 12 | from cutlass.cute.nvgpu import cpasync |
| 13 | from quack.compile_utils import make_fake_tensor |
| 14 | |
| 15 | from vllm.cute_utils import _bf16x2_mul, cvt |
| 16 | from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs |
| 17 | from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import ( |
| 18 | CuTeDSLLaunchSpec, |
| 19 | VllmCuTeDSLJitKernel, |
| 20 | ) |
| 21 | |
| 22 | |
| 23 | class DequantGatherKCacheKernel( |
| 24 | VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"] |
| 25 | ): |
| 26 | # DSv4 geometry. The record format (how much of the row is quantized and |
| 27 | # how wide a scale group is) comes from the compile key: V4 quantizes the |
| 28 | # 448 NoPE dims in groups of 64 and keeps a bf16 RoPE tail, V4.1 quantizes |
| 29 | # all 512 dims in groups of 32 (MXFP8) and has no tail. |
| 30 | head_dim = 512 |
| 31 | num_warps = 4 |
| 32 | tb_size = num_warps * 32 |
| 33 | num_stages = 4 |
| 34 | |
| 35 | @dataclass(frozen=True) |
| 36 | class CompileKey: |
| 37 | block_size: int |
| 38 | has_gather_lens: bool |
| 39 | fp8_dim: int = 448 |
| 40 | group_size: int = 64 # dims per scale |
| 41 | |
| 42 | @staticmethod |
| 43 | def kernel(compile_key: CompileKey) -> Any: |
| 44 | head_dim = DequantGatherKCacheKernel.head_dim |
| 45 | fp8_dim = compile_key.fp8_dim |
| 46 | group_size = compile_key.group_size |
| 47 | num_warps = DequantGatherKCacheKernel.num_warps |
| 48 | tb_size = DequantGatherKCacheKernel.tb_size |
| 49 | num_stages = DequantGatherKCacheKernel.num_stages |
| 50 | bf16_dim = head_dim - fp8_dim |
| 51 | data_dim = fp8_dim + bf16_dim * 2 |
| 52 | scale_bytes = head_dim // group_size |
| 53 | block_size = compile_key.block_size |
| 54 | |
| 55 | @cute.jit |
| 56 | def load_g2s( |
| 57 | k_data_slice: cute.Tensor, |
| 58 | k_scale: cute.Tensor, |
| 59 | block_table: cute.Tensor, |
| 60 | s_kdata_slice: cute.Tensor, |
| 61 | s_kscale: cute.Tensor, |
| 62 | req_id, |
| 63 | pos, |
| 64 | lane_id, |
| 65 | stage_id, |
| 66 | ): |
| 67 | op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) |
| 68 | cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) |
| 69 | cp_scale_atom = cute.make_copy_atom( |
| 70 | cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=scale_bytes * 8 |
| 71 | ) |
| 72 | page_id = block_table[req_id, pos // block_size] |
| 73 | block_offset = pos % block_size |
| 74 | |
| 75 | # Load the first 512 bytes (32x16B). |
| 76 | idx = lane_id |
| 77 | src = k_data_slice[page_id, block_offset, (None, idx)] |
| 78 | cute.copy( |
| 79 | cp16_atom, |
| 80 | cute.recast_tensor(src, Uint32), |
| 81 | s_kdata_slice[(None, idx), stage_id], |
| 82 | ) |
| 83 | |
| 84 | # Load the tail 64 bytes. |
| 85 | idx += 32 |
| 86 | if idx < cutlass.const_expr(data_dim // 16): |
| 87 | src = k_data_slice[page_id, block_offset, (None, idx)] |
| 88 | cute.copy( |
| 89 | cp16_atom, |
| 90 | cute.recast_tensor(src, Uint32), |
| 91 | s_kdata_slice[(None, idx), stage_id], |
| 92 | ) |
| 93 | elif idx == cutlass.const_expr(data_dim // 16): |
| 94 | cute.copy( |
| 95 | cp_scale_atom, |
| 96 | k_scale[page_id, block_offset, None], |
| 97 | s_kscale[None, stage_id], |
| 98 | ) |
| 99 | |
| 100 | @cute.kernel |
| 101 | def device_kernel( |
| 102 | out: cute.Tensor, |
| 103 | k_data: cute.Tensor, |
| 104 | k_scale: cute.Tensor, |
| 105 | seq_lens: cute.Tensor, |
| 106 | gather_lens: cute.Tensor | None, |
| 107 | block_table: cute.Tensor, |
| 108 | offset: Int32, |
| 109 | ): |
| 110 | req_id, worker_id, _ = cute.arch.block_idx() |
| 111 | tid, _, _ = cute.arch.thread_idx() |
| 112 | warp_id = cute.arch.make_warp_uniform(tid // 32) |
| 113 | lane_id = tid % 32 |
| 114 | |
| 115 | _, num_workers, _ = cute.arch.grid_dim() |
| 116 | |
| 117 | # Prepare smem. |
| 118 | smem = cutlass.utils.SmemAllocator() |
| 119 | s_kdata = smem.allocate_tensor( |
| 120 | Uint32, |
| 121 | cute.make_layout((data_dim // 4, num_warps, num_stages)), |
| 122 | byte_alignment=16, |
| 123 | )[None, warp_id, None] |
| 124 | s_kscale = smem.allocate_tensor( |
| 125 | Uint8, |
| 126 | cute.make_layout((scale_bytes, num_warps, num_stages)), |
| 127 | byte_alignment=scale_bytes, |
| 128 | )[None, warp_id, None] |
| 129 | |
| 130 | # Prepare for 16B cp.async, also for BF16 smem loads later. |
| 131 | k_data_slice = cute.logical_divide(k_data, (None, None, 16)) |
| 132 | s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None)) |
| 133 | |
| 134 | # Load FP8 elems in 8B units, so once dequantized, they are 16B units. |
| 135 | s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None)) |
| 136 | |
| 137 | # 16B st.global. |
| 138 | out_slice = cute.logical_divide(out, (None, None, 8)) |
| 139 | |
| 140 | cp_op = cute.nvgpu.CopyUniversalOp() |
| 141 | cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64) |
| 142 | cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128) |
| 143 | |
| 144 | seq_len = seq_lens[req_id] |
| 145 | gather_len = seq_len |
| 146 | if cutlass.const_expr(gather_lens is not None): |
| 147 | gather_len = gather_lens[req_id] # type: ignore[index] |
| 148 | start_pos = seq_len - gather_len |
| 149 | |
| 150 | # Start prefetch. |
| 151 | for i in cutlass.range_constexpr(num_stages - 1): |
| 152 | next_pos = ( |
| 153 | start_pos |
| 154 | + worker_id * num_warps |
| 155 | + warp_id |
| 156 | + i * num_workers * num_warps |
| 157 | ) |
| 158 | if next_pos < seq_len: |
| 159 | load_g2s( |
| 160 | k_data_slice, |
| 161 | k_scale, |
| 162 | block_table, |
| 163 | s_kdata_16B_slice, |
| 164 | s_kscale, |
| 165 | req_id, |
| 166 | next_pos, |
| 167 | lane_id, |
| 168 | i, |
| 169 | ) |
| 170 | cute.arch.cp_async_commit_group() |
| 171 | prefetch_stage = num_stages - 1 |
| 172 | compute_stage = 0 |
| 173 | |
| 174 | # Main loop. |
| 175 | for i in range( |
| 176 | worker_id * num_warps + warp_id, |
| 177 | gather_len, |
| 178 | num_workers * num_warps, |
| 179 | ): |
| 180 | pos = start_pos + i |
| 181 | |
| 182 | # Prefetch next stage. |
| 183 | next_pos = pos + num_workers * num_warps * (num_stages - 1) |
| 184 | if next_pos < seq_len: |
| 185 | load_g2s( |
| 186 | k_data_slice, |
| 187 | k_scale, |
| 188 | block_table, |
| 189 | s_kdata_16B_slice, |
| 190 | s_kscale, |
| 191 | req_id, |
| 192 | next_pos, |
| 193 | lane_id, |
| 194 | prefetch_stage, |
| 195 | ) |
| 196 | prefetch_stage = (prefetch_stage + 1) % num_stages |
| 197 | cute.arch.cp_async_commit_group() |
| 198 | |
| 199 | # Wait for gmem->smem to finish. |
| 200 | cute.arch.cp_async_wait_group(num_stages - 1) |
| 201 | cute.arch.sync_warp() |
| 202 | |
| 203 | # There are 512 elems per token. As a warp, data0 holds the first |
| 204 | # 256 elems and data1 holds the second 256 elems, i.e. each thread |
| 205 | # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as |
| 206 | # contiguous 16B global stores. On Blackwell, this might not be |
| 207 | # necessary as we have 32B global stores, but doing it this way |
| 208 | # does not seem to be slower. |
| 209 | data0 = cute.make_rmem_tensor((2,), Uint32) |
| 210 | data1 = cute.make_rmem_tensor((2,), Uint32) |
| 211 | cute.copy( |
| 212 | cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0 |
| 213 | ) |
| 214 | cute.copy( |
| 215 | cp8_atom, |
| 216 | s_kdata_8B_slice[(None, lane_id + 32), compute_stage], |
| 217 | data1, |
| 218 | ) |
| 219 | |
| 220 | # Convert to bf16x2 via bit manipulation. FP8 scales are per 64 |
| 221 | # elements. An 8-element chunk advances the scale index by |
| 222 | # chunk_id * 8 // group_size. |
| 223 | scale0_u32 = Uint32(s_kscale[lane_id * 8 // group_size, compute_stage]) |
| 224 | scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7)) |
| 225 | scale1_u32 = Uint32( |
| 226 | s_kscale[(lane_id + 32) * 8 // group_size, compute_stage] |
| 227 | ) |
| 228 | scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7)) |
| 229 | |
| 230 | # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2 |
| 231 | # (CUDA 13.2). |
| 232 | dequant0 = cute.make_rmem_tensor(4, Uint32) |
| 233 | dequant1 = cute.make_rmem_tensor(4, Uint32) |
| 234 | for j in cutlass.range_constexpr(2): |
| 235 | tmp0 = cvt.fp8x4_to_bf16x4(data0[j]) |
| 236 | tmp1 = cvt.fp8x4_to_bf16x4(data1[j]) |
| 237 | |
| 238 | # BF16 multiply is safe because the scales are exact powers of 2. |
| 239 | dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2) |
| 240 | dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2) |
| 241 | dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2) |
| 242 | dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2) |
| 243 | |
| 244 | # Last 64 elems are BF16 tail, corresponds to dequant1 of last |
| 245 | # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + |
| 246 | # 8x 16B for BF16. The V4.1 record is fully quantized and has |
| 247 | # no tail. |
| 248 | if cutlass.const_expr(bf16_dim > 0): # noqa: SIM102 |
| 249 | if lane_id + 32 >= fp8_dim // 8: |
| 250 | idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 |
| 251 | cute.copy( |
| 252 | cp16_atom, |
| 253 | s_kdata_16B_slice[(None, idx), compute_stage], |
| 254 | dequant1, |
| 255 | ) |
| 256 | |
| 257 | # Store two 16B BF16 chunks per lane: first half, then second half. |
| 258 | dst = out_slice[req_id, offset + i, (None, lane_id)] |
| 259 | cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32)) |
| 260 | |
| 261 | dst = out_slice[req_id, offset + i, (None, lane_id + 32)] |
| 262 | cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32)) |
| 263 | |
| 264 | compute_stage = (compute_stage + 1) % num_stages |
| 265 | |
| 266 | @cute.jit |
| 267 | def host_entrypoint( |
| 268 | out: cute.Tensor, |
| 269 | k_cache: cute.Tensor, |
| 270 | seq_lens: cute.Tensor, |
| 271 | gather_lens: cute.Tensor | None, |
| 272 | block_table: cute.Tensor, |
| 273 | offset: Int32, |
| 274 | stream: CUstream, |
| 275 | ): |
| 276 | # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] |
| 277 | # block is actually a concat of |
| 278 | # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes]. |
| 279 | k_data = cute.make_tensor( |
| 280 | k_cache.iterator, |
| 281 | layout=cute.make_layout( |
| 282 | (k_cache.shape[0], block_size, data_dim), |
| 283 | stride=(k_cache.stride[0], data_dim, 1), |
| 284 | ), |
| 285 | ) |
| 286 | # k_data_slice: [num_blocks, block_size, (16, data_dim/16)] |
| 287 | # s_kdata_slice: [(4, data_dim/16), num_stages] |
| 288 | k_scale = cute.make_tensor( |
| 289 | k_cache.iterator + (block_size * data_dim), |
| 290 | layout=cute.make_layout( |
| 291 | (k_cache.shape[0], block_size, scale_bytes), |
| 292 | stride=(k_cache.stride[0], scale_bytes, 1), |
| 293 | ), |
| 294 | ) |
| 295 | |
| 296 | num_reqs = out.shape[0] |
| 297 | max_rows = out.shape[1] - offset |
| 298 | num_workers = cutlass.max( |
| 299 | 1, |
| 300 | cutlass.min(cute.ceil_div(max_rows, 4), cute.ceil_div(8192, num_reqs)), |
| 301 | ) |
| 302 | grid = (num_reqs, num_workers, 1) |
| 303 | device_kernel( |
| 304 | out, |
| 305 | k_data, |
| 306 | k_scale, |
| 307 | seq_lens, |
| 308 | gather_lens, |
| 309 | block_table, |
| 310 | offset, |
| 311 | ).launch(grid=grid, block=(tb_size, 1, 1), stream=stream) |
| 312 | |
| 313 | return host_entrypoint |
| 314 | |
| 315 | def dispatch( # type: ignore[override] |
| 316 | self, |
| 317 | *, |
| 318 | block_size: int, |
| 319 | has_gather_lens: bool, |
| 320 | fp8_dim: int = 448, |
| 321 | group_size: int = 64, |
| 322 | ) -> CompileKey: |
| 323 | return self.CompileKey( |
| 324 | block_size=block_size, |
| 325 | has_gather_lens=has_gather_lens, |
| 326 | fp8_dim=fp8_dim, |
| 327 | group_size=group_size, |
| 328 | ) |
| 329 | |
| 330 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 331 | block_size = vllm_config.cache_config.block_size |
| 332 | compress_ratios = vllm_config.model_config.hf_config.compress_ratios |
| 333 | return self._trace_dispatch(self.dispatch)( |
| 334 | zip_inputs( |
| 335 | dict( |
| 336 | block_size=block_size, |
| 337 | has_gather_lens=True, |
| 338 | enabled=True, |
| 339 | ), |
| 340 | dict( |
| 341 | block_size=max(1, block_size // 4), |
| 342 | has_gather_lens=False, |
| 343 | enabled=4 in compress_ratios, |
| 344 | ), |
| 345 | dict( |
| 346 | block_size=max(1, block_size // 128), |
| 347 | has_gather_lens=False, |
| 348 | enabled=128 in compress_ratios, |
| 349 | ), |
| 350 | ), |
| 351 | valid_block_size=block_size > 0, |
| 352 | _when=lambda *, enabled, valid_block_size: enabled and valid_block_size, |
| 353 | ) |
| 354 | |
| 355 | def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: |
| 356 | num_reqs = cute.sym_int() |
| 357 | head_dim = self.head_dim |
| 358 | fp8_dim = compile_key.fp8_dim |
| 359 | block_size = compile_key.block_size |
| 360 | head_bytes = ( |
| 361 | fp8_dim + (head_dim - fp8_dim) * 2 + head_dim // compile_key.group_size |
| 362 | ) |
| 363 | |
| 364 | out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) |
| 365 | k_cache = cute.runtime.make_fake_tensor( |
| 366 | Uint8, |
| 367 | (cute.sym_int(), block_size, head_bytes), |
| 368 | stride=(cute.sym_int64(divisibility=32), head_bytes, 1), |
| 369 | assumed_align=32, |
| 370 | ) |
| 371 | seq_lens = make_fake_tensor(Int32, (num_reqs,)) |
| 372 | gather_lens = ( |
| 373 | make_fake_tensor(Int32, (num_reqs,)) |
| 374 | if compile_key.has_gather_lens |
| 375 | else None |
| 376 | ) |
| 377 | block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int())) |
| 378 | |
| 379 | return ( |
| 380 | out, |
| 381 | k_cache, |
| 382 | seq_lens, |
| 383 | gather_lens, |
| 384 | block_table, |
| 385 | Int32(0), |
| 386 | ) |
| 387 | |
| 388 | @kernel_launcher |
| 389 | def __call__( |
| 390 | self, |
| 391 | *, |
| 392 | out: torch.Tensor, |
| 393 | k_cache: torch.Tensor, |
| 394 | seq_lens: torch.Tensor, |
| 395 | gather_lens: torch.Tensor | None, |
| 396 | block_table: torch.Tensor, |
| 397 | block_size: int, |
| 398 | offset: int, |
| 399 | fp8_dim: int = 448, |
| 400 | group_size: int = 64, |
| 401 | ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]: |
| 402 | compile_key = self.dispatch( |
| 403 | block_size=block_size, |
| 404 | has_gather_lens=gather_lens is not None, |
| 405 | fp8_dim=fp8_dim, |
| 406 | group_size=group_size, |
| 407 | ) |
| 408 | launch_args = ( |
| 409 | out, |
| 410 | k_cache, |
| 411 | seq_lens, |
| 412 | gather_lens, |
| 413 | block_table, |
| 414 | offset, |
| 415 | ) |
| 416 | return compile_key, launch_args |
| 417 | |
| 418 | |
| 419 | _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL = DequantGatherKCacheKernel() |