A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 07:00 UTC
[bugfix] correct local_chunk_len for DCP in reorg_kvcache with long context (#28526)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Qiu committed
Nov 14, 2025 at 03:29 UTC
968060c15adc0b68a76d37db00acf1273a23b829
1 file changed
+25
-4
vllm/v1/attention/backends/mla/common.py
+25
-4
@@ -337,6 +337,7 @@ class MLACommonPrefillMetadata:
337
local_context_lens_allranks: list[list[int]] | None = None
338
padded_local_cu_seq_lens: torch.Tensor | None = None
339
cu_seq_lens_lst: list[list[int]] | None = None
340
+ chunk_size: int | None = None
341
342
block_table: torch.Tensor
343
query_start_loc: torch.Tensor
@@ -902,6 +903,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
903
device, non_blocking=True
904
),
905
cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
906
+ chunk_size=padded_local_max_context_chunk_across_ranks,
907
)
908
else:
909
chunked_context_metadata = chunked_context_metadata_cls(
@@ -986,6 +988,8 @@ def reorg_kvcache(
988
local_context_lens_allranks: list[list[int]],
989
sum_seq_len: int,
990
max_seq_len: int,
991
+ chunk_size: int,
992
+ chunk_idx: int,
993
toks: int,
994
) -> tuple[torch.Tensor, torch.Tensor]:
995
"""
@@ -1001,6 +1005,9 @@ def reorg_kvcache(
1005
local_context_lens_allranks: local context lengths on each CP rank.
1006
sum_seq_len: the sum of cp_chunk_seq_lens_lst.
1007
max_seq_len: the max value of cp_chunk_seq_lens_lst.
1008
+ chunk_size: the local padded max context chunk from
1009
+ chunked_context_metadata building.
1010
+ chunk_idx: chunk idx of chunked_prefill.
1011
toks: the number of tokens for local gather cache.
1012
"""
1013
kv_c_segments = []
@@ -1012,20 +1019,31 @@ def reorg_kvcache(
1019
):
1020
cur_seq_len = 0
1021
for rank, local_context_len in enumerate(local_context_lens):
1015
- if local_context_len != 0:
1022
+ # Note(qcs): We split the context into multiple chunks,
1023
+ # depending on the size of the workspace.
1024
+ # local_context in dcp0: |-----------------|
1025
+ # local_context in dcp1: |--------------|
1026
+ # n*padded_local_chunk: |-----|-----|-----|
1027
+ # local_chunk_len in dcp1: |-----|-----|--|
1028
+ # so we need update the last chunk length in dcp1.
1029
+ local_chunk_len = min(
1030
+ max(0, local_context_len - chunk_idx * chunk_size),
1031
+ padded_local_chunk_seq_len,
1032
+ )
1033
+ if local_chunk_len != 0:
1034
kv_c_segment = allgatered_kv_c_normed[
1035
rank * toks + src_token_idx : rank * toks
1036
+ src_token_idx
1019
- + local_context_len
1037
+ + local_chunk_len
1038
]
1039
k_pe_segment = allgatered_k_pe[
1040
rank * toks + src_token_idx : rank * toks
1041
+ src_token_idx
1024
- + local_context_len
1042
+ + local_chunk_len
1043
]
1044
kv_c_segments.append(kv_c_segment)
1045
k_pe_segments.append(k_pe_segment)
1028
- cur_seq_len += local_context_len
1046
+ cur_seq_len += local_chunk_len
1047
max_seq_len_check = max(max_seq_len_check, cur_seq_len)
1048
src_token_idx += padded_local_chunk_seq_len
1049
reorganized_kv_c_normed = torch.cat(kv_c_segments, dim=0)
@@ -1676,6 +1694,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1694
assert prefill_metadata.chunked_context.local_context_lens_allranks is not None
1695
assert prefill_metadata.chunked_context.padded_local_cu_seq_lens is not None
1696
assert prefill_metadata.chunked_context.cu_seq_lens_lst is not None
1697
+ assert prefill_metadata.chunked_context.chunk_size is not None
1698
1699
output = None
1700
iters = len(prefill_metadata.chunked_context.seq_tot)
@@ -1725,6 +1744,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1744
local_context_lens_allranks=prefill_metadata.chunked_context.local_context_lens_allranks,
1745
sum_seq_len=prefill_metadata.chunked_context.cu_seq_lens_lst[i][-1],
1746
max_seq_len=prefill_metadata.chunked_context.max_seq_lens[i],
1747
+ chunk_size=prefill_metadata.chunked_context.chunk_size,
1748
+ chunk_idx=i,
1749
toks=toks,
1750
)
1751