chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Align CP shared-KV prefetch with the attention overlap window

Browse Source 
Launching CP shared-KV prefetch from MLA prepare or the indexer can make
next-layer prefix work overlap current-layer MQA/materialization instead of
the attention window. Centralize the launch in the NSA backend after
current-layer materialization and before attention, and leave the early
indexer hook inert so the call site cannot regress silently.

The accompanying notes capture the draft-as-forward-layer follow-up plan and
the latest OOM diagnosis: the observed 178945-token failure matches the CP
in-seq MQA logits allocation, so the follow-up fix is q-dimension chunking in
_get_topk_ragged_with_cp(), not a max-prefetch-size gate.

Constraint: CP shared-KV prefetch must avoid overlapping current-layer MQA/materialization and must not add silent fallback behavior.
Rejected: Limit maximum prefetch size | hides the CP logits peak and can reduce cache/prefetch effectiveness.
Rejected: Keep prepare/indexer launch sites | they place next-layer collectives in the wrong overlap window.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce early CP prefetch launch without checking Nsight overlap and CP MQA memory peaks.
Tested: Local git diff --check and py_compile for touched Python files.
Tested: Remote container py_compile plus targeted pytest: 3 passed, 5 warnings.
Not-tested: Full ETE under production traffic after this commit.
Not-tested: CP in-seq MQA logits chunking; documented as follow-up.
This commit is contained in:
laoyao0822
2026-06-01 01:09:43 +08:00
parent 4342de0463
commit 46be97adc0
6 changed files with 206 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -20,7 +20,6 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_log,
cp_shared_kv_mla_prefetch_log_enabled,
cp_shared_kv_mla_prefetch_should_log_layer,
cp_shared_kv_should_prefetch_next_layer,
filter_owned_logical_locs,
get_or_build_shared_paged_buffer_slot_remap,
is_current_only_extend_batch,
@@ -478,36 +477,16 @@ class Indexer(MultiPlatformOp):
forward_batch: ForwardBatch,
layer_id: int,
) -> None:
next_layer_id = layer_id + 1
index_prefetcher = getattr(
forward_batch, "cp_shared_kv_index_prefetcher", None
)
if (
cp_shared_kv_mla_prefetch_log_enabled()
and cp_shared_kv_mla_prefetch_should_log_layer(next_layer_id)
):
layout = getattr(forward_batch, "cp_shared_kv_layout", None)
seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
cp_shared_kv_mla_prefetch_log(
"index_start_request cp_rank=%s cp_size=%s layer=%s "
"next_layer=%s has_index=%s uses_cp_shared_kv=%s "
"seq_lens_cpu_len=%s",
getattr(layout, "cp_rank", None),
getattr(layout, "cp_size", None),
layer_id,
next_layer_id,
index_prefetcher is not None,
getattr(forward_batch, "uses_cp_shared_kv", None),
len(seq_lens_cpu) if seq_lens_cpu is not None else None,
)
if index_prefetcher is None:
return
if not cp_shared_kv_should_prefetch_next_layer(forward_batch, layer_id):
return
index_prefetcher.start_next_layer_prefix(
next_layer_id=next_layer_id,
token_to_kv_pool=forward_batch.token_to_kv_pool,
)
"""Deprecated early prefetch hook.
Next-layer CP shared-KV prefetch must be launched by the attention
backend after current-layer index/MLA materialization is complete and
immediately before the attention kernel. Starting here overlaps the
next-layer collective with current-layer MLA materialization/all-reduce,
which is the wrong dependency window.
"""
return
def _filter_shared_index_write(
self,
@@ -1901,7 +1880,6 @@ class Indexer(MultiPlatformOp):
topk=self.index_topk,
layer_id=layer_id,
)
self._maybe_start_next_layer_index_prefetch(forward_batch, layer_id)
return topk_result
def forward_npu(

46
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -98,6 +98,43 @@ global_workspace_buffer = None
_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip
def _maybe_start_cp_shared_kv_attention_prefetch(
forward_batch: ForwardBatch,
layer_id: int,
) -> None:
"""Launch next-layer CP shared-KV prefetch in the attention overlap window.
This hook intentionally runs after current-layer index/MLA cache
materialization has finished and immediately before the attention kernel.
Earlier hooks can make the next-layer collective overlap current-layer KV
materialization/reduce instead of attention, which can serialize or contend
with the current layer's required cache work.
"""
token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
if token_to_kv_pool is None:
return
if not cp_shared_kv_should_prefetch_next_layer(forward_batch, layer_id):
return
next_layer_id = int(layer_id) + 1
index_prefetcher = getattr(
forward_batch, "cp_shared_kv_index_prefetcher", None
)
if index_prefetcher is not None:
index_prefetcher.start_next_layer_prefix(
next_layer_id=next_layer_id,
token_to_kv_pool=token_to_kv_pool,
)
mla_prefetcher = getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None)
if mla_prefetcher is not None:
mla_prefetcher.start_next_layer_prefix(
next_layer_id=next_layer_id,
token_to_kv_pool=token_to_kv_pool,
)
@dataclass(frozen=True)
class NSAFlashMLAMetadata:
"""Metadata only needed by FlashMLA"""
@@ -2070,16 +2107,11 @@ class NativeSparseAttnBackend(
else None,
tuple(out_cache_loc.shape) if out_cache_loc is not None else None,
)
if mla_prefetcher is not None and cp_shared_kv_should_prefetch_next_layer(
forward_batch, layer.layer_id
):
mla_prefetcher.start_next_layer_prefix(
next_layer_id=layer.layer_id + 1,
token_to_kv_pool=forward_batch.token_to_kv_pool,
)
else:
mla_prefetcher = None
_maybe_start_cp_shared_kv_attention_prefetch(forward_batch, layer.layer_id)
index_prefetcher = getattr(
forward_batch, "cp_shared_kv_index_prefetcher", None
)

44
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
View File

@@ -14,7 +14,6 @@ from sglang.srt.layers.attention.nsa.utils import (
nsa_use_prefill_cp,
)
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_should_prefetch_next_layer,
should_reuse_current_extend_kv,
try_tai_fused_mla_store,
)
@@ -95,47 +94,6 @@ class DeepseekMLAForwardMixin:
get_global_server_args().flashinfer_mla_disable_ragged
)
def _maybe_start_cp_shared_next_layer_prefetch(
self: DeepseekV2AttentionMLA,
forward_batch: ForwardBatch,
) -> None:
"""Launch next-layer CP shared-KV prefix prefetch before rank-skewed work.
The original Phase8 hook launched prefetch from the indexer/backend after
current-layer MQA/materialize work. Those paths can be rank-skewed in
in-seq CP, so rank 0 may enqueue the next-layer collective early while
other ranks reach it too late to overlap. Starting from layer prepare
keeps the launch point before the current layer's per-rank MQA/topk and
materialize imbalance. Existing later hooks remain as fallbacks and
become no-ops via the prefetchers' already-started guard.
"""
token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
if token_to_kv_pool is None:
return
if not cp_shared_kv_should_prefetch_next_layer(
forward_batch, self.layer_id
):
return
next_layer_id = int(self.layer_id) + 1
index_prefetcher = getattr(
forward_batch, "cp_shared_kv_index_prefetcher", None
)
if index_prefetcher is not None:
index_prefetcher.start_next_layer_prefix(
next_layer_id=next_layer_id,
token_to_kv_pool=token_to_kv_pool,
)
mla_prefetcher = getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None)
if mla_prefetcher is not None:
mla_prefetcher.start_next_layer_prefix(
next_layer_id=next_layer_id,
token_to_kv_pool=token_to_kv_pool,
)
def forward_absorb_prepare(
self: DeepseekV2AttentionMLA,
positions: torch.Tensor,
@@ -146,8 +104,6 @@ class DeepseekMLAForwardMixin:
):
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
self._maybe_start_cp_shared_next_layer_prefetch(forward_batch)
q_lora = None
topk_indices = None
if self.q_lora_rank is not None: