Keep Phase8 prefetch on the deferred-consume path

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
@@ -257,21 +257,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_log_enabled,
cp_shared_kv_mla_prefetch_should_log_layer,
cp_shared_kv_mla_prefetch_wait_after_attention_enabled,
)
envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.clear()
envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.clear()
self.assertFalse(cp_shared_kv_mla_prefetch_log_enabled())
self.assertFalse(cp_shared_kv_mla_prefetch_wait_after_attention_enabled())
self.assertFalse(cp_shared_kv_mla_prefetch_should_log_layer(1))
self.assertTrue(cp_shared_kv_mla_prefetch_should_log_layer(2))
self.assertFalse(cp_shared_kv_mla_prefetch_should_log_layer(3))
with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True):
self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled())
with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.override(True):
self.assertTrue(cp_shared_kv_mla_prefetch_wait_after_attention_enabled())
def test_token_range_materialize_uses_tai_kernel_when_enabled(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime