feat(dcp, flashinfer): support FULL_DECODE_ONLY for flashinfer with DCP>1 by pisceskkk · Pull Request #36503 · vllm-project/vllm

pisceskkk · 2026-03-09T13:37:23Z

Purpose

Allocate persistent buffers for DCP-related variables to enable FULL_DECODE_ONLY support.

Test Plan

bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m Qwen/Qwen2.5-1.5B-Instruct/ -b "auto" -l 1000 -f 5 -t 4

Test Result

TP4DCP1

vllm ({'pretrained': 'Qwen/Qwen2.5-1.5B-Instruct', 'tensor_parallel_size': 4, 'add_bos_token': True, 'max_model_len': 4096, 'decode_context_parallel_size': 1, 'attention_backend': 'FLASHINFER'}), gen_kwargs: ({}), limit: 1000.0, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.606|±  |0.0155|
|     |       |strict-match    |     5|exact_match|↑  |0.552|±  |0.0157|

TP4DCP2(4ba38c8)

vllm ({'pretrained': '/data/weight/Qwen2.5-1.5B-Instruct', 'tensor_parallel_size': 4, 'decode_context_parallel_size': 2, 'prefill_context_parallel_size': 1, 'add_bos_token': True, 'max_model_len': 4096, 'attention_backend': 'FLASHINFER'}), gen_kwargs: ({}), limit: None, num_fewshot: 5, batch_size: 512
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|_  |0.5914|_  |0.0135|
|     |       |strict-match    |     5|exact_match|_  |0.5368|_  |0.0137|

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

gemini-code-assist

Code Review

This pull request introduces support for FULL_DECODE_ONLY in FlashInfer with DCP by allocating persistent buffers for DCP-related variables. The changes involve initializing these buffers in the FlashInferBackend and using them in the forward pass to avoid on-the-fly memory allocations. While the buffer management logic appears sound, I've identified a critical issue in vllm/v1/attention/backends/flashinfer.py where a variable is calculated and then immediately overwritten with a different, simpler logic. This is not only confusing and inefficient but also potentially a bug as the new calculation ignores DCP-specific parameters, which could lead to incorrect behavior. I've recommended refactoring this part to make the logic explicit and correct.

ehfd · 2026-03-12T16:07:37Z

@cjackal This is the fix for DCP on Qwen3 and Qwen3.5, could you test together with #36480? (Cherry-picking both)

cjackal · 2026-03-12T18:36:32Z

@cjackal This is the fix for DCP on Qwen3 and Qwen3.5, could you test together with #36480? (Cherry-picking both)

I think #36070 (flashattn counterpart of this PR) is relevant as well for SM90. It seems to work for short inputs (like gsm8k), but for long inputs I still see infinite repetition / gibberish.

pisceskkk · 2026-03-16T07:23:54Z

@cjackal This is the fix for DCP on Qwen3 and Qwen3.5, could you test together with #36480? (Cherry-picking both)

I think #36070 (flashattn counterpart of this PR) is relevant as well for SM90. It seems to work for short inputs (like gsm8k), but for long inputs I still see infinite repetition / gibberish.

Sorry for the late reply. I initially suspected there might be an issue with chunked prefill, but testing with chunked prefill didn’t show noticeable drop in accuracy. Could you provide a reproducible script and output samples?

test results with long_prefill_token_threshold=49:

vllm ({'pretrained': 'qwen/Qwen2.5-1.5B-Instruct', 'tensor_parallel_size': 4, 'decode_context_parallel_size': 2, 'long_prefill_token_threshold': 49, 'add_bos_token': True, 'max_model_len': 4096, 'attention_backend': 'FLASHINFER'}), gen_kwargs: ({}), limit: None, num_fewshot: 5, batch_size: 512
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5762|±  |0.0136|
|     |       |strict-match    |     5|exact_match|↑  |0.5527|±  |0.0137|

There is indeed a slight precision drop, but it appears to be an edge case. Looking into it.

cjackal · 2026-03-16T15:01:30Z

Sorry for the late reply. I initially suspected there might be an issue with chunked prefill, but testing with chunked prefill didn’t show noticeable drop in accuracy. Could you provide a reproducible script and output samples?

I think there was a miscommunication - what I've mentioned before is for FA3 as my main computing pool is H100 and the default attention backend for H100 is FA3. I'd like to check the FI backend on H100 in any case, it would be great if DCP is back to functioning well with FI. Will come back with some benchmark results.

pisceskkk · 2026-03-17T11:03:05Z

