Add NVFP4 KV 4-over-6 scale search (#45187)
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Wei-Ming Chen committed
Aug 7, 2026 at 13:22 UTC
1c94e8dc7da43d12c8686f05cc5d2eab3fd9528f
11 files changed
+284
-42
csrc/libtorch_stable/cache_kernels.cu
+4
-3
@@ -770,16 +770,17 @@ void reshape_and_cache_flash(
770
key.get_device_index());
771
const cudaStream_t stream = get_current_cuda_stream();
772
773
- if (kv_cache_dtype == "nvfp4") {
773
+ if (kv_cache_dtype == "nvfp4" || kv_cache_dtype == "nvfp4_4over6") {
774
#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
775
// NVFP4 dispatch is compiled separately for SM100+.
776
extern void reshape_and_cache_nvfp4_dispatch(
777
torch::stable::Tensor & key, torch::stable::Tensor & value,
778
torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache,
779
torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale,
780
- torch::stable::Tensor & v_scale);
780
+ torch::stable::Tensor & v_scale, const std::string& kv_cache_dtype);
781
reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
782
- slot_mapping, k_scale, v_scale);
782
+ slot_mapping, k_scale, v_scale,
783
+ kv_cache_dtype);
784
return;
785
#else
786
STD_TORCH_CHECK(
csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu
+167
-21
@@ -20,8 +20,17 @@
20
#include "libtorch_stable/dispatch_utils.h"
21
#include "libtorch_stable/torch_utils.h"
22
23
+#include <cmath>
24
+#include <string>
25
+
26
namespace vllm {
27
28
+// The cache layout is identical; this only selects the store-time scale search.
29
+enum class NVFP4KVScaleSearch {
30
+ DEFAULT,
31
+ FOUR_OVER_SIX,
32
+};
33
+
34
// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
35
// The swizzle pattern for HND layout is:
36
// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
@@ -38,6 +47,118 @@ __device__ __forceinline__ int swizzle_scale_offset(int t, int s,
47
return swizzled_t * scale_dim + swizzled_s;
48
}
49
50
+__device__ __forceinline__ float round_to_nearest_e2m1(float x) {
51
+ const float ax = fabsf(x);
52
+ float q;
53
+ // Match cvt.rn.satfinite.e2m1x2.f32, including ties-to-even boundaries.
54
+ if (ax <= 0.25f) {
55
+ q = 0.0f;
56
+ } else if (ax < 0.75f) {
57
+ q = 0.5f;
58
+ } else if (ax <= 1.25f) {
59
+ q = 1.0f;
60
+ } else if (ax < 1.75f) {
61
+ q = 1.5f;
62
+ } else if (ax <= 2.5f) {
63
+ q = 2.0f;
64
+ } else if (ax < 3.5f) {
65
+ q = 3.0f;
66
+ } else if (ax <= 5.0f) {
67
+ q = 4.0f;
68
+ } else {
69
+ q = 6.0f;
70
+ }
71
+ return copysignf(q, x);
72
+}
73
+
74
+__device__ __forceinline__ void nvfp4_candidate_scale(float SFScaleVal,
75
+ float vecMax,
76
+ float denominator,
77
+ uint8_t* fp8_sf,
78
+ float* outputScale) {
79
+ float SFValue =
80
+ SFScaleVal * (vecMax * reciprocal_approximate_ftz(denominator));
81
+ __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
82
+ reinterpret_cast<__nv_fp8_e4m3&>(*fp8_sf) = tmp;
83
+ SFValue = float(tmp);
84
+ *outputScale = SFValue != 0.0f
85
+ ? reciprocal_approximate_ftz(
86
+ SFValue * reciprocal_approximate_ftz(SFScaleVal))
87
+ : 0.0f;
88
+}
89
+
90
+template <class Type, int CVT_FP4_NUM_THREADS_PER_SF>
91
+__device__ __forceinline__ float nvfp4_reconstruction_error(
92
+ PackedVec<Type, CVT_FP4_PACK16>& vec, float outputScale) {
93
+ float error = 0.0f;
94
+ const float dequantScale =
95
+ outputScale != 0.0f ? reciprocal_approximate_ftz(outputScale) : 0.0f;
96
+
97
+#pragma unroll
98
+ for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
99
+ float2 vals = cast_to_float2(vec.elts[i]);
100
+ if (outputScale == 0.0f) {
101
+ error += vals.x * vals.x + vals.y * vals.y;
102
+ } else {
103
+ float qx = round_to_nearest_e2m1(vals.x * outputScale);
104
+ float qy = round_to_nearest_e2m1(vals.y * outputScale);
105
+ float dx = qx * dequantScale - vals.x;
106
+ float dy = qy * dequantScale - vals.y;
107
+ error += dx * dx + dy * dy;
108
+ }
109
+ }
110
+
111
+ if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) {
112
+ error += __shfl_xor_sync(0xffffffffu, error, 1);
113
+ }
114
+ return error;
115
+}
116
+
117
+template <class Type, int CVT_FP4_NUM_THREADS_PER_SF>
118
+__device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4_4over6(
119
+ PackedVec<Type, CVT_FP4_PACK16>& vec, float SFScaleVal, uint8_t* SFout) {
120
+ auto localMax = __habs2(vec.elts[0]);
121
+
122
+#pragma unroll
123
+ for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
124
+ localMax = __hmax2(localMax, __habs2(vec.elts[i]));
125
+ }
126
+
127
+ if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) {
128
+ localMax = __hmax2(__shfl_xor_sync(0xffffffffu, localMax, 1), localMax);
129
+ }
130
+ const float vecMax = float(__hmax(localMax.x, localMax.y));
131
+
132
+ uint8_t sf6;
133
+ uint8_t sf4;
134
+ float outputScale6;
135
+ float outputScale4;
136
+ nvfp4_candidate_scale(SFScaleVal, vecMax, 6.0f, &sf6, &outputScale6);
137
+ nvfp4_candidate_scale(SFScaleVal, vecMax, 4.0f, &sf4, &outputScale4);
138
+
139
+ const float err6 =
140
+ nvfp4_reconstruction_error<Type, CVT_FP4_NUM_THREADS_PER_SF>(
141
+ vec, outputScale6);
142
+ const float err4 =
143
+ nvfp4_reconstruction_error<Type, CVT_FP4_NUM_THREADS_PER_SF>(
144
+ vec, outputScale4);
145
+ const bool use4 = err4 < err6;
146
+ const float outputScale = use4 ? outputScale4 : outputScale6;
147
+
148
+ if (SFout) *SFout = use4 ? sf4 : sf6;
149
+
150
+ float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
151
+
152
+#pragma unroll
153
+ for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
154
+ fp2Vals[i] = cast_to_float2(vec.elts[i]);
155
+ fp2Vals[i].x *= outputScale;
156
+ fp2Vals[i].y *= outputScale;
157
+ }
158
+
159
+ return pack_fp4(fp2Vals);
160
+}
161
+
162
// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
163
//
164
// Takes separate data and scale cache pointers for K and V.
@@ -45,7 +166,7 @@ __device__ __forceinline__ int swizzle_scale_offset(int t, int s,
166
//
167
// Threading: one CUDA block per token, threads process heads and
168
// groups of 16 elements within each head.
48
-template <typename scalar_t>
169
+template <typename scalar_t, NVFP4KVScaleSearch SCALE_SEARCH>
170
__global__ void reshape_and_cache_nvfp4_kernel(
171
const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
172
const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
@@ -129,8 +250,14 @@ __global__ void reshape_and_cache_nvfp4_kernel(
250
uint8_t sf_val;
251
uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
252
132
- fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
133
- in_vec, global_scale, sf_out_ptr);
253
+ fp4_packed_t packed;
254
+ if constexpr (SCALE_SEARCH == NVFP4KVScaleSearch::FOUR_OVER_SIX) {
255
+ packed = cvt_warp_fp16_to_fp4_4over6<CudaType, THREADS_PER_SF>(
256
+ in_vec, global_scale, sf_out_ptr);
257
+ } else {
258
+ packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
259
+ in_vec, global_scale, sf_out_ptr);
260
+ }
261
262
// Write packed FP4 data to data cache
263
uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
@@ -181,13 +308,11 @@ __global__ void reshape_and_cache_nvfp4_kernel(
308
// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
309
// Each KV side contains both data and scale:
310
// page = [K_data | K_scale | V_data | V_scale]
184
-void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
185
- torch::stable::Tensor& value,
186
- torch::stable::Tensor& key_cache,
187
- torch::stable::Tensor& value_cache,
188
- torch::stable::Tensor& slot_mapping,
189
- torch::stable::Tensor& k_scale,
190
- torch::stable::Tensor& v_scale) {
311
+void reshape_and_cache_nvfp4_dispatch(
312
+ torch::stable::Tensor& key, torch::stable::Tensor& value,
313
+ torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
314
+ torch::stable::Tensor& slot_mapping, torch::stable::Tensor& k_scale,
315
+ torch::stable::Tensor& v_scale, const std::string& kv_cache_dtype) {
316
int num_tokens = slot_mapping.size(0);
317
int num_heads = key.size(1);
318
int head_size = key.size(2);
@@ -262,16 +387,37 @@ void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
387
388
VLLM_STABLE_DISPATCH_HALF_TYPES(
389
key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
265
- vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
266
- <<<grid, block, 0, stream>>>(
267
- key.const_data_ptr<scalar_t>(),
268
- value.const_data_ptr<scalar_t>(),
269
- key_cache.mutable_data_ptr<uint8_t>(),
270
- value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
271
- value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
272
- k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0),
273
- num_heads, head_size, block_size, data_block_stride,
274
- data_head_stride, data_block_offset_stride, scale_block_stride,
275
- scale_head_stride, scale_block_offset_stride);
390
+ if (kv_cache_dtype == "nvfp4") {
391
+ vllm::reshape_and_cache_nvfp4_kernel<
392
+ scalar_t, vllm::NVFP4KVScaleSearch::DEFAULT>
393
+ <<<grid, block, 0, stream>>>(
394
+ key.const_data_ptr<scalar_t>(),
395
+ value.const_data_ptr<scalar_t>(),
396
+ key_cache.mutable_data_ptr<uint8_t>(),
397
+ value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
398
+ value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
399
+ k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0),
400
+ num_heads, head_size, block_size, data_block_stride,
401
+ data_head_stride, data_block_offset_stride,
402
+ scale_block_stride, scale_head_stride,
403
+ scale_block_offset_stride);
404
+ } else if (kv_cache_dtype == "nvfp4_4over6") {
405
+ vllm::reshape_and_cache_nvfp4_kernel<
406
+ scalar_t, vllm::NVFP4KVScaleSearch::FOUR_OVER_SIX>
407
+ <<<grid, block, 0, stream>>>(
408
+ key.const_data_ptr<scalar_t>(),
409
+ value.const_data_ptr<scalar_t>(),
410
+ key_cache.mutable_data_ptr<uint8_t>(),
411
+ value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
412
+ value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
413
+ k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0),
414
+ num_heads, head_size, block_size, data_block_stride,
415
+ data_head_stride, data_block_offset_stride,
416
+ scale_block_stride, scale_head_stride,
417
+ scale_block_offset_stride);
418
+ } else {
419
+ STD_TORCH_CHECK(false,
420
+ "Unsupported NVFP4 KV cache dtype: ", kv_cache_dtype);
421
+ }
422
});
423
}
tests/kernels/attention/test_attention_selector.py
+1
@@ -528,6 +528,7 @@ def test_non_causal_autoselect_backend():
528
"fp8_ds_mla",
529
"fp8_inc",
530
"nvfp4",
531
+ "nvfp4_4over6",
532
"fp8_per_token_head",
533
"int8_per_token_head",
534
],
tests/kernels/attention/test_cache.py
+67
@@ -426,6 +426,73 @@ def test_reshape_and_cache_flash(
426
torch.testing.assert_close(value_cache_compact, cloned_value_cache)
427
428
429
+@torch.inference_mode()
430
+def test_nvfp4_4over6_selects_lower_error_scale(
431
+ kv_cache_factory_flashinfer,
432
+) -> None:
433
+ if not current_platform.has_device_capability(100):
434
+ pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")
435
+
436
+ device = CUDA_DEVICES[0]
437
+ block_size = 16
438
+ head_size = 64
439
+ source = torch.tensor(
440
+ [6.0] + [4.5] * 15,
441
+ dtype=torch.bfloat16,
442
+ device=device,
443
+ ).repeat(1, 1, head_size // 16)
444
+ slot_mapping = torch.zeros(1, dtype=torch.long, device=device)
445
+ scale = torch.ones(1, dtype=torch.float32, device=device)
446
+
447
+ def quantize_and_dequantize(cache_dtype: str) -> torch.Tensor:
448
+ key_caches, value_caches = kv_cache_factory_flashinfer(
449
+ 1,
450
+ block_size,
451
+ 1,
452
+ 1,
453
+ head_size,
454
+ cache_dtype,
455
+ source.dtype,
456
+ device=device,
457
+ cache_layout="NHD",
458
+ )
459
+ key_cache, value_cache = key_caches[0], value_caches[0]
460
+ key_cache.zero_()
461
+ value_cache.zero_()
462
+
463
+ ops.reshape_and_cache_flash(
464
+ source,
465
+ source,
466
+ key_cache,
467
+ value_cache,
468
+ slot_mapping,
469
+ cache_dtype,
470
+ scale,
471
+ scale,
472
+ )
473
+
474
+ from tests.kernels.quantization.nvfp4_utils import (
475
+ dequant_nvfp4_kv_cache,
476
+ )
477
+
478
+ data, block_scales = nvfp4_split_data_scale(key_cache)
479
+ return dequant_nvfp4_kv_cache(
480
+ data.permute(0, 2, 1, 3),
481
+ block_scales.permute(0, 2, 1, 3),
482
+ scale.item(),
483
+ head_size,
484
+ block_size,
485
+ )[0, 0, 0]
486
+
487
+ default = quantize_and_dequantize("nvfp4")
488
+ four_over_six = quantize_and_dequantize("nvfp4_4over6")
489
+ source_f32 = source[0, 0].float()
490
+
491
+ default_mse = (default - source_f32).square().mean()
492
+ four_over_six_mse = (four_over_six - source_f32).square().mean()
493
+ assert four_over_six_mse < default_mse
494
+
495
+
496
@pytest.mark.parametrize("dtype", DTYPES)
497
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
498
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
tests/utils_/test_torch_utils.py
+14
@@ -8,12 +8,26 @@ from vllm.utils.torch_utils import (
8
available_cpu_count,
9
common_broadcastable_dtype,
10
current_stream,
11
+ get_kv_cache_torch_dtype,
12
is_lossless_cast,
13
+ is_quantized_kv_cache,
14
set_torch_threads_for_runtime,
15
startup_omp_num_threads,
16
)
17
18
19
+def test_nvfp4_4over6_cache_dtype() -> None:
20
+ from vllm.config.cache import CacheConfig
21
+ from vllm.v1.kv_cache_interface import KVQuantMode, get_kv_quant_mode
22
+
23
+ cache_config = CacheConfig(cache_dtype="nvfp4_4over6")
24
+
25
+ assert cache_config.cache_dtype == "nvfp4_4over6"
26
+ assert get_kv_cache_torch_dtype(cache_config.cache_dtype) == torch.uint8
27
+ assert is_quantized_kv_cache(cache_config.cache_dtype)
28
+ assert get_kv_quant_mode(cache_config.cache_dtype) == KVQuantMode.NVFP4
29
+
30
+
31
@pytest.mark.parametrize(
32
("src_dtype", "tgt_dtype", "expected_result"),
33
[
vllm/config/cache.py
+3
@@ -33,6 +33,7 @@ CacheDType = Literal[
33
"int8_per_token_head",
34
"fp8_per_token_head",
35
"nvfp4",
36
+ "nvfp4_4over6",
37
]
38
MambaDType = Literal["auto", "float32", "float16", "bfloat16"]
39
MambaCacheMode = Literal["all", "align", "none"]
@@ -80,6 +81,8 @@ class CacheConfig:
81
Some models (namely DeepSeekV3.2) default to fp8, set to bfloat16 to use
82
bfloat16 instead, this is an invalid option for models that do not default
83
to fp8.
84
+ "nvfp4_4over6" uses the NVFP4 layout and selects between max/6 and max/4
85
+ scales per 16 values by minimizing squared reconstruction error.
86
"""
87
is_attention_free: bool = False
88
"""Whether the model is attention-free. This is primarily set in
vllm/config/vllm.py
+4
-1
@@ -2428,7 +2428,10 @@ class VllmConfig:
2428
def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig":
2429
if self.model_config is None:
2430
return self
2431
- if self.cache_config.cache_dtype == "nvfp4" and self.model_config.use_mla:
2431
+ if (
2432
+ self.cache_config.cache_dtype.startswith("nvfp4")
2433
+ and self.model_config.use_mla
2434
+ ):
2435
raise ValueError(
2436
"nvfp4 KV cache is not supported with MLA (Multi-head Latent "
2437
"Attention) backends. Please use a different --kv-cache-dtype "
vllm/model_executor/layers/attention/attention.py
+5
-2
@@ -459,7 +459,8 @@ class Attention(nn.Module, AttentionLayerBase):
459
if (
460
self.impl.supports_quant_query_input
461
and (
462
- self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4"
462
+ self.kv_cache_dtype.startswith("fp8")
463
+ or self.kv_cache_dtype.startswith("nvfp4")
464
)
465
and not self.kv_cache_dtype.endswith("per_token_head")
466
):
@@ -502,7 +503,9 @@ class Attention(nn.Module, AttentionLayerBase):
503
# which reduces overheads during decoding.
504
# Otherwise queries are quantized using custom ops
505
# which causes decoding overheads
505
- assert self.kv_cache_dtype in {"fp8", "fp8_e4m3", "nvfp4"}
506
+ assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"} or (
507
+ self.kv_cache_dtype.startswith("nvfp4")
508
+ )
509
510
# check if query quantization is supported
511
if self.impl.supports_quant_query_input:
vllm/utils/torch_utils.py
+3
-2
@@ -50,6 +50,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {
50
"turboquant_k3v4_nc": torch.uint8,
51
"turboquant_3bit_nc": torch.uint8,
52
"nvfp4": torch.uint8,
53
+ "nvfp4_4over6": torch.uint8,
54
}
55
56
TORCH_DTYPE_TO_NUMPY_DTYPE = {
@@ -77,7 +78,7 @@ def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
78
return (
79
kv_cache_dtype.startswith("fp8")
80
or kv_cache_dtype.endswith("per_token_head")
80
- or kv_cache_dtype == "nvfp4"
81
+ or kv_cache_dtype.startswith("nvfp4")
82
)
83
84
@@ -608,7 +609,7 @@ def create_kv_caches_with_random_flash(
609
value_caches: list[torch.Tensor] = []
610
611
for _ in range(num_layers):
611
- if cache_dtype == "nvfp4":
612
+ if isinstance(cache_dtype, str) and cache_dtype.startswith("nvfp4"):
613
# Full page dim: fp4 data + fp8 block scales per head.
614
# Per page layout: [K_data | K_scale | V_data | V_scale]
615
# Returns [:, 0] and [:, 1] like all other dtypes.
vllm/v1/attention/backends/flashinfer.py
+15
-12
@@ -350,6 +350,7 @@ class FlashInferBackend(AttentionBackend):
350
"fp8_e4m3",
351
"fp8_e5m2",
352
"nvfp4",
353
+ "nvfp4_4over6",
354
]
355
356
@staticmethod
@@ -402,7 +403,7 @@ class FlashInferBackend(AttentionBackend):
403
head_size: int,
404
cache_dtype_str: str = "auto",
405
) -> tuple[int, ...]:
405
- if cache_dtype_str == "nvfp4":
406
+ if cache_dtype_str.startswith("nvfp4"):
407
full_dim = nvfp4_kv_cache_full_dim(head_size)
408
return (num_blocks, 2 * num_kv_heads, block_size, full_dim)
409
# Pack K and V in the content dim (B, H, N, 2*hs).
@@ -438,14 +439,14 @@ class FlashInferBackend(AttentionBackend):
439
return torch.float8_e4m3fn
440
elif kv_cache_dtype == "fp8_e5m2":
441
return torch.float8_e5m2
441
- elif kv_cache_dtype == "nvfp4":
442
+ elif kv_cache_dtype.startswith("nvfp4"):
443
return torch.uint8
444
else:
445
raise ValueError(f"Unrecognized dtype: {kv_cache_dtype}")
446
447
@classmethod
448
def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
448
- if kv_cache_dtype == "nvfp4":
449
+ if kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4"):
450
return (
451
current_platform.is_device_capability_family(100)
452
and supports_trtllm_attention(is_prefill=True)
@@ -709,7 +710,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
710
self.cache_dtype = self.cache_config.cache_dtype
711
# Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
712
# storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
712
- self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4"
713
+ self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4")
714
if self.is_kvcache_nvfp4:
715
if (
716
force_use_trtllm_attention() is False
@@ -717,12 +718,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
718
or not supports_trtllm_attention(is_prefill=False)
719
):
720
raise ValueError(
720
- "--kv-cache-dtype nvfp4 requires the SM100 trtllm-gen "
721
+ f"--kv-cache-dtype {self.cache_dtype} requires the "
722
+ "SM100 trtllm-gen "
723
"FlashInfer path."
724
)
723
- # For NVFP4, kv_cache_dtype stays as the string "nvfp4"
724
- # which is passed to FlashInferImpl
725
- self.kv_cache_dtype = self.cache_dtype
725
+ # The scale search only affects the store kernel. FlashInfer
726
+ # reads both variants using the same NVFP4 layout.
727
+ self.kv_cache_dtype = "nvfp4"
728
else:
729
self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer(
730
self.cache_dtype
@@ -875,7 +877,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
877
) or current_platform.is_device_capability_family(100):
878
return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype)
879
return self.model_config.dtype
878
- if cache_dtype == "nvfp4":
880
+ if cache_dtype.startswith("nvfp4"):
881
return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3")
882
return self.kv_cache_spec.dtype
883
@@ -1558,8 +1560,9 @@ class FlashInferImpl(AttentionImpl):
1560
self.window_left = (
1561
self.sliding_window[0] if self.sliding_window is not None else -1
1562
)
1561
- self.kv_cache_dtype = kv_cache_dtype
1562
- self.is_kvcache_nvfp4 = kv_cache_dtype == "nvfp4"
1563
+ self.cache_dtype = kv_cache_dtype
1564
+ self.is_kvcache_nvfp4 = kv_cache_dtype.startswith("nvfp4")
1565
+ self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype
1566
self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0
1567
self.logits_soft_cap = logits_soft_cap
1568
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
@@ -2284,7 +2287,7 @@ class FlashInferImpl(AttentionImpl):
2287
k_cache,
2288
v_cache,
2289
slot_mapping,
2287
- self.kv_cache_dtype,
2290
+ self.cache_dtype,
2291
layer._k_scale,
2292
layer._v_scale,
2293
)
vllm/v1/kv_cache_interface.py
+1
-1
@@ -73,7 +73,7 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
73
return KVQuantMode.INT8_PER_TOKEN_HEAD
74
if kv_cache_dtype == "fp8_per_token_head":
75
return KVQuantMode.FP8_PER_TOKEN_HEAD
76
- if kv_cache_dtype == "nvfp4":
76
+ if kv_cache_dtype.startswith("nvfp4"):
77
return KVQuantMode.NVFP4
78
if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"):
79
return KVQuantMode.TURBOQUANT