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 4, 10:00 UTC
Add nvfp4 kv cache support (#40177)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
sychen52 committed
Apr 30, 2026 at 21:55 UTC
947138b6c22f9a4751b63b9aa75a2bc4b42835e9
8 files changed
+503
-96
docs/design/attention_backends.md
+1
-1
@@ -169,7 +169,7 @@ Priority is **1 = highest** (tried first).
169
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
170
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
171
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
172
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
172
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
173
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
174
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
175
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+146
-50
@@ -5,12 +5,17 @@ import pytest
5
import torch
6
7
from tests.kernels.quantization.nvfp4_utils import (
8
+ dequant_nvfp4_kv_cache,
9
dequantize_nvfp4_to_dtype,
10
get_nvfp4_global_scale,
11
)
12
from vllm.platforms import current_platform
13
from vllm.utils.math_utils import round_up
13
-from vllm.utils.torch_utils import set_random_seed
14
+from vllm.utils.torch_utils import (
15
+ nvfp4_kv_cache_full_dim,
16
+ nvfp4_kv_cache_split_views,
17
+ set_random_seed,
18
+)
19
20
if not current_platform.is_device_capability_family(100):
21
pytest.skip(
@@ -33,6 +38,117 @@ def to_float8(x, dtype=torch.float8_e4m3fn):
38
return x_scl_sat.to(dtype), scale.float().reciprocal()
39
40
41
+def build_paged_kv_metadata(
42
+ seq_lens: torch.Tensor,
43
+ block_tables: torch.Tensor,
44
+ block_size: int,
45
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
46
+ """Build paged-KV indptr/indices/last_page_lens from seq_lens + block_tables."""
47
+ kv_indptr = [0]
48
+ kv_indices = []
49
+ kv_last_page_lens = []
50
+ for i in range(len(seq_lens)):
51
+ sl = int(seq_lens[i])
52
+ assert sl > 0
53
+ nb = (sl + block_size - 1) // block_size
54
+ kv_indices.extend(block_tables[i, :nb].tolist())
55
+ kv_indptr.append(kv_indptr[-1] + nb)
56
+ kv_last_page_lens.append(sl % block_size or block_size)
57
+ return (
58
+ torch.tensor(kv_indptr, dtype=torch.int32),
59
+ torch.tensor(kv_indices, dtype=torch.int32),
60
+ torch.tensor(kv_last_page_lens, dtype=torch.int32),
61
+ )
62
+
63
+
64
+def make_nvfp4_kv_cache(
65
+ kv_bf16_hnd: torch.Tensor, block_size: int, head_size: int
66
+) -> tuple:
67
+ """Quantize bf16 KV cache to nvfp4 via reshape_and_cache_flash.
68
+
69
+ Returns (k_data, v_data), (k_scales, v_scales), kv_scale, ref_kv_bf16.
70
+ """
71
+ num_blocks, _, num_kv_heads, _, _ = kv_bf16_hnd.shape
72
+ kv_scale_val = (kv_bf16_hnd.abs().amax() / 448.0).item()
73
+ kv_scale_tensor = torch.tensor(
74
+ kv_scale_val, dtype=torch.float32, device=kv_bf16_hnd.device
75
+ )
76
+
77
+ # Allocate in HND physical order, permute to NHD logical order.
78
+ # hnd_order swaps dims 2↔3; it is its own inverse.
79
+ full_dim = nvfp4_kv_cache_full_dim(head_size)
80
+ hnd_order = (0, 1, 3, 2, 4)
81
+ kv_cache = torch.zeros(
82
+ (num_blocks, 2, num_kv_heads, block_size, full_dim),
83
+ dtype=torch.uint8,
84
+ device=kv_bf16_hnd.device,
85
+ ).permute(*hnd_order)
86
+
87
+ # Flatten NHD [N, T, H, D] → token tensors [N*T, H, D] for the kernel.
88
+ num_tokens = num_blocks * block_size
89
+ k_tokens = (
90
+ kv_bf16_hnd[:, 0]
91
+ .permute(0, 2, 1, 3)
92
+ .reshape(num_tokens, num_kv_heads, head_size)
93
+ )
94
+ v_tokens = (
95
+ kv_bf16_hnd[:, 1]
96
+ .permute(0, 2, 1, 3)
97
+ .reshape(num_tokens, num_kv_heads, head_size)
98
+ )
99
+ slot_mapping = torch.arange(num_tokens, dtype=torch.long, device=kv_bf16_hnd.device)
100
+
101
+ # reshape_and_cache_flash: kernel receives kv_cache[:, 0] and [:, 1]
102
+ # (full K/V buffers containing both data and scale).
103
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
104
+ k_tokens,
105
+ v_tokens,
106
+ kv_cache[:, 0],
107
+ kv_cache[:, 1],
108
+ slot_mapping,
109
+ "nvfp4",
110
+ kv_scale_tensor,
111
+ kv_scale_tensor,
112
+ )
113
+
114
+ # Split in HND order for trtllm kernel (expects HND numTokensPerPage).
115
+ kv_cache_hnd = kv_cache.permute(*hnd_order)
116
+ (k_data, v_data), (k_scales, v_scales) = nvfp4_kv_cache_split_views(kv_cache_hnd)
117
+
118
+ # Dequantize for the FA2 reference baseline.
119
+ ref_k = dequant_nvfp4_kv_cache(
120
+ k_data, k_scales, kv_scale_val, head_size, block_size
121
+ ).to(torch.bfloat16)
122
+ ref_v = dequant_nvfp4_kv_cache(
123
+ v_data, v_scales, kv_scale_val, head_size, block_size
124
+ ).to(torch.bfloat16)
125
+ ref_kv_bf16 = torch.stack([ref_k, ref_v], dim=1) # [N, 2, H, T, D]
126
+
127
+ return (k_data, v_data), (k_scales, v_scales), kv_scale_val, ref_kv_bf16
128
+
129
+
130
+def make_quantized_kv_cache(
131
+ kv_cache: torch.Tensor,
132
+ kv_quant_dtype: torch.dtype,
133
+ block_size: int,
134
+ head_size: int,
135
+) -> tuple:
136
+ """Quantize kv_cache based on dtype. Returns (kv_cache, kv_cache_sf,
137
+ kv_scale, ref_kv_cache, is_nvfp4_kv)."""
138
+ is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE
139
+ if is_nvfp4_kv:
140
+ data, scales, kv_scale, ref = make_nvfp4_kv_cache(
141
+ kv_cache, block_size, head_size
142
+ )
143
+ return data, scales, kv_scale, ref, True
144
+ elif kv_quant_dtype == FP8_DTYPE:
145
+ kv_fp8, kv_scale = to_float8(kv_cache)
146
+ ref = kv_fp8.to(kv_cache.dtype) * kv_scale
147
+ return kv_fp8, None, kv_scale, ref, False
148
+ else:
149
+ return kv_cache, None, 1.0, kv_cache, False
150
+
151
+
152
DTYPE = [torch.bfloat16]
153
QUANT_DTYPES = [
154
# (q_quant_dtype, kv_quant_dtype, o_quant_dtype)
@@ -41,6 +157,7 @@ QUANT_DTYPES = [
157
(FP8_DTYPE, FP8_DTYPE, None),
158
(FP8_DTYPE, FP8_DTYPE, FP8_DTYPE),
159
(FP8_DTYPE, FP8_DTYPE, FP4_DTYPE),
160
+ (FP8_DTYPE, FP4_DTYPE, FP8_DTYPE), # nvfp4 KV cache
161
]
162
BATCH_SIZE = [4, 12]
163
MAX_SEQ_LENS = [(1024, 4096)]
@@ -127,35 +244,19 @@ def test_flashinfer_trtllm_decode_with_baseline(
244
max_seq_len = torch.max(seq_lens).item()
245
246
kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
130
- if kv_quant_dtype == FP8_DTYPE:
131
- kv_cache, kv_scale = to_float8(kv_cache)
132
- ref_kv_cache = kv_cache.to(dtype) * kv_scale
133
- else:
134
- kv_scale = 1.0
135
- ref_kv_cache = kv_cache
247
+ kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, is_nvfp4_kv = (
248
+ make_quantized_kv_cache(kv_cache, kv_quant_dtype, block_size, head_size)
249
+ )
250
+
251
k_scale = v_scale = kv_scale
252
253
max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
254
block_tables = torch.randint(
255
0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32
256
)
142
- kv_indptr = [0]
143
- kv_indices = []
144
- kv_last_page_lens = []
145
- for i in range(batch_size):
146
- seq_len = seq_lens[i]
147
- assert seq_len > 0
148
- num_blocks = (seq_len + block_size - 1) // block_size
149
- kv_indices.extend(block_tables[i, :num_blocks])
150
- kv_indptr.append(kv_indptr[-1] + num_blocks)
151
- kv_last_page_len = seq_len % block_size
152
- if kv_last_page_len == 0:
153
- kv_last_page_len = block_size
154
- kv_last_page_lens.append(kv_last_page_len)
155
-
156
- kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
157
- kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
158
- kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
257
+ kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata(
258
+ seq_lens, block_tables, block_size
259
+ )
260
workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
261
262
# Baseline Decode
@@ -225,6 +326,7 @@ def test_flashinfer_trtllm_decode_with_baseline(
326
sinks=sinks,
327
o_sf_scale=o_sf_scale_float,
328
out=output_trtllm,
329
+ kv_cache_sf=kv_cache_sf,
330
)
331
if o_quant_dtype == FP8_DTYPE:
332
output_trtllm = output_trtllm.to(dtype) * o_scale
@@ -237,7 +339,9 @@ def test_flashinfer_trtllm_decode_with_baseline(
339
)
340
output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2])
341
240
- if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
342
+ if is_nvfp4_kv:
343
+ rtol, atol = 1.0, 1.0 # nvfp4 has higher quantization error
344
+ elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
345
rtol, atol = 7e-2, 9e-2
346
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
347
rtol, atol = 3e-2, 4e-2
@@ -287,7 +391,12 @@ def test_flashinfer_trtllm_prefill_with_baseline(
391
kv_quant_dtype = kv_quant_dtype or dtype
392
o_quant_dtype = o_quant_dtype or dtype
393
290
- if q_quant_dtype != kv_quant_dtype:
394
+ # FP8 Q + nvfp4 KV is the required combination for the nvfp4 KV path.
395
+ # All other mixed Q/KV dtype combinations are unsupported.
396
+ is_nvfp4_kv = kv_quant_dtype == FP4_DTYPE
397
+ if q_quant_dtype != kv_quant_dtype and not (
398
+ q_quant_dtype == FP8_DTYPE and is_nvfp4_kv
399
+ ):
400
pytest.skip("Skipped mixed QKV dtypes for prefill")
401
402
max_q_len, max_kv_len = max_seq_lens
@@ -329,35 +438,19 @@ def test_flashinfer_trtllm_prefill_with_baseline(
438
max_seq_len = torch.max(seq_lens).item()
439
440
kv_cache = torch.randn(kv_cache_shape, dtype=dtype)
332
- if kv_quant_dtype == FP8_DTYPE:
333
- kv_cache, kv_scale = to_float8(kv_cache)
334
- ref_kv_cache = kv_cache.to(dtype) * kv_scale
335
- else:
336
- kv_scale = 1.0
337
- ref_kv_cache = kv_cache
441
+ kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, is_nvfp4_kv = (
442
+ make_quantized_kv_cache(kv_cache, kv_quant_dtype, block_size, head_size)
443
+ )
444
+
445
k_scale = v_scale = kv_scale
446
447
max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
448
block_tables = torch.randint(
449
0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32
450
)
344
- kv_indptr = [0]
345
- kv_indices = []
346
- kv_last_page_lens = []
347
- for i in range(batch_size):
348
- seq_len = seq_lens[i]
349
- assert seq_len > 0
350
- num_blocks = (seq_len + block_size - 1) // block_size
351
- kv_indices.extend(block_tables[i, :num_blocks])
352
- kv_indptr.append(kv_indptr[-1] + num_blocks)
353
- kv_last_page_len = seq_len % block_size
354
- if kv_last_page_len == 0:
355
- kv_last_page_len = block_size
356
- kv_last_page_lens.append(kv_last_page_len)
357
-
358
- kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32)
359
- kv_indices = torch.tensor(kv_indices, dtype=torch.int32)
360
- kv_last_page_lens = torch.tensor(kv_last_page_lens, dtype=torch.int32)
451
+ kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata(
452
+ seq_lens, block_tables, block_size
453
+ )
454
workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)
455
456
# Baseline Prefill
@@ -431,6 +524,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(
524
sinks=sinks,
525
o_sf_scale=o_sf_scale_float,
526
out=output_trtllm,
527
+ kv_cache_sf=kv_cache_sf,
528
)
529
if o_quant_dtype == FP8_DTYPE:
530
output_trtllm = output_trtllm.to(dtype) * o_scale
@@ -443,7 +537,9 @@ def test_flashinfer_trtllm_prefill_with_baseline(
537
)
538
output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2])
539
446
- if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
540
+ if is_nvfp4_kv:
541
+ rtol, atol = 1.0, 1.5 # nvfp4 has higher quantization error
542
+ elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:
543
rtol, atol = 3e-1, 4e-1
544
elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:
545
rtol, atol = 4e-2, 6e-2
tests/v1/attention/test_trtllm_attention_integration.py
+198
-23
@@ -17,13 +17,13 @@ from tests.v1.attention.utils import (
17
from vllm.config import set_current_vllm_config
18
from vllm.platforms import current_platform
19
from vllm.utils.math_utils import cdiv
20
-from vllm.utils.torch_utils import set_random_seed
20
+from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim, set_random_seed
21
from vllm.v1.attention.backends.utils import (
22
PerLayerParameters,
23
get_kv_cache_layout,
24
set_kv_cache_layout,
25
)
26
-from vllm.v1.kv_cache_interface import FullAttentionSpec
26
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
27
28
if not current_platform.is_device_capability_family(100):
29
pytest.skip(
@@ -53,6 +53,7 @@ class MockAttentionLayer:
53
54
55
MODEL = "Qwen/Qwen2.5-0.5B"
56
+MODEL_NVFP4 = "Qwen/Qwen3-4B" # nvfp4 needs head_dim >= 128 (or 80)
57
BLOCK_SIZE = 16
58
NUM_GPU_BLOCKS = 8192
59
DEVICE_TYPE = current_platform.device_type
@@ -169,19 +170,129 @@ def _create_hnd_kv_cache(
170
return kv_cache
171
172
172
-def _run_trtllm_integration(batch_spec):
173
+def _create_nvfp4_hnd_kv_cache(
174
+ k_contexts,
175
+ v_contexts,
176
+ block_size,
177
+ num_kv_heads,
178
+ head_size,
179
+ dtype,
180
+ device,
181
+ num_blocks,
182
+ common_attn_metadata,
183
+ kv_scale_val,
184
+):
185
+ """Create an nvfp4 KV cache by quantizing bf16 context via
186
+ reshape_and_cache_flash, using the same block-table layout as
187
+ _create_hnd_kv_cache.
188
+
189
+ The returned tensor is dtype ``uint8`` with shape
190
+ ``(num_blocks, 2, block_size, num_kv_heads, full_dim)`` in logical
191
+ (NHD) order, but physically permuted to HND layout via stride order
192
+ ``(0, 1, 3, 2, 4)`` (i.e. ``num_kv_heads`` before ``block_size``).
193
+
194
+ The last dimension ``full_dim = head_size // 2 + head_size // 16``
195
+ packs two regions contiguously:
196
+ - **FP4 data** (``head_size // 2`` bytes): pairs of E2M1 values,
197
+ two per byte.
198
+ - **FP8 block scales** (``head_size // 16`` bytes): one E4M3
199
+ scale per 16-element block.
200
+
201
+ Dimension 1 indexes K (``[:, 0]``) and V (``[:, 1]``).
202
+
203
+ Args:
204
+ k_contexts: List of key context tensors, one per sequence.
205
+ v_contexts: List of value context tensors, one per sequence.
206
+ block_size: Number of tokens per cache block.
207
+ num_kv_heads: Number of key/value heads.
208
+ head_size: Head dimension (must be divisible by 16).
209
+ dtype: Source data type for the bf16 intermediate cache.
210
+ device: Target device.
211
+ num_blocks: Total number of blocks to allocate.
212
+ common_attn_metadata: Metadata containing block tables and
213
+ sequence lengths.
214
+ kv_scale_val: Scalar float used as both k_scale and v_scale
215
+ during quantization.
216
+
217
+ Returns:
218
+ ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided).
219
+ """
220
+ # First create a bf16 HND cache so block tables are populated.
221
+ bf16_cache = _create_hnd_kv_cache(
222
+ k_contexts,
223
+ v_contexts,
224
+ block_size,
225
+ num_kv_heads,
226
+ head_size,
227
+ dtype,
228
+ device,
229
+ num_blocks,
230
+ common_attn_metadata,
231
+ )
232
+
233
+ # Allocate nvfp4 cache: same shape but with full_dim (data + scale).
234
+ full_dim = nvfp4_kv_cache_full_dim(head_size)
235
+ hnd_order = (0, 1, 3, 2, 4)
236
+ nvfp4_cache = torch.zeros(
237
+ (num_blocks, 2, num_kv_heads, block_size, full_dim),
238
+ dtype=torch.uint8,
239
+ device=device,
240
+ ).permute(*hnd_order)
241
+
242
+ # Flatten bf16 context into tokens and quantize via reshape_and_cache_flash.
243
+ # bf16_cache is (num_blocks, 2, block_size, num_kv_heads, head_size) logical
244
+ # with HND physical strides.
245
+ block_table = common_attn_metadata.block_table_tensor
246
+ seq_lens = common_attn_metadata.seq_lens.cpu()
247
+ query_lens = (
248
+ common_attn_metadata.query_start_loc_cpu[1:]
249
+ - common_attn_metadata.query_start_loc_cpu[:-1]
250
+ )
251
+ kv_scale_t = torch.tensor(kv_scale_val, dtype=torch.float32, device=device)
252
+
253
+ for i in range(len(k_contexts)):
254
+ ctx_len = int(seq_lens[i]) - int(query_lens[i])
255
+ if ctx_len == 0:
256
+ continue
257
+ # Gather context tokens from the bf16 cache using block table.
258
+ n_ctx_blocks = (ctx_len + block_size - 1) // block_size
259
+ blocks = block_table[i, :n_ctx_blocks]
260
+ # bf16_cache[:, kv_idx] is (num_blocks, block_size, num_kv_heads, head_size)
261
+ k_ctx = bf16_cache[blocks, 0].reshape(-1, num_kv_heads, head_size)[:ctx_len]
262
+ v_ctx = bf16_cache[blocks, 1].reshape(-1, num_kv_heads, head_size)[:ctx_len]
263
+ # Build slot mapping for these context tokens.
264
+ token_offsets = torch.arange(ctx_len, device=device)
265
+ block_indices = token_offsets // block_size
266
+ intra_offsets = token_offsets % block_size
267
+ slots = block_table[i, block_indices] * block_size + intra_offsets
268
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
269
+ k_ctx,
270
+ v_ctx,
271
+ nvfp4_cache[:, 0],
272
+ nvfp4_cache[:, 1],
273
+ slots,
274
+ "nvfp4",
275
+ kv_scale_t,
276
+ kv_scale_t,
277
+ )
278
+
279
+ return nvfp4_cache
280
+
281
+
282
+def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL):
283
"""Run TRTLLM attention through the full FlashInfer pipeline
284
and compare against an SDPA reference."""
285
set_random_seed(42)
286
device = torch.device(f"{DEVICE_TYPE}:0")
287
288
vllm_config = create_vllm_config(
179
- model_name=MODEL,
289
+ model_name=model_name,
290
max_model_len=max(batch_spec.seq_lens),
291
block_size=BLOCK_SIZE,
292
num_gpu_blocks=NUM_GPU_BLOCKS,
293
)
294
vllm_config.attention_config.use_trtllm_attention = True
295
+ vllm_config.cache_config.cache_dtype = kv_cache_dtype
296
297
num_q_heads = vllm_config.model_config.get_num_attention_heads(
298
vllm_config.parallel_config
@@ -248,28 +359,51 @@ def _run_trtllm_integration(batch_spec):
359
common_attn_metadata = create_common_attn_metadata(batch_spec, BLOCK_SIZE, device)
360
361
# 2. Create HND KV cache
251
- kv_cache = _create_hnd_kv_cache(
252
- k_contexts,
253
- v_contexts,
254
- BLOCK_SIZE,
255
- num_kv_heads,
256
- head_size,
257
- dtype,
258
- device,
259
- NUM_GPU_BLOCKS,
260
- common_attn_metadata,
261
- )
362
+ is_nvfp4 = kv_cache_dtype == "nvfp4"
363
+ if is_nvfp4:
364
+ # Compute a global scale from the context data.
365
+ all_ctx = torch.cat(k_contexts + v_contexts, dim=0)
366
+ kv_scale_val = (all_ctx.abs().amax() / 448.0).item()
367
+ kv_cache = _create_nvfp4_hnd_kv_cache(
368
+ k_contexts,
369
+ v_contexts,
370
+ BLOCK_SIZE,
371
+ num_kv_heads,
372
+ head_size,
373
+ dtype,
374
+ device,
375
+ NUM_GPU_BLOCKS,
376
+ common_attn_metadata,
377
+ kv_scale_val,
378
+ )
379
+ else:
380
+ kv_scale_val = 1.0
381
+ kv_cache = _create_hnd_kv_cache(
382
+ k_contexts,
383
+ v_contexts,
384
+ BLOCK_SIZE,
385
+ num_kv_heads,
386
+ head_size,
387
+ dtype,
388
+ device,
389
+ NUM_GPU_BLOCKS,
390
+ common_attn_metadata,
391
+ )
392
393
# 3. Run through FlashInfer with TRTLLM enabled
394
set_kv_cache_layout("HND")
395
get_kv_cache_layout.cache_clear()
396
397
try:
398
+ is_nvfp4 = kv_cache_dtype == "nvfp4"
399
+ kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE
400
+ spec_dtype = torch.uint8 if is_nvfp4 else dtype
401
kv_cache_spec = FullAttentionSpec(
402
block_size=BLOCK_SIZE,
403
num_kv_heads=num_kv_heads,
404
head_size=head_size,
272
- dtype=dtype,
405
+ dtype=spec_dtype,
406
+ kv_quant_mode=kv_quant_mode,
407
)
408
layer_names = ["test_layer_0"]
409
@@ -312,10 +446,20 @@ def _run_trtllm_integration(batch_spec):
446
num_kv_heads=num_kv_heads,
447
alibi_slopes=None,
448
sliding_window=None,
315
- kv_cache_dtype="auto",
449
+ kv_cache_dtype=kv_cache_dtype,
450
)
451
452
mock_layer = MockAttentionLayer(device)
453
+ if is_nvfp4:
454
+ # For nvfp4, k_scale/v_scale are the global quantization
455
+ # scales (amax/448) used by reshape_and_cache_flash.
456
+ kv_scale_t = torch.tensor(
457
+ kv_scale_val, dtype=torch.float32, device=device
458
+ )
459
+ mock_layer._k_scale = kv_scale_t
460
+ mock_layer._v_scale = kv_scale_t
461
+ mock_layer._k_scale_float = kv_scale_val
462
+ mock_layer._v_scale_float = kv_scale_val
463
output = torch.empty_like(query_vllm)
464
465
impl.do_kv_cache_update(
@@ -326,6 +470,23 @@ def _run_trtllm_integration(batch_spec):
470
attn_metadata.slot_mapping,
471
)
472
473
+ # nvfp4 trtllm kernel requires FP8 queries. In the real
474
+ # pipeline the attention layer handles this; here we
475
+ # quantize manually.
476
+ if is_nvfp4:
477
+ finfo = torch.finfo(torch.float8_e4m3fn)
478
+ q_amax = query_vllm.abs().amax().clamp(min=1e-12)
479
+ q_s = (finfo.max / q_amax * 0.1).item()
480
+ query_vllm = (
481
+ (query_vllm * q_s)
482
+ .clamp(finfo.min, finfo.max)
483
+ .to(torch.float8_e4m3fn)
484
+ )
485
+ mock_layer._q_scale = torch.tensor(
486
+ 1.0 / q_s, dtype=torch.float32, device=device
487
+ )
488
+ mock_layer._q_scale_float = 1.0 / q_s
489
+
490
output = impl.forward(
491
mock_layer,
492
query_vllm,
@@ -337,12 +498,11 @@ def _run_trtllm_integration(batch_spec):
498
)
499
500
# 4. Compare against SDPA reference
340
- torch.testing.assert_close(
341
- output,
342
- sdpa_output,
343
- atol=1e-2,
344
- rtol=1e-2,
345
- )
501
+ if is_nvfp4:
502
+ atol, rtol = 1.0, 1.0 # nvfp4 has higher quantization error
503
+ else:
504
+ atol, rtol = 1e-2, 1e-2
505
+ torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol)
506
507
finally:
508
set_kv_cache_layout(None)
@@ -359,3 +519,18 @@ def test_trtllm_gen_full_attention_integration(batch_spec_name: str):
519
MetadataBuilder.build() -> FlashInferImpl.forward() pipeline,
520
with real TRTLLM kernels on Blackwell."""
521
_run_trtllm_integration(BATCH_SPECS[batch_spec_name])
522
+
523
+
524
+@pytest.mark.parametrize(
525
+ "batch_spec_name",
526
+ list(BATCH_SPECS.keys()),
527
+)
528
+@torch.inference_mode()
529
+def test_trtllm_gen_nvfp4_kv_integration(batch_spec_name: str):
530
+ """Test TRTLLM attention with nvfp4 KV cache through the full
531
+ FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline."""
532
+ _run_trtllm_integration(
533
+ BATCH_SPECS[batch_spec_name],
534
+ kv_cache_dtype="nvfp4",
535
+ model_name=MODEL_NVFP4,
536
+ )
tools/pre_commit/generate_attention_backend_docs.py
+15
@@ -788,6 +788,11 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
788
if not trtllm_compute_cap:
789
return {}
790
791
+ # KV cache dtypes that only work with a dedicated kernel (e.g. nvfp4
792
+ # requires the SM100 NVFP4 MHA kernel) and should not appear in the
793
+ # generic attention-backend feature matrix.
794
+ kernel_only_kv_dtypes = ["nvfp4"]
795
+
796
return {
797
"native": {
798
# Native FlashInfer: everything except SM100
@@ -798,6 +803,7 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
803
"compute_capability": trtllm_compute_cap,
804
"supports_sink": True,
805
},
806
+ "exclude_kv_dtypes": kernel_only_kv_dtypes,
807
}
808
809
@@ -963,6 +969,15 @@ def _expand_flashinfer_variants(
969
native["supports_sink"] = fi_features["native"]["supports_sink"]
970
native["compute_capability"] = f"{min_cc}.x-9.x"
971
972
+ # Remove KV dtypes only supported by SM100 kernels (e.g. nvfp4)
973
+ exclude = fi_features.get("exclude_kv_dtypes", [])
974
+ if exclude:
975
+ native["kv_cache_dtypes"] = ", ".join(
976
+ d
977
+ for d in (d.strip() for d in native["kv_cache_dtypes"].split(","))
978
+ if d not in exclude
979
+ )
980
+
981
# Create TRTLLM entry
982
trtllm = backend.copy()
983
trtllm["version"] = "TRTLLM†"
vllm/config/vllm.py
+12
@@ -1884,6 +1884,18 @@ class VllmConfig:
1884
"in the middle of a mm input"
1885
)
1886
1887
+ @model_validator(mode="after")
1888
+ def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig":
1889
+ if self.model_config is None:
1890
+ return self
1891
+ if self.cache_config.cache_dtype == "nvfp4" and self.model_config.use_mla:
1892
+ raise ValueError(
1893
+ "nvfp4 KV cache is not supported with MLA (Multi-head Latent "
1894
+ "Attention) backends. Please use a different --kv-cache-dtype "
1895
+ "(e.g., 'fp8' or 'auto') for MLA models such as DeepSeek."
1896
+ )
1897
+ return self
1898
+
1899
@model_validator(mode="after")
1900
def validate_mamba_block_size(self) -> "VllmConfig":
1901
if self.model_config is None:
vllm/model_executor/layers/quantization/modelopt.py
+7
-7
@@ -114,12 +114,12 @@ QUANT_ALGOS = [
114
# MIXED_PRECISION,
115
"MIXED_PRECISION",
116
]
117
-KV_CACHE_QUANT_ALGOS = ["FP8"]
117
+KV_CACHE_QUANT_ALGOS = ["FP8", "NVFP4"]
118
119
120
-class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):
120
+class ModelOptKVCacheMethod(BaseKVCacheMethod):
121
"""
122
- Supports loading kv-cache scaling factors from FP8 checkpoints.
122
+ Supports loading kv-cache scaling factors from FP8 or NVFP4 checkpoints.
123
"""
124
125
def __init__(self, quant_config: "ModelOptQuantConfigBase"):
@@ -995,7 +995,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
995
996
ModelOptFp8Config.LinearMethodCls = ModelOptFp8LinearMethod
997
ModelOptFp8Config.FusedMoEMethodCls = ModelOptFp8MoEMethod
998
-ModelOptFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
998
+ModelOptFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod
999
1000
1001
class ModelOptNvFp4Config(ModelOptQuantConfigBase):
@@ -1488,7 +1488,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1488
1489
ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod
1490
ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE
1491
-ModelOptNvFp4Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
1491
+ModelOptNvFp4Config.KVCacheMethodCls = ModelOptKVCacheMethod
1492
1493
1494
class ModelOptMxFp8Config(ModelOptQuantConfigBase):
@@ -2018,7 +2018,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
2018
# Register the method classes for ModelOptMxFp8Config
2019
ModelOptMxFp8Config.LinearMethodCls = ModelOptMxFp8LinearMethod
2020
ModelOptMxFp8Config.FusedMoEMethodCls = ModelOptMxFp8FusedMoE
2021
-ModelOptMxFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
2021
+ModelOptMxFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod
2022
2023
2024
class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
@@ -2166,7 +2166,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2166
# KV-cache quantization
2167
if isinstance(layer, Attention):
2168
if self.kv_cache_quant_method:
2169
- return ModelOptFp8KVCacheMethod(self)
2169
+ return ModelOptKVCacheMethod(self)
2170
return None
2171
2172
# Excluded layers
vllm/v1/attention/backends/flashinfer.py
+114
-15
@@ -332,6 +332,7 @@ class FlashInferBackend(AttentionBackend):
332
"fp8",
333
"fp8_e4m3",
334
"fp8_e5m2",
335
+ "nvfp4",
336
]
337
338
@staticmethod
@@ -388,13 +389,15 @@ class FlashInferBackend(AttentionBackend):
389
return stride_order
390
391
@staticmethod
391
- def get_fp8_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:
392
+ def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:
393
if kv_cache_dtype in ("fp8", "fp8_e4m3"):
394
return torch.float8_e4m3fn
395
elif kv_cache_dtype == "fp8_e5m2":
396
return torch.float8_e5m2
397
+ elif kv_cache_dtype == "nvfp4":
398
+ return torch.uint8
399
else:
397
- raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
400
+ raise ValueError(f"Unrecognized dtype: {kv_cache_dtype}")
401
402
@classmethod
403
def get_supported_head_sizes(cls) -> list[int]:
@@ -622,9 +625,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
625
# For NVFP4, kv_cache_dtype stays as the string "nvfp4"
626
# which is passed to FlashInferImpl
627
self.kv_cache_dtype = self.cache_dtype
625
- raise NotImplementedError("nvfp4 KV cache is not yet supported")
628
else:
627
- self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
629
+ self.kv_cache_dtype = FlashInferBackend.get_dtype_for_flashinfer(
630
self.cache_dtype
631
)
632
else:
@@ -645,7 +647,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
647
):
648
if self.is_kvcache_nvfp4:
649
# NVFP4 KV cache uses FP8 quantized queries
648
- self.q_data_type = FlashInferBackend.get_fp8_dtype_for_flashinfer(
650
+ self.q_data_type = FlashInferBackend.get_dtype_for_flashinfer(
651
"fp8_e4m3"
652
)
653
else:
@@ -765,8 +767,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
767
dcp_a2a=self.dcp_a2a,
768
)
769
else:
770
+ # NVFP4 KV cache requires the trtllm-gen backend inside
771
+ # the wrapper; fa2/fa3 do not support nvfp4.
772
+ backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
773
self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
769
- self._get_workspace_buffer(), get_kv_cache_layout()
774
+ self._get_workspace_buffer(),
775
+ get_kv_cache_layout(),
776
+ backend=backend,
777
)
778
assert self._prefill_wrapper is not None
779
return self._prefill_wrapper
@@ -786,6 +793,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
793
paged_kv_indptr = None
794
paged_kv_indices = None
795
paged_kv_last_page_len = None
796
+ # NVFP4 KV cache requires the trtllm-gen backend inside
797
+ # the wrapper; fa2/fa3 do not support nvfp4.
798
+ backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
799
decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
800
self._get_workspace_buffer(),
801
get_kv_cache_layout(),
@@ -797,6 +807,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
807
# at least as good as cuda cores for all attention ops in latest
808
# gpus.
809
use_tensor_cores=True,
810
+ backend=backend,
811
)
812
813
# save the decode wrapper
@@ -1148,6 +1159,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1159
prefill_wrapper,
1160
BatchPrefillWithPagedKVCacheWrapper,
1161
)
1162
+ # NVFP4 trtllm kernel only supports FP8 output;
1163
+ # use FP8 o_data_type so the wrapper matches the
1164
+ # FP8 output buffer allocated in forward().
1165
+ o_dtype = (
1166
+ FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype
1167
+ )
1168
prefill_wrapper.plan(
1169
qo_indptr=qo_indptr_prefill_cpu,
1170
paged_kv_indptr=paged_kv_indptr_prefill_cpu,
@@ -1163,7 +1180,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1180
logits_soft_cap=self.logits_soft_cap,
1181
q_data_type=self.q_data_type,
1182
kv_data_type=self.kv_cache_dtype,
1166
- o_data_type=self.model_config.dtype,
1183
+ o_data_type=o_dtype,
1184
fixed_split_size=self.prefill_fixed_split_size,
1185
disable_split_kv=self.disable_split_kv,
1186
)
@@ -1197,6 +1214,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1214
# Use the persistent buffer with padding length,
1215
# instead of the same address but chunked version
1216
# in atten_metadata when using cudagraph.
1217
+ # NVFP4 trtllm kernel only supports FP8 output;
1218
+ # use FP8 o_data_type so the wrapper matches the
1219
+ # FP8 output buffer allocated in forward().
1220
+ o_dtype = (
1221
+ FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype
1222
+ )
1223
fast_plan_decode(
1224
decode_wrapper,
1225
indptr_cpu=self.paged_kv_indptr.cpu[: num_input_tokens + 1],
@@ -1215,7 +1238,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1238
logits_soft_cap=self.logits_soft_cap,
1239
q_data_type=self.q_data_type,
1240
kv_data_type=self.kv_cache_dtype,
1218
- o_data_type=self.model_config.dtype,
1241
+ o_data_type=o_dtype,
1242
fixed_split_size=self.decode_fixed_split_size,
1243
disable_split_kv=self.disable_split_kv,
1244
)
@@ -1300,6 +1323,17 @@ class FlashInferImpl(AttentionImpl):
1323
self.bmm2_scale: float | None = None
1324
self.o_sf_scale: float | None = None
1325
1326
+ # Pre-allocated FP8 output buffer for NVFP4 without fused output quant.
1327
+ if self.is_kvcache_nvfp4 and vllm_config is not None:
1328
+ max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
1329
+ self._nvfp4_fp8_out = torch.empty(
1330
+ (max_num_tokens, num_heads, head_size),
1331
+ dtype=FP8_DTYPE,
1332
+ device="cuda",
1333
+ )
1334
+ else:
1335
+ self._nvfp4_fp8_out = None
1336
+
1337
dcp_a2a = (
1338
vllm_config is not None
1339
and vllm_config.parallel_config.decode_context_parallel_size > 1
@@ -1420,7 +1454,7 @@ class FlashInferImpl(AttentionImpl):
1454
if self.kv_sharing_target_layer_name is None and is_quantized_kv_cache(
1455
self.kv_cache_dtype
1456
):
1423
- torch_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
1457
+ torch_dtype = FlashInferBackend.get_dtype_for_flashinfer(
1458
self.kv_cache_dtype
1459
)
1460
kv_cache = kv_cache.view(torch_dtype)
@@ -1500,13 +1534,37 @@ class FlashInferImpl(AttentionImpl):
1534
)
1535
assert prefill_wrapper._sm_scale == self.scale
1536
assert prefill_wrapper._causal
1537
+
1538
+ if self.is_kvcache_nvfp4:
1539
+ kv_cache_permute = nvfp4_kv_data
1540
+ kv_cache_sf = (
1541
+ nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
1542
+ )
1543
+
1544
+ # NVFP4 trtllm kernel only supports FP8 output.
1545
+ # Use a pre-allocated FP8 buffer and dequantize
1546
+ # afterwards.
1547
+ needs_fp8_out_prefill = (
1548
+ self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
1549
+ )
1550
+ if needs_fp8_out_prefill:
1551
+ out_prefill = self._nvfp4_fp8_out[:num_prefill_tokens]
1552
+ else:
1553
+ out_prefill = output[num_decode_tokens:]
1554
+
1555
prefill_wrapper.run(
1556
prefill_query,
1557
kv_cache_permute,
1558
k_scale=layer._k_scale_float,
1559
v_scale=layer._v_scale_float,
1508
- out=output[num_decode_tokens:],
1560
+ out=out_prefill,
1561
+ kv_cache_sf=kv_cache_sf,
1562
)
1563
+
1564
+ if needs_fp8_out_prefill:
1565
+ output[
1566
+ num_decode_tokens : num_decode_tokens + num_prefill_tokens
1567
+ ].copy_(out_prefill.to(output.dtype))
1568
else:
1569
assert isinstance(attn_metadata.prefill, TRTLLMPrefill)
1570
# prefill_query may be non-contiguous or have degenerate strides
@@ -1537,6 +1595,12 @@ class FlashInferImpl(AttentionImpl):
1595
assert self.o_sf_scale is None
1596
out = output[num_decode_tokens:]
1597
1598
+ # NVFP4 trtllm kernel only supports FP8 output.
1599
+ # Use a pre-allocated FP8 buffer and dequantize afterwards.
1600
+ needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
1601
+ if needs_fp8_out:
1602
+ out = self._nvfp4_fp8_out[:num_prefill_tokens]
1603
+
1604
prefill_kv_block_scales = None
1605
if self.is_kvcache_nvfp4:
1606
# NVFP4 trtllm-gen kernel requires FP8 query.
@@ -1547,7 +1611,7 @@ class FlashInferImpl(AttentionImpl):
1611
)
1612
mock_kv_cache = nvfp4_kv_data
1613
mock_block_table = block_tables_prefill
1550
- prefill_kv_block_scales = nvfp4_kv_block_scales # noqa: F841
1614
+ prefill_kv_block_scales = nvfp4_kv_block_scales
1615
elif (
1616
attn_metadata.q_data_type != FP8_DTYPE
1617
and self.kv_cache_dtype.startswith("fp8")
@@ -1598,8 +1662,14 @@ class FlashInferImpl(AttentionImpl):
1662
sinks=self.sinks,
1663
o_sf_scale=self.o_sf_scale,
1664
out=out,
1665
+ kv_cache_sf=prefill_kv_block_scales,
1666
)
1667
1668
+ if needs_fp8_out:
1669
+ output[
1670
+ num_decode_tokens : num_decode_tokens + num_prefill_tokens
1671
+ ].copy_(out[:num_prefill_tokens].to(output.dtype))
1672
+
1673
if num_decode_tokens > 0:
1674
decode_query = query[:num_decode_tokens]
1675
assert decode_query.shape[0] == num_decode_tokens
@@ -1612,6 +1682,18 @@ class FlashInferImpl(AttentionImpl):
1682
assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap or 0.0)
1683
assert decode_wrapper._sm_scale == self.scale
1684
1685
+ if self.is_kvcache_nvfp4:
1686
+ kv_cache_permute = nvfp4_kv_data
1687
+ kv_cache_sf = nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
1688
+
1689
+ # NVFP4 kernel only supports FP8 output.
1690
+ # Use a pre-allocated FP8 buffer and dequantize afterwards.
1691
+ needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
1692
+ if needs_fp8_out:
1693
+ out_decode = self._nvfp4_fp8_out[:num_decode_tokens]
1694
+ else:
1695
+ out_decode = output[:num_decode_tokens]
1696
+
1697
if use_dcp:
1698
decode_query = get_dcp_group().all_gather(
1699
decode_query.contiguous(), dim=-2
@@ -1630,6 +1712,7 @@ class FlashInferImpl(AttentionImpl):
1712
out=output_tmp,
1713
lse=lse,
1714
return_lse=True,
1715
+ kv_cache_sf=kv_cache_sf,
1716
)
1717
output[:num_decode_tokens] = self.dcp_combine(
1718
output_tmp,
@@ -1642,8 +1725,12 @@ class FlashInferImpl(AttentionImpl):
1725
kv_cache_permute,
1726
k_scale=layer._k_scale_float,
1727
v_scale=layer._v_scale_float,
1645
- out=output[:num_decode_tokens],
1728
+ out=out_decode,
1729
+ kv_cache_sf=kv_cache_sf,
1730
)
1731
+
1732
+ if needs_fp8_out:
1733
+ output[:num_decode_tokens].copy_(out_decode.to(output.dtype))
1734
else:
1735
# decode_query may be non-contiguous or have degenerate strides
1736
assert isinstance(attn_metadata.decode, TRTLLMDecode)
@@ -1686,6 +1773,12 @@ class FlashInferImpl(AttentionImpl):
1773
assert self.o_sf_scale is None
1774
out = output[:num_decode_tokens]
1775
1776
+ # NVFP4 trtllm kernel only supports FP8 output.
1777
+ # Use a pre-allocated FP8 buffer and dequantize afterwards.
1778
+ needs_fp8_out = self.is_kvcache_nvfp4 and output.dtype != FP8_DTYPE
1779
+ if needs_fp8_out:
1780
+ out = self._nvfp4_fp8_out[:num_decode_tokens]
1781
+
1782
if num_decode_tokens % attn_metadata.num_decodes != 0:
1783
# This gets triggered when the dummy_run forces
1784
# attention to be initialized with q_len = 0
@@ -1695,9 +1788,9 @@ class FlashInferImpl(AttentionImpl):
1788
1789
trtllm_batch_decode_with_kv_cache(
1790
query=decode_query,
1698
- kv_cache=nvfp4_kv_data
1699
- if self.is_kvcache_nvfp4
1700
- else kv_cache_permute,
1791
+ kv_cache=(
1792
+ nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_permute
1793
+ ),
1794
workspace_buffer=workspace_buffer,
1795
block_tables=block_tables_decode,
1796
seq_lens=seq_lens_decode,
@@ -1709,7 +1802,13 @@ class FlashInferImpl(AttentionImpl):
1802
o_sf_scale=self.o_sf_scale,
1803
out=out,
1804
q_len_per_req=q_len_per_req,
1805
+ kv_cache_sf=(
1806
+ nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None
1807
+ ),
1808
)
1809
+
1810
+ if needs_fp8_out:
1811
+ output[:num_decode_tokens].copy_(out.to(output.dtype))
1812
return output_padded
1813
1814
def do_kv_cache_update(
vllm/v1/kv_cache_interface.py
+10
@@ -151,6 +151,16 @@ class AttentionSpec(KVCacheSpec):
151
152
@property
153
def real_page_size_bytes(self) -> int:
154
+ if self.kv_quant_mode.is_nvfp4:
155
+ # Packed layout: fp4 data + fp8 block scales per head.
156
+ full_dim = nvfp4_kv_cache_full_dim(self.head_size)
157
+ return (
158
+ 2
159
+ * self.block_size
160
+ * self.num_kv_heads
161
+ * full_dim
162
+ * get_dtype_size(self.dtype)
163
+ )
164
return (
165
2
166
* self.block_size