iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add NVFP4 KV 4-over-6 scale search (#45187)

Browse Source 
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
This commit is contained in:
Wei-Ming Chen
2026-08-07 13:22:27 -07:00
committed by GitHub
parent 70456e5e6f
commit 1c94e8dc7d
11 changed files with 284 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/cache_kernels.cu
+4 -3
View File
@@ -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(
csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu
+167 -21
View File
@@ -20,8 +20,17 @@
#include "libtorch_stable/dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
#include <cmath>
#include <string>
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 <class Type, int CVT_FP4_NUM_THREADS_PER_SF>
__device__ __forceinline__ float nvfp4_reconstruction_error(
PackedVec<Type, CVT_FP4_PACK16>& 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 <class Type, int CVT_FP4_NUM_THREADS_PER_SF>
__device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4_4over6(
PackedVec<Type, CVT_FP4_PACK16>& 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<Type, CVT_FP4_NUM_THREADS_PER_SF>(
vec, outputScale6);
const float err4 =
nvfp4_reconstruction_error<Type, CVT_FP4_NUM_THREADS_PER_SF>(
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 <typename scalar_t>
template <typename scalar_t, NVFP4KVScaleSearch SCALE_SEARCH>
__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<CudaType, THREADS_PER_SF>(
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<CudaType, THREADS_PER_SF>(
in_vec, global_scale, sf_out_ptr);
} else {
packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
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<scalar_t>
<<<grid, block, 0, stream>>>(
key.const_data_ptr<scalar_t>(),
value.const_data_ptr<scalar_t>(),
key_cache.mutable_data_ptr<uint8_t>(),
value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
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>
<<<grid, block, 0, stream>>>(
key.const_data_ptr<scalar_t>(),
value.const_data_ptr<scalar_t>(),
key_cache.mutable_data_ptr<uint8_t>(),
value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
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>
<<<grid, block, 0, stream>>>(
key.const_data_ptr<scalar_t>(),
value.const_data_ptr<scalar_t>(),
key_cache.mutable_data_ptr<uint8_t>(),
value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
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);
}
});
}
tests/kernels/attention/test_attention_selector.py
+1
View File
@@ -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",
],
tests/kernels/attention/test_cache.py
+67
View File
@@ -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)
tests/utils_/test_torch_utils.py
+14
View File
@@ -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"),
[
vllm/config/cache.py
+3
View File
@@ -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
vllm/config/vllm.py
+4 -1
View File
@@ -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 "
vllm/model_executor/layers/attention/attention.py
+5 -2
View File
@@ -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:
vllm/utils/torch_utils.py
+3 -2
View File
@@ -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.
vllm/v1/attention/backends/flashinfer.py
+15 -12
View File
@@ -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,
)
vllm/v1/kv_cache_interface.py
+1 -1
View File
@@ -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