feat: enable trtllm-gen attn speculative decoding verify by decode #1453

yyihuang · 2025-08-11T03:31:15Z

📌 Description

decode with q_len > 1

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

gemini-code-assist

Summary of Changes

Hello @yyihuang, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces support for multi-token prediction (MTP) within the trtllm-gen attention mechanism. My changes primarily involve modifying the expected shape of the query tensor to accommodate multiple tokens per request, moving away from a single-token assumption. This enables more efficient speculative decoding by allowing the attention operation to process several tokens simultaneously. I've also updated the relevant test cases to validate this new functionality across various query lengths.

Highlights

Core Logic Change for Multi-Token Prediction: I've updated the _paged_run and trtllm_batch_decode_with_kv_cache functions in flashinfer/decode.py to remove the unsqueeze(1) call on the query tensor. This change allows the query tensor to accept a multi-token dimension, moving from a [B, 1, H, D] shape to [B, S, H, D], where S represents the sequence length or number of tokens per request. This is crucial for enabling multi-token prediction (MTP) within the trtllm-gen attention mechanism.
Docstring Update for Query Tensor Shape: To support the new multi-token query shape, I've updated the docstring for the query parameter in trtllm_batch_decode_with_kv_cache to clearly state its new expected dimensions: [batch_size, q_len_per_request, num_heads, head_dim]. This ensures clarity and correctness for future usage and development.
Expanded Test Coverage for Multi-Token Queries: I've modified tests/test_trtllm_gen_decode.py to include q_len_per_request as a parameter in test_trtllm_batch_decode_fmha. This allows the test suite to validate the trtllm-gen attention mechanism with varying numbers of tokens per request, ensuring the new multi-token prediction capabilities work as expected across different scenarios.
Test Adaptation for Multi-Token Output: In the test test_trtllm_batch_decode_fmha, I've adjusted the q tensor creation and removed the output.squeeze(1) operation. These changes align the test's data handling with the new multi-token query tensor shape, ensuring that the test correctly processes and verifies outputs from the multi-token prediction enabled attention.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request enables speculative decoding for trtllm-gen attention by allowing the query tensor to have a sequence length dimension. The changes in flashinfer/decode.py are consistent with this goal. However, the corresponding test file tests/test_trtllm_gen_decode.py has a critical issue where the validation logic is disabled, making the test ineffective. Additionally, some debugging code has been left in the test file which should be removed.

gemini-code-assist · 2025-08-11T03:32:46Z

tests/test_trtllm_gen_decode.py

+    # output_ref = wrapper.run(ref_q, ref_kv_cache) # todo(Yingyi): fix mtp here
+    # tmp
+    output_ref = output


The reference implementation call to wrapper.run is commented out, and output_ref is assigned the value of output. This effectively disables the test's validation logic, as it compares the output against itself. The test will pass as long as the function doesn't crash, but it doesn't verify the correctness of the computation. This should be fixed before merging. The todo comment indicates this is a known issue that needs to be addressed.

gemini-code-assist · 2025-08-11T03:32:46Z

tests/test_trtllm_gen_decode.py

+if __name__ == "__main__":
+    test_trtllm_batch_decode_fmha(
+        kv_layout="HND",
+        batch_size=4,
+        q_len_per_request=2,
+        page_size=32,
+        num_kv_heads=2,
+        head_grp_size=1,
+        window_left=-1,
+        q_dtype="half",
+        o_dtype="half",
+        kv_cache_dtype="half",
+    )


This if __name__ == "__main__": block appears to be temporary debugging code. It should be removed before merging to keep the test file clean.

…m-attn-mtp

yzh119 · 2025-08-31T05:52:04Z

tests/test_trtllm_gen_attention.py

@@ -196,6 +204,51 @@ def get_last_page_len(seq_lens, page_size):
    return kv_last_page_len


