From bcca76e7dbca608979deea55f589c7316393bc52 Mon Sep 17 00:00:00 2001 From: Raya Elena Solano Date: Tue, 8 Sep 2026 19:13:39 +0200 Subject: [PATCH] [Test] Dequantize NVFP4 KV cache scales in the layout the kernel writes (#55908) Signed-off-by: Raya Elena Solano Co-authored-by: mgoin --- .buildkite/test_areas/kernels.yaml | 4 +++ tests/kernels/attention/test_cache.py | 21 +++++++++--- .../test_flashinfer_trtllm_attention.py | 2 +- tests/kernels/quantization/nvfp4_utils.py | 32 ++++++++++++------- 4 files changed, 42 insertions(+), 17 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4a19bc1795..2b43938060 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -312,6 +312,7 @@ steps: # optional: true source_file_dependencies: - csrc/quantization/fp4/ + - csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu - csrc/libtorch_stable/quantization/fp4/ - csrc/attention/mla/ - csrc/quantization/cutlass_w8a8/moe/ @@ -345,6 +346,8 @@ steps: - csrc/libtorch_stable/torch_bindings.cpp - tests/kernels/attention/test_flashinfer_mla_decode.py - tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py + - tests/kernels/attention/test_cache.py + - tests/kernels/quantization/nvfp4_utils.py - tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - tests/kernels/mamba/test_gdn_prefill_cutedsl.py - tests/kernels/test_bf16x3_router_gemm_cutedsl.py @@ -360,6 +363,7 @@ steps: case "$$BUILDKITE_PARALLEL_JOB" in 0) pytest -v -s tests/kernels/attention/test_attention_selector.py + pytest -v -s tests/kernels/attention/test_cache.py -k nvfp4 pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py pytest -v -s tests/kernels/test_top_k_per_row.py diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 69744cd7b4..6fabcaa789 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -356,22 +356,34 @@ def test_reshape_and_cache_flash( dequant_nvfp4_kv_cache, ) - def dequant_nvfp4_cache(data_cache, scale_cache, global_scale): + def dequant_nvfp4_cache(data_cache, scale_cache, global_scale, swizzled_scales): # data_cache: [B, N, H, data_dim] logical view (layout in strides) # scale_cache: [B, N, H, scale_dim] logical view (layout in strides) # Permute to [B, H, N, dim] for the dequant utility. data_hnd = data_cache.permute(0, 2, 1, 3) scale_hnd = scale_cache.permute(0, 2, 1, 3) result_hnd = dequant_nvfp4_kv_cache( - data_hnd, scale_hnd, global_scale, head_size, block_size + data_hnd, + scale_hnd, + global_scale, + head_size, + block_size, + swizzled_scales=swizzled_scales, ) return result_hnd.permute(0, 2, 1, 3) # back to [B, N, H, dim] + # The kernel writes K scales linearly on every arch and V scales in + # the 4x4 swizzle of the SM100 trtllm-gen reader. The FlashInfer + # reader used on SM12x expects linear V scales as well. + v_scales_swizzled = not current_platform.is_device_capability_family(120) result_key_cache = dequant_nvfp4_cache( - nvfp4_key_data, key_scale_cache, k_scale.item() + nvfp4_key_data, key_scale_cache, k_scale.item(), swizzled_scales=False ) result_value_cache = dequant_nvfp4_cache( - nvfp4_value_data, value_scale_cache, v_scale.item() + nvfp4_value_data, + value_scale_cache, + v_scale.item(), + swizzled_scales=v_scales_swizzled, ) # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping. @@ -484,6 +496,7 @@ def test_nvfp4_4over6_selects_lower_error_scale( scale.item(), head_size, block_size, + swizzled_scales=False, )[0, 0, 0] default = quantize_and_dequantize("nvfp4") diff --git a/tests/kernels/attention/test_flashinfer_trtllm_attention.py b/tests/kernels/attention/test_flashinfer_trtllm_attention.py index 55d11f99ba..48eb58b5e1 100644 --- a/tests/kernels/attention/test_flashinfer_trtllm_attention.py +++ b/tests/kernels/attention/test_flashinfer_trtllm_attention.py @@ -117,7 +117,7 @@ def make_nvfp4_kv_cache( # Dequantize for the FA2 reference baseline. ref_k = dequant_nvfp4_kv_cache( - k_data, k_scales, kv_scale_val, head_size, block_size + k_data, k_scales, kv_scale_val, head_size, block_size, swizzled_scales=False ).to(torch.bfloat16) ref_v = dequant_nvfp4_kv_cache( v_data, v_scales, kv_scale_val, head_size, block_size diff --git a/tests/kernels/quantization/nvfp4_utils.py b/tests/kernels/quantization/nvfp4_utils.py index 21d0bebf15..e0120301f7 100644 --- a/tests/kernels/quantization/nvfp4_utils.py +++ b/tests/kernels/quantization/nvfp4_utils.py @@ -97,8 +97,9 @@ def dequant_nvfp4_kv_cache( global_scale: float, head_size: int, block_size: int, + swizzled_scales: bool = True, ) -> torch.Tensor: - """Dequantize an NVFP4 KV cache with 4x4-swizzled block scales. + """Dequantize an NVFP4 KV cache. The input must be in HND layout so that the last two dims are (block_size, last_dim). For NHD caches, permute to HND first. @@ -110,6 +111,10 @@ def dequant_nvfp4_kv_cache( global_scale: checkpoint dequant scale (k_scale or v_scale). head_size: head dimension. block_size: page size. + swizzled_scales: whether ``block_scale`` is in the 4x4-swizzled + layout of the SM100 trtllm-gen kernel. ``reshape_and_cache_flash`` + writes K scales linearly and V scales swizzled, so pass ``False`` + for K and ``True`` for V. Returns: [..., num_heads, block_size, head_size] float32. @@ -118,18 +123,21 @@ def dequant_nvfp4_kv_cache( scale_dim = head_size // 16 fp4_packed = fp4_data - sf_swizzled = block_scale.view(torch.uint8) - - # Unswizzle 4x4 block scales on (block_size, scale_dim) plane. - # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S] - batch_shape = sf_swizzled.shape[:-2] + sf_bytes = block_scale.view(torch.uint8) + batch_shape = sf_bytes.shape[:-2] T, S = block_size, scale_dim - sg = S // 4 - sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4) - ndim = sf_reshape.ndim - # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg) - perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2] - sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S) + + if swizzled_scales: + # Unswizzle 4x4 block scales on (block_size, scale_dim) plane. + # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S] + sg = S // 4 + sf_reshape = sf_bytes.reshape(*batch_shape, T // 4, 4, sg, 4) + ndim = sf_reshape.ndim + # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg) + perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2] + sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S) + else: + sf_linear = sf_bytes sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32) # Unpack fp4