| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import random |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck |
| 10 | from vllm import _custom_ops as ops |
| 11 | from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.torch_utils import nvfp4_split_data_scale, set_random_seed |
| 14 | |
| 15 | pytestmark = pytest.mark.skip_global_cleanup |
| 16 | |
| 17 | COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")] |
| 18 | DTYPES = [torch.bfloat16, torch.float] |
| 19 | NUM_TOKENS = [42] # Arbitrary values for testing |
| 20 | NUM_LAYERS = [1] # Arbitrary values for testing |
| 21 | NUM_HEADS = [8] # Arbitrary values for testing |
| 22 | HEAD_SIZES = [64, 80, 256] |
| 23 | BLOCK_SIZES = [8, 16, 32] |
| 24 | CACHE_LAYOUTS = ["NHD", "HND"] |
| 25 | KV_SCALE_TYPES = ["tensor", "attn_head"] |
| 26 | |
| 27 | # Parameters for MLA tests. |
| 28 | KV_LORA_RANKS = [256, 512] |
| 29 | QK_ROPE_HEAD_DIMS = [64] |
| 30 | NUM_TOKENS_MLA = [42] |
| 31 | BLOCK_SIZES_MLA = [16] |
| 32 | NUM_BLOCKS_MLA = [8] |
| 33 | |
| 34 | # Arbitrary values for testing |
| 35 | # don't make it too large. e.g. [1024, 36000] will OOM |
| 36 | NUM_BLOCKS = [1024, 10000] |
| 37 | |
| 38 | NUM_MAPPINGS = [256] # Arbitrary values for testing |
| 39 | SEEDS = [0] |
| 40 | CUDA_DEVICES = [ |
| 41 | f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2) |
| 42 | ] |
| 43 | |
| 44 | # We assume fp8 is always enabled for testing. |
| 45 | KV_CACHE_DTYPE = ["auto", "fp8"] |
| 46 | |
| 47 | RESHAPE_FLASH_IMPLEMENTATIONS = ["cuda", "triton"] |
| 48 | |
| 49 | |
| 50 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS) |
| 51 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 52 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 53 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 54 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS) |
| 55 | @pytest.mark.parametrize("dtype", DTYPES) |
| 56 | @pytest.mark.parametrize("seed", SEEDS) |
| 57 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 58 | @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) |
| 59 | @torch.inference_mode() |
| 60 | def test_reshape_and_cache( |
| 61 | kv_cache_factory, |
| 62 | num_tokens: int, |
| 63 | num_heads: int, |
| 64 | head_size: int, |
| 65 | block_size: int, |
| 66 | num_blocks: int, |
| 67 | dtype: torch.dtype, |
| 68 | seed: int, |
| 69 | device: str, |
| 70 | kv_cache_dtype: str, |
| 71 | ) -> None: |
| 72 | if kv_cache_dtype == "fp8" and head_size % 16: |
| 73 | pytest.skip() |
| 74 | set_random_seed(seed) |
| 75 | torch.set_default_device(device) |
| 76 | torch.accelerator.set_device_index(device) |
| 77 | # Create a random slot mapping. |
| 78 | num_slots = block_size * num_blocks |
| 79 | slot_mapping_lst = random.sample(range(num_slots), num_tokens) |
| 80 | slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long) |
| 81 | |
| 82 | qkv = torch.randn(num_tokens, 3, num_heads, head_size, dtype=dtype) |
| 83 | _, key, value = qkv.unbind(dim=1) |
| 84 | |
| 85 | # Create the KV caches. |
| 86 | key_caches, value_caches = kv_cache_factory( |
| 87 | num_blocks, |
| 88 | block_size, |
| 89 | 1, |
| 90 | num_heads, |
| 91 | head_size, |
| 92 | kv_cache_dtype, |
| 93 | dtype, |
| 94 | seed, |
| 95 | device, |
| 96 | ) |
| 97 | key_cache, value_cache = key_caches[0], value_caches[0] |
| 98 | |
| 99 | # Using default kv_scale |
| 100 | k_scale = (key.amax() / 64.0).to(torch.float32) |
| 101 | v_scale = (value.amax() / 64.0).to(torch.float32) |
| 102 | |
| 103 | # Clone the KV caches. |
| 104 | if kv_cache_dtype == "fp8": |
| 105 | cloned_key_cache = torch.empty_like(key_cache, dtype=torch.float16) |
| 106 | ops.convert_fp8(cloned_key_cache, key_cache, k_scale.item()) |
| 107 | cloned_value_cache = torch.empty_like(value_cache, dtype=torch.float16) |
| 108 | ops.convert_fp8(cloned_value_cache, value_cache, v_scale.item()) |
| 109 | else: |
| 110 | cloned_key_cache = key_cache.clone() |
| 111 | cloned_value_cache = value_cache.clone() |
| 112 | |
| 113 | # Call the reshape_and_cache kernel. |
| 114 | opcheck( |
| 115 | torch.ops._C_cache_ops.reshape_and_cache, |
| 116 | ( |
| 117 | key, |
| 118 | value, |
| 119 | key_cache, |
| 120 | value_cache, |
| 121 | slot_mapping, |
| 122 | kv_cache_dtype, |
| 123 | k_scale, |
| 124 | v_scale, |
| 125 | ), |
| 126 | cond=(head_size == HEAD_SIZES[0]), |
| 127 | ) |
| 128 | ops.reshape_and_cache( |
| 129 | key, |
| 130 | value, |
| 131 | key_cache, |
| 132 | value_cache, |
| 133 | slot_mapping, |
| 134 | kv_cache_dtype, |
| 135 | k_scale, |
| 136 | v_scale, |
| 137 | ) |
| 138 | |
| 139 | if kv_cache_dtype == "fp8": |
| 140 | result_key_cache = torch.empty_like(key_cache, dtype=torch.float16) |
| 141 | ops.convert_fp8(result_key_cache, key_cache, k_scale.item()) |
| 142 | result_value_cache = torch.empty_like(value_cache, dtype=torch.float16) |
| 143 | ops.convert_fp8(result_value_cache, value_cache, v_scale.item()) |
| 144 | |
| 145 | # Run the reference implementation. |
| 146 | reshaped_key = key.reshape(num_tokens, *key_cache[0, :, :, 0, :].shape) |
| 147 | block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor") |
| 148 | block_indices_lst = block_indices.cpu().tolist() |
| 149 | block_offsets = slot_mapping % block_size |
| 150 | block_offsets_lst = block_offsets.cpu().tolist() |
| 151 | for i in range(num_tokens): |
| 152 | block_idx = block_indices_lst[i] |
| 153 | block_offset = block_offsets_lst[i] |
| 154 | cloned_key_cache[block_idx, :, :, block_offset, :] = reshaped_key[i] |
| 155 | cloned_value_cache[block_idx, :, :, block_offset] = value[i] |
| 156 | |
| 157 | if kv_cache_dtype == "fp8": |
| 158 | torch.testing.assert_close( |
| 159 | result_key_cache, cloned_key_cache, atol=0.001, rtol=0.1 |
| 160 | ) |
| 161 | torch.testing.assert_close( |
| 162 | result_value_cache, cloned_value_cache, atol=0.001, rtol=0.1 |
| 163 | ) |
| 164 | else: |
| 165 | torch.testing.assert_close(key_cache, cloned_key_cache) |
| 166 | torch.testing.assert_close(value_cache, cloned_value_cache) |
| 167 | |
| 168 | |
| 169 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS) |
| 170 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 171 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 172 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 173 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS) |
| 174 | @pytest.mark.parametrize("dtype", DTYPES) |
| 175 | @pytest.mark.parametrize("seed", SEEDS) |
| 176 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 177 | @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE + ["nvfp4"]) |
| 178 | @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS) |
| 179 | @pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES) |
| 180 | @pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS) |
| 181 | @torch.inference_mode() |
| 182 | def test_reshape_and_cache_flash( |
| 183 | kv_cache_factory_flashinfer, |
| 184 | num_tokens: int, |
| 185 | num_heads: int, |
| 186 | head_size: int, |
| 187 | block_size: int, |
| 188 | num_blocks: int, |
| 189 | dtype: torch.dtype, |
| 190 | seed: int, |
| 191 | device: str, |
| 192 | kv_cache_dtype: str, |
| 193 | kv_cache_layout: str, |
| 194 | kv_scale_type: str, |
| 195 | implementation: str, |
| 196 | ) -> None: |
| 197 | set_random_seed(seed) |
| 198 | torch.set_default_device(device) |
| 199 | torch.accelerator.set_device_index(device) |
| 200 | assert implementation in ["cuda", "triton"] |
| 201 | if implementation == "triton" and kv_cache_layout == "HND": |
| 202 | pytest.skip("Triton implementation only supports NHD layout.") |
| 203 | |
| 204 | if kv_scale_type == "attn_head" and implementation != "cuda": |
| 205 | pytest.skip("Only CUDA implementation supports attn_head scaling.") |
| 206 | |
| 207 | if kv_cache_dtype == "nvfp4": |
| 208 | if not current_platform.has_device_capability(100): |
| 209 | pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).") |
| 210 | if implementation != "cuda": |
| 211 | pytest.skip("NVFP4 only supports CUDA implementation.") |
| 212 | if kv_scale_type != "tensor": |
| 213 | pytest.skip("NVFP4 only supports per-tensor scaling.") |
| 214 | if head_size % 16 != 0: |
| 215 | pytest.skip("NVFP4 requires head_size divisible by 16.") |
| 216 | if (head_size // 16) % 4 != 0: |
| 217 | pytest.skip( |
| 218 | "NVFP4 requires (head_size // 16) divisible by 4 " |
| 219 | "for 4x4 block scale swizzle." |
| 220 | ) |
| 221 | if block_size % 4 != 0: |
| 222 | pytest.skip("NVFP4 requires block_size divisible by 4.") |
| 223 | if dtype not in (torch.float16, torch.bfloat16): |
| 224 | pytest.skip("NVFP4 quantization only supports fp16/bf16 input.") |
| 225 | |
| 226 | # fp8 conversion requires continugous memory buffer. Reduce the number of |
| 227 | # blocks and tokens to consume less memory. |
| 228 | num_tokens = num_tokens // 2 |
| 229 | num_blocks = num_blocks // 2 |
| 230 | # Create a random slot mapping. |
| 231 | num_slots = block_size * num_blocks |
| 232 | slot_mapping_lst = random.sample(range(num_slots), num_tokens) |
| 233 | slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device) |
| 234 | qkv = torch.randn(num_tokens, 3, num_heads, head_size, dtype=dtype, device=device) |
| 235 | _, key, value = qkv.unbind(dim=1) |
| 236 | |
| 237 | # Create the KV caches. |
| 238 | key_caches, value_caches = kv_cache_factory_flashinfer( |
| 239 | num_blocks, |
| 240 | block_size, |
| 241 | 1, |
| 242 | num_heads, |
| 243 | head_size, |
| 244 | kv_cache_dtype, |
| 245 | dtype, |
| 246 | device=device, |
| 247 | cache_layout=kv_cache_layout, |
| 248 | ) |
| 249 | key_cache, value_cache = key_caches[0], value_caches[0] |
| 250 | del key_caches |
| 251 | del value_caches |
| 252 | |
| 253 | # For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes. |
| 254 | # Split views are still needed for dequant verification. |
| 255 | key_scale_cache = None |
| 256 | value_scale_cache = None |
| 257 | nvfp4_key_data = None |
| 258 | nvfp4_value_data = None |
| 259 | if kv_cache_dtype == "nvfp4": |
| 260 | nvfp4_key_data, key_scale_cache = nvfp4_split_data_scale(key_cache) |
| 261 | nvfp4_value_data, value_scale_cache = nvfp4_split_data_scale(value_cache) |
| 262 | |
| 263 | if kv_cache_dtype == "nvfp4": |
| 264 | # Global scale = amax / 448 (per-tensor) |
| 265 | k_scale = (key.abs().amax() / 448.0).to(torch.float32) |
| 266 | v_scale = (value.abs().amax() / 448.0).to(torch.float32) |
| 267 | elif kv_scale_type == "tensor": |
| 268 | k_scale = (key.amax() / 64.0).to(torch.float32) |
| 269 | v_scale = (value.amax() / 64.0).to(torch.float32) |
| 270 | else: # "attn_head" |
| 271 | k_scale = (key.amax(dim=(0, 2)) / 64.0).to(torch.float32) |
| 272 | v_scale = (value.amax(dim=(0, 2)) / 64.0).to(torch.float32) |
| 273 | |
| 274 | def permute_and_compact(x): |
| 275 | y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3) |
| 276 | return y.contiguous() |
| 277 | |
| 278 | if kv_cache_dtype != "nvfp4": |
| 279 | key_cache_compact = permute_and_compact(key_cache) |
| 280 | value_cache_compact = permute_and_compact(value_cache) |
| 281 | |
| 282 | def convert_fp8_local(output, input, scale, kv_dtype): |
| 283 | fp8_input = input.view(current_platform.fp8_dtype()) |
| 284 | if scale.numel() == 1: # per-tensor |
| 285 | result = scaled_dequantize( |
| 286 | fp8_input.flatten(0, 2), scale, group_shape=None, out_dtype=output.dtype |
| 287 | ).reshape(*input.shape) |
| 288 | else: # per-head: broadcast scale along the head dimension |
| 289 | # Original code uses dim 2 for NHD, dim 1 for HND |
| 290 | if kv_cache_layout == "NHD": |
| 291 | result = fp8_input.to(output.dtype) * scale.view(1, 1, -1, 1) |
| 292 | else: |
| 293 | result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1) |
| 294 | output.copy_(result) |
| 295 | |
| 296 | # Clone the KV caches (for non-nvfp4, used as reference baseline). |
| 297 | if kv_cache_dtype == "fp8": |
| 298 | cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16) |
| 299 | convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype) |
| 300 | cloned_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16) |
| 301 | convert_fp8_local( |
| 302 | cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype |
| 303 | ) |
| 304 | elif kv_cache_dtype != "nvfp4": |
| 305 | cloned_key_cache = key_cache_compact.clone() |
| 306 | cloned_value_cache = value_cache_compact.clone() |
| 307 | |
| 308 | # Call the reshape_and_cache kernel. |
| 309 | if implementation == "cuda": |
| 310 | if kv_cache_dtype != "nvfp4": |
| 311 | opcheck( |
| 312 | torch.ops._C_cache_ops.reshape_and_cache_flash, |
| 313 | ( |
| 314 | key, |
| 315 | value, |
| 316 | key_cache, |
| 317 | value_cache, |
| 318 | slot_mapping, |
| 319 | kv_cache_dtype, |
| 320 | k_scale, |
| 321 | v_scale, |
| 322 | ), |
| 323 | cond=(head_size == HEAD_SIZES[0]), |
| 324 | ) |
| 325 | ops.reshape_and_cache_flash( |
| 326 | key, |
| 327 | value, |
| 328 | key_cache, |
| 329 | value_cache, |
| 330 | slot_mapping, |
| 331 | kv_cache_dtype, |
| 332 | k_scale, |
| 333 | v_scale, |
| 334 | ) |
| 335 | elif implementation == "triton": |
| 336 | from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( |
| 337 | triton_reshape_and_cache_flash, |
| 338 | ) |
| 339 | |
| 340 | triton_reshape_and_cache_flash( |
| 341 | key, |
| 342 | value, |
| 343 | key_cache, |
| 344 | value_cache, |
| 345 | slot_mapping, |
| 346 | kv_cache_dtype, |
| 347 | k_scale, |
| 348 | v_scale, |
| 349 | ) |
| 350 | |
| 351 | if kv_cache_dtype == "nvfp4": |
| 352 | # Verify NVFP4 by dequantizing the entire cache and comparing |
| 353 | # the written positions against original bf16 values. |
| 354 | # Same pattern as FP8: dequant whole cache, then extract and compare. |
| 355 | from tests.kernels.quantization.nvfp4_utils import ( |
| 356 | dequant_nvfp4_kv_cache, |
| 357 | ) |
| 358 | |
| 359 | def dequant_nvfp4_cache(data_cache, scale_cache, global_scale, swizzled_scales): |
| 360 | # data_cache: [B, N, H, data_dim] logical view (layout in strides) |
| 361 | # scale_cache: [B, N, H, scale_dim] logical view (layout in strides) |
| 362 | # Permute to [B, H, N, dim] for the dequant utility. |
| 363 | data_hnd = data_cache.permute(0, 2, 1, 3) |
| 364 | scale_hnd = scale_cache.permute(0, 2, 1, 3) |
| 365 | result_hnd = dequant_nvfp4_kv_cache( |
| 366 | data_hnd, |
| 367 | scale_hnd, |
| 368 | global_scale, |
| 369 | head_size, |
| 370 | block_size, |
| 371 | swizzled_scales=swizzled_scales, |
| 372 | ) |
| 373 | return result_hnd.permute(0, 2, 1, 3) # back to [B, N, H, dim] |
| 374 | |
| 375 | # The kernel writes K scales linearly on every arch and V scales in |
| 376 | # the 4x4 swizzle of the SM100 trtllm-gen reader. The FlashInfer |
| 377 | # reader used on SM12x expects linear V scales as well. |
| 378 | v_scales_swizzled = not current_platform.is_device_capability_family(120) |
| 379 | result_key_cache = dequant_nvfp4_cache( |
| 380 | nvfp4_key_data, key_scale_cache, k_scale.item(), swizzled_scales=False |
| 381 | ) |
| 382 | result_value_cache = dequant_nvfp4_cache( |
| 383 | nvfp4_value_data, |
| 384 | value_scale_cache, |
| 385 | v_scale.item(), |
| 386 | swizzled_scales=v_scales_swizzled, |
| 387 | ) |
| 388 | |
| 389 | # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping. |
| 390 | num_slots = num_blocks * block_size |
| 391 | result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size) |
| 392 | result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size) |
| 393 | |
| 394 | torch.testing.assert_close( |
| 395 | result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5 |
| 396 | ) |
| 397 | torch.testing.assert_close( |
| 398 | result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5 |
| 399 | ) |
| 400 | return |
| 401 | |
| 402 | key_cache_compact = permute_and_compact(key_cache) |
| 403 | value_cache_compact = permute_and_compact(value_cache) |
| 404 | |
| 405 | if kv_cache_dtype == "fp8": |
| 406 | result_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16) |
| 407 | convert_fp8_local(result_key_cache, key_cache_compact, k_scale, kv_cache_dtype) |
| 408 | result_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16) |
| 409 | convert_fp8_local( |
| 410 | result_value_cache, |
| 411 | value_cache_compact, |
| 412 | v_scale, |
| 413 | kv_cache_dtype, |
| 414 | ) |
| 415 | |
| 416 | # Run the reference implementation. |
| 417 | block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor") |
| 418 | block_indices_lst = block_indices.cpu().tolist() |
| 419 | block_offsets = slot_mapping % block_size |
| 420 | block_offsets_lst = block_offsets.cpu().tolist() |
| 421 | for i in range(num_tokens): |
| 422 | block_idx = block_indices_lst[i] |
| 423 | block_offset = block_offsets_lst[i] |
| 424 | if kv_cache_layout == "NHD": |
| 425 | cloned_key_cache[block_idx, block_offset, :, :] = key[i] |
| 426 | cloned_value_cache[block_idx, block_offset, :, :] = value[i] |
| 427 | else: |
| 428 | cloned_key_cache[block_idx, :, block_offset, :] = key[i] |
| 429 | cloned_value_cache[block_idx, :, block_offset, :] = value[i] |
| 430 | |
| 431 | if kv_cache_dtype == "fp8": |
| 432 | torch.testing.assert_close( |
| 433 | result_key_cache, cloned_key_cache, atol=0.001, rtol=0.1 |
| 434 | ) |
| 435 | torch.testing.assert_close( |
| 436 | result_value_cache, cloned_value_cache, atol=0.001, rtol=0.1 |
| 437 | ) |
| 438 | else: |
| 439 | torch.testing.assert_close(key_cache_compact, cloned_key_cache) |
| 440 | torch.testing.assert_close(value_cache_compact, cloned_value_cache) |
| 441 | |
| 442 | |
| 443 | @torch.inference_mode() |
| 444 | def test_nvfp4_4over6_selects_lower_error_scale( |
| 445 | kv_cache_factory_flashinfer, |
| 446 | ) -> None: |
| 447 | if not current_platform.has_device_capability(100): |
| 448 | pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).") |
| 449 | |
| 450 | device = CUDA_DEVICES[0] |
| 451 | block_size = 16 |
| 452 | head_size = 64 |
| 453 | source = torch.tensor( |
| 454 | [6.0] + [4.5] * 15, |
| 455 | dtype=torch.bfloat16, |
| 456 | device=device, |
| 457 | ).repeat(1, 1, head_size // 16) |
| 458 | slot_mapping = torch.zeros(1, dtype=torch.long, device=device) |
| 459 | scale = torch.ones(1, dtype=torch.float32, device=device) |
| 460 | |
| 461 | def quantize_and_dequantize(cache_dtype: str) -> torch.Tensor: |
| 462 | key_caches, value_caches = kv_cache_factory_flashinfer( |
| 463 | 1, |
| 464 | block_size, |
| 465 | 1, |
| 466 | 1, |
| 467 | head_size, |
| 468 | cache_dtype, |
| 469 | source.dtype, |
| 470 | device=device, |
| 471 | cache_layout="NHD", |
| 472 | ) |
| 473 | key_cache, value_cache = key_caches[0], value_caches[0] |
| 474 | key_cache.zero_() |
| 475 | value_cache.zero_() |
| 476 | |
| 477 | ops.reshape_and_cache_flash( |
| 478 | source, |
| 479 | source, |
| 480 | key_cache, |
| 481 | value_cache, |
| 482 | slot_mapping, |
| 483 | cache_dtype, |
| 484 | scale, |
| 485 | scale, |
| 486 | ) |
| 487 | |
| 488 | from tests.kernels.quantization.nvfp4_utils import ( |
| 489 | dequant_nvfp4_kv_cache, |
| 490 | ) |
| 491 | |
| 492 | data, block_scales = nvfp4_split_data_scale(key_cache) |
| 493 | return dequant_nvfp4_kv_cache( |
| 494 | data.permute(0, 2, 1, 3), |
| 495 | block_scales.permute(0, 2, 1, 3), |
| 496 | scale.item(), |
| 497 | head_size, |
| 498 | block_size, |
| 499 | swizzled_scales=False, |
| 500 | )[0, 0, 0] |
| 501 | |
| 502 | default = quantize_and_dequantize("nvfp4") |
| 503 | four_over_six = quantize_and_dequantize("nvfp4_4over6") |
| 504 | source_f32 = source[0, 0].float() |
| 505 | |
| 506 | default_mse = (default - source_f32).square().mean() |
| 507 | four_over_six_mse = (four_over_six - source_f32).square().mean() |
| 508 | assert four_over_six_mse < default_mse |
| 509 | |
| 510 | |
| 511 | @pytest.mark.parametrize("dtype", DTYPES) |
| 512 | @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) |
| 513 | @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS) |
| 514 | @pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS) |
| 515 | @torch.inference_mode() |
| 516 | def test_reshape_and_cache_flash_unaligned_rows( |
| 517 | kv_cache_factory_flashinfer, |
| 518 | dtype: torch.dtype, |
| 519 | kv_cache_dtype: str, |
| 520 | kv_cache_layout: str, |
| 521 | implementation: str, |
| 522 | ) -> None: |
| 523 | """Regression test for https://github.com/vllm-project/vllm/issues/41257. |
| 524 | |
| 525 | head_size=46 with num_heads=13 places KV-cache rows at byte offsets |
| 526 | that are not a multiple of the vector width (NHD row pitch |
| 527 | 13*46*itemsize, HND head pitch 46*itemsize), unlike HEAD_SIZES above |
| 528 | which are all 16-byte multiples. The CUDA kernel used to issue |
| 529 | vectorized stores to those rows -> CUDA misaligned address. |
| 530 | """ |
| 531 | test_reshape_and_cache_flash( |
| 532 | kv_cache_factory_flashinfer, |
| 533 | num_tokens=42, |
| 534 | num_heads=13, |
| 535 | head_size=46, |
| 536 | block_size=16, |
| 537 | num_blocks=128, |
| 538 | dtype=dtype, |
| 539 | seed=0, |
| 540 | device=CUDA_DEVICES[0], |
| 541 | kv_cache_dtype=kv_cache_dtype, |
| 542 | kv_cache_layout=kv_cache_layout, |
| 543 | kv_scale_type="tensor", |
| 544 | implementation=implementation, |
| 545 | ) |
| 546 | |
| 547 | |
| 548 | @pytest.mark.parametrize("direction", COPYING_DIRECTION) |
| 549 | @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS) |
| 550 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 551 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 552 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 553 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS) |
| 554 | @pytest.mark.parametrize("dtype", DTYPES) |
| 555 | @pytest.mark.parametrize("seed", SEEDS) |
| 556 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 557 | @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) |
| 558 | @torch.inference_mode() |
| 559 | def test_swap_blocks( |
| 560 | kv_cache_factory, |
| 561 | direction: tuple[str, str], |
| 562 | num_mappings: int, |
| 563 | num_heads: int, |
| 564 | head_size: int, |
| 565 | block_size: int, |
| 566 | num_blocks: int, |
| 567 | dtype: torch.dtype, |
| 568 | seed: int, |
| 569 | device: str, |
| 570 | kv_cache_dtype: str, |
| 571 | ) -> None: |
| 572 | if kv_cache_dtype == "fp8" and "cpu" in direction: |
| 573 | pytest.skip() |
| 574 | if kv_cache_dtype == "fp8" and head_size % 16: |
| 575 | pytest.skip() |
| 576 | |
| 577 | set_random_seed(seed) |
| 578 | |
| 579 | src_device = device if direction[0] == "cuda" else "cpu" |
| 580 | dst_device = device if direction[1] == "cuda" else "cpu" |
| 581 | |
| 582 | src_blocks = random.sample(range(num_blocks), num_mappings) |
| 583 | # For the same device, mapping must not overlap |
| 584 | if src_device == dst_device: |
| 585 | remaining_blocks = list(set(range(num_blocks)) - set(src_blocks)) |
| 586 | dst_blocks = random.sample(remaining_blocks, num_mappings) |
| 587 | else: |
| 588 | dst_blocks = random.sample(range(num_blocks), num_mappings) |
| 589 | |
| 590 | block_mapping = list(zip(src_blocks, dst_blocks)) |
| 591 | block_mapping_tensor = torch.tensor( |
| 592 | block_mapping, dtype=torch.int64, device="cpu" |
| 593 | ).view(-1, 2) |
| 594 | |
| 595 | # Create the KV caches on the first device. |
| 596 | src_key_caches, src_value_caches = kv_cache_factory( |
| 597 | num_blocks, |
| 598 | block_size, |
| 599 | 1, |
| 600 | num_heads, |
| 601 | head_size, |
| 602 | kv_cache_dtype, |
| 603 | dtype, |
| 604 | seed, |
| 605 | src_device, |
| 606 | ) |
| 607 | |
| 608 | # Create the KV caches on the second device. |
| 609 | dist_key_caches, dist_value_caches = kv_cache_factory( |
| 610 | num_blocks, |
| 611 | block_size, |
| 612 | 1, |
| 613 | num_heads, |
| 614 | head_size, |
| 615 | kv_cache_dtype, |
| 616 | dtype, |
| 617 | seed, |
| 618 | dst_device, |
| 619 | ) |
| 620 | |
| 621 | src_key_caches_clone = src_key_caches[0].clone() |
| 622 | src_value_caches_clone = src_value_caches[0].clone() |
| 623 | |
| 624 | # Call the swap_blocks kernel. |
| 625 | do_opcheck = head_size == HEAD_SIZES[0] |
| 626 | src_cache = src_key_caches[0] |
| 627 | block_size_in_bytes = src_cache.element_size() * src_cache.stride(0) |
| 628 | opcheck( |
| 629 | torch.ops._C_cache_ops.swap_blocks, |
| 630 | ( |
| 631 | src_key_caches[0], |
| 632 | dist_key_caches[0], |
| 633 | block_size_in_bytes, |
| 634 | block_mapping_tensor, |
| 635 | ), |
| 636 | cond=do_opcheck, |
| 637 | ) |
| 638 | opcheck( |
| 639 | torch.ops._C_cache_ops.swap_blocks, |
| 640 | ( |
| 641 | src_value_caches[0], |
| 642 | dist_value_caches[0], |
| 643 | block_size_in_bytes, |
| 644 | block_mapping_tensor, |
| 645 | ), |
| 646 | cond=do_opcheck, |
| 647 | ) |
| 648 | |
| 649 | ops.swap_blocks( |
| 650 | src_key_caches[0], |
| 651 | dist_key_caches[0], |
| 652 | block_size_in_bytes, |
| 653 | block_mapping_tensor, |
| 654 | ) |
| 655 | ops.swap_blocks( |
| 656 | src_value_caches[0], |
| 657 | dist_value_caches[0], |
| 658 | block_size_in_bytes, |
| 659 | block_mapping_tensor, |
| 660 | ) |
| 661 | |
| 662 | src_indices = block_mapping_tensor[:, 0].to(src_key_caches_clone.device) |
| 663 | dst_indices = block_mapping_tensor[:, 1].to(dist_key_caches[0].device) |
| 664 | |
| 665 | torch.testing.assert_close( |
| 666 | src_key_caches_clone.index_select(0, src_indices).cpu(), |
| 667 | dist_key_caches[0].index_select(0, dst_indices).cpu(), |
| 668 | ) |
| 669 | torch.testing.assert_close( |
| 670 | src_value_caches_clone.index_select(0, src_indices).cpu(), |
| 671 | dist_value_caches[0].index_select(0, dst_indices).cpu(), |
| 672 | ) |
| 673 | |
| 674 | |
| 675 | @pytest.mark.parametrize("num_heads", NUM_HEADS) |
| 676 | @pytest.mark.parametrize("head_size", HEAD_SIZES) |
| 677 | @pytest.mark.parametrize("block_size", BLOCK_SIZES) |
| 678 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS) |
| 679 | @pytest.mark.parametrize("dtype", DTYPES) |
| 680 | @pytest.mark.parametrize("seed", SEEDS) |
| 681 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 682 | @torch.inference_mode() |
| 683 | def test_fp8_e4m3_conversion( |
| 684 | num_heads: int, |
| 685 | head_size: int, |
| 686 | block_size: int, |
| 687 | num_blocks: int, |
| 688 | dtype: torch.dtype, |
| 689 | seed: int, |
| 690 | device: str, |
| 691 | ) -> None: |
| 692 | set_random_seed(seed) |
| 693 | |
| 694 | low = -224.0 |
| 695 | high = 224.0 |
| 696 | shape = (num_blocks, num_heads, head_size, block_size) |
| 697 | cache = torch.empty(shape, dtype=dtype, device=device) |
| 698 | cache.uniform_(low, high) |
| 699 | |
| 700 | cache_fp8 = torch.empty_like(cache, dtype=torch.uint8) |
| 701 | ops.convert_fp8(cache_fp8, cache) |
| 702 | |
| 703 | converted_cache = torch.empty_like(cache) |
| 704 | ops.convert_fp8(converted_cache, cache_fp8) |
| 705 | |
| 706 | torch.testing.assert_close(cache, converted_cache, atol=0.001, rtol=0.1) |
| 707 | |
| 708 | |
| 709 | def _create_mla_cache( |
| 710 | num_blocks: int, |
| 711 | block_size: int, |
| 712 | entry_size: int, |
| 713 | dtype: torch.dtype, |
| 714 | kv_cache_dtype: str, |
| 715 | device: str, |
| 716 | ) -> torch.Tensor: |
| 717 | cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype |
| 718 | return torch.zeros( |
| 719 | num_blocks, block_size, entry_size, dtype=cache_dtype, device=device |
| 720 | ) |
| 721 | |
| 722 | |
| 723 | def _fill_mla_cache(cache: torch.Tensor, kv_cache_dtype: str): |
| 724 | rand_dtype = torch.float16 if kv_cache_dtype == "fp8" else cache.dtype |
| 725 | |
| 726 | vals = torch.randn(*cache.shape, device=cache.device, dtype=rand_dtype) |
| 727 | if kv_cache_dtype == "fp8": |
| 728 | temp = torch.zeros_like(cache) |
| 729 | ops.convert_fp8(temp, vals, 1.0, kv_dtype=kv_cache_dtype) |
| 730 | vals = temp |
| 731 | cache.copy_(vals) |
| 732 | |
| 733 | |
| 734 | @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS) |
| 735 | @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS) |
| 736 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA) |
| 737 | @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA) |
| 738 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA) |
| 739 | @pytest.mark.parametrize("dtype", DTYPES) |
| 740 | @pytest.mark.parametrize("seed", SEEDS) |
| 741 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 742 | @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) |
| 743 | @torch.inference_mode() |
| 744 | def test_concat_and_cache_mla( |
| 745 | kv_lora_rank: int, |
| 746 | qk_rope_head_dim: int, |
| 747 | num_tokens: int, |
| 748 | block_size: int, |
| 749 | num_blocks: int, |
| 750 | dtype: torch.dtype, |
| 751 | seed: int, |
| 752 | device: str, |
| 753 | kv_cache_dtype: str, |
| 754 | ) -> None: |
| 755 | set_random_seed(seed) |
| 756 | torch.set_default_device(device) |
| 757 | torch.accelerator.set_device_index(device) |
| 758 | |
| 759 | total_slots = num_blocks * block_size |
| 760 | slot_mapping_lst = random.sample(range(total_slots), num_tokens) |
| 761 | slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device) |
| 762 | |
| 763 | kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device) |
| 764 | k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device) |
| 765 | entry_size = kv_lora_rank + qk_rope_head_dim |
| 766 | |
| 767 | scale = torch.tensor(0.1, dtype=torch.float32, device=device) |
| 768 | kv_cache = _create_mla_cache( |
| 769 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 770 | ) |
| 771 | ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device) |
| 772 | |
| 773 | for i in range(num_tokens): |
| 774 | slot = slot_mapping[i].item() |
| 775 | block_idx = slot // block_size |
| 776 | block_offset = slot % block_size |
| 777 | ref_temp[block_idx, block_offset, :kv_lora_rank] = kv_c[i] |
| 778 | ref_temp[block_idx, block_offset, kv_lora_rank:] = k_pe[i] |
| 779 | |
| 780 | if kv_cache_dtype == "fp8": |
| 781 | ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype) |
| 782 | ops.convert_fp8(ref_kv_cache, ref_temp, scale.item(), kv_dtype=kv_cache_dtype) |
| 783 | else: |
| 784 | ref_kv_cache = ref_temp |
| 785 | |
| 786 | opcheck( |
| 787 | torch.ops._C_cache_ops.concat_and_cache_mla, |
| 788 | (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale), |
| 789 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 790 | ) |
| 791 | |
| 792 | ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale) |
| 793 | |
| 794 | if kv_cache_dtype == "fp8": |
| 795 | result_temp = torch.empty_like(kv_cache, dtype=torch.float16) |
| 796 | ops.convert_fp8( |
| 797 | result_temp, kv_cache.contiguous(), scale.item(), kv_dtype=kv_cache_dtype |
| 798 | ) |
| 799 | expected_temp = torch.empty_like(ref_kv_cache, dtype=torch.float16) |
| 800 | ops.convert_fp8( |
| 801 | expected_temp, ref_kv_cache, scale.item(), kv_dtype=kv_cache_dtype |
| 802 | ) |
| 803 | torch.testing.assert_close(result_temp, expected_temp, atol=0.001, rtol=0.1) |
| 804 | else: |
| 805 | torch.testing.assert_close(kv_cache, ref_kv_cache) |
| 806 | |
| 807 | |
| 808 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 809 | @pytest.mark.parametrize("shared_slot_mapping", [False, True]) |
| 810 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e5m2"]) |
| 811 | @torch.inference_mode() |
| 812 | def test_concat_and_cache_mla_grouped( |
| 813 | device: str, |
| 814 | shared_slot_mapping: bool, |
| 815 | kv_cache_dtype: str, |
| 816 | ) -> None: |
| 817 | if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm(): |
| 818 | pytest.skip("fp8_e5m2 KV cache is not supported on ROCm/HIP") |
| 819 | |
| 820 | set_random_seed(0) |
| 821 | torch.set_default_device(device) |
| 822 | torch.accelerator.set_device_index(device) |
| 823 | |
| 824 | num_layers = 5 |
| 825 | num_tokens = 42 |
| 826 | num_blocks = 8 |
| 827 | block_size = 16 |
| 828 | kv_lora_rank = 512 |
| 829 | qk_rope_head_dim = 64 |
| 830 | entry_size = kv_lora_rank + qk_rope_head_dim |
| 831 | total_slots = num_blocks * block_size |
| 832 | |
| 833 | kv_c = torch.randn(num_layers, num_tokens, kv_lora_rank, dtype=torch.bfloat16) |
| 834 | k_pe = torch.randn(num_layers, num_tokens, qk_rope_head_dim, dtype=torch.bfloat16) |
| 835 | kv_caches = torch.zeros( |
| 836 | num_layers, |
| 837 | num_blocks, |
| 838 | block_size, |
| 839 | entry_size, |
| 840 | dtype=torch.bfloat16 if kv_cache_dtype == "auto" else torch.uint8, |
| 841 | ) |
| 842 | reference = torch.zeros_like(kv_caches) |
| 843 | scales = torch.linspace(0.05, 0.25, num_layers, dtype=torch.float32) |
| 844 | |
| 845 | slot_mapping = torch.stack( |
| 846 | [torch.randperm(total_slots)[:num_tokens] for _ in range(num_layers)] |
| 847 | ) |
| 848 | slot_mapping[:, -1] = -1 |
| 849 | if shared_slot_mapping: |
| 850 | slot_mapping = slot_mapping[:1].expand(num_layers, -1) |
| 851 | |
| 852 | for layer_idx in range(num_layers): |
| 853 | ops.concat_and_cache_mla( |
| 854 | kv_c[layer_idx], |
| 855 | k_pe[layer_idx], |
| 856 | reference[layer_idx], |
| 857 | slot_mapping[layer_idx], |
| 858 | kv_cache_dtype, |
| 859 | scales[layer_idx], |
| 860 | ) |
| 861 | |
| 862 | cache_ptrs = torch.tensor( |
| 863 | [kv_caches[layer_idx].data_ptr() for layer_idx in range(num_layers)], |
| 864 | dtype=torch.int64, |
| 865 | ) |
| 866 | ref_cache = kv_caches[0] |
| 867 | |
| 868 | def run_grouped(kv_scales: torch.Tensor | None) -> None: |
| 869 | ops.concat_and_cache_mla_grouped( |
| 870 | kv_c, |
| 871 | k_pe, |
| 872 | cache_ptrs, |
| 873 | slot_mapping, |
| 874 | ref_cache.size(1), |
| 875 | ref_cache.stride(0), |
| 876 | ref_cache.stride(1), |
| 877 | kv_scales, |
| 878 | kv_cache_dtype, |
| 879 | ) |
| 880 | |
| 881 | run_grouped(None if kv_cache_dtype == "auto" else scales) |
| 882 | |
| 883 | torch.testing.assert_close(kv_caches, reference, rtol=0, atol=0) |
| 884 | |
| 885 | if kv_cache_dtype == "fp8" and not shared_slot_mapping: |
| 886 | noncontiguous_scales = torch.ones(num_layers, 2)[:, 0] |
| 887 | assert not noncontiguous_scales.is_contiguous() |
| 888 | for invalid_scales, error in ( |
| 889 | (scales.cpu(), "same CUDA device"), |
| 890 | (noncontiguous_scales, "must be contiguous"), |
| 891 | ): |
| 892 | with pytest.raises(RuntimeError, match=error): |
| 893 | run_grouped(invalid_scales) |
| 894 | |
| 895 | |
| 896 | @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS) |
| 897 | @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS) |
| 898 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA) |
| 899 | @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA) |
| 900 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA) |
| 901 | @pytest.mark.parametrize("dtype", DTYPES) |
| 902 | @pytest.mark.parametrize("seed", SEEDS) |
| 903 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 904 | @torch.inference_mode() |
| 905 | def test_concat_and_cache_ds_mla( |
| 906 | kv_lora_rank: int, |
| 907 | qk_rope_head_dim: int, |
| 908 | num_tokens: int, |
| 909 | block_size: int, |
| 910 | num_blocks: int, |
| 911 | dtype: torch.dtype, |
| 912 | seed: int, |
| 913 | device: str, |
| 914 | ) -> None: |
| 915 | if current_platform.is_rocm(): |
| 916 | pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm") |
| 917 | if dtype.itemsize != 2: |
| 918 | pytest.skip("ds_mla only supports 16-bit input") |
| 919 | if kv_lora_rank != 512: |
| 920 | pytest.skip("fp8_ds_mla requires kv_lora_rank == 512") |
| 921 | kv_cache_dtype = "fp8_ds_mla" |
| 922 | set_random_seed(seed) |
| 923 | torch.set_default_device(device) |
| 924 | torch.accelerator.set_device_index(device) |
| 925 | |
| 926 | total_slots = num_blocks * block_size |
| 927 | slot_mapping_lst = random.sample(range(total_slots), num_tokens) |
| 928 | slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device) |
| 929 | |
| 930 | kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device) |
| 931 | k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device) |
| 932 | entry_size = kv_lora_rank + (4 * 4) + (2 * qk_rope_head_dim) |
| 933 | |
| 934 | scale = torch.tensor(1.0, dtype=torch.float32, device=device) |
| 935 | kv_cache = _create_mla_cache( |
| 936 | num_blocks, |
| 937 | block_size, |
| 938 | entry_size, |
| 939 | dtype=torch.uint8, |
| 940 | kv_cache_dtype=kv_cache_dtype, |
| 941 | device=device, |
| 942 | ) |
| 943 | |
| 944 | ref_cache = torch.zeros_like(kv_cache, dtype=kv_cache.dtype) |
| 945 | tile_data = torch.zeros(128, dtype=dtype, device=device) |
| 946 | |
| 947 | for i in range(num_tokens): |
| 948 | slot = slot_mapping[i].item() |
| 949 | block_idx = slot // block_size |
| 950 | block_offset = slot % block_size |
| 951 | |
| 952 | ref_cache_slice = ref_cache[block_idx, block_offset] |
| 953 | ref_cache_16bit = ref_cache_slice.view(dtype) |
| 954 | ref_cache_32bit = ref_cache_slice.view(torch.float32) |
| 955 | |
| 956 | kv_c_data = kv_c[i] |
| 957 | num_tiles = kv_lora_rank // 128 |
| 958 | for tile_idx in range(num_tiles): |
| 959 | tile_start = tile_idx * 128 |
| 960 | tile_end = (tile_idx + 1) * 128 |
| 961 | tile_data[:] = kv_c_data[tile_start:tile_end] |
| 962 | |
| 963 | # Using torch's amax() gives different results, so this must be |
| 964 | # manually computed. |
| 965 | tile_data_float = tile_data.to(torch.float32) |
| 966 | manual_max = abs(tile_data_float[0]) |
| 967 | for j in range(1, 128): |
| 968 | manual_max = max(manual_max, abs(tile_data_float[j])) |
| 969 | raw_scale = torch.clamp(manual_max / 448.0, min=1e-4) |
| 970 | tile_scale = torch.exp2(torch.ceil(torch.log2(raw_scale))) |
| 971 | |
| 972 | ref_cache_32bit[kv_lora_rank // 4 + tile_idx] = tile_scale |
| 973 | |
| 974 | ops.convert_fp8( |
| 975 | ref_cache_slice[tile_start:tile_end], |
| 976 | tile_data, |
| 977 | tile_scale.item(), |
| 978 | kv_dtype="fp8", |
| 979 | ) |
| 980 | |
| 981 | for j in range(qk_rope_head_dim): |
| 982 | ref_cache_16bit[kv_lora_rank // 2 + 8 + j] = k_pe[i, j] |
| 983 | |
| 984 | opcheck( |
| 985 | torch.ops._C_cache_ops.concat_and_cache_mla, |
| 986 | (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale), |
| 987 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 988 | ) |
| 989 | |
| 990 | ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale) |
| 991 | |
| 992 | for i in range(num_tokens): |
| 993 | slot = slot_mapping[i].item() |
| 994 | block_idx = slot // block_size |
| 995 | block_offset = slot % block_size |
| 996 | kv_cache_slice = kv_cache[block_idx, block_offset] |
| 997 | ref_cache_slice = ref_cache[block_idx, block_offset] |
| 998 | |
| 999 | kv_nope = kv_cache_slice[:kv_lora_rank] |
| 1000 | ref_nope = ref_cache_slice[:kv_lora_rank] |
| 1001 | kv_scales = kv_cache_slice.view(torch.float32)[ |
| 1002 | kv_lora_rank // 4 : kv_lora_rank // 4 + 4 |
| 1003 | ] |
| 1004 | ref_scales = ref_cache_slice.view(torch.float32)[ |
| 1005 | kv_lora_rank // 4 : kv_lora_rank // 4 + 4 |
| 1006 | ] |
| 1007 | kv_rope = kv_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :] |
| 1008 | ref_rope = ref_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :] |
| 1009 | |
| 1010 | torch.testing.assert_close(kv_nope, ref_nope, atol=0.001, rtol=0.1) |
| 1011 | torch.testing.assert_close(kv_scales, ref_scales, atol=0, rtol=0) |
| 1012 | torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1) |
| 1013 | |
| 1014 | |
| 1015 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1016 | @pytest.mark.parametrize("block_size", [64, 256]) |
| 1017 | @torch.inference_mode() |
| 1018 | def test_concat_and_cache_ds_mla_nope(device: str, block_size: int) -> None: |
| 1019 | """NoPE matches zero RoPE, clears valid tails, and preserves unused slots.""" |
| 1020 | dtype = torch.bfloat16 |
| 1021 | if current_platform.is_rocm(): |
| 1022 | pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm") |
| 1023 | num_tokens, num_blocks = 3, 2 |
| 1024 | set_random_seed(0) |
| 1025 | torch.set_default_device(device) |
| 1026 | torch.accelerator.set_device_index(device) |
| 1027 | |
| 1028 | slot_mapping = torch.tensor( |
| 1029 | [block_size - 1, block_size, -1], dtype=torch.long, device=device |
| 1030 | ) |
| 1031 | kv_c = torch.randn(num_tokens, 512, dtype=dtype, device=device) |
| 1032 | k_pe = torch.empty(num_tokens, 0, dtype=dtype, device=device) |
| 1033 | zero_k_pe = torch.zeros(num_tokens, 64, dtype=dtype, device=device) |
| 1034 | scale = torch.tensor(1.0, dtype=torch.float32, device=device) |
| 1035 | kv_cache = torch.full( |
| 1036 | (num_blocks, block_size, 656), 0xFF, dtype=torch.uint8, device=device |
| 1037 | ) |
| 1038 | ref_cache = kv_cache.clone() |
| 1039 | |
| 1040 | opcheck( |
| 1041 | torch.ops._C_cache_ops.concat_and_cache_mla, |
| 1042 | (kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale), |
| 1043 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 1044 | ) |
| 1045 | kv_cache.fill_(0xFF) |
| 1046 | ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale) |
| 1047 | ops.concat_and_cache_mla( |
| 1048 | kv_c, zero_k_pe, ref_cache, slot_mapping, "fp8_ds_mla", scale |
| 1049 | ) |
| 1050 | |
| 1051 | torch.testing.assert_close(kv_cache, ref_cache, atol=0, rtol=0) |
| 1052 | rows = kv_cache.view(num_blocks * block_size, 656) |
| 1053 | assert (rows[slot_mapping[:2], 528:] == 0).all() |
| 1054 | untouched = torch.ones(num_blocks * block_size, dtype=torch.bool, device=device) |
| 1055 | untouched[slot_mapping[:2]] = False |
| 1056 | assert (rows[untouched] == 0xFF).all() |
| 1057 | |
| 1058 | |
| 1059 | # Bytes per token for the nvfp4_ds_mla cache layout (see flashmla_sparse.py). |
| 1060 | NVFP4_DS_MLA_ENTRY_SIZE = 352 |
| 1061 | |
| 1062 | # e2m1 magnitude table; the sign lives in bit 3 of each 4-bit code. |
| 1063 | E2M1_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] |
| 1064 | |
| 1065 | |
| 1066 | def _unpack_e2m1(packed: torch.Tensor) -> torch.Tensor: |
| 1067 | """Unpack uint8 bytes holding two e2m1 codes each into float32 values. |
| 1068 | |
| 1069 | The low nibble holds the even element, the high nibble the odd one. |
| 1070 | """ |
| 1071 | low = packed & 0x0F |
| 1072 | high = packed >> 4 |
| 1073 | codes = torch.stack([low, high], dim=-1).reshape(-1).long() |
| 1074 | table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=packed.device) |
| 1075 | magnitude = table[codes & 0x7] |
| 1076 | sign = torch.where((codes & 0x8) != 0, -1.0, 1.0) |
| 1077 | return sign * magnitude |
| 1078 | |
| 1079 | |
| 1080 | def _ref_nvfp4_sf_bytes(vals: torch.Tensor, divisor: float) -> torch.Tensor: |
| 1081 | """Reference e4m3 scale-factor bytes: e4m3(max(amax_per_16 / divisor, 2^-9)). |
| 1082 | |
| 1083 | In element-block order (byte s scales elements [16s, 16s + 16)). |
| 1084 | """ |
| 1085 | tiles = vals.to(torch.float32).reshape(-1, 16) |
| 1086 | sf = torch.clamp(tiles.abs().amax(dim=1) / divisor, min=2.0**-9) |
| 1087 | return sf.to(torch.float8_e4m3fn).view(torch.uint8) |
| 1088 | |
| 1089 | |
| 1090 | def _nvfp4_sf_element_order(wire_sf_bytes: torch.Tensor) -> torch.Tensor: |
| 1091 | """Stored NVFP4 scale-factor byte order -> element-block order. |
| 1092 | |
| 1093 | The 32 NoPE scale-factor bytes of an nvfp4_ds_mla entry are stored |
| 1094 | permuted (an 8x4 -> 4x8 transpose): the scale for element block s lives at |
| 1095 | byte 8 * (s & 3) + (s >> 2), so that the 8 scales one FlashMLA dequant |
| 1096 | thread needs are contiguous. This undoes that permutation. |
| 1097 | """ |
| 1098 | num_sf = wire_sf_bytes.shape[-1] |
| 1099 | return wire_sf_bytes.unflatten(-1, (4, num_sf // 4)).transpose(-1, -2).flatten(-2) |
| 1100 | |
| 1101 | |
| 1102 | @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS) |
| 1103 | @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS) |
| 1104 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA) |
| 1105 | @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA) |
| 1106 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA) |
| 1107 | @pytest.mark.parametrize("seed", SEEDS) |
| 1108 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1109 | @torch.inference_mode() |
| 1110 | def test_concat_and_cache_nvfp4_ds_mla( |
| 1111 | kv_lora_rank: int, |
| 1112 | qk_rope_head_dim: int, |
| 1113 | num_tokens: int, |
| 1114 | block_size: int, |
| 1115 | num_blocks: int, |
| 1116 | seed: int, |
| 1117 | device: str, |
| 1118 | ) -> None: |
| 1119 | if current_platform.is_rocm(): |
| 1120 | pytest.skip("concat_and_cache_mla doesn't support NVFP4 DS-MLA on ROCm") |
| 1121 | device_capability = current_platform.get_device_capability() |
| 1122 | if device_capability is None or device_capability.major != 10: |
| 1123 | pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x") |
| 1124 | if kv_lora_rank != 512: |
| 1125 | pytest.skip("The NVFP4 DS-MLA layout requires kv_lora_rank == 512") |
| 1126 | dtype = torch.bfloat16 |
| 1127 | set_random_seed(seed) |
| 1128 | torch.set_default_device(device) |
| 1129 | torch.accelerator.set_device_index(device) |
| 1130 | |
| 1131 | kv_cache_dtype = "nvfp4_ds_mla" |
| 1132 | entry_size = NVFP4_DS_MLA_ENTRY_SIZE |
| 1133 | |
| 1134 | total_slots = num_blocks * block_size |
| 1135 | slot_mapping_lst = random.sample(range(total_slots), num_tokens) |
| 1136 | slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device) |
| 1137 | |
| 1138 | kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device) |
| 1139 | k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device) |
| 1140 | |
| 1141 | scale = torch.tensor(1.0, dtype=torch.float32, device=device) |
| 1142 | kv_cache = _create_mla_cache( |
| 1143 | num_blocks, |
| 1144 | block_size, |
| 1145 | entry_size, |
| 1146 | dtype=torch.uint8, |
| 1147 | kv_cache_dtype=kv_cache_dtype, |
| 1148 | device=device, |
| 1149 | ) |
| 1150 | |
| 1151 | opcheck( |
| 1152 | torch.ops._C_cache_ops.concat_and_cache_mla, |
| 1153 | (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale), |
| 1154 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 1155 | ) |
| 1156 | |
| 1157 | ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale) |
| 1158 | |
| 1159 | # Byte layout of a single cache entry: |
| 1160 | # [0, 256) 512 e2m1 NoPE values packed 2/byte |
| 1161 | # [256, 320) 64 unscaled e4m3 RoPE values |
| 1162 | # [nope_sf_off, 352) 32 e4m3 NoPE SFs (one per 16 elements), stored |
| 1163 | # permuted (see _nvfp4_sf_element_order) |
| 1164 | nope_bytes = kv_lora_rank // 2 |
| 1165 | rope_bytes = qk_rope_head_dim |
| 1166 | nope_sf_off = nope_bytes + rope_bytes |
| 1167 | num_nope_sf = kv_lora_rank // 16 |
| 1168 | assert nope_sf_off + num_nope_sf == entry_size |
| 1169 | |
| 1170 | for i in range(num_tokens): |
| 1171 | slot = slot_mapping[i].item() |
| 1172 | block_idx = slot // block_size |
| 1173 | block_offset = slot % block_size |
| 1174 | entry = kv_cache[block_idx, block_offset] |
| 1175 | |
| 1176 | kv_c_ref = kv_c[i].to(torch.float32) |
| 1177 | k_pe_ref = k_pe[i].to(torch.float32) |
| 1178 | |
| 1179 | # Scale-factor bytes must match e4m3(max(amax_per_16 / divisor, 2^-9)) |
| 1180 | # computed in python, allowing a 1-ulp e4m3 difference (all SFs are |
| 1181 | # positive, so e4m3 bit patterns are monotonic and adjacent codes |
| 1182 | # differ by 1). |
| 1183 | nope_sf_bytes = _nvfp4_sf_element_order( |
| 1184 | entry[nope_sf_off : nope_sf_off + num_nope_sf] |
| 1185 | ) |
| 1186 | ref_nope_sf_bytes = _ref_nvfp4_sf_bytes(kv_c_ref, 6.0) |
| 1187 | assert ( |
| 1188 | (nope_sf_bytes.to(torch.int16) - ref_nope_sf_bytes.to(torch.int16)).abs() |
| 1189 | <= 1 |
| 1190 | ).all() |
| 1191 | |
| 1192 | # Dequantize with the *stored* scale factors: x = float(code) * float(sf). |
| 1193 | nope_sf = nope_sf_bytes.view(torch.float8_e4m3fn).to(torch.float32) |
| 1194 | nope_sf_per_elem = nope_sf.repeat_interleave(16) |
| 1195 | nope_vals = _unpack_e2m1(entry[:nope_bytes]) * nope_sf_per_elem |
| 1196 | # e2m1 error bound: the widest grid spacing is 2 (between codes 4 and |
| 1197 | # 6), so round-to-nearest is off by at most 1.0*sf; the e4m3 rounding |
| 1198 | # of the SF itself adds at most ~2^-4 relative, so 1.5*sf is safe for |
| 1199 | # either encode convention (quantized or unquantized SF). |
| 1200 | assert ((nope_vals - kv_c_ref).abs() <= 1.5 * nope_sf_per_elem).all() |
| 1201 | |
| 1202 | # Plain unscaled e4m3: ~2^-4 relative rounding, with a small absolute |
| 1203 | # floor for values in the subnormal range. |
| 1204 | rope_payload = entry[nope_bytes:nope_sf_off] |
| 1205 | rope_vals = rope_payload.view(torch.float8_e4m3fn).to(torch.float32) |
| 1206 | rope_tol = k_pe_ref.abs() * 2.0**-3 + 2.0**-9 |
| 1207 | assert ((rope_vals - k_pe_ref).abs() <= rope_tol).all() |
| 1208 | |
| 1209 | |
| 1210 | @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS) |
| 1211 | @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS) |
| 1212 | @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA) |
| 1213 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA) |
| 1214 | @pytest.mark.parametrize("dtype", DTYPES) |
| 1215 | @pytest.mark.parametrize("seed", SEEDS) |
| 1216 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1217 | @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) |
| 1218 | @torch.inference_mode() |
| 1219 | def test_swap_blocks_mla( |
| 1220 | kv_lora_rank: int, |
| 1221 | qk_rope_head_dim: int, |
| 1222 | block_size: int, |
| 1223 | num_blocks: int, |
| 1224 | dtype: torch.dtype, |
| 1225 | seed: int, |
| 1226 | device: str, |
| 1227 | kv_cache_dtype: str, |
| 1228 | ) -> None: |
| 1229 | set_random_seed(seed) |
| 1230 | torch.set_default_device(device) |
| 1231 | torch.accelerator.set_device_index(device) |
| 1232 | |
| 1233 | entry_size = kv_lora_rank + qk_rope_head_dim |
| 1234 | |
| 1235 | src_cache = _create_mla_cache( |
| 1236 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 1237 | ) |
| 1238 | dst_cache = _create_mla_cache( |
| 1239 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 1240 | ) |
| 1241 | |
| 1242 | _fill_mla_cache(src_cache, kv_cache_dtype) |
| 1243 | _fill_mla_cache(dst_cache, kv_cache_dtype) |
| 1244 | |
| 1245 | src_cache_clone = src_cache.clone() |
| 1246 | |
| 1247 | num_mappings = min(2, num_blocks // 2) |
| 1248 | src_blocks = random.sample(range(num_blocks), num_mappings) |
| 1249 | remaining_blocks = list(set(range(num_blocks)) - set(src_blocks)) |
| 1250 | dst_blocks = random.sample(remaining_blocks, num_mappings) |
| 1251 | block_mapping = list(zip(src_blocks, dst_blocks)) |
| 1252 | block_mapping_tensor = torch.tensor( |
| 1253 | block_mapping, dtype=torch.int64, device="cpu" |
| 1254 | ).view(-1, 2) |
| 1255 | |
| 1256 | block_size_in_bytes = src_cache.element_size() * src_cache.stride(0) |
| 1257 | opcheck( |
| 1258 | torch.ops._C_cache_ops.swap_blocks, |
| 1259 | (src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor), |
| 1260 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 1261 | ) |
| 1262 | |
| 1263 | ops.swap_blocks(src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor) |
| 1264 | |
| 1265 | for src, dst in block_mapping: |
| 1266 | torch.testing.assert_close( |
| 1267 | src_cache_clone[src].cpu(), |
| 1268 | dst_cache[dst].cpu(), |
| 1269 | msg=f"Block {src} from src should have been swapped to block " |
| 1270 | f"{dst} in dst_cache.", |
| 1271 | ) |
| 1272 | |
| 1273 | |
| 1274 | @pytest.mark.parametrize("kv_lora_rank", [512]) |
| 1275 | @pytest.mark.parametrize("qk_rope_head_dim", [64]) |
| 1276 | @pytest.mark.parametrize("block_size", [16]) |
| 1277 | @pytest.mark.parametrize("num_blocks", [1024]) |
| 1278 | @pytest.mark.parametrize("max_seq_len", [512]) |
| 1279 | @pytest.mark.parametrize("batch_size", [8]) |
| 1280 | @pytest.mark.parametrize("dtype", [torch.float32]) |
| 1281 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 1282 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1283 | @torch.inference_mode() |
| 1284 | def test_gather_and_maybe_dequant_cache_mla( |
| 1285 | kv_lora_rank, |
| 1286 | qk_rope_head_dim, |
| 1287 | block_size, |
| 1288 | num_blocks, |
| 1289 | max_seq_len, |
| 1290 | batch_size, |
| 1291 | dtype, |
| 1292 | kv_cache_dtype, |
| 1293 | device, |
| 1294 | ): |
| 1295 | entry_size = kv_lora_rank + qk_rope_head_dim |
| 1296 | scale = torch.tensor(0.1, dtype=torch.float32, device=device) |
| 1297 | src_cache = _create_mla_cache( |
| 1298 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 1299 | ) |
| 1300 | _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) |
| 1301 | |
| 1302 | seq_len_tensor = torch.randint( |
| 1303 | max_seq_len, max_seq_len + 1, (batch_size,), device=device |
| 1304 | ) |
| 1305 | |
| 1306 | total_tokens = seq_len_tensor.sum() |
| 1307 | cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device) |
| 1308 | cu_seq_lens[0] = 0 |
| 1309 | cu_seq_lens[1:] = seq_len_tensor.cumsum(dim=0).to(dtype=torch.int32) |
| 1310 | token_to_seq = torch.arange(0, batch_size, dtype=torch.int32, device=device) |
| 1311 | token_to_seq = torch.repeat_interleave(token_to_seq, seq_len_tensor) |
| 1312 | print("seq_len_tensor", seq_len_tensor) |
| 1313 | |
| 1314 | tot_blocks_tensor = (seq_len_tensor + block_size - 1) // block_size |
| 1315 | block_table = torch.empty( |
| 1316 | (batch_size, num_blocks), dtype=torch.int32, device=device |
| 1317 | ) |
| 1318 | |
| 1319 | for b in range(batch_size): |
| 1320 | perm = torch.randperm(num_blocks, device=device) |
| 1321 | block_table[b, :] = perm |
| 1322 | |
| 1323 | dst = torch.zeros((total_tokens, entry_size), dtype=dtype, device=device) |
| 1324 | |
| 1325 | expected_batches = [] |
| 1326 | for b in range(batch_size): |
| 1327 | s = seq_len_tensor[b] |
| 1328 | if s == 0: |
| 1329 | continue |
| 1330 | tot = tot_blocks_tensor[b] |
| 1331 | blocks = block_table[b, :tot].tolist() |
| 1332 | |
| 1333 | gathered_rows = [] |
| 1334 | for i in range(tot - 1): |
| 1335 | block_data = src_cache[blocks[i]] |
| 1336 | if kv_cache_dtype == "fp8": |
| 1337 | dequantized_block = torch.empty_like(block_data, dtype=dtype) |
| 1338 | ops.convert_fp8(dequantized_block, block_data, scale.item()) |
| 1339 | gathered_rows.append(dequantized_block) |
| 1340 | else: |
| 1341 | gathered_rows.append(block_data) |
| 1342 | remaining = s - (tot - 1) * block_size |
| 1343 | last_block_data = src_cache[blocks[-1], :remaining, :] |
| 1344 | if kv_cache_dtype == "fp8": |
| 1345 | dequantized_last_block = torch.empty_like(last_block_data, dtype=dtype) |
| 1346 | ops.convert_fp8(dequantized_last_block, last_block_data, scale.item()) |
| 1347 | gathered_rows.append(dequantized_last_block) |
| 1348 | else: |
| 1349 | gathered_rows.append(last_block_data) |
| 1350 | |
| 1351 | batch_expected = torch.cat(gathered_rows, dim=0) |
| 1352 | expected_batches.append(batch_expected) |
| 1353 | expected = torch.cat(expected_batches, dim=0) |
| 1354 | |
| 1355 | opcheck( |
| 1356 | torch.ops._C_cache_ops.gather_and_maybe_dequant_cache, |
| 1357 | ( |
| 1358 | src_cache, |
| 1359 | dst, |
| 1360 | block_table, |
| 1361 | cu_seq_lens, |
| 1362 | token_to_seq, |
| 1363 | total_tokens, |
| 1364 | kv_cache_dtype, |
| 1365 | scale, |
| 1366 | None, |
| 1367 | ), |
| 1368 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 1369 | ) |
| 1370 | |
| 1371 | ops.gather_and_maybe_dequant_cache( |
| 1372 | src_cache, |
| 1373 | dst, |
| 1374 | block_table, |
| 1375 | cu_seq_lens, |
| 1376 | token_to_seq, |
| 1377 | total_tokens, |
| 1378 | kv_cache_dtype, |
| 1379 | scale, |
| 1380 | None, |
| 1381 | ) |
| 1382 | torch.testing.assert_close(dst, expected) |
| 1383 | |
| 1384 | |
| 1385 | @pytest.mark.parametrize("kv_lora_rank", [512]) |
| 1386 | @pytest.mark.parametrize("qk_rope_head_dim", [64]) |
| 1387 | @pytest.mark.parametrize("block_size", [16]) |
| 1388 | @pytest.mark.parametrize("num_blocks", [128]) |
| 1389 | @pytest.mark.parametrize("dtype", [torch.float32]) |
| 1390 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 1391 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1392 | @torch.inference_mode() |
| 1393 | def test_gather_and_maybe_dequant_cache_mla_with_seq_starts( |
| 1394 | kv_lora_rank, |
| 1395 | qk_rope_head_dim, |
| 1396 | block_size, |
| 1397 | num_blocks, |
| 1398 | dtype, |
| 1399 | kv_cache_dtype, |
| 1400 | device, |
| 1401 | ): |
| 1402 | entry_size = kv_lora_rank + qk_rope_head_dim |
| 1403 | scale = torch.tensor(0.1, dtype=torch.float32, device=device) |
| 1404 | src_cache = _create_mla_cache( |
| 1405 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 1406 | ) |
| 1407 | _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) |
| 1408 | |
| 1409 | seq_starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device) |
| 1410 | seq_lens = torch.tensor([20, 10, 16], dtype=torch.int32, device=device) |
| 1411 | batch_size = seq_lens.shape[0] |
| 1412 | total_tokens = seq_lens.sum().item() |
| 1413 | cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device) |
| 1414 | cu_seq_lens[0] = 0 |
| 1415 | cu_seq_lens[1:] = seq_lens.cumsum(dim=0) |
| 1416 | token_to_seq = torch.repeat_interleave( |
| 1417 | torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens |
| 1418 | ) |
| 1419 | |
| 1420 | block_table = torch.empty( |
| 1421 | (batch_size, num_blocks), dtype=torch.int32, device=device |
| 1422 | ) |
| 1423 | for b in range(batch_size): |
| 1424 | block_table[b, :] = torch.randperm(num_blocks, device=device) |
| 1425 | |
| 1426 | if kv_cache_dtype == "fp8": |
| 1427 | dequant_src_cache = torch.empty_like(src_cache, dtype=dtype) |
| 1428 | ops.convert_fp8(dequant_src_cache, src_cache, scale.item()) |
| 1429 | else: |
| 1430 | dequant_src_cache = src_cache |
| 1431 | |
| 1432 | expected_rows = [] |
| 1433 | for b in range(batch_size): |
| 1434 | start = seq_starts[b].item() |
| 1435 | length = seq_lens[b].item() |
| 1436 | for offset in range(start, start + length): |
| 1437 | block_id = block_table[b, offset // block_size] |
| 1438 | slot = offset % block_size |
| 1439 | expected_rows.append(dequant_src_cache[block_id, slot]) |
| 1440 | expected = torch.stack(expected_rows) |
| 1441 | |
| 1442 | dst = torch.zeros((total_tokens, entry_size), dtype=dtype, device=device) |
| 1443 | ops.gather_and_maybe_dequant_cache( |
| 1444 | src_cache, |
| 1445 | dst, |
| 1446 | block_table, |
| 1447 | cu_seq_lens, |
| 1448 | token_to_seq, |
| 1449 | total_tokens, |
| 1450 | kv_cache_dtype, |
| 1451 | scale, |
| 1452 | seq_starts, |
| 1453 | ) |
| 1454 | torch.testing.assert_close(dst, expected) |
| 1455 | |
| 1456 | |
| 1457 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 1458 | @pytest.mark.parametrize("use_seq_starts", [False, True]) |
| 1459 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1460 | @torch.inference_mode() |
| 1461 | def test_gather_and_maybe_dequant_cache_mla_large_uneven_sequences( |
| 1462 | kv_cache_dtype, |
| 1463 | use_seq_starts, |
| 1464 | device, |
| 1465 | ): |
| 1466 | block_size = 64 |
| 1467 | entry_size = 576 |
| 1468 | num_blocks = 1024 |
| 1469 | dtype = torch.bfloat16 |
| 1470 | scale = torch.tensor(0.1, dtype=torch.float32, device=device) |
| 1471 | src_cache = _create_mla_cache( |
| 1472 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 1473 | ) |
| 1474 | _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) |
| 1475 | |
| 1476 | starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device) |
| 1477 | seq_starts = starts if use_seq_starts else None |
| 1478 | seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device) |
| 1479 | batch_size = seq_lens.shape[0] |
| 1480 | total_tokens = seq_lens.sum().item() |
| 1481 | cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) |
| 1482 | cu_seq_lens[1:] = seq_lens.cumsum(dim=0) |
| 1483 | token_to_seq = torch.repeat_interleave( |
| 1484 | torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens |
| 1485 | ) |
| 1486 | block_table = torch.stack( |
| 1487 | [torch.randperm(num_blocks, device=device) for _ in range(batch_size)] |
| 1488 | ).to(torch.int32) |
| 1489 | |
| 1490 | if kv_cache_dtype == "fp8": |
| 1491 | dequant_src_cache = torch.empty_like(src_cache, dtype=dtype) |
| 1492 | ops.convert_fp8(dequant_src_cache, src_cache, scale.item()) |
| 1493 | else: |
| 1494 | dequant_src_cache = src_cache |
| 1495 | expected_batches = [] |
| 1496 | for req_id in range(batch_size): |
| 1497 | start = starts[req_id].item() if use_seq_starts else 0 |
| 1498 | source_tokens = torch.arange( |
| 1499 | start, start + seq_lens[req_id].item(), device=device |
| 1500 | ) |
| 1501 | physical_blocks = block_table[req_id, source_tokens // block_size].long() |
| 1502 | expected_batches.append( |
| 1503 | dequant_src_cache[physical_blocks, source_tokens % block_size] |
| 1504 | ) |
| 1505 | expected = torch.cat(expected_batches) |
| 1506 | dst = torch.empty_like(expected) |
| 1507 | |
| 1508 | ops.gather_and_maybe_dequant_cache( |
| 1509 | src_cache, |
| 1510 | dst, |
| 1511 | block_table, |
| 1512 | cu_seq_lens, |
| 1513 | token_to_seq, |
| 1514 | total_tokens, |
| 1515 | kv_cache_dtype, |
| 1516 | scale, |
| 1517 | seq_starts, |
| 1518 | ) |
| 1519 | torch.testing.assert_close(dst, expected) |
| 1520 | |
| 1521 | |
| 1522 | @pytest.mark.parametrize("kv_lora_rank", [512]) |
| 1523 | @pytest.mark.parametrize("qk_rope_head_dim", [64]) |
| 1524 | @pytest.mark.parametrize("block_size", [16]) |
| 1525 | @pytest.mark.parametrize("num_blocks", [1024]) |
| 1526 | @pytest.mark.parametrize("max_seq_len", [512]) |
| 1527 | @pytest.mark.parametrize("batch_size", [8]) |
| 1528 | @pytest.mark.parametrize("dtype", [torch.float32]) |
| 1529 | @pytest.mark.parametrize( |
| 1530 | "kv_cache_dtype", ["auto"] |
| 1531 | ) # You can also test "fp8" if needed. |
| 1532 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1533 | @torch.inference_mode() |
| 1534 | def test_cp_gather_cache_mla( |
| 1535 | kv_lora_rank, |
| 1536 | qk_rope_head_dim, |
| 1537 | block_size, |
| 1538 | num_blocks, |
| 1539 | max_seq_len, |
| 1540 | batch_size, |
| 1541 | dtype, |
| 1542 | kv_cache_dtype, |
| 1543 | device, |
| 1544 | ): |
| 1545 | entry_size = kv_lora_rank + qk_rope_head_dim |
| 1546 | src_cache = _create_mla_cache( |
| 1547 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 1548 | ) |
| 1549 | _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) |
| 1550 | |
| 1551 | seq_len_tensor = torch.randint(0, max_seq_len + 1, (batch_size,), device=device) |
| 1552 | |
| 1553 | total_tokens = seq_len_tensor.sum() |
| 1554 | cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device) |
| 1555 | cu_seq_lens[0] = 0 |
| 1556 | cu_seq_lens[1:] = seq_len_tensor.cumsum(dim=0).to(dtype=torch.int32) |
| 1557 | print("seq_len_tensor", seq_len_tensor) |
| 1558 | |
| 1559 | tot_blocks_tensor = (seq_len_tensor + block_size - 1) // block_size |
| 1560 | block_table = torch.empty( |
| 1561 | (batch_size, num_blocks), dtype=torch.int32, device=device |
| 1562 | ) |
| 1563 | |
| 1564 | for b in range(batch_size): |
| 1565 | perm = torch.randperm(num_blocks, device=device) |
| 1566 | block_table[b, :] = perm |
| 1567 | |
| 1568 | dst = torch.zeros((total_tokens, entry_size), dtype=src_cache.dtype, device=device) |
| 1569 | |
| 1570 | expected_batches = [] |
| 1571 | for b in range(batch_size): |
| 1572 | s = seq_len_tensor[b] |
| 1573 | if s == 0: |
| 1574 | continue |
| 1575 | tot = tot_blocks_tensor[b] |
| 1576 | blocks = block_table[b, :tot].tolist() |
| 1577 | |
| 1578 | gathered_rows = [] |
| 1579 | for i in range(tot - 1): |
| 1580 | gathered_rows.append(src_cache[blocks[i]]) |
| 1581 | remaining = s - (tot - 1) * block_size |
| 1582 | gathered_rows.append(src_cache[blocks[-1], :remaining, :]) |
| 1583 | |
| 1584 | batch_expected = torch.cat(gathered_rows, dim=0) |
| 1585 | expected_batches.append(batch_expected) |
| 1586 | expected = torch.cat(expected_batches, dim=0) |
| 1587 | |
| 1588 | opcheck( |
| 1589 | torch.ops._C_cache_ops.cp_gather_cache, |
| 1590 | (src_cache, dst, block_table, cu_seq_lens, batch_size, None), |
| 1591 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 1592 | ) |
| 1593 | |
| 1594 | ops.cp_gather_cache(src_cache, dst, block_table, cu_seq_lens, batch_size) |
| 1595 | torch.testing.assert_close(dst, expected) |
| 1596 | |
| 1597 | |
| 1598 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 1599 | @pytest.mark.parametrize("unaligned", [False, True]) |
| 1600 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1601 | @torch.inference_mode() |
| 1602 | def test_cp_gather_cache_mla_with_seq_starts( |
| 1603 | kv_cache_dtype, |
| 1604 | unaligned, |
| 1605 | device, |
| 1606 | ): |
| 1607 | block_size = 16 |
| 1608 | entry_size = 576 |
| 1609 | num_blocks = 128 |
| 1610 | src_storage = _create_mla_cache( |
| 1611 | num_blocks, |
| 1612 | block_size, |
| 1613 | entry_size + int(unaligned), |
| 1614 | torch.bfloat16, |
| 1615 | kv_cache_dtype, |
| 1616 | device, |
| 1617 | ) |
| 1618 | src_cache = src_storage[..., 1:] if unaligned else src_storage |
| 1619 | _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) |
| 1620 | |
| 1621 | seq_starts = torch.tensor([3, 17, 5, 31], dtype=torch.int32, device=device) |
| 1622 | seq_lens = torch.tensor([0, 1, 63, 4], dtype=torch.int32, device=device) |
| 1623 | batch_size = seq_lens.shape[0] |
| 1624 | cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) |
| 1625 | cu_seq_lens[1:] = seq_lens.cumsum(dim=0) |
| 1626 | block_table = torch.stack( |
| 1627 | [torch.randperm(num_blocks, device=device) for _ in range(batch_size)] |
| 1628 | ).to(torch.int32) |
| 1629 | |
| 1630 | expected_rows = [] |
| 1631 | for req_id in range(batch_size): |
| 1632 | start = seq_starts[req_id].item() |
| 1633 | for source_token in range(start, start + seq_lens[req_id].item()): |
| 1634 | block_id = block_table[req_id, source_token // block_size] |
| 1635 | expected_rows.append(src_cache[block_id, source_token % block_size]) |
| 1636 | expected = torch.stack(expected_rows) |
| 1637 | if unaligned: |
| 1638 | dst_storage = torch.empty( |
| 1639 | (expected.shape[0], entry_size + 1), |
| 1640 | dtype=expected.dtype, |
| 1641 | device=device, |
| 1642 | ) |
| 1643 | dst = dst_storage[:, 1:] |
| 1644 | else: |
| 1645 | dst = torch.empty_like(expected) |
| 1646 | |
| 1647 | ops.cp_gather_cache( |
| 1648 | src_cache, |
| 1649 | dst, |
| 1650 | block_table, |
| 1651 | cu_seq_lens, |
| 1652 | batch_size, |
| 1653 | seq_starts, |
| 1654 | ) |
| 1655 | torch.testing.assert_close(dst, expected) |
| 1656 | |
| 1657 | |
| 1658 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 1659 | @pytest.mark.parametrize("use_seq_starts", [False, True]) |
| 1660 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 1661 | @torch.inference_mode() |
| 1662 | def test_cp_gather_cache_mla_large_uneven_sequences( |
| 1663 | kv_cache_dtype, |
| 1664 | use_seq_starts, |
| 1665 | device, |
| 1666 | ): |
| 1667 | block_size = 64 |
| 1668 | entry_size = 576 |
| 1669 | num_blocks = 1024 |
| 1670 | src_cache = _create_mla_cache( |
| 1671 | num_blocks, |
| 1672 | block_size, |
| 1673 | entry_size, |
| 1674 | torch.bfloat16, |
| 1675 | kv_cache_dtype, |
| 1676 | device, |
| 1677 | ) |
| 1678 | _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) |
| 1679 | |
| 1680 | starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device) |
| 1681 | seq_starts = starts if use_seq_starts else None |
| 1682 | seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device) |
| 1683 | batch_size = seq_lens.shape[0] |
| 1684 | cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) |
| 1685 | cu_seq_lens[1:] = seq_lens.cumsum(dim=0) |
| 1686 | block_table = torch.stack( |
| 1687 | [torch.randperm(num_blocks, device=device) for _ in range(batch_size)] |
| 1688 | ).to(torch.int32) |
| 1689 | |
| 1690 | expected_batches = [] |
| 1691 | for req_id in range(batch_size): |
| 1692 | start = starts[req_id].item() if use_seq_starts else 0 |
| 1693 | source_tokens = torch.arange( |
| 1694 | start, start + seq_lens[req_id].item(), device=device |
| 1695 | ) |
| 1696 | physical_blocks = block_table[req_id, source_tokens // block_size].long() |
| 1697 | expected_batches.append(src_cache[physical_blocks, source_tokens % block_size]) |
| 1698 | expected = torch.cat(expected_batches) |
| 1699 | dst = torch.empty_like(expected) |
| 1700 | |
| 1701 | ops.cp_gather_cache( |
| 1702 | src_cache, |
| 1703 | dst, |
| 1704 | block_table, |
| 1705 | cu_seq_lens, |
| 1706 | batch_size, |
| 1707 | seq_starts, |
| 1708 | ) |
| 1709 | torch.testing.assert_close(dst, expected) |
| 1710 | |
| 1711 | |
| 1712 | @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS) |
| 1713 | @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS) |
| 1714 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA) |
| 1715 | @pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA) |
| 1716 | @pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA) |
| 1717 | @pytest.mark.parametrize("dtype", DTYPES) |
| 1718 | @pytest.mark.parametrize("seed", SEEDS) |
| 1719 | @pytest.mark.cpu_model |
| 1720 | @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only") |
| 1721 | @torch.inference_mode() |
| 1722 | def test_concat_and_cache_mla_cpu( |
| 1723 | kv_lora_rank: int, |
| 1724 | qk_rope_head_dim: int, |
| 1725 | num_tokens: int, |
| 1726 | block_size: int, |
| 1727 | num_blocks: int, |
| 1728 | dtype: torch.dtype, |
| 1729 | seed: int, |
| 1730 | ) -> None: |
| 1731 | device = "cpu" |
| 1732 | kv_cache_dtype = "auto" |
| 1733 | set_random_seed(seed) |
| 1734 | torch.set_default_device(device) |
| 1735 | |
| 1736 | total_slots = num_blocks * block_size |
| 1737 | slot_mapping_lst = random.sample(range(total_slots), num_tokens) |
| 1738 | slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device) |
| 1739 | |
| 1740 | kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device) |
| 1741 | k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device) |
| 1742 | entry_size = kv_lora_rank + qk_rope_head_dim |
| 1743 | |
| 1744 | scale = torch.tensor(0.1, dtype=torch.float32, device=device) |
| 1745 | kv_cache = _create_mla_cache( |
| 1746 | num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device |
| 1747 | ) |
| 1748 | ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device) |
| 1749 | |
| 1750 | for i in range(num_tokens): |
| 1751 | slot = slot_mapping[i].item() |
| 1752 | block_idx = slot // block_size |
| 1753 | block_offset = slot % block_size |
| 1754 | ref_temp[block_idx, block_offset, :kv_lora_rank] = kv_c[i] |
| 1755 | ref_temp[block_idx, block_offset, kv_lora_rank:] = k_pe[i] |
| 1756 | |
| 1757 | if kv_cache_dtype == "fp8": |
| 1758 | ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype) |
| 1759 | ops.convert_fp8(ref_kv_cache, ref_temp, scale.item(), kv_dtype=kv_cache_dtype) |
| 1760 | else: |
| 1761 | ref_kv_cache = ref_temp |
| 1762 | |
| 1763 | opcheck( |
| 1764 | torch.ops._C_cache_ops.concat_and_cache_mla, |
| 1765 | (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale), |
| 1766 | test_utils=DEFAULT_OPCHECK_TEST_UTILS, |
| 1767 | ) |
| 1768 | |
| 1769 | ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale) |
| 1770 | torch.testing.assert_close(kv_cache, ref_kv_cache) |