A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 01:00 UTC
Add nvfp4 support to reshape_and_cache_flash (#37332)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
sychen52 committed
Apr 17, 2026 at 07:28 UTC
6b2b7bd0ebd43ef756632d2142ce974929f05d8f
10 files changed
+679
-51
CMakeLists.txt
+14
@@ -923,6 +923,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
923
SRCS "${SRCS}"
924
CUDA_ARCHS "${FP4_ARCHS}")
925
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
926
+ # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
927
+ # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
928
+ set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
929
+ set_gencode_flags_for_srcs(
930
+ SRCS "${NVFP4_KV_SRC}"
931
+ CUDA_ARCHS "${FP4_ARCHS}")
932
+ target_sources(_C PRIVATE ${NVFP4_KV_SRC})
933
+ target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
934
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
935
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
936
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -949,6 +957,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
957
SRCS "${SRCS}"
958
CUDA_ARCHS "${FP4_ARCHS}")
959
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
960
+ set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
961
+ set_gencode_flags_for_srcs(
962
+ SRCS "${NVFP4_KV_SRC}"
963
+ CUDA_ARCHS "${FP4_ARCHS}")
964
+ target_sources(_C PRIVATE ${NVFP4_KV_SRC})
965
+ target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
966
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
967
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
968
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
csrc/cache_kernels.cu
+22
-2
@@ -724,6 +724,28 @@ void reshape_and_cache_flash(
724
int num_tokens = slot_mapping.size(0);
725
int num_heads = key.size(1);
726
int head_size = key.size(2);
727
+
728
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
729
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
730
+
731
+ if (kv_cache_dtype == "nvfp4") {
732
+#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
733
+ // NVFP4 dispatch is compiled separately for SM100+.
734
+ extern void reshape_and_cache_nvfp4_dispatch(
735
+ torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
736
+ torch::Tensor & value_cache, torch::Tensor & slot_mapping,
737
+ torch::Tensor & k_scale, torch::Tensor & v_scale);
738
+ reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
739
+ slot_mapping, k_scale, v_scale);
740
+ return;
741
+#else
742
+ TORCH_CHECK(false,
743
+ "NVFP4 KV cache requires SM100+ (Blackwell). "
744
+ "Please rebuild vllm with a Blackwell-compatible CUDA target.");
745
+#endif
746
+ }
747
+
748
+ // Original FP8/auto path.
749
int block_size = key_cache.size(1);
750
751
int64_t key_stride = key.stride(0);
@@ -741,8 +763,6 @@ void reshape_and_cache_flash(
763
764
dim3 grid(num_tokens);
765
dim3 block(std::min(num_heads * head_size, 512));
744
- const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
745
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
766
767
DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
768
CALL_RESHAPE_AND_CACHE_FLASH);
csrc/nvfp4_kv_cache_kernels.cu
new
+275
@@ -0,0 +1,275 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+// NVFP4 KV cache store kernel.
5
+// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
6
+// into the paged KV cache.
7
+//
8
+// Per page layout: [K_data | K_scale | V_data | V_scale]
9
+// Both data and scale regions are contiguous per head, enabling direct
10
+// TMA descriptor use.
11
+//
12
+// Reuses device functions from nvfp4_utils.cuh:
13
+// - cvt_warp_fp16_to_fp4() for bf16 → fp4 quantization + block scale
14
+// - pack_fp4() for packing float pairs to fp4
15
+// - reciprocal_approximate_ftz() for fast reciprocal
16
+
17
+#define NVFP4_ENABLE_ELTS16 1
18
+#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
19
+
20
+#include <ATen/cuda/CUDAContext.h>
21
+#include <c10/cuda/CUDAGuard.h>
22
+#include <torch/all.h>
23
+
24
+#include "dispatch_utils.h"
25
+
26
+namespace vllm {
27
+
28
+// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
29
+// The swizzle pattern for HND layout is:
30
+// [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
31
+// where T = block_size (page_size), S = scale_dim = head_size // 16.
32
+//
33
+// For a linear (t, s) position, the swizzled position is:
34
+// swizzled_t = (t / 4) * 4 + (s / (S / 4))
35
+// swizzled_s = (s % (S / 4)) * 4 + (t % 4)
36
+__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
37
+ int scale_dim) {
38
+ int s_group = scale_dim / 4;
39
+ int swizzled_t = (t / 4) * 4 + (s / s_group);
40
+ int swizzled_s = (s % s_group) * 4 + (t % 4);
41
+ return swizzled_t * scale_dim + swizzled_s;
42
+}
43
+
44
+// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
45
+//
46
+// Takes separate data and scale cache pointers for K and V.
47
+// Within each KV side, data and scale are separate contiguous regions.
48
+//
49
+// Threading: one CUDA block per token, threads process heads and
50
+// groups of 16 elements within each head.
51
+template <typename scalar_t>
52
+__global__ void reshape_and_cache_nvfp4_kernel(
53
+ const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size]
54
+ const scalar_t* __restrict__ value, // [num_tokens, num_heads, head_size]
55
+ uint8_t* __restrict__ key_data_cache, // data region for K
56
+ uint8_t* __restrict__ value_data_cache, // data region for V
57
+ uint8_t* __restrict__ key_scale_cache, // scale region for K
58
+ uint8_t* __restrict__ value_scale_cache, // scale region for V
59
+ const int64_t* __restrict__ slot_mapping, // [num_actual_tokens]
60
+ const float* __restrict__ k_scale_ptr, // pointer to checkpoint k_scale
61
+ const float* __restrict__ v_scale_ptr, // pointer to checkpoint v_scale
62
+ const int64_t key_stride, // key.stride(0) in elements
63
+ const int64_t value_stride, // value.stride(0) in elements
64
+ const int num_heads, const int head_size, const int block_size,
65
+ const int64_t data_block_stride, // data cache stride for dim 0
66
+ const int64_t data_head_stride, // data cache stride for heads
67
+ const int64_t data_block_offset_stride, // data cache stride for tokens
68
+ const int64_t scale_block_stride, // scale cache stride for dim 0
69
+ const int64_t scale_head_stride, // scale cache stride for heads
70
+ const int64_t scale_block_offset_stride // scale cache stride for tokens
71
+) {
72
+ using CudaType = typename CUDATypeConverter<scalar_t>::Type;
73
+ using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
74
+
75
+ static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD; // 16 or 8
76
+ static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
77
+
78
+ const int64_t token_idx = blockIdx.x;
79
+ const int64_t slot_idx = slot_mapping[token_idx];
80
+ if (slot_idx < 0) return;
81
+
82
+ const int64_t block_idx = slot_idx / block_size;
83
+ const int block_offset = static_cast<int>(slot_idx % block_size);
84
+
85
+ const int scale_dim = head_size / 16;
86
+ const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
87
+
88
+ const int total_groups = num_heads * groups_per_head;
89
+ const int tid = threadIdx.x;
90
+ const int num_thread_groups = blockDim.x / THREADS_PER_SF;
91
+ const int tg_id = tid / THREADS_PER_SF;
92
+ const int tg_lane = tid % THREADS_PER_SF;
93
+
94
+ // Process both K (kv=0) and V (kv=1)
95
+#pragma unroll
96
+ for (int kv = 0; kv < 2; kv++) {
97
+ const scalar_t* __restrict__ src = (kv == 0) ? key : value;
98
+ const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
99
+ const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
100
+ uint8_t* __restrict__ data_cache =
101
+ (kv == 0) ? key_data_cache : value_data_cache;
102
+ uint8_t* __restrict__ sc_cache =
103
+ (kv == 0) ? key_scale_cache : value_scale_cache;
104
+
105
+ // Source pointer for this token (use actual stride, not assumed contiguous)
106
+ const CudaType* __restrict__ token_src =
107
+ reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
108
+
109
+ // Destination bases in data and scale caches for this token's block
110
+ uint8_t* __restrict__ data_block =
111
+ data_cache + block_idx * data_block_stride;
112
+ uint8_t* __restrict__ scale_block =
113
+ sc_cache + block_idx * scale_block_stride;
114
+
115
+ for (int g = tg_id; g < total_groups; g += num_thread_groups) {
116
+ const int head = g / groups_per_head;
117
+ const int group_in_head = g % groups_per_head;
118
+
119
+ // Load 16 (or 8) bf16 elements from source
120
+ PVec in_vec;
121
+ const CudaType* __restrict__ src_ptr =
122
+ token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
123
+ tg_lane * ELTS;
124
+
125
+#pragma unroll
126
+ for (int i = 0; i < ELTS / 2; i++) {
127
+ in_vec.elts[i] = reinterpret_cast<
128
+ const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
129
+ }
130
+
131
+ // Quantize: produces packed fp4 and writes scale factor.
132
+ uint8_t sf_val;
133
+ uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
134
+
135
+ fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
136
+ in_vec, global_scale, sf_out_ptr);
137
+
138
+ // Write packed FP4 data to data cache
139
+ uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
140
+ block_offset * data_block_offset_stride;
141
+
142
+#if CVT_FP4_PACK16
143
+ {
144
+ // 16 elements → 8 bytes (u32x2)
145
+ int data_byte_offset = group_in_head * 8;
146
+ reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
147
+ (uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
148
+ }
149
+#else
150
+ {
151
+ // 8 elements → 4 bytes (uint32_t)
152
+ int data_byte_offset =
153
+ group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
154
+ reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
155
+ }
156
+#endif
157
+
158
+ // Write block scale to scale cache.
159
+ // K (kv==0): linear layout (no swizzle).
160
+ // V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
161
+ if (sf_out_ptr != nullptr) {
162
+ int scale_idx = group_in_head;
163
+ uint8_t* __restrict__ scale_dst;
164
+ if (kv == 0) {
165
+ scale_dst = scale_block + head * scale_head_stride +
166
+ block_offset * scale_block_offset_stride + scale_idx;
167
+ } else {
168
+ int swizzled_offset =
169
+ swizzle_scale_offset(block_offset, scale_idx, scale_dim);
170
+ int swizzled_t = swizzled_offset / scale_dim;
171
+ int swizzled_s = swizzled_offset % scale_dim;
172
+ scale_dst = scale_block + head * scale_head_stride +
173
+ swizzled_t * scale_block_offset_stride + swizzled_s;
174
+ }
175
+ *scale_dst = sf_val;
176
+ }
177
+ }
178
+ }
179
+}
180
+
181
+} // namespace vllm
182
+
183
+// Non-template entry point callable from cache_kernels.cu.
184
+// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
185
+// Each KV side contains both data and scale:
186
+// page = [K_data | K_scale | V_data | V_scale]
187
+void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
188
+ torch::Tensor& key_cache,
189
+ torch::Tensor& value_cache,
190
+ torch::Tensor& slot_mapping,
191
+ torch::Tensor& k_scale,
192
+ torch::Tensor& v_scale) {
193
+ int num_tokens = slot_mapping.size(0);
194
+ int num_heads = key.size(1);
195
+ int head_size = key.size(2);
196
+ int data_dim = head_size / 2;
197
+ int scale_dim = head_size / 16;
198
+ int full_dim = data_dim + scale_dim;
199
+
200
+ // key_cache is kv_cache[:, 0] with shape
201
+ // [num_blocks, block_size, num_heads, full_dim] in logical order.
202
+ // Strides encode the physical layout (HND or NHD).
203
+ TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
204
+ TORCH_CHECK(key_cache.size(3) == full_dim,
205
+ "key_cache last dim must be data_dim + scale_dim, got ",
206
+ key_cache.size(3), " expected ", full_dim);
207
+
208
+ int block_size = key_cache.size(1);
209
+
210
+ TORCH_CHECK(head_size % 16 == 0,
211
+ "head_size must be divisible by 16 for NVFP4 KV cache");
212
+ TORCH_CHECK(block_size % 4 == 0,
213
+ "block_size must be divisible by 4 for NVFP4 KV cache swizzle");
214
+
215
+ // Detect physical layout from strides (based on full_dim).
216
+ // HND: head stride > block_offset stride.
217
+ bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
218
+
219
+ int64_t data_block_stride = key_cache.stride(0); // page_bytes
220
+ int64_t data_head_stride, data_block_offset_stride;
221
+ if (is_hnd) {
222
+ data_head_stride = (int64_t)block_size * data_dim;
223
+ data_block_offset_stride = data_dim;
224
+ } else {
225
+ data_head_stride = data_dim;
226
+ data_block_offset_stride = (int64_t)num_heads * data_dim;
227
+ }
228
+
229
+ // Page layout: [K_data | K_scale | V_data | V_scale]
230
+ // Scale follows data within each KV side.
231
+ int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
232
+
233
+ uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
234
+ uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
235
+
236
+ // Scale strides: same page stride, inner strides from layout.
237
+ int64_t scale_block_stride = data_block_stride;
238
+ int64_t scale_head_stride, scale_block_offset_stride;
239
+ if (is_hnd) {
240
+ scale_head_stride = (int64_t)block_size * scale_dim;
241
+ scale_block_offset_stride = scale_dim;
242
+ } else {
243
+ scale_head_stride = scale_dim;
244
+ scale_block_offset_stride = (int64_t)num_heads * scale_dim;
245
+ }
246
+
247
+ const float* k_scale_ptr = k_scale.data_ptr<float>();
248
+ const float* v_scale_ptr = v_scale.data_ptr<float>();
249
+
250
+ int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
251
+ int total_groups = num_heads * groups_per_head;
252
+ constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
253
+ int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
254
+ num_threads = ((num_threads + 31) / 32) * 32;
255
+
256
+ dim3 grid(num_tokens);
257
+ dim3 block(num_threads);
258
+
259
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
260
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
261
+
262
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(
263
+ key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
264
+ vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
265
+ <<<grid, block, 0, stream>>>(
266
+ key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
267
+ key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
268
+ key_scale_ptr, value_scale_ptr,
269
+ slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
270
+ key.stride(0), value.stride(0), num_heads, head_size,
271
+ block_size, data_block_stride, data_head_stride,
272
+ data_block_offset_stride, scale_block_stride, scale_head_stride,
273
+ scale_block_offset_stride);
274
+ });
275
+}
tests/kernels/attention/test_cache.py
+99
-21
@@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
10
from vllm import _custom_ops as ops
11
from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
12
from vllm.platforms import current_platform
13
-from vllm.utils.torch_utils import set_random_seed
13
+from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed
14
15
COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
16
DTYPES = [torch.bfloat16, torch.float]
@@ -172,7 +172,7 @@ def test_reshape_and_cache(
172
@pytest.mark.parametrize("dtype", DTYPES)
173
@pytest.mark.parametrize("seed", SEEDS)
174
@pytest.mark.parametrize("device", CUDA_DEVICES)
175
-@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
175
+@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE + ["nvfp4"])
176
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
177
@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
178
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@@ -202,6 +202,25 @@ def test_reshape_and_cache_flash(
202
if kv_scale_type == "attn_head" and implementation != "cuda":
203
pytest.skip("Only CUDA implementation supports attn_head scaling.")
204
205
+ if kv_cache_dtype == "nvfp4":
206
+ if not current_platform.has_device_capability(100):
207
+ pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")
208
+ if implementation != "cuda":
209
+ pytest.skip("NVFP4 only supports CUDA implementation.")
210
+ if kv_scale_type != "tensor":
211
+ pytest.skip("NVFP4 only supports per-tensor scaling.")
212
+ if head_size % 16 != 0:
213
+ pytest.skip("NVFP4 requires head_size divisible by 16.")
214
+ if (head_size // 16) % 4 != 0:
215
+ pytest.skip(
216
+ "NVFP4 requires (head_size // 16) divisible by 4 "
217
+ "for 4x4 block scale swizzle."
218
+ )
219
+ if block_size % 4 != 0:
220
+ pytest.skip("NVFP4 requires block_size divisible by 4.")
221
+ if dtype not in (torch.float16, torch.bfloat16):
222
+ pytest.skip("NVFP4 quantization only supports fp16/bf16 input.")
223
+
224
# fp8 conversion requires continugous memory buffer. Reduce the number of
225
# blocks and tokens to consume less memory.
226
num_tokens = num_tokens // 2
@@ -229,7 +248,23 @@ def test_reshape_and_cache_flash(
248
del key_caches
249
del value_caches
250
232
- if kv_scale_type == "tensor":
251
+ # For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes.
252
+ # Split views are still needed for dequant verification.
253
+ key_scale_cache = None
254
+ value_scale_cache = None
255
+ nvfp4_key_data = None
256
+ nvfp4_value_data = None
257
+ if kv_cache_dtype == "nvfp4":
258
+ (nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache)
259
+ (nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views(
260
+ value_cache
261
+ )
262
+
263
+ if kv_cache_dtype == "nvfp4":
264
+ # Global scale = amax / 448 (per-tensor)
265
+ k_scale = (key.abs().amax() / 448.0).to(torch.float32)
266
+ v_scale = (value.abs().amax() / 448.0).to(torch.float32)
267
+ elif kv_scale_type == "tensor":
268
k_scale = (key.amax() / 64.0).to(torch.float32)
269
v_scale = (value.amax() / 64.0).to(torch.float32)
270
else: # "attn_head"
@@ -240,8 +275,9 @@ def test_reshape_and_cache_flash(
275
y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
276
return y.contiguous()
277
243
- key_cache_compact = permute_and_compact(key_cache)
244
- value_cache_compact = permute_and_compact(value_cache)
278
+ if kv_cache_dtype != "nvfp4":
279
+ key_cache_compact = permute_and_compact(key_cache)
280
+ value_cache_compact = permute_and_compact(value_cache)
281
282
def convert_fp8_local(output, input, scale, kv_dtype):
283
fp8_input = input.view(current_platform.fp8_dtype())
@@ -257,7 +293,7 @@ def test_reshape_and_cache_flash(
293
result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
294
output.copy_(result)
295
260
- # Clone the KV caches.
296
+ # Clone the KV caches (for non-nvfp4, used as reference baseline).
297
if kv_cache_dtype == "fp8":
298
cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
299
convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
@@ -265,25 +301,27 @@ def test_reshape_and_cache_flash(
301
convert_fp8_local(
302
cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
303
)
268
- else:
304
+ elif kv_cache_dtype != "nvfp4":
305
cloned_key_cache = key_cache_compact.clone()
306
cloned_value_cache = value_cache_compact.clone()
307
+
308
# Call the reshape_and_cache kernel.
309
if implementation == "cuda":
273
- opcheck(
274
- torch.ops._C_cache_ops.reshape_and_cache_flash,
275
- (
276
- key,
277
- value,
278
- key_cache,
279
- value_cache,
280
- slot_mapping,
281
- kv_cache_dtype,
282
- k_scale,
283
- v_scale,
284
- ),
285
- cond=(head_size == HEAD_SIZES[0]),
286
- )
310
+ if kv_cache_dtype != "nvfp4":
311
+ opcheck(
312
+ torch.ops._C_cache_ops.reshape_and_cache_flash,
313
+ (
314
+ key,
315
+ value,
316
+ key_cache,
317
+ value_cache,
318
+ slot_mapping,
319
+ kv_cache_dtype,
320
+ k_scale,
321
+ v_scale,
322
+ ),
323
+ cond=(head_size == HEAD_SIZES[0]),
324
+ )
325
ops.reshape_and_cache_flash(
326
key,
327
value,
@@ -309,6 +347,46 @@ def test_reshape_and_cache_flash(
347
k_scale,
348
v_scale,
349
)
350
+
351
+ if kv_cache_dtype == "nvfp4":
352
+ # Verify NVFP4 by dequantizing the entire cache and comparing
353
+ # the written positions against original bf16 values.
354
+ # Same pattern as FP8: dequant whole cache, then extract and compare.
355
+ from tests.kernels.quantization.nvfp4_utils import (
356
+ dequant_nvfp4_kv_cache,
357
+ )
358
+
359
+ def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
360
+ # data_cache: [N, T, H, data_dim] NHD (contiguous inner dims)
361
+ # scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
362
+ # Permute to HND layout for the dequant utility.
363
+ data_hnd = data_cache.permute(0, 2, 1, 3)
364
+ scale_hnd = scale_cache.permute(0, 2, 1, 3)
365
+ result_hnd = dequant_nvfp4_kv_cache(
366
+ data_hnd, scale_hnd, global_scale, head_size, block_size
367
+ )
368
+ return result_hnd.permute(0, 2, 1, 3) # back to [N, T, H, D]
369
+
370
+ result_key_cache = dequant_nvfp4_cache_nhd(
371
+ nvfp4_key_data, key_scale_cache, k_scale.item()
372
+ )
373
+ result_value_cache = dequant_nvfp4_cache_nhd(
374
+ nvfp4_value_data, value_scale_cache, v_scale.item()
375
+ )
376
+
377
+ # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
378
+ num_slots = num_blocks * block_size
379
+ result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
380
+ result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)
381
+
382
+ torch.testing.assert_close(
383
+ result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
384
+ )
385
+ torch.testing.assert_close(
386
+ result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5
387
+ )
388
+ return
389
+
390
key_cache_compact = permute_and_compact(key_cache)
391
value_cache_compact = permute_and_compact(value_cache)
392
tests/kernels/quantization/nvfp4_utils.py
+54
@@ -88,6 +88,60 @@ def break_fp4_bytes(a, dtype):
88
return values.reshape(m, n * 2).to(dtype=dtype)
89
90
91
+def dequant_nvfp4_kv_cache(
92
+ fp4_data: torch.Tensor,
93
+ block_scale: torch.Tensor,
94
+ global_scale: float,
95
+ head_size: int,
96
+ block_size: int,
97
+) -> torch.Tensor:
98
+ """Dequantize an NVFP4 KV cache with 4x4-swizzled block scales.
99
+
100
+ The input must be in HND layout so that the last two dims are
101
+ (block_size, last_dim). For NHD caches, permute to HND first.
102
+
103
+ Args:
104
+ fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4.
105
+ block_scale: [..., num_heads, block_size, head_size//16] fp8 block
106
+ scales (as uint8 or float8_e4m3fn).
107
+ global_scale: checkpoint dequant scale (k_scale or v_scale).
108
+ head_size: head dimension.
109
+ block_size: page size.
110
+
111
+ Returns:
112
+ [..., num_heads, block_size, head_size] float32.
113
+ """
114
+ data_dim = head_size // 2
115
+ scale_dim = head_size // 16
116
+
117
+ fp4_packed = fp4_data
118
+ sf_swizzled = block_scale.view(torch.uint8)
119
+
120
+ # Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
121
+ # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
122
+ batch_shape = sf_swizzled.shape[:-2]
123
+ T, S = block_size, scale_dim
124
+ sg = S // 4
125
+ sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
126
+ ndim = sf_reshape.ndim
127
+ # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
128
+ perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
129
+ sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
130
+ sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32)
131
+
132
+ # Unpack fp4
133
+ shape = fp4_packed.shape # [..., T, data_dim]
134
+ fp4_flat = fp4_packed.reshape(-1, data_dim)
135
+ fp4_vals = break_fp4_bytes(fp4_flat, torch.float32)
136
+ fp4_vals = fp4_vals.reshape(*shape[:-1], head_size)
137
+
138
+ # Dequant: fp4_val * block_scale * global_scale per 16-element group
139
+ return (
140
+ fp4_vals.reshape(*shape[:-1], scale_dim, 16)
141
+ * (sf_f32 * global_scale).unsqueeze(-1)
142
+ ).reshape(*shape[:-1], head_size)
143
+
144
+
145
def get_nvfp4_global_scale(a: torch.Tensor):
146
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
147
vllm/config/cache.py
+1
@@ -30,6 +30,7 @@ CacheDType = Literal[
30
"turboquant_3bit_nc",
31
"int8_per_token_head",
32
"fp8_per_token_head",
33
+ "nvfp4",
34
]
35
MambaDType = Literal["auto", "float32", "float16"]
36
MambaCacheMode = Literal["all", "align", "none"]
vllm/model_executor/layers/attention/attention.py
+4
-2
@@ -387,7 +387,9 @@ class Attention(nn.Module, AttentionLayerBase):
387
self.query_quant = None
388
if (
389
self.impl.supports_quant_query_input
390
- and self.kv_cache_dtype.startswith("fp8")
390
+ and (
391
+ self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4"
392
+ )
393
and not self.kv_cache_dtype.endswith("per_token_head")
394
):
395
is_per_head = (
@@ -492,7 +494,7 @@ class Attention(nn.Module, AttentionLayerBase):
494
# which reduces overheads during decoding.
495
# Otherwise queries are quantized using custom ops
496
# which causes decoding overheads
495
- assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"}
497
+ assert self.kv_cache_dtype in {"fp8", "fp8_e4m3", "nvfp4"}
498
499
# check if query quantization is supported
500
if self.impl.supports_quant_query_input:
vllm/utils/torch_utils.py
+122
-12
@@ -46,6 +46,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {
46
"turboquant_4bit_nc": torch.uint8,
47
"turboquant_k3v4_nc": torch.uint8,
48
"turboquant_3bit_nc": torch.uint8,
49
+ "nvfp4": torch.uint8,
50
}
51
52
TORCH_DTYPE_TO_NUMPY_DTYPE = {
@@ -59,17 +60,19 @@ TORCH_DTYPE_TO_NUMPY_DTYPE = {
60
61
62
MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {
62
- # TODO: Add more modelopt kv cache dtype
63
- # mappings here when it supported by some attention backend
64
- # (for example supports nvfp4).
63
"fp8": "fp8_e4m3",
64
+ "nvfp4": "nvfp4",
65
}
66
67
T = TypeVar("T")
68
69
70
def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
72
- return kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head")
71
+ return (
72
+ kv_cache_dtype.startswith("fp8")
73
+ or kv_cache_dtype.endswith("per_token_head")
74
+ or kv_cache_dtype == "nvfp4"
75
+ )
76
77
78
def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
@@ -299,6 +302,8 @@ def get_kv_cache_quant_algo_string(quant_cfg: dict[str, Any]) -> str | None:
302
and kv_algo.get("type") == "float"
303
):
304
kv_algo = "fp8"
305
+ elif kv_algo.get("num_bits") == 4 and kv_algo.get("type") == "float":
306
+ kv_algo = "nvfp4"
307
else:
308
# Unknown/unsupported format - return "auto" as safe fallback
309
logger.warning(
@@ -375,6 +380,95 @@ def set_random_seed(seed: int | None) -> None:
380
current_platform.manual_seed_all(seed)
381
382
383
+def nvfp4_kv_cache_full_dim(head_size: int) -> int:
384
+ """Packed last dim for NVFP4 KV cache: fp4 data + fp8 block scales."""
385
+ return head_size // 2 + head_size // 16
386
+
387
+
388
+def _nvfp4_split_data_scale(
389
+ kv_side: torch.Tensor,
390
+) -> tuple[torch.Tensor, torch.Tensor]:
391
+ """Split a single NVFP4 KV-side buffer into data and scale views.
392
+
393
+ The input is a 4D tensor for one KV side (K or V) whose last
394
+ dimension is ``full_dim = data_dim + scale_dim``. The physical
395
+ layout within each side is [data | scale], both packed contiguously.
396
+
397
+ Args:
398
+ kv_side: 4D uint8 tensor with shape
399
+ ``(num_pages, dim_1, dim_2, full_dim)``.
400
+ May be in any permutation order (NHD or HND).
401
+
402
+ Returns:
403
+ ``(data, scale)`` where
404
+ ``data`` is a uint8 view with shape
405
+ ``(num_pages, dim_1, dim_2, data_dim)``.
406
+ ``scale`` is a float8_e4m3fn view with shape
407
+ ``(num_pages, dim_1, dim_2, scale_dim)``.
408
+ """
409
+ num_pages = kv_side.shape[0]
410
+ dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2]
411
+ full_dim = kv_side.shape[3]
412
+ data_dim = full_dim * 8 // 9
413
+ scale_dim = full_dim - data_dim
414
+
415
+ data_per_kv = dim_1 * dim_2 * data_dim
416
+ page_bytes = kv_side.stride(0)
417
+
418
+ # Derive inner strides from the kv_side strides, scaling by the
419
+ # ratio of the target dim to full_dim. This preserves the physical
420
+ # layout (NHD vs HND) encoded in the input tensor's strides.
421
+ s1 = kv_side.stride(1) * data_dim // full_dim
422
+ s2 = kv_side.stride(2) * data_dim // full_dim
423
+ data_shape = (num_pages, dim_1, dim_2, data_dim)
424
+ data_strides = (page_bytes, s1, s2, 1)
425
+
426
+ s1_s = kv_side.stride(1) * scale_dim // full_dim
427
+ s2_s = kv_side.stride(2) * scale_dim // full_dim
428
+ scale_shape = (num_pages, dim_1, dim_2, scale_dim)
429
+ scale_strides = (page_bytes, s1_s, s2_s, 1)
430
+
431
+ base = kv_side.storage_offset()
432
+ data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base)
433
+ scale = torch.as_strided(
434
+ kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv
435
+ ).view(torch.float8_e4m3fn)
436
+
437
+ return data, scale
438
+
439
+
440
+def nvfp4_kv_cache_split_views(kv_cache: torch.Tensor) -> tuple[tuple, tuple]:
441
+ """Split an NVFP4 KV cache tensor into data and scale views.
442
+
443
+ Accepts either a 5D tensor ``(num_pages, 2, dim_2, dim_3, full_dim)``
444
+ or a 4D single-side tensor ``(num_pages, dim_2, dim_3, full_dim)``.
445
+
446
+ Per-page layout: [K_data | K_scale | V_data | V_scale].
447
+ Each KV side is self-contained (data followed by its scale), so the
448
+ 5D case simply splits each side independently.
449
+
450
+ The returned views are in the same dim order as the input (NHD or
451
+ HND), so callers get views matching whichever order they passed in.
452
+
453
+ Args:
454
+ kv_cache: 5D or 4D uint8 tensor where the last dimension is
455
+ ``full_dim = data_dim + scale_dim = 9 * head_size / 16``.
456
+
457
+ Returns:
458
+ For 5D input:
459
+ ``(k_data, v_data), (k_scale, v_scale)``
460
+ For 4D input (single KV side):
461
+ ``(data,), (scale,)``
462
+ """
463
+ if kv_cache.dim() == 4:
464
+ data, scale = _nvfp4_split_data_scale(kv_cache)
465
+ return (data,), (scale,)
466
+
467
+ k_data, k_scale = _nvfp4_split_data_scale(kv_cache[:, 0])
468
+ v_data, v_scale = _nvfp4_split_data_scale(kv_cache[:, 1])
469
+ return (k_data, v_data), (k_scale, v_scale)
470
+
471
+
472
def create_kv_caches_with_random_flash(
473
num_blocks: int,
474
block_size: int,
@@ -401,15 +495,31 @@ def create_kv_caches_with_random_flash(
495
value_caches: list[torch.Tensor] = []
496
497
for _ in range(num_layers):
404
- key_value_cache = torch.empty(
405
- size=kv_cache_allocation_shape, dtype=dtype, device=device
406
- ).permute(*stride_order)
407
- if cache_dtype in ["auto", "half", "bfloat16", "float"]:
408
- key_value_cache.uniform_(-scale, scale)
409
- elif cache_dtype == "fp8":
410
- _generate_random_fp8(key_value_cache, -scale, scale)
498
+ if cache_dtype == "nvfp4":
499
+ # Full page dim: fp4 data + fp8 block scales per head.
500
+ # Per page layout: [K_data | K_scale | V_data | V_scale]
501
+ # Returns [:, 0] and [:, 1] like all other dtypes.
502
+ full_dim = nvfp4_kv_cache_full_dim(head_size)
503
+ nvfp4_shape = (num_blocks, 2, block_size, num_heads, full_dim)
504
+ nvfp4_phys = tuple(nvfp4_shape[i] for i in stride_order)
505
+ inv = [stride_order.index(i) for i in range(len(stride_order))]
506
+ key_value_cache = torch.randint(
507
+ 0,
508
+ 256,
509
+ nvfp4_phys,
510
+ dtype=dtype,
511
+ device=device,
512
+ ).permute(*inv)
513
else:
412
- raise ValueError(f"Does not support key cache of type {cache_dtype}")
514
+ key_value_cache = torch.empty(
515
+ size=kv_cache_allocation_shape, dtype=dtype, device=device
516
+ ).permute(*stride_order)
517
+ if cache_dtype in ["auto", "half", "bfloat16", "float"]:
518
+ key_value_cache.uniform_(-scale, scale)
519
+ elif cache_dtype == "fp8":
520
+ _generate_random_fp8(key_value_cache, -scale, scale)
521
+ else:
522
+ raise ValueError(f"Does not support key cache of type {cache_dtype}")
523
key_caches.append(key_value_cache[:, 0])
524
value_caches.append(key_value_cache[:, 1])
525
return key_caches, value_caches
vllm/v1/attention/backends/flashinfer.py
+61
-11
@@ -42,7 +42,12 @@ from vllm.utils.flashinfer import (
42
)
43
from vllm.utils.math_utils import cdiv
44
from vllm.utils.platform_utils import is_pin_memory_available
45
-from vllm.utils.torch_utils import is_quantized_kv_cache, is_strictly_contiguous
45
+from vllm.utils.torch_utils import (
46
+ is_quantized_kv_cache,
47
+ is_strictly_contiguous,
48
+ nvfp4_kv_cache_full_dim,
49
+ nvfp4_kv_cache_split_views,
50
+)
51
from vllm.v1.attention.backend import (
52
AttentionBackend,
53
AttentionCGSupport,
@@ -355,6 +360,10 @@ class FlashInferBackend(AttentionBackend):
360
head_size: int,
361
cache_dtype_str: str = "auto",
362
) -> tuple[int, ...]:
363
+ if cache_dtype_str == "nvfp4":
364
+ # Packed layout: fp4 data + fp8 block scales in last dim
365
+ last_dim = nvfp4_kv_cache_full_dim(head_size)
366
+ return (num_blocks, 2, block_size, num_kv_heads, last_dim)
367
return (num_blocks, 2, block_size, num_kv_heads, head_size)
368
369
@staticmethod
@@ -608,11 +617,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
617
self.cache_dtype = self.cache_config.cache_dtype
618
# Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
619
# storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
611
- self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
612
- self.cache_dtype
613
- )
620
+ self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4"
621
+ if self.is_kvcache_nvfp4:
622
+ # For NVFP4, kv_cache_dtype stays as the string "nvfp4"
623
+ # which is passed to FlashInferImpl
624
+ self.kv_cache_dtype = self.cache_dtype
625
+ raise NotImplementedError("nvfp4 KV cache is not yet supported")
626
+ else:
627
+ self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
628
+ self.cache_dtype
629
+ )
630
else:
631
self.cache_dtype = "auto"
632
+ self.is_kvcache_nvfp4 = False
633
assert self.kv_cache_spec.dtype == self.model_config.dtype
634
self.kv_cache_dtype = self.kv_cache_spec.dtype
635
@@ -626,7 +643,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
643
can_use_trtllm
644
and not vllm_config.attention_config.disable_flashinfer_q_quantization
645
):
629
- self.q_data_type = self.kv_cache_dtype
646
+ if self.is_kvcache_nvfp4:
647
+ # NVFP4 KV cache uses FP8 quantized queries
648
+ self.q_data_type = FlashInferBackend.get_fp8_dtype_for_flashinfer(
649
+ "fp8_e4m3"
650
+ )
651
+ else:
652
+ self.q_data_type = self.kv_cache_dtype
653
else:
654
self.q_data_type = self.model_config.dtype
655
@@ -1228,6 +1251,8 @@ class FlashInferImpl(AttentionImpl):
1251
self.sliding_window[0] if self.sliding_window is not None else -1
1252
)
1253
self.kv_cache_dtype = kv_cache_dtype
1254
+ self.is_kvcache_nvfp4 = kv_cache_dtype == "nvfp4"
1255
+ self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0
1256
self.logits_soft_cap = logits_soft_cap
1257
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
1258
@@ -1406,7 +1431,16 @@ class FlashInferImpl(AttentionImpl):
1431
num_prefill_tokens = attn_metadata.num_prefill_tokens
1432
1433
stride_order = FlashInferBackend.get_kv_cache_stride_order()
1409
- kv_cache_permute = kv_cache.permute(*stride_order)
1434
+ kv_cache_permute = kv_cache.permute(*stride_order) # HND and contiguous
1435
+
1436
+ # For NVFP4, the kv_cache last dim is full_dim (data + scale packed).
1437
+ # Split into correctly-strided data and scale views.
1438
+ nvfp4_kv_data = None
1439
+ nvfp4_kv_block_scales = None
1440
+ if self.is_kvcache_nvfp4:
1441
+ nvfp4_kv_data, nvfp4_kv_block_scales = nvfp4_kv_cache_split_views(
1442
+ kv_cache_permute
1443
+ )
1444
1445
use_dcp = self.dcp_world_size > 1
1446
@@ -1490,8 +1524,20 @@ class FlashInferImpl(AttentionImpl):
1524
assert self.o_sf_scale is None
1525
out = output[num_decode_tokens:]
1526
1493
- if attn_metadata.q_data_type != FP8_DTYPE and is_quantized_kv_cache(
1494
- self.kv_cache_dtype
1527
+ prefill_kv_block_scales = None
1528
+ if self.is_kvcache_nvfp4:
1529
+ # NVFP4 trtllm-gen kernel requires FP8 query.
1530
+ assert attn_metadata.q_data_type == FP8_DTYPE, (
1531
+ "NVFP4 KV cache requires FP8 quantized queries for "
1532
+ "trtllm-gen prefill. Set "
1533
+ "disable_flashinfer_q_quantization=False."
1534
+ )
1535
+ mock_kv_cache = nvfp4_kv_data
1536
+ mock_block_table = block_tables_prefill
1537
+ prefill_kv_block_scales = nvfp4_kv_block_scales # noqa: F841
1538
+ elif (
1539
+ attn_metadata.q_data_type != FP8_DTYPE
1540
+ and self.kv_cache_dtype.startswith("fp8")
1541
):
1542
# TRTLLM prefill attention does not support BF16 Q
1543
# and fp8 kv cache. So to enable prefill attention
@@ -1636,7 +1682,9 @@ class FlashInferImpl(AttentionImpl):
1682
1683
trtllm_batch_decode_with_kv_cache(
1684
query=decode_query,
1639
- kv_cache=kv_cache_permute,
1685
+ kv_cache=nvfp4_kv_data
1686
+ if self.is_kvcache_nvfp4
1687
+ else kv_cache_permute,
1688
workspace_buffer=workspace_buffer,
1689
block_tables=block_tables_decode,
1690
seq_lens=seq_lens_decode,
@@ -1667,11 +1715,13 @@ class FlashInferImpl(AttentionImpl):
1715
# and value[:num_actual_tokens] because the reshape_and_cache_flash
1716
# op uses the slot_mapping's shape to determine the number of
1717
# actual tokens.
1718
+ k_cache = kv_cache[:, 0]
1719
+ v_cache = kv_cache[:, 1]
1720
torch.ops._C_cache_ops.reshape_and_cache_flash(
1721
key,
1722
value,
1673
- kv_cache[:, 0],
1674
- kv_cache[:, 1],
1723
+ k_cache,
1724
+ v_cache,
1725
slot_mapping,
1726
self.kv_cache_dtype,
1727
layer._k_scale,
vllm/v1/kv_cache_interface.py
+27
-3
@@ -17,7 +17,7 @@ from vllm.logger import init_logger
17
if TYPE_CHECKING:
18
from vllm.config import VllmConfig
19
from vllm.utils.math_utils import cdiv
20
-from vllm.utils.torch_utils import get_dtype_size
20
+from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
21
22
logger = init_logger(__name__)
23
@@ -38,11 +38,20 @@ class KVQuantMode(IntEnum):
38
FP8_PER_TENSOR = 1 # per-tensor scales (current fp8 path)
39
INT8_PER_TOKEN_HEAD = 2 # per-token-head dynamic scales for int8
40
FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8
41
+ NVFP4 = 4 # packed fp4 data + fp8 block scales
42
43
@property
44
def is_per_token_head(self) -> bool:
45
"""True for any per-token-head quantization mode."""
45
- return self >= 2
46
+ return self in (
47
+ KVQuantMode.INT8_PER_TOKEN_HEAD,
48
+ KVQuantMode.FP8_PER_TOKEN_HEAD,
49
+ )
50
+
51
+ @property
52
+ def is_nvfp4(self) -> bool:
53
+ """True for NVFP4 packed quantization mode."""
54
+ return self == KVQuantMode.NVFP4
55
56
57
def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
@@ -51,7 +60,9 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
60
return KVQuantMode.INT8_PER_TOKEN_HEAD
61
if kv_cache_dtype == "fp8_per_token_head":
62
return KVQuantMode.FP8_PER_TOKEN_HEAD
54
- if kv_cache_dtype.startswith("fp8"):
63
+ if kv_cache_dtype == "nvfp4":
64
+ return KVQuantMode.NVFP4
65
+ if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
66
return KVQuantMode.FP8_PER_TENSOR
67
return KVQuantMode.NONE
68
@@ -237,6 +248,19 @@ class FullAttentionSpec(AttentionSpec):
248
249
@property
250
def real_page_size_bytes(self) -> int:
251
+ if self.kv_quant_mode.is_nvfp4:
252
+ # Packed layout per head: fp4 data + fp8 block scales.
253
+ # fp4 data: head_size//2 bytes (2 fp4 values per byte)
254
+ # fp8 block scale: head_size//16 bytes (1 scale per 16 elements)
255
+ last_dim = nvfp4_kv_cache_full_dim(
256
+ self.head_size
257
+ ) + nvfp4_kv_cache_full_dim(self.head_size_v)
258
+ return (
259
+ self.block_size
260
+ * self.num_kv_heads
261
+ * last_dim
262
+ * get_dtype_size(self.dtype)
263
+ )
264
return (
265
self.block_size
266
* self.num_kv_heads