From 1c94e8dc7da43d12c8686f05cc5d2eab3fd9528f Mon Sep 17 00:00:00 2001 From: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:22:27 -0700 Subject: [PATCH] Add NVFP4 KV 4-over-6 scale search (#45187) Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> --- csrc/libtorch_stable/cache_kernels.cu | 7 +- .../libtorch_stable/nvfp4_kv_cache_kernels.cu | 188 ++++++++++++++++-- .../attention/test_attention_selector.py | 1 + tests/kernels/attention/test_cache.py | 67 +++++++ tests/utils_/test_torch_utils.py | 14 ++ vllm/config/cache.py | 3 + vllm/config/vllm.py | 5 +- .../layers/attention/attention.py | 7 +- vllm/utils/torch_utils.py | 5 +- vllm/v1/attention/backends/flashinfer.py | 27 +-- vllm/v1/kv_cache_interface.py | 2 +- 11 files changed, 284 insertions(+), 42 deletions(-) diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index 83716d6fc7..b16f405a36 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -770,16 +770,17 @@ void reshape_and_cache_flash( key.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); - if (kv_cache_dtype == "nvfp4") { + if (kv_cache_dtype == "nvfp4" || kv_cache_dtype == "nvfp4_4over6") { #if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120) // NVFP4 dispatch is compiled separately for SM100+. extern void reshape_and_cache_nvfp4_dispatch( torch::stable::Tensor & key, torch::stable::Tensor & value, torch::stable::Tensor & key_cache, torch::stable::Tensor & value_cache, torch::stable::Tensor & slot_mapping, torch::stable::Tensor & k_scale, - torch::stable::Tensor & v_scale); + torch::stable::Tensor & v_scale, const std::string& kv_cache_dtype); reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache, - slot_mapping, k_scale, v_scale); + slot_mapping, k_scale, v_scale, + kv_cache_dtype); return; #else STD_TORCH_CHECK( diff --git a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu b/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu index 693d55f231..bb11e474a6 100644 --- a/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu +++ b/csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu @@ -20,8 +20,17 @@ #include "libtorch_stable/dispatch_utils.h" #include "libtorch_stable/torch_utils.h" +#include +#include + namespace vllm { +// The cache layout is identical; this only selects the store-time scale search. +enum class NVFP4KVScaleSearch { + DEFAULT, + FOUR_OVER_SIX, +}; + // Compute swizzled scale offset for SM100 trtllm-gen MHA kernel. // The swizzle pattern for HND layout is: // [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, return swizzled_t * scale_dim + swizzled_s; } +__device__ __forceinline__ float round_to_nearest_e2m1(float x) { + const float ax = fabsf(x); + float q; + // Match cvt.rn.satfinite.e2m1x2.f32, including ties-to-even boundaries. + if (ax <= 0.25f) { + q = 0.0f; + } else if (ax < 0.75f) { + q = 0.5f; + } else if (ax <= 1.25f) { + q = 1.0f; + } else if (ax < 1.75f) { + q = 1.5f; + } else if (ax <= 2.5f) { + q = 2.0f; + } else if (ax < 3.5f) { + q = 3.0f; + } else if (ax <= 5.0f) { + q = 4.0f; + } else { + q = 6.0f; + } + return copysignf(q, x); +} + +__device__ __forceinline__ void nvfp4_candidate_scale(float SFScaleVal, + float vecMax, + float denominator, + uint8_t* fp8_sf, + float* outputScale) { + float SFValue = + SFScaleVal * (vecMax * reciprocal_approximate_ftz(denominator)); + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); + reinterpret_cast<__nv_fp8_e4m3&>(*fp8_sf) = tmp; + SFValue = float(tmp); + *outputScale = SFValue != 0.0f + ? reciprocal_approximate_ftz( + SFValue * reciprocal_approximate_ftz(SFScaleVal)) + : 0.0f; +} + +template +__device__ __forceinline__ float nvfp4_reconstruction_error( + PackedVec& vec, float outputScale) { + float error = 0.0f; + const float dequantScale = + outputScale != 0.0f ? reciprocal_approximate_ftz(outputScale) : 0.0f; + +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + float2 vals = cast_to_float2(vec.elts[i]); + if (outputScale == 0.0f) { + error += vals.x * vals.x + vals.y * vals.y; + } else { + float qx = round_to_nearest_e2m1(vals.x * outputScale); + float qy = round_to_nearest_e2m1(vals.y * outputScale); + float dx = qx * dequantScale - vals.x; + float dy = qy * dequantScale - vals.y; + error += dx * dx + dy * dy; + } + } + + if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) { + error += __shfl_xor_sync(0xffffffffu, error, 1); + } + return error; +} + +template +__device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4_4over6( + PackedVec& vec, float SFScaleVal, uint8_t* SFout) { + auto localMax = __habs2(vec.elts[0]); + +#pragma unroll + for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + localMax = __hmax2(localMax, __habs2(vec.elts[i])); + } + + if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) { + localMax = __hmax2(__shfl_xor_sync(0xffffffffu, localMax, 1), localMax); + } + const float vecMax = float(__hmax(localMax.x, localMax.y)); + + uint8_t sf6; + uint8_t sf4; + float outputScale6; + float outputScale4; + nvfp4_candidate_scale(SFScaleVal, vecMax, 6.0f, &sf6, &outputScale6); + nvfp4_candidate_scale(SFScaleVal, vecMax, 4.0f, &sf4, &outputScale4); + + const float err6 = + nvfp4_reconstruction_error( + vec, outputScale6); + const float err4 = + nvfp4_reconstruction_error( + vec, outputScale4); + const bool use4 = err4 < err6; + const float outputScale = use4 ? outputScale4 : outputScale6; + + if (SFout) *SFout = use4 ? sf4 : sf6; + + float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; + +#pragma unroll + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { + fp2Vals[i] = cast_to_float2(vec.elts[i]); + fp2Vals[i].x *= outputScale; + fp2Vals[i].y *= outputScale; + } + + return pack_fp4(fp2Vals); +} + // Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache. // // 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, // // Threading: one CUDA block per token, threads process heads and // groups of 16 elements within each head. -template +template __global__ void reshape_and_cache_nvfp4_kernel( const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size] @@ -129,8 +250,14 @@ __global__ void reshape_and_cache_nvfp4_kernel( uint8_t sf_val; uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr; - fp4_packed_t packed = cvt_warp_fp16_to_fp4( - in_vec, global_scale, sf_out_ptr); + fp4_packed_t packed; + if constexpr (SCALE_SEARCH == NVFP4KVScaleSearch::FOUR_OVER_SIX) { + packed = cvt_warp_fp16_to_fp4_4over6( + in_vec, global_scale, sf_out_ptr); + } else { + packed = cvt_warp_fp16_to_fp4( + in_vec, global_scale, sf_out_ptr); + } // Write packed FP4 data to data cache uint8_t* __restrict__ data_dst = data_block + head * data_head_stride + @@ -181,13 +308,11 @@ __global__ void reshape_and_cache_nvfp4_kernel( // Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1]. // Each KV side contains both data and scale: // page = [K_data | K_scale | V_data | V_scale] -void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key, - torch::stable::Tensor& value, - torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, - torch::stable::Tensor& slot_mapping, - torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale) { +void reshape_and_cache_nvfp4_dispatch( + torch::stable::Tensor& key, torch::stable::Tensor& value, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + torch::stable::Tensor& slot_mapping, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const std::string& kv_cache_dtype) { int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); @@ -262,16 +387,37 @@ void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key, VLLM_STABLE_DISPATCH_HALF_TYPES( key.scalar_type(), "reshape_and_cache_nvfp4", [&] { - vllm::reshape_and_cache_nvfp4_kernel - <<>>( - key.const_data_ptr(), - value.const_data_ptr(), - key_cache.mutable_data_ptr(), - value_cache.mutable_data_ptr(), key_scale_ptr, - value_scale_ptr, slot_mapping.const_data_ptr(), - k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0), - num_heads, head_size, block_size, data_block_stride, - data_head_stride, data_block_offset_stride, scale_block_stride, - scale_head_stride, scale_block_offset_stride); + if (kv_cache_dtype == "nvfp4") { + vllm::reshape_and_cache_nvfp4_kernel< + scalar_t, vllm::NVFP4KVScaleSearch::DEFAULT> + <<>>( + key.const_data_ptr(), + value.const_data_ptr(), + key_cache.mutable_data_ptr(), + value_cache.mutable_data_ptr(), key_scale_ptr, + value_scale_ptr, slot_mapping.const_data_ptr(), + k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0), + num_heads, head_size, block_size, data_block_stride, + data_head_stride, data_block_offset_stride, + scale_block_stride, scale_head_stride, + scale_block_offset_stride); + } else if (kv_cache_dtype == "nvfp4_4over6") { + vllm::reshape_and_cache_nvfp4_kernel< + scalar_t, vllm::NVFP4KVScaleSearch::FOUR_OVER_SIX> + <<>>( + key.const_data_ptr(), + value.const_data_ptr(), + key_cache.mutable_data_ptr(), + value_cache.mutable_data_ptr(), key_scale_ptr, + value_scale_ptr, slot_mapping.const_data_ptr(), + k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0), + num_heads, head_size, block_size, data_block_stride, + data_head_stride, data_block_offset_stride, + scale_block_stride, scale_head_stride, + scale_block_offset_stride); + } else { + STD_TORCH_CHECK(false, + "Unsupported NVFP4 KV cache dtype: ", kv_cache_dtype); + } }); } diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 587f75a38d..e2f5c7f5d1 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/tests/kernels/attention/test_attention_selector.py @@ -528,6 +528,7 @@ def test_non_causal_autoselect_backend(): "fp8_ds_mla", "fp8_inc", "nvfp4", + "nvfp4_4over6", "fp8_per_token_head", "int8_per_token_head", ], diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 6fe355f945..c48a1d2338 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -426,6 +426,73 @@ def test_reshape_and_cache_flash( torch.testing.assert_close(value_cache_compact, cloned_value_cache) +@torch.inference_mode() +def test_nvfp4_4over6_selects_lower_error_scale( + kv_cache_factory_flashinfer, +) -> None: + if not current_platform.has_device_capability(100): + pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).") + + device = CUDA_DEVICES[0] + block_size = 16 + head_size = 64 + source = torch.tensor( + [6.0] + [4.5] * 15, + dtype=torch.bfloat16, + device=device, + ).repeat(1, 1, head_size // 16) + slot_mapping = torch.zeros(1, dtype=torch.long, device=device) + scale = torch.ones(1, dtype=torch.float32, device=device) + + def quantize_and_dequantize(cache_dtype: str) -> torch.Tensor: + key_caches, value_caches = kv_cache_factory_flashinfer( + 1, + block_size, + 1, + 1, + head_size, + cache_dtype, + source.dtype, + device=device, + cache_layout="NHD", + ) + key_cache, value_cache = key_caches[0], value_caches[0] + key_cache.zero_() + value_cache.zero_() + + ops.reshape_and_cache_flash( + source, + source, + key_cache, + value_cache, + slot_mapping, + cache_dtype, + scale, + scale, + ) + + from tests.kernels.quantization.nvfp4_utils import ( + dequant_nvfp4_kv_cache, + ) + + data, block_scales = nvfp4_split_data_scale(key_cache) + return dequant_nvfp4_kv_cache( + data.permute(0, 2, 1, 3), + block_scales.permute(0, 2, 1, 3), + scale.item(), + head_size, + block_size, + )[0, 0, 0] + + default = quantize_and_dequantize("nvfp4") + four_over_six = quantize_and_dequantize("nvfp4_4over6") + source_f32 = source[0, 0].float() + + default_mse = (default - source_f32).square().mean() + four_over_six_mse = (four_over_six - source_f32).square().mean() + assert four_over_six_mse < default_mse + + @pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE) @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS) diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index ddaadfd123..6d63fc5db4 100644 --- a/tests/utils_/test_torch_utils.py +++ b/tests/utils_/test_torch_utils.py @@ -8,12 +8,26 @@ from vllm.utils.torch_utils import ( available_cpu_count, common_broadcastable_dtype, current_stream, + get_kv_cache_torch_dtype, is_lossless_cast, + is_quantized_kv_cache, set_torch_threads_for_runtime, startup_omp_num_threads, ) +def test_nvfp4_4over6_cache_dtype() -> None: + from vllm.config.cache import CacheConfig + from vllm.v1.kv_cache_interface import KVQuantMode, get_kv_quant_mode + + cache_config = CacheConfig(cache_dtype="nvfp4_4over6") + + assert cache_config.cache_dtype == "nvfp4_4over6" + assert get_kv_cache_torch_dtype(cache_config.cache_dtype) == torch.uint8 + assert is_quantized_kv_cache(cache_config.cache_dtype) + assert get_kv_quant_mode(cache_config.cache_dtype) == KVQuantMode.NVFP4 + + @pytest.mark.parametrize( ("src_dtype", "tgt_dtype", "expected_result"), [ diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 5b3821f8e5..654e9a5358 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -33,6 +33,7 @@ CacheDType = Literal[ "int8_per_token_head", "fp8_per_token_head", "nvfp4", + "nvfp4_4over6", ] MambaDType = Literal["auto", "float32", "float16", "bfloat16"] MambaCacheMode = Literal["all", "align", "none"] @@ -80,6 +81,8 @@ class CacheConfig: Some models (namely DeepSeekV3.2) default to fp8, set to bfloat16 to use bfloat16 instead, this is an invalid option for models that do not default to fp8. + "nvfp4_4over6" uses the NVFP4 layout and selects between max/6 and max/4 + scales per 16 values by minimizing squared reconstruction error. """ is_attention_free: bool = False """Whether the model is attention-free. This is primarily set in diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index e4c8ea5289..986a3f7af3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2428,7 +2428,10 @@ class VllmConfig: def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig": if self.model_config is None: return self - if self.cache_config.cache_dtype == "nvfp4" and self.model_config.use_mla: + if ( + self.cache_config.cache_dtype.startswith("nvfp4") + and self.model_config.use_mla + ): raise ValueError( "nvfp4 KV cache is not supported with MLA (Multi-head Latent " "Attention) backends. Please use a different --kv-cache-dtype " diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 4db42e0ecf..8eb8aaa866 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -459,7 +459,8 @@ class Attention(nn.Module, AttentionLayerBase): if ( self.impl.supports_quant_query_input and ( - self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4" + self.kv_cache_dtype.startswith("fp8") + or self.kv_cache_dtype.startswith("nvfp4") ) and not self.kv_cache_dtype.endswith("per_token_head") ): @@ -502,7 +503,9 @@ class Attention(nn.Module, AttentionLayerBase): # which reduces overheads during decoding. # Otherwise queries are quantized using custom ops # which causes decoding overheads - assert self.kv_cache_dtype in {"fp8", "fp8_e4m3", "nvfp4"} + assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"} or ( + self.kv_cache_dtype.startswith("nvfp4") + ) # check if query quantization is supported if self.impl.supports_quant_query_input: diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 16c092c949..0a188875b6 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -50,6 +50,7 @@ STR_DTYPE_TO_TORCH_DTYPE = { "turboquant_k3v4_nc": torch.uint8, "turboquant_3bit_nc": torch.uint8, "nvfp4": torch.uint8, + "nvfp4_4over6": torch.uint8, } TORCH_DTYPE_TO_NUMPY_DTYPE = { @@ -77,7 +78,7 @@ def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: return ( kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head") - or kv_cache_dtype == "nvfp4" + or kv_cache_dtype.startswith("nvfp4") ) @@ -608,7 +609,7 @@ def create_kv_caches_with_random_flash( value_caches: list[torch.Tensor] = [] for _ in range(num_layers): - if cache_dtype == "nvfp4": + if isinstance(cache_dtype, str) and cache_dtype.startswith("nvfp4"): # Full page dim: fp4 data + fp8 block scales per head. # Per page layout: [K_data | K_scale | V_data | V_scale] # Returns [:, 0] and [:, 1] like all other dtypes. diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 5251e25e5e..8e70440e77 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -350,6 +350,7 @@ class FlashInferBackend(AttentionBackend): "fp8_e4m3", "fp8_e5m2", "nvfp4", + "nvfp4_4over6", ] @staticmethod @@ -402,7 +403,7 @@ class FlashInferBackend(AttentionBackend): head_size: int, cache_dtype_str: str = "auto", ) -> tuple[int, ...]: - if cache_dtype_str == "nvfp4": + if cache_dtype_str.startswith("nvfp4"): full_dim = nvfp4_kv_cache_full_dim(head_size) return (num_blocks, 2 * num_kv_heads, block_size, full_dim) # Pack K and V in the content dim (B, H, N, 2*hs). @@ -438,14 +439,14 @@ class FlashInferBackend(AttentionBackend): return torch.float8_e4m3fn elif kv_cache_dtype == "fp8_e5m2": return torch.float8_e5m2 - elif kv_cache_dtype == "nvfp4": + elif kv_cache_dtype.startswith("nvfp4"): return torch.uint8 else: raise ValueError(f"Unrecognized dtype: {kv_cache_dtype}") @classmethod def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: - if kv_cache_dtype == "nvfp4": + if kv_cache_dtype is not None and kv_cache_dtype.startswith("nvfp4"): return ( current_platform.is_device_capability_family(100) and supports_trtllm_attention(is_prefill=True) @@ -709,7 +710,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): self.cache_dtype = self.cache_config.cache_dtype # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3) - self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4" + self.is_kvcache_nvfp4 = self.cache_dtype.startswith("nvfp4") if self.is_kvcache_nvfp4: if ( force_use_trtllm_attention() is False @@ -717,12 +718,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): or not supports_trtllm_attention(is_prefill=False) ): raise ValueError( - "--kv-cache-dtype nvfp4 requires the SM100 trtllm-gen " + f"--kv-cache-dtype {self.cache_dtype} requires the " + "SM100 trtllm-gen " "FlashInfer path." ) - # For NVFP4, kv_cache_dtype stays as the string "nvfp4" - # which is passed to FlashInferImpl - self.kv_cache_dtype = self.cache_dtype + # The scale search only affects the store kernel. FlashInfer + # reads both variants using the same NVFP4 layout. + self.kv_cache_dtype = "nvfp4" else: self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer( self.cache_dtype @@ -875,7 +877,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): ) or current_platform.is_device_capability_family(100): return FlashInferBackend.get_dtype_for_flashinfer(cache_dtype) return self.model_config.dtype - if cache_dtype == "nvfp4": + if cache_dtype.startswith("nvfp4"): return FlashInferBackend.get_dtype_for_flashinfer("fp8_e4m3") return self.kv_cache_spec.dtype @@ -1558,8 +1560,9 @@ class FlashInferImpl(AttentionImpl): self.window_left = ( self.sliding_window[0] if self.sliding_window is not None else -1 ) - self.kv_cache_dtype = kv_cache_dtype - self.is_kvcache_nvfp4 = kv_cache_dtype == "nvfp4" + self.cache_dtype = kv_cache_dtype + self.is_kvcache_nvfp4 = kv_cache_dtype.startswith("nvfp4") + self.kv_cache_dtype = "nvfp4" if self.is_kvcache_nvfp4 else kv_cache_dtype self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0 self.logits_soft_cap = logits_soft_cap self.kv_sharing_target_layer_name = kv_sharing_target_layer_name @@ -2284,7 +2287,7 @@ class FlashInferImpl(AttentionImpl): k_cache, v_cache, slot_mapping, - self.kv_cache_dtype, + self.cache_dtype, layer._k_scale, layer._v_scale, ) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index beb06b08f0..9021006572 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -73,7 +73,7 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: return KVQuantMode.INT8_PER_TOKEN_HEAD if kv_cache_dtype == "fp8_per_token_head": return KVQuantMode.FP8_PER_TOKEN_HEAD - if kv_cache_dtype == "nvfp4": + if kv_cache_dtype.startswith("nvfp4"): return KVQuantMode.NVFP4 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"): return KVQuantMode.TURBOQUANT