| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit test for fused_indexer_q_rope_quant. |
| 4 | |
| 5 | Compares the fused Triton kernel against the unfused reference flow used by |
| 6 | the DeepseekV4 indexer in model_tracking: |
| 7 | q_rot = ops.rotary_embedding(positions, q, None, head_dim, cos_sin_cache, |
| 8 | is_neox_style=False, |
| 9 | rope_dim_offset=head_dim - rope_dim) |
| 10 | q_fp8, q_scale = per_token_group_quant_fp8(q_rot, head_dim, use_ue8m0=True) |
| 11 | weights_out = weights * q_scale * softmax_scale * head_scale |
| 12 | |
| 13 | Expects bit-exact equality on both q_fp8 and weights_out. |
| 14 | """ |
| 15 | |
| 16 | import contextlib |
| 17 | import importlib |
| 18 | from unittest import mock |
| 19 | |
| 20 | import pytest |
| 21 | import torch |
| 22 | |
| 23 | from vllm import _custom_ops as ops |
| 24 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 25 | per_token_group_quant_fp8, |
| 26 | ) |
| 27 | from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant |
| 28 | from vllm.platforms import current_platform |
| 29 | from vllm.utils.import_utils import has_cutedsl |
| 30 | |
| 31 | HEAD_DIM = 128 |
| 32 | ROPE_DIM = 64 |
| 33 | MAX_POS = 4096 |
| 34 | |
| 35 | CUTEDSL_MODULE = "vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl" |
| 36 | |
| 37 | |
| 38 | def quantize_to_mxfp4( |
| 39 | x: torch.Tensor, |
| 40 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 41 | """Reference MXFP4 quantization. |
| 42 | |
| 43 | Args: |
| 44 | x: [..., head_dim] where head_dim is divisible by 32 |
| 45 | Returns: |
| 46 | packed: [..., head_dim//2] uint8 2 E2M1 nibbles/byte, low nibble = even index |
| 47 | scales: [..., head_dim//32] uint8 1 ue8m0 byte |
| 48 | |
| 49 | """ |
| 50 | MXFP4_BLOCK_SIZE = 32 |
| 51 | orig_shape = x.shape |
| 52 | head_dim = orig_shape[-1] |
| 53 | n_blocks = head_dim // MXFP4_BLOCK_SIZE |
| 54 | |
| 55 | x_f32 = x.float().reshape(-1, n_blocks, MXFP4_BLOCK_SIZE) |
| 56 | |
| 57 | # Per-block ue8m0 scale: 2^ceil(log2(amax / 6.0)), stored as byte = exp + 127 |
| 58 | # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163 |
| 59 | amax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=6 * (2**-126)) |
| 60 | log2_ratio = (amax * (1.0 / 6.0)).log2().ceil().clamp(-127.0, 127.0) |
| 61 | scale = log2_ratio.exp2() |
| 62 | ue8m0 = (log2_ratio + 127.0).to(torch.uint8) # [*, n_blocks] |
| 63 | |
| 64 | # E2M1 round-to-nearest-even: midpoints round to the even code. |
| 65 | # E2M1 values: [0.00, 0.50, 1.00, 1.50, 2.00, 3.00, 4.00, 6.00] |
| 66 | # boundaries: [ 0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00] |
| 67 | x_scaled = (x_f32 / scale).clamp(-6.0, 6.0) |
| 68 | abs_x = x_scaled.abs() |
| 69 | code = torch.zeros_like(abs_x, dtype=torch.int32) |
| 70 | code = torch.where(abs_x > 0.25, 1, code) |
| 71 | code = torch.where(abs_x >= 0.75, 2, code) |
| 72 | code = torch.where(abs_x > 1.25, 3, code) |
| 73 | code = torch.where(abs_x >= 1.75, 4, code) |
| 74 | code = torch.where(abs_x > 2.5, 5, code) |
| 75 | code = torch.where(abs_x >= 3.5, 6, code) |
| 76 | code = torch.where(abs_x > 5.0, 7, code) |
| 77 | sign = ((x_scaled.view(torch.int32) >> 31) & 1).to(torch.uint8) |
| 78 | nibble = code.to(torch.uint8) | (sign << 3) |
| 79 | |
| 80 | # Pack: even-index element → low nibble, odd-index → high nibble |
| 81 | nibble_flat = nibble.reshape(-1, head_dim) |
| 82 | packed = (nibble_flat[:, 0::2] | (nibble_flat[:, 1::2] << 4)).contiguous() |
| 83 | packed = packed.reshape(*orig_shape[:-1], head_dim // 2) |
| 84 | |
| 85 | scales = ue8m0.view(*orig_shape[:-1], n_blocks) |
| 86 | return packed, scales |
| 87 | |
| 88 | |
| 89 | def _reference( |
| 90 | positions: torch.Tensor, |
| 91 | q: torch.Tensor, |
| 92 | cos_sin_cache: torch.Tensor, |
| 93 | weights: torch.Tensor, |
| 94 | softmax_scale: float, |
| 95 | head_scale: float, |
| 96 | n_head: int, |
| 97 | use_fp4: bool = False, |
| 98 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 99 | q_rot = q.clone() |
| 100 | ops.rotary_embedding( |
| 101 | positions, |
| 102 | q_rot, |
| 103 | None, |
| 104 | HEAD_DIM, |
| 105 | cos_sin_cache, |
| 106 | False, # is_neox_style=False → GPT-J interleaved |
| 107 | HEAD_DIM - ROPE_DIM, # rope_dim_offset → rotate the tail |
| 108 | False, |
| 109 | ) |
| 110 | |
| 111 | if use_fp4: |
| 112 | q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, n_head, HEAD_DIM)) |
| 113 | # Pack 4 ue8m0 bytes into 1 int32 |
| 114 | q_scale = ue8m0.view(torch.int32).squeeze(-1) |
| 115 | # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar) |
| 116 | weights_out = weights.to(torch.float32) * softmax_scale * head_scale |
| 117 | return (q_packed, q_scale), weights_out |
| 118 | |
| 119 | else: |
| 120 | q_fp8, q_scale = per_token_group_quant_fp8( |
| 121 | q_rot.view(-1, HEAD_DIM).contiguous(), |
| 122 | HEAD_DIM, |
| 123 | use_ue8m0=True, |
| 124 | ) |
| 125 | q_fp8 = q_fp8.view(-1, n_head, HEAD_DIM) |
| 126 | q_scale = q_scale.view(-1, n_head) |
| 127 | |
| 128 | weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale |
| 129 | return q_fp8, weights_out |
| 130 | |
| 131 | |
| 132 | @pytest.mark.parametrize("num_tokens", [1, 7, 32, 257, 1023]) |
| 133 | @pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16]) |
| 134 | @pytest.mark.parametrize( |
| 135 | "use_fp4", |
| 136 | [ |
| 137 | False, |
| 138 | pytest.param( |
| 139 | True, |
| 140 | marks=pytest.mark.skipif( |
| 141 | not ( |
| 142 | current_platform.is_cuda() |
| 143 | and current_platform.is_device_capability_family(100) |
| 144 | ), |
| 145 | reason="MXFP4 indexer cache requires an SM100-family GPU", |
| 146 | ), |
| 147 | ), |
| 148 | ], |
| 149 | ) |
| 150 | @pytest.mark.parametrize("use_cutedsl", [False, True]) |
| 151 | @pytest.mark.parametrize("n_head", [32, 64]) |
| 152 | @torch.inference_mode() |
| 153 | def test_fused_indexer_q_rope_quant_matches_unfused( |
| 154 | num_tokens, cache_dtype, use_fp4, use_cutedsl, n_head |
| 155 | ): |
| 156 | if use_cutedsl and not has_cutedsl(): |
| 157 | pytest.skip("cutedsl (cutlass) not installed") |
| 158 | |
| 159 | device = "cuda" |
| 160 | torch.manual_seed(0) |
| 161 | |
| 162 | q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device) |
| 163 | positions = torch.randint( |
| 164 | 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device |
| 165 | ) |
| 166 | cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=cache_dtype, device=device) |
| 167 | weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device) |
| 168 | softmax_scale = HEAD_DIM**-0.5 |
| 169 | # head_scale must be an exact power of two for bit-exactness to be a fair |
| 170 | # expectation: the triton kernel multiplies by softmax_scale and head_scale |
| 171 | # separately (two roundings) while the cutedsl kernels fold them into one |
| 172 | # fp32 scalar first (one rounding). Those agree only when one factor is a |
| 173 | # pow2. n_head**-0.5 is a pow2 at 64 heads but not at 32. |
| 174 | head_scale = 0.125 |
| 175 | |
| 176 | q_quant_ref, weights_ref = _reference( |
| 177 | positions, |
| 178 | q, |
| 179 | cos_sin_cache, |
| 180 | weights, |
| 181 | softmax_scale, |
| 182 | head_scale, |
| 183 | n_head, |
| 184 | use_fp4, |
| 185 | ) |
| 186 | # use_cutedsl=False: force the triton path even when cutedsl is installed |
| 187 | # by patching the dispatcher's has_cutedsl() binding to return False. |
| 188 | cutedsl_patch = ( |
| 189 | mock.patch( |
| 190 | "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl", |
| 191 | return_value=False, |
| 192 | ) |
| 193 | if not use_cutedsl |
| 194 | else contextlib.nullcontext() |
| 195 | ) |
| 196 | with cutedsl_patch: |
| 197 | q_quant_fused, weights_fused = fused_indexer_q_rope_quant( |
| 198 | positions, |
| 199 | q.clone(), |
| 200 | cos_sin_cache, |
| 201 | weights, |
| 202 | softmax_scale, |
| 203 | head_scale, |
| 204 | use_fp4, |
| 205 | ) |
| 206 | |
| 207 | if use_fp4: |
| 208 | q_quant_ref, q_scale_ref = q_quant_ref |
| 209 | q_quant_fused, q_scale_fused = q_quant_fused |
| 210 | |
| 211 | assert torch.equal(q_scale_ref, q_scale_fused), ( |
| 212 | f"q_scale mismatch: " |
| 213 | f"{(q_scale_ref != q_scale_fused).sum().item()} " |
| 214 | f"/ {q_scale_ref.numel()} bytes differ" |
| 215 | ) |
| 216 | |
| 217 | # fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8. |
| 218 | ref_bits = q_quant_ref.view(torch.int8) |
| 219 | fused_bits = q_quant_fused.view(torch.int8) |
| 220 | assert torch.equal(ref_bits, fused_bits), ( |
| 221 | f"q_quant_fused mismatch: " |
| 222 | f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ" |
| 223 | ) |
| 224 | |
| 225 | assert weights_fused.dtype == torch.float32 |
| 226 | assert torch.equal(weights_ref, weights_fused), ( |
| 227 | f"weights mismatch: max abs diff " |
| 228 | f"{(weights_ref - weights_fused).abs().max().item()}" |
| 229 | ) |
| 230 | |
| 231 | |
| 232 | @pytest.mark.skipif(not has_cutedsl(), reason="cutedsl (cutlass) not installed") |
| 233 | @pytest.mark.parametrize( |
| 234 | "use_fp4", |
| 235 | [ |
| 236 | False, |
| 237 | pytest.param( |
| 238 | True, |
| 239 | marks=pytest.mark.skipif( |
| 240 | not ( |
| 241 | current_platform.is_cuda() |
| 242 | and current_platform.is_device_capability_family(100) |
| 243 | ), |
| 244 | reason="MXFP4 indexer cache requires an SM100-family GPU", |
| 245 | ), |
| 246 | ), |
| 247 | ], |
| 248 | ) |
| 249 | @pytest.mark.parametrize("n_head", [32, 64]) |
| 250 | @torch.inference_mode() |
| 251 | def test_cutedsl_indexer_q_writes_stay_within_num_tokens(use_fp4, n_head): |
| 252 | """The cutedsl kernels must not store past ``num_tokens``. |
| 253 | |
| 254 | The launcher rounds ``num_tokens * threads_per_token`` up to whole 128-thread |
| 255 | blocks, so the trailing block is only partially in range unless |
| 256 | ``threads_per_token`` is itself a multiple of the block size. With 32 heads |
| 257 | and coarsen=4 (picked for >= 512 tokens) a token needs 64 threads, so an odd |
| 258 | token count leaves half of the last block addressing token ``num_tokens``. |
| 259 | Every output gets one extra sentinel row that the kernel must leave alone. |
| 260 | """ |
| 261 | mod = importlib.import_module(CUTEDSL_MODULE) |
| 262 | device = "cuda" |
| 263 | torch.manual_seed(0) |
| 264 | num_tokens = 935 |
| 265 | |
| 266 | sentinel_u8, sentinel_f32 = 0x55, -12345.0 |
| 267 | # positions/q are padded too, so an out-of-bounds *load* stays addressable |
| 268 | # and the assertions below are about stores only. |
| 269 | positions = torch.zeros(num_tokens + 1, dtype=torch.int64, device=device) |
| 270 | q = torch.randn( |
| 271 | num_tokens + 1, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device |
| 272 | ) |
| 273 | cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device) |
| 274 | weights = torch.randn(num_tokens + 1, n_head, dtype=torch.bfloat16, device=device) |
| 275 | weights_out = torch.full( |
| 276 | (num_tokens + 1, n_head), sentinel_f32, dtype=torch.float32, device=device |
| 277 | ) |
| 278 | outputs = {"weights_out": weights_out} |
| 279 | |
| 280 | if use_fp4: |
| 281 | outputs["q_packed"] = torch.full( |
| 282 | (num_tokens + 1, n_head, HEAD_DIM // 2), |
| 283 | sentinel_u8, |
| 284 | dtype=torch.uint8, |
| 285 | device=device, |
| 286 | ) |
| 287 | outputs["q_scale"] = torch.full( |
| 288 | (num_tokens + 1, n_head, HEAD_DIM // 32), |
| 289 | sentinel_u8, |
| 290 | dtype=torch.uint8, |
| 291 | device=device, |
| 292 | ) |
| 293 | mod._INDEXER_Q_MXFP4_KERNEL( |
| 294 | positions=positions[:num_tokens], |
| 295 | q=q[:num_tokens], |
| 296 | cos_sin_cache=cos_sin_cache, |
| 297 | weights=weights[:num_tokens], |
| 298 | weights_softmax_scale=1.0, |
| 299 | weights_head_scale=1.0, |
| 300 | q_packed=outputs["q_packed"][:num_tokens], |
| 301 | q_scale=outputs["q_scale"][:num_tokens], |
| 302 | weights_out=weights_out[:num_tokens], |
| 303 | ) |
| 304 | else: |
| 305 | outputs["q_fp8"] = torch.full( |
| 306 | (num_tokens + 1, n_head, HEAD_DIM), |
| 307 | sentinel_u8, |
| 308 | dtype=torch.uint8, |
| 309 | device=device, |
| 310 | ) |
| 311 | mod._INDEXER_Q_FP8_KERNEL( |
| 312 | positions=positions[:num_tokens], |
| 313 | q=q[:num_tokens], |
| 314 | cos_sin_cache=cos_sin_cache, |
| 315 | weights=weights[:num_tokens], |
| 316 | weights_softmax_scale=1.0, |
| 317 | weights_head_scale=1.0, |
| 318 | q_fp8=outputs["q_fp8"][:num_tokens].view(torch.float8_e4m3fn), |
| 319 | weights_out=weights_out[:num_tokens], |
| 320 | ) |
| 321 | torch.accelerator.synchronize() |
| 322 | |
| 323 | clobbered = {} |
| 324 | for name, tensor in outputs.items(): |
| 325 | sentinel = sentinel_f32 if tensor.dtype == torch.float32 else sentinel_u8 |
| 326 | n_changed = (tensor[num_tokens] != sentinel).sum().item() |
| 327 | if n_changed: |
| 328 | clobbered[name] = n_changed |
| 329 | assert not clobbered, ( |
| 330 | f"kernel wrote past token {num_tokens - 1} into the sentinel row: {clobbered}" |
| 331 | ) |
| 332 | |
| 333 | |
| 334 | def _indexer_k_reference( |
| 335 | k_pre, positions, cos_sin_cache, rms_norm_weight, compress_ratio, use_fp4 |
| 336 | ): |
| 337 | k = k_pre.float() |
| 338 | k = k * torch.rsqrt(k.square().mean(dim=-1, keepdim=True) + 1e-20) |
| 339 | k = (k * rms_norm_weight.float()).to(torch.bfloat16).float() |
| 340 | group_positions = positions // compress_ratio * compress_ratio |
| 341 | cos, sin = cos_sin_cache[group_positions].float().chunk(2, dim=-1) |
| 342 | even, odd = k[:, 64::2], k[:, 65::2] |
| 343 | rotated = torch.stack((even * cos - odd * sin, odd * cos + even * sin), dim=-1) |
| 344 | k = torch.cat((k[:, :64], rotated.flatten(1)), dim=-1).to(torch.bfloat16) |
| 345 | if use_fp4: |
| 346 | return quantize_to_mxfp4(k) |
| 347 | exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448 |
| 348 | scale = exponent.log2().ceil().exp2() |
| 349 | values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn) |
| 350 | return values.view(torch.uint8), scale.view(torch.uint8) |
| 351 | |
| 352 | |
| 353 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel") |
| 354 | @pytest.mark.parametrize("num_tokens", [1, 17, 257]) |
| 355 | @pytest.mark.parametrize("compress_ratio", [1, 2]) |
| 356 | @pytest.mark.parametrize("use_fp4", [False, True]) |
| 357 | @pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16]) |
| 358 | @torch.inference_mode() |
| 359 | def test_indexer_k_inserts_only_valid_groups_into_padded_pages( |
| 360 | num_tokens, compress_ratio, use_fp4, cache_dtype |
| 361 | ): |
| 362 | """Insert group ends and preserve skipped slots, graph rows, and page padding.""" |
| 363 | from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( |
| 364 | indexer_k_norm_rope_store, |
| 365 | ) |
| 366 | |
| 367 | if use_fp4 and not current_platform.is_device_capability_family(100): |
| 368 | pytest.skip("MXFP4 indexer cache requires an SM100-family GPU") |
| 369 | torch.manual_seed(17) |
| 370 | device = "cuda" |
| 371 | k_pre = torch.randn(num_tokens + 3, 128, device=device, dtype=torch.bfloat16) |
| 372 | norm_weight = torch.randn(128, device=device, dtype=torch.bfloat16) |
| 373 | positions = torch.arange(7, num_tokens + 10, device=device) |
| 374 | slots = torch.randperm(2 * num_tokens, device=device)[:num_tokens].clone() |
| 375 | slots[5::7] = -1 |
| 376 | valid = (slots >= 0) & ((positions[:num_tokens] + 1) % compress_ratio == 0) |
| 377 | k_pre[3::7] = 0 |
| 378 | k_pre[8::11] *= 1e-10 |
| 379 | k_pre[:num_tokens][~valid] = float("nan") |
| 380 | k_pre[num_tokens:] = float("nan") |
| 381 | positions[num_tokens:] = MAX_POS + 10 |
| 382 | angles = torch.randn(MAX_POS, 32, device=device) |
| 383 | cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).to(cache_dtype) |
| 384 | block_size = 16 |
| 385 | value_bytes = 64 if use_fp4 else 128 |
| 386 | page_bytes = block_size * (value_bytes + 4) |
| 387 | num_blocks = (2 * num_tokens + block_size - 1) // block_size + 1 |
| 388 | storage = torch.full( |
| 389 | (num_blocks, page_bytes + 128), 0xA5, device=device, dtype=torch.uint8 |
| 390 | ) |
| 391 | cache = storage.as_strided( |
| 392 | (num_blocks, block_size, value_bytes + 4), |
| 393 | (storage.stride(0), value_bytes + 4, 1), |
| 394 | ) |
| 395 | expected = storage.clone() |
| 396 | values, scales = _indexer_k_reference( |
| 397 | k_pre[:num_tokens][valid], |
| 398 | positions[:num_tokens][valid], |
| 399 | cos_sin, |
| 400 | norm_weight, |
| 401 | compress_ratio, |
| 402 | use_fp4, |
| 403 | ) |
| 404 | block_ids = slots[valid] // block_size |
| 405 | block_offsets = slots[valid] % block_size |
| 406 | value_offsets = block_offsets[:, None] * value_bytes + torch.arange( |
| 407 | value_bytes, device=device |
| 408 | ) |
| 409 | scale_offsets = block_size * value_bytes + block_offsets[:, None] * 4 |
| 410 | scale_offsets = scale_offsets + torch.arange(4, device=device) |
| 411 | expected[block_ids[:, None], value_offsets] = values |
| 412 | expected[block_ids[:, None], scale_offsets] = scales |
| 413 | |
| 414 | indexer_k_norm_rope_store( |
| 415 | k_pre, |
| 416 | positions, |
| 417 | cos_sin, |
| 418 | norm_weight, |
| 419 | 1e-20, |
| 420 | cache, |
| 421 | slots, |
| 422 | compress_ratio, |
| 423 | use_fp4, |
| 424 | ) |
| 425 | # Fused RoPE can change zero signs without changing their numerical values. |
| 426 | actual_values = storage[block_ids[:, None], value_offsets] |
| 427 | if use_fp4: |
| 428 | for shift in (0, 4): |
| 429 | zero_mask = (actual_values >> shift) & 7 == 0 |
| 430 | actual_values &= ~(zero_mask.to(torch.uint8) << (shift + 3)) |
| 431 | zero_mask = (values >> shift) & 7 == 0 |
| 432 | values &= ~(zero_mask.to(torch.uint8) << (shift + 3)) |
| 433 | else: |
| 434 | actual_values[actual_values == 128] = 0 |
| 435 | values[values == 128] = 0 |
| 436 | storage[block_ids[:, None], value_offsets] = actual_values |
| 437 | expected[block_ids[:, None], value_offsets] = values |
| 438 | torch.testing.assert_close(storage, expected, rtol=0, atol=0) |
| 439 | |
| 440 | |
| 441 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA kernel") |
| 442 | @pytest.mark.parametrize("use_fp4", [False, True]) |
| 443 | @torch.inference_mode() |
| 444 | def test_indexer_k_cuda_graph_replay_reads_current_projection(use_fp4): |
| 445 | """Replay must consume updated keys and slot mapping through the captured API.""" |
| 446 | from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( |
| 447 | indexer_k_norm_rope_store, |
| 448 | ) |
| 449 | |
| 450 | if use_fp4 and not current_platform.is_device_capability_family(100): |
| 451 | pytest.skip("MXFP4 indexer cache requires an SM100-family GPU") |
| 452 | torch.manual_seed(18) |
| 453 | k_pre = torch.randn(19, 128, device="cuda", dtype=torch.bfloat16) |
| 454 | norm_weight = torch.ones(128, device="cuda", dtype=torch.bfloat16) |
| 455 | positions = torch.arange(19, device="cuda") |
| 456 | cos_sin = torch.zeros(32, 64, device="cuda") |
| 457 | cos_sin[:, :32] = 1 |
| 458 | slots = torch.arange(17, device="cuda") |
| 459 | cache = torch.full( |
| 460 | (2, 16, (64 if use_fp4 else 128) + 4), |
| 461 | 0xA5, |
| 462 | device="cuda", |
| 463 | dtype=torch.uint8, |
| 464 | ) |
| 465 | |
| 466 | def run(): |
| 467 | indexer_k_norm_rope_store( |
| 468 | k_pre, |
| 469 | positions, |
| 470 | cos_sin, |
| 471 | norm_weight, |
| 472 | 1e-20, |
| 473 | cache, |
| 474 | slots, |
| 475 | 2, |
| 476 | use_fp4, |
| 477 | ) |
| 478 | |
| 479 | run() |
| 480 | graph = torch.cuda.CUDAGraph() |
| 481 | with torch.cuda.graph(graph): |
| 482 | run() |
| 483 | k_pre.neg_() |
| 484 | slots.add_(4) |
| 485 | cache.fill_(0xA5) |
| 486 | graph.replay() |
| 487 | replayed = cache.clone() |
| 488 | cache.fill_(0xA5) |
| 489 | run() |
| 490 | torch.testing.assert_close(cache, replayed, rtol=0, atol=0) |
| 491 | |
| 492 | |
| 493 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm cache layout") |
| 494 | @pytest.mark.parametrize("block_size", [16, 64, 128]) |
| 495 | @pytest.mark.parametrize("compress_ratio", [1, 2]) |
| 496 | @torch.inference_mode() |
| 497 | def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_ratio): |
| 498 | """ROCm reads the indexer K cache back with the 16x16-tiled value layout. |
| 499 | |
| 500 | ``cp_gather_indexer_k_quant_cache_triton`` (prefill) and aiter's |
| 501 | ``Preshuffle=True`` paged kernel (decode) both pick that layout whenever |
| 502 | ``block_size > 1``, so a row-major store would feed the indexer permuted |
| 503 | key bytes and randomize its top-k. Assert the write/read pair is exact. |
| 504 | """ |
| 505 | from vllm.models.deepseek_v41.common.ops.indexer_k_store import ( |
| 506 | indexer_k_norm_rope_store, |
| 507 | ) |
| 508 | from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( |
| 509 | cp_gather_indexer_k_quant_cache_triton, |
| 510 | ) |
| 511 | |
| 512 | torch.manual_seed(0) |
| 513 | device = "cuda" |
| 514 | num_tokens = 3 * block_size + 5 |
| 515 | num_blocks = num_tokens // block_size + 2 |
| 516 | k_pre = torch.randn(num_tokens, HEAD_DIM, device=device, dtype=torch.bfloat16) |
| 517 | norm_weight = torch.randn(HEAD_DIM, device=device, dtype=torch.bfloat16) |
| 518 | positions = torch.arange(num_tokens, device=device) |
| 519 | angles = torch.randn(MAX_POS, ROPE_DIM // 2, device=device) |
| 520 | cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1).float() |
| 521 | cache = torch.zeros( |
| 522 | num_blocks, block_size, HEAD_DIM + 4, device=device, dtype=torch.uint8 |
| 523 | ) |
| 524 | slots = torch.arange(num_tokens, device=device) |
| 525 | |
| 526 | indexer_k_norm_rope_store( |
| 527 | k_pre, |
| 528 | positions, |
| 529 | cos_sin, |
| 530 | norm_weight, |
| 531 | 1e-20, |
| 532 | cache, |
| 533 | slots, |
| 534 | compress_ratio, |
| 535 | False, |
| 536 | ) |
| 537 | |
| 538 | k_fp8 = torch.zeros(num_tokens, HEAD_DIM, device=device, dtype=torch.float8_e4m3fn) |
| 539 | k_scale = torch.zeros(num_tokens, 4, device=device, dtype=torch.uint8) |
| 540 | cp_gather_indexer_k_quant_cache_triton( |
| 541 | cache, |
| 542 | k_fp8, |
| 543 | k_scale, |
| 544 | torch.arange(num_blocks, device=device, dtype=torch.int32).view(1, -1), |
| 545 | torch.tensor([0, num_tokens], device=device, dtype=torch.int32), |
| 546 | torch.zeros(num_tokens, device=device, dtype=torch.int32), |
| 547 | ) |
| 548 | |
| 549 | expected_values, expected_scale = _indexer_k_reference( |
| 550 | k_pre, positions, cos_sin, norm_weight, compress_ratio, False |
| 551 | ) |
| 552 | emitted = (positions + 1) % compress_ratio == 0 |
| 553 | gathered = k_fp8.view(torch.uint8) |
| 554 | gathered = torch.where(gathered == 128, gathered.new_zeros(()), gathered) |
| 555 | expected_values = torch.where( |
| 556 | expected_values == 128, expected_values.new_zeros(()), expected_values |
| 557 | ) |
| 558 | torch.testing.assert_close( |
| 559 | gathered[emitted], expected_values[emitted], rtol=0, atol=0 |
| 560 | ) |
| 561 | torch.testing.assert_close( |
| 562 | k_scale[emitted], expected_scale[emitted], rtol=0, atol=0 |
| 563 | ) |
| 564 | |
| 565 | |
| 566 | @pytest.mark.parametrize("use_cutedsl", [False, True]) |
| 567 | @pytest.mark.skipif( |
| 568 | not ( |
| 569 | current_platform.is_cuda() and current_platform.is_device_capability_family(100) |
| 570 | ), |
| 571 | reason="MXFP4 indexer cache requires an SM100-family GPU", |
| 572 | ) |
| 573 | @torch.inference_mode() |
| 574 | def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): |
| 575 | """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's |
| 576 | sparse MQA-logits kernels take) instead of fp32; the values are the fp32 |
| 577 | result rounded once, so the scoring kernel needs no cast.""" |
| 578 | if use_cutedsl and not has_cutedsl(): |
| 579 | pytest.skip("cutedsl (cutlass) not installed") |
| 580 | |
| 581 | device = "cuda" |
| 582 | torch.manual_seed(0) |
| 583 | num_tokens, n_head = 257, 32 |
| 584 | q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device) |
| 585 | positions = torch.randint( |
| 586 | 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device |
| 587 | ) |
| 588 | cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device) |
| 589 | weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device) |
| 590 | softmax_scale, head_scale = HEAD_DIM**-0.5, 0.125 # pow2 head_scale, see above |
| 591 | |
| 592 | _, weights_ref = _reference( |
| 593 | positions, q, cos_sin_cache, weights, softmax_scale, head_scale, n_head, True |
| 594 | ) |
| 595 | cutedsl_patch = ( |
| 596 | mock.patch( |
| 597 | "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl", |
| 598 | return_value=False, |
| 599 | ) |
| 600 | if not use_cutedsl |
| 601 | else contextlib.nullcontext() |
| 602 | ) |
| 603 | with cutedsl_patch: |
| 604 | _, weights_fused = fused_indexer_q_rope_quant( |
| 605 | positions, |
| 606 | q.clone(), |
| 607 | cos_sin_cache, |
| 608 | weights, |
| 609 | softmax_scale, |
| 610 | head_scale, |
| 611 | use_fp4=True, |
| 612 | weights_out_dtype=torch.bfloat16, |
| 613 | ) |
| 614 | assert weights_fused.dtype == torch.bfloat16 |
| 615 | assert torch.equal(weights_fused, weights_ref.to(torch.bfloat16)) |