@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6a9c69fa85

[Attention][Spec Decode] Support varlen trtllm-gen decode for adaptive verification (#52157)

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>

Guan-Ming Chiu committed Aug 25, 2026 at 02:43 UTC 6a9c69fa851389dcf1ee5d3a2363e27af665d26d

2 files changed +138 -7

tests/kernels/attention/test_flashinfer_trtllm_attention.py

+93

@@ -358,6 +358,99 @@ def test_flashinfer_trtllm_decode_with_baseline(

358 )

359

360

361 +@pytest.mark.parametrize("quant_dtypes", [(None, None, None), (None, FP8_DTYPE, None)])

362 +@pytest.mark.parametrize("q_lens_cfg", [[1, 3, 2, 4], [2, 2, 1, 5, 3]])

363 +@torch.inference_mode

364 +def test_flashinfer_trtllm_decode_varlen(

365 + quant_dtypes: tuple[torch.dtype | None, torch.dtype | None, torch.dtype | None],

366 + q_lens_cfg: list[int],

367 +) -> None:

368 + """Ragged decode via cum_seq_lens_q must match the varlen baseline."""

369 + dtype = torch.bfloat16

370 + num_qo_heads, num_kv_heads = 40, 8

371 + head_size = 128

372 + block_size = 16

373 + max_kv_len = 4096

374 + torch.set_default_device("cuda")

375 + set_random_seed(42)

376 +

377 + _, kv_quant_dtype, _ = quant_dtypes

378 + kv_quant_dtype = kv_quant_dtype or dtype

379 + sm_scale = float(1.0 / (head_size**0.5))

380 + batch_size = len(q_lens_cfg)

381 +

382 + q_lens = torch.tensor(q_lens_cfg, dtype=torch.int32)

383 + q_indptr = torch.cat(

384 + [

385 + torch.tensor([0], dtype=torch.int32),

386 + torch.cumsum(q_lens, dim=0, dtype=torch.int32),

387 + ]

388 + )

389 + query = torch.randn(int(q_lens.sum()), num_qo_heads, head_size, dtype=dtype)

390 +

391 + kv_lens = torch.randint(1, max_kv_len, (batch_size,), dtype=torch.int32)

392 + kv_lens[-1] = max_kv_len

393 + seq_lens = kv_lens + q_lens

394 + max_seq_len = int(seq_lens.max())

395 +

396 + kv_cache_shape = (NUM_BLOCKS, 2, num_kv_heads, block_size, head_size)

397 + kv_cache = torch.randn(kv_cache_shape, dtype=dtype)

398 + kv_cache, kv_cache_sf, kv_scale, ref_kv_cache, _ = make_quantized_kv_cache(

399 + kv_cache, kv_quant_dtype, block_size, head_size

400 + )

401 +

402 + max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size

403 + block_tables = torch.randint(

404 + 0, NUM_BLOCKS, (batch_size, max_num_blocks_per_seq), dtype=torch.int32

405 + )

406 + kv_indptr, kv_indices, kv_last_page_lens = build_paged_kv_metadata(

407 + seq_lens, block_tables, block_size

408 + )

409 + workspace_buffer = torch.zeros(128 * 1024 * 1024, dtype=torch.int8)

410 +

411 + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(

412 + float_workspace_buffer=workspace_buffer, kv_layout="HND", backend="fa2"

413 + )

414 + wrapper.plan(

415 + qo_indptr=q_indptr,

416 + paged_kv_indptr=kv_indptr,

417 + paged_kv_indices=kv_indices,

418 + paged_kv_last_page_len=kv_last_page_lens,

419 + num_qo_heads=num_qo_heads,

420 + num_kv_heads=num_kv_heads,

421 + head_dim_qk=head_size,

422 + page_size=block_size,

423 + causal=True,

424 + sm_scale=sm_scale,

425 + window_left=-1,

426 + q_data_type=dtype,

427 + kv_data_type=dtype,

428 + )

429 + output = torch.empty(query.shape, dtype=dtype)

430 + wrapper.run(query, ref_kv_cache, out=output)

431 +

432 + output_trtllm = torch.empty(query.shape, dtype=dtype)

433 + flashinfer.decode.trtllm_batch_decode_with_kv_cache(

434 + query=query,

435 + kv_cache=kv_cache,

436 + workspace_buffer=workspace_buffer,

437 + block_tables=block_tables,

438 + seq_lens=seq_lens,

439 + max_seq_len=max_seq_len,

440 + bmm1_scale=kv_scale * sm_scale,

441 + bmm2_scale=kv_scale,

442 + window_left=-1,

443 + out=output_trtllm,

444 + q_len_per_req=None,

445 + max_q_len=int(q_lens.max()),

446 + cum_seq_lens_q=q_indptr,

447 + kv_cache_sf=kv_cache_sf,

448 + )

449 +

450 + rtol, atol = (4e-2, 6e-2) if kv_quant_dtype == FP8_DTYPE else (1e-2, 1e-2)

451 + torch.testing.assert_close(output, output_trtllm, atol=atol, rtol=rtol)

452 +

453 +

454 @pytest.mark.parametrize("dtype", DTYPE)

455 @pytest.mark.parametrize("quant_dtypes", QUANT_DTYPES)

456 @pytest.mark.parametrize("batch_size", BATCH_SIZE)

vllm/v1/attention/backends/flashinfer.py

+45 -7

@@ -450,9 +450,19 @@ class FlashInferBackend(AttentionBackend):

450

451 @classmethod

452 def supports_device_cpu_query_lens_mismatch(cls) -> bool:

453 - # The wrappers are planned from qo_indptr_cpu, so the CPU query offsets

454 - # have to be the ones the kernel runs on.

455 - return False

453 + # Only trtllm-gen decode (SM100 family) reads query offsets from the

454 + # device tensor; XQA and wrapper decode plan from qo_indptr_cpu, and

455 + # under DCP spec queries are routed to the prefill path.

456 + if not current_platform.is_device_capability_family(100):

457 + return False

458 + vllm_config = get_current_vllm_config_or_none()

459 + if vllm_config is None or vllm_config.model_config is None:

460 + return True

461 + pc = vllm_config.parallel_config

462 + mc = vllm_config.model_config

463 + return pc.decode_context_parallel_size == 1 and can_use_trtllm_attention(

464 + mc.get_num_attention_heads(pc), mc.get_num_kv_heads(pc), is_prefill=False

465 + )

466

467 @classmethod

468 def supports_sliding_window(cls) -> bool:

@@ -834,6 +844,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

844 current_platform.is_device_capability_family(120)

845 and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA

846 )

