Skip to content

[Performance]: Cache group-invariant DSA-CP drafting metadata across kv-cache groups #14916

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
Sep 3, 2026
Merged

[Performance]: Cache group-invariant DSA-CP drafting metadata across kv-cache groups #14916

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
221 changes: 157 additions & 64 deletions vllm_ascend/attention/context_parallel/dsa_cp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import torch_npu
from vllm.config import VllmConfig, get_current_vllm_config
from vllm.distributed import get_pcp_group, get_tp_group
from vllm.logger import logger
from vllm.triton_utils import HAS_TRITON, triton
from vllm.v1.attention.backend import AttentionCGSupport, AttentionImplBase, AttentionMetadataBuilder
from vllm.v1.kv_cache_interface import AttentionSpec
Expand Down Expand Up @@ -395,11 +396,49 @@ def build_for_drafting(
assert self.compressor_ratio <= 1, "vLLM-Ascend only support SWA-layer for Deepseek-V4 now."
num_reqs = common_attn_metadata.num_reqs
num_input_tokens = common_attn_metadata.num_input_tokens
num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
common_attn_metadata,
decode_threshold=self.decode_threshold,
treat_short_extends_as_decodes=False,
)
# Cross-kv-cache-group metadata cache. The spec-decode proposer passes
# one dict shared by all attn groups of the same decode step. Everything
# cached below only depends on group-invariant fields of
# common_attn_metadata (positions / seq_lens / query_start_loc), never
# on the per-group block_table / slot_mapping, so it is safe to reuse
# across groups. When no shared dict is provided (e.g. the per-step
# update path in attn_update_stack_num_spec_norm, where each group call
# receives an already-updated common_attn_metadata) the metadata cache
# stays empty and the behavior is unchanged.
shared_metadata_cache = kwargs.get("common_ratio_to_sas_metadata")
metadata_cache = (
shared_metadata_cache.setdefault(("draft_for_drafting", draft_index), {})
if shared_metadata_cache is not None
else None
)

if metadata_cache is not None and "num_decodes" in metadata_cache:
logger.debug("[dsa_cp] Cross-group drafting metadata cache hit for draft step %d.", draft_index)
num_decodes = metadata_cache["num_decodes"]
num_prefills = metadata_cache["num_prefills"]
num_decode_tokens = metadata_cache["num_decode_tokens"]
input_positions = metadata_cache["input_positions"]
cos = metadata_cache["cos"]
sin = metadata_cache["sin"]
else:
num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
common_attn_metadata,
decode_threshold=self.decode_threshold,
treat_short_extends_as_decodes=False,
)
input_positions = common_attn_metadata.positions[:num_input_tokens].long()
# Draft steps update positions independently. Reusing the global RoPE
# cache can let later draft steps overwrite step-0 metadata.
cos, sin = get_cos_and_sin_dsa(input_positions, use_cache=False)
if metadata_cache is not None:
metadata_cache.update(
num_decodes=num_decodes,
num_prefills=num_prefills,
num_decode_tokens=num_decode_tokens,
input_positions=input_positions,
cos=cos,
sin=sin,
)

