| 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 dataclasses |
| 6 | import unittest.mock |
| 7 | from functools import partial |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | from torch.nn.attention.flex_attention import create_block_mask, flex_attention |
| 12 | |
| 13 | from tests.kernels.attention.test_trtllm_kvfp8_dequant import make_random_kv_cache |
| 14 | from tests.v1.attention.test_attention_backends import create_and_prepopulate_kv_cache |
| 15 | from tests.v1.attention.utils import ( |
| 16 | BatchSpec, |
| 17 | create_common_attn_metadata, |
| 18 | create_vllm_config, |
| 19 | ) |
| 20 | from vllm.config import set_current_vllm_config |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.utils.torch_utils import ( |
| 23 | is_strictly_contiguous, |
| 24 | nvfp4_kv_cache_full_dim, |
| 25 | set_random_seed, |
| 26 | ) |
| 27 | from vllm.v1.attention.backends.utils import PerLayerParameters |
| 28 | from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode |
| 29 | |
| 30 | if not current_platform.is_device_capability_family(100): |
| 31 | pytest.skip( |
| 32 | "TRTLLM integration tests require NVIDIA Blackwell (SM100).", |
| 33 | allow_module_level=True, |
| 34 | ) |
| 35 | |
| 36 | from vllm.v1.attention.backends.flashinfer import ( # noqa: E402 |
| 37 | FlashInferDecodeKernel, |
| 38 | FlashInferImpl, |
| 39 | FlashInferMetadataBuilder, |
| 40 | FlashInferTrtllmAPIDecode, |
| 41 | TRTLLMPrefill, |
| 42 | trtllm_prefill_attn_kvfp8_dequant, |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | class MockAttentionLayer: |
| 47 | """Minimal mock of an attention layer for testing.""" |
| 48 | |
| 49 | def __init__(self, device: torch.device): |
| 50 | self._q_scale = torch.tensor(2.0, device=device) |
| 51 | self._k_scale = torch.tensor(3.0, device=device) |
| 52 | self._v_scale = torch.tensor(4.0, device=device) |
| 53 | self._q_scale_float = 2.0 |
| 54 | self._k_scale_float = 3.0 |
| 55 | self._v_scale_float = 4.0 |
| 56 | self._o_scale_float = None |
| 57 | |
| 58 | |
| 59 | MODEL = "Qwen/Qwen2.5-0.5B" |
| 60 | MODEL_NVFP4 = "Qwen/Qwen3-4B" # nvfp4 needs head_dim >= 128 (or 80) |
| 61 | BLOCK_SIZE = 16 |
| 62 | NUM_GPU_BLOCKS = 8192 |
| 63 | DEVICE_TYPE = current_platform.device_type |
| 64 | |
| 65 | BATCH_SPECS = { |
| 66 | "decode_only": BatchSpec( |
| 67 | seq_lens=[128, 256, 512], |
| 68 | query_lens=[1, 1, 1], |
| 69 | ), |
| 70 | "prefill_only": BatchSpec( |
| 71 | seq_lens=[64, 128, 256], |
| 72 | query_lens=[16, 32, 16], |
| 73 | ), |
| 74 | "mixed": BatchSpec( |
| 75 | seq_lens=[128, 256, 512, 128], |
| 76 | query_lens=[1, 1, 8, 16], |
| 77 | ), |
| 78 | } |
| 79 | |
| 80 | |
| 81 | def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls): |
| 82 | head_size = vllm_config.model_config.get_head_size() |
| 83 | return { |
| 84 | name: PerLayerParameters( |
| 85 | window_left=-1, |
| 86 | logits_soft_cap=0.0, |
| 87 | sm_scale=1.0 / (head_size**0.5), |
| 88 | ) |
| 89 | for name in layer_names |
| 90 | } |
| 91 | |
| 92 | |
| 93 | @pytest.mark.parametrize( |
| 94 | "block_size,num_prefills,has_host_lengths,quantize_query", |
| 95 | [ |
| 96 | (16, 1, True, False), |
| 97 | (64, 2, True, False), |
| 98 | (128, 2, True, False), |
| 99 | (16, 2, False, False), |
| 100 | (16, 2, True, True), |
| 101 | ], |
| 102 | ) |
| 103 | @torch.inference_mode() |
| 104 | def test_prefill_dequant_scratch_ignores_unused_capacity( |
| 105 | block_size, num_prefills, has_host_lengths, quantize_query, monkeypatch |
| 106 | ): |
| 107 | """Bound scratch by attended prefill KV, including cached prefixes.""" |
| 108 | device = torch.device("cuda") |
| 109 | config = create_vllm_config(model_name=MODEL, max_model_len=32768) |
| 110 | config.cache_config.cache_dtype = "fp8" |
| 111 | config.cache_config.kv_cache_layout = "LBHNC" |
| 112 | config.attention_config.use_trtllm_attention = True |
| 113 | config.attention_config.disable_flashinfer_q_quantization = not quantize_query |
| 114 | monkeypatch.setattr( |
| 115 | "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters", |
| 116 | _mock_get_per_layer_parameters, |
| 117 | ) |
| 118 | spec = FullAttentionSpec( |
| 119 | block_size=block_size, |
| 120 | num_kv_heads=2, |
| 121 | head_size=64, |
| 122 | dtype=torch.uint8, |
| 123 | kv_quant_mode=KVQuantMode.FP8_PER_TENSOR, |
| 124 | ) |
| 125 | kv_cache, scale = make_random_kv_cache(32, 2, block_size, 64) |
| 126 | batch = BatchSpec( |
| 127 | seq_lens=[1024] + [17, 65][-num_prefills:], |
| 128 | query_lens=[1] + [2, 8][-num_prefills:], |
| 129 | ) |
| 130 | common = create_common_attn_metadata(batch, block_size, device) |
| 131 | if not has_host_lengths: |
| 132 | common.seq_lens_cpu_upper_bound = None |
| 133 | live_pages = ((65 if has_host_lengths else 1024) + block_size - 1) // block_size |
| 134 | with set_current_vllm_config(config): |
| 135 | builder = FlashInferMetadataBuilder(spec, ["layer.0"], config, device) |
| 136 | for capacity in (64, 256): |
| 137 | # Unused columns can contain nonzero page IDs. |
| 138 | common.block_table_tensor = torch.randint( |
| 139 | 1, 32, (batch.batch_size, capacity), dtype=torch.int32, device=device |
| 140 | ) |
| 141 | metadata = builder.build(0, common) |
| 142 | assert isinstance(metadata.prefill, TRTLLMPrefill) |
| 143 | table = metadata.prefill.block_tables |
| 144 | assert is_strictly_contiguous(table) |
| 145 | assert (metadata.q_data_type_prefill == current_platform.fp8_dtype()) == ( |
| 146 | quantize_query |
| 147 | ) |
| 148 | if not quantize_query: |
| 149 | converted, _ = trtllm_prefill_attn_kvfp8_dequant( |
| 150 | kv_cache, table, scale, scale, metadata.q_data_type_prefill |
| 151 | ) |
| 152 | assert converted.shape[0] == num_prefills * live_pages + 1 |
| 153 | width = capacity if quantize_query else live_pages |
| 154 | torch.testing.assert_close(table, common.block_table_tensor[1:, :width]) |
| 155 | |
| 156 | |
| 157 | def _create_nvfp4_hnd_kv_cache( |
| 158 | k_contexts, |
| 159 | v_contexts, |
| 160 | block_size, |
| 161 | num_kv_heads, |
| 162 | head_size, |
| 163 | dtype, |
| 164 | device, |
| 165 | num_blocks, |
| 166 | common_attn_metadata, |
| 167 | kv_scale_val, |
| 168 | ): |
| 169 | """Create an nvfp4 KV cache by quantizing bf16 context via |
| 170 | reshape_and_cache_flash, using the same block-table layout as |
| 171 | create_and_prepopulate_kv_cache. |
| 172 | |
| 173 | The returned tensor is dtype ``uint8`` with head-group layout |
| 174 | ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)`` |
| 175 | where K heads occupy the first ``num_kv_heads`` heads and V heads the second. |
| 176 | Each ``full_dim = head_size // 2 + head_size // 16`` block packs two regions: |
| 177 | - **FP4 data** (``head_size // 2`` bytes): pairs of E2M1 values, |
| 178 | two per byte. |
| 179 | - **FP8 block scales** (``head_size // 16`` bytes): one E4M3 |
| 180 | scale per 16-element block. |
| 181 | |
| 182 | Args: |
| 183 | k_contexts: List of key context tensors, one per sequence. |
| 184 | v_contexts: List of value context tensors, one per sequence. |
| 185 | block_size: Number of tokens per cache block. |
| 186 | num_kv_heads: Number of key/value heads. |
| 187 | head_size: Head dimension (must be divisible by 16). |
| 188 | dtype: Source data type for the bf16 intermediate cache. |
| 189 | device: Target device. |
| 190 | num_blocks: Total number of blocks to allocate. |
| 191 | common_attn_metadata: Metadata containing block tables and |
| 192 | sequence lengths. |
| 193 | kv_scale_val: Scalar float used as both k_scale and v_scale |
| 194 | during quantization. |
| 195 | |
| 196 | Returns: |
| 197 | ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, LBHNC-strided). |
| 198 | |
| 199 | """ |
| 200 | # First create a bf16 cache so block tables are populated. |
| 201 | bf16_cache = create_and_prepopulate_kv_cache( |
| 202 | k_contexts, |
| 203 | v_contexts, |
| 204 | block_size, |
| 205 | num_kv_heads, |
| 206 | head_size, |
| 207 | dtype, |
| 208 | device, |
| 209 | num_blocks, |
| 210 | common_attn_metadata, |
| 211 | layout=KVCacheLayout.LBHNC, |
| 212 | ) |
| 213 | |
| 214 | # (num_blocks, 2 * num_kv_heads, block_size, full_dim) — K heads first, then V heads |
| 215 | full_dim = nvfp4_kv_cache_full_dim(head_size) |
| 216 | nvfp4_cache = torch.zeros( |
| 217 | (num_blocks, 2 * num_kv_heads, block_size, full_dim), |
| 218 | dtype=torch.uint8, |
| 219 | device=device, |
| 220 | ) |
| 221 | k_cache, v_cache = nvfp4_cache.split(num_kv_heads, dim=1) |
| 222 | |
| 223 | # Flatten bf16 context into tokens and quantize via reshape_and_cache_flash. |
| 224 | # bf16_cache is (B, H, N, 2*hs); split K/V on the content dim. |
| 225 | block_table = common_attn_metadata.block_table_tensor |
| 226 | seq_lens = common_attn_metadata.seq_lens.cpu() |
| 227 | query_lens = ( |
| 228 | common_attn_metadata.query_start_loc_cpu[1:] |
| 229 | - common_attn_metadata.query_start_loc_cpu[:-1] |
| 230 | ) |
| 231 | kv_scale_t = torch.tensor(kv_scale_val, dtype=torch.float32, device=device) |
| 232 | |
| 233 | for i in range(len(k_contexts)): |
| 234 | ctx_len = int(seq_lens[i]) - int(query_lens[i]) |
| 235 | if ctx_len == 0: |
| 236 | continue |
| 237 | # Gather context tokens from the bf16 cache using block table. |
| 238 | n_ctx_blocks = (ctx_len + block_size - 1) // block_size |
| 239 | blocks = block_table[i, :n_ctx_blocks] |
| 240 | k_ctx = ( |
| 241 | bf16_cache[blocks, :, :, :head_size] |
| 242 | .transpose(1, 2) |
| 243 | .reshape(-1, num_kv_heads, head_size)[:ctx_len] |
| 244 | ) |
| 245 | v_ctx = ( |
| 246 | bf16_cache[blocks, :, :, head_size:] |
| 247 | .transpose(1, 2) |
| 248 | .reshape(-1, num_kv_heads, head_size)[:ctx_len] |
| 249 | ) |
| 250 | # Build slot mapping for these context tokens. |
| 251 | token_offsets = torch.arange(ctx_len, device=device) |
| 252 | block_indices = token_offsets // block_size |
| 253 | intra_offsets = token_offsets % block_size |
| 254 | slots = block_table[i, block_indices] * block_size + intra_offsets |
| 255 | # reshape_and_cache_flash expects (B, N, H, D) cache views. |
| 256 | torch.ops._C_cache_ops.reshape_and_cache_flash( |
| 257 | k_ctx, |
| 258 | v_ctx, |
| 259 | k_cache.transpose(1, 2), |
| 260 | v_cache.transpose(1, 2), |
| 261 | slots, |
| 262 | "nvfp4", |
| 263 | kv_scale_t, |
| 264 | kv_scale_t, |
| 265 | ) |
| 266 | |
| 267 | return nvfp4_cache |
| 268 | |
| 269 | |
| 270 | def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL): |
| 271 | """Run TRTLLM attention through the full FlashInfer pipeline |
| 272 | and compare against an SDPA reference.""" |
| 273 | set_random_seed(42) |
| 274 | device = torch.device(f"{DEVICE_TYPE}:0") |
| 275 | |
| 276 | vllm_config = create_vllm_config( |
| 277 | model_name=model_name, |
| 278 | max_model_len=max(batch_spec.seq_lens), |
| 279 | block_size=BLOCK_SIZE, |
| 280 | num_gpu_blocks=NUM_GPU_BLOCKS, |
| 281 | ) |
| 282 | vllm_config.attention_config.use_trtllm_attention = True |
| 283 | vllm_config.cache_config.cache_dtype = kv_cache_dtype |
| 284 | |
| 285 | num_q_heads = vllm_config.model_config.get_num_attention_heads( |
| 286 | vllm_config.parallel_config |
| 287 | ) |
| 288 | num_kv_heads = vllm_config.model_config.get_num_kv_heads( |
| 289 | vllm_config.parallel_config |
| 290 | ) |
| 291 | head_size = vllm_config.model_config.get_head_size() |
| 292 | dtype = vllm_config.model_config.dtype |
| 293 | scale = 1.0 / (head_size**0.5) |
| 294 | |
| 295 | # 1. Generate data and compute SDPA reference |
| 296 | all_q, all_k, all_v = [], [], [] |
| 297 | all_sdpa_out = [] |
| 298 | k_contexts, v_contexts = [], [] |
| 299 | |
| 300 | for i in range(batch_spec.batch_size): |
| 301 | s_len = batch_spec.seq_lens[i] |
| 302 | q_len = batch_spec.query_lens[i] |
| 303 | ctx_len = s_len - q_len |
| 304 | |
| 305 | q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device) |
| 306 | k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device) |
| 307 | v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device) |
| 308 | |
| 309 | # SDPA reference (N=1, H, L, D) |
| 310 | q_sdpa = q.unsqueeze(0).transpose(1, 2) |
| 311 | k_sdpa = k_full.unsqueeze(0).transpose(1, 2) |
| 312 | v_sdpa = v_full.unsqueeze(0).transpose(1, 2) |
| 313 | |
| 314 | if num_q_heads != num_kv_heads: |
| 315 | repeats = num_q_heads // num_kv_heads |
| 316 | k_sdpa = k_sdpa.repeat_interleave(repeats, dim=1) |
| 317 | v_sdpa = v_sdpa.repeat_interleave(repeats, dim=1) |
| 318 | |
| 319 | def causal_mask_mod(b, h, q_idx, kv_idx, *, context_len): |
| 320 | return (q_idx + context_len) >= kv_idx |
| 321 | |
| 322 | mask_fn = partial(causal_mask_mod, context_len=ctx_len) |
| 323 | block_mask = create_block_mask( |
| 324 | mask_fn, B=None, H=None, Q_LEN=q_len, KV_LEN=s_len, device=device |
| 325 | ) |
| 326 | sdpa_out = flex_attention( |
| 327 | q_sdpa, |
| 328 | k_sdpa, |
| 329 | v_sdpa, |
| 330 | block_mask=block_mask, |
| 331 | scale=scale, |
| 332 | enable_gqa=True, |
| 333 | ) |
| 334 | all_sdpa_out.append(sdpa_out.transpose(1, 2).squeeze(0)) |
| 335 | |
| 336 | all_q.append(q) |
| 337 | all_k.append(k_full[ctx_len:]) |
| 338 | all_v.append(v_full[ctx_len:]) |
| 339 | k_contexts.append(k_full[:ctx_len]) |
| 340 | v_contexts.append(v_full[:ctx_len]) |
| 341 | |
| 342 | query_vllm = torch.cat(all_q, dim=0) |
| 343 | key_vllm = torch.cat(all_k, dim=0) |
| 344 | value_vllm = torch.cat(all_v, dim=0) |
| 345 | sdpa_output = torch.cat(all_sdpa_out, dim=0) |
| 346 | |
| 347 | common_attn_metadata = create_common_attn_metadata(batch_spec, BLOCK_SIZE, device) |
| 348 | |
| 349 | # 2. Create HND KV cache |
| 350 | is_nvfp4 = kv_cache_dtype == "nvfp4" |
| 351 | if is_nvfp4: |
| 352 | # Compute a global scale from the context data. |
| 353 | all_ctx = torch.cat(k_contexts + v_contexts, dim=0) |
| 354 | kv_scale_val = (all_ctx.abs().amax() / 448.0).item() |
| 355 | kv_cache = _create_nvfp4_hnd_kv_cache( |
| 356 | k_contexts, |
| 357 | v_contexts, |
| 358 | BLOCK_SIZE, |
| 359 | num_kv_heads, |
| 360 | head_size, |
| 361 | dtype, |
| 362 | device, |
| 363 | NUM_GPU_BLOCKS, |
| 364 | common_attn_metadata, |
| 365 | kv_scale_val, |
| 366 | ) |
| 367 | else: |
| 368 | kv_scale_val = 1.0 |
| 369 | kv_cache = create_and_prepopulate_kv_cache( |
| 370 | k_contexts, |
| 371 | v_contexts, |
| 372 | BLOCK_SIZE, |
| 373 | num_kv_heads, |
| 374 | head_size, |
| 375 | dtype, |
| 376 | device, |
| 377 | NUM_GPU_BLOCKS, |
| 378 | common_attn_metadata, |
| 379 | layout=KVCacheLayout.LBHNC, |
| 380 | ) |
| 381 | |
| 382 | # 3. Run through FlashInfer with TRTLLM enabled |
| 383 | vllm_config.cache_config.kv_cache_layout = "LBHNC" |
| 384 | |
| 385 | is_nvfp4 = kv_cache_dtype == "nvfp4" |
| 386 | kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE |
| 387 | spec_dtype = torch.uint8 if is_nvfp4 else dtype |
| 388 | kv_cache_spec = FullAttentionSpec( |
| 389 | block_size=BLOCK_SIZE, |
| 390 | num_kv_heads=num_kv_heads, |
| 391 | head_size=head_size, |
| 392 | dtype=spec_dtype, |
| 393 | kv_quant_mode=kv_quant_mode, |
| 394 | ) |
| 395 | layer_names = ["test_layer_0"] |
| 396 | |
| 397 | with ( |
| 398 | set_current_vllm_config(vllm_config), |
| 399 | unittest.mock.patch( |
| 400 | "vllm.utils.flashinfer.supports_trtllm_attention", |
| 401 | return_value=True, |
| 402 | ), |
| 403 | unittest.mock.patch( |
| 404 | "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters", |
| 405 | _mock_get_per_layer_parameters, |
| 406 | ), |
| 407 | ): |
| 408 | builder = FlashInferMetadataBuilder( |
| 409 | kv_cache_spec, layer_names, vllm_config, device |
| 410 | ) |
| 411 | attn_metadata = builder.build( |
| 412 | common_prefix_len=0, |
| 413 | common_attn_metadata=common_attn_metadata, |
| 414 | ) |
| 415 | |
| 416 | # Verify the correct TRTLLM metadata types were produced. |
| 417 | has_prefills = any(ql > 1 for ql in batch_spec.query_lens) |
| 418 | has_decodes = any(ql == 1 for ql in batch_spec.query_lens) |
| 419 | |
| 420 | if has_prefills: |
| 421 | assert isinstance(attn_metadata.prefill, TRTLLMPrefill), ( |
| 422 | f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}" |
| 423 | ) |
| 424 | if has_decodes: |
| 425 | assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), ( |
| 426 | f"Expected FlashInferTrtllmAPIDecode, got {type(attn_metadata.decode)}" |
| 427 | ) |
| 428 | assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN |
| 429 | |
| 430 | impl = FlashInferImpl( |
| 431 | num_heads=num_q_heads, |
| 432 | head_size=head_size, |
| 433 | scale=scale, |
| 434 | num_kv_heads=num_kv_heads, |
| 435 | alibi_slopes=None, |
| 436 | sliding_window=None, |
| 437 | kv_cache_dtype=kv_cache_dtype, |
| 438 | ) |
| 439 | |
| 440 | mock_layer = MockAttentionLayer(device) |
| 441 | if is_nvfp4: |
| 442 | # For nvfp4, k_scale/v_scale are the global quantization |
| 443 | # scales (amax/448) used by reshape_and_cache_flash. |
| 444 | kv_scale_t = torch.tensor(kv_scale_val, dtype=torch.float32, device=device) |
| 445 | mock_layer._k_scale = kv_scale_t |
| 446 | mock_layer._v_scale = kv_scale_t |
| 447 | mock_layer._k_scale_float = kv_scale_val |
| 448 | mock_layer._v_scale_float = kv_scale_val |
| 449 | output = torch.empty_like(query_vllm) |
| 450 | |
| 451 | impl.do_kv_cache_update( |
| 452 | mock_layer, |
| 453 | key_vllm, |
| 454 | value_vllm, |
| 455 | kv_cache, |
| 456 | attn_metadata.slot_mapping, |
| 457 | ) |
| 458 | |
| 459 | # nvfp4 trtllm kernel requires FP8 queries. In the real |
| 460 | # pipeline the attention layer handles this; here we |
| 461 | # quantize manually. |
| 462 | if is_nvfp4: |
| 463 | finfo = torch.finfo(torch.float8_e4m3fn) |
| 464 | q_amax = query_vllm.abs().amax().clamp(min=1e-12) |
| 465 | q_s = (finfo.max / q_amax * 0.1).item() |
| 466 | query_vllm = ( |
| 467 | (query_vllm * q_s).clamp(finfo.min, finfo.max).to(torch.float8_e4m3fn) |
| 468 | ) |
| 469 | mock_layer._q_scale = torch.tensor( |
| 470 | 1.0 / q_s, dtype=torch.float32, device=device |
| 471 | ) |
| 472 | mock_layer._q_scale_float = 1.0 / q_s |
| 473 | |
| 474 | output = impl.forward( |
| 475 | mock_layer, |
| 476 | query_vllm, |
| 477 | key_vllm, |
| 478 | value_vllm, |
| 479 | kv_cache, |
| 480 | attn_metadata, |
| 481 | output=output, |
| 482 | ) |
| 483 | |
| 484 | # 4. Compare against SDPA reference |
| 485 | if is_nvfp4: |
| 486 | atol, rtol = 1.0, 1.0 # nvfp4 has higher quantization error |
| 487 | else: |
| 488 | atol, rtol = 1e-2, 1e-2 |
| 489 | torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol) |
| 490 | |
| 491 | |
| 492 | @pytest.mark.parametrize( |
| 493 | "batch_spec_name", |
| 494 | list(BATCH_SPECS.keys()), |
| 495 | ) |
| 496 | @torch.inference_mode() |
| 497 | def test_trtllm_gen_full_attention_integration(batch_spec_name: str): |
| 498 | """Test TRTLLM gen-full attention through the full FlashInfer |
| 499 | MetadataBuilder.build() -> FlashInferImpl.forward() pipeline, |
| 500 | with real TRTLLM kernels on Blackwell.""" |
| 501 | _run_trtllm_integration(BATCH_SPECS[batch_spec_name]) |
| 502 | |
| 503 | |
| 504 | @pytest.mark.parametrize( |
| 505 | "batch_spec_name", |
| 506 | list(BATCH_SPECS.keys()), |
| 507 | ) |
| 508 | @torch.inference_mode() |
| 509 | def test_trtllm_gen_nvfp4_kv_integration(batch_spec_name: str): |
| 510 | """Test TRTLLM attention with nvfp4 KV cache through the full |
| 511 | FlashInfer MetadataBuilder.build() -> FlashInferImpl.forward() pipeline.""" |
| 512 | _run_trtllm_integration( |
| 513 | BATCH_SPECS[batch_spec_name], |
| 514 | kv_cache_dtype="nvfp4", |
| 515 | model_name=MODEL_NVFP4, |
| 516 | ) |
| 517 | |
| 518 | |
| 519 | @torch.inference_mode() |
| 520 | def test_trtllm_gen_draft_decode_metadata_tracks_in_place_seq_lens( |
| 521 | monkeypatch: pytest.MonkeyPatch, |
| 522 | ): |
| 523 | """Fused draft steps reuse the step-1 build after advancing seq_lens in |
| 524 | place, so its decode metadata must match the split path's step-2 rebuild.""" |
| 525 | monkeypatch.setattr( |
| 526 | "vllm.utils.flashinfer.supports_trtllm_attention", |
| 527 | unittest.mock.Mock(return_value=True), |
| 528 | ) |
| 529 | monkeypatch.setattr( |
| 530 | "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters", |
| 531 | _mock_get_per_layer_parameters, |
| 532 | ) |
| 533 | vllm_config = create_vllm_config( |
| 534 | model_name=MODEL, block_size=BLOCK_SIZE, num_gpu_blocks=NUM_GPU_BLOCKS |
| 535 | ) |
| 536 | model_config = vllm_config.model_config |
| 537 | kv_cache_spec = FullAttentionSpec( |
| 538 | block_size=BLOCK_SIZE, |
| 539 | num_kv_heads=model_config.get_num_kv_heads(vllm_config.parallel_config), |
| 540 | head_size=model_config.get_head_size(), |
| 541 | dtype=model_config.dtype, |
| 542 | ) |
| 543 | device = torch.device(f"{DEVICE_TYPE}:0") |
| 544 | with set_current_vllm_config(vllm_config): |
| 545 | builder = FlashInferMetadataBuilder( |
| 546 | kv_cache_spec, ["test_layer_0"], vllm_config, device |
| 547 | ) |
| 548 | assert builder.supports_draft_decode_metadata_update |
| 549 | |
| 550 | step1 = create_common_attn_metadata(BATCH_SPECS["decode_only"], BLOCK_SIZE, device) |
| 551 | reused = builder.build(common_prefix_len=0, common_attn_metadata=step1) |
| 552 | step1.seq_lens.add_(1) |
| 553 | step2 = dataclasses.replace(step1, seq_lens_cpu_upper_bound=step1.seq_lens.cpu()) |
| 554 | rebuilt = builder.build(common_prefix_len=0, common_attn_metadata=step2) |
| 555 | |
| 556 | assert isinstance(reused.decode, FlashInferTrtllmAPIDecode) |
| 557 | for field in dataclasses.fields(reused.decode): |
| 558 | value = getattr(reused.decode, field.name) |
| 559 | expected = getattr(rebuilt.decode, field.name) |
| 560 | if isinstance(value, torch.Tensor): |
| 561 | assert torch.equal(value, expected), field.name |
| 562 | else: |
| 563 | assert value == expected, field.name |