diff --git a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md new file mode 100644 index 000000000..e1667c12f --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md @@ -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 / 重复 +``` + +同时 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//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`. diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 20b1424e4..796bf92b6 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -2,11 +2,14 @@ from __future__ import annotations import asyncio import dataclasses +import hashlib +import importlib.util import logging import threading import time from collections import defaultdict from functools import cache +from pathlib import Path from typing import Dict, List, Optional, Set, Tuple, Union import numpy as np @@ -44,6 +47,75 @@ from sglang.srt.utils.network import ( logger = logging.getLogger(__name__) +_RUNTIME_FINGERPRINT_MODULES = ( + "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", +) + + +@cache +def _get_runtime_python_root() -> Optional[Path]: + spec = importlib.util.find_spec("sglang") + origin = getattr(spec, "origin", None) if spec is not None else None + if origin is None: + return None + path = Path(origin).resolve() + # ...//python/sglang/__init__.py + for parent in path.parents: + if parent.name == "python": + return parent + return path.parent.parent + + +@cache +def get_runtime_source_root() -> Optional[str]: + python_root = _get_runtime_python_root() + if python_root is None: + return None + return str(python_root.parent) + + +@cache +def get_runtime_source_fingerprint() -> Optional[str]: + """Return a stable fingerprint for code that defines the P/D KV contract. + + Prefill and decode may run from different mounted source trees in long-lived + containers. CP shared-KV + FP8 + page-tail handling is not wire-compatible + across our recent development snapshots, so decode must fail fast when the + source contract differs instead of silently consuming malformed KV. + + Do not import the target modules while fingerprinting: some package + ``__init__`` files import the connection modules themselves. Build paths + from the already-imported ``sglang`` package root instead. + """ + + python_root = _get_runtime_python_root() + if python_root is None: + return None + + digest = hashlib.sha256() + for module_name in _RUNTIME_FINGERPRINT_MODULES: + origin = python_root.joinpath(*module_name.split(".")).with_suffix(".py") + digest.update(module_name.encode("utf-8")) + digest.update(b"\0") + try: + digest.update(origin.read_bytes()) + except OSError as exc: + logger.warning( + "Failed to read source for disaggregation runtime fingerprint: " + "module=%s origin=%s error=%s", + module_name, + origin, + exc, + ) + return None + digest.update(b"\0") + return digest.hexdigest() + + @dataclasses.dataclass class PrefillServerInfo: # Topology fields (fetched from bootstrap server) @@ -54,6 +126,8 @@ class PrefillServerInfo: page_size: Optional[int] kv_cache_dtype: Optional[str] follow_bootstrap_room: bool + runtime_source_fingerprint: Optional[str] = None + runtime_source_root: Optional[str] = None # Pre-computed rank mapping (set by try_ensure_parallel_info on decode side) target_tp_rank: Optional[int] = None @@ -73,6 +147,16 @@ class PrefillServerInfo: str(self.kv_cache_dtype) if self.kv_cache_dtype is not None else None ) self.follow_bootstrap_room = bool(self.follow_bootstrap_room) + self.runtime_source_fingerprint = ( + str(self.runtime_source_fingerprint) + if self.runtime_source_fingerprint is not None + else None + ) + self.runtime_source_root = ( + str(self.runtime_source_root) + if self.runtime_source_root is not None + else None + ) @dataclasses.dataclass @@ -236,6 +320,22 @@ class CommonKVManager(BaseKVManager): f"Both servers must use the same --kv-cache-dtype value." ) + local_source_fingerprint = get_runtime_source_fingerprint() + if ( + info.runtime_source_fingerprint is not None + and local_source_fingerprint is not None + and info.runtime_source_fingerprint != local_source_fingerprint + ): + raise RuntimeError( + "Runtime source fingerprint mismatch between prefill and decode. " + "This can corrupt disaggregated KV transfer contracts. " + f"prefill_fingerprint={info.runtime_source_fingerprint}, " + f"decode_fingerprint={local_source_fingerprint}, " + f"prefill_source_root={info.runtime_source_root}, " + f"decode_source_root={get_runtime_source_root()}. " + "Restart both prefill and decode from the same source tree." + ) + self._resolve_rank_mapping(info) self.prefill_info_table[bootstrap_addr] = info logger.debug(f"Prefill parallel info for [{bootstrap_addr}]: {info}") @@ -348,6 +448,8 @@ class CommonKVManager(BaseKVManager): "page_size": self.kv_args.page_size, "kv_cache_dtype": self.server_args.kv_cache_dtype, "load_balance_method": self.server_args.load_balance_method, + "runtime_source_fingerprint": get_runtime_source_fingerprint(), + "runtime_source_root": get_runtime_source_root(), } try: @@ -668,6 +770,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): self.page_size = None self.kv_cache_dtype: Optional[str] = None self.follow_bootstrap_room: Optional[bool] = None + self.runtime_source_fingerprint: Optional[str] = get_runtime_source_fingerprint() + self.runtime_source_root: Optional[str] = get_runtime_source_root() self.prefill_port_table: Dict[ int, Dict[int, Dict[int, Dict[int, PrefillRankInfo]]] ] = {} @@ -734,6 +838,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): rank_port = int(data["rank_port"]) page_size = int(data["page_size"]) kv_cache_dtype = data["kv_cache_dtype"] + runtime_source_fingerprint = data.get("runtime_source_fingerprint") + runtime_source_root = data.get("runtime_source_root") if self.attn_tp_size is None: self.attn_tp_size = attn_tp_size @@ -753,6 +859,15 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): if self.kv_cache_dtype is None and kv_cache_dtype is not None: self.kv_cache_dtype = kv_cache_dtype + if ( + self.runtime_source_fingerprint is None + and runtime_source_fingerprint is not None + ): + self.runtime_source_fingerprint = str(runtime_source_fingerprint) + + if self.runtime_source_root is None and runtime_source_root is not None: + self.runtime_source_root = str(runtime_source_root) + if self.follow_bootstrap_room is None: load_balance_method = data.get( "load_balance_method", "follow_bootstrap_room" @@ -822,6 +937,8 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): if self.follow_bootstrap_room is not None else True ), + runtime_source_fingerprint=self.runtime_source_fingerprint, + runtime_source_root=self.runtime_source_root, ) return web.json_response(dataclasses.asdict(info), status=200) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index b3469e407..a8d6a86a0 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -68,6 +68,7 @@ from sglang.srt.distributed.parallel_state import get_pp_group from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.nsa.utils import ( cp_all_gather_rerange_output, + cp_split_and_rebuild_data, get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, get_cp_shared_kv_local_physical_out_cache_loc, @@ -132,6 +133,65 @@ class BatchTopKQueryLengths: compute_token_count: int +def _current_index_k_for_tai(current_index_k: torch.Tensor) -> torch.Tensor: + """Return the current-index K tensor in the uint8 ABI expected by tai-kernel.""" + + if current_index_k.dtype == torch.uint8: + return current_index_k + fp8_dtypes = { + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + getattr(torch, "float8_e4m3fnuz", None), + getattr(torch, "float8_e5m2fnuz", None), + ) + if dtype is not None + } + if current_index_k.dtype in fp8_dtypes: + return current_index_k.view(torch.uint8) + return current_index_k + + +def _build_current_index_request_bases(forward_batch: ForwardBatch) -> List[int]: + """Build request bases for flattened current-index K/S rows. + + Without CP shared-KV batch metadata the legacy current-index tensor is a + global flattened `[req0 current, req1 current, ...]` buffer, so global + extend offsets are correct. With CP shared-KV batch metadata the tensor is + CP-rank-local valid rows, so request bases must come from the local valid + owner-lane plan instead of global extend lengths. + """ + + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + if extend_seq_lens_cpu is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=current_index_missing_extend_lens" + ) + + batch_plan = get_cp_shared_kv_batch_plan(forward_batch) + if batch_plan is not None: + offsets = getattr(batch_plan, "request_valid_rank_local_offsets", None) + if offsets is None: + offsets = getattr(batch_plan, "request_rank_local_offsets", None) + if offsets is None or len(offsets) != len(extend_seq_lens_cpu): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=current_index_missing_local_bases " + f"batch_size={len(extend_seq_lens_cpu)} " + f"offsets={offsets}" + ) + return [int(x) for x in offsets] + + current_req_offsets = [] + current_cursor = 0 + for extend_len in extend_seq_lens_cpu: + current_req_offsets.append(current_cursor) + current_cursor += int(extend_len) + return current_req_offsets + + def _select_batch_topk_query_lengths( *, cp_metadata, @@ -503,29 +563,33 @@ class Indexer(MultiPlatformOp): f"current_locs_shape={tuple(current_locs.shape)}" ) else: - current_locs = forward_batch.out_cache_loc - valid_current_rows = current_extend_kv_rows_for_reuse( - forward_batch, - current_index_kv[0], - current_index_kv[1], - ) - if valid_current_rows is None: + current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) + if current_locs is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " + "CP shared KV index partial-current compose requires local " + "valid out_cache_loc. " + f"cp_rank={layout.cp_rank} layer_id={layer_id} " + f"prefix_lens={prefix_lens} extend_lens={extend_lens} " + f"current_k_shape={tuple(current_index_kv[0].shape)} " + f"current_scale_shape={tuple(current_index_kv[1].shape)}" + ) + valid_current_rows = int(current_locs.numel()) + if ( + int(current_index_kv[0].shape[0]) != valid_current_rows + or int(current_index_kv[1].shape[0]) != valid_current_rows + ): raise RuntimeError( "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " "CP shared KV index partial-current compose received " - "current_index_kv that does not satisfy current reuse " - "metadata. " + "current_index_kv rows that do not match local valid " + "out_cache_loc. " f"cp_rank={layout.cp_rank} layer_id={layer_id} " f"prefix_lens={prefix_lens} extend_lens={extend_lens} " f"current_k_shape={tuple(current_index_kv[0].shape)} " f"current_scale_shape={tuple(current_index_kv[1].shape)} " f"out_cache_loc_shape={tuple(current_locs.shape)}" ) - current_locs = current_locs[:valid_current_rows] - current_index_kv = ( - current_index_kv[0][:valid_current_rows], - current_index_kv[1][:valid_current_rows], - ) prefix_slot_span = None if len(prefix_lens_cpu) == 1: prefix_pages = int(prefix_lens_cpu[0]) // page_size @@ -1333,13 +1397,11 @@ class Indexer(MultiPlatformOp): if cp_index is not None: current_req_offsets: Optional[List[int]] = None if current_index_kv is not None: - current_req_offsets = [] - current_cursor = 0 - for extend_len in forward_batch.extend_seq_lens_cpu: - current_req_offsets.append(current_cursor) - current_cursor += int(extend_len) + current_req_offsets = _build_current_index_request_bases( + forward_batch + ) - segment_records: List[Tuple[int, int, int, int, int, int, int]] = [] + segment_records: List[Tuple[int, int, int, int, int, int, int, int]] = [] batch_idx_list = [] kv_lens_list = [] q_starts_list = [] @@ -1374,6 +1436,7 @@ class Indexer(MultiPlatformOp): kv_len_i, k_cursor, q_cursor, + int(pre_chunk_offset), ) ) batch_idx_list.append(batch_idx) @@ -1455,7 +1518,7 @@ class Indexer(MultiPlatformOp): int(current_index_kv[0].reshape(current_index_kv[0].shape[0], -1).shape[1]), ) tai_current_prepared = try_tai_prepare_cp_mqa_current_index_batch( - current_index_k=current_index_kv[0], + current_index_k=_current_index_k_for_tai(current_index_kv[0]), current_index_scale=current_index_kv[1], current_bases=torch.tensor( current_bases_list, dtype=torch.int32, device=descriptor_device @@ -1516,6 +1579,7 @@ class Indexer(MultiPlatformOp): _kv_len_i, segment_k_base, _segment_q_base, + pre_chunk_offset, ) in segment_records: if current_index_kv is None: k_fp8 = index_buf_accessor.GetK.execute( @@ -2437,19 +2501,22 @@ class Indexer(MultiPlatformOp): int(current_locs.numel()) if current_locs is not None else None ) else: - valid_current_rows = current_extend_kv_rows_for_reuse( + global_valid_current_rows = current_extend_kv_rows_for_reuse( forward_batch, key ) - current_key = ( - key[:valid_current_rows] - if valid_current_rows is not None - else None - ) - current_locs = ( - forward_batch.out_cache_loc[:valid_current_rows] - if valid_current_rows is not None - else None + current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) + valid_current_rows = ( + int(current_locs.numel()) if current_locs is not None else None ) + current_key = None + if global_valid_current_rows is not None and valid_current_rows is not None: + if int(local_key.shape[0]) == valid_current_rows: + current_key = local_key + else: + current_key = cp_split_and_rebuild_data( + forward_batch, + key[: int(global_valid_current_rows)].contiguous(), + ) if ( valid_current_rows is not None and current_key is not None diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index bbc42ea0d..0308afeef 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -1558,7 +1558,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): ) result = torch.cat( [input_list[i] for i in metadata.zigzag_index], dim=0 - ).view(-1, input_.shape[-1]) + ).view(-1, *input_.shape[1:]) return result diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index e809d041e..a2c4104ea 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -50,6 +50,7 @@ from sglang.srt.layers.attention.nsa.transform_index import ( ) from sglang.srt.layers.attention.nsa.utils import ( can_nsa_prefill_cp_round_robin_split, + cp_split_and_rebuild_data, compute_nsa_seqlens, get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, @@ -1883,11 +1884,41 @@ class NativeSparseAttnBackend( assert current_kv_rows_for_reuse is not None valid_current_rows = int(current_kv_rows_for_reuse) if not compute_padding_current: - current_k_nope = k[:valid_current_rows] - current_k_rope = k_rope[:valid_current_rows] - current_locs_for_reuse = forward_batch.out_cache_loc[ - :valid_current_rows - ] + current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc( + forward_batch + ) + if current_locs_for_reuse is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_current_reuse_locs] " + "CP shared KV MLA current reuse requires local " + "out_cache_loc in the prefill CP path." + ) + local_current_rows = int(current_locs_for_reuse.numel()) + if int(k.shape[0]) == local_current_rows and int( + k_rope.shape[0] + ) == local_current_rows: + current_k_nope = k + current_k_rope = k_rope + else: + current_k_nope = cp_split_and_rebuild_data( + forward_batch, k[:valid_current_rows].contiguous() + ) + current_k_rope = cp_split_and_rebuild_data( + forward_batch, k_rope[:valid_current_rows].contiguous() + ) + if ( + int(current_k_nope.shape[0]) != local_current_rows + or int(current_k_rope.shape[0]) != local_current_rows + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_current_reuse_shape] " + "CP shared KV MLA current rows do not match local " + "out_cache_loc after CP split. " + f"k_rows={int(current_k_nope.shape[0])} " + f"k_rope_rows={int(current_k_rope.shape[0])} " + f"local_locs={local_current_rows} " + f"valid_current_rows={valid_current_rows}" + ) assert current_k_nope is not None assert current_k_rope is not None assert current_locs_for_reuse is not None diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index aaf9e0cce..b1044dd9a 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -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() diff --git a/test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py b/test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py new file mode 100644 index 000000000..f6174d19f --- /dev/null +++ b/test/registered/unit/disaggregation/test_common_conn_runtime_fingerprint.py @@ -0,0 +1,93 @@ +import types +import unittest +from unittest.mock import Mock, patch + +from sglang.srt.disaggregation.common import conn as common_conn + + +class TestDisaggregationRuntimeFingerprint(unittest.TestCase): + def _manager(self): + mgr = common_conn.CommonKVManager.__new__(common_conn.CommonKVManager) + mgr.prefill_info_table = {} + mgr.kv_args = types.SimpleNamespace(page_size=64, engine_rank=0) + mgr.server_args = types.SimpleNamespace(kv_cache_dtype="fp8_e4m3") + mgr._resolve_rank_mapping = Mock() + return mgr + + def _response(self, payload, status_code=200): + response = Mock() + response.status_code = status_code + response.json.return_value = payload + response.text = "" + return response + + def _payload(self, **overrides): + payload = { + "attn_tp_size": 8, + "attn_cp_size": 8, + "dp_size": 1, + "pp_size": 1, + "page_size": 64, + "kv_cache_dtype": "fp8_e4m3", + "follow_bootstrap_room": True, + } + payload.update(overrides) + return payload + + def test_try_ensure_parallel_info_rejects_runtime_source_fingerprint_mismatch(self): + mgr = self._manager() + with patch.object( + common_conn.requests, + "get", + return_value=self._response( + self._payload( + runtime_source_fingerprint="prefill-source-hash", + runtime_source_root="/sgl-workspace/sglang-tai", + ) + ), + ), patch.object( + common_conn, + "get_runtime_source_fingerprint", + return_value="decode-source-hash", + ), patch.object( + common_conn, + "get_runtime_source_root", + return_value="/sgl-workspace/sglang", + ): + with self.assertRaisesRegex( + RuntimeError, + "Runtime source fingerprint mismatch", + ): + mgr.try_ensure_parallel_info("g0034:18992") + + mgr._resolve_rank_mapping.assert_not_called() + self.assertEqual(mgr.prefill_info_table, {}) + + def test_try_ensure_parallel_info_accepts_matching_runtime_source_fingerprint(self): + mgr = self._manager() + with patch.object( + common_conn.requests, + "get", + return_value=self._response( + self._payload( + runtime_source_fingerprint="same-source-hash", + runtime_source_root="/sgl-workspace/sglang-tai", + ) + ), + ), patch.object( + common_conn, + "get_runtime_source_fingerprint", + return_value="same-source-hash", + ), patch.object( + common_conn, + "get_runtime_source_root", + return_value="/sgl-workspace/sglang-tai", + ): + self.assertTrue(mgr.try_ensure_parallel_info("g0034:18992")) + + mgr._resolve_rank_mapping.assert_called_once() + self.assertIn("g0034:18992", mgr.prefill_info_table) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index cd5f90fec..acb316552 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -1322,6 +1322,28 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2))) + def test_cp_split_and_rebuild_data_preserves_non_token_dimensions(self): + import torch + + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_extend_lens=[4, 9], + request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]], + request_zigzag_indices=[[0, 3], [0, 3]], + ) + ) + tensor = torch.arange(13 * 2 * 3, dtype=torch.float32).view(13, 2, 3) + + with patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ): + local = cp_split_and_rebuild_data(forward_batch, tensor) + + self.assertEqual(tuple(local.shape), (8, 2, 3)) + self.assertTrue(torch.equal(local[0], tensor[0])) + def test_cp_split_and_rebuild_data_uses_compute_padding_rows(self): import torch @@ -3309,6 +3331,122 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10]) self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14]) + def test_indexer_ragged_cp_index_current_batch_uses_cp_local_bases_and_uint8_k( + self, + ): + import contextlib + import torch + from types import SimpleNamespace + from unittest.mock import patch + + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + indexer._with_real_sm_count = lambda: contextlib.nullcontext() + + prepare_calls = [] + + def fake_prepare(**kwargs): + prepare_calls.append(kwargs) + total_kv_len = int(kwargs["total_kv_len"]) + return ( + torch.zeros((total_kv_len, 1), dtype=torch.uint8), + torch.zeros((total_kv_len,), dtype=torch.float32), + torch.tensor([0, 0, 3, 3, 3], dtype=torch.int32), + torch.tensor([2, 3, 2, 3, 4], dtype=torch.int32), + ) + + class Metadata: + def get_page_table_64(self): + raise AssertionError("current cp_index path must not materialize index pages") + + def topk_transform(self, logits, topk, **kwargs): + return ( + torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) + .view(-1, 1) + .repeat(1, topk) + ) + + forward_batch = SimpleNamespace( + token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1), + seq_lens_cpu=torch.tensor([5, 7], dtype=torch.int64), + extend_seq_lens_cpu=[5, 7], + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + batch_plan=SimpleNamespace( + request_rank_local_offsets=[0, 2], + request_valid_rank_local_offsets=[0, 2], + ), + ), + ) + current_index_k = ( + torch.arange(6, dtype=torch.float32) + .to(torch.float8_e4m3fn) + .view(6, 1) + ) + current_index_scale = torch.arange(6, dtype=torch.float32).view(6, 1) + current_index_kv = (current_index_k, current_index_scale) + + with patch.object( + nsa_indexer, + "try_tai_prepare_cp_mqa_current_index_batch", + side_effect=fake_prepare, + create=True, + ), patch.object( + nsa_indexer, + "deep_gemm", + SimpleNamespace( + fp8_mqa_logits=lambda q_fp8, kv_fp8, weights, ks, ke, clean_logits=False: torch.zeros( + (int(q_fp8.shape[0]), 8), dtype=torch.float32 + ) + ), + ): + result = Indexer._get_topk_ragged_with_cp( + indexer, + forward_batch, + layer_id=7, + q_fp8=torch.empty((5, 1), dtype=torch.float32), + weights=torch.empty((5, 1, 1), dtype=torch.float32), + metadata=Metadata(), + kv_len=0, + actual_seq_q=5, + cp_index=[(0, 1, 3), (1, 1, 4)], + current_index_kv=current_index_kv, + ) + + self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]) + self.assertEqual(len(prepare_calls), 1) + call = prepare_calls[0] + self.assertEqual(call["current_index_k"].dtype, torch.uint8) + self.assertEqual(call["current_bases"].tolist(), [0, 2]) + self.assertEqual(call["kv_lens"].tolist(), [3, 4]) + self.assertEqual(call["q_lens"].tolist(), [2, 3]) + + def test_eagle_capture_for_decode_clears_cp_local_hidden_marker(self): + import torch + + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.speculative.eagle_worker import EAGLEWorker + + worker = object.__new__(EAGLEWorker) + worker.topk = 1 + draft_input = EagleDraftInput( + hidden_states=torch.full((4, 2), -1.0), + cp_local_hidden_states=True, + ) + draft_output_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + logits_output = LogitsProcessorOutput( + next_token_logits=torch.tensor([[0.1, 0.9], [0.7, 0.3]]), + hidden_states=draft_output_hidden, + ) + + worker.capture_for_decode(logits_output, draft_input) + + self.assertIs(draft_input.hidden_states, draft_output_hidden) + self.assertFalse(draft_input.cp_local_hidden_states) + def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self): import torch diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 4a4f60f18..f536ccdd1 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -729,7 +729,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertIn("valid_current_rows=int(current_kv_rows_for_reuse)", body_source) self.assertIn("pack_current_mla_kv_for_reuse", body_source) - self.assertIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source) + self.assertIn("get_cp_shared_kv_local_out_cache_loc", body_source) + self.assertNotIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source) def test_flashmla_kv_current_only_reuse_keeps_page_slot_layout(self): from pathlib import Path @@ -755,12 +756,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): Path(__file__).resolve().parents[4] / "python/sglang/srt/layers/attention/nsa_backend.py" ).read_text() + method_start = source.index(" current_locs_for_reuse = None") branch_start = source.index("eagle_draft_mla_branch = \"partial_current_sync\"") branch_end = source.index("forward_partial_current_sync_compose", branch_start) branch_source = source[branch_start:branch_end] + local_locs_source = source[method_start:branch_end] self.assertIn("build_batch_prefix_slot_span", source) self.assertIn("prefix_slot_span=", branch_source) + self.assertIn("get_cp_shared_kv_local_out_cache_loc", local_locs_source) + self.assertNotIn("current_locs = forward_batch.out_cache_loc", branch_source) self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source) def test_nsa_backend_topk_transform_uses_effective_forward_impl(self): @@ -1504,12 +1509,18 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): Path(__file__).resolve().parents[4] / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py" ).read_text() - branch_start = source.index(" if current_index_kv is not None:") + method_start = source.index(" def _maybe_materialize_shared_index_buffer") + branch_start = source.index( + " if current_index_kv is not None:", method_start + ) branch_end = source.index(" return materialized, dense_pages", branch_start) branch_source = source[branch_start:branch_end] + branch_compact = "".join(branch_source.split()) self.assertIn("build_batch_prefix_slot_span", source) self.assertIn("prefix_slot_span=", branch_source) + self.assertIn("get_cp_shared_kv_local_out_cache_loc", branch_source) + self.assertNotIn("current_locs=forward_batch.out_cache_loc", branch_compact) self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source) def test_ipc_page_descriptor_builder_maps_slots_to_owner_physical_pages(self): @@ -2141,7 +2152,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prepare_source = source[start:end] self.assertIn("valid_current_rows", prepare_source) - self.assertIn("key[:valid_current_rows]", prepare_source) + self.assertIn("get_cp_shared_kv_local_out_cache_loc", prepare_source) + self.assertIn("cp_split_and_rebuild_data", prepare_source) + self.assertNotIn("forward_batch.out_cache_loc[:valid_current_rows]", prepare_source) self.assertNotIn( "key.shape[0] == forward_batch.out_cache_loc.numel()", prepare_source, @@ -2852,6 +2865,91 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): load_kernel.assert_not_called() self.assertTrue(any("debug_enabled" in message for message in logs.output)) + @unittest.skipIf(not torch.cuda.is_available(), "CUDA required") + def test_fp8_fused_mla_store_matches_sglang_fallback_for_cp_owned_tail_pages(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.layers.attention.nsa.quant_k_cache import ( + quantize_k_cache_separate, + ) + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton + + try: + runtime._load_tai_fused_mla_store_kernel() + except Exception as exc: + self.skipTest(f"TAI fused MLA store kernel unavailable: {exc}") + + page_size = 64 + cp_size = 8 + rank = 7 + n_tokens = page_size * 3 + 17 + layout = CpSharedKVLayout( + page_size=page_size, + cp_size=cp_size, + cp_rank=rank, + ) + local_token_ids = torch.arange(n_tokens, device="cuda", dtype=torch.long) + local_pages = torch.div(local_token_ids, page_size, rounding_mode="floor") + logical_pages = 1 + rank + cp_size * local_pages + logical_locs = ( + logical_pages * page_size + torch.remainder(local_token_ids, page_size) + ).to(torch.int64) + physical_locs = layout.logical_locs_to_physical(logical_locs) + capacity_tokens = int(physical_locs.max().item()) + page_size + 1 + + torch.manual_seed(20260604) + k_nope = ( + torch.randn((n_tokens, 1, 512), device="cuda", dtype=torch.bfloat16) + * 2.0 + ) + 0.25 + latent_cache = torch.randn( + (n_tokens, 576), device="cuda", dtype=torch.bfloat16 + ) + k_rope = latent_cache[:, 512:].unsqueeze(1) + self.assertFalse(k_rope.is_contiguous()) + + expected = torch.zeros( + (capacity_tokens, 1, 656), dtype=torch.uint8, device="cuda" + ) + nope_part, rope_part = quantize_k_cache_separate(k_nope, k_rope) + set_mla_kv_buffer_triton(expected, physical_locs, nope_part, rope_part) + + class FakePool: + def __init__(self): + self.nsa_kv_cache_store_fp8 = True + self.page_size = page_size + self.start_layer = 0 + self.kv_buffer = [ + torch.zeros( + (capacity_tokens, 1, 656), + dtype=torch.uint8, + device="cuda", + ) + ] + + pool = FakePool() + layer = SimpleNamespace(layer_id=0) + with patch.object( + runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=True + ), patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=False): + used = runtime.try_tai_fused_mla_store( + token_to_kv_pool=pool, + layer=layer, + layout=layout, + logical_locs=logical_locs, + k_nope=k_nope, + k_rope=k_rope, + ) + + torch.cuda.synchronize() + self.assertTrue(used) + torch.testing.assert_close( + pool.kv_buffer[0], + expected, + atol=0, + rtol=0, + ) + def test_token_range_materialize_uses_tai_kernel_when_enabled(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout