diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md index 9ac67d172..8deec955d 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md @@ -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,不是本次进程退出原因。 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 6f426868e..b3469e407 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index d7186dbb2..bbc42ea0d 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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: diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 32a88b927..7ff22e258 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 572a76f6a..7843f9ed6 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -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 diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 35ce3ea41..aaf9e0cce 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -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, diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 9eae9249f..cd5f90fec 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -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