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

Stabilize CP shared-KV batch padding semantics

Browse Source 
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
2026-06-04 07:25:11 +08:00
parent 02af370e87
commit 3d6007246b
7 changed files with 1403 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

623
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md
View File

@@ -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，不是本次进程退出原因。

170
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -3,6 +3,7 @@ from __future__ import annotations
import contextlib
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
import torch
@@ -120,6 +121,120 @@ def _compute_contiguous_valid_cp_query_count(
return max(0, min(actual_seq_q, valid_count))
@dataclass(frozen=True)
class BatchTopKQueryLengths:
request_seq_q_prev: List[int]
request_seq_q_next: List[int]
request_valid_seq_q_prev: List[int]
request_valid_seq_q_next: List[int]
uses_compute_query_rows: bool
valid_token_count: int
compute_token_count: int
def _select_batch_topk_query_lengths(
*,
cp_metadata,
batch_plan,
batch_size: int,
q_tokens: int,
weights_tokens: int,
layer_id: Optional[int] = None,
) -> BatchTopKQueryLengths:
"""Select CP top-k segment lengths that match actual q/weight rows."""
def metadata_list(name: str, fallback_name: Optional[str] = None) -> List[int]:
values = getattr(cp_metadata, name, None)
if values is None and batch_plan is not None:
values = getattr(batch_plan, name, None)
if values is None and fallback_name is not None:
values = getattr(cp_metadata, fallback_name, None)
if values is None and batch_plan is not None:
values = getattr(batch_plan, fallback_name, None)
return [int(x) for x in (values or [])]
request_compute_seq_q_prev = metadata_list(
"request_compute_seq_q_prev",
fallback_name="request_actual_seq_q_prev",
)
request_compute_seq_q_next = metadata_list(
"request_compute_seq_q_next",
fallback_name="request_actual_seq_q_next",
)
request_valid_seq_q_prev = metadata_list(
"request_valid_seq_q_prev",
fallback_name="request_valid_actual_seq_q_prev",
)
request_valid_seq_q_next = metadata_list(
"request_valid_seq_q_next",
fallback_name="request_valid_actual_seq_q_next",
)
if not request_valid_seq_q_prev:
request_valid_seq_q_prev = metadata_list("request_actual_seq_q_prev")
if not request_valid_seq_q_next:
request_valid_seq_q_next = metadata_list("request_actual_seq_q_next")
if not (
len(request_compute_seq_q_prev) == batch_size
and len(request_compute_seq_q_next) == batch_size
and len(request_valid_seq_q_prev) == batch_size
and len(request_valid_seq_q_next) == batch_size
):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_gt1_index_query_metadata_incomplete "
f"batch_size={batch_size} layer_id={layer_id} "
f"compute_q_prev={request_compute_seq_q_prev} "
f"compute_q_next={request_compute_seq_q_next} "
f"valid_q_prev={request_valid_seq_q_prev} "
f"valid_q_next={request_valid_seq_q_next}"
)
valid_token_count = sum(request_valid_seq_q_prev) + sum(request_valid_seq_q_next)
compute_token_count = sum(request_compute_seq_q_prev) + sum(
request_compute_seq_q_next
)
q_tokens = int(q_tokens)
weights_tokens = int(weights_tokens)
if q_tokens != weights_tokens:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_gt1_index_q_weight_length_mismatch "
f"batch_size={batch_size} layer_id={layer_id} q_tokens={q_tokens} "
f"weights_tokens={weights_tokens} valid_tokens={valid_token_count} "
f"compute_tokens={compute_token_count}"
)
if q_tokens == valid_token_count:
return BatchTopKQueryLengths(
request_seq_q_prev=request_valid_seq_q_prev,
request_seq_q_next=request_valid_seq_q_next,
request_valid_seq_q_prev=request_valid_seq_q_prev,
request_valid_seq_q_next=request_valid_seq_q_next,
uses_compute_query_rows=False,
valid_token_count=valid_token_count,
compute_token_count=compute_token_count,
)
if q_tokens == compute_token_count:
return BatchTopKQueryLengths(
request_seq_q_prev=request_compute_seq_q_prev,
request_seq_q_next=request_compute_seq_q_next,
request_valid_seq_q_prev=request_valid_seq_q_prev,
request_valid_seq_q_next=request_valid_seq_q_next,
uses_compute_query_rows=True,
valid_token_count=valid_token_count,
compute_token_count=compute_token_count,
)
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_gt1_index_q_length_mismatch "
f"batch_size={batch_size} layer_id={layer_id} q_tokens={q_tokens} "
f"weights_tokens={weights_tokens} valid_tokens={valid_token_count} "
f"compute_tokens={compute_token_count}"
)
def _log_cp_shared_kv_index_prefetch_fallback(
reason: str,
message: str,
@@ -1747,20 +1862,6 @@ class Indexer(MultiPlatformOp):
assert cp_metadata is not None
batch_size = int(getattr(cp_metadata, "batch_size", 1) or 1)
batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
compute_padding_enabled = bool(
getattr(cp_metadata, "compute_padding_enabled", False)
or bool(getattr(batch_plan, "compute_padding_enabled", False))
)
def metadata_list(name: str, fallback_name: Optional[str] = None) -> List[int]:
values = getattr(cp_metadata, name, None)
if values is None and batch_plan is not None:
values = getattr(batch_plan, name, None)
if values is None and fallback_name is not None:
values = getattr(cp_metadata, fallback_name, None)
if values is None and batch_plan is not None:
values = getattr(batch_plan, fallback_name, None)
return list(values or [])
request_kv_len_prev = list(getattr(cp_metadata, "request_kv_len_prev", []) or [])
request_kv_len_next = list(getattr(cp_metadata, "request_kv_len_next", []) or [])
@@ -1772,35 +1873,20 @@ class Indexer(MultiPlatformOp):
request_kv_len_next = list(
getattr(batch_plan, "request_kv_len_next", []) or []
)
request_actual_seq_q_prev = metadata_list(
"request_compute_seq_q_prev"
if compute_padding_enabled
else "request_actual_seq_q_prev",
fallback_name="request_actual_seq_q_prev",
query_lengths = _select_batch_topk_query_lengths(
cp_metadata=cp_metadata,
batch_plan=batch_plan,
batch_size=batch_size,
q_tokens=int(q_fp8.shape[0]),
weights_tokens=int(weights.shape[0]),
layer_id=layer_id,
)
request_actual_seq_q_next = metadata_list(
"request_compute_seq_q_next"
if compute_padding_enabled
else "request_actual_seq_q_next",
fallback_name="request_actual_seq_q_next",
)
request_valid_seq_q_prev = metadata_list(
"request_valid_seq_q_prev",
fallback_name="request_valid_actual_seq_q_prev",
)
request_valid_seq_q_next = metadata_list(
"request_valid_seq_q_next",
fallback_name="request_valid_actual_seq_q_next",
)
if not compute_padding_enabled:
# Older bs>1 metadata did not have explicit valid-q aliases because
# actual q length was also the valid q length. Keep that path
# compatible while compute-padding remains fail-fast if valid
# lengths are missing.
if not request_valid_seq_q_prev:
request_valid_seq_q_prev = request_actual_seq_q_prev
if not request_valid_seq_q_next:
request_valid_seq_q_next = request_actual_seq_q_next
request_actual_seq_q_prev = query_lengths.request_seq_q_prev
request_actual_seq_q_next = query_lengths.request_seq_q_next
request_valid_seq_q_prev = query_lengths.request_valid_seq_q_prev
request_valid_seq_q_next = query_lengths.request_valid_seq_q_next
if not (
len(request_kv_len_prev) == batch_size
and len(request_kv_len_next) == batch_size

117
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -2,7 +2,7 @@
import logging
from dataclasses import dataclass
from itertools import accumulate
from typing import TYPE_CHECKING, List, Tuple, Union
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
@@ -106,6 +106,23 @@ def _is_cp_shared_kv_forward_batch(forward_batch: "ForwardBatch") -> bool:
return bool(getattr(forward_batch, "uses_cp_shared_kv", False))
def _is_cp_shared_kv_draft_extend(forward_batch: "ForwardBatch") -> bool:
"""Return whether EAGLE/NextN draft extend should keep CP-local semantics."""
if not _is_cp_shared_kv_forward_batch(forward_batch):
return False
if not envs.SGLANG_CP_DRAFT_SHARED_KV.get():
return False
forward_mode = getattr(forward_batch, "forward_mode", None)
is_draft_extend = getattr(forward_mode, "is_draft_extend", None)
if not callable(is_draft_extend):
return False
try:
return bool(is_draft_extend(include_v2=True))
except TypeError:
return bool(is_draft_extend())
def _fail_if_cp_shared_kv_round_robin(
forward_batch: "ForwardBatch",
*,
@@ -751,12 +768,25 @@ def get_cp_shared_kv_batch_plan(forward_batch: "ForwardBatch"):
return None
def _get_forward_batch_static_padded_tokens(
forward_batch: "ForwardBatch",
) -> Optional[int]:
static_padded_tokens = getattr(forward_batch, "extend_num_tokens", None)
if static_padded_tokens is None:
return None
try:
return int(static_padded_tokens)
except (TypeError, ValueError):
return None
def split_tensor_by_cp_batch_plan(
tensor: torch.Tensor,
plan,
*,
mode: str = "data",
split_kind: str = "compute",
static_padded_tokens: Optional[int] = None,
) -> torch.Tensor:
"""Split a flattened batch tensor by per-request in-seq CP plan.
@@ -802,11 +832,42 @@ def split_tensor_by_cp_batch_plan(
)
expected_tokens = sum(int(x) for x in request_extend_lens)
if int(tensor.shape[0]) != expected_tokens:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch] "
f"input tokens={int(tensor.shape[0])} expected={expected_tokens}"
input_tokens = int(tensor.shape[0])
if input_tokens != expected_tokens:
max_compute_tokens = (
sum(int(x) for x in request_target_lens)
if split_kind == "compute" and compute_padding_enabled
else expected_tokens
)
max_static_tokens = (
int(static_padded_tokens)
if static_padded_tokens is not None
else expected_tokens
)
max_allowed_tokens = max(max_compute_tokens, max_static_tokens)
if (
split_kind == "compute"
and input_tokens > expected_tokens
and input_tokens <= max_allowed_tokens
and (
compute_padding_enabled
or (
static_padded_tokens is not None
and int(static_padded_tokens) > expected_tokens
)
)
):
tensor = tensor[:expected_tokens]
else:
expected_detail = f"expected={expected_tokens}"
if split_kind == "compute" and compute_padding_enabled:
expected_detail += f" max_compute={max_compute_tokens}"
if split_kind == "compute" and static_padded_tokens is not None:
expected_detail += f" static_padded={int(static_padded_tokens)}"
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_input_len_mismatch] "
f"input tokens={input_tokens} {expected_detail}"
)
local_chunks = []
request_tensors = torch.split(tensor, [int(x) for x in request_extend_lens], dim=0)
@@ -1462,11 +1523,15 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
min_extend_token_count = 1
else:
min_extend_token_count = cp_size
is_context_parallel_extend = (
forward_batch.forward_mode.is_context_parallel_extend()
or _is_cp_shared_kv_draft_extend(forward_batch)
)
if (
cur_cp_seq_len != 0
and cp_size > 1
and use_nsa
and forward_batch.forward_mode.is_context_parallel_extend()
and is_context_parallel_extend
and is_nsa_enable_prefill_cp()
and extend_token_count >= min_extend_token_count
):
@@ -1523,6 +1588,7 @@ def _cp_split_and_rebuild_batch_in_seq(forward_batch, input_: torch.Tensor):
input_,
get_cp_shared_kv_batch_plan(forward_batch),
mode="1d" if input_.dim() == 1 else "data",
static_padded_tokens=_get_forward_batch_static_padded_tokens(forward_batch),
)
@@ -1623,12 +1689,31 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
mismatch_reason = "split_out_cache_len_mismatch"
out_cache_tokens = int(out_cache_loc.numel())
if split_tokens != out_cache_tokens:
raise_cp_shared_kv_direct_write_error(
mismatch_reason,
"split_list tokens=%s out_cache_loc tokens=%s",
split_tokens,
out_cache_tokens,
)
static_padded_tokens = _get_forward_batch_static_padded_tokens(forward_batch)
if (
static_padded_tokens is not None
and out_cache_tokens > split_tokens
and out_cache_tokens <= static_padded_tokens
):
# Model-runner/speculative warmup can append one global block of
# static padding rows after the valid flattened batch. These rows
# may carry dummy cache locs, but CP shared-KV direct-write is a
# valid-token operation: never split/write dummy compute rows.
#
# Do not use CP compute-padding metadata as an implicit allowance
# here. Page-tail/owner-lane compute padding is produced by
# split_tensor_by_cp_batch_plan() after valid input rows are split;
# only forward_batch.extend_num_tokens proves that out_cache_loc
# already contains global trailing static padding rows.
out_cache_loc = out_cache_loc[:split_tokens]
else:
raise_cp_shared_kv_direct_write_error(
mismatch_reason,
"split_list tokens=%s out_cache_loc tokens=%s static_padded=%s",
split_tokens,
out_cache_tokens,
static_padded_tokens,
)
if batch_plan is not None:
local_out_cache_loc = split_tensor_by_cp_batch_plan(
@@ -1724,6 +1809,7 @@ def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
positions,
get_cp_shared_kv_batch_plan(forward_batch),
mode="position",
static_padded_tokens=_get_forward_batch_static_padded_tokens(forward_batch),
)
position_id_list = list(
@@ -1833,10 +1919,15 @@ def nsa_cp_round_robin_split_q_seqs(
def nsa_use_prefill_cp(forward_batch, nsa_enable_prefill_cp=None):
if nsa_enable_prefill_cp is None:
nsa_enable_prefill_cp = is_nsa_enable_prefill_cp()
forward_mode = getattr(forward_batch, "forward_mode", None)
is_context_parallel_extend = (
forward_mode is not None
and forward_mode.is_context_parallel_extend()
) or _is_cp_shared_kv_draft_extend(forward_batch)
if (
forward_batch.nsa_cp_metadata is not None
and nsa_enable_prefill_cp
and forward_batch.forward_mode.is_context_parallel_extend()
and is_context_parallel_extend
):
return True
else:

25
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -1002,9 +1002,28 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
spec_info.accept_length = self._pad_tensor_to_size(
spec_info.accept_length, bs
)
spec_info.hidden_states = self._pad_tensor_to_size(
spec_info.hidden_states, num_tokens
)
# EAGLE/NextN draft extend can receive a CP-local hidden side
# channel captured by the target model before CP output collect.
# This block is already guarded by spec_info.is_draft_input().
# prepare_mlp_sync_batch may temporarily rewrite draft forward
# modes to EXTEND while static DP padding is prepared, so the
# contract must be carried by EagleDraftInput rather than
# inferred from forward_mode or tensor length.
keep_cp_local_hidden = getattr(spec_info, "cp_local_hidden_states", False)
if not keep_cp_local_hidden:
if spec_info.hidden_states.shape[0] > num_tokens:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST]"
"[draft_hidden_static_padding_mismatch] "
"EAGLE draft hidden_states is larger than the static "
"MLP-sync token count but is not marked as CP-local. "
f"hidden_tokens={spec_info.hidden_states.shape[0]} "
f"num_tokens={num_tokens} "
f"forward_mode={self.forward_mode}"
)
spec_info.hidden_states = self._pad_tensor_to_size(
spec_info.hidden_states, num_tokens
)
def prepare_attn_tp_scatter_input(self, model_runner: ModelRunner):
from sglang.srt.layers.communicator import get_attn_tp_context

4
python/sglang/srt/speculative/eagle_info.py
View File

@@ -622,6 +622,10 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
topk_index: torch.Tensor = None
# shape: (b, hidden_size)
hidden_states: torch.Tensor = None
# True when hidden_states is a CP-local side channel captured by the target
# model before CP output collect. This is a semantic marker; consumers must
# not infer it from tensor length alone.
cp_local_hidden_states: bool = False
capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL
# Inputs for extend

7
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -334,6 +334,9 @@ class EAGLEWorker(TpModelWorker):
if logits_output.draft_hidden_states is not None
else logits_output.hidden_states
)
cp_local_draft_hidden_states = (
logits_output.draft_hidden_states is not None
)
if (
envs.SGLANG_CP_DRAFT_SHARED_KV.get()
and draft_hidden_states is None
@@ -347,6 +350,7 @@ class EAGLEWorker(TpModelWorker):
next_token_ids,
seq_lens_cpu,
logits_output.mm_input_embeds,
cp_local_hidden_states=cp_local_draft_hidden_states,
)
return GenerationBatchResult(
logits_output=logits_output,
@@ -935,6 +939,8 @@ class EAGLEWorker(TpModelWorker):
next_token_ids: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
mm_input_embeds: Optional[torch.Tensor] = None,
*,
cp_local_hidden_states: bool = False,
):
"""Run draft model extend. This API modifies the states of the batch.
@@ -945,6 +951,7 @@ class EAGLEWorker(TpModelWorker):
"""
batch.spec_info = EagleDraftInput(
hidden_states=hidden_states,
cp_local_hidden_states=cp_local_hidden_states,
verified_id=next_token_ids,
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,

517
test/registered/unit/layers/test_nsa_cp_utils.py
View File

@@ -27,6 +27,7 @@ from sglang.srt.layers.attention.nsa.utils import (
get_cp_shared_kv_local_out_cache_loc,
get_cp_shared_kv_local_physical_out_cache_loc,
get_cp_local_embedding_padded_token_count,
nsa_use_prefill_cp,
pad_cp_local_input_ids_for_embedding,
prepare_input_dp_with_cp_dsa,
select_cp_local_valid_rows_for_cache_write,
@@ -34,7 +35,9 @@ from sglang.srt.layers.attention.nsa.utils import (
split_in_seq_cp_local_pair,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.models.deepseek_nextn import DeepseekModelNextN
from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
@@ -326,6 +329,45 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
):
self.assertTrue(can_cp_split(64, 8, True, forward_batch))
def test_can_cp_split_enables_cp_draft_shared_kv_draft_extend(self):
class DraftMode:
def is_context_parallel_extend(self, include_draft_extend_v2=False):
return False
def is_draft_extend(self, include_v2=False):
return True
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_seq_lens_cpu=[4],
extend_prefix_lens_cpu=[0],
token_to_kv_pool=SimpleNamespace(page_size=64),
forward_mode=DraftMode(),
)
with (
patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
return_value=True,
),
):
self.assertTrue(can_cp_split(8, 8, True, forward_batch))
def test_nsa_use_prefill_cp_enables_cp_draft_shared_kv_draft_extend(self):
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
forward_mode=ForwardMode.DRAFT_EXTEND,
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=1),
)
with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}):
self.assertTrue(nsa_use_prefill_cp(forward_batch, True))
def test_can_cp_split_uses_compute_padding_per_request_for_batched_tiny_suffix(
self,
):
@@ -512,8 +554,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(plan.request_last_token_local_offset, [0])
# Compatibility aliases for cache/page accounting stay valid-token
# based. Query-length metadata is split separately below: attention and
# top-k consume compute rows, cache/current paths consume valid rows.
# based. Query-length metadata exposes both valid and compute rows:
# consumers must choose the view that matches their actual q layout.
self.assertEqual(plan.request_split_lists, plan.request_valid_split_lists)
self.assertEqual(plan.request_padded_pages, plan.request_valid_padded_pages)
self.assertEqual(plan.request_actual_seq_q_prev, [64])
@@ -523,6 +565,67 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(plan.request_compute_seq_q_prev, [64])
self.assertEqual(plan.request_compute_seq_q_next, [0])
def test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias(self):
from sglang.srt.layers.attention.nsa.nsa_indexer import (
_select_batch_topk_query_lengths,
)
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[40387],
prefix_lens=[0],
page_size=64,
cp_size=8,
cp_rank=0,
)
valid_local_rows = (
plan.request_valid_seq_q_prev[0] + plan.request_valid_seq_q_next[0]
)
compute_local_rows = (
plan.request_compute_seq_q_prev[0] + plan.request_compute_seq_q_next[0]
)
self.assertEqual(valid_local_rows, 4995)
self.assertEqual(compute_local_rows, 5056)
selected = _select_batch_topk_query_lengths(
cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan),
batch_plan=plan,
batch_size=1,
q_tokens=valid_local_rows,
weights_tokens=valid_local_rows,
)
self.assertFalse(selected.uses_compute_query_rows)
self.assertEqual(selected.request_seq_q_prev, plan.request_valid_seq_q_prev)
self.assertEqual(selected.request_seq_q_next, plan.request_valid_seq_q_next)
self.assertEqual(
selected.request_valid_seq_q_prev, plan.request_valid_seq_q_prev
)
self.assertEqual(
selected.request_valid_seq_q_next, plan.request_valid_seq_q_next
)
selected_compute = _select_batch_topk_query_lengths(
cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan),
batch_plan=plan,
batch_size=1,
q_tokens=compute_local_rows,
weights_tokens=compute_local_rows,
)
self.assertTrue(selected_compute.uses_compute_query_rows)
self.assertEqual(
selected_compute.request_seq_q_prev, plan.request_compute_seq_q_prev
)
self.assertEqual(
selected_compute.request_seq_q_next, plan.request_compute_seq_q_next
)
self.assertEqual(
selected_compute.request_valid_seq_q_prev, plan.request_valid_seq_q_prev
)
self.assertEqual(
selected_compute.request_valid_seq_q_next, plan.request_valid_seq_q_next
)
def test_batch_plan_compute_padding_is_per_request_not_batch_total(self):
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[65, 1024],
@@ -752,6 +855,166 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
local_num_tokens=8,
)
def test_cp_draft_padding_keeps_local_hidden_when_static_tokens_are_shorter(self):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
cp_local_hidden_states=True,
)
forward_batch = ForwardBatch(
forward_mode=ForwardMode.DRAFT_EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=True,
)
with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}):
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
self.assertTrue(
torch.equal(
forward_batch.hidden_states_backup,
torch.ones((64, 2), dtype=torch.float32),
)
)
def test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible(
self,
):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
cp_local_hidden_states=True,
)
forward_batch = ForwardBatch(
forward_mode=ForwardMode.DRAFT_EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=False,
)
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
self.assertTrue(
torch.equal(
forward_batch.hidden_states_backup,
torch.ones((64, 2), dtype=torch.float32),
)
)
def test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite(
self,
):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
cp_local_hidden_states=True,
)
forward_batch = ForwardBatch(
# prepare_mlp_sync_batch can temporarily rewrite draft extend to
# EXTEND while static DP padding is being prepared. The draft
# side-channel contract must therefore be carried by spec_info, not
# inferred from forward_mode.
forward_mode=ForwardMode.EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=True,
)
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
self.assertTrue(
torch.equal(
forward_batch.hidden_states_backup,
torch.ones((64, 2), dtype=torch.float32),
)
)
def test_cp_draft_padding_rejects_unmarked_oversized_hidden(self):
import torch
class FakeAttnBackend:
def get_cuda_graph_seq_len_fill_value(self):
return 0
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
spec_info = EagleDraftInput(
hidden_states=torch.ones((64, 2), dtype=torch.float32),
verified_id=torch.tensor([1], dtype=torch.int64),
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
)
forward_batch = ForwardBatch(
forward_mode=ForwardMode.DRAFT_EXTEND,
batch_size=1,
input_ids=torch.arange(8, dtype=torch.int64),
req_pool_indices=torch.tensor([0], dtype=torch.int64),
seq_lens=torch.tensor([4], dtype=torch.int32),
out_cache_loc=torch.arange(8, dtype=torch.int64),
seq_lens_sum=4,
positions=torch.arange(8, dtype=torch.int64),
lora_ids=[None],
spec_info=spec_info,
uses_cp_shared_kv=False,
)
with self.assertRaisesRegex(
RuntimeError,
r"\[CP_SHARED_KV_FAIL_FAST\]\[draft_hidden_static_padding_mismatch\]",
):
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
def test_full_rerange_fails_fast_for_batch_metadata(self):
import torch
@@ -1087,6 +1350,90 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local[0].tolist(), [128.0, 129.0])
self.assertTrue(torch.equal(local[1:], torch.zeros((63, 2))))
def test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows(self):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
forward_batch = SimpleNamespace(
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
)
)
tensor = torch.arange(8 * 2, dtype=torch.float32).view(8, 2)
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(local.shape, (4, 2))
self.assertTrue(torch.equal(local, torch.zeros((4, 2))))
def test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding(
self,
):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[7],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
self.assertFalse(plan.compute_padding_enabled)
forward_batch = SimpleNamespace(
extend_num_tokens=8,
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
),
)
tensor = torch.arange(8 * 2, dtype=torch.float32).view(8, 2)
expected = split_tensor_by_cp_batch_plan(
tensor[:7],
plan,
mode="data",
)
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.assertTrue(torch.equal(local, expected))
def test_cp_split_valid_kind_rejects_trailing_padding_rows(self):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
with self.assertRaisesRegex(
RuntimeError,
"batch_gt1_split_input_len_mismatch",
):
split_tensor_by_cp_batch_plan(
torch.arange(8),
plan,
mode="1d",
split_kind="valid",
)
def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self):
import torch
@@ -1186,6 +1533,41 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local.shape, (64,))
self.assertEqual(local.tolist(), list(range(40384, 40448)))
def test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding(
self,
):
import torch
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[7],
prefix_lens=[0],
page_size=4,
cp_size=2,
cp_rank=1,
)
self.assertFalse(plan.compute_padding_enabled)
forward_batch = SimpleNamespace(
extend_num_tokens=8,
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
),
)
positions = torch.arange(8, dtype=torch.int32)
expected = split_tensor_by_cp_batch_plan(
positions[:7],
plan,
mode="position",
)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local = cp_split_and_rebuild_position(forward_batch, positions)
self.assertTrue(torch.equal(local, expected))
def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self):
from types import SimpleNamespace
@@ -1299,6 +1681,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_num_tokens=8,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
@@ -1368,6 +1751,136 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertEqual(local_locs.tolist(), [2 * page_size])
def test_local_out_cache_loc_ignores_trailing_static_padding_locs(self):
import torch
page_size = 4
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=page_size,
cp_size=2,
cp_rank=0,
)
valid_locs = torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64)
static_padding_locs = torch.arange(
99 * page_size, 100 * page_size, dtype=torch.int64
)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_num_tokens=8,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
cp_rank=0,
),
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
page_aligned=True,
page_size=page_size,
extend_prefix_len=0,
),
out_cache_loc=torch.cat((valid_locs, static_padding_locs)),
)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
self.assertEqual(local_locs.tolist(), valid_locs.tolist())
def test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding(
self,
):
import torch
page_size = 4
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[7],
prefix_lens=[0],
page_size=page_size,
cp_size=2,
cp_rank=1,
)
self.assertFalse(plan.compute_padding_enabled)
valid_locs = torch.cat(
(
torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64),
torch.arange(2 * page_size, 2 * page_size + 3, dtype=torch.int64),
)
)
static_padding_locs = torch.tensor([99 * page_size], dtype=torch.int64)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_num_tokens=8,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
cp_rank=1,
),
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
page_aligned=True,
page_size=page_size,
extend_prefix_len=0,
),
out_cache_loc=torch.cat((valid_locs, static_padding_locs)),
)
with patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
):
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
self.assertEqual(
local_locs.tolist(),
[2 * page_size, 2 * page_size + 1, 2 * page_size + 2],
)
def test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding(
self,
):
import torch
page_size = 4
plan = build_batch_page_aligned_in_seq_split_plan(
extend_lens=[4],
prefix_lens=[0],
page_size=page_size,
cp_size=2,
cp_rank=0,
)
self.assertTrue(plan.compute_padding_enabled)
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
cp_shared_kv_layout=CpSharedKVLayout(
page_size=page_size,
cp_size=2,
cp_rank=0,
),
nsa_cp_metadata=NSAContextParallelMetadata(
batch_size=1,
batch_plan=plan,
page_aligned=True,
page_size=page_size,
extend_prefix_len=0,
),
out_cache_loc=torch.cat(
(
torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64),
torch.arange(99 * page_size, 100 * page_size, dtype=torch.int64),
)
),
)
with self.assertRaisesRegex(RuntimeError, "static_padded=None"):
get_cp_shared_kv_local_out_cache_loc(forward_batch)
def test_batch_local_physical_out_cache_loc_reuses_layer_invariant_plan(self):
import torch
from types import SimpleNamespace