+def assert_close_with_mismatch_tolerance(


consider moving this to conftests

yzh119

Would be great if you can work on the remaining item I mentioned in later PRs.

yzh119 · 2025-08-31T07:52:29Z

tests/test_trtllm_gen_attention.py

-                output.float(), output_wrapper.float(), rtol=1e-1, atol=1e-1
-            )
+            # todo(Yingyi): fix precision issue with this test
+            if not (


I still have deep concern about the special handling of precision here, a more fundamental solution could be a PrecisionManager to centrialize the handling of tolerance for different data types.

Also, let's add a flashinfer.testing.assert_allclose() function where there could be arguments such as max mismatched elements, it will be fallback to torch.testing.assert_allclose when mismatched elements is not specified.

yyihuang · 2025-08-31T23:14:50Z

should unblock sgl-project/sglang#9471 now

init

12945f3

gemini-code-assist bot reviewed Aug 11, 2025

View reviewed changes

yyihuang self-assigned this Aug 11, 2025

gemini-code-assist bot reviewed Aug 11, 2025

View reviewed changes

yyihuang assigned yzh119 Aug 11, 2025

yyihuang added 11 commits August 11, 2025 01:35

upd workspace

e1be6a8

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

99db15a

…m-attn-mtp

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

76dd13a

…m-attn-mtp

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

b9d59ad

…m-attn-mtp

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

6c3e713

…m-attn-mtp

stash

14079df

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

3701205

…m-attn-mtp

stash

33712dc

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

4e3e5c9

…m-attn-mtp

revert stash

2560dab

upd

f6ad3b1

zhyncs self-assigned this Aug 18, 2025

zhyncs added the priority: high label Aug 18, 2025

zhyncs marked this pull request as ready for review August 19, 2025 02:49

yyihuang and others added 7 commits August 19, 2025 01:01

upd

a83029e

upd

4830e47

ruff

fe90f24

fix

b15c72e

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

2a76af2

…m-attn-mtp

stash

45dd56e

Merge branch 'main' of github.com:flashinfer-ai/flashinfer into trtll…

53a6855

…m-attn-mtp

yyihuang marked this pull request as draft August 30, 2025 22:13

yyihuang added 3 commits August 30, 2025 19:01

fix

7952789

cleanup

8ac93a6

upd

7e3039a

yyihuang added 2 commits August 30, 2025 19:20

upd

3089139

fix

c6d462a

yyihuang marked this pull request as ready for review August 31, 2025 03:22

yyihuang requested a review from yzh119 August 31, 2025 03:22

yzh119 reviewed Aug 31, 2025

View reviewed changes

upd

57f8338

yyihuang requested a review from yzh119 August 31, 2025 05:59

yzh119 approved these changes Aug 31, 2025

View reviewed changes

yzh119 merged commit efe55ee into flashinfer-ai:main Aug 31, 2025
2 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: enable trtllm-gen attn speculative decoding verify by decode #1453

feat: enable trtllm-gen attn speculative decoding verify by decode #1453

yyihuang commented Aug 11, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Aug 11, 2025

Uh oh!

gemini-code-assist bot Aug 11, 2025

Uh oh!

yzh119 Aug 31, 2025

Uh oh!

yyihuang Aug 31, 2025

Uh oh!

yzh119 left a comment

Uh oh!

yzh119 Aug 31, 2025

Uh oh!

Uh oh!

yyihuang commented Aug 31, 2025 •

edited

Loading

Uh oh!

Uh oh!

		@@ -196,6 +204,51 @@ def get_last_page_len(seq_lens, page_size):
		return kv_last_page_len


		def assert_close_with_mismatch_tolerance(

feat: enable trtllm-gen attn speculative decoding verify by decode #1453

feat: enable trtllm-gen attn speculative decoding verify by decode #1453

Conversation

yyihuang commented Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Aug 11, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 Aug 31, 2025

Choose a reason for hiding this comment

Uh oh!

yyihuang Aug 31, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 left a comment

Choose a reason for hiding this comment

Uh oh!

yzh119 Aug 31, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yyihuang commented Aug 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

yyihuang commented Aug 11, 2025 •

edited

Loading

yyihuang commented Aug 31, 2025 •

edited

Loading