847 + # Adaptive verification trims drafts on device, so decode query lengths

848 + # must come from the device qo_indptr; only trtllm-gen supports that

849 + # (the selector already rejects the other configurations).

850 + self.use_trtllm_gen_varlen_decode = (

851 + speculative_config is not None

852 + and speculative_config.enable_adaptive_verification

853 + and self.flashinfer_trtllm_api_decode_kernel

854 + == FlashInferDecodeKernel.TRTLLM_GEN

855 + and not self.use_dcp

856 + )

857 supports_spec_as_decode = (

858 self.flashinfer_trtllm_api_decode_kernel

859 == FlashInferDecodeKernel.TRTLLM_GEN

@@ -1619,7 +1639,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1639 if num_decodes > 0:

1640 if decode_with_flashinfer_trtllm_api:

1641 assert self.flashinfer_trtllm_api_decode_kernel is not None

1622 - if not self.use_dedicated_xqa:

1642 + if not self.use_dedicated_xqa and not self.use_trtllm_gen_varlen_decode:

1643 assert num_decode_tokens % num_decodes == 0, (

1644 "XQA/trtllm-gen decode requires uniform query lengths "

1645 f"per request. Got {num_decode_tokens=} and {num_decodes=}."

@@ -1648,6 +1668,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1668 num_decodes,

1669 bool(causal),

1670 )

1671 + elif self.use_trtllm_gen_varlen_decode:

1672 + # CPU lens are only an upper bound, device qo_indptr is

1673 + # the truth.

1674 + decode_q_lens = (

1675 + qo_indptr_cpu[1 : num_decodes + 1] - qo_indptr_cpu[:num_decodes]

1676 + )

1677 + max_q_upper = int(decode_q_lens.max())

1678 + if max_q_upper > 1:

1679 + q_len_per_req = max_q_upper

1680 + q_cu_seq_lens = qo_indptr[: num_decodes + 1]

1681 attn_metadata.decode = FlashInferTrtllmAPIDecode(

1682 kernel=self.flashinfer_trtllm_api_decode_kernel,

1683 block_tables=block_table_tensor[:num_decodes],

@@ -2382,7 +2412,7 @@ class FlashInferImpl(AttentionImpl):

2412 bmm1_scale = self.get_xqa_bmm1_scale(

2413 layer, attn_metadata.q_data_type_decode

2414 )

2385 - q_len_per_req = attn_metadata.decode.q_len_per_req

2415 + q_len_per_req: int | None = attn_metadata.decode.q_len_per_req

2416

2417 flashinfer_xqa_batch_decode_with_kv_cache(

2418 query=decode_query,

@@ -2421,14 +2451,20 @@ class FlashInferImpl(AttentionImpl):

2451 if needs_fp8_out:

2452 out = self._nvfp4_fp8_out[:num_decode_tokens]

2453

2424 - if num_decode_tokens % attn_metadata.num_decodes != 0:

2454 + q_cu_seq_lens = attn_metadata.decode.q_cu_seq_lens

2455 + max_q_len = None

2456 + q_len_per_req = None

2457 + if q_cu_seq_lens is not None:

2458 + # q_len_per_req must stay None for the kernel's varlen path.

2459 + max_q_len = attn_metadata.decode.q_len_per_req

2460 + elif num_decode_tokens % attn_metadata.num_decodes != 0:

2461 # This gets triggered when the dummy_run forces

2462 # attention to be initialized with q_len = 0

2463 q_len_per_req = 1

2464 else:

2465 q_len_per_req = num_decode_tokens // attn_metadata.num_decodes

2466

2431 - if decode_with_xqa and q_len_per_req > 1:

2467 + if decode_with_xqa and q_len_per_req is not None and q_len_per_req > 1:

2468 raise NotImplementedError(

2469 "FlashInfer XQA speculative decode is not wired in vLLM yet."

2470 )

@@ -2472,6 +2508,8 @@ class FlashInferImpl(AttentionImpl):

2508 kv_layout=get_flashinfer_layout_string(self.kv_cache_layout),

2509 backend=attn_metadata.decode.kernel.value,

2510 q_len_per_req=q_len_per_req,

2511 + max_q_len=max_q_len,

2512 + cum_seq_lens_q=q_cu_seq_lens,

2513 kv_cache_sf=(

2514 nvfp4_kv_block_scales if self.is_kvcache_nvfp4 else None

2515 ),