[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>

.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

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")

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

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