[Perf][Attention] Bound FlashInfer prefill dequantization scratch (#57918)
Signed-off-by: Yueh-Ting Chen <yueh.ting.chen@gmail.com>
Yueh-Ting (eop) Chen committed
Sep 24, 2026 at 22:42 UTC
26f49e336a1498300a442a0d95911c11f5284500
2 files changed
+94
-2
tests/v1/attention/test_trtllm_attention_integration.py
+71
-1
@@ -9,6 +9,7 @@ import pytest
9
import torch
10
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
11
12
+from tests.kernels.attention.test_trtllm_kvfp8_dequant import make_random_kv_cache
13
from tests.v1.attention.test_attention_backends import create_and_prepopulate_kv_cache
14
from tests.v1.attention.utils import (
15
BatchSpec,
@@ -17,7 +18,11 @@ from tests.v1.attention.utils import (
18
)
19
from vllm.config import set_current_vllm_config
20
from vllm.platforms import current_platform
20
-from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim, set_random_seed
21
+from vllm.utils.torch_utils import (
22
+ is_strictly_contiguous,
23
+ nvfp4_kv_cache_full_dim,
24
+ set_random_seed,
25
+)
26
from vllm.v1.attention.backends.utils import PerLayerParameters
27
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode
28
@@ -33,6 +38,7 @@ from vllm.v1.attention.backends.flashinfer import ( # noqa: E402
38
FlashInferMetadataBuilder,
39
FlashInferTrtllmAPIDecode,
40
TRTLLMPrefill,
41
+ trtllm_prefill_attn_kvfp8_dequant,
42
)
43
44
@@ -83,6 +89,70 @@ def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
89
}
90
91
92
+@pytest.mark.parametrize(
93
+ "block_size,num_prefills,has_host_lengths,quantize_query",
94
+ [
95
+ (16, 1, True, False),
96
+ (64, 2, True, False),
97
+ (128, 2, True, False),
98
+ (16, 2, False, False),
99
+ (16, 2, True, True),
100
+ ],
101
+)
102
+@torch.inference_mode()
103
+def test_prefill_dequant_scratch_ignores_unused_capacity(
104
+ block_size, num_prefills, has_host_lengths, quantize_query, monkeypatch
105
+):
106
+ """Bound scratch by attended prefill KV, including cached prefixes."""
107
+ device = torch.device("cuda")
108
+ config = create_vllm_config(model_name=MODEL, max_model_len=32768)
109
+ config.cache_config.cache_dtype = "fp8"
110
+ config.cache_config.kv_cache_layout = "LBHNC"
111
+ config.attention_config.use_trtllm_attention = True
112
+ config.attention_config.disable_flashinfer_q_quantization = not quantize_query
113
+ monkeypatch.setattr(
114
+ "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
115
+ _mock_get_per_layer_parameters,
116
+ )
117
+ spec = FullAttentionSpec(
118
+ block_size=block_size,
119
+ num_kv_heads=2,
120
+ head_size=64,
121
+ dtype=torch.uint8,
122
+ kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
123
+ )
124
+ kv_cache, scale = make_random_kv_cache(32, 2, block_size, 64)
125
+ batch = BatchSpec(
126
+ seq_lens=[1024] + [17, 65][-num_prefills:],
127
+ query_lens=[1] + [2, 8][-num_prefills:],
128
+ )
129
+ common = create_common_attn_metadata(batch, block_size, device)
130
+ if not has_host_lengths:
131
+ common.seq_lens_cpu_upper_bound = None
132
+ live_pages = ((65 if has_host_lengths else 1024) + block_size - 1) // block_size
133
+ with set_current_vllm_config(config):
134
+ builder = FlashInferMetadataBuilder(spec, ["layer.0"], config, device)
135
+ for capacity in (64, 256):
136
+ # Unused columns can contain nonzero page IDs.
137
+ common.block_table_tensor = torch.randint(
138
+ 1, 32, (batch.batch_size, capacity), dtype=torch.int32, device=device
139
+ )
140
+ metadata = builder.build(0, common)
141
+ assert isinstance(metadata.prefill, TRTLLMPrefill)
142
+ table = metadata.prefill.block_tables
143
+ assert is_strictly_contiguous(table)
144
+ assert (metadata.q_data_type_prefill == current_platform.fp8_dtype()) == (
145
+ quantize_query
146
+ )
147
+ if not quantize_query:
148
+ converted, _ = trtllm_prefill_attn_kvfp8_dequant(
149
+ kv_cache, table, scale, scale, metadata.q_data_type_prefill
150
+ )
151
+ assert converted.shape[0] == num_prefills * live_pages + 1
152
+ width = capacity if quantize_query else live_pages
153
+ torch.testing.assert_close(table, common.block_table_tensor[1:, :width])
154
+
155
+
156
def _create_nvfp4_hnd_kv_cache(
157
k_contexts,
158
v_contexts,
vllm/v1/attention/backends/flashinfer.py
+23
-1
@@ -1563,8 +1563,30 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1563
qo_indptr_prefill_cpu[1:] - qo_indptr_prefill_cpu[:-1]
1564
)
1565
max_q_len_prefill = int(query_lens_prefill_cpu.max().item())
1566
+ prefill_block_tables = block_table_tensor[prefill_start:]
1567
+ if (
1568
+ self.q_data_type_prefill != FP8_DTYPE
1569
+ and self.cache_dtype.startswith("fp8")
1570
+ ):
1571
+ seq_lens_cpu_upper_bound = (
1572
+ common_attn_metadata.seq_lens_cpu_upper_bound
1573
+ )
1574
+ max_prefill_seq_len = max_seq_len
1575
+ if seq_lens_cpu_upper_bound is not None:
1576
+ max_prefill_seq_len = int(
1577
+ seq_lens_cpu_upper_bound[prefill_start:num_reqs]
1578
+ .max()
1579
+ .item()
1580
+ )
1581
+ # Dequantization allocates one page per table entry and
1582
+ # indexes rows by their logical width.
1583
+ prefill_block_tables = canonicalize_singleton_dim_strides(
1584
+ prefill_block_tables[
1585
+ :, : cdiv(max_prefill_seq_len, page_size)
1586
+ ].contiguous()
1587
+ )
1588
attn_metadata.prefill = TRTLLMPrefill(
1567
- block_tables=block_table_tensor[prefill_start:],
1589
+ block_tables=prefill_block_tables,
1590
seq_lens=prefill_seq_lens,
1591
cum_seq_lens_q=qo_indptr_prefill_gpu,
1592
cum_seq_lens_kv=paged_kv_indptr_prefill_gpu,