Stabilize CP shared-KV batch padding semantics
CP shared-KV bs>1 exposed three distinct padding domains: valid cache rows, CP page-tail compute rows, and MLP-sync flattened static padding. The previous implementation mixed these domains in direct-write and index top-k paths, so real requests failed when q/out_cache_loc lengths matched valid rows while metadata aliases described compute rows.\n\nThis change makes compute split strip only proven flattened static padding, keeps valid cache writes strict except for extend_num_tokens-proven static tails, marks CP-local EAGLE draft hidden state explicitly, and selects NSA index top-k query metadata by the actual q/weight row count.\n\nConstraint: CP shared-KV cache writes must never persist dummy page-tail or MLP static padding rows.\nConstraint: EAGLE draft hidden state can be CP-local before full CP metadata is visible in prepare_mlp_sync_batch.\nRejected: Use compute_padding_enabled as direct-write truncation proof | it silently accepts unknown out_cache_loc tails.\nRejected: Always consume compute q metadata in index top-k | actual q/weights can be valid-only after CP split.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not collapse valid rows, CP compute padding, and MLP static padding into one length condition; use explicit provenance.\nTested: remote py_compile for touched NSA files\nTested: remote targeted CP shared-KV padding/top-k regressions\nTested: remote pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -k 'not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel' => 228 passed, 1 deselected, 5 warnings, 2 subtests passed\nNot-tested: full ETE replay after the final index top-k fix\nNot-tested: TAI current-index fast path dtype fallback
This commit is contained in:
laoyao0822
@@ -566,6 +566,16 @@ target path 正确后，再恢复 EAGLE/draft，并做远端 ETE/perf 验证。
|
||||
- bs>1 时不允许 silent fallback：如果 CP draft shared-KV 已开启但 spec hidden、embedding pad metadata、input embeds 形状不满足 batch fast path，直接
|
||||
`[CP_SHARED_KV_FAIL_FAST][draft_batch_gt1_*]` 报错。bs=1 兼容 fallback 暂时保留。
|
||||
- scheduler 的 bs>1 admission gate 仍未打开；打开前必须完成下面 ETE 场景，尤其是 EAGLE accept length 与 output len。
|
||||
- 2026-06-04 远端启动失败记录：
|
||||
- 症状：pd warmup 阶段 EAGLE target forward 在 `cp_split_and_rebuild_data()` 抛
|
||||
`[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch] input tokens=8 expected=4`。
|
||||
- 根因：启动/warmup 或 speculative buffer 会把 `input_ids` 做 batch 尾部静态 padding；
|
||||
CP batch plan 的 `request_extend_lens` 仍只描述真实 valid rows。`split_kind="compute"` 应该丢弃尾部静态 padding，
|
||||
再由 CP split helper 自己生成 dummy compute rows；不能把 pad token embedding 当成 dummy compute data。
|
||||
- 合同订正：`split_kind="compute"` 允许两类尾部截断：
|
||||
1) CP compute padding 上限内的 dummy rows；
|
||||
2) `prepare_mlp_sync_batch()` 产生、并由 `forward_batch.extend_num_tokens` 显式给出上限的全局 static padding rows。
|
||||
`split_kind="valid"` 仍 fail-fast，避免 direct-write/out_cache_loc 静默丢 token。
|
||||
|
||||
### ETE 验证场景
|
||||
|
||||
@@ -993,3 +1003,616 @@ PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
|
||||
=> 117 passed, 5 warnings
|
||||
```
|
||||
|
||||
## 20. 2026-06-04 direct-write out_cache_loc 尾部 static padding 修正
|
||||
|
||||
### 失败现象
|
||||
|
||||
远端 warmup/EAGLE target forward 进入 NSA indexer direct-write 后失败：
|
||||
|
||||
```text
|
||||
[CP_SHARED_KV_FAIL_FAST][direct_write]
|
||||
reason=batch_split_out_cache_len_mismatch
|
||||
split_list tokens=4 out_cache_loc tokens=8
|
||||
```
|
||||
|
||||
栈在：
|
||||
|
||||
```text
|
||||
forward_mla.py -> nsa_indexer.py::_store_cp_shared_local_index_k_cache
|
||||
-> get_cp_shared_kv_local_out_cache_loc()
|
||||
```
|
||||
|
||||
这是前一轮 `cp_split_and_rebuild_data()` input-token mismatch 的同类问题，但发生在 cache write loc 边界：
|
||||
|
||||
- `CPSharedKVBatchPlan.request_extend_lens` 描述 valid rows；
|
||||
- model-runner/speculative warmup 可能在 flattened batch 尾部追加一段全局 static padding rows；
|
||||
- `out_cache_loc` 也可能携带这段尾部 padding loc；
|
||||
- CP shared-KV direct-write 是 valid-token 写入，不能把 dummy compute rows 写进 KV/index cache。
|
||||
|
||||
### 合同修正
|
||||
|
||||
保持 `split_tensor_by_cp_batch_plan(..., split_kind="valid")` 严格：
|
||||
|
||||
- valid split helper 仍然拒绝输入长度超过 `sum(request_extend_lens)`；
|
||||
- direct-write 边界负责先剥离已知的 **全局尾部 static padding locs**；
|
||||
- 只有在 `batch_plan.compute_padding_enabled=True`，且
|
||||
`valid_tokens < out_cache_loc.numel() <= sum(request_compute_padded_tokens)` 时允许截尾；
|
||||
- 其它 mismatch 继续 fail-fast，不做 silent fallback。
|
||||
|
||||
这样保证：
|
||||
|
||||
1. compute path 仍可吃到 bs>1 padding 后的堆叠计算；
|
||||
2. valid cache write 不会写 dummy rows；
|
||||
3. 如果未来改成 request 内部 interleaved padding layout，当前截尾规则会 fail-fast，而不是错误写 cache。
|
||||
|
||||
### 验证
|
||||
|
||||
新增回归：
|
||||
|
||||
- `test_local_out_cache_loc_ignores_trailing_static_padding_locs`
|
||||
- 构造 `valid_locs=4`，`out_cache_loc=8`，后 4 个为 static padding loc；
|
||||
- 验证 local direct-write 只返回前 4 个 valid loc。
|
||||
|
||||
远端 targeted 验证：
|
||||
|
||||
```text
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_ignores_trailing_static_padding_locs \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_valid_kind_rejects_trailing_padding_rows \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_uses_valid_rows_under_compute_padding
|
||||
=> 4 passed, 5 warnings
|
||||
```
|
||||
|
||||
远端相关文件验证：
|
||||
|
||||
```text
|
||||
python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py
|
||||
=> 73 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
|
||||
=> 37 passed, 3 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 107 passed, 1 deselected, 5 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
未完成验证：
|
||||
|
||||
- `test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel`
|
||||
在远端 installed-kernel self-test 处卡住；这是 TAI kernel self-test/import 路径，
|
||||
与本次 `get_cp_shared_kv_local_out_cache_loc()` 尾部 loc 截断不是同一逻辑路径，
|
||||
但后续需要单独处理，避免完整 runtime suite 长时间挂住。
|
||||
|
||||
## 21. 2026-06-04 warmup 卡死与 EAGLE draft compute padding 修正
|
||||
|
||||
### 现象
|
||||
|
||||
远端 warmup 初看卡在：
|
||||
|
||||
```text
|
||||
Start of pd disaggregation warmup ...
|
||||
[CacheCtrl-write] submit_write_cp_per_layer registered ...
|
||||
[HiCache-write] prepared CP per-layer backup before forward ...
|
||||
CUTE_DSL WARNING Unexpected error during package walk: cutlass.cute.experimental
|
||||
```
|
||||
|
||||
`py-spy` 看到 scheduler rank0 实际阻塞在：
|
||||
|
||||
```text
|
||||
torch/utils/file_baton.py:50 wait
|
||||
torch/utils/cpp_extension.py load
|
||||
tai_kernel/nsa_prefill/_extension_loader.py load_tai_kernel_extension
|
||||
in_seq_all_gather_rerange_cuda
|
||||
```
|
||||
|
||||
本次不是 HiCache ack 卡死，而是 TAI cpp extension JIT cache 目录存在 stale `lock`：
|
||||
|
||||
```text
|
||||
/root/.cache/tai-kernel/ops/aa52c7047416/lock
|
||||
```
|
||||
|
||||
同目录 `.so` 已存在且没有活跃 `ninja/nvcc`，移除 stale lock 后 warmup 继续。
|
||||
|
||||
### 继续暴露的真实错误
|
||||
|
||||
移除 stale lock 后，warmup 在 EAGLE draft extend 失败：
|
||||
|
||||
```text
|
||||
RuntimeError: Trying to create tensor with negative dimension -56: [-56, 6144]
|
||||
|
||||
forward_batch_info.py:_pad_inputs_to_size
|
||||
spec_info.hidden_states = self._pad_tensor_to_size(spec_info.hidden_states, num_tokens)
|
||||
```
|
||||
|
||||
根因：
|
||||
|
||||
- target 侧 `capture_draft_hidden_states=True` 在 `DeepseekV2Model.forward()` 的 CP output collect 之前抓取 hidden；
|
||||
- compute padding 后这个 side-channel 是 CP-local compute rows，例如 warmup 单请求 `hidden_states.shape[0] == 64`；
|
||||
- model-runner/speculative warmup 的 draft input 仍带全局 static padding token 数，例如 `num_tokens == 8`；
|
||||
- `ForwardBatch._pad_inputs_to_size()` 试图把 64 行 hidden pad 到 8 行，导致负维度；
|
||||
- 更深一层问题是 `ForwardMode.DRAFT_EXTEND` 没被视为 `context_parallel_extend`，所以 DeepSeek NextN draft 模型不会进入 CP-local draft fast path，CP-local side-channel 与非 CP draft input 语义不一致。
|
||||
|
||||
### 合同修正
|
||||
|
||||
EAGLE/NextN draft 在 `SGLANG_CP_DRAFT_SHARED_KV=1` 且 `uses_cp_shared_kv=True` 时必须跟 target 保持 CP-local 语义：
|
||||
|
||||
1. `can_cp_split()` 对 CP shared-KV draft extend 返回 true，使 draft 模型也构建 NSA CP metadata；
|
||||
2. `nsa_use_prefill_cp()` 对 CP shared-KV draft extend 返回 true，使 `DeepseekModelNextN` 进入 local draft path；
|
||||
3. `_pad_inputs_to_size()` 遇到 CP shared-KV draft hidden rows 已经大于全局 static padded `num_tokens` 时，不再尝试缩短/重 pad hidden；保留原始 CP-local rows，让 draft 模型按 CP split 后的 local input 使用。
|
||||
|
||||
这不是打开新的 scheduler bs>1 行为；它修正的是 EAGLE draft 在 CP shared-KV + compute padding 下的已有语义。
|
||||
|
||||
### 回归
|
||||
|
||||
新增：
|
||||
|
||||
- `test_can_cp_split_enables_cp_draft_shared_kv_draft_extend`
|
||||
- `test_nsa_use_prefill_cp_enables_cp_draft_shared_kv_draft_extend`
|
||||
- `test_cp_draft_padding_keeps_local_hidden_when_static_tokens_are_shorter`
|
||||
|
||||
远端 RED：三个测试在修复前分别失败为 `can_cp_split=False`、`nsa_use_prefill_cp=False`、负维度 `-56`。
|
||||
|
||||
远端 GREEN：
|
||||
|
||||
```text
|
||||
python -m py_compile \
|
||||
python/sglang/srt/layers/attention/nsa/utils.py \
|
||||
python/sglang/srt/model_executor/forward_batch_info.py
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py
|
||||
=> 76 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 144 passed, 1 deselected, 3 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
未验证：
|
||||
|
||||
- 最新代码还没重新跑完整 ETE warmup；需要用户重启服务后再看是否越过 draft extend。
|
||||
- stale TAI JIT lock 是环境/loader 问题，不应由业务逻辑修复掩盖；后续应在 TAI loader 层处理 stale lock 或超时诊断。
|
||||
|
||||
### 2026-06-04 订正：上一版 guard 仍过窄
|
||||
|
||||
远端新一轮 warmup 仍在相同位置失败：
|
||||
|
||||
```text
|
||||
RuntimeError: Trying to create tensor with negative dimension -56: [-56, 6144]
|
||||
|
||||
forward_batch_info.py:prepare_mlp_sync_batch
|
||||
forward_batch_info.py:_pad_inputs_to_size
|
||||
spec_info.hidden_states = self._pad_tensor_to_size(spec_info.hidden_states, num_tokens)
|
||||
```
|
||||
|
||||
新证据：
|
||||
|
||||
- 远端实际加载的 `forward_batch_info.py` 已包含上一版 guard；
|
||||
- 因此不是同步问题，也不是 HiCache ack；
|
||||
- 失败仍说明 `spec_info.hidden_states.shape[0] == 64` 且
|
||||
MLP sync static `num_tokens == 8`；
|
||||
- 上一版 guard 依赖 `uses_cp_shared_kv && SGLANG_CP_DRAFT_SHARED_KV`
|
||||
才保留 oversized draft hidden。这个前提过强：
|
||||
`prepare_mlp_sync_batch()` 发生在 `attn_backend.init_forward_metadata()`
|
||||
之前，draft 的 CP metadata/flag 传播不能作为此处是否允许普通 padding 的唯一依据。
|
||||
|
||||
上一版修正存在的问题：
|
||||
|
||||
- 仅通过 `spec_info.hidden_states.shape[0] > num_tokens` 推断 CP-local hidden
|
||||
过于宽泛；
|
||||
- 这个条件只能证明“普通 padding 会失败”，不能证明该 hidden 的语义一定是
|
||||
CP-local draft hidden；
|
||||
- 如果未来出现其它 draft path 产生 oversized hidden，长度判断会把错误 silently
|
||||
推迟到后面的模型计算，排查成本更高。
|
||||
|
||||
最终合同：
|
||||
|
||||
- CP-local draft hidden 必须由源头显式标记：
|
||||
`EagleDraftInput.cp_local_hidden_states=True`；
|
||||
- 只有 target 侧实际返回 `logits_output.draft_hidden_states` 时，EAGLE worker
|
||||
才把该 marker 传给 `EagleDraftInput`；
|
||||
- `_pad_inputs_to_size()` 只信这个 semantic marker，不再根据长度猜测；
|
||||
- 如果没有 marker 但 `hidden_states.shape[0] > num_tokens`，直接 fail-fast：
|
||||
`[CP_SHARED_KV_FAIL_FAST][draft_hidden_static_padding_mismatch]`；
|
||||
- 标记为 CP-local 的 hidden 保持原 tensor，后续由
|
||||
`DeepseekV3ForCausalLMNextN.forward()` 的 CP-local draft path 校验/消费。
|
||||
|
||||
新增回归：
|
||||
|
||||
- `test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible`
|
||||
- 构造 `cp_local_hidden_states=True` 且 CP metadata/env 尚不可见的最小 draft batch；
|
||||
- `hidden_states.shape[0]=64`，`num_tokens=8`；
|
||||
- 修复前远端 RED：`EagleDraftInput` 不接受该 marker；
|
||||
- 修复后 GREEN：保留 `(64, hidden)` hidden side-channel。
|
||||
- `test_cp_draft_padding_rejects_unmarked_oversized_hidden`
|
||||
- 构造未标记 oversized hidden；
|
||||
- 修复前远端 RED：旧长度判断会静默保留；
|
||||
- 修复后 GREEN：fail-fast，不允许靠长度误判 hidden 语义。
|
||||
|
||||
远端验证：
|
||||
|
||||
```text
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_rejects_unmarked_oversized_hidden
|
||||
=> 2 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py
|
||||
=> 78 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 144 passed, 1 deselected, 3 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
### 2026-06-04 订正：draft hidden marker 不应依赖 forward_mode
|
||||
|
||||
远端最新失败栈已经从负维度变为我们新增的 fail-fast：
|
||||
|
||||
```text
|
||||
RuntimeError: [CP_SHARED_KV_FAIL_FAST][draft_hidden_static_padding_mismatch]
|
||||
hidden_tokens=64 num_tokens=8 forward_mode=1
|
||||
|
||||
scheduler.run_batch
|
||||
model_worker.forward_batch_generation
|
||||
eagle_worker.py:forward_draft_extend
|
||||
draft_model_runner.forward
|
||||
model_runner.py:_forward_raw
|
||||
forward_batch.prepare_mlp_sync_batch
|
||||
forward_batch_info.py:_pad_inputs_to_size
|
||||
```
|
||||
|
||||
根因订正：
|
||||
|
||||
- `ForwardMode.EXTEND == 1`，所以 fail-fast 中的 `forward_mode=1` 不是
|
||||
`DRAFT_EXTEND`；
|
||||
- 这不是说明 EAGLE draft 语义丢失，而是
|
||||
`prepare_mlp_sync_batch()` 在 `is_extend_in_batch + DP max padding` 路径下会临时把
|
||||
draft/verify/decode/idle 等 mode 改写为 `ForwardMode.EXTEND`，用于复用 extend 静态
|
||||
MLP sync padding；
|
||||
- 因此 `_pad_inputs_to_size()` 不能用 `forward_mode.is_draft_extend()` 判断是否允许保留
|
||||
CP-local draft hidden；
|
||||
- 正确的语义边界是：当前代码块已经由 `spec_info.is_draft_input()` 保护，是否 CP-local
|
||||
只应由 `EagleDraftInput.cp_local_hidden_states` 这个显式 marker 决定。
|
||||
|
||||
修正：
|
||||
|
||||
- `keep_cp_local_hidden = getattr(spec_info, "cp_local_hidden_states", False)`；
|
||||
- 保留未标记 oversized hidden 的 fail-fast，避免再次回到长度推断；
|
||||
- 新增回归覆盖实际运行形态：`forward_mode=ForwardMode.EXTEND` 且
|
||||
`EagleDraftInput.cp_local_hidden_states=True` 时，`hidden_states=(64, hidden)` 在
|
||||
`num_tokens=8` 的静态 padding 下必须保持不变。
|
||||
|
||||
远端 RED/GREEN 证据：
|
||||
|
||||
```text
|
||||
# RED: 修复前新增测试复现远端 fail-fast
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite
|
||||
=> FAILED with [CP_SHARED_KV_FAIL_FAST][draft_hidden_static_padding_mismatch]
|
||||
|
||||
# GREEN: 修复后 marker 测试 + 未标记 oversized 保护同时通过
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_draft_padding_rejects_unmarked_oversized_hidden
|
||||
=> 2 passed, 5 warnings
|
||||
|
||||
# 相关套件
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 223 passed, 1 deselected, 5 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
未验证：
|
||||
|
||||
- 需要用户重启远端服务并打 warmup/ETE 流量确认已经越过 EAGLE draft extend；
|
||||
- 如果后续仍失败，应优先看新栈，不要回到 forward_mode 判定或长度推断。
|
||||
|
||||
### 2026-06-04 订正：普通 MLP sync static padding 也会进入 CP compute split
|
||||
|
||||
远端 warmup 已通过后，第一条真实请求在 target forward 失败：
|
||||
|
||||
```text
|
||||
RuntimeError: [CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch]
|
||||
input tokens=40392 expected=40387
|
||||
|
||||
model_runner.forward_extend
|
||||
DeepseekV2Model.forward
|
||||
cp_split_and_rebuild_data
|
||||
split_tensor_by_cp_batch_plan
|
||||
```
|
||||
|
||||
关键信号：
|
||||
|
||||
- 请求日志为 `40391 input + 200000 new`；
|
||||
- CP plan valid extend rows 为 `40387`；
|
||||
- 进入 `cp_split_and_rebuild_data()` 的 hidden rows 为 `40392`；
|
||||
- `40392` 是按 `attn_tp_size/attn_cp_size` 对齐后的长度，差值 5 是
|
||||
`ForwardBatch.prepare_mlp_sync_batch()` 添加的全局尾部 static padding；
|
||||
- 这次不是 tiny compute padding：`plan.compute_padding_enabled=False` 时也会发生。
|
||||
|
||||
上一版遗漏：
|
||||
|
||||
- `split_tensor_by_cp_batch_plan(split_kind="compute")` 只在
|
||||
`compute_padding_enabled=True` 时允许剥离尾部 padding；
|
||||
- 但 MLP sync static padding 与 CP compute padding 是两类 padding：
|
||||
- CP compute padding：为 page/owner-lane compute 形态补 dummy rows；
|
||||
- MLP sync static padding：为 DP/TP/CP collective buffer 对齐，在 flattened batch 尾部补 rows；
|
||||
- 二者不能混为一谈。即使不需要 CP compute padding，也必须允许 compute split 剥离已知的
|
||||
全局尾部 static padding。
|
||||
|
||||
修正合同：
|
||||
|
||||
- `split_kind="compute"` 可以剥离尾部 static padding，但必须由调用方显式传入
|
||||
`static_padded_tokens` 上限；
|
||||
- `_cp_split_and_rebuild_batch_in_seq()` 从 `forward_batch.extend_num_tokens` 传入该上限；
|
||||
这是 `prepare_mlp_sync_batch()` 后的本地 padded token count；
|
||||
- `split_kind="valid"` 仍不接受任何尾部 padding，direct-write/out_cache_loc 路径保持严格；
|
||||
- 如果 input 超过 `static_padded_tokens` 或 compute padded 上限，继续 fail-fast。
|
||||
|
||||
新增回归：
|
||||
|
||||
- `test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding`
|
||||
- 构造 `extend_len=7,page_size=4,cp_size=2`，此时 `compute_padding_enabled=False`；
|
||||
- 输入 tensor 有 8 行，模拟 MLP sync static padding；
|
||||
- 修复前 RED：`input tokens=8 expected=7`；
|
||||
- 修复后 GREEN：剥离第 8 行，仅按前 7 个 valid rows 做 CP split。
|
||||
|
||||
远端验证：
|
||||
|
||||
```text
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_valid_kind_rejects_trailing_padding_rows \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows
|
||||
=> 3 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 224 passed, 1 deselected, 5 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
未验证：
|
||||
|
||||
- 需要重启远端服务并重放第一条真实请求，确认已经越过
|
||||
`input tokens=40392 expected=40387` 这一栈；
|
||||
- 如果后续还有失败，应优先看新栈，不应回退到放宽 `valid` split。
|
||||
|
||||
### 2026-06-04 订正：position split 也必须接收 MLP sync static padding 上限
|
||||
|
||||
上一轮修复 data/1d compute split 后，远端再次失败，但栈前进到 position：
|
||||
|
||||
```text
|
||||
RuntimeError: [CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch]
|
||||
input tokens=40392 expected=40387
|
||||
|
||||
DeepseekV2Model.forward
|
||||
positions = cp_split_and_rebuild_position(forward_batch, positions)
|
||||
split_tensor_by_cp_batch_plan
|
||||
```
|
||||
|
||||
根因：
|
||||
|
||||
- `cp_split_and_rebuild_data()` 和 `cp_split_and_rebuild_1d()` 通过
|
||||
`_cp_split_and_rebuild_batch_in_seq()` 传入了 `forward_batch.extend_num_tokens`；
|
||||
- `cp_split_and_rebuild_position()` 是独立 wrapper，仍直接调用
|
||||
`split_tensor_by_cp_batch_plan(..., mode="position")`，没有传
|
||||
`static_padded_tokens`；
|
||||
- 因此同一批 MLP sync 尾部 static padding 在 hidden data 上已被剥离，但 positions 上仍
|
||||
fail-fast。
|
||||
|
||||
修正：
|
||||
|
||||
- `cp_split_and_rebuild_position()` 也传入
|
||||
`static_padded_tokens=getattr(forward_batch, "extend_num_tokens", None)`；
|
||||
- runtime 调用点复核：
|
||||
- `_cp_split_and_rebuild_batch_in_seq()`：compute data/1d，已传 static 上限；
|
||||
- `cp_split_and_rebuild_position()`：compute position，本次补齐；
|
||||
- `get_cp_shared_kv_local_out_cache_loc()`：valid cache write，仍保持严格，不传 static 上限。
|
||||
|
||||
新增回归：
|
||||
|
||||
- `test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding`
|
||||
- 构造 `extend_len=7,page_size=4,cp_size=2`，`compute_padding_enabled=False`；
|
||||
- 输入 positions 长度 8，模拟 MLP sync static padding；
|
||||
- 修复前 RED：`input tokens=8 expected=7`；
|
||||
- 修复后 GREEN：position split 与先截断到 7 行再 split 的结果一致。
|
||||
|
||||
远端验证：
|
||||
|
||||
```text
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_valid_kind_rejects_trailing_padding_rows
|
||||
=> 3 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 225 passed, 1 deselected, 5 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
未验证：
|
||||
|
||||
- 需要重启服务重放第一条请求，确认已经越过 `cp_split_and_rebuild_position()`。
|
||||
|
||||
### 2026-06-04 订正：direct-write 也会收到 MLP sync static padding，不能用 compute padding 兜底
|
||||
|
||||
远端再次失败，栈前进到 NSA indexer cache write：
|
||||
|
||||
```text
|
||||
RuntimeError: [CP_SHARED_KV_FAIL_FAST][direct_write]
|
||||
reason=batch_split_out_cache_len_mismatch
|
||||
split_list tokens=40387 out_cache_loc tokens=40392
|
||||
|
||||
nsa_indexer.py::_store_cp_shared_local_index_k_cache
|
||||
get_cp_shared_kv_local_out_cache_loc()
|
||||
```
|
||||
|
||||
这次失败与 data/position split 属于同一类，但边界不同：
|
||||
|
||||
- `ForwardBatch.prepare_mlp_sync_batch()` 会把 `input_ids / positions / out_cache_loc`
|
||||
一起 pad 到 `extend_num_tokens`；
|
||||
- `CPSharedKVBatchPlan.request_extend_lens` 仍只描述 valid token rows；
|
||||
- direct-write 的 `out_cache_loc` 是 cache 写入地址，只能对应 valid rows；
|
||||
- 因此 direct-write 边界必须先剥离全局尾部 static padding locs，再进入
|
||||
`split_kind="valid"`；
|
||||
- 这与 CP page-tail/owner-lane compute padding 不是一回事。
|
||||
|
||||
修正后的三层合同：
|
||||
|
||||
1. **valid rows**：`sum(request_extend_lens)`，唯一允许写入 KV/index cache 的 rows；
|
||||
2. **CP compute padding / page tail**：按 request/page/owner-lane 补 dummy rows，
|
||||
只由 `split_tensor_by_cp_batch_plan(split_kind="compute")` 在 split 内部产生；
|
||||
3. **MLP sync static padding**：`prepare_mlp_sync_batch()` 在 flattened batch 尾部追加，
|
||||
上限由 `forward_batch.extend_num_tokens` 显式给出。
|
||||
|
||||
direct-write 规则：
|
||||
|
||||
- 如果 `out_cache_loc.numel() == valid_tokens`，直接进入 strict valid split；
|
||||
- 如果 `out_cache_loc.numel() > valid_tokens`，只有
|
||||
`out_cache_loc.numel() <= forward_batch.extend_num_tokens` 时才截尾；
|
||||
- 不再把 `batch_plan.compute_padding_enabled` 当成 direct-write 截尾依据；
|
||||
- 如果没有 `extend_num_tokens` 证明这段尾部来自 MLP static padding，即使
|
||||
`compute_padding_enabled=True` 也 fail-fast。
|
||||
|
||||
同类调用点复核：
|
||||
|
||||
- `_cp_split_and_rebuild_batch_in_seq()`：data/1d compute split，传入
|
||||
`extend_num_tokens`，允许剥离 MLP static padding；
|
||||
- `cp_split_and_rebuild_position()`：position compute split，传入
|
||||
`extend_num_tokens`，允许剥离 MLP static padding；
|
||||
- `get_cp_shared_kv_local_out_cache_loc()`：valid cache write，先用
|
||||
`extend_num_tokens` 截掉静态尾部，再调用 strict valid split；
|
||||
- `select_cp_local_valid_rows_for_cache_write()`：只负责从 local compute rows 中剥离
|
||||
CP compute padding，不处理 flattened MLP static padding；该 padding 必须在 CP split 前处理。
|
||||
|
||||
新增回归：
|
||||
|
||||
- `test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding`
|
||||
- `extend_len=7,page_size=4,cp_size=2`，`compute_padding_enabled=False`；
|
||||
- `out_cache_loc` 长度 8，模拟 MLP sync static padding；
|
||||
- 修复前 RED：`split_list tokens=7 out_cache_loc tokens=8`；
|
||||
- 修复后只返回 rank-local valid tail locs。
|
||||
- `test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding`
|
||||
- 验证 direct-write 不再因为 `compute_padding_enabled=True` 静默截掉未知尾部 loc。
|
||||
|
||||
远端验证：
|
||||
|
||||
```text
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_local_out_cache_loc_ignores_trailing_static_padding_locs \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding
|
||||
=> 5 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 227 passed, 1 deselected, 5 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
仍需 ETE 验证：
|
||||
|
||||
- 需要重启远端服务并重放第一条真实请求，确认已经越过
|
||||
`get_cp_shared_kv_local_out_cache_loc()` 的 mismatch；
|
||||
- 如果后续还有失败，应继续按新栈定位，不要再把三类 padding 混成一个长度条件。
|
||||
|
||||
### 2026-06-04 订正：index top-k batch path 必须按实际 q rows 选择 valid/compute metadata
|
||||
|
||||
重启后远端越过 direct-write mismatch，但在 NSA indexer top-k 阶段失败：
|
||||
|
||||
```text
|
||||
RuntimeError: [CP_SHARED_KV_FAIL_FAST][index_topk]
|
||||
reason=batch_gt1_index_q_length_mismatch
|
||||
batch_size=1 layer_id=0 cursor=5056 q_tokens=4995 weights_tokens=4995
|
||||
```
|
||||
|
||||
关键观察：
|
||||
|
||||
- 失败请求仍是 page-tail 场景；本 rank valid q rows 为 `4995`，compute/page-tail rows 为 `5056`；
|
||||
- `_get_topk_in_seq_cp_pair_batch()` 旧逻辑直接选择 `request_compute_seq_q_*`
|
||||
或 `request_actual_seq_q_*`，而 `request_actual_seq_q_*` 当前也是 compute alias；
|
||||
- 但 `forward_mla.py` 中 indexer 的 q/weights 来自已经 CP split 后的 `hidden_states/q_lora` 投影，
|
||||
实际输入是 valid rows `4995`，不是 compute-padded rows `5056`；
|
||||
- 因此 top-k batch path 不能从 metadata 名字推断 q layout，必须用实际
|
||||
`q_fp8.shape[0] / weights.shape[0]` 选择 valid 或 compute 视图。
|
||||
|
||||
修正合同：
|
||||
|
||||
- CP batch plan 同时暴露 valid 和 compute q metadata；
|
||||
- index top-k 在 runtime 入口选择：
|
||||
- 如果 `q_tokens == sum(request_valid_seq_q_prev/next)`，按 valid segments 消费；
|
||||
- 如果 `q_tokens == sum(request_compute_seq_q_prev/next)`，按 compute segments 消费，
|
||||
但仍用 valid segment length 过滤 dummy page-tail rows；
|
||||
- 如果两者都不匹配，fail-fast 并同时打印 valid/compute 期望长度；
|
||||
- 这避免把 page-tail dummy rows 强行加到 indexer q/weights 上，也避免未来真的传入
|
||||
compute-padded q 时丢失支持。
|
||||
|
||||
新增实现：
|
||||
|
||||
- `BatchTopKQueryLengths`
|
||||
- `_select_batch_topk_query_lengths()`
|
||||
|
||||
新增回归：
|
||||
|
||||
- `test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias`
|
||||
- 复现线上同型：`extend_len=40387,page_size=64,cp_size=8,cp_rank=0`；
|
||||
- valid local rows = `4995`，compute local rows = `5056`；
|
||||
- 修复前 RED：helper 不存在 / 原 runtime 只能按 compute cursor；
|
||||
- 修复后 GREEN：q rows 为 4995 时选择 valid metadata；q rows 为 5056 时仍选择 compute metadata。
|
||||
|
||||
同类路径复核：
|
||||
|
||||
- `nsa_backend.forward_extend()` 不直接读取 `request_compute_seq_q_*`；它只把
|
||||
`topk_indices` pad 到当前 q 行数，因此由 indexer 返回 rows 与 q rows 对齐即可；
|
||||
- `select_cp_local_valid_rows_for_cache_write()` 仍只处理 cache write 的 compute rows -> valid rows，
|
||||
不参与 top-k metadata 选择。
|
||||
|
||||
远端验证：
|
||||
|
||||
```text
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py::TestNSAInSeqCPUtils::test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias
|
||||
=> 1 passed, 5 warnings
|
||||
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_cp_utils.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
|
||||
-k "not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel"
|
||||
=> 228 passed, 1 deselected, 5 warnings, 2 subtests passed
|
||||
```
|
||||
|
||||
仍需 ETE 验证：
|
||||
|
||||
- 需要重启远端服务并重放真实请求，确认已经越过
|
||||
`batch_gt1_index_q_length_mismatch`；
|
||||
- 日志中的 `[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] current_index_k must be uint8`
|
||||
是另一个性能 fast-path dtype 问题，本次未修；它当前是 warning fallback，不是本次进程退出原因。
|
||||
|
||||
Reference in New Issue
Block a user