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 7, 09:00 UTC
TRTLLM gen-full attn Test Coverage (#34986)
Signed-off-by: Anshika Ojha <anshikao@nvidia.com> Co-authored-by: Anshika Ojha <anshikao@gb-nvl-059-compute09.nvidia.com>
ojhaanshika committed
Mar 3, 2026 at 08:35 UTC
e05cb3b93e5db3afd510189651a128018c31c251
2 files changed
+556
tests/kernels/attention/test_use_trtllm_attention.py
new
+196
@@ -0,0 +1,196 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from unittest.mock import patch
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.utils.flashinfer import (
10
+ can_use_trtllm_attention,
11
+ supports_trtllm_attention,
12
+ use_trtllm_attention,
13
+)
14
+
15
+MODEL_CONFIGS = {
16
+ "Llama-3-70B": dict(num_qo_heads=64, num_kv_heads=8),
17
+ "Llama-3-8B": dict(num_qo_heads=32, num_kv_heads=8),
18
+ "Qwen2.5-0.5B": dict(num_qo_heads=14, num_kv_heads=2),
19
+ "Mistral-7B": dict(num_qo_heads=32, num_kv_heads=8),
20
+ "Gemma-2-9B": dict(num_qo_heads=8, num_kv_heads=4),
21
+ "Falcon-40B": dict(num_qo_heads=128, num_kv_heads=8),
22
+}
23
+
24
+
25
+def get_config(model: str) -> dict:
26
+ """Return the attention config for a model."""
27
+ return MODEL_CONFIGS[model]
28
+
29
+
30
+DEFAULT_KWARGS = dict(
31
+ **get_config("Llama-3-70B"),
32
+ num_tokens=128,
33
+ max_seq_len=4096,
34
+ dcp_world_size=1,
35
+ kv_cache_dtype="auto",
36
+ q_dtype=torch.bfloat16,
37
+ is_prefill=False,
38
+ force_use_trtllm=None,
39
+ has_sinks=False,
40
+ has_spec=False,
41
+)
42
+
43
+
44
+def _call(**overrides) -> bool:
45
+ kwargs = {**DEFAULT_KWARGS, **overrides}
46
+ return use_trtllm_attention(**kwargs)
47
+
48
+
49
+@pytest.fixture(autouse=True)
50
+def _clear_supports_cache():
51
+ """Clear functools.cache to ensure each test runs independently."""
52
+ supports_trtllm_attention.cache_clear()
53
+
54
+
55
+# supports_trtllm_attention
56
+
57
+
58
+@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=True)
59
+def test_supports_batch_invariant_disables(_mock):
60
+ assert supports_trtllm_attention() is False
61
+
62
+
63
+@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
64
+@patch(
65
+ "vllm.utils.flashinfer.current_platform.is_device_capability_family",
66
+ return_value=True,
67
+)
68
+@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
69
+def test_supports_sm100_with_artifactory(_art, _cap, _bi):
70
+ assert supports_trtllm_attention() is True
71
+
72
+
73
+@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
74
+@patch(
75
+ "vllm.utils.flashinfer.current_platform.is_device_capability_family",
76
+ return_value=False,
77
+)
78
+def test_supports_non_sm100_platform(_cap, _bi):
79
+ assert supports_trtllm_attention() is False
80
+
81
+
82
+@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
83
+@patch(
84
+ "vllm.utils.flashinfer.current_platform.is_device_capability_family",
85
+ return_value=True,
86
+)
87
+@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)
88
+def test_supports_sm100_without_artifactory(_art, _cap, _bi):
89
+ assert supports_trtllm_attention() is False
90
+
91
+
92
+# can_use_trtllm_attention
93
+
94
+
95
+@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=False)
96
+def test_can_use_force_disabled(_mock):
97
+ cfg = get_config("Llama-3-70B")
98
+ assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
99
+
100
+
101
+@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
102
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
103
+def test_can_use_compatible_heads(_sup, _force):
104
+ cfg = get_config("Llama-3-70B")
105
+ assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is True
106
+
107
+
108
+@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
109
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
110
+def test_can_use_incompatible_heads(_sup, _force):
111
+ assert can_use_trtllm_attention(40, 6) is False
112
+
113
+
114
+@pytest.mark.parametrize("model", list(MODEL_CONFIGS.keys()))
115
+@patch("vllm.utils.flashinfer.force_use_trtllm_attention", return_value=None)
116
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
117
+def test_can_use_platform_unsupported(_sup, _force, model):
118
+ cfg = get_config(model)
119
+ assert can_use_trtllm_attention(cfg["num_qo_heads"], cfg["num_kv_heads"]) is False
120
+
121
+
122
+# use_trtllm_attention
123
+
124
+
125
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
126
+def test_use_force_off(_mock):
127
+ assert _call(force_use_trtllm=False) is False
128
+
129
+
130
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
131
+def test_use_dcp_fallback(_mock):
132
+ assert _call(dcp_world_size=2) is False
133
+
134
+
135
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
136
+def test_use_platform_unsupported(_mock):
137
+ assert _call() is False
138
+
139
+
140
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=False)
141
+def test_use_platform_unsupported_force_on_still_false(_mock):
142
+ assert _call(force_use_trtllm=True) is False
143
+
144
+
145
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
146
+def test_use_incompatible_heads(_mock):
147
+ assert _call(num_qo_heads=40, num_kv_heads=6) is False
148
+
149
+
150
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
151
+def test_use_incompatible_heads_force_on_still_false(_mock):
152
+ assert _call(num_qo_heads=40, num_kv_heads=6, force_use_trtllm=True) is False
153
+
154
+
155
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
156
+def test_use_spec_decode_enables(_mock):
157
+ assert _call(has_spec=True, is_prefill=False) is True
158
+
159
+
160
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
161
+@patch(
162
+ "vllm.utils.flashinfer.current_platform.fp8_dtype",
163
+ return_value=torch.float8_e4m3fn,
164
+)
165
+def test_use_fp8_query_forces_trtllm(_fp8, _sup):
166
+ assert _call(q_dtype=torch.float8_e4m3fn) is True
167
+
168
+
169
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
170
+def test_use_sinks_force_trtllm(_mock):
171
+ assert _call(has_sinks=True) is True
172
+
173
+
174
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
175
+def test_use_auto_prefill_kv_auto(_mock):
176
+ assert _call(is_prefill=True, kv_cache_dtype="auto") is True
177
+
178
+
179
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
180
+def test_use_auto_prefill_kv_fp8(_mock):
181
+ assert _call(is_prefill=True, kv_cache_dtype="fp8") is False
182
+
183
+
184
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
185
+def test_use_auto_decode_small_batch(_mock):
186
+ assert _call(is_prefill=False, num_tokens=128, kv_cache_dtype="auto") is True
187
+
188
+
189
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
190
+def test_use_auto_decode_large_batch(_mock):
191
+ assert _call(is_prefill=False, num_tokens=512, kv_cache_dtype="auto") is False
192
+
193
+
194
+@patch("vllm.utils.flashinfer.supports_trtllm_attention", return_value=True)
195
+def test_use_force_on(_mock):
196
+ assert _call(force_use_trtllm=True) is True
tests/v1/attention/test_trtllm_attention_integration.py
new
+360
@@ -0,0 +1,360 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Integration tests for TRTLLM gen-full attention through FlashInfer."""
4
+
5
+import unittest.mock
6
+from functools import partial
7
+
8
+import pytest
9
+import torch
10
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
11
+
12
+from tests.v1.attention.utils import (
13
+ BatchSpec,
14
+ create_common_attn_metadata,
15
+ create_vllm_config,
16
+)
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
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
27
+
28
+if not current_platform.is_device_capability_family(100):
29
+ pytest.skip(
30
+ "TRTLLM integration tests require NVIDIA Blackwell (SM100).",
31
+ allow_module_level=True,
32
+ )
33
+
34
+from vllm.v1.attention.backends.flashinfer import ( # noqa: E402
35
+ FlashInferImpl,
36
+ FlashInferMetadataBuilder,
37
+ TRTLLMDecode,
38
+ TRTLLMPrefill,
39
+)
40
+
41
+
42
+class MockAttentionLayer:
43
+ """Minimal mock of an attention layer for testing."""
44
+
45
+ def __init__(self, device: torch.device):
46
+ self._q_scale = torch.tensor(1.0, device=device)
47
+ self._k_scale = torch.tensor(1.0, device=device)
48
+ self._v_scale = torch.tensor(1.0, device=device)
49
+ self._q_scale_float = 1.0
50
+ self._k_scale_float = 1.0
51
+ self._v_scale_float = 1.0
52
+ self._o_scale_float = None
53
+
54
+
55
+MODEL = "Qwen/Qwen2.5-0.5B"
56
+BLOCK_SIZE = 16
57
+NUM_GPU_BLOCKS = 8192
58
+
59
+BATCH_SPECS = {
60
+ "decode_only": BatchSpec(
61
+ seq_lens=[128, 256, 512],
62
+ query_lens=[1, 1, 1],
63
+ ),
64
+ "prefill_only": BatchSpec(
65
+ seq_lens=[64, 128, 256],
66
+ query_lens=[16, 32, 16],
67
+ ),
68
+ "mixed": BatchSpec(
69
+ seq_lens=[128, 256, 512, 128],
70
+ query_lens=[1, 1, 8, 16],
71
+ ),
72
+}
73
+
74
+
75
+def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
76
+ head_size = vllm_config.model_config.get_head_size()
77
+ return {
78
+ name: PerLayerParameters(
79
+ window_left=-1,
80
+ logits_soft_cap=0.0,
81
+ sm_scale=1.0 / (head_size**0.5),
82
+ )
83
+ for name in layer_names
84
+ }
85
+
86
+
87
+def _create_hnd_kv_cache(
88
+ k_contexts,
89
+ v_contexts,
90
+ block_size,
91
+ num_kv_heads,
92
+ head_size,
93
+ dtype,
94
+ device,
95
+ num_blocks,
96
+ common_attn_metadata,
97
+):
98
+ """Create and populate a KV cache with HND-compatible strides.
99
+
100
+ The returned tensor has logical shape
101
+ (num_blocks, 2, block_size, num_kv_heads, head_size) but is physically
102
+ laid out as (num_blocks, 2, num_kv_heads, block_size, head_size) so that
103
+ ``kv_cache.permute(0, 1, 3, 2, 4)`` yields a contiguous HND view.
104
+ """
105
+ seq_lens = common_attn_metadata.seq_lens.cpu()
106
+ query_lens = (
107
+ common_attn_metadata.query_start_loc_cpu[1:]
108
+ - common_attn_metadata.query_start_loc_cpu[:-1]
109
+ )
110
+ block_table = common_attn_metadata.block_table_tensor
111
+ slot_mapping = common_attn_metadata.slot_mapping
112
+ batch_size = len(k_contexts)
113
+
114
+ # Build cache in (2, num_blocks, block_size, num_kv_heads, head_size)
115
+ # then convert to HND format (same approach as test_attention_backends.py).
116
+ kv_cache_raw = torch.zeros(
117
+ 2,
118
+ num_blocks,
119
+ block_size,
120
+ num_kv_heads,
121
+ head_size,
122
+ dtype=dtype,
123
+ device=device,
124
+ )
125
+ kv_cache_flat = kv_cache_raw.view(2, -1, num_kv_heads, head_size)
126
+
127
+ start_block_idx = 1
128
+ for i in range(batch_size):
129
+ k_ctx, v_ctx = k_contexts[i], v_contexts[i]
130
+ start = start_block_idx * block_size
131
+ end = start + k_ctx.shape[0]
132
+ kv_cache_flat[0, start:end] = k_ctx
133
+ kv_cache_flat[1, start:end] = v_ctx
134
+ start_block_idx += cdiv(int(seq_lens[i]), block_size)
135
+
136
+ blocks_end = start_block_idx
137
+
138
+ # Randomly permute blocks (starting from block 1; block 0 is null).
139
+ perm = torch.randperm(blocks_end - 1) + 1
140
+ inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
141
+ inv_perm[1:] = torch.argsort(perm) + 1
142
+ kv_cache_raw[:, 1:blocks_end] = kv_cache_raw[:, perm]
143
+
144
+ # Build block table.
145
+ start_block_idx = 1
146
+ for i in range(batch_size):
147
+ n_blocks = cdiv(int(seq_lens[i]), block_size)
148
+ block_table[i, :n_blocks] = inv_perm[
149
+ start_block_idx : start_block_idx + n_blocks
150
+ ]
151
+ start_block_idx += n_blocks
152
+
153
+ # Build slot mapping that is consistent with the block table.
154
+ for i in range(batch_size):
155
+ ctx_len = int(seq_lens[i]) - int(query_lens[i])
156
+ token_offsets = torch.arange(int(query_lens[i])) + ctx_len
157
+ block_indices = token_offsets // block_size
158
+ intra_block_offsets = token_offsets % block_size
159
+ start = common_attn_metadata.query_start_loc_cpu[i]
160
+ end = common_attn_metadata.query_start_loc_cpu[i + 1]
161
+ slot_mapping[start:end] = block_table[
162
+ i, block_indices
163
+ ] * block_size + intra_block_offsets.to(device)
164
+
165
+ # Transpose to FlashInfer logical shape then make HND-strided.
166
+ kv_cache = kv_cache_raw.transpose(0, 1)
167
+ kv_cache = kv_cache.transpose(2, 3).contiguous().transpose(2, 3)
168
+ return kv_cache
169
+
170
+
171
+def _run_trtllm_integration(batch_spec):
172
+ """Run TRTLLM attention through the full FlashInfer pipeline
173
+ and compare against an SDPA reference."""
174
+ set_random_seed(42)
175
+ device = torch.device("cuda:0")
176
+
177
+ vllm_config = create_vllm_config(
178
+ model_name=MODEL,
179
+ max_model_len=max(batch_spec.seq_lens),
180
+ block_size=BLOCK_SIZE,
181
+ num_gpu_blocks=NUM_GPU_BLOCKS,
182
+ )
183
+ vllm_config.attention_config.use_trtllm_attention = True
184
+
185
+ num_q_heads = vllm_config.model_config.get_num_attention_heads(
186
+ vllm_config.parallel_config
187
+ )
188
+ num_kv_heads = vllm_config.model_config.get_num_kv_heads(
189
+ vllm_config.parallel_config
190
+ )
191
+ head_size = vllm_config.model_config.get_head_size()
192
+ dtype = vllm_config.model_config.dtype
193
+ scale = 1.0 / (head_size**0.5)
194
+
195
+ # 1. Generate data and compute SDPA reference
196
+ all_q, all_k, all_v = [], [], []
197
+ all_sdpa_out = []
198
+ k_contexts, v_contexts = [], []
199
+
200
+ for i in range(batch_spec.batch_size):
201
+ s_len = batch_spec.seq_lens[i]
202
+ q_len = batch_spec.query_lens[i]
203
+ ctx_len = s_len - q_len
204
+
205
+ q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device)
206
+ k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
207
+ v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
208
+
209
+ # SDPA reference (N=1, H, L, D)
210
+ q_sdpa = q.unsqueeze(0).transpose(1, 2)
211
+ k_sdpa = k_full.unsqueeze(0).transpose(1, 2)
212
+ v_sdpa = v_full.unsqueeze(0).transpose(1, 2)
213
+
214
+ if num_q_heads != num_kv_heads:
215
+ repeats = num_q_heads // num_kv_heads
216
+ k_sdpa = k_sdpa.repeat_interleave(repeats, dim=1)
217
+ v_sdpa = v_sdpa.repeat_interleave(repeats, dim=1)
218
+
219
+ def causal_mask_mod(b, h, q_idx, kv_idx, *, context_len):
220
+ return (q_idx + context_len) >= kv_idx
221
+
222
+ mask_fn = partial(causal_mask_mod, context_len=ctx_len)
223
+ block_mask = create_block_mask(
224
+ mask_fn, B=None, H=None, Q_LEN=q_len, KV_LEN=s_len, device=device
225
+ )
226
+ sdpa_out = flex_attention(
227
+ q_sdpa,
228
+ k_sdpa,
229
+ v_sdpa,
230
+ block_mask=block_mask,
231
+ scale=scale,
232
+ enable_gqa=True,
233
+ )
234
+ all_sdpa_out.append(sdpa_out.transpose(1, 2).squeeze(0))
235
+
236
+ all_q.append(q)
237
+ all_k.append(k_full[ctx_len:])
238
+ all_v.append(v_full[ctx_len:])
239
+ k_contexts.append(k_full[:ctx_len])
240
+ v_contexts.append(v_full[:ctx_len])
241
+
242
+ query_vllm = torch.cat(all_q, dim=0)
243
+ key_vllm = torch.cat(all_k, dim=0)
244
+ value_vllm = torch.cat(all_v, dim=0)
245
+ sdpa_output = torch.cat(all_sdpa_out, dim=0)
246
+
247
+ common_attn_metadata = create_common_attn_metadata(batch_spec, BLOCK_SIZE, device)
248
+
249
+ # 2. Create HND KV cache
250
+ kv_cache = _create_hnd_kv_cache(
251
+ k_contexts,
252
+ v_contexts,
253
+ BLOCK_SIZE,
254
+ num_kv_heads,
255
+ head_size,
256
+ dtype,
257
+ device,
258
+ NUM_GPU_BLOCKS,
259
+ common_attn_metadata,
260
+ )
261
+
262
+ # 3. Run through FlashInfer with TRTLLM enabled
263
+ set_kv_cache_layout("HND")
264
+ get_kv_cache_layout.cache_clear()
265
+
266
+ try:
267
+ kv_cache_spec = FullAttentionSpec(
268
+ block_size=BLOCK_SIZE,
269
+ num_kv_heads=num_kv_heads,
270
+ head_size=head_size,
271
+ dtype=dtype,
272
+ )
273
+ layer_names = ["test_layer_0"]
274
+
275
+ with (
276
+ set_current_vllm_config(vllm_config),
277
+ unittest.mock.patch(
278
+ "vllm.utils.flashinfer.supports_trtllm_attention",
279
+ return_value=True,
280
+ ),
281
+ unittest.mock.patch(
282
+ "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
283
+ _mock_get_per_layer_parameters,
284
+ ),
285
+ ):
286
+ builder = FlashInferMetadataBuilder(
287
+ kv_cache_spec, layer_names, vllm_config, device
288
+ )
289
+ attn_metadata = builder.build(
290
+ common_prefix_len=0,
291
+ common_attn_metadata=common_attn_metadata,
292
+ )
293
+
294
+ # Verify the correct TRTLLM metadata types were produced.
295
+ has_prefills = any(ql > 1 for ql in batch_spec.query_lens)
296
+ has_decodes = any(ql == 1 for ql in batch_spec.query_lens)
297
+
298
+ if has_prefills:
299
+ assert isinstance(attn_metadata.prefill, TRTLLMPrefill), (
300
+ f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
301
+ )
302
+ if has_decodes:
303
+ assert isinstance(attn_metadata.decode, TRTLLMDecode), (
304
+ f"Expected TRTLLMDecode, got {type(attn_metadata.decode)}"
305
+ )
306
+
307
+ impl = FlashInferImpl(
308
+ num_heads=num_q_heads,
309
+ head_size=head_size,
310
+ scale=scale,
311
+ num_kv_heads=num_kv_heads,
312
+ alibi_slopes=None,
313
+ sliding_window=None,
314
+ kv_cache_dtype="auto",
315
+ )
316
+
317
+ mock_layer = MockAttentionLayer(device)
318
+ output = torch.empty_like(query_vllm)
319
+
320
+ impl.do_kv_cache_update(
321
+ mock_layer,
322
+ key_vllm,
323
+ value_vllm,
324
+ kv_cache,
325
+ attn_metadata.slot_mapping,
326
+ )
327
+
328
+ output = impl.forward(
329
+ mock_layer,
330
+ query_vllm,
331
+ key_vllm,
332
+ value_vllm,
333
+ kv_cache,
334
+ attn_metadata,
335
+ output=output,
336
+ )
337
+
338
+ # 4. Compare against SDPA reference
339
+ torch.testing.assert_close(
340
+ output,
341
+ sdpa_output,
342
+ atol=1e-2,
343
+ rtol=1e-2,
344
+ )
345
+
346
+ finally:
347
+ set_kv_cache_layout(None)
348
+ get_kv_cache_layout.cache_clear()
349
+
350
+
351
+@pytest.mark.parametrize(
352
+ "batch_spec_name",
353
+ list(BATCH_SPECS.keys()),
354
+)
355
+@torch.inference_mode()
356
+def test_trtllm_gen_full_attention_integration(batch_spec_name: str):
357
+ """Test TRTLLM gen-full attention through the full FlashInfer
358
+ MetadataBuilder.build() -> FlashInferImpl.forward() pipeline,
359
+ with real TRTLLM kernels on Blackwell."""
360
+ _run_trtllm_integration(BATCH_SPECS[batch_spec_name])