fix some bugs and now the accuracy with 'fully' chunked prefill looks better:

vllm ({'pretrained': 'qwen/Qwen2.5-1.5B-Instruct', 'tensor_parallel_size': 4, 'decode_context_parallel_size': 2, 'long_prefill_token_threshold': 49, 'add_bos_token': True, 'max_model_len': 4096, 'attention_backend': 'FLASHINFER'}), gen_kwargs: ({}), limit: None, num_fewshot: 5, batch_size: 512
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5951|±  |0.0135|
|     |       |strict-match    |     5|exact_match|↑  |0.5421|±  |0.0137|

BTW, I noticed FlashInfer's LSE uses base-2, but merge_attn_states always sticks to natural logs. I put in a temporary scaling fix in f10579e.
@LucasWilkinson Do you think it's worth adding a base selection or a scaling factor to merge_attn_states like 5931320
?

LucasWilkinson

overall looks pretty good to me! thanks for the contribution! i have a couple questions

LucasWilkinson · 2026-04-02T04:01:27Z

        if self.use_dcp:
            assert seq_lens_cpu is not None
+            # Keep the shared metadata immutable for mixed decode/prefill batches.
+            seq_lens_cpu = seq_lens_cpu.clone()


why does it need to be cloned, i see it says it needs to be immutable but im confused where the mutation is coming from, async scheduling?

Looks like that was leftover code from debugging the accuracy issue. I just confirmed that removing it doesn't change anything. The real cause for the accuracy issue was the LSE base value.

Tasks Version Filter n-shot Metric Value Stderr

gsm8k 3 flexible-extract 5 exact_match ↑ 0.600 ± 0.0155

strict-match 5 exact_match ↑ 0.546 ± 0.0158

can we remove it then please?

I remember deleting this line of code, but it might still be there due to a rebase error. deleted

LucasWilkinson · 2026-04-02T04:02:09Z

