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