[Test] Dequantize NVFP4 KV cache scales in the layout the kernel writes (#55908)
Signed-off-by: Raya Elena Solano <raya.solano@mbinf.de> Co-authored-by: mgoin <mgoin64@gmail.com>
Raya Elena Solano committed
Sep 8, 2026 at 19:13 UTC
bcca76e7dbca608979deea55f589c7316393bc52
4 files changed
+42
-17
.buildkite/test_areas/kernels.yaml
+4
@@ -312,6 +312,7 @@ steps:
312
# optional: true
313
source_file_dependencies:
314
- csrc/quantization/fp4/
315
+ - csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu
316
- csrc/libtorch_stable/quantization/fp4/
317
- csrc/attention/mla/
318
- csrc/quantization/cutlass_w8a8/moe/
@@ -345,6 +346,8 @@ steps:
346
- csrc/libtorch_stable/torch_bindings.cpp
347
- tests/kernels/attention/test_flashinfer_mla_decode.py
348
- tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
349
+ - tests/kernels/attention/test_cache.py
350
+ - tests/kernels/quantization/nvfp4_utils.py
351
- tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
352
- tests/kernels/mamba/test_gdn_prefill_cutedsl.py
353
- tests/kernels/test_bf16x3_router_gemm_cutedsl.py
@@ -360,6 +363,7 @@ steps:
363
case "$$BUILDKITE_PARALLEL_JOB" in
364
0)
365
pytest -v -s tests/kernels/attention/test_attention_selector.py
366
+ pytest -v -s tests/kernels/attention/test_cache.py -k nvfp4
367
pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
368
pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
369
pytest -v -s tests/kernels/test_top_k_per_row.py
tests/kernels/attention/test_cache.py
+17
-4
@@ -356,22 +356,34 @@ def test_reshape_and_cache_flash(
356
dequant_nvfp4_kv_cache,
357
)
358
359
- def dequant_nvfp4_cache(data_cache, scale_cache, global_scale):
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, scale_hnd, global_scale, head_size, block_size
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(
371
- nvfp4_key_data, key_scale_cache, k_scale.item()
380
+ nvfp4_key_data, key_scale_cache, k_scale.item(), swizzled_scales=False
381
)
382
result_value_cache = dequant_nvfp4_cache(
374
- nvfp4_value_data, value_scale_cache, v_scale.item()
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.
@@ -484,6 +496,7 @@ def test_nvfp4_4over6_selects_lower_error_scale(
496
scale.item(),
497
head_size,
498
block_size,
499
+ swizzled_scales=False,
500
)[0, 0, 0]
501
502
default = quantize_and_dequantize("nvfp4")
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+1
-1
@@ -117,7 +117,7 @@ def make_nvfp4_kv_cache(
117
118
# Dequantize for the FA2 reference baseline.
119
ref_k = dequant_nvfp4_kv_cache(
120
- k_data, k_scales, kv_scale_val, head_size, block_size
120
+ k_data, k_scales, kv_scale_val, head_size, block_size, swizzled_scales=False
121
).to(torch.bfloat16)
122
ref_v = dequant_nvfp4_kv_cache(
123
v_data, v_scales, kv_scale_val, head_size, block_size
tests/kernels/quantization/nvfp4_utils.py
+20
-12
@@ -97,8 +97,9 @@ def dequant_nvfp4_kv_cache(
97
global_scale: float,
98
head_size: int,
99
block_size: int,
100
+ swizzled_scales: bool = True,
101
) -> torch.Tensor:
101
- """Dequantize an NVFP4 KV cache with 4x4-swizzled block scales.
102
+ """Dequantize an NVFP4 KV cache.
103
104
The input must be in HND layout so that the last two dims are
105
(block_size, last_dim). For NHD caches, permute to HND first.
@@ -110,6 +111,10 @@ def dequant_nvfp4_kv_cache(
111
global_scale: checkpoint dequant scale (k_scale or v_scale).
112
head_size: head dimension.
113
block_size: page size.
114
+ swizzled_scales: whether ``block_scale`` is in the 4x4-swizzled
115
+ layout of the SM100 trtllm-gen kernel. ``reshape_and_cache_flash``
116
+ writes K scales linearly and V scales swizzled, so pass ``False``
117
+ for K and ``True`` for V.
118
119
Returns:
120
[..., num_heads, block_size, head_size] float32.
@@ -118,18 +123,21 @@ def dequant_nvfp4_kv_cache(
123
scale_dim = head_size // 16
124
125
fp4_packed = fp4_data
121
- sf_swizzled = block_scale.view(torch.uint8)
122
-
123
- # Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
124
- # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
125
- batch_shape = sf_swizzled.shape[:-2]
126
+ sf_bytes = block_scale.view(torch.uint8)
127
+ batch_shape = sf_bytes.shape[:-2]
128
T, S = block_size, scale_dim
127
- sg = S // 4
128
- sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
129
- ndim = sf_reshape.ndim
130
- # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
131
- perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
132
- sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
129
+
130
+ if swizzled_scales:
131
+ # Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
132
+ # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
133
+ sg = S // 4
134
+ sf_reshape = sf_bytes.reshape(*batch_shape, T // 4, 4, sg, 4)
135
+ ndim = sf_reshape.ndim
136
+ # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
137
+ perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
138
+ sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
139
+ else:
140
+ sf_linear = sf_bytes
141
sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32)
142
143
# Unpack fp4