TRTLLM gen-full attn spec decode & FP8 KV dequant tests

[ojhaanshika]

Purpose

Add two new integration tests for TRTLLM gen-full attention covering previously untested code paths in FlashInferImpl.forward():

1. Speculative decode (multi-token decode): Exercises the q_len_per_req > 1 path (line 1583 of flashinfer.py) by injecting a SpeculativeConfig that sets reorder_batch_threshold=4, causing split_decodes_and_prefills to classify multi-token requests as decode. This path had zero test coverage at any level.

2. FP8 KV cache with bf16 queries (prefill dequant): Exercises the trtllm_prefill_attn_kvfp8_dequant fallback path (lines 1464-1478 of flashinfer.py) where FP8 KV pages are dequantized to bf16 before the TRTLLM prefill kernel. The kernel-level tests explicitly skip mixed Q/KV dtypes for prefill, so this path also had zero coverage.

Combined TRTLLM-relevant coverage increased

File	Before	After	Change
flashinfer.py	356/641 (56%)	375/641 (59%)	+19 lines
utils/flashinfer.py	110/252 (44%)	137/252 (54%)	+27 lines

Note: This PR depends on #<34986> which adds the base integration test file (test_trtllm_attention_integration.py). Both PRs modify the same file — this one appends two new test functions.

Test Plan

pytest tests/v1/attention/test_trtllm_attention_integration.py -v

Requires Blackwell (SM100) GPU.

Test Result

tests/v1/attention/test_trtllm_attention_integration.py::test_trtllm_gen_full_attention_integration[decode_only] PASSED [ 20%]
tests/v1/attention/test_trtllm_attention_integration.py::test_trtllm_gen_full_attention_integration[prefill_only] PASSED [ 40%]
tests/v1/attention/test_trtllm_attention_integration.py::test_trtllm_gen_full_attention_integration[mixed] PASSED [ 60%]
tests/v1/attention/test_trtllm_attention_integration.py::test_trtllm_spec_decode_integration PASSED [ 80%]
tests/v1/attention/test_trtllm_attention_integration.py::test_trtllm_fp8_kv_dequant_integration PASSED [100%]
5 passed, 20 warnings in 23.32s

[gemini-code-assist]

Code Review

This pull request adds two new integration tests for TRTLLM gen-full attention, covering speculative decoding and FP8 KV cache dequantization. The tests are comprehensive and correctly use a reference implementation for verification. My main feedback is on improving the maintainability of the new test file by refactoring duplicated code.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ojhaanshika.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[jasonlizhengjian]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds two valuable integration tests for TRTLLM gen-full attention, covering speculative decoding and FP8 KV cache dequantization, which previously had no test coverage. The tests are well-structured and increase coverage as intended. My review identifies a significant opportunity for improvement by refactoring duplicated code between the two new test functions to enhance maintainability.

[ojhaanshika]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces two new integration tests for the TRTLLM attention backend in FlashInferImpl. The first test covers speculative decoding with multi-token decode requests, exercising a previously untested code path. The second test covers the scenario of an FP8 KV cache with bfloat16 queries, which triggers the FP8 dequantization path for prefill operations. The changes also include refactoring existing test code to improve modularity and reuse by extracting common setup logic into helper functions. The new tests are well-structured and correctly set up the specific conditions to target the new code paths, comparing the results against a reference SDPA implementation to ensure correctness. My review did not find any critical or high-severity issues.