self.num_decodes = num_decodes
self.num_prefills = num_prefills
Expand All @@ -410,12 +449,14 @@ def build_for_drafting(
self.seq_lens_cpu = common_attn_metadata._seq_lens_cpu[:num_reqs]
elif common_attn_metadata.seq_lens_cpu is not None:
self.seq_lens_cpu = common_attn_metadata.seq_lens_cpu[:num_reqs]
elif metadata_cache is not None and "seq_lens_cpu" in metadata_cache:
# Reuse the device->host copy built by a sibling kv-cache group
# instead of synchronizing again.
self.seq_lens_cpu = metadata_cache["seq_lens_cpu"]
else:
self.seq_lens_cpu = self.seq_lens.cpu()
input_positions = common_attn_metadata.positions[:num_input_tokens].long()
# Draft steps update positions independently. Reusing the global RoPE
# cache can let later draft steps overwrite step-0 metadata.
cos, sin = get_cos_and_sin_dsa(input_positions, use_cache=False)
if metadata_cache is not None:
metadata_cache["seq_lens_cpu"] = self.seq_lens_cpu

slot_mapping = common_attn_metadata.slot_mapping[:num_input_tokens]

Expand All @@ -432,6 +473,7 @@ def build_for_drafting(
num_input_tokens=num_input_tokens,
cos=cos,
sin=sin,
metadata_cache=metadata_cache,
)

return self.metadata_cls( # type: ignore
Expand Down Expand Up @@ -460,6 +502,7 @@ def build_req_metadata_for_drafting(
num_input_tokens: int,
cos: RopeDataProxy,
sin: RopeDataProxy,
metadata_cache: dict | None = None,
) -> AscendDSAReqMetadata:
"""Build DSA-CP metadata for one draft step."""
num_reqs = common_attn_metadata.num_reqs
Expand All @@ -469,39 +512,70 @@ def build_req_metadata_for_drafting(
is_noncausal = not common_attn_metadata.causal
has_prefill = self.num_prefills > 0

(
local_start,
local_end_with_pad,
tokens_per_rank,
num_tokens_pad,
local_query_start_loc,
local_seq_lens,
) = self._build_local_token_metadata(
num_reqs=num_reqs,
num_input_tokens=num_input_tokens,
query_start_loc=query_start_loc,
seq_lens=self.seq_lens[:num_reqs],
local_query_start_loc=self.spec_local_query_start_loc[draft_index - 1],
local_seq_lens=self.spec_local_seq_lens[draft_index - 1],
is_noncausal=is_noncausal,
)
local_query_start_loc = local_query_start_loc.clone()
local_seq_lens = local_seq_lens.clone()
local_cos = cos.pad_to(num_tokens_pad)[local_start:local_end_with_pad]
local_sin = sin.pad_to(num_tokens_pad)[local_start:local_end_with_pad]

_, _, _, _, local_query_start_loc_cpu, local_seq_lens_cpu = self._build_local_token_metadata(
num_reqs=num_reqs,
num_input_tokens=num_input_tokens,
query_start_loc=query_start_loc_cpu,
seq_lens=self.seq_lens_cpu[:num_reqs],
is_noncausal=is_noncausal,
)
local_seq_lens_q_cpu = local_query_start_loc_cpu[1 : num_reqs + 1] - local_query_start_loc_cpu[:num_reqs]
max_local_query_len = max(1, int(local_seq_lens_q_cpu.max().item()))
max_local_seq_lens = max(1, int(local_seq_lens_cpu.max().item()))
if metadata_cache is not None and "local_query_start_loc" in metadata_cache:
# Hit: local token metadata computed by a sibling kv-cache group
# of the same draft step. Values are identical across groups, so
# the cached (cloned) tensors are reused directly.
local_start = metadata_cache["local_start"]
local_end_with_pad = metadata_cache["local_end_with_pad"]
tokens_per_rank = metadata_cache["tokens_per_rank"]
num_tokens_pad = metadata_cache["num_tokens_pad"]
local_query_start_loc = metadata_cache["local_query_start_loc"]
local_seq_lens = metadata_cache["local_seq_lens"]
max_local_query_len = metadata_cache["max_local_query_len"]
max_local_seq_lens = metadata_cache["max_local_seq_lens"]
local_cos = metadata_cache["local_cos"]
local_sin = metadata_cache["local_sin"]
start_pos = metadata_cache["start_pos"]
else:
(
local_start,
local_end_with_pad,
tokens_per_rank,
num_tokens_pad,
local_query_start_loc,
local_seq_lens,
) = self._build_local_token_metadata(
num_reqs=num_reqs,
num_input_tokens=num_input_tokens,
query_start_loc=query_start_loc,
seq_lens=self.seq_lens[:num_reqs],
local_query_start_loc=self.spec_local_query_start_loc[draft_index - 1],
local_seq_lens=self.spec_local_seq_lens[draft_index - 1],
is_noncausal=is_noncausal,
)
local_query_start_loc = local_query_start_loc.clone()
local_seq_lens = local_seq_lens.clone()
local_cos = cos.pad_to(num_tokens_pad)[local_start:local_end_with_pad]
local_sin = sin.pad_to(num_tokens_pad)[local_start:local_end_with_pad]

start_pos = self.seq_lens[:num_reqs] - seq_lens_q
_, _, _, _, local_query_start_loc_cpu, local_seq_lens_cpu = self._build_local_token_metadata(
num_reqs=num_reqs,
num_input_tokens=num_input_tokens,
query_start_loc=query_start_loc_cpu,
seq_lens=self.seq_lens_cpu[:num_reqs],
is_noncausal=is_noncausal,
)
local_seq_lens_q_cpu = local_query_start_loc_cpu[1 : num_reqs + 1] - local_query_start_loc_cpu[:num_reqs]
max_local_query_len = max(1, int(local_seq_lens_q_cpu.max().item()))
max_local_seq_lens = max(1, int(local_seq_lens_cpu.max().item()))

start_pos = self.seq_lens[:num_reqs] - seq_lens_q

if metadata_cache is not None:
metadata_cache.update(
local_start=local_start,
local_end_with_pad=local_end_with_pad,
tokens_per_rank=tokens_per_rank,
num_tokens_pad=num_tokens_pad,
local_query_start_loc=local_query_start_loc,
local_seq_lens=local_seq_lens,
max_local_query_len=max_local_query_len,
max_local_seq_lens=max_local_seq_lens,
local_cos=local_cos,
local_sin=local_sin,
start_pos=start_pos,
)

dspark_swa_indices = None
ori_win_left, ori_win_right = self.model_config.hf_config.sliding_window - 1, 0
Expand Down Expand Up @@ -531,14 +605,15 @@ def build_req_metadata_for_drafting(
slot_mapping = self.spec_slot_mapping[draft_index - 1][: self.num_actual_tokens]

num_heads = self.model_config.hf_config.num_attention_heads
head_dim = self.model_config.get_head_size()
kv_plan = get_dsa_attn_kv_plan(self.vllm_config)
metadata_op = kv_plan.get_dsa_sparse_attn_metadata_op()
metadata_kwargs = kv_plan.get_dsa_sparse_attn_metadata_kwargs(self.seqused_q.device)
cu_seqlens_ori_kv = (
local_query_start_loc
if has_prefill
else DeviceOperator.get_dsa_decode_cu_seqlens_ori_kv(
None,
metadata_cache,
"draft_cu_seqlens_ori_kv",
local_seq_lens,
num_reqs,
Expand All @@ -549,28 +624,46 @@ def build_req_metadata_for_drafting(
cu_seqlens_cmp_kv = (
None if has_prefill else DeviceOperator.get_dsa_decode_cu_seqlens_cmp_kv(self.cu_seqlens_cmp_kv)
)
sas_metadata = metadata_op(
**metadata_kwargs,
num_heads_q=num_heads,
num_heads_kv=1,
head_dim=self.model_config.get_head_size(),
cu_seqlens_q=local_query_start_loc,
cu_seqlens_ori_kv=cu_seqlens_ori_kv,
cu_seqlens_cmp_kv=cu_seqlens_cmp_kv,
seqused_q=self.seqused_q,
seqused_kv=local_seq_lens,
max_seqlen_q=max_local_query_len,
max_seqlen_kv=max_local_seq_lens,
batch_size=num_reqs,
cmp_ratio=1,
ori_mask_mode=4,
ori_win_left=ori_win_left,
ori_win_right=ori_win_right,
layout_q="TND",
layout_kv=_dsa_layout_kv(self.vllm_config),
has_ori_kv=True,
has_cmp_kv=False,
)
if (
metadata_cache is not None
and metadata_cache.get("sas_num_heads") == num_heads
and metadata_cache.get("sas_head_dim") == head_dim
and "sas_metadata" in metadata_cache
):
# The SAS metadata op only consumes group-invariant inputs
# (local cu_seqlens / seq_lens / window config), so the tensor
# built by a sibling kv-cache group of the same draft step is
# reused instead of re-running the device metadata kernel.
sas_metadata = metadata_cache["sas_metadata"]
else:
sas_metadata = metadata_op(
**metadata_kwargs,
num_heads_q=num_heads,
num_heads_kv=1,
head_dim=head_dim,
cu_seqlens_q=local_query_start_loc,
cu_seqlens_ori_kv=cu_seqlens_ori_kv,
cu_seqlens_cmp_kv=cu_seqlens_cmp_kv,
seqused_q=self.seqused_q,
seqused_kv=local_seq_lens,
max_seqlen_q=max_local_query_len,
max_seqlen_kv=max_local_seq_lens,
batch_size=num_reqs,
cmp_ratio=1,
ori_mask_mode=4,
ori_win_left=ori_win_left,
ori_win_right=ori_win_right,
layout_q="TND",
layout_kv=_dsa_layout_kv(self.vllm_config),
has_ori_kv=True,
has_cmp_kv=False,
)
if metadata_cache is not None:
metadata_cache.update(
sas_num_heads=num_heads,
sas_head_dim=head_dim,
sas_metadata=sas_metadata,
)

cp_metadata = DSACPMetadata(
local_query_start_loc=local_query_start_loc,
Expand Down
17 changes: 13 additions & 4 deletions vllm_ascend/spec_decode/llm_base_proposer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2356,13 +2356,22 @@ def build_draft_attn_metadata(
# FIXME(woosuk): The below two ops cause synchronization. Optimize.
assert len(self.draft_attn_groups) > 0
per_layer_attn_metadata: dict[str, Any] = {}
# One DSA cache dict shared by all attn groups within this decode step.
# DSpark draft layers span multiple kv-cache groups; every group gets
# its own build_for_drafting call but most of the drafted metadata
# (RoPE tables, local token metadata, SAS metadata) only depends
# on group-invariant fields of common_attn_metadata, so it is computed
# by the first group and reused by the remaining ones. Only the DSA
# metadata builder (used by dspark) understands this kwarg; the generic
# GQA builder used by dflash/eagle does not accept **kwargs, so leave
# the cache empty in the non-compression path to avoid passing an
# unexpected argument.
shared_dsa_draft_cache: dict = dict(common_ratio_to_sas_metadata=dict()) if self.use_compress else {}
for attn_group in self.draft_attn_groups:
builder = attn_group.get_metadata_builder()
extra_attn_metadata_args: dict = {}
extra_attn_metadata_args: dict = dict(shared_dsa_draft_cache)
if self.use_compress:
extra_attn_metadata_args = dict(
common_ratio_to_sas_metadata=dict(),
)
extra_attn_metadata_args["block_size"] = attn_group.kv_cache_spec.block_size
if self.method == "dspark":
gid = attn_group.kv_cache_group_id
common_attn_metadata = copy.copy(common_attn_metadata)
Expand Down
Loading