[DCP] Support Decode Context Parallel (DCP) for GQA with Flashinfer #25438
Conversation
gjc0824
requested review from
DarkLight1337,
ProExpertProg,
WoosukKwon,
hmellor,
houseroad,
mgoin,
robertgshaw2-redhat,
simon-mo,
tlrmchlsmth,
yewentao256,
youkaichao and
ywang96
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces Decode Context Parallel (DCP) support for Grouped-Query Attention (GQA) with the FlashInfer backend, which is a valuable enhancement for distributed inference performance. The changes are comprehensive, covering configuration validation, modifications to the attention backend to support DCP-specific logic like query head gathering and LSE-based output correction, and the implementation of a custom attention mask for prefills. The addition of tests for a GQA model using the new functionality is also a great inclusion. The overall implementation is well-executed. I have a couple of suggestions to enhance code quality by addressing a dynamically assigned attribute and removing duplicated code.
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run You ask your reviewers to trigger select CI tests on top of Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. 🚀 |
gjc0824
force-pushed
the
dcp-gqa-flashinfer
branch
9 times, most recently
from
540c862 to
b9e9b41
Compare
| continue | ||
| K = ((rightmost - r) // p) + 1 | ||
| j = torch.arange(K) | ||
| t = torch.arange(Q) |
There was a problem hiding this comment.
nit: we generally avoid single character variable names; theyre ok though if there is supporting comment, can you please add comments explaining what the mask looks like and how it is constructed?
There was a problem hiding this comment.
Thank you for your review. We have added the comment about mask examples and algorithm explanation after vectorized improvements.
| torch.int64).tolist() | ||
| r = self.dcp_rank | ||
| p = self.dcp_world_size | ||
| for i in range(num_prefills): |
There was a problem hiding this comment.
nit: is there a way we can vectorize this loop or replace it with a triton kernel? ideally we avoid python loops as they can be very slow and create GPU bubbles
There was a problem hiding this comment.
Thank you for your valuable review. We have vectorized the "num_prefills" loop to avoid GPU bubbles. Looking forward to your further review.
There was a problem hiding this comment.
if self.dcp_world_size > 1:
# init custom mask for interleave kv cache
# |-------total_lens----------|
# |--context_lens--|--q_lens--|
# Example: dcp_size=2, dcp_rank=0
# For a SINGLE prefill seq, q_lens=3, total_lens=5
# k_lens on RANK1 is (5 - 1 - 0) // 2 + 1 = 3
# mask.shape = [q_lens, k_lens] = [3,3]
# mask [[True, True, False],
# [True, True, False],
# [True, True, True]]
dcp_rank = self.dcp_rank
dcp_size = self.dcp_world_size
q_lens = (qo_indptr_cpu[1:] - qo_indptr_cpu[:-1]).to(
dtype=torch.int64, device=self.device)
total_lens = seq_lens_cpu[prefill_start:prefill_start +
num_prefills].to(dtype=torch.int64,
device=self.device)
context_lens = total_lens - q_lens
# max indices for global sequences
max_indices = total_lens - 1
# if max_indices are smaller than dcp_rank,
# current rank has no kv cache, is invalid,
# the mask is skipped
valid = (max_indices >= dcp_rank)
assert torch.any(valid), "There is no valid sequence"
# local kv lens on current dcp_rank
k_lens = torch.div(max_indices - dcp_rank,
dcp_size,
rounding_mode="floor") + 1
k_lens = torch.where(
valid,
k_lens,
torch.zeros_like(k_lens))
# vectorize operation
# obtain the max length of all prefill reqs
max_q = int(q_lens[valid].max().item())
max_k = int(k_lens[valid].max().item())
# generate local q and k indices
q_indices = torch.arange(max_q, device=self.device)
k_indices = torch.arange(max_k, device=self.device)
# valid q and k indices of each reqs
valid_q = valid[:, None] & \
(q_indices[None, :] < q_lens[:, None])
valid_k = valid[:, None] & \
(k_indices[None, :] < k_lens[:, None])
# where global q_indices >= global k_indices,
# the mask is True
# global q_indices = context_lens + local q_indices
# global k_indices = local k_indcies * dcp_size + dcp_rank
# ====> local k_indcies must be smaller or equal k_upper
# k_upper=(context_lens + local q_indices - dcp_rank) // dcp_size
k_upper = torch.div(
context_lens[:, None] + q_indices - dcp_rank,
dcp_size, rounding_mode="floor")
k_upper = torch.where(
valid_q,
torch.clamp(k_upper, min=-1),
k_upper.new_full(k_upper.shape, -1))
mask = (k_indices[None, None, :] <= k_upper[:, :, None]) \
& (k_upper[:, :, None] >= 0)
valid_positions = valid_q[:, :, None] & valid_k[:, None, :]
# flashinfer backend needs flattened format
custom_mask = torch.masked_select(mask, valid_positions)
|
Apologies for the delayed review! left a couple nits; overall its looking pretty good though |
|
This pull request has merge conflicts that must be resolved before it can be |
When handling contexts for chunked prefill, we split contexts into chunks based on the workspace. However, the recent refactoring of the |
gjc0824
force-pushed
the
dcp-gqa-flashinfer
branch
from
9111f48 to
433340e
Compare
Signed-off-by: Jingchun Gao <[email protected]>
gjc0824
force-pushed
the
dcp-gqa-flashinfer
branch
from
433340e to
c172a1a
Compare
github-project-automation
bot
moved this from In progress
to Ready
in gpt-oss Issues & Enhancements
|
This pull request has merge conflicts that must be resolved before it can be |
LucasWilkinson
removed
documentation
Signed-off-by: QiuChunshuo <[email protected]>
pisceskkk
force-pushed
the
dcp-gqa-flashinfer
branch
from
0b360f5 to
25ece95
Compare
Signed-off-by: QiuChunshuo <[email protected]>
4 participants
Purpose
This PR adds Decode Context Parallel (DCP) support for GQA follwing PR #23734 and PR #24864. Current implementation based on FlashInfer Attention.
FlashInfer inserts the current query KV into the cache before computation. Each query then attends to both its own KV and the context KV on the local device, with LSE applied to correct the attention outputs.
Test Plan
Qwen/Qwen3-235B-A22B
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.