Stabilize CP shared-KV prefetch around draft cache hits

Cache-hit EAGLE/NextN draft extends can enter the draft DeepEP MoE immediately after CP shared-KV attention. The partial current-reuse path is kept for target layers, but draft cache-hit suffixes now use full materialization until draft has an explicit same-layer reuse contract. Next-layer MLA/index prefetch is also gated by the actual model depth, so the single-layer draft model does not enqueue unused next-layer async work.

The temporary stage traces used to isolate the hang are removed. The retained draft current-reuse fallback is a bounded warning because it changes the runtime path intentionally.

Constraint: EAGLE/NextN has one executable draft layer and mirrors target KV state.

Rejected: Keep partial current reuse for draft cache-hit suffixes | reproduced hangs at draft layer0 before DeepEP MoE completion.

Rejected: Keep temporary stage traces | useful for diagnosis but too noisy for normal runs.

Confidence: medium

Scope-risk: moderate

Directive: Do not re-enable draft cache-hit partial current reuse without an explicit draft same-layer reuse contract and ETE validation with CP shared KV + HiCache + EAGLE.

Tested: py_compile on edited Python files; git diff --check; temp trace grep returned no matches.

Not-tested: Local targeted pytest is blocked by missing pybase64 in this environment; full ETE after log cleanup not run.
This commit is contained in:
laoyao0822
2026-05-29 00:33:41 +08:00
parent 26c792939d
commit c3fc3ff752
7 changed files with 808 additions and 71 deletions
+21 -7
View File
@@ -1680,14 +1680,28 @@ class DeepseekV2DecoderLayer(nn.Module):
quant_format,
)
hidden_states = self.self_attn(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
zero_allocator=zero_allocator,
llama_4_scaling=llama_4_scaling,
layer_scatter_modes=self.layer_scatter_modes,
previous_cp_shared_kv_num_model_layers = getattr(
forward_batch, "cp_shared_kv_num_model_layers", None
)
forward_batch.cp_shared_kv_num_model_layers = (
1 if self.is_nextn else self.config.num_hidden_layers
)
try:
hidden_states = self.self_attn(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
zero_allocator=zero_allocator,
llama_4_scaling=llama_4_scaling,
layer_scatter_modes=self.layer_scatter_modes,
)
finally:
if previous_cp_shared_kv_num_model_layers is None:
delattr(forward_batch, "cp_shared_kv_num_model_layers")
else:
forward_batch.cp_shared_kv_num_model_layers = (
previous_cp_shared_kv_num_model_layers
)
hidden_states, residual = self.layer_communicator.prepare_mlp(
hidden_states, residual, forward_batch