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

Keep Phase8 prefetch on the deferred-consume path

Browse Source 
Phase8 only gains useful overlap when the next-layer MLA prefix prefetch is allowed to run until the next layer actually consumes the prefetched buffer. The old wait-after-attention switch let runtime configuration collapse the optimization back into current-layer tail latency, so the prefetch path now has one wait policy and the documentation records the implemented behavior.

Constraint: Phase8 should keep the production environment surface minimal while preserving the existing enable and debug-log knobs
Rejected: SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION | it reintroduced current-layer synchronous waiting and made profiling behavior depend on a nonessential policy knob
Confidence: medium
Scope-risk: narrow
Directive: Do not add another Phase8 wait policy knob without first proving the added policy improves end-to-end prefill latency under CP shared KV
Tested: Python AST parse for touched Python files
Tested: git diff --check
Not-tested: Full pytest and remote server integration were not run in this commit
This commit is contained in:
laoyao0822
2026-05-03 03:47:31 +08:00
parent bc23a81884
commit 9fec89ba09
5 changed files with 77 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/environ.py
View File

@@ -207,7 +207,6 @@ class Envs:
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION = EnvBool(False)
SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)

15
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
View File

@@ -14,7 +14,6 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_enabled,
cp_shared_kv_mla_prefetch_log,
cp_shared_kv_mla_prefetch_should_log_layer,
cp_shared_kv_mla_prefetch_wait_after_attention_enabled,
filter_locs_mappable_to_physical_pool,
materialize_local_token_kv_page_slots_into,
remap_logical_locs_to_slot_dense_locs_optimized,
@@ -393,24 +392,12 @@ class CpSharedKVMlaPrefetcher:
)
def wait_attention_window(self) -> None:
if not cp_shared_kv_mla_prefetch_wait_after_attention_enabled():
handle = self.pending_attention_handle
if handle is not None:
self._log_next_layer(
handle.layer_id,
"attention_wait_deferred next_layer=%s",
handle.layer_id,
)
return
handle = self.pending_attention_handle
self.pending_attention_handle = None
if handle is None:
return
torch.cuda.current_stream().wait_event(handle.event)
self._log_next_layer(
handle.layer_id,
"attention_wait next_layer=%s",
"attention_wait_deferred next_layer=%s",
handle.layer_id,
)

4
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -37,10 +37,6 @@ def cp_shared_kv_mla_prefetch_log_enabled() -> bool:
return envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.get()
def cp_shared_kv_mla_prefetch_wait_after_attention_enabled() -> bool:
return envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.get()
def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None:
if cp_shared_kv_mla_prefetch_log_enabled():
logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args)