Prevent stale CP shared-KV contracts from corrupting prefill
CP shared-KV now uses CP-local current rows consistently across MLA/index current reuse, passes fp8 current-index K through the tai-kernel uint8 ABI, and clears the transient EAGLE CP-local hidden marker after draft capture. The disaggregation bootstrap also fingerprints the runtime source contract so prefill/decode mismatches fail fast instead of silently exchanging incompatible KV metadata. Constraint: CP shared-KV batch paths flatten current K/V rows in CP-rank-local valid order, not global request order. Constraint: tai-kernel current-index prepare validates current_index_k as uint8 bytes for fp8 payloads. Rejected: Keep using global extend offsets for bs>1 current-index reuse | corrupts request-local bases once current_index_kv is CP-local. Rejected: Infer CP-local EAGLE hidden semantics from tensor shape | static padding and bs>1 can make shape-based inference unsafe. Confidence: medium Scope-risk: moderate Directive: Do not reintroduce forward_batch.out_cache_loc slicing in CP shared-KV current reuse without verifying CP-local owner-lane layout. Tested: Remote container py_compile for touched runtime/test files. Tested: Remote PYTHONPATH=python pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py (198 passed, 2 subtests passed). Not-tested: Full remote ETE traffic after this commit; accept length and garbage-output recovery still require a fresh prefill/decode run. Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
@@ -1072,6 +1072,7 @@ class EAGLEWorker(TpModelWorker):
|
||||
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
|
||||
draft_input.topk_p, draft_input.topk_index = fast_topk(probs, self.topk, dim=-1)
|
||||
draft_input.hidden_states = logits_output.hidden_states
|
||||
draft_input.cp_local_hidden_states = False
|
||||
|
||||
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
|
||||
monkey_patch_torch_reductions()
|
||||
|
||||
Reference in New Issue
Block a user