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

Prevent stale CP shared-KV contracts from corrupting prefill

Browse Source 
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:
laoyao0822
2026-06-04 20:22:29 +08:00
parent 3d6007246b
commit f50e2b1e00
9 changed files with 1463 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

878
docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md Normal file
View File

@@ -0,0 +1,878 @@
# NSA Prefill CP shared-KV：accept len=1 / 输出乱码排查 Ledger
日期：2026-06-04
目标：记录已经检查过的证据、排除项和当前假设，避免后续在日志/代码压缩后重复阅读同一批路径。
维护规则：
- 后续每确认一个新事实，先追加到本文，再继续排查。
- 只记录有日志、代码位置或测试输出支撑的事实；推测必须放在“假设”段落，并标注待验证。
- 已排除项不要在后续轮次重新当作主线，除非出现新的相反证据。
- 修复前必须先把“哪个 logical/physical/padded 合同被破坏”定位清楚。
## 0. 当前现象
用户在 g0033 经 router 调用：
```text
python summarize_story.py
api_url=http://g0034:17000/v1/chat/completions
model=GLM5
prompt_tokens=40387
completion_tokens=113
HTTP_OK seconds=5.79
输出大量无意义 token / 重复 </think>
```
同时 decode 日志显示 accept len 掉到约 1.0。
## 1. 当前远端运行配置已确认
时间：2026-06-04 07:51 CST 左右。
prefill 进程在 `g0034`：
```text
python -m sglang.launch_server ...
--disaggregation-mode prefill
--tp-size 8 --ep-size 8
--enable-nsa-prefill-context-parallel
--nsa-prefill-cp-mode in-seq-split
--attn-cp-size 8
--enable-nsa-prefill-cp-shared-kv
--speculative-algorithm EAGLE
--speculative-num-steps 3
--speculative-num-draft-tokens 4
--enable-hierarchical-cache
--hicache-size 150
--hicache-io-backend direct
--hicache-mem-layout page_first_direct
--kv-cache-dtype fp8_e4m3
--enable-cp-shared-kv-prefill-bs-gt1
--cp-shared-kv-prefill-max-batch-requests 5
--cp-shared-kv-prefill-max-total-extend-tokens 65536
--chunked-prefill-size 65536 # 命令里有 -1 和 65536，ServerArgs 生效为 65536
```
关键环境变量由用户命令提供：
```text
SGLANG_EAGLE_ACCEPT_DEBUG=1
SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=1
SGLANG_CP_DRAFT_SHARED_KV=1
SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE=1
SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=1
SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1
SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1
SGLANG_CP_SHARED_KV_CURRENT_REUSE=1
SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1
```
decode 在 g0035/g0036，日志文件当前为：
```text
/mnt/beegfs/cjy/log/decode0_20260603_220439.log
/mnt/beegfs/cjy/log/decode1_20260603_220442.log
```
prefill 最新日志：
```text
/mnt/beegfs/cjy/log/sglang_cp_hicache_20260603_232444.log
```
## 2. 本次坏请求对应日志证据
### 2.1 prefill 请求与 cache 写入
prefill 日志中坏请求：
```text
[2026-06-03 23:28:20] Requested token count (40391 input + 200000 new) exceeds ... Truncating max_new_tokens.
```
随后 HiCache backup：
```text
[23:28:21 ATTN_CP0..7] [CacheCtrl-write] submit_write_cp_per_layer registered: node_id=4 logical_len=40386 layers=78 draft=True
[23:28:21 ATTN_CP0..7] [HiCache-write] prepared CP per-layer backup before forward: ... logical_len=40386 owned_positions=5056
[23:28:22 ATTN_CP0..7] [CacheCtrl-write] submit_write_cp_layer final ack: node_id=4 logical_len=40386 layers=78 draft=True
```
注意：日志中的 request input 是 40391，但 HiCache node logical_len 是 40386。该差异尚未定性，不能重复忽略。
### 2.2 prefill 实际 forward token 数
```text
[23:28:22 ATTN_CP0..7] Prefill batch, #new-seq: 1, #new-token: 40448, #cached-token: 0
```
解释：`40448 = ceil(40386 / 64) * 64`，这是 page granularity / static padding 后的物理计算 token 数。它本身不是 bug 证据；bug 应该出在某条路径把物理 padding 当成 valid/logical token 消费。
### 2.3 index current reuse fallback 警告
同一请求中出现大量：
```text
[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] reason=current_batch_kernel_failed
error=current_index_k must be uint8 segments=2 total_kv_len=43072..43456 total_q_count=5056
```
已知事实：
- 这些 warning 主要出现在 CP1-CP7；CP0 的长请求 warning 未观察到。
- 对 `L=40386, page_size=64, cp_size=8`：
- padded tokens = 40448
- CP1-CP7 的 valid local rows = 5056，因此 `total_q_count=5056` 对 CP1-CP7 本身不证明 tail 泄漏。
- CP0 valid local rows 应为 4994，full-page compute rows为 5056。
- 所以该 warning 目前更像 dtype/perf fallback：TAI current-index prepare 期望 `uint8`，而调用侧可能传入 `float8_e4m3fn` view；不应直接把它当成乱码根因。
### 2.4 decode accept len
decode1 当前请求：
```text
[23:28:24 DP15] Decode batch, #running-req: 1, #token: 40448, accept len: 1.10, accept rate: 0.28
[23:28:25 DP15] Decode batch, #running-req: 1, #token: 40512, accept len: 1.00, accept rate: 0.25
```
说明：`#token=40448` 来自 allocator/token usage 统计，不一定等于 logical seq len；不能仅凭这条判定 decode 使用了 padding token，但它与 prefill 物理 page padding 数一致，需要继续追踪消费者是否使用了 `len(kv_loc)`。
## 3. 已排除/不要重复误判的点
### 3.1 不是“服务没同步/旧代码”导致的直接结论
之前已通过 md5 对比确认远端关键文件与本地一致：
```text
nsa_indexer.py md5 efe7bcc118fbd4cc84c6017aac818228
utils.py md5 b953059a79a7e19cd12154098d995ada
```
当前进程启动时间在最新 commit 之后，且日志中已经出现最新 fail-fast/fallback 字符串。
### 3.2 早期 500 / OUTPUT_ZERO 主要是 prefill bootstrap 不可用
decode 日志中 22:45-23:18 大量：
```text
Could not fetch prefill parallel info from g0034:18992
Decode transfer failed ... Losing connection with prefill instance
[OUTPUT_ZERO_DEBUG] ... finish_reason abort
```
这些发生在 prefill 重启/不可用阶段，不是当前 23:28 HTTP 200 但输出乱码的主证据。后续不要把这些旧 500 当成当前乱码根因。
### 3.3 warmup 输出也不正常，但不等同本次根因
decode warmup response 有：
```text
text='! of the of the of the'
spec_accept_rate=0.25
spec_accept_length=2.0
```
这说明模型/EAGLE 在短 warmup 上也偏异常，但短 warmup 与 40k prompt + CP shared KV + FP8 + HiCache 的路径不同。可作为背景信号，不能单独定位 root cause。
## 4. 已阅读的关键代码路径
本节用于防止重复读同一批入口。后续如果继续深入某个函数，需要在对应条目下补“确认/排除/风险”。
### 4.1 CP metadata / split / local cache loc
文件：`python/sglang/srt/layers/attention/nsa/utils.py`
已看函数/对象：
- `CPSharedKVBatchPlan`
- `build_batch_page_aligned_in_seq_split_plan()`
- `split_tensor_by_cp_batch_plan()`
- `_get_cp_local_valid_row_indices_cache()`
- `select_cp_local_valid_rows_for_cache_write()`
- `get_cp_shared_kv_local_out_cache_loc()`
当前理解：
- 普通 page-tail padding 与 compute-padding 是两个概念。
- 现在 `compute_padding_enabled` 只代表“为了避免 tiny/zero lane，把 compute padded 到至少 cp_size pages”等更强 padding；长请求尾页 padding 不一定使它为 true。
- 因此任何只在 `compute_padding_enabled=True` 时才裁 valid rows 的路径，都需要审查是否会在普通尾页下泄漏 padding。
当前状态：
- 已确认：`40448` 这类 full-page token 数可以是合法物理覆盖。
- 未确认：所有 current-only reuse 路径是否都使用 CP-local valid rows，而不是 global valid rows 或 physical padded rows。
### 4.2 NSA index current reuse/top-k
文件：`python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
已看函数：
- `_maybe_materialize_shared_index_buffer()`
- `_get_topk_ragged_with_cp()`
- `_get_topk_in_seq_cp_pair_batch()`
- `forward_cuda()` 中 current index reuse / cache write 逻辑
关键点：
- current-only 时 `_get_topk_in_seq_cp_pair_batch()` 不 materialize prefix，而是把 `current_index_kv_for_topk` 传给 `_get_topk_ragged_with_cp()`。
- `_select_batch_topk_query_lengths()` 已用于根据 `q_fp8/weights` 实际行数选择 valid/compute metadata，修复过 40392 vs 40387、5056 vs 4995 类 fail-fast。
- 现在仍需确认 current-only 下 `kv_len`、`seq_len`、`logical_kv_limit` 与 current rows 是否完全一致。
当前状态：
- 已确认：FP8 下 `tai_index_mqa_prepare` warning 是 `current_index_k must be uint8`，至少包含 dtype fast-path 不匹配。
- 未确认：fallback 后 Python 拼接路径是否仍严格保持 index/top-k 的 logical order 与 CP owner-lane order。
### 4.3 MLA current/partial-current reuse
文件：`python/sglang/srt/layers/attention/nsa_backend.py`
已看区域：
- `init_forward_metadata()` / `_transform_table_1_to_real()`
- MLA current-only / partial-current reuse around `_forward_flashmla_kv`
当前高优先级怀疑：
- current-only MLA reuse 通过 `materialize_prefix_and_reuse_current_kv_page_slots(prefix_pages=0)` 构造 page-slot dense KV。
- 如果 `logical_page_table_1` 覆盖的 logical loc 数量大于 `current_locs_for_reuse` 覆盖数量，就会出现部分 logical page/loc 没有 current KV 填充，但 attention 仍可能消费。
- 这类问题对 FP8 更敏感，因为 FP8 cache 必须严格 page-slot aligned，不能 compact。
当前状态：
- 未确认：`logical_page_table_1` 的消费范围是否只覆盖 valid logical loc。
- 未确认：FP8 page-slot compose 与 BF16 page-slot compose 是否走完全相同的 logical/page-table 合同。
追加确认（2026-06-04）：
- `forward_mla.py::_maybe_write_cp_shared_local_mla_kv()` 在 `rebuild_cp_kv_cache()` 之前写 persistent CP shared-KV。
- 当 `should_reuse_current_extend_kv()` 为真时，prefill attention 本轮可以通过 `rebuild_cp_kv_cache()` 得到 all-gather 后的 current KV；这条路径与写入给 decode/HiCache/transfer 使用的 persistent KV 不是同一个消费者。
- 因此，`prefill HTTP 200` 不能证明 persistent KV 写入正确；当前乱码更可能是 persistent KV 写入或 transfer 映射损坏，decode 消费后 EAGLE accept len 掉到 1。
### 4.4 runtime helpers
文件：`python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
已看函数：
- `is_current_only_extend_batch()`
- `should_reuse_current_extend_kv()`
- `current_extend_kv_rows_for_reuse()`
- `build_current_loc_remap()`
- `build_shared_token_kv_slot_remap()`
- `materialize_prefix_and_reuse_current_kv_page_slots()`
- `fill_current_kv_page_slots_and_remap_locs()`
当前需要验证：
- `current_extend_kv_rows_for_reuse()` 返回 `sum(extend_seq_lens_cpu)`，它不一定等于当前 attention page table 需要覆盖的 logical loc 数。
- `materialize_prefix_and_reuse_current_kv_page_slots()` 中 current loc remap 如果存在 missing logical loc，应 fail-fast 或禁用 current-only，而不能继续构造不完整 KV。
当前状态：
- 已确认：helper 层有能力表达 `prefix_pages=0` 的 current-only page-slot compose。
- 风险：`current_extend_kv_rows_for_reuse()` 是 token-count 视角；如果调用侧需要 CP-local valid rows，必须在调用侧显式传入/选择 CP-local loc。
### 4.5 padding 进入 forward 的位置
文件：`python/sglang/srt/model_executor/forward_batch_info.py`
已看函数：
- `prepare_mlp_sync_batch()`
- `_pad_inputs_to_size()`
当前理解：
- MLP sync / attn_tp / cp 对齐会把 `forward_batch.extend_num_tokens` 和张量行数 pad 到更大数，例如 40392、40448。
- CPU logical fields（`extend_seq_lens_cpu`, `seq_lens_cpu`）不应该被 padding 改写。
- 所有 cache write / index topk / MLA attention 需要明确区分：valid rows、compute rows、static-padded rows、page-padded rows。
当前状态：
- 已确认：`forward_batch.extend_num_tokens` 不能当作 request logical length。
- 待查：哪些 consumer 仍然隐式使用张量第一维或 `out_cache_loc.numel()` 作为 logical row 数。
### 4.6 disagg transfer / decode prealloc
文件：
- `python/sglang/srt/disaggregation/prefill.py`：`send_kv_chunk()`
- `python/sglang/srt/disaggregation/decode.py`：`_pre_alloc()`
已知：
- prefill transfer 的 `end_idx` 基于 `len(req.fill_ids)` / `len(req.origin_input_ids)`，再转 page indices；最后一个 partial page 仍会按 page 转移。
- decode `_pre_alloc()` 的 `fill_len` 基于 logical input/output ids，但 paged allocation 返回的 `kv_loc` 可能包含 full tail page locs。
- 需要继续检查 decode attention 是否始终用 logical `seq_lens`，不要用 `len(kv_loc)`/page padded loc 数作为真实上下文长度。
当前状态：
- 已排除：早期 decode bootstrap 失败导致的 `OUTPUT_ZERO` 不是当前 HTTP 200 乱码主线。
- 待查：decode 端 EAGLE accept 计算是否使用 logical seq len 还是 padded kv loc 数。
### 4.7 FP8 CP shared-KV persistent write
文件：
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py`
- `python/sglang/srt/layers/attention/nsa/quant_k_cache.py`
- `python/sglang/srt/mem_cache/memory_pool.py`
- `tai-kernel/python/tai_kernel/nsa_prefill/cp_shared_mla_store.py`
- `tai-kernel/python/tai_kernel/nsa_prefill/csrc/cp_shared_mla_store.cu`
已确认：
- 运行参数开启 `--kv-cache-dtype fp8_e4m3` 且 `SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=1`。
- persistent MLA 写入优先走 `try_tai_fused_mla_store()`；失败才回到 SGLang fallback：`logical_locs -> physical_locs -> set_mla_kv_buffer()`。
- TAI FP8 fused store 对 `kv_buffer.dtype == torch.uint8` 走 fused quant store，row bytes 为 `512 FP8 + 16 scale + 128 rope bf16 = 656`。
- TAI 侧已有自洽测试，但目前看到的测试主要比较 TAI kernel 与 TAI 自己的 reference；还缺“TAI fused store 与 SGLang fallback 写同一个 raw FP8 pool 后逐字节等价”的集成单测。
高优先级待验证：
- 对相同的 CP logical locs、同样的 `k_nope/k_rope`，TAI fused store 和 SGLang fallback 写出的 `kv_buffer[layer]` 必须完全一致。
- 若该测试失败，root cause 在 persistent write；若通过，则转向 `SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1` 的 prefill->decode physical/logical 映射。
## 5. 当前最强假设（待验证）
### H1：persistent CP shared-KV 写入或 transfer 映射损坏（当前主线）
证据：
- 坏请求 `#cached-token=0`，prefill attention 可以主要依赖 current all-gather KV；但是 decode 只能消费 prefill 写入并 transfer 过来的 persistent KV。
- prefill 在同一请求中完成 `submit_write_cp_per_layer` / `final ack`，随后 decode accept len 立即掉到约 1。
- 乱码而非 shape crash，更符合“KV 内容/位置不一致”。
- FP8 下 persistent write 走 TAI fused store，尚未有与 SGLang fallback 的集成等价测试。
最小验证方向：
- 先补 CUDA 集成单测：同一批 CP logical locs 下，比较 TAI fused FP8 MLA store 与 SGLang fallback `set_mla_kv_buffer()` 的 raw buffer 字节。
- 若等价，再补 transfer 映射单测：CP rank 发送必须从 source physical locs 读，decode destination 必须按 logical page/loc 还原。
### H2：current-only MLA/index 复用的 logical length 与 current rows 覆盖不一致（降级为次线）
证据：
- 用户请求 usage prompt_tokens=40387；prefill log input=40391；HiCache logical_len=40386；forward #new-token=40448。
- 多个长度同时存在，说明当前路径同时持有 tokenizer/request length、cache logical length、page padded length、static padded length。
- 乱码 + accept len=1 比 crash 更像 KV 内容/位置不一致，而不是简单 shape mismatch。
最小验证方向：
- 在 current-only MLA/index reuse 点验证：`logical_page_table_1` 中每个会被 attention/topk 消费的 loc 都必须由 `current_locs_for_reuse` 覆盖。
- 如果未覆盖，应该 fail-fast 并打印：`seq_lens_cpu`、`extend_seq_lens_cpu`、`extend_prefix_lens_cpu`、`out_cache_loc.numel`、`page_table_1.shape/max`、`valid_current_rows`。
追加确认（2026-06-04）：
- 非 compute-padding 的 current-only index reuse 使用 all-gather 后 global `key/current_index_kv`，`_get_topk_ragged_with_cp()` 也按 per-request full valid rows 切片；这看起来是有意设计，不是当前 40k 单请求乱码的首要证据。
- MLA current-only 在 attention 前也会通过 `rebuild_cp_kv_cache()` 取得 full current KV；prefill 正常返回不能证明 persistent KV 正确。
### H3：FP8 index current TAI fallback 是性能问题，不是当前首要正确性根因
证据：
- warning 错误为 `current_index_k must be uint8`，fallback 到 Python slice/cat。
- CP1-CP7 的 `total_q_count=5056` 对该长度是正常 local valid rows。
- 仍需后续修复 dtype，让 fast path 吃到 FP8 uint8 raw view，但不应先把它当成乱码根因。
### H4：decode 端可能把 full-tail page loc 当作 logical seq 消费
证据：
- decode batch `#token=40448` 与 prefill physical page padded token 数一致。
- 这条目前只是风险，不是结论。需要继续追 `seq_lens` / page table 初始化是否用 logical fill_len。
## 6. 下一步不重复的执行清单
1. 先验证 H1：给 current-only MLA/index compose 增加窄 fail-fast 或 debug-only consistency check，优先单测复现，不直接大面积加日志。
2. 如果 H1 成立，修 source：统一当前请求的 logical/current rows 合同；不要用 silent fallback。
3. 如果 H1 不成立，再追 H3：decode transfer/prealloc 后的 `req_to_token_pool`、`seq_lens`、attention page table 是否 logical/padded 混用。
4. H2 单独排期：TAI current-index prepare dtype 应该传 uint8 raw storage，避免 FP8 下永远 fallback。
## 7. 禁止重复误区
- 不要再把 CP1-CP7 的 `total_q_count=5056` 单独当成错误；它对 40386/40391 这类长度可能是正常的。
- 不要再把 22:45-23:18 的 bootstrap 500 当成 23:28 HTTP 200 乱码的根因。
- 不要因为 `#new-token=40448` 直接判断 bug；必须定位到哪个 consumer 错把 physical padding 当 valid/logical 使用。
- 不要在没有证据前关闭 EAGLE/FP8/HiCache 作为“修复”；这只能绕过问题，不能解释输出乱码。
## 8. 当前排查边界
本 ledger 只覆盖当前 `accept len≈1 + HTTP 200 输出乱码` 主线。以下问题需要单独建账，不应混在本主线里：
- `tai_index_mqa_prepare` FP8 dtype fast-path fallback 的性能修复。
- HiCache / Mooncake transfer 带宽低、page 分散、allocator 连续性优化。
- L1/L2 free-room eviction 策略调参。
- bs>1 scheduler 组 batch 策略和 tiny extend compute-padding 性能模型。
如果后续证据证明其中某项直接导致输出乱码，再把它移动回主线并补充证据。
## 7. 2026-06-04 继续排查：新增证据
### 7.1 FP8 fused MLA persistent store 等价性单测
已新增并在远端 CUDA 环境验证：
```text
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_fused_mla_store_matches_sglang_fallback_for_cp_owned_tail_pages
1 passed, 3 warnings in 6.08s
```
测试内容：构造 `fp8_e4m3` packed MLA cache，包含 CP owned partial tail page，比较 `try_tai_fused_mla_store()` 与 SGLang fallback `quantize_k_cache_separate + set_mla_kv_buffer_triton()` 写出的 raw uint8 buffer。
结论：在已覆盖的 CP-owned tail-page 场景下，TAI fused FP8 MLA persistent store 与 SGLang fallback 等价。后续不要把“TAI fused MLA store 本身写错 tail page”作为最高优先级假设，除非出现更强反证。
### 7.2 prefill 与 decode 实际运行代码路径不一致
已用 `sha256sum` 比较本地、g0034 prefill、g0035/g0036 decode 容器内关键文件：
- 本地 `/root/sglang-work/sglang-dev` 与 `g0034:/sgl-workspace/sglang-tai` 完全一致。
- `g0035/g0036:/sgl-workspace/sglang-tai` 也与本地一致。
- 但 decode 日志栈路径是 `/sgl-workspace/sglang/...`，而 `g0035/g0036:/sgl-workspace/sglang` 的关键文件 hash 与本地不同：
- `nsa_backend.py`
- `nsa_indexer.py`
- `utils.py`
- `memory_pool.py`
- `disaggregation/mooncake/conn.py`
证据摘录：
```text
local nsa_backend.py: 6664f004024b8849...
g0034 /sgl-workspace/sglang-tai nsa_backend.py: 6664f004024b8849...
g0035 /sgl-workspace/sglang nsa_backend.py: 197624526cace147...
g0036 /sgl-workspace/sglang nsa_backend.py: 197624526cace147...
```
结论：当前线上 prefill 和 decode 很可能不是同一份 CP shared-KV/FP8/page-tail 合同。由于本次乱码发生在 prefill HTTP 200 后、decode 消费 transferred KV 时，代码版本不一致现在是最高优先级根因候选。
后续动作：
1. 先不要继续在 prefill 侧反复排查已验证的 store 路径。
2. 需要确认 decode 实际 bind mount 与启动目录，确保 decode 使用与 prefill 相同的代码。
3. 修复/同步后需要重启 decode 才能验证 accept len/output 是否恢复；仅 scp 文件不会影响已运行 Python 进程。
### 7.3 decode 容器挂载已确认：`/sgl-workspace/sglang` 是镜像内旧代码
`docker inspect sglang-glm5-dev-2`：
```text
g0035: /mnt/beegfs/cjy/sglang-dev -> /sgl-workspace/sglang-tai
g0036: /mnt/beegfs/cjy/sglang-dev -> /sgl-workspace/sglang-tai
```
没有 `/sgl-workspace/sglang` 的 bind mount。但 decode 栈路径是：
```text
/sgl-workspace/sglang/python/sglang/...
```
结论：decode 当前实际从容器镜像内旧代码启动；`/mnt/beegfs/cjy/sglang-dev` 的同步不会影响当前 decode 启动目录。prefill 与 decode 的 CP shared-KV/FP8/page-tail/transfer 合同不一致不是理论风险，而是当前运行事实。
正确修复方向：decode 启动脚本/工作目录必须切到 `/sgl-workspace/sglang-tai`，或在重启前显式更新 `/sgl-workspace/sglang`。仅修改 prefill 或 cjy host path 不能修复当前 accept len=1/输出乱码。
### 7.4 已加入代码级防线：prefill/decode runtime source fingerprint fail-fast
变更：`python/sglang/srt/disaggregation/common/conn.py`
- bootstrap parallel info 新增：
- `runtime_source_fingerprint`
- `runtime_source_root`
- fingerprint 覆盖当前 CP shared-KV / FP8 / page-tail 合同的关键源码：
- `sglang.srt.layers.attention.nsa_backend`
- `sglang.srt.layers.attention.nsa.nsa_indexer`
- `sglang.srt.layers.attention.nsa.utils`
- `sglang.srt.mem_cache.memory_pool`
- `sglang.srt.disaggregation.mooncake.conn`
- decode 拉取 prefill parallel info 时，如果 prefill fingerprint 与 decode 本地 fingerprint 不一致，直接 `RuntimeError`，避免继续运行到 silent KV corruption / accept len=1 / 输出乱码。
新增单测：
```text
test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py
```
RED：远端初跑 2 failed，原因是 `common.conn` 尚无 `get_runtime_source_fingerprint`。
GREEN：实现后远端验证：
```text
PYTHONPATH=python python -m pytest -q \
test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py \
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_fused_mla_store_matches_sglang_fallback_for_cp_owned_tail_pages
3 passed, 3 warnings in 6.13s
```
注意：这个防线会阻止后续同类错配静默运行；它不会改变已经启动的旧 decode 进程。当前远端要恢复输出质量，仍需要让 decode 进程从 `/sgl-workspace/sglang-tai` 或已同步的同源代码重启。
### 7.5 相关 transfer mapping 回归验证
远端验证：
```text
PYTHONPATH=python python -m pytest -q \
test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py \
test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py \
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_fused_mla_store_matches_sglang_fallback_for_cp_owned_tail_pages
10 passed, 3 warnings in 6.00s
```
覆盖范围：
- runtime source fingerprint mismatch fail-fast；
- CP shared-KV transfer 的 logical page -> physical source page / logical destination position 映射；
- valid tail page 不按 cp_size 过度 padding；
- FP8 fused MLA persistent store 与 fallback 等价。
### 7.6 fingerprint 实现修正：避免 import 副作用
实现细节已修正：fingerprint 不再对每个目标模块调用 `importlib.find_spec(module)`，避免触发 `mooncake.__init__` 这类父包 import。现在只对已加载的 `sglang` 包定位 `python/` 根目录，再按模块名拼出源码路径读取 bytes。
远端实际检查：
```text
get_runtime_source_root() -> /sgl-workspace/sglang-tai
len(get_runtime_source_fingerprint()) -> 64
fingerprint prefix -> 0c0e736fa6de304f
```
回归：同 7.5 的 10 条相关测试继续通过。
### 7.7 g0035/g0036 挂载代码已具备 fingerprint 防线，但当前 decode 启动目录仍旧
远端检查：
```text
g0035 /sgl-workspace/sglang-tai fingerprint: 0c0e736fa6de304f
g0036 /sgl-workspace/sglang-tai fingerprint: 0c0e736fa6de304f
```
同时 `/sgl-workspace/sglang/python/sglang/srt/disaggregation/common/conn.py` 中没有 `get_runtime_source_fingerprint`，说明当前 decode 日志路径对应的旧目录仍不具备新防线。
结论：下一轮 ETE 必须确保 decode 启动工作目录/PYTHONPATH 指向 `/sgl-workspace/sglang-tai`。如果仍从 `/sgl-workspace/sglang` 启动，既不会使用最新修复，也不会触发 fingerprint fail-fast。
### 7.8 当前 decode 主进程 cwd 已确认
通过容器内 `/proc/<pid>/cwd`：
```text
g0035 decode PID=15207 cwd=/sgl-workspace/sglang
g0036 decode PID=14525 cwd=/sgl-workspace/sglang
```
这与日志栈路径一致，排除“只是日志引用旧路径”的可能。当前 accept len=1 / 乱码验证环境确实是 prefill 新代码 + decode 旧镜像代码混跑。
### 7.9 纠正：cwd 不是根因判据，必须以 Python 实际 import source 为准
用户指出：`python -m sglang.launch_server` 从 `/sgl-workspace/sglang` 启动并不必然有问题；如果安装位置或 `PYTHONPATH` 指向正确代码，cwd 可以无影响。
修正：前文 7.3/7.8 中“cwd=/sgl-workspace/sglang 因此一定旧代码”的结论过强。当前可作为证据的是：
- traceback/code object path 显示 decode 运行栈来自 `/sgl-workspace/sglang/python/sglang/...`；
- `/sgl-workspace/sglang` 与 `/sgl-workspace/sglang-tai` 的关键文件 hash 不一致；
- 但仍需确认运行进程 `PYTHONPATH`、editable install、`sglang.__file__` 的实际来源。
后续排查必须以 import source/安装位置为准，不再单独用 cwd 下结论。
### 7.10 再次修正：当前 Python install source 实际一致
检查方式：在 g0034/g0035/g0036 容器内分别从 `/sgl-workspace/sglang-tai` 和 `/sgl-workspace/sglang` 执行：
```python
import sglang
import sglang.srt.layers.attention.nsa_backend as nb
from sglang.srt.disaggregation.common.conn import get_runtime_source_fingerprint, get_runtime_source_root
```
结果均为：
```text
sglang.__file__ = /sgl-workspace/sglang-tai/python/sglang/__init__.py
nsa_backend.__file__ = /sgl-workspace/sglang-tai/python/sglang/srt/layers/attention/nsa_backend.py
source_root = /sgl-workspace/sglang-tai
fingerprint prefix = 0c0e736fa6de304f
```
结论：当前环境下 cwd 不影响 `python -m` 的 import source；安装位置/PYTHONPATH 已指向 `/sgl-workspace/sglang-tai`。7.2/7.3/7.8 中将 `/sgl-workspace/sglang` 目录 hash 差异作为根因的判断不成立，只能保留 fingerprint fail-fast 作为未来防线。
当前根因需要回到 CP shared-KV / FP8 / bs>1 / page-tail / decode transfer 消费路径本身继续排查。
## 2026-06-04 note: why decode can regress even when CP is disabled
The current decode process does not enable CP shared-KV (`SGLANG_CP_DRAFT_SHARED_KV` is absent and decode server args show `enable_nsa_prefill_context_parallel=False`, `enable_nsa_prefill_cp_shared_kv=False`). That does **not** isolate it from recent CP-driven changes because several touched files are shared by normal decode:
- `speculative/eagle_worker.py` runs for any EAGLE decode. The recent `cp_local_hidden_states` marker is derived from `logits_output.draft_hidden_states is not None`, while the actual CP draft path is gated by `_can_use_cp_draft_shared_kv(batch)`. If `draft_hidden_states` is populated outside CP, decode may incorrectly treat hidden states as CP-local and skip static padding.
- `model_executor/forward_batch_info.py::_pad_inputs_to_size()` is generic static-padding code used by EAGLE/DP/MLP-sync, not a CP-only path.
- `models/deepseek_v2.py` always has the hidden-state capture hook (`capture_draft_hidden_states`) and can populate `logits_output.draft_hidden_states` when requested by EAGLE worker.
- Decode also consumes prefill-produced KV through disaggregation transfer, so prefill CP corruption can show up in decode even if decode itself is not CP-enabled. However, the latest evidence shows decode warmup itself already emits abnormal text, so the current primary suspicion is decode-side EAGLE/FP8/static-padding behavior rather than only prefill HiCache.
Current hypothesis to test next: `cp_local_draft_hidden_states` should be true only when `_can_use_cp_draft_shared_kv(batch)` is true, not merely when `logits_output.draft_hidden_states` exists.
## 2026-06-04 prefill KV corruption lead: current reuse uses global locs in CP path
New primary lead after re-centering on prefill-produced KV:
- Decode may not enable CP, but disagg decode consumes prefill-produced KV. Bad prefill CP shared-KV writes/composes can directly cause low EAGLE accept length and garbage decode text.
- In `nsa_backend.py`, MLA partial/current reuse has two branches:
- compute-padding branch uses `select_cp_local_valid_rows_for_cache_write()` and `get_cp_shared_kv_local_out_cache_loc()`.
- non-compute-padding branch still uses `k[:valid_current_rows]`, `k_rope[:valid_current_rows]`, and `forward_batch.out_cache_loc[:valid_current_rows]`.
- In CP shared-KV, the local current rows and local locs must come from the same CP split contract. Using the global `out_cache_loc` prefix can mismatch this rank's local rows and produce a bad page-slot compose.
- Existing unit test explicitly expected the old global loc prefix; changed it to RED against this bug.
- Separate code defect found: `_cp_split_and_rebuild_batch_in_seq()` passes `input_` twice to `split_tensor_by_cp_batch_plan`, inconsistent with the current signature. This should be fixed in the same CP split consistency pass.
## 2026-06-04 fix applied: CP-local current reuse for MLA and index
Applied local fixes for the prefill-produced KV corruption hypothesis:
- `nsa_backend.py` MLA partial/current reuse no longer uses `forward_batch.out_cache_loc[:valid_current_rows]` in the non-compute-padding branch. It now obtains `get_cp_shared_kv_local_out_cache_loc()` and ensures current KV rows are CP-local, using `cp_split_and_rebuild_data()` when the source tensors are still gathered/global.
- `nsa_indexer.py` index current reuse now follows the same contract: local locs + local key rows. The old non-compute-padding path used global `out_cache_loc` prefix and global `key[:valid_current_rows]`.
- `cp_split_and_rebuild_data()` now preserves all non-token dimensions (`view(-1, *shape[1:])`) so MLA/index KV tensors do not collapse from 3D to 2D when CP-split.
Verification so far:
- RED observed for MLA source-gate test before code change.
- RED observed for index source-gate test before code change.
- Local py_compile passed for touched code.
- Local targeted pytest passed for both source-gate tests.
- Remote py_compile and targeted pytest pending after sync.
Remote verification after sync:
- `python -m py_compile` passed for `nsa_indexer.py`, `nsa_backend.py`, `nsa/utils.py`, and the runtime test file in `g0034` container.
- Remote targeted pytest passed:
- `test_index_current_reuse_prepare_accepts_padded_out_cache_loc`
- `test_mla_current_reuse_gate_accepts_padded_out_cache_loc`
- `test_fp8_fused_mla_store_matches_sglang_fallback_for_cp_owned_tail_pages`
- Remote minimal script verified `cp_split_and_rebuild_data()` preserves 3D shape `(tokens, heads, dim)`.
- Full `test_nsa_cp_utils.py` import is currently blocked by remote `sgl_kernel` missing `gelu_tanh_and_mul`; keep using targeted/minimal verification until the extension install is corrected.
Additional same-class fix:
- `_maybe_materialize_shared_index_buffer()` also had a non-compute-padding partial-current compose branch using global `forward_batch.out_cache_loc` and slicing `current_index_kv` by global rows. After `forward_cuda()` now prepares CP-local `current_index_kv`, this compose function must consume CP-local locs as well.
- Added a stricter source-gate test anchored inside `_maybe_materialize_shared_index_buffer()`; observed RED before the fix and GREEN after the fix.
## 2026-06-04 full prefill-flow audit pass: confirmed flow boundaries
Checked the end-to-end prefill path again from scheduler to decode handoff.
Code evidence:
- `disaggregation/prefill.py::event_loop_normal_disagg_prefill()` runs `recv -> bootstrap_queue -> get_new_batch_prefill/process_chunk -> run_batch -> process_batch_result -> process_inflight_queue`.
- `process_batch_result_disagg_prefill()` appends the first target token, inserts/locks radix via `tree_cache.cache_unfinished_req(req)`, queues KV transfer, and copies EAGLE handoff tensors from `batch.spec_info` into the request before `send_kv_chunk(last_chunk=True)`.
- `send_kv_chunk()` transfers page ids derived from `req_to_token_pool.req_to_token[start:end]`; non-last chunks floor to page boundary, last chunks use `len(origin_input_ids)`.
- `CPSharedPagedTokenToKVPoolAllocator` explicitly stores **logical** page ids in scheduler/radix/`req_to_token_pool`; physical translation is deferred to `CpSharedKVLayout`. Therefore the Mooncake `filter_kv_pages_for_cp_shared_kv()` path is internally consistent **if** no physical page loc leaks into `req_to_token_pool`.
Remote evidence from `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260604_003716.log`:
- Prefill was launched with `enable_cp_shared_kv_prefill_bs_gt1=True`, `kv_cache_dtype='fp8_e4m3'`, `hicache_io_backend='direct'`, `hicache_mem_layout='page_first_direct'`, `chunked_prefill_size=65536`.
- For the 40k prompt, prefill logged `logical_len=40386` but `#new-token: 40448`, matching page-rounded compute/cache coverage.
- The hot path still emits repeated `[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] ... error=current_index_k must be uint8`, so current-index TAI fast path is not being used.
Conclusion: the transfer page filtering path is not the first suspect by itself. The stronger correctness suspects remain the CP shared-KV current/index materialization and the EAGLE handoff tensors produced by prefill.
## 2026-06-04 high-risk finding: EAGLE prefill handoff reads CP-local hidden as if it were batch-major
This is a separate high-risk accept-length lead from KV/index corruption.
Code evidence:
- `speculative/eagle_worker.py::forward_target_extend()` sets `model_worker_batch.capture_draft_hidden_states = True` when CP draft shared-KV is enabled. The target model then captures `logits_output.draft_hidden_states` before CP output collection.
- `models/deepseek_v2.py` captures `forward_batch.draft_hidden_states = hidden_states` at line-level before `cp_collect_last_token_hidden()` / `cp_all_gather_rerange_output()`. Under NSA prefill CP, that tensor is CP-local token-major, not batch-major last-token hidden.
- `disaggregation/prefill.py::process_batch_result_disagg_prefill()` currently stores `req.hidden_states_tensor = batch.spec_info.hidden_states[i].cpu().clone()`. That assumes `batch.spec_info.hidden_states` is `[batch, hidden]`.
- In the CP-local draft path, `batch.spec_info.hidden_states` can be `[local_tokens, hidden]`. For `bs=1`, indexing `[0]` stores the first local token hidden, not the request's last-token hidden. For `bs>1`, this is even less likely to match the correct per-request handoff row.
- Decode-side `decode_schedule_batch_mixin.py` stacks `req.hidden_states_tensor` into `EagleDraftInput(hidden_states=[bs, hidden])`, so the wrong prefill row becomes the draft model side-channel.
Impact:
- This can directly explain low EAGLE accept length, because the draft model receives the wrong prefill hidden state.
- It does **not** by itself prove target KV corruption; target verification should reject bad draft tokens. Therefore it explains accept-len collapse more strongly than garbage target text.
Required fix direction:
- Prefill EAGLE handoff must explicitly collect the per-request last-token hidden from the CP-local draft hidden tensor before copying it to `req.hidden_states_tensor`.
- It should use the same owner/local-offset contract as `cp_collect_last_token_hidden()` / batch metadata, not `hidden_states[i]`.
- Add a unit test where `hidden_states` is token-major CP-local and `bs=1/bs>1`; verify the stored hidden row equals the true request last token, not row 0.
## 2026-06-04 high-risk finding: current-only index top-k uses global request offsets on CP-local current_index_kv
Code evidence:
- `nsa_indexer.py::forward_cuda()` now prepares `current_index_kv` from CP-local current rows using `get_cp_shared_kv_local_out_cache_loc()` and either local `key` rows or `cp_split_and_rebuild_data()`.
- `_get_topk_in_seq_cp_pair_batch()` passes `current_index_kv` directly to `_get_topk_ragged_with_cp()` when `current_only_batch=True`.
- `_get_topk_ragged_with_cp()` builds `current_req_offsets` by accumulating global `forward_batch.extend_seq_lens_cpu`.
- It then passes these global offsets as `current_bases` to the TAI current-index prepare path and uses the same offsets in the Python fallback slice.
Why this is wrong:
- After the recent CP-local current-index fix, `current_index_kv` is no longer global flattened request rows. It is local flattened rows in CP split order.
- For `bs>1`, request 1+ local base is not the sum of previous global extend lengths. It must come from batch-plan CP-local request offsets (`request_rank_local_offsets` / valid split metadata), and it must account for compute padding only where the local source tensor includes compute-padded rows.
Additional concrete code bug:
- In the Python fallback loop, the check `if pre_chunk_offset != 0` references `pre_chunk_offset` from the previous segment-building loop, not the current `segment_record`. This can test the wrong request/segment and is unsafe even before performance tuning.
Impact:
- This can corrupt index top-k for current-only cache-miss/prefill bs>1 paths.
- Remote fallback logs show this path is active: `[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] reason=current_batch_kernel_failed ... segments=2 total_kv_len≈43k total_q_count=5056`.
- This is currently one of the strongest candidates for bad prefill-produced KV/index state under bs>1.
Required fix direction:
- Build a per-request CP-local current row base vector from `CPSharedKVBatchPlan` and pass that to both TAI and Python current-index top-k paths.
- Include `pre_chunk_offset` in `segment_records` instead of reading the loop variable after the fact.
- Add a unit test with two requests whose global extend lengths differ and whose CP-local row counts differ; assert `_get_topk_ragged_with_cp()` slices from the local base, not global base.
## 2026-06-04 performance/correctness finding: current-index TAI fallback is caused by dtype contract mismatch
Code evidence:
- `tai-kernel/python/tai_kernel/nsa_prefill/cp_index_mqa_prepare.py::_validate_current_batch_inputs()` requires `current_index_k.dtype == torch.uint8`.
- `nsa_indexer.py::forward_cuda()` stores `current_k_fp8` as a `torch.float8_e4m3fn` tensor in `current_index_kv`.
- `cp_shared_kv_runtime.py::try_tai_prepare_cp_mqa_current_index_batch()` passes `current_index_k` through directly, so the TAI path rejects it with `current_index_k must be uint8`.
Remote evidence:
- Latest prefill logs contain repeated `current_index_k must be uint8` warnings across ranks and layers.
Impact:
- This is definitely a hot-path performance regression.
- Because the fallback path currently has the local/global offset issue above, this dtype mismatch also increases exposure to the buggy Python fallback.
Required fix direction:
- Either pass `current_index_k.view(torch.uint8)` into the TAI wrapper, or update the TAI wrapper to accept float8 and reinterpret internally.
- Keep the Python fallback as fail-fast/diagnostic for CUDA hot path until local-base correctness is fixed.
## 2026-06-04 HiCache/chunked prefill stale-tail status
Code evidence:
- `hiradix_cache.py::_build_prepare_write_backup_candidate_for_req()` floors chunked-prefill backup end to a completed page and skips backup when no full page is available.
- `_cp_prune_stale_tail_after_page_floor()` still raises `HiCachePendingBackupSplit` if a stale tail subtree has lock refs, host refs, or pending host write state. It drains finished write acks before declaring blocked, but it still fails if the state remains unprunable.
- `match_prefix()` catches `HiCachePendingBackupSplit` and returns a deferred node in match results, but insertion/cache-unfinished paths can still surface this if chunked insertion attempts to prune while backup remains pending.
Impact:
- The previous chunked-prefill crash is mitigated for some cases by page-floor backup, but not eliminated by construction.
- Risk remains when chunked prefill extends a node whose stale tail is still protected or backup-pending.
Required fix direction:
- Insert/cache-unfinished should defer and retry this request, not crash the scheduler, when stale-tail prune is blocked by in-flight backup.
- A regression test should simulate: first chunk creates sub-page valid tail, async backup pending, next chunk extends past page boundary; expected result is deferred insertion/rollback, not uncaught `HiCachePendingBackupSplit`.
## 2026-06-04 prefetch timing audit status
Code evidence:
- `nsa_backend.py::_maybe_start_cp_shared_kv_attention_prefetch()` is now called immediately before the current layer attention kernel.
- Both index and MLA prefetchers materialize the next-layer prefix on the current stream first, then enqueue the CP all-reduce on a shared prefetch stream.
- `wait_attention_window()` is deliberately deferred; the real wait happens when the next layer consumes the prefetched handle. This matches the intended "next layer consume waits" model.
- Draft/EAGLE inputs do not create MLA/index prefetchers (`disable_draft_prefetch = cp_shared_kv_is_draft_input(forward_batch)`).
Risks:
- The prefix materialize step still allocates and fills a dense buffer before the async all-reduce. For large prefixes this is non-trivial CPU/GPU launch overhead and can still contend with current-layer work if the hook is placed too early or if the dense buffer allocation is expensive.
- The async collective is still an all-reduce over dense materialized prefix rows/pages. Even when overlapped, it consumes communication resources and can interfere with other CP/TP collectives.
- Because wait is deferred, a bad handle/key mismatch should fail at consume. Missing consume-side fail-fast coverage remains important.
Current conclusion:
- Prefetch timing is not the strongest correctness suspect for garbage output in the latest logs. It remains a performance and hang-risk area, especially under long prefixes and bs>1, but the latest concrete correctness evidence points more directly at EAGLE handoff and current-index top-k local/global offsets.
## 2026-06-04 P0/P1 repair plan
Scope for this repair pass:
1. P0 EAGLE prefill handoff correctness
- Add a small helper that converts CP-local draft hidden states to the per-request hidden rows expected by disaggregated decode.
- Use batch-plan last-token owner/local-offset metadata when available; otherwise reuse the single-request in-seq last-token owner calculation.
- Keep non-CP/full hidden behavior unchanged.
- Fail fast if CP-local hidden is marked but required CP metadata is missing or points outside the local hidden tensor.
2. P0 current-only index top-k bs>1 local base correctness
- Stop deriving `current_req_offsets` from global `extend_seq_lens_cpu` when `current_index_kv` is CP-local.
- Derive per-request local current bases from `CPSharedKVBatchPlan` rank-local valid metadata.
- Store `pre_chunk_offset` in each segment record so fallback checks the segment it is actually processing.
3. P1 TAI current-index dtype contract
- Pass the current-index K tensor to tai-kernel as uint8 bytes, matching `tai_kernel.nsa_prefill.cp_index_mqa_prepare` validation.
- Keep the returned buffer interpreted as fp8 at the callsite.
Tests to add before production changes:
- EAGLE CP-local handoff selects the true last-token local row for bs=1 and bs>1, not `hidden_states[i]`.
- Current-index top-k local base builder returns CP-local bases for bs>1 and does not use global extend offsets.
- TAI wrapper receives uint8 current-index K for fp8 input.
## 2026-06-04 P0 refinement: EAGLE hidden handoff marker must be reset after draft forward
Additional code-flow evidence:
- Target `DeepseekV3ForCausalLM.forward_core()` stores `forward_batch.draft_hidden_states` before CP output collect so the draft model can consume CP-local target hidden rows.
- `EagleWorker.forward_draft_extend()` marks the temporary `EagleDraftInput.hidden_states` as `cp_local_hidden_states=True` while the draft model consumes that CP-local side channel.
- `DeepseekModelNextN.forward()` then collects the draft model output hidden via `cp_collect_last_token_hidden()` when using CP-local draft input, so `logits_output.hidden_states` from the draft forward is expected to be batch-major `[bs, hidden]`.
- `EagleWorker.capture_for_decode()` overwrites `draft_input.hidden_states` with this draft output, but currently leaves `draft_input.cp_local_hidden_states=True`.
Correction to the earlier suspect:
- The prefill handoff does not directly copy the original target CP-local hidden side channel after draft forward; it copies `batch.spec_info.hidden_states` after `capture_for_decode()`.
- The remaining P0 risk is the stale semantic marker: after `capture_for_decode()`, the tensor is no longer CP-local, so leaving `cp_local_hidden_states=True` can make later static-padding / draft-input consumers treat a batch-major hidden tensor as CP-local.
Required fix direction:
- Reset `EagleDraftInput.cp_local_hidden_states=False` immediately after `capture_for_decode()` replaces `hidden_states` with draft output hidden.
- Add a unit test that starts with `cp_local_hidden_states=True`, calls `capture_for_decode()`, and verifies the hidden tensor is replaced by draft output while the CP-local marker is cleared.
## 2026-06-04 P0/P1 implementation status
Implemented in this pass:
1. EAGLE draft hidden marker cleanup
- `EagleWorker.capture_for_decode()` now clears `EagleDraftInput.cp_local_hidden_states` after replacing the temporary target CP-local side-channel with the draft model output hidden.
- This keeps the contract precise: CP-local marker is only valid before/during draft model consumption, not after decode handoff hidden is batch-major.
2. Current-index top-k CP-local bases
- `_get_topk_ragged_with_cp()` now derives current-index request bases via `_build_current_index_request_bases()`.
- Legacy/no-plan current tensors still use global extend offsets.
- CP shared-KV batch-plan current tensors use `request_valid_rank_local_offsets` first, then `request_rank_local_offsets`, because the current-index K/S buffers are flattened CP-rank-local valid rows.
- Missing local bases under a batch plan fail fast instead of silently using global offsets.
3. TAI current-index dtype ABI
- The TAI current-index batch prepare path now passes fp8 K tensors as `uint8` views via `_current_index_k_for_tai()`.
- Existing uint8 callers are unchanged; unknown dtypes still pass through to the wrapper, which keeps its own validation/fallback behavior.
Regression coverage added:
- `test_indexer_ragged_cp_index_current_batch_uses_cp_local_bases_and_uint8_k`
- RED before fix: TAI received `torch.float8_e4m3fn` and global bases `[0, 5]`.
- GREEN after fix: TAI receives `torch.uint8` and CP-local bases `[0, 2]`.
- `test_eagle_capture_for_decode_clears_cp_local_hidden_marker`
- RED before fix: marker stayed `True` after capture.
- GREEN after fix: marker is cleared while hidden points to the draft output tensor.
Verification:
- Remote container targeted RED was observed before production changes: both new tests failed for the expected reasons.
- Remote container after fix:
- `python -m py_compile python/sglang/srt/layers/attention/nsa/nsa_indexer.py python/sglang/srt/speculative/eagle_worker.py test/registered/unit/layers/test_nsa_cp_utils.py`
- targeted current-index/EAGLE tests: `2 passed`
- current-index related subset: `6 passed`
- full `test/registered/unit/layers/test_nsa_cp_utils.py`: `87 passed`
Remaining runtime risk:
- This fixes two concrete CP shared-KV correctness hazards, but it does not prove the remote ETE garbage-output issue is fully resolved. The next ETE run should specifically check that the `current_index_k must be uint8` fallback warning disappears and that EAGLE accept length recovers above the previous ~1-2 range.
Post-cleanup verification note:
- `_current_index_k_for_tai()` was made robust to PyTorch builds that may not expose every fp8 dtype symbol.
- Re-ran remote full `test/registered/unit/layers/test_nsa_cp_utils.py`: `87 passed`.