Skip to content

[0.13.0][cherry-pick][bugfix] fix the complex and potentially problematic generate_kv_idx. #5955

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yiz-liu merged 1 commit into from
Jan 21, 2026
Merged

[0.13.0][cherry-pick][bugfix] fix the complex and potentially problematic generate_kv_idx. #5955

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions vllm_ascend/attention/context_parallel/attention_cp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,10 +161,13 @@ def build(
(len(local_context_lens_allranks)),
dtype=torch.int32,
device=self.device)
cp_kv_recover_idx_for_chunk = common_long_seq_metadata.cp_kv_recover_idx_for_chunk
kv_inverse_idx_for_chunk = torch.argsort(
cp_kv_recover_idx_for_chunk.to(torch.float32)
) if cp_kv_recover_idx_for_chunk is not None else None
common_long_seq_metadata.
pcp_allgather_restore_idx[pcp_size *
num_decode_tokens:].to(
torch.float32))
cp_kv_recover_idx_for_chunk = torch.argsort(
kv_inverse_idx_for_chunk)

batch_chunk_seq_mask = (
local_context_lens_allranks[:, self.pcp_rank,
Expand Down
2 changes: 0 additions & 2 deletions vllm_ascend/attention/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,6 @@ def enable_cp():
class AscendPrefillContextParallelMetadata:
pcp_allgather_restore_idx: torch.Tensor = None

cp_kv_recover_idx_for_chunk: torch.Tensor = None

num_actual_tokens_pcp_padded: int = 0

num_computed_tokens_of_pcp_dcp: Optional[list[list[list[int]]]] = None
Expand Down
3 changes: 0 additions & 3 deletions vllm_ascend/worker/model_runner_v1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -553,9 +553,6 @@ def _prepare_inputs(
self.num_spec_tokens)

if self.pcp_size > 1:
if not self.vllm_config.model_config.use_mla:
self.pcp_manager.generate_kv_idx(scheduler_output,
self.input_batch)
num_scheduled_tokens[:
num_reqs], position_pcp = self.pcp_manager.update_tokens_for_pcp(
num_scheduled_tokens[:num_reqs],
Expand Down
48 changes: 1 addition & 47 deletions vllm_ascend/worker/pcp_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,11 @@
# Adapted from vllm-project/vllm/vllm/worker/worker.py
#

from typing import TYPE_CHECKING, List
from typing import TYPE_CHECKING

import numpy as np
import torch
from vllm.config import VllmConfig
from vllm.utils.math_utils import cdiv
from vllm.v1.utils import CpuGpuBuffer

if TYPE_CHECKING:
Expand Down Expand Up @@ -86,9 +85,6 @@ def __init__(
)
self.num_actual_tokens_pcp_padded = 0
self.pcp_unpad_mask_cpu = self.pcp_unpad_mask_cpu_tensor.numpy()
self.cp_kv_recover_idx_for_chunk: List[List[int]] = [
[] for _ in range(self.pcp_world_size)
]
self.full_indices = list(
range(self.max_num_tokens * self.pcp_world_size *
self.dcp_world_size + self.pcp_world_size *
Expand Down Expand Up @@ -563,47 +559,6 @@ def _get_cp_local_seq_lens(
[-1, pcp_world_size, dcp_world_size])
return dcp_local_seq_lens

def generate_kv_idx(self, scheduler_output, input_batch):
if not self.pcp_world_size > 1:
return
self.cp_kv_recover_idx_for_chunk = [[]
for _ in range(self.pcp_world_size)
]

for i, req_id in enumerate(input_batch.req_ids):
num_scheduled_token = scheduler_output.num_scheduled_tokens[req_id]
is_prefill = num_scheduled_token > self.decode_threshold
if is_prefill:
num_cp_padded_scheduled_tokens = cdiv(
num_scheduled_token,
2 * self.pcp_world_size) * (2 * self.pcp_world_size)
chunk_size = num_cp_padded_scheduled_tokens // (
2 * self.pcp_world_size)
num_added_recover_tokens = len(
self.cp_kv_recover_idx_for_chunk[0]) * self.pcp_world_size
for rank in range(self.pcp_world_size):
self.cp_kv_recover_idx_for_chunk[rank].extend(
self.full_indices[rank * chunk_size +
num_added_recover_tokens:(rank + 1) *
chunk_size +
num_added_recover_tokens])
self.cp_kv_recover_idx_for_chunk[rank].extend(
self.full_indices[num_cp_padded_scheduled_tokens -
(rank + 1) * chunk_size +
num_added_recover_tokens:
num_cp_padded_scheduled_tokens -
rank * chunk_size +
num_added_recover_tokens])

cp_kv_recover_idx_for_chunk = torch.from_numpy(
np.concatenate(
self.cp_kv_recover_idx_for_chunk)).to(device=self.device)
cp_kv_recover_idx_for_chunk.copy_(torch.tensor(
np.array(self.cp_kv_recover_idx_for_chunk).flatten().tolist()),
non_blocking=True)
self.cp_kv_recover_idx_for_chunk = cp_kv_recover_idx_for_chunk.to(
torch.float32).argsort().to(torch.int32)

def generate_pcp_metadata(self, total_num_scheduled_tokens, query_lens,
input_batch, num_scheduled_tokens):
from vllm_ascend.attention.utils import \
Expand Down Expand Up @@ -774,7 +729,6 @@ def generate_pcp_metadata(self, total_num_scheduled_tokens, query_lens,
}
long_seq_metadata.pcp_allgather_restore_idx = self.pcp_allgather_restore_idx.gpu[:
num_actual_tokens_pcp_padded]
long_seq_metadata.cp_kv_recover_idx_for_chunk = self.cp_kv_recover_idx_for_chunk
long_seq_metadata.q_head_idx_tensor = self.q_head_idx_tensor
long_seq_metadata.q_tail_idx_tensor = self.q_tail_idx_tensor
long_seq_metadata.q_full_idx = self.q_full_idx
Expand Down