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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user