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 6, 14:30 UTC
Support non-contiguous KV cache in TRTLLM fp8 dequant kernel (#36867)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com> Signed-off-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Co-authored-by: Pavani Majety <pavanimajety@gmail.com>
Vadim Gimpelson committed
Mar 17, 2026 at 04:48 UTC
6c1cfbad325067c4afa12c87992f45a58ce0614b
2 files changed
+491
-26
tests/kernels/attention/test_trtllm_kvfp8_dequant.py
new
+434
@@ -0,0 +1,434 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
5
+
6
+Tests both contiguous and non-contiguous (cross-layer unified) KV cache
7
+layouts against a pure-PyTorch reference implementation.
8
+"""
9
+
10
+import pytest
11
+import torch
12
+
13
+from vllm.platforms import current_platform
14
+
15
+FP8_DTYPE = current_platform.fp8_dtype()
16
+
17
+NUM_BLOCKS = 128
18
+
19
+
20
+def to_float8(x, dtype=None):
21
+ if dtype is None:
22
+ dtype = FP8_DTYPE
23
+ finfo = torch.finfo(dtype)
24
+ min_val, max_val = x.aminmax()
25
+ amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
26
+ scale = finfo.max / amax * 0.1
27
+ x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
28
+ return x_scl_sat.to(dtype), scale.float().reciprocal()
29
+
30
+
31
+def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
32
+ """Create a standard contiguous fp8 KV cache (HND layout)."""
33
+ raw = torch.randn(
34
+ num_blocks,
35
+ 2,
36
+ num_kv_heads,
37
+ block_size,
38
+ head_size,
39
+ dtype=torch.bfloat16,
40
+ device="cuda",
41
+ )
42
+ kv_cache, scale = to_float8(raw)
43
+ return kv_cache, scale
44
+
45
+
46
+def make_cross_layer_kv_cache(
47
+ num_blocks,
48
+ num_kv_heads,
49
+ block_size,
50
+ head_size,
51
+ num_layers=4,
52
+):
53
+ """
54
+ Create a non-contiguous per-layer view mimicking cross-layer allocation.
55
+
56
+ Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
57
+ Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size)
58
+ with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
59
+ """
60
+ raw = torch.randn(
61
+ num_blocks,
62
+ 2,
63
+ num_kv_heads,
64
+ num_layers,
65
+ block_size,
66
+ head_size,
67
+ dtype=torch.bfloat16,
68
+ device="cuda",
69
+ )
70
+ fp8_full, scale = to_float8(raw)
71
+ layer_view = fp8_full[:, :, :, 0, :, :]
72
+ assert not layer_view.is_contiguous(), (
73
+ f"Expected non-contiguous view, got strides {layer_view.stride()}"
74
+ )
75
+ return layer_view, scale
76
+
77
+
78
+def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
79
+ """Pure PyTorch reference: gather pages and dequantize fp8 -> dequant_dtype."""
80
+ batch_size, num_pages_per_seq = block_tables.shape
81
+ s = kv_cache.shape
82
+ out = torch.zeros(
83
+ batch_size * num_pages_per_seq + 1,
84
+ s[1],
85
+ s[2],
86
+ s[3],
87
+ s[4],
88
+ dtype=dequant_dtype,
89
+ device=kv_cache.device,
90
+ )
91
+ for b in range(batch_size):
92
+ for p in range(num_pages_per_seq):
93
+ page_idx = block_tables[b, p].item()
94
+ if page_idx <= 0:
95
+ continue
96
+ mock_idx = b * num_pages_per_seq + p + 1
97
+ out[mock_idx, 0] = (kv_cache[page_idx, 0].float() * k_scale.item()).to(
98
+ dequant_dtype
99
+ )
100
+ out[mock_idx, 1] = (kv_cache[page_idx, 1].float() * v_scale.item()).to(
101
+ dequant_dtype
102
+ )
103
+ return out
104
+
105
+
106
+@pytest.mark.parametrize("num_kv_heads", [1, 8])
107
+@pytest.mark.parametrize("head_size", [64, 128])
108
+@pytest.mark.parametrize("block_size", [16, 32])
109
+@pytest.mark.parametrize("batch_size", [1, 4])
110
+@pytest.mark.parametrize("num_pages_per_seq", [3, 8])
111
+@pytest.mark.parametrize("contiguous", [True, False])
112
+@torch.inference_mode()
113
+def test_trtllm_kvfp8_dequant(
114
+ num_kv_heads: int,
115
+ head_size: int,
116
+ block_size: int,
117
+ batch_size: int,
118
+ num_pages_per_seq: int,
119
+ contiguous: bool,
120
+):
121
+ from vllm.v1.attention.backends.flashinfer import (
122
+ trtllm_prefill_attn_kvfp8_dequant,
123
+ )
124
+
125
+ torch.set_default_device("cuda")
126
+
127
+ if contiguous:
128
+ kv_cache, scale = make_contiguous_kv_cache(
129
+ NUM_BLOCKS,
130
+ num_kv_heads,
131
+ block_size,
132
+ head_size,
133
+ )
134
+ else:
135
+ kv_cache, scale = make_cross_layer_kv_cache(
136
+ NUM_BLOCKS,
137
+ num_kv_heads,
138
+ block_size,
139
+ head_size,
140
+ )
141
+
142
+ k_scale = scale.clone()
143
+ v_scale = scale.clone()
144
+
145
+ block_tables = torch.randint(
146
+ 1,
147
+ NUM_BLOCKS,
148
+ (batch_size, num_pages_per_seq),
149
+ dtype=torch.int32,
150
+ )
151
+
152
+ mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
153
+ kv_cache,
154
+ block_tables,
155
+ k_scale,
156
+ v_scale,
157
+ torch.bfloat16,
158
+ )
159
+
160
+ ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
161
+
162
+ expected_bt = torch.arange(
163
+ 1,
164
+ batch_size * num_pages_per_seq + 1,
165
+ dtype=torch.int32,
166
+ device="cuda",
167
+ ).reshape(batch_size, num_pages_per_seq)
168
+ torch.testing.assert_close(mock_block_table, expected_bt)
169
+
170
+ # Page 0 is padding (never written), compare only pages 1+
171
+ torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
172
+
173
+
174
+@torch.inference_mode()
175
+def test_block_tables_with_zero_pages():
176
+ """Pages with index <= 0 must be skipped (early return in kernel)."""
177
+ from vllm.v1.attention.backends.flashinfer import (
178
+ trtllm_prefill_attn_kvfp8_dequant,
179
+ )
180
+
181
+ torch.set_default_device("cuda")
182
+ num_kv_heads, block_size, head_size = 8, 16, 64
183
+
184
+ kv_cache, scale = make_contiguous_kv_cache(
185
+ NUM_BLOCKS,
186
+ num_kv_heads,
187
+ block_size,
188
+ head_size,
189
+ )
190
+ k_scale = v_scale = scale.clone()
191
+
192
+ # Mix of valid pages and zeros (padding)
193
+ block_tables = torch.tensor(
194
+ [[5, 0, 10], [0, 0, 0], [3, 7, 0]],
195
+ dtype=torch.int32,
196
+ device="cuda",
197
+ )
198
+
199
+ mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
200
+ kv_cache,
201
+ block_tables,
202
+ k_scale,
203
+ v_scale,
204
+ torch.bfloat16,
205
+ )
206
+ ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
207
+
208
+ # Only compare pages that were actually written (non-zero page indices)
209
+ for b in range(block_tables.shape[0]):
210
+ for p in range(block_tables.shape[1]):
211
+ if block_tables[b, p].item() > 0:
212
+ idx = b * block_tables.shape[1] + p + 1
213
+ torch.testing.assert_close(
214
+ mock_kv_cache[idx],
215
+ ref[idx],
216
+ atol=1e-3,
217
+ rtol=1e-3,
218
+ )
219
+
220
+
221
+@torch.inference_mode()
222
+def test_all_zero_block_tables():
223
+ """All-zero block_tables: kernel should write nothing."""
224
+ from vllm.v1.attention.backends.flashinfer import (
225
+ trtllm_prefill_attn_kvfp8_dequant,
226
+ )
227
+
228
+ torch.set_default_device("cuda")
229
+ num_kv_heads, block_size, head_size = 4, 16, 64
230
+
231
+ kv_cache, scale = make_contiguous_kv_cache(
232
+ NUM_BLOCKS,
233
+ num_kv_heads,
234
+ block_size,
235
+ head_size,
236
+ )
237
+ k_scale = v_scale = scale.clone()
238
+
239
+ block_tables = torch.zeros(2, 4, dtype=torch.int32, device="cuda")
240
+
241
+ # Should not crash even though no pages are valid
242
+ mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
243
+ kv_cache,
244
+ block_tables,
245
+ k_scale,
246
+ v_scale,
247
+ torch.bfloat16,
248
+ )
249
+ assert mock_kv_cache.shape[0] == 2 * 4 + 1
250
+ assert mock_block_table.shape == (2, 4)
251
+
252
+
253
+@torch.inference_mode()
254
+def test_different_k_v_scales():
255
+ """Verify K and V are dequantized with independent scales."""
256
+ from vllm.v1.attention.backends.flashinfer import (
257
+ trtllm_prefill_attn_kvfp8_dequant,
258
+ )
259
+
260
+ torch.set_default_device("cuda")
261
+ num_kv_heads, block_size, head_size = 8, 16, 64
262
+
263
+ kv_cache, _ = make_contiguous_kv_cache(
264
+ NUM_BLOCKS,
265
+ num_kv_heads,
266
+ block_size,
267
+ head_size,
268
+ )
269
+ k_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
270
+ v_scale = torch.tensor([2.0], dtype=torch.float32, device="cuda")
271
+
272
+ block_tables = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
273
+
274
+ mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
275
+ kv_cache,
276
+ block_tables,
277
+ k_scale,
278
+ v_scale,
279
+ torch.bfloat16,
280
+ )
281
+ ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
282
+
283
+ torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
284
+
285
+
286
+@torch.inference_mode()
287
+def test_single_page_per_seq():
288
+ """Minimum grid dim 1 = 1 page per sequence."""
289
+ from vllm.v1.attention.backends.flashinfer import (
290
+ trtllm_prefill_attn_kvfp8_dequant,
291
+ )
292
+
293
+ torch.set_default_device("cuda")
294
+ num_kv_heads, block_size, head_size = 8, 16, 128
295
+
296
+ kv_cache, scale = make_contiguous_kv_cache(
297
+ NUM_BLOCKS,
298
+ num_kv_heads,
299
+ block_size,
300
+ head_size,
301
+ )
302
+ k_scale = v_scale = scale.clone()
303
+
304
+ block_tables = torch.tensor([[5], [10], [20]], dtype=torch.int32, device="cuda")
305
+
306
+ mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
307
+ kv_cache,
308
+ block_tables,
309
+ k_scale,
310
+ v_scale,
311
+ torch.bfloat16,
312
+ )
313
+ ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
314
+
315
+ torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
316
+
317
+
318
+@torch.inference_mode()
319
+def test_large_page_indices():
320
+ """Page indices near the top of the buffer stress offset arithmetic."""
321
+ from vllm.v1.attention.backends.flashinfer import (
322
+ trtllm_prefill_attn_kvfp8_dequant,
323
+ )
324
+
325
+ torch.set_default_device("cuda")
326
+ num_kv_heads, block_size, head_size = 8, 16, 128
327
+ large_num_blocks = 32768
328
+
329
+ kv_cache, scale = make_contiguous_kv_cache(
330
+ large_num_blocks,
331
+ num_kv_heads,
332
+ block_size,
333
+ head_size,
334
+ )
335
+ k_scale = v_scale = scale.clone()
336
+
337
+ # Use page indices near the top of the buffer
338
+ block_tables = torch.tensor(
339
+ [[large_num_blocks - 1, large_num_blocks - 2, 1]],
340
+ dtype=torch.int32,
341
+ device="cuda",
342
+ )
343
+
344
+ mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
345
+ kv_cache,
346
+ block_tables,
347
+ k_scale,
348
+ v_scale,
349
+ torch.bfloat16,
350
+ )
351
+ ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
352
+
353
+ torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
354
+
355
+
356
+@torch.inference_mode()
357
+def test_large_block_size():
358
+ """block_size=64 -> HEAD_STRIDE=8192, large tl.arange per thread block."""
359
+ from vllm.v1.attention.backends.flashinfer import (
360
+ trtllm_prefill_attn_kvfp8_dequant,
361
+ )
362
+
363
+ torch.set_default_device("cuda")
364
+ num_kv_heads, block_size, head_size = 4, 64, 128
365
+
366
+ kv_cache, scale = make_contiguous_kv_cache(
367
+ NUM_BLOCKS,
368
+ num_kv_heads,
369
+ block_size,
370
+ head_size,
371
+ )
372
+ k_scale = v_scale = scale.clone()
373
+
374
+ block_tables = torch.randint(
375
+ 1,
376
+ NUM_BLOCKS,
377
+ (2, 4),
378
+ dtype=torch.int32,
379
+ device="cuda",
380
+ )
381
+
382
+ mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
383
+ kv_cache,
384
+ block_tables,
385
+ k_scale,
386
+ v_scale,
387
+ torch.bfloat16,
388
+ )
389
+ ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
390
+
391
+ torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
392
+
393
+
394
+@torch.inference_mode()
395
+def test_cross_layer_many_layers():
396
+ """
397
+ Non-contiguous with 36 layers -- matches real gpt-oss-120b.
398
+ Strides are far from contiguous (factor of 36 in the gaps).
399
+ """
400
+ from vllm.v1.attention.backends.flashinfer import (
401
+ trtllm_prefill_attn_kvfp8_dequant,
402
+ )
403
+
404
+ torch.set_default_device("cuda")
405
+ num_kv_heads, block_size, head_size = 8, 16, 64
406
+ num_layers = 36
407
+
408
+ kv_cache, scale = make_cross_layer_kv_cache(
409
+ NUM_BLOCKS,
410
+ num_kv_heads,
411
+ block_size,
412
+ head_size,
413
+ num_layers=num_layers,
414
+ )
415
+ k_scale = v_scale = scale.clone()
416
+
417
+ block_tables = torch.randint(
418
+ 1,
419
+ NUM_BLOCKS,
420
+ (4, 6),
421
+ dtype=torch.int32,
422
+ device="cuda",
423
+ )
424
+
425
+ mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
426
+ kv_cache,
427
+ block_tables,
428
+ k_scale,
429
+ v_scale,
430
+ torch.bfloat16,
431
+ )
432
+ ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
433
+
434
+ torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
vllm/v1/attention/backends/flashinfer.py
+57
-26
@@ -96,8 +96,13 @@ def _trtllm_prefill_attn_kvfp8_dequant(
96
mock_kv_cache_ptr,
97
k_scale_ptr,
98
v_scale_ptr,
99
- K_CACHE_STRIDE: tl.constexpr,
100
- KV_CACHE_STRIDE: tl.constexpr,
99
+ src_stride_page,
100
+ src_stride_kv,
101
+ src_stride_head,
102
+ DST_K_CACHE_STRIDE: tl.constexpr,
103
+ DST_KV_CACHE_STRIDE: tl.constexpr,
104
+ HEAD_STRIDE: tl.constexpr,
105
+ NUM_KV_HEADS: tl.constexpr,
106
):
107
batch_idx = tl.program_id(0).to(tl.int64)
108
mock_block_table_idx = tl.program_id(1).to(tl.int64)
@@ -108,31 +113,42 @@ def _trtllm_prefill_attn_kvfp8_dequant(
113
return
114
dequant_dtype = mock_kv_cache_ptr.dtype.element_ty
115
111
- # Dequantize K
116
k_scale_val = tl.load(k_scale_ptr)
113
- offset = orig_page_num * KV_CACHE_STRIDE + tl.arange(0, K_CACHE_STRIDE)
114
- fp8_vals = tl.load(kv_cache_ptr + offset)
115
- dequantized_vals = fp8_vals.to(tl.float32) * k_scale_val
116
- mock_cache_offset = (
117
- batch_idx * block_table_stride + mock_block_table_idx + 1
118
- ) * KV_CACHE_STRIDE + tl.arange(0, K_CACHE_STRIDE)
119
- dequantized_vals = dequantized_vals.to(dequant_dtype)
120
- tl.store(mock_kv_cache_ptr + mock_cache_offset, dequantized_vals)
121
-
122
- # Dequantize V
117
v_scale_val = tl.load(v_scale_ptr)
124
- offset = (
125
- orig_page_num * KV_CACHE_STRIDE + K_CACHE_STRIDE + tl.arange(0, K_CACHE_STRIDE)
126
- )
127
- fp8_vals = tl.load(kv_cache_ptr + offset)
128
- dequantized_vals = fp8_vals.to(tl.float32) * v_scale_val
129
- mock_cache_offset = (
130
- (batch_idx * block_table_stride + mock_block_table_idx + 1) * KV_CACHE_STRIDE
131
- + K_CACHE_STRIDE
132
- + tl.arange(0, K_CACHE_STRIDE)
133
- )
134
- dequantized_vals = dequantized_vals.to(dequant_dtype)
135
- tl.store(mock_kv_cache_ptr + mock_cache_offset, dequantized_vals)
118
+
119
+ mock_page_idx = batch_idx * block_table_stride + mock_block_table_idx + 1
120
+ head_offsets = tl.arange(0, HEAD_STRIDE)
121
+
122
+ for h in range(NUM_KV_HEADS):
123
+ h_off = tl.cast(h, tl.int64)
124
+
125
+ # Read K from source (supports non-contiguous page/kv/head strides)
126
+ src_k = orig_page_num * src_stride_page + h_off * src_stride_head + head_offsets
127
+ fp8_k = tl.load(kv_cache_ptr + src_k)
128
+ dequant_k = (fp8_k.to(tl.float32) * k_scale_val).to(dequant_dtype)
129
+
130
+ # Write K to contiguous mock cache
131
+ dst_k = mock_page_idx * DST_KV_CACHE_STRIDE + h * HEAD_STRIDE + head_offsets
132
+ tl.store(mock_kv_cache_ptr + dst_k, dequant_k)
133
+
134
+ # Read V from source (offset by src_stride_kv for the V half)
135
+ src_v = (
136
+ orig_page_num * src_stride_page
137
+ + src_stride_kv
138
+ + h_off * src_stride_head
139
+ + head_offsets
140
+ )
141
+ fp8_v = tl.load(kv_cache_ptr + src_v)
142
+ dequant_v = (fp8_v.to(tl.float32) * v_scale_val).to(dequant_dtype)
143
+
144
+ # Write V to contiguous mock cache
145
+ dst_v = (
146
+ mock_page_idx * DST_KV_CACHE_STRIDE
147
+ + DST_K_CACHE_STRIDE
148
+ + h * HEAD_STRIDE
149
+ + head_offsets
150
+ )
151
+ tl.store(mock_kv_cache_ptr + dst_v, dequant_v)
152
153
154
def trtllm_prefill_attn_kvfp8_dequant(
@@ -146,8 +162,18 @@ def trtllm_prefill_attn_kvfp8_dequant(
162
s = kv_cache.shape
163
assert s[1] == 2
164
assert dequant_dtype in (torch.bfloat16, torch.float16)
149
- k_cache_stride = s[2] * s[3] * s[4]
165
+
166
+ num_kv_heads, block_size, head_size = s[2], s[3], s[4]
167
+ head_stride = block_size * head_size
168
+ k_cache_stride = num_kv_heads * head_stride
169
kv_cache_stride = k_cache_stride * s[1]
170
+
171
+ strides = kv_cache.stride()
172
+ assert strides[3] == head_size and strides[4] == 1, (
173
+ "For kv cache layouts, (block_size, head_size) "
174
+ f"dimensions must be contiguous, got strides {strides}"
175
+ )
176
+
177
new_s = (batch_size * num_of_page_per_token + 1, s[1], s[2], s[3], s[4])
178
# mock kv cache contains just the pages needed by this prefill
179
mock_kv_cache = torch.empty(new_s, dtype=dequant_dtype, device=kv_cache.device)
@@ -166,8 +192,13 @@ def trtllm_prefill_attn_kvfp8_dequant(
192
mock_kv_cache,
193
k_scale,
194
v_scale,
195
+ strides[0],
196
+ strides[1],
197
+ strides[2],
198
k_cache_stride,
199
kv_cache_stride,
200
+ head_stride,
201
+ num_kv_heads,
202
)
203
return mock_kv_cache, mock_block_table
204