-                        (decode_query.size(0), decode_query.size(1)),
-                        dtype=torch.float32,
-                        device=decode_query.device,
+                    decode_query_buffer.copy_(


why is this copy necessary?

These buffers need to be assigned in-place. Do you mean using the decode_query_buffer[:] = get_dcp_group(...) syntax?

i guess i dont understand why we can't use the output of get_dcp_group().all_gather directly?

Ah, yes, I had some misunderstandings about the persistence of tensors in CUDAGraph. This value indeed does not need to be persisted. It has been corrected.

Copilot

Pull request overview

Adds persistent buffer allocation to enable FULL_DECODE_ONLY when using FlashInfer with decode_context_parallel_size > 1, and adjusts LSE handling for downstream attention-state merging.

Changes:

Allocate reusable decode buffers via the workspace manager for DCP all-gathered queries / outputs / LSE.
Convert FlashInfer LSE to the expected log base for merge_attn_states.
Update DCP/CP interleave sizing and adjust how CPU seq-lens are materialized for DCP/cascade paths.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-02T07:34:22Z

 logger = init_logger(__name__)

 trtllm_gen_workspace_buffer = None
+LOG2_E_SCALE = math.log(2.0)


LOG2_E_SCALE is ambiguous/misleading: math.log(2.0) is ln(2), which is specifically a base-conversion factor (log2 -> ln). Consider renaming to something unambiguous like LN_2 or LOG2_TO_LN to make the intent (and direction) clear and avoid future mistakes.

Copilot · 2026-04-02T07:34:23Z

            self.dcp_world_size = 1
            self.dcp_rank = 0
-            self.dcp_kv_cache_interleave_size = 1
+            self.cp_kv_cache_interleave_size = 1


The naming switches from dcp_* to cp_* while the surrounding fields remain dcp_world_size, dcp_rank, use_dcp, etc. This inconsistency makes it harder to reason about whether the interleave size is tied to DCP specifically or a broader CP concept. Either align the field name with the DCP terminology used in this class, or add a short clarifying comment explaining why the config uses cp_* while the runtime concept is DCP.

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

LucasWilkinson · 2026-05-10T04:18:24Z

-            if seq_lens_np is not None
-            else None
+        seq_lens_cpu = (
+            common_attn_metadata.seq_lens.cpu() if needs_seq_lens_cpu else None


why is this change needed? if its just refactoring I think we should break that into a separate PR

I moved seq_lens_np to the bottom because the code snippet below might modify seq_lens_cpu, so I need to assign a value to seq_lens_np after that. However, this modification really shouldn't be included in this PR. I'll open a new PR to complete this modification.

Open a new PR for this modifcation: #42205

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

mergify · 2026-05-10T05:12:34Z

Hi @pisceskkk, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

mergify · 2026-05-10T05:20:42Z

Hi @pisceskkk, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

LucasWilkinson · 2026-05-12T03:42:58Z

        if self.sinks is not None and self.sinks.dtype != torch.float32:
            self.sinks = self.sinks.to(torch.float32)

+    def _get_dcp_decode_buffers(


can we just inline this?

yes, have changed.

LucasWilkinson · 2026-05-12T03:52:27Z

+        num_decode_tokens: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert self.dcp_world_size > 1
+        total_num_heads = self.num_heads * self.dcp_world_size


is this always true? what about say TP=8 and DCP=2; our is this not truely total_num_heads just total_num_heads within DCP group?

I think this is correct, right? The output should always maintain the same shape as decode_query.

https://github.com/pisceskkk/vllm/blob/13672b4b6fa59535377528338947a6ee9395c575/vllm/v1/attention/backends/flashinfer.py#L1742-L1744

decode_query = get_dcp_group().all_gather( decode_query.contiguous(), dim=-2 )

our is this not truely total_num_heads just total_num_heads within DCP group?

Sorry, I misunderstood the meaning. This variable indeed should not be called "total_num_heads." I have changed the name to ensure it aligns with its intended meaning.

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

mergify Bot added nvidia v1 labels Mar 9, 2026

github-project-automation Bot added this to NVIDIA Mar 9, 2026

gemini-code-assist Bot reviewed Mar 9, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/flashinfer.py Outdated

pisceskkk force-pushed the dcp/fdo branch 2 times, most recently from 73b42dc to aab57d7 Compare March 9, 2026 13:50

pisceskkk marked this pull request as ready for review March 9, 2026 13:50

pisceskkk requested review from mgoin and pavanimajety as code owners March 9, 2026 13:50

pisceskkk mentioned this pull request Mar 9, 2026

bugfix(dcp, gdn): disabling DCP semantics for linear-attention KV/state groups #36480

Closed

5 tasks

LucasWilkinson reviewed Apr 2, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings April 2, 2026 07:27

pisceskkk requested a review from vadiklyutiy as a code owner April 2, 2026 07:27

Copilot AI reviewed Apr 2, 2026

View reviewed changes

Copilot started reviewing on behalf of pisceskkk April 2, 2026 07:37 View session

pisceskkk mentioned this pull request May 8, 2026

feat(nixl,dcp): Supports MLA DCP for PD disaggregation with nixl connector #38433

Open

feat(dcp, flashinfer): add persistent buffers for dcp

5814179

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

pisceskkk force-pushed the dcp/fdo branch from 4ba38c8 to 999cb02 Compare May 8, 2026 03:55

LucasWilkinson reviewed May 10, 2026

View reviewed changes

refactor(flashinfer): use workspace manager for dcp decode buffers

1213bba

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

pisceskkk force-pushed the dcp/fdo branch 2 times, most recently from fc81fb2 to 7b65dea Compare May 10, 2026 05:05

revert unnecessary modification

39e1c0d

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

pisceskkk force-pushed the dcp/fdo branch from 7b65dea to 39e1c0d Compare May 10, 2026 05:05

lint

13672b4

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

LucasWilkinson reviewed May 12, 2026

View reviewed changes

pisceskkk added 2 commits May 12, 2026 11:54

refactor: inline _get_dcp_decode_buffers

a9ab052

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

fix: correct the var name of total_num_heads within DCP group

476378f

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

GirasoleY mentioned this pull request May 28, 2026

[Bugfix][MLA] Fix LSE log-base mismatch in DCP + FlashInfer MLA decode #43919

Closed

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.600	±	0.0155
		strict-match	5	exact_match	↑	0.546	±	0.0158

Uh oh!

Uh oh!

Conversation

pisceskkk commented Mar 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

ehfd commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cjackal commented Mar 12, 2026

Uh oh!

pisceskkk commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cjackal commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pisceskkk commented Mar 17, 2026

Uh oh!

LucasWilkinson left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI Apr 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented May 10, 2026

Uh oh!

mergify Bot commented May 10, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pisceskkk commented Mar 9, 2026 •

edited

Loading

ehfd commented Mar 12, 2026 •

edited

Loading

pisceskkk commented Mar 16, 2026 •

edited

Loading

cjackal commented Mar 16, 2026 •

edited

Loading