From c8510938b9e446bec2c934b658fea8e524c7361f Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Fri, 5 Jun 2026 06:25:06 +0800
Subject: [PATCH] Repair CP current-slot composition instead of disabling reuse

CP shared KV current-only and partial-current paths were previously unsafe because owner-local current rows could be filled into dense page slots without synchronizing the current slots across CP ranks. This preserves the fast path and fixes the missing synchronization at the actual contract boundary: prefix slots remain materialized by the existing IPC/collective path, while freshly produced current slots are composed from current forward tensors and reduced only over the current slot ranges.\n\nThe change also keeps page-tail compute padding explicit, restores current-only index compose, fixes TAI index descriptor dtype consistency, and records the investigation ledger to avoid repeating discarded hypotheses.\n\nConstraint: CP shared KV must preserve current reuse and TAI materialize fast paths; blanket disable is not an acceptable fix.\nConstraint: Prefix slots must not be reduced twice after IPC/prefix materialization.\nRejected: Disable current-only MLA/index reuse | hides the owner-local composition bug and regresses the intended fast path.\nRejected: Disable TAI materialize globally | avoids symptoms without proving the byte/layout contract.\nConfidence: medium\nScope-risk: broad\nDirective: Do not remove current-slot range synchronization unless replacing it with an equivalent owner-aware P2P/IPC gather contract.\nTested: Local py_compile for touched files.\nTested: Remote g0034 py_compile for touched runtime/test files.\nTested: Remote test_cp_shared_kv_runtime.py: 111 passed, 5 warnings, 2 subtests passed.\nTested: Remote targeted current-slot compose tests and direct current-only index compose script.\nNot-tested: Full ETE output quality and decode accept len after restart.\nNot-tested: Full test_nsa_cp_utils.py collection on remote due incomplete installed sgl_kernel import/fake-op environment.
---
 ...accept_len1_garbage_output_debug_ledger.md | 371 ++++++++++++++++++
 .../attention/nsa/cp_shared_kv_prefetch.py    |  29 ++
 .../attention/nsa/cp_shared_kv_runtime.py     | 186 ++++++++-
 .../srt/layers/attention/nsa/nsa_indexer.py   |  20 +-
 .../sglang/srt/layers/attention/nsa/utils.py  |   6 +-
 .../srt/layers/attention/nsa_backend.py       |  19 +
 .../unit/layers/test_nsa_cp_utils.py          | 215 +++++++---
 .../mem_cache/test_cp_shared_kv_runtime.py    | 115 +++++-
 8 files changed, 898 insertions(+), 63 deletions(-)

diff --git a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md
index c9098253d..7254dbecf 100644
--- a/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md
+++ b/docs/advanced_features/nsa_prefill_cp_accept_len1_garbage_output_debug_ledger.md
@@ -1011,3 +1011,374 @@ Verification status:
 - Local `py_compile` passed for `disaggregation/utils.py`, `disaggregation/mooncake/conn.py`, and the updated unit test.
 - Local pytest is blocked by missing runtime dependency `orjson` in this environment.
 - Remote RED/GREEN execution could not be completed in this pass because `ssh g0034` timed out.
+
+## 2026-06-05 当前现象修正：decode 已结束但输出仍乱码/重复，accept len=4 不是恢复正确性
+
+远端最新日志：
+
+- prefill: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260604_181200.log`
+- decode0: `/mnt/beegfs/cjy/log/decode0_20260604_181227.log`
+- decode1: `/mnt/beegfs/cjy/log/decode1_20260604_181237.log`
+
+新增事实：
+
+1. 用户确认：请求能返回结束，但内容依然重复、无意义。这说明当前主症状不是“scheduler 不停止”，而是模型 token 分布/上下文状态已经损坏。
+2. decode1 在 18:17:52-18:19:29 持续显示 `accept len: 4.00, accept rate: 1.00`，并且 `#transfer-req: 0`、`#retracted-req: 0`。这不是之前的 accept collapse；它表示 draft 4 个 token 基本全部被 target 接受。若输出乱码，说明 target 验证路径本身也在错误上下文上同意了错误 token。
+3. decode warmup 本身已经异常：`prompt_tokens=4` 的 warmup 输出固定为 `! of the of the of the`，`completion_tokens=8`，`spec_accept_rate=0.25`。这发生在 40k cache-hit 之前，因此不能继续把问题限定为长 prompt/HiCache cache-hit。
+4. prefill warmup 同一阶段输出是 `text=' /'`，而 decode warmup 输出重复 `of the`。这把范围收窄到“prefill 产生并交给 decode 的 persistent KV / disagg transfer / decode FP8 KV/EAGLE 消费路径”，而不是 prefill 直接 forward 的 logits 路径。
+5. 日志中未看到 `CP_SHARED_KV_FAIL_FAST` transfer page-count mismatch。说明最新 page-count fail-fast 没抓到这个问题；如果仍是 transfer 问题，更可能是内容/order/scale 级别，而不是页数截断。
+6. decode 启动时所有 DP rank 打印：`Using FP8 KV cache but no scaling factors provided. Defaulting to scaling factors of 1.0.` 该风险仍未排除。FP8 KV 下若 scale/packed layout/transfer bytes 不一致，target 与 draft 都可能在同一个错误 KV 上稳定地产生重复 token。
+
+当前重新排序的假设：
+
+- H1（最高）：CP prefill persistent KV 在 FP8 下写入/transfer/decode 接收的内容、顺序或 scale 合同不一致。page-count fail-fast 只能排除页数截断，不能排除每页内容错位、FP8 scale 行错位、或 decode 端按不同 layout 解释。
+- H2：decode FP8 KV cache / EAGLE 在没有 scale factors 时本身就会产生退化输出；需要用同一 disagg 配置但禁用 FP8 KV 或禁用 EAGLE 做最小分叉验证。
+- H3：长 prompt/HiCache cache-hit 不是当前首要根因，因为 4-token decode warmup 已异常。
+
+禁止重复误区：
+
+- 不要再把 `accept len=4.00` 当成“恢复正常”；它只能说明 speculative accept 高，不代表 token 语义正确。
+- 不要只盯 40k cache-hit；短 warmup 已经提供更小复现面。
+- 不要继续只检查 page-count mismatch；当前 fail-fast 未触发，下一步应查 FP8 packed KV 的 byte layout、scale、source/destination page order。
+
+## 2026-06-05 cache-hit 口径核对：当前日志里的首个 40k prefill 请求未命中
+
+用户追问“这次第一次请求为什么会 cache hit”。对当前已检查的 prefill 日志 `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260604_181200.log`：
+
+- server warmup 请求在 18:15:17，`prompt_tokens=4`，prefill batch 显示 `#cached-token: 0`，但会插入/backup 一个 `logical_len=3` 的小 radix/HiCache node。
+- 40k 长请求在 18:16:41，prefill batch 行 1020-1027 显示 `#new-token: 40448, #cached-token: 0`。
+- 因此这份日志中，首个长业务请求没有 prefill cache hit。若客户端或另一轮日志显示“第一次请求 cache hit”，需要区分：
+  1. 该 prefill 进程是否真的重启/清 cache；radix/HiCache 是进程内状态，client/router 重启不会清。
+  2. replay 的“seq 第一条”不等于 scheduler 实际第一条；并发/重试可能先插入相同 prefix。
+  3. `cached_tokens` 客户端字段与 prefill log 的 `#cached-token` 可能不是同一条请求/同一统计口径。
+
+下一步如果再次出现“首条命中”，应先抓对应 log 文件和请求时间，核对该请求之前是否已有同 prompt/prefix 的 `POST /v1/chat` 或 warmup/replay 请求插入。
+
+## 2026-06-05 当前最高置信根因：bs=1 batch-plan 在非 page-aligned tail 下误判未启用 compute padding
+
+新增代码证据：
+
+- `python/sglang/srt/layers/attention/nsa/utils.py:2439-2457` 在 CP shared-KV 下即使 `batch_size=1` 也会构造 `CPSharedKVBatchPlan`。
+- `python/sglang/srt/layers/attention/nsa/utils.py:604-676` 同时构造两套 split：
+  - `request_split_lists` / valid split：只覆盖真实 token，tail page 的最后一个 segment 可能小于整页。
+  - `request_compute_split_lists` / compute split：以 page 为最小单位，tail page 按整页参与 attention/index compute。
+- `python/sglang/srt/layers/attention/nsa/utils.py:718-723` 旧判定用 `request_compute_padded_tokens != request_valid_padded_tokens` 决定 `compute_padding_enabled`。这只能发现“为了覆盖至少 cp_size 个 page 而额外 pad”的场景，无法发现“总 page 数相同但最后一个 segment 从 tail token 扩到整页”的场景。
+- 远端最小复现：
+  - `extend_len=40387, page_size=64, cp_size=8, cp_rank=0`
+  - valid local rows = `2560 + 2435 = 4995`
+  - compute local rows = `2560 + 2496 = 5056`
+  - 旧代码却给出 `compute_padding_enabled=False`
+  - `40387` 正是当前业务请求的 `prompt_tokens`，因此这是当前乱码输出的强相关输入类。
+
+影响链路：
+
+1. `split_tensor_by_cp_batch_plan(..., split_kind="compute")` 在 `compute_padding_enabled=False` 时会选 valid split，返回 4995 行。
+2. batch plan 内的 `request_actual_seq_q_prev/next`、`request_compute_seq_q_prev/next` 已按 compute split 记录 5056 行的语义。
+3. 于是同一个 forward 内存在“tensor 实际 CP-local 行数是 valid rows，但 metadata/attention/index 部分消费者认为是 compute rows”的不一致。
+4. index top-k 路径已有 `_select_batch_topk_query_lengths()` 的局部自适应逻辑，能在某些地方避免 fail-fast；但 attention/MLA/cache write 仍可能在不一致 metadata 下写入或读取错误 CP-local row，导致 prefill 生成错误 KV/logits，decode 在错误上下文上高 accept 但输出乱码。
+
+结论：
+
+- 这比 FP8 scale、Mooncake page-count mismatch 更贴近当前现象：它只在 non-page-aligned 长 prompt 下触发，能解释为什么 `prompt_tokens=40387` 的请求输出不说人话，同时 fail-fast page-count 没有触发。
+- 修复方向不是扩大 fallback，而是让 `compute_padding_enabled` 表示“valid split 与 compute split 是否不同”，而不是只比较 padded-token 总量。
+- 回归测试必须覆盖 `40387/64/CP8`：期望 `compute_padding_enabled=True`，compute split 返回 5056 local rows，cache write valid selector 过滤回 4995 rows。
+
+## 2026-06-05 修复记录：compute_padding_enabled 改为 split-shape 语义
+
+代码修复：
+
+- `build_batch_page_aligned_in_seq_split_plan()` 的 `compute_padding_enabled` 从“比较 valid/compute padded-token 总量”改为“比较 valid/compute split list 是否完全一致”。
+- 这样 `extend_len=40387,page_size=64,cp_size=8` 会正确进入 compute-padding 路径：
+  - CP rank0 compute local rows = 5056。
+  - cache write valid selector 继续过滤回 valid local rows = 4995。
+- 保留 page-aligned、valid split 与 compute split 完全相同的场景为 `compute_padding_enabled=False`；对应 static padding 测试已改为 `extend_len=8,page_size=4,cp_size=2`，避免继续把 non-page-aligned tail 误认为“无 compute padding”。
+
+验证：
+
+- 本地 `py_compile` 通过：`utils.py`、`test_nsa_cp_utils.py`。
+- 本地 pytest 仍被环境缺失 `orjson` 阻塞。
+- 远端 `g0034` 容器内通过：
+  - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py`
+  - 4 个针对性回归测试：`4 passed`
+  - 全量 `test/registered/unit/layers/test_nsa_cp_utils.py`：`87 passed`
+
+剩余验证缺口：
+
+- 还需要用户用新 prefill 进程重新打同一个 40387-token 请求，确认输出是否恢复；当前修复解释了现象，但 ETE 仍必须验证。
+- FP8 scale / Mooncake 内容顺序问题仍是次级假设；如果该修复后仍乱码，下一轮应继续查 FP8 packed KV byte layout 与 transfer page order，而不是回到 page-count 或 cache-hit 口径。
+
+## 2026-06-05 再次乱码后的新定位：current-only index current reuse TAI fast path 会静默补零
+
+最新事实：
+
+- compute-padding split-shape 修复后，用户重新验证仍然输出乱码/重复，说明上一轮修复只解决了一个真实不一致，但不是唯一根因。
+- 最新 prefill 日志中的 40k 首请求仍是 cache miss（`#cached-token: 0`），所以当前错误必须能发生在 current-only extend 路径。
+
+新增代码证据：
+
+1. `nsa_indexer.py::_can_reuse_current_index_kv()` 在 CP shared-KV + current reuse 开启时，对 current-only extend 也返回 true。
+2. `nsa_indexer.py::forward_cuda()` 随后用 `select_cp_local_valid_rows_for_cache_write()` 准备 `current_index_kv`。在 CP shared-KV batch plan 下，这个 tensor 是当前 CP rank 的 owner-lane compact rows，不是 request 的全局 dense rows。
+3. `nsa_indexer.py::_get_topk_in_seq_cp_pair_batch()` 对 `current_only_batch` 不调用 `_maybe_materialize_shared_index_buffer()`，而是把 compact `current_index_kv` 直接传入 `_get_topk_ragged_with_cp()`。
+4. `_get_topk_ragged_with_cp()` 的 TAI current-index batch descriptor 使用 `kv_lens_list=end_seq_position`。对首个 40k current-only 请求，segment 的 `kv_len` 是全局序列长度量级；但 `current_bases_list` 来自本 rank local offset，`current_index_kv` 行数只有本 rank owner rows（例如 rank0 约 4995/5056 行）。
+5. `tai-kernel/python/tai_kernel/nsa_prefill/cp_index_mqa_prepare.py::_prepare_cp_mqa_current_kv_batch_kernel` 的 mask 是 `src_tokens < CURRENT_ROWS`。当 `kv_len` 大于 compact local rows 时，超出的全局 K 行被静默写成 0，而不是 fail-fast。
+
+结论：
+
+- 这解释了“uint8 修复后不再 fallback/报错，但模型仍不说人话”：以前 TAI current-index path 因 dtype 不匹配失败并回退；修复 dtype 后快路径开始成功执行，但它把 CP-local compact current K 当成 global dense K 使用，导致 top-k 的 index K 大量补零。
+- 该问题能直接污染 target index/top-k，进而让 target attention 在错误上下文下生成/接受错误 token；这比 decode 端 EAGLE 单独异常更符合 `accept len` 高但输出乱码的现象。
+
+当前修复策略：
+
+- 暂时禁用 **CP shared-KV current-only 的 compact current-index reuse**，并打印醒目的 fallback warning。
+- current-only index top-k 改回 full shared index materialize 路径：先写入各 CP owner page，再通过 shared paged buffer materialize/all-reduce 得到全局 dense index K。
+- partial-current（有 page-aligned prefix）仍保留现有 prefix+current compose 路径；它有 prefix page contract 和 current loc shape 校验，另行验证。
+
+后续如果要重新打开 current-only fast path，必须新增 descriptor，显式表达“每个 CP rank 的 local owner rows 如何映射到全局 token positions”，不能继续用 `current_base + token_offsets` 当 global dense row。
+
+## 2026-06-05 修复记录：禁用 current-only compact index current reuse
+
+代码修复：
+
+- `Indexer._can_reuse_current_index_kv()` 在 current-only extend 下返回 false，并打印一次 `[CP_SHARED_KV_FALLBACK][index_current_reuse] reason=current_only_compact_unsupported`。
+- `_get_topk_in_seq_cp_pair()` 与 `_get_topk_in_seq_cp_pair_batch()` 增加防御：即使调用者误传 compact `current_index_kv`，current-only CP shared-KV index top-k 也会忽略它，改走 full shared index materialize，并打印一次 `current_only_compact_ignored`。
+- partial-current index reuse 不变：有正的 page-aligned prefix 时仍使用 `_maybe_materialize_shared_index_buffer(..., current_index_kv=...)` 的 prefix+current compose。
+
+为什么这次修复针对当前乱码：
+
+- current-only cache-miss 40k 请求会先写各 CP owner page 的 index K，再做 full shared materialize；不会再把本 rank compact rows 当作全局 dense rows 喂给 TAI current-index top-k kernel。
+- 这避免了 TAI kernel 用 `src_tokens < CURRENT_ROWS` 静默把大部分 global K rows 填 0 的问题。
+
+验证：
+
+- 本地 `py_compile` 通过：`nsa_indexer.py`、`test_nsa_cp_utils.py`、`test_cp_shared_kv_runtime.py`。
+- 本地 pytest 仍被缺失 `orjson` 阻塞。
+- 远端 `g0034` 容器内通过：
+  - `python -m py_compile python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
+  - 针对性 3 个回归：`3 passed`
+  - `test_nsa_cp_utils.py` + `test_cp_shared_kv_runtime.py` 全量：`196 passed, 2 subtests passed`
+
+剩余验证缺口：
+
+- 还需要用重新启动后的 prefill/route 流量验证 40387-token 请求输出是否恢复。
+- 如果仍乱码，下一轮应继续查 FP8 MLA/current reuse 或 transfer 内容顺序；但当前 index current-only 静默补零已经是明确 bug，不能再回到该路径。
+
+## 2026-06-05 bs>1 改造后 bs=1 也坏的回归链路
+
+最新用户反馈：index current-only gate 生效后，输出仍然不说人话，并且 decode accept len 很低。最新日志确认：
+
+- prefill 首个 40k 请求仍是 cache miss：`#new-token: 40448, #cached-token: 0`。
+- decode 端该请求 accept len 约 `1.0~1.5`，不是高 accept 误接收，而是 draft/target 在错误上下文下明显不一致。
+- `index_current_reuse` 的 current-only compact gate 已在 warmup 生效，说明远端代码不是旧版本。
+
+新的关键差异：
+
+1. bs>1 改造后，`utils.py` 的 CP metadata 入口只要 `uses_cp_shared_kv` 且存在 `extend_seq_lens_cpu/extend_prefix_lens_cpu`，即使 `batch_size=1` 也会构造 `CPSharedKVBatchPlan`，不再走旧 scalar bs=1 path。
+2. commit `50d000870` 去掉了 current reuse 的 `batch_size == 1` 保守限制，并把 current-only extend 直接判定为可 reuse。
+3. 旧 scalar bs=1 下，`current_extend_kv_rows_for_reuse()` 用 global `extend_len` 校验 tensor rows；CP split 后的 `k/k_rope` 通常只有本 rank local rows，因此 current reuse 不会启用，会走 full shared materialize。
+4. 新 batch/compute-padding path 下，`nsa_backend.py` 先用 `select_cp_local_valid_rows_for_cache_write()` 得到本 rank owner-local current rows，再用 `local_out_cache_loc.numel()` 校验，因此 current-only MLA reuse 会被启用。
+
+当前最高置信破坏点：MLA/attention current-only page-slot compose。
+
+- `materialize_prefix_and_reuse_current_kv_page_slots(prefix_pages=0)` 会创建 dense slot buffer。
+- prefix 为 0 时，不会 materialize/reduce 任何 prefix rows。
+- 随后 `fill_current_kv_page_slots_and_remap_locs()` 只把本 rank owner-local current rows 填进 dense slot。
+- 目前该函数不会对 current/suffix page slots 做 all-reduce/all-gather。
+- 因此 attention 可能看到：本 rank owner pages 有值，其他 rank owner pages 是 0/stale，直接污染 prefill KV/logits。
+
+这解释了为什么“之前 bs=1 可以，现在 bs>1 后同样 bs=1 错”：不是 scheduler 真组了 bs>1，而是 bs>1 metadata/current-reuse 改造改变了 bs=1 的 fast path，使其开始走 owner-local current reuse。
+
+下一步验证：
+
+- 增加单测证明 `materialize_prefix_and_reuse_current_kv_page_slots(prefix_pages=0)` 对 cp_size>1 且 current_locs 只含本 rank owner rows 时，不会聚合其他 owner pages。
+- 修复方向二选一：
+  1. 保守 correctness：CP shared-KV current-only MLA reuse 先禁用，回到 full shared materialize。
+  2. 正式 fast path：填 current rows 后，对 current page slots 执行 owner-aware 聚合，仅 reduce/gather current slots，不能重复 reduce 已经 reduced 的 prefix slots。
+
+短期为了恢复正确性，优先做 1；随后再单独实现 2 以恢复性能。
+
+## 2026-06-05 修复记录：current-only MLA/current reuse 回退到 full shared materialize
+
+短期 correctness 修复：
+
+- `should_reuse_current_extend_kv()` 在 `uses_cp_shared_kv=True` 且 current-only extend 时返回 false。
+- 新增 warning：`[CP_SHARED_KV_FALLBACK][current_reuse] reason=current_only_owner_local_unsupported`。
+- 这样 current-only cache-miss 请求不会再进入 MLA current page-slot compose，而是走 full shared materialize：先从 owner-local persistent pool materialize owned pages，再通过 collective/IPC 得到每个 rank 完整 dense KV view。
+
+为什么这能解释并修复 bs>1 后 bs=1 回归：
+
+- bs>1 改造让 bs=1 也构造 `CPSharedKVBatchPlan`，并在 compute-padding path 中把 current KV rows 收窄为本 rank owner-local compact rows。
+- current-only MLA reuse 没有 current page-slot 聚合，attention 会消费不完整 dense KV。
+- 禁用 CP shared current-only reuse 后，bs=1 cache miss 回到旧的 full materialize 消费语义，避免 owner-local compact rows 被当作 global dense rows。
+
+保留内容：
+
+- `is_current_only_extend_batch()` 本身不变，它仍是描述性 helper。
+- 非 CP shared-KV 的 current-only reuse 不受影响。
+- partial-current reuse 暂未关闭；如果后续 cache-hit 仍有异常，需要继续实现 current suffix page-slot 聚合或临时关闭 partial-current。
+
+验证：
+
+- 本地 `py_compile` 通过：`cp_shared_kv_runtime.py`、`test_cp_shared_kv_runtime.py`。
+- 远端 `g0034` 容器内通过：
+  - 新增 targeted test：`1 passed`
+  - `test_cp_shared_kv_runtime.py`：`110 passed, 5 warnings, 2 subtests passed`
+  - `test_nsa_cp_utils.py test_cp_shared_kv_runtime.py`：`197 passed, 5 warnings, 2 subtests passed`
+
+剩余验证缺口：
+
+- 需要用户重启 prefill 后重新打 40387-token 请求，确认输出和 decode accept len 是否恢复。
+- 若 current-only 恢复但 cache-hit partial-current 仍坏，下一步应实现“fill current rows 后只聚合 current page slots”的正式 fast path，而不是继续让 owner-local compact current rows直接给 attention 消费。
+
+## 8. 2026-06-05 继续排查：current-only MLA 禁用已生效但乱码仍存在
+
+新增证据来自：
+
+- prefill：`/mnt/beegfs/cjy/log/sglang_cp_hicache_20260604_200924.log`
+- decode：`/mnt/beegfs/cjy/log/decode0_20260604_200912.log`
+
+确认事实：
+
+1. 最新 prefill 进程确实使用了当前代码：日志 885-952 多次出现
+   `[CP_SHARED_KV_FALLBACK][current_reuse] reason=current_only_owner_local_unsupported`。
+   因此“current-only MLA/current reuse 仍在启用”不是最新复现的主因。
+2. decode 侧 warmup 在 prefill 20:12:48 ready 之前（20:11:20）已经输出异常：
+   `!!!!!!!??!??!` / `! of the of the of the`，并且 spec accept rate 只有 0.13-0.25。
+   这说明当前坏现象至少不只来自长 prompt HiCache 命中或 Mooncake transfer；decode 自身/共享代码路径/FP8 KV cache 也需要被纳入排查。
+3. prefill warmup 对 4-token request 的 CP shared-KV path 中，`current_only_owner_local_unsupported` 生效后仍反复触发
+   `[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] reason=batch_kernel_failed ... error=descriptor tensors must share dtype segments=1 total_kv_len=4 total_q_count=4`。
+   这说明 index MQA prepare fast path 在 FP8/当前 dtype 组合下没有工作，但尚未证明 fallback 本身会破坏语义。
+4. 长 40k 请求仍有 `tai_index_mqa_prepare` fallback：`segments=2 total_kv_len=43072..43456 total_q_count=5056`。
+   当前不能再把问题只归咎到 current-only reuse；下一步需要检查 full materialize、persistent KV direct write、FP8 KV store/load、decode flashmla_kv 的共同合同。
+5. `omx explore` 在本机因 Cargo.lock version 4 解析失败，后续本轮改用普通 grep/sed；该失败不是模型路径证据。
+
+当前高优先级假设（待验证）：
+
+- H1：FP8 KV cache 的 persistent store 或 full materialize page-slot compose 写入/读取的 dtype/layout 与 flashmla_kv decode 消费合同不一致。支持证据是 decode warmup 已异常，且当前运行使用 `kv_cache_dtype='fp8_e4m3'`、decode backend `flashmla_kv`。
+- H2：`tai_index_mqa_prepare` fast path 因 descriptor dtype 不一致失败后，Python per-segment fallback 在 compute padding / page-tail / FP8 下构造的 K/S 或 top-k order 不等价。支持证据是所有最新复现都有该 warning；反证是该 warning 可能只是性能 fallback，需要单测或最小 CPU/GPU 对比证明。
+- H3：CP shared-KV full materialize 对 current-only prefix=0 读 persistent pool，而 persistent pool direct write 在当前流/TAI fused store/FP8 packed layout 下未完成或写错。支持证据是 `should_reuse_current_extend_kv=False` 后会从当前 k/k_rope fast path退到读 pool。
+
+已排除/降级：
+
+- 最新复现不再优先怀疑 current-only MLA/current reuse 本身，因为对应 warning 已证明它被禁用。
+- 最新复现不应只盯长 prompt/HiCache/transfer；decode warmup 异常把问题范围扩大到 decode 或共享 FP8 attention 路径。
+
+## 2026-06-05 继续排查：最新 ETE 仍乱码后的新增事实
+
+最新远端日志：
+
+- prefill: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260604_200924.log`
+- decode: `/mnt/beegfs/cjy/log/decode0_20260604_200912.log`
+
+新增证据：
+
+1. 最新 prefill 仍在 40k 首次请求上走 cache miss / current-only 场景：
+   - `Requested token count (40391 input + 200000 new)`；
+   - `Prefill batch, #new-seq: 1, #new-token: 40448, #cached-token: 0`；
+   - decode 后续 `#token: 40448/40512, accept len: 1.43/1.38`。
+
+2. `current-only MLA/current reuse` 禁用逻辑确实已生效：日志出现
+   `[CP_SHARED_KV_FALLBACK][current_reuse] reason=current_only_owner_local_unsupported ... Falling back to full shared materialize`。
+   因此“旧 current-only reuse 还在直接 corrupt”的假设降低优先级。
+
+3. 远端仍有大量 `tai_index_mqa_prepare` fallback：
+   `error=descriptor tensors must share dtype segments=2 total_kv_len=43072..43456 total_q_count=5056`。
+   代码检查确认：`nsa_indexer.py::_get_topk_ragged_with_cp()` 的 non-current batch fast path把 `batch_indices` 构造成 `int64`，而 `kv_lens/q_starts/q_lens/k_bases/q_bases` 是 `int32`；TAI C++ kernel 要求这些 descriptor tensors 同 dtype。
+   这解释了该 warning。它是明确性能/覆盖问题，但 fallback 路径仍应语义正确，暂不单独作为“乱码”的根因结论。
+
+4. TAI IPC slot-dense materialize 的 dummy page offset 已检查：
+   - `tai-kernel/python/tai_kernel/nsa_prefill/ipc.py` 文档声明 slot `i` 写 dense page `i+1`；
+   - `kvcacheio_lf_pf.cu::ipc_materialize_peer_pages_slot_dense_*` 实现使用 `dense_page = slot + 1`，并 zero page 0。
+   因此“IPC kernel 把 page 0/dummy page 写错导致整体 shift”的假设已排除。
+
+5. 新的更强 correctness 风险：`SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1` 下，`materialize_shared_token_kv_buffer()` 如果 `_try_tai_ipc_materialize_token_kv_page_slots_into()` 成功，会跳过 `_all_reduce_materialized_buffer()`。
+   对 prefix/stable cache page 这可能成立；但对当前 cache miss 的 current-only 请求，MLA direct write 刚把本层 owner-local KV 写入各 rank 的 pool，随后同一层 attention 立即通过 IPC 读取 peer rank 的 pool。这里没有 per-layer cross-rank event/barrier 保证 peer rank 的 direct-write kernel 已完成。
+   旧 all-reduce materialize 至少通过 collective 提供跨 rank stream/order 同步；IPC peer read 跳过 collective 后可能读到未写完或旧 KV。这与“prefill HTTP 200、decode 消费 persistent KV 后 accept len≈1/输出乱码”高度匹配。
+
+当前优先假设：
+
+- H1': 当前-layer fresh KV 的 CP full materialize 不能使用无同步的 TAI IPC peer read。需要先在 current-only/fresh-current 场景禁用 IPC materialize，回到 local materialize + all-reduce；后续如果要保留 IPC，需要设计显式跨 rank write-complete 同步或恢复正确 current-reuse dense compose。
+
+待验证/修复：
+
+- 给 `materialize_shared_token_kv_buffer()` 增加 `allow_ipc_materialize` 控制；在 MLA full materialize callsite 对 `is_current_only_extend_batch(forward_batch)` 传 `False`。
+- 同时修正 `tai_index_mqa_prepare` non-current batch descriptor dtype：`batch_indices` 应与其他 descriptor 同 dtype，优先 `int32`。
+- 补单测：current-only MLA full materialize callsite 必须禁止 IPC；non-current TAI batch descriptor dtype 必须统一。
+
+## 2026-06-05 修正：不做 blanket disable，必须修 fresh-current 合同
+
+用户明确指出：不能盲目禁止 fast path；必须找到问题并尝试修复。
+
+纠正当前方向：
+
+- 不把 `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE` 全局关掉，也不简单在 current-only 下禁止 IPC 作为最终方案。
+- 真正需要验证的是 fresh current KV 的 source-of-truth：同一层 forward 里刚由 direct-write 写入 persistent pool 的 current KV，是否可以立即作为 full-materialize 的稳定读源。
+- 若问题成立，正确修复应是 owner-aware current compose：prefix/stable pages 仍可用 IPC/collective materialize，当前 extend rows 从本层已有的 CP-local `k/k_rope` 源张量按 logical loc 填入 dense slot，并在必要时只同步缺失 owner pages；不能简单退回慢路径。
+
+下一步检查：
+
+1. 阅读 `materialize_prefix_and_reuse_current_kv_page_slots()` 与 `fill_current_token_kv_page_slots_and_remap_locs()`，确认它是否已经具备 owner-local current row 填充合同。
+2. 确认当前为什么禁用了 current-only reuse：是缺少全 rank current rows，还是调用侧传错 loc/base。
+3. 构造一个最小单测验证：current-only CP shared KV 不应从 persistent `kv_cache` 读取 current pages；应从 current source tensor 填入 dense page slots。这个测试先 RED，再修。
+
+## 2026-06-05 修正方向：不禁用 current fast path，修 owner-local current compose 后的同步合同
+
+用户指出不能继续靠盲目禁用定位。重新审视后，当前修复方向调整为恢复 current-only fast path，并修正它的真实缺口：
+
+- current-only / partial-current compose 的 source-of-truth 应该是当前 forward 产生的 owner-local `k/k_rope` 或 index K，而不是刚写入 persistent pool 的 fresh current pages。
+- 但是 owner-local current rows 填入 slot-dense buffer 后，原实现没有对 current slots 做 CP 同步；每个 rank 只能看到自己 owner lane 的 current pages，其它 owner pages 仍是 0/stale。
+- 这会同时影响 MLA attention current reuse、index current reuse，以及 prefetch consume-prefix-with-current 路径。
+
+本轮代码修复：
+
+1. 新增 `build_batch_current_slot_spans()`：用 CPU `prefix_lens_cpu/extend_lens_cpu` 计算 current page slot spans，避免扫描 CUDA page table。
+2. `materialize_prefix_and_reuse_current_kv_page_slots()`：填入 current rows 后，只对 current token slot rows 做 `_all_reduce_materialized_buffer_range()`；prefix rows 不重复 reduce。
+3. `materialize_prefix_and_reuse_current_index_page_slots()`：同样只对 current index page slots 做 range reduce。
+4. `CpSharedKVMlaPrefetcher.consume_prefix_with_current()` 与 `CpSharedKVIndexPrefetcher.consume_prefix_with_current()`：prefetch 已经完成 prefix reduce/IPC 后，consume 时填 current rows，再对 suffix/current slots 做同步。
+5. `Indexer._can_reuse_current_index_kv()` 不再 blanket 禁用 current-only；`_maybe_materialize_shared_index_buffer()` 允许 zero-prefix current-only compose，并把 `current_index_kv` 传入 compose。
+6. 修复 TAI non-current index batch descriptor dtype：`batch_indices` 从 `int64` 改为 `int32`，与 `kv_lens/q_starts/q_lens/k_bases/q_bases` 保持一致，避免 `descriptor tensors must share dtype` fast-path fallback。
+
+为什么这不是盲目禁用：
+
+- 没有关 `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE`。
+- 没有关 current-only MLA/index reuse。
+- 修复的是 owner-local current rows 填入 dense slot 后缺少跨 CP rank 聚合这一明确合同漏洞。
+
+验证：
+
+- 本地 `py_compile` 通过：`cp_shared_kv_runtime.py`、`nsa_indexer.py`、`nsa_backend.py`、`cp_shared_kv_prefetch.py`、相关单测文件。
+- 远端 `g0034` 容器内 `py_compile` 通过同一组文件。
+- 远端 targeted runtime tests 通过：current slot span、token current compose reduce、index current compose reduce、owner-local current rows gate。
+- 远端完整 `test_cp_shared_kv_runtime.py` 通过：`111 passed, 5 warnings, 2 subtests passed`。
+- `test_nsa_cp_utils.py` 当前远端 pytest collection 被环境里的 `sgl_kernel`/fake-op 注册问题阻塞；已用直接 Python 脚本验证：
+  - TAI shared batch descriptor `batch_indices.dtype == torch.int32`。
+  - current-only index path 会把 `current_index_kv` 交给 materialize/compose，而不是丢弃后读 persistent full materialize。
+
+剩余验证缺口：
+
+- 需要用户重启 prefill 并重新打业务请求，确认 40387-token 请求输出与 decode accept len 是否恢复。
+- 如果仍乱码，下一步继续查 FP8 packed MLA KV compose/transfer byte layout，而不是回退到禁用 current fast path。
+
+### 2026-06-05 追加验证
+
+本轮继续验证不采用 blanket disable 的修复：
+
+- 远端 `g0034` 容器内 `py_compile` 再次通过：
+  `cp_shared_kv_runtime.py`、`nsa_indexer.py`、`nsa_backend.py`、`cp_shared_kv_prefetch.py`、`test_cp_shared_kv_runtime.py`、`test_nsa_cp_utils.py`。
+- 远端 `test_cp_shared_kv_runtime.py` 完整通过：
+  `111 passed, 5 warnings, 2 subtests passed`。
+- 远端 targeted tests 通过：
+  - `test_batch_current_slot_spans_follow_prefix_and_extend_pages`
+  - `test_materialize_prefix_current_token_kv_uses_ipc_and_reduces_current_slots`
+  - `test_materialize_prefix_current_index_uses_ipc_and_reduces_current_slots`
+  - `test_cp_shared_current_only_reuse_accepts_owner_local_current_rows`
+- 新增 direct 脚本验证 `Indexer._maybe_materialize_shared_index_buffer()` 对 current-only index compose 的合同：
+  - `prefix_pages == 0`
+  - `current_slot_spans == [(0, 2)]`
+  - `current_index_k/current_index_scale` 直接来自当前 forward source tensor
+  - 输出 `CURRENT_ONLY_INDEX_MATERIALIZE_OK`
+
+注意：
+
+- `test_nsa_cp_utils.py` 与 runtime tests 混跑时，远端 pytest collection 仍可能被当前环境的 `sgl_kernel` 安装不完整阻塞：
+  `cannot import name 'gelu_tanh_and_mul' from 'sgl_kernel'`。
+  这是测试环境 import 问题，不作为本轮语义验证失败；关键合同用 targeted/direct 验证覆盖。
diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
index d5b73cf74..e738654fd 100644
--- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
+++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
@@ -772,6 +772,21 @@ class CpSharedKVMlaPrefetcher:
             page_size=self.page_size,
             mask_non_current_in_current_pages=True,
         )
+        if self.layout.cp_size > 1 and self.prefix_pages < self.total_slots:
+            current_rows = slot_range_to_token_slice(
+                self.page_size,
+                self.prefix_pages,
+                self.total_slots,
+            )
+            _all_reduce_materialized_buffer_range(
+                mixed_kv_cache,
+                self.layout.cp_size,
+                current_rows.start,
+                current_rows.stop,
+                nvtx_source="mla.prefetch_current",
+                nvtx_layer_id=layer_id,
+                nvtx_cp_rank=self.layout.cp_rank,
+            )
         remap_ms = _cpu_timing_ms(remap_cpu)
         total_ms = _cpu_timing_ms(consume_cpu)
         self._log_layer(
@@ -1551,6 +1566,20 @@ class CpSharedKVIndexPrefetcher:
             page_size=page_size,
             index_head_dim=index_head_dim,
         )
+        if self.layout.cp_size > 1 and self.prefix_pages < self.total_slots:
+            current_pages = slot_range_to_page_slice(
+                self.prefix_pages,
+                self.total_slots,
+            )
+            _all_reduce_materialized_buffer_range(
+                dense_page_buffer,
+                self.layout.cp_size,
+                current_pages.start,
+                current_pages.stop,
+                nvtx_source="index.prefetch_current",
+                nvtx_layer_id=layer_id,
+                nvtx_cp_rank=self.layout.cp_rank,
+            )
         remap_ms = _cpu_timing_ms(remap_cpu)
         total_ms = _cpu_timing_ms(consume_cpu)
         self._log_layer(
diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
index 2d3532c21..63c3f3916 100644
--- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
@@ -10,6 +10,7 @@ import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.nsa.utils import (
+    get_cp_shared_kv_local_out_cache_loc,
     log_cp_draft_shared_kv_debug,
 )  # noqa: F401
 from sglang.srt.layers.dp_attention import get_attention_cp_group
@@ -2203,18 +2204,32 @@ def current_extend_kv_rows_for_reuse(
     extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
     if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) == 0:
         return None
-    valid_current_rows = sum(int(x) for x in extend_seq_lens_cpu)
-    if valid_current_rows <= 0:
+    global_current_rows = sum(int(x) for x in extend_seq_lens_cpu)
+    if global_current_rows <= 0:
         return None
 
     out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
-    if out_cache_loc is None or int(out_cache_loc.numel()) < valid_current_rows:
+    if out_cache_loc is None or int(out_cache_loc.numel()) < global_current_rows:
         return None
 
+    tensor_rows: list[int] = []
     for tensor in current_kv_tensors:
-        if tensor is None or int(tensor.shape[0]) < valid_current_rows:
+        if tensor is None:
             return None
-    return valid_current_rows
+        tensor_rows.append(int(tensor.shape[0]))
+
+    if all(rows >= global_current_rows for rows in tensor_rows):
+        return global_current_rows
+
+    if getattr(forward_batch, "uses_cp_shared_kv", False):
+        local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+        if local_out_cache_loc is None:
+            return None
+        local_current_rows = int(local_out_cache_loc.numel())
+        if local_current_rows > 0 and all(rows >= local_current_rows for rows in tensor_rows):
+            return local_current_rows
+
+    return None
 
 
 def build_batch_prefix_slot_span(
@@ -2295,6 +2310,108 @@ def build_batch_prefix_slot_span(
     return (start_slot, end_slot)
 
 
+def _merge_slot_spans(spans: list[tuple[int, int]]) -> list[tuple[int, int]]:
+    normalized = sorted(
+        (int(start), int(end)) for start, end in spans if int(end) > int(start)
+    )
+    if not normalized:
+        return []
+    merged: list[tuple[int, int]] = []
+    cur_start, cur_end = normalized[0]
+    for start, end in normalized[1:]:
+        if start <= cur_end:
+            cur_end = max(cur_end, end)
+        else:
+            merged.append((cur_start, cur_end))
+            cur_start, cur_end = start, end
+    merged.append((cur_start, cur_end))
+    return merged
+
+
+def build_batch_current_slot_spans(
+    *,
+    logical_pages: torch.Tensor,
+    prefix_lens_cpu,
+    extend_lens_cpu=None,
+    page_size: int,
+) -> list[tuple[int, int]]:
+    """Return flattened page-table slot spans that contain current extend pages."""
+
+    if prefix_lens_cpu is None:
+        raise ValueError("prefix_lens_cpu is required for current slot spans")
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    prefix_lens = [int(x) for x in prefix_lens_cpu]
+    if not prefix_lens:
+        return []
+
+    if logical_pages.dim() == 1:
+        if len(prefix_lens) != 1:
+            raise ValueError(
+                "1D logical_pages can only describe one request for current slot spans: "
+                f"batch_size={len(prefix_lens)} logical_pages_shape={tuple(logical_pages.shape)}"
+            )
+        pages_per_request = int(logical_pages.numel())
+    else:
+        if int(logical_pages.shape[0]) < len(prefix_lens):
+            raise ValueError(
+                "logical_pages has fewer rows than prefix_lens_cpu: "
+                f"rows={int(logical_pages.shape[0])} batch_size={len(prefix_lens)}"
+            )
+        pages_per_request = int(
+            logical_pages.reshape(logical_pages.shape[0], -1).shape[1]
+        )
+
+    if extend_lens_cpu is None:
+        extend_lens: list[int | None] = [None] * len(prefix_lens)
+    else:
+        extend_lens = [int(x) for x in extend_lens_cpu]
+        if len(extend_lens) != len(prefix_lens):
+            raise ValueError(
+                "extend_lens_cpu length must match prefix_lens_cpu: "
+                f"extend={len(extend_lens)} prefix={len(prefix_lens)}"
+            )
+
+    spans: list[tuple[int, int]] = []
+    for req_id, prefix_len in enumerate(prefix_lens):
+        if prefix_len < 0:
+            raise ValueError(
+                f"prefix_lens_cpu contains negative prefix length: req_id={req_id} "
+                f"prefix_len={prefix_len}"
+            )
+        if prefix_len % page_size != 0:
+            raise ValueError(
+                "CP shared KV current slot spans require page-aligned prefixes: "
+                f"req_id={req_id} prefix_len={prefix_len} page_size={page_size}"
+            )
+        prefix_pages = prefix_len // page_size
+        if prefix_pages > pages_per_request:
+            raise ValueError(
+                "prefix pages exceed per-request logical page-table width: "
+                f"req_id={req_id} prefix_pages={prefix_pages} "
+                f"pages_per_request={pages_per_request}"
+            )
+
+        extend_len = extend_lens[req_id]
+        if extend_len is None:
+            end_pages = pages_per_request
+        else:
+            if extend_len < 0:
+                raise ValueError(
+                    f"extend_lens_cpu contains negative extend length: req_id={req_id} "
+                    f"extend_len={extend_len}"
+                )
+            end_pages = (prefix_len + extend_len + page_size - 1) // page_size
+            end_pages = min(end_pages, pages_per_request)
+        if end_pages <= prefix_pages:
+            continue
+        req_start = req_id * pages_per_request
+        spans.append((req_start + prefix_pages, req_start + end_pages))
+
+    return _merge_slot_spans(spans)
+
+
 def current_loc_remap_fast_path_args(
     forward_batch,
 ) -> tuple[int | None, int | None]:
@@ -3285,6 +3402,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
     page_size: int,
     prefix_pages: int,
     prefix_slot_span: tuple[int, int] | None = None,
+    current_slot_spans: list[tuple[int, int]] | None = None,
     layer_id: int | None = None,
     nvtx_source: str = "mla.partial_current_sync",
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -3379,6 +3497,35 @@ def materialize_prefix_and_reuse_current_kv_page_slots(
         page_size=page_size,
         mask_non_current_in_current_pages=True,
     )
+    if layout.cp_size > 1:
+        if current_slot_spans is None:
+            if prefix_slot_span is not None:
+                raise ValueError(
+                    "CP shared KV batched current compose requires explicit "
+                    "current_slot_spans to avoid reducing prefix slots twice."
+                )
+            current_slot_spans = (
+                [(int(prefix_pages), total_slots)]
+                if int(prefix_pages) < total_slots
+                else []
+            )
+        for current_start_slot, current_end_slot in _merge_slot_spans(
+            current_slot_spans
+        ):
+            current_rows = slot_range_to_token_slice(
+                page_size,
+                current_start_slot,
+                current_end_slot,
+            )
+            _all_reduce_materialized_buffer_range(
+                mixed_kv_cache,
+                layout.cp_size,
+                current_rows.start,
+                current_rows.stop,
+                nvtx_source=f"{nvtx_source}.current",
+                nvtx_layer_id=layer_id,
+                nvtx_cp_rank=layout.cp_rank,
+            )
     return mixed_kv_cache, mixed_locs
 
 
@@ -3394,6 +3541,7 @@ def materialize_prefix_and_reuse_current_index_page_slots(
     index_head_dim: int,
     prefix_pages: int,
     prefix_slot_span: tuple[int, int] | None = None,
+    current_slot_spans: list[tuple[int, int]] | None = None,
     layer_id: int | None = None,
     nvtx_source: str = "index.partial_current_sync",
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -3462,6 +3610,34 @@ def materialize_prefix_and_reuse_current_index_page_slots(
         page_size=page_size,
         index_head_dim=index_head_dim,
     )
+    if layout.cp_size > 1:
+        if current_slot_spans is None:
+            if prefix_slot_span is not None:
+                raise ValueError(
+                    "CP shared KV batched index current compose requires explicit "
+                    "current_slot_spans to avoid reducing prefix slots twice."
+                )
+            current_slot_spans = (
+                [(int(prefix_pages), total_slots)]
+                if int(prefix_pages) < total_slots
+                else []
+            )
+        for current_start_slot, current_end_slot in _merge_slot_spans(
+            current_slot_spans
+        ):
+            current_pages = slot_range_to_page_slice(
+                current_start_slot,
+                current_end_slot,
+            )
+            _all_reduce_materialized_buffer_range(
+                dense_page_buffer,
+                layout.cp_size,
+                current_pages.start,
+                current_pages.stop,
+                nvtx_source=f"{nvtx_source}.current",
+                nvtx_layer_id=layer_id,
+                nvtx_cp_rank=layout.cp_rank,
+            )
     return dense_page_buffer, slot_remap.dense_pages
 
 
diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
index a8d6a86a0..3964a19d1 100644
--- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
@@ -16,6 +16,7 @@ from sglang.jit_kernel.fused_store_index_cache import (
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.nsa import index_buf_accessor
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+    build_batch_current_slot_spans,
     build_batch_prefix_slot_span,
     cp_shared_kv_debug_enabled,
     cp_shared_kv_debug_log,
@@ -524,12 +525,12 @@ class Indexer(MultiPlatformOp):
                 prefix_lens_valid = all(
                     int(prefix_len) >= 0 and int(prefix_len) % page_size == 0
                     for prefix_len in prefix_lens_cpu
-                ) and any(int(prefix_len) > 0 for prefix_len in prefix_lens_cpu)
+                )
             if not prefix_lens_valid:
                 raise RuntimeError(
                     "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
-                    "CP shared KV index partial-current compose requires "
-                    "positive page-aligned prefix pages. "
+                    "CP shared KV index current compose requires page-aligned "
+                    "prefix metadata. "
                     f"cp_rank={layout.cp_rank} layer_id={layer_id} "
                     f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
                     f"logical_page_table_shape={tuple(logical_page_table.shape)} "
@@ -591,6 +592,12 @@ class Indexer(MultiPlatformOp):
                         f"out_cache_loc_shape={tuple(current_locs.shape)}"
                     )
             prefix_slot_span = None
+            current_slot_spans = build_batch_current_slot_spans(
+                logical_pages=logical_page_table,
+                prefix_lens_cpu=prefix_lens_cpu,
+                extend_lens_cpu=extend_lens_cpu,
+                page_size=page_size,
+            )
             if len(prefix_lens_cpu) == 1:
                 prefix_pages = int(prefix_lens_cpu[0]) // page_size
             else:
@@ -655,6 +662,7 @@ class Indexer(MultiPlatformOp):
                     index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
                     prefix_pages=prefix_pages,
                     prefix_slot_span=prefix_slot_span,
+                    current_slot_spans=current_slot_spans,
                     layer_id=layer_id,
                 )
             )
@@ -1456,7 +1464,7 @@ class Indexer(MultiPlatformOp):
                     index_buffer=index_buffer,
                     block_tables=block_tables,
                     batch_indices=torch.tensor(
-                        batch_idx_list, dtype=torch.int64, device=descriptor_device
+                        batch_idx_list, dtype=torch.int32, device=descriptor_device
                     ),
                     kv_lens=torch.tensor(
                         kv_lens_list, dtype=torch.int32, device=descriptor_device
@@ -1861,7 +1869,7 @@ class Indexer(MultiPlatformOp):
         shared_block_tables = None
         current_index_kv_for_topk = current_index_kv
         current_only_batch = is_current_only_extend_batch(forward_batch)
-        if current_index_kv is not None and not current_only_batch:
+        if current_index_kv is not None:
             current_index_kv_for_topk = None
             shared_block_tables = metadata.get_page_table_64()
             shared_index_buffer, shared_block_tables = (
@@ -1973,7 +1981,7 @@ class Indexer(MultiPlatformOp):
         shared_block_tables = None
         current_index_kv_for_topk = current_index_kv
         current_only_batch = is_current_only_extend_batch(forward_batch)
-        if current_index_kv is not None and not current_only_batch:
+        if current_index_kv is not None:
             current_index_kv_for_topk = None
             shared_block_tables = metadata.get_page_table_64()
             shared_index_buffer, shared_block_tables = (
diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py
index 0308afeef..c073c7a0c 100644
--- a/python/sglang/srt/layers/attention/nsa/utils.py
+++ b/python/sglang/srt/layers/attention/nsa/utils.py
@@ -716,9 +716,9 @@ def build_batch_page_aligned_in_seq_split_plan(
         flat_segment_request_ids=flat_segment_request_ids,
         flat_segment_offsets=flat_segment_offsets,
         compute_padding_enabled=any(
-            compute_tokens != valid_tokens
-            for compute_tokens, valid_tokens in zip(
-                request_compute_padded_tokens, request_valid_padded_tokens
+            list(compute_split) != list(valid_split)
+            for compute_split, valid_split in zip(
+                request_compute_split_lists, request_valid_split_lists
             )
         ),
         request_valid_split_lists=request_valid_split_lists,
diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py
index a2c4104ea..db3910fdf 100644
--- a/python/sglang/srt/layers/attention/nsa_backend.py
+++ b/python/sglang/srt/layers/attention/nsa_backend.py
@@ -15,6 +15,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
     CpSharedKVMlaPrefetcher,
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
+    build_batch_current_slot_spans,
     build_batch_prefix_slot_span,
     build_current_loc_remap,
     cp_shared_kv_debug_enabled,
@@ -1971,6 +1972,16 @@ class NativeSparseAttnBackend(
                         layout=forward_batch.cp_shared_kv_layout,
                         page_size=page_size,
                     )
+                    current_slot_spans = build_batch_current_slot_spans(
+                        logical_pages=metadata.real_page_table,
+                        prefix_lens_cpu=getattr(
+                            forward_batch, "extend_prefix_lens_cpu", None
+                        ),
+                        extend_lens_cpu=getattr(
+                            forward_batch, "extend_seq_lens_cpu", None
+                        ),
+                        page_size=page_size,
+                    )
                     kv_cache, page_table_1 = (
                         materialize_prefix_and_reuse_current_kv_page_slots(
                             kv_cache=kv_cache,
@@ -1981,6 +1992,7 @@ class NativeSparseAttnBackend(
                             layout=forward_batch.cp_shared_kv_layout,
                             page_size=page_size,
                             prefix_pages=0,
+                            current_slot_spans=current_slot_spans,
                             layer_id=layer.layer_id,
                             nvtx_source="mla.current_only_page_slots",
                         )
@@ -2078,6 +2090,12 @@ class NativeSparseAttnBackend(
                                 f"page_size={page_size}"
                             )
                         prefix_slot_span = None
+                        current_slot_spans = build_batch_current_slot_spans(
+                            logical_pages=metadata.real_page_table,
+                            prefix_lens_cpu=prefix_lens_cpu,
+                            extend_lens_cpu=extend_lens_cpu,
+                            page_size=page_size,
+                        )
                         if len(prefix_lens_cpu) == 1:
                             prefix_pages = int(prefix_lens_cpu[0]) // page_size
                         else:
@@ -2105,6 +2123,7 @@ class NativeSparseAttnBackend(
                                 page_size=page_size,
                                 prefix_pages=prefix_pages,
                                 prefix_slot_span=prefix_slot_span,
+                                current_slot_spans=current_slot_spans,
                                 layer_id=layer.layer_id,
                             )
                         )
diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py
index acb316552..8ecc11b81 100644
--- a/test/registered/unit/layers/test_nsa_cp_utils.py
+++ b/test/registered/unit/layers/test_nsa_cp_utils.py
@@ -566,6 +566,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertEqual(plan.request_compute_seq_q_next, [0])
 
     def test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias(self):
+        import torch
+
         from sglang.srt.layers.attention.nsa.nsa_indexer import (
             _select_batch_topk_query_lengths,
         )
@@ -585,6 +587,29 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         )
         self.assertEqual(valid_local_rows, 4995)
         self.assertEqual(compute_local_rows, 5056)
+        self.assertTrue(plan.compute_padding_enabled)
+
+        local_compute = split_tensor_by_cp_batch_plan(
+            torch.arange(40387, dtype=torch.int64),
+            plan,
+            mode="1d",
+        )
+        self.assertEqual(local_compute.numel(), compute_local_rows)
+        self.assertEqual(local_compute[:2560].tolist(), list(range(2560)))
+        self.assertEqual(local_compute[2560:4995].tolist(), list(range(37952, 40387)))
+        self.assertEqual(local_compute[4995:].tolist(), [0] * 61)
+
+        local_valid = select_cp_local_valid_rows_for_cache_write(
+            SimpleNamespace(
+                nsa_cp_metadata=NSAContextParallelMetadata(
+                    batch_size=1,
+                    batch_plan=plan,
+                )
+            ),
+            local_compute,
+        )
+        self.assertEqual(local_valid.numel(), valid_local_rows)
+        self.assertEqual(local_valid.tolist(), local_compute[:valid_local_rows].tolist())
 
         selected = _select_batch_topk_query_lengths(
             cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan),
@@ -1405,7 +1430,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         import torch
 
         plan = build_batch_page_aligned_in_seq_split_plan(
-            extend_lens=[7],
+            extend_lens=[8],
             prefix_lens=[0],
             page_size=4,
             cp_size=2,
@@ -1413,15 +1438,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         )
         self.assertFalse(plan.compute_padding_enabled)
         forward_batch = SimpleNamespace(
-            extend_num_tokens=8,
+            extend_num_tokens=9,
             nsa_cp_metadata=NSAContextParallelMetadata(
                 batch_size=1,
                 batch_plan=plan,
             ),
         )
-        tensor = torch.arange(8 * 2, dtype=torch.float32).view(8, 2)
+        tensor = torch.arange(9 * 2, dtype=torch.float32).view(9, 2)
         expected = split_tensor_by_cp_batch_plan(
-            tensor[:7],
+            tensor[:8],
             plan,
             mode="data",
         )
@@ -1561,7 +1586,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         import torch
 
         plan = build_batch_page_aligned_in_seq_split_plan(
-            extend_lens=[7],
+            extend_lens=[8],
             prefix_lens=[0],
             page_size=4,
             cp_size=2,
@@ -1569,15 +1594,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         )
         self.assertFalse(plan.compute_padding_enabled)
         forward_batch = SimpleNamespace(
-            extend_num_tokens=8,
+            extend_num_tokens=9,
             nsa_cp_metadata=NSAContextParallelMetadata(
                 batch_size=1,
                 batch_plan=plan,
             ),
         )
-        positions = torch.arange(8, dtype=torch.int32)
+        positions = torch.arange(9, dtype=torch.int32)
         expected = split_tensor_by_cp_batch_plan(
-            positions[:7],
+            positions[:8],
             plan,
             mode="position",
         )
@@ -1821,23 +1846,17 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         page_size = 4
         plan = build_batch_page_aligned_in_seq_split_plan(
-            extend_lens=[7],
+            extend_lens=[8],
             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,
+            extend_num_tokens=9,
             cp_shared_kv_layout=CpSharedKVLayout(
                 page_size=page_size,
                 cp_size=2,
@@ -1850,7 +1869,12 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                 page_size=page_size,
                 extend_prefix_len=0,
             ),
-            out_cache_loc=torch.cat((valid_locs, static_padding_locs)),
+            out_cache_loc=torch.cat(
+                (
+                    torch.arange(1 * page_size, 3 * page_size, dtype=torch.int64),
+                    static_padding_locs,
+                )
+            ),
         )
 
         with patch(
@@ -1861,7 +1885,12 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         self.assertEqual(
             local_locs.tolist(),
-            [2 * page_size, 2 * page_size + 1, 2 * page_size + 2],
+            [
+                2 * page_size,
+                2 * page_size + 1,
+                2 * page_size + 2,
+                2 * page_size + 3,
+            ],
         )
 
     def test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding(
@@ -2502,13 +2531,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                 return logical_pages
 
             def get_page_table_1(self):
-                return torch.empty((2, 1000), dtype=torch.int32)
-
-            def get_page_table_1(self):
-                return torch.empty((2, 1000), dtype=torch.int32)
-
-            def get_page_table_1(self):
-                return torch.empty((2, 1000), dtype=torch.int32)
+                return torch.empty((1, 1000), dtype=torch.int32)
 
         def fake_materialize(forward_batch, layer_id, logical_page_table):
             materialize_calls.append((layer_id, logical_page_table))
@@ -2943,7 +2966,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
         )
 
-    def test_indexer_in_seq_cp_pair_batch_reuses_current_index_without_materialize(self):
+    def test_indexer_in_seq_cp_pair_batch_composes_current_only_index_reuse(self):
         import torch
 
         from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
@@ -2951,6 +2974,10 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         indexer = object.__new__(Indexer)
         indexer.index_topk = 2
         current_index_kv = (torch.arange(7), torch.arange(7))
+        logical_pages = torch.tensor([[1, 2], [3, 4]], dtype=torch.int32)
+        materialized_index = torch.tensor([17], dtype=torch.int32)
+        dense_pages = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
+        materialize_calls = []
         topk_calls = []
 
         class Mode:
@@ -2959,10 +2986,25 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         class Metadata:
             def get_page_table_64(self):
-                raise AssertionError("current-only index reuse should not read page table")
+                return logical_pages
 
-        def fake_materialize(*args, **kwargs):
-            raise AssertionError("current-only index reuse should not materialize")
+            def get_page_table_1(self):
+                return torch.empty((2, 512), dtype=torch.int32)
+
+        def fake_materialize(
+            forward_batch,
+            layer_id,
+            logical_page_table,
+            current_index_kv=None,
+        ):
+            materialize_calls.append(
+                {
+                    "layer_id": layer_id,
+                    "logical_page_table": logical_page_table,
+                    "current_index_kv": current_index_kv,
+                }
+            )
+            return materialized_index, dense_pages
 
         def fake_get_topk(
             forward_batch,
@@ -3022,15 +3064,16 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             current_index_kv=current_index_kv,
         )
 
+        self.assertEqual(len(materialize_calls), 1)
+        self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages)
+        self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv)
         self.assertEqual(len(topk_calls), 1)
+        self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls))
         self.assertTrue(
-            all(call["current_index_kv"] is current_index_kv for call in topk_calls)
+            all(call["shared_index_buffer"] is materialized_index for call in topk_calls)
         )
         self.assertTrue(
-            all(call["shared_index_buffer"] is None for call in topk_calls)
-        )
-        self.assertTrue(
-            all(call["shared_block_tables"] is None for call in topk_calls)
+            all(call["shared_block_tables"] is dense_pages for call in topk_calls)
         )
         self.assertEqual(topk_calls[0]["batch_idx"], 0)
         self.assertEqual(topk_calls[0]["actual_seq_q"], 7)
@@ -3045,6 +3088,69 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
         )
 
+    def test_indexer_shared_index_materialize_accepts_current_only_compose(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        page_buffer = torch.zeros((8, 32), dtype=torch.uint8)
+        logical_pages = torch.tensor([[1, 2]], dtype=torch.int64)
+        current_index_kv = (
+            torch.zeros((2, 4), dtype=torch.uint8),
+            torch.zeros((2, 1), dtype=torch.float32),
+        )
+        compose_calls = []
+
+        class Pool:
+            page_size = 4
+            index_head_dim = 4
+
+            def get_index_k_with_scale_buffer(self, layer_id):
+                return page_buffer
+
+        class Mode:
+            def is_extend_without_speculative(self):
+                return True
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
+            token_to_kv_pool=Pool(),
+            cp_shared_kv_index_prefetcher=None,
+            forward_mode=Mode(),
+            extend_prefix_lens_cpu=[0],
+            extend_seq_lens_cpu=[5],
+            seq_lens_cpu=torch.tensor([5], dtype=torch.int64),
+            cp_local_out_cache_loc=torch.tensor([4, 5], dtype=torch.int64),
+        )
+
+        def fake_compose(**kwargs):
+            compose_calls.append(kwargs)
+            return torch.empty((3, 32), dtype=torch.uint8), torch.tensor([[1, 2]])
+
+        with patch.object(
+            nsa_indexer,
+            "materialize_prefix_and_reuse_current_index_page_slots",
+            side_effect=fake_compose,
+        ):
+            materialized, dense_pages = Indexer._maybe_materialize_shared_index_buffer(
+                indexer,
+                forward_batch,
+                layer_id=7,
+                logical_page_table=logical_pages,
+                current_index_kv=current_index_kv,
+            )
+
+        self.assertEqual(len(compose_calls), 1)
+        self.assertEqual(compose_calls[0]["prefix_pages"], 0)
+        self.assertEqual(compose_calls[0]["current_slot_spans"], [(0, 2)])
+        self.assertIs(compose_calls[0]["current_index_k"], current_index_kv[0])
+        self.assertIs(compose_calls[0]["current_index_scale"], current_index_kv[1])
+        self.assertEqual(list(materialized.shape), [3, 32])
+        self.assertEqual(dense_pages.tolist(), [[1, 2]])
+
     def test_indexer_ragged_cp_index_current_batch_does_not_materialize(self):
         import contextlib
         import torch
@@ -3215,7 +3321,9 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         call = prepare_calls[0]
         self.assertIs(call["index_buffer"], shared_index_buffer)
         self.assertIs(call["block_tables"], shared_block_tables)
+        self.assertEqual(call["batch_indices"].dtype, torch.int32)
         self.assertEqual(call["batch_indices"].tolist(), [0, 0, 1, 1])
+        self.assertEqual(call["kv_lens"].dtype, torch.int32)
         self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4])
         self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1])
         self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3])
@@ -3447,23 +3555,40 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertIs(draft_input.hidden_states, draft_output_hidden)
         self.assertFalse(draft_input.cp_local_hidden_states)
 
-    def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
+    def test_indexer_in_seq_cp_pair_composes_current_only_index_reuse(self):
         import torch
 
         from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
 
         indexer = object.__new__(Indexer)
         current_index_kv = (torch.tensor([1]), torch.tensor([2]))
+        logical_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        materialized_index = torch.tensor([7], dtype=torch.int32)
+        dense_pages = torch.tensor([[3, 4]], dtype=torch.int32)
         materialize_calls = []
         topk_calls = []
 
         class Metadata:
             def get_page_table_64(self):
-                raise AssertionError("current index reuse should not read page table")
+                return logical_pages
 
-        def fake_materialize(forward_batch, layer_id, logical_page_table):
-            materialize_calls.append((layer_id, logical_page_table))
-            raise AssertionError("current index reuse should not materialize")
+            def get_page_table_1(self):
+                return torch.empty((1, 512), dtype=torch.int32)
+
+        def fake_materialize(
+            forward_batch,
+            layer_id,
+            logical_page_table,
+            current_index_kv=None,
+        ):
+            materialize_calls.append(
+                {
+                    "layer_id": layer_id,
+                    "logical_page_table": logical_page_table,
+                    "current_index_kv": current_index_kv,
+                }
+            )
+            return materialized_index, dense_pages
 
         def fake_get_topk(
             forward_batch,
@@ -3527,12 +3652,14 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             current_index_kv=current_index_kv,
         )
 
-        self.assertEqual(materialize_calls, [])
+        self.assertEqual(len(materialize_calls), 1)
+        self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages)
+        self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv)
         self.assertEqual(len(topk_calls), 2)
-        self.assertIs(topk_calls[0]["current_index_kv"], current_index_kv)
-        self.assertIs(topk_calls[1]["current_index_kv"], current_index_kv)
-        self.assertIsNone(topk_calls[0]["shared_index_buffer"])
-        self.assertIsNone(topk_calls[1]["shared_block_tables"])
+        self.assertIsNone(topk_calls[0]["current_index_kv"])
+        self.assertIsNone(topk_calls[1]["current_index_kv"])
+        self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index)
+        self.assertIs(topk_calls[1]["shared_block_tables"], dense_pages)
         self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3])
         self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2])
         self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
index f536ccdd1..58e6d46c0 100644
--- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
@@ -1030,6 +1030,43 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         self.assertIn("[CP_SHARED_KV_FALLBACK][current_reuse]", joined)
         self.assertIn("prefix_extend_seq_len_mismatch_req_0", joined)
 
+    def test_cp_shared_current_only_reuse_accepts_owner_local_current_rows(
+        self,
+    ):
+        from sglang.srt.environ import envs
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            forward_mode=_FakeExtendForwardMode(),
+            batch_size=1,
+            extend_prefix_lens_cpu=[0],
+            extend_seq_lens_cpu=[40387],
+            seq_lens_cpu=torch.tensor([40387], dtype=torch.int32),
+            out_cache_loc=torch.arange(40392, dtype=torch.int64),
+            cp_local_out_cache_loc=torch.arange(5056, dtype=torch.int64),
+        )
+        local_k = torch.empty((5056, 1), dtype=torch.float32)
+        local_rope = torch.empty((5056, 1), dtype=torch.float32)
+
+        with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True):
+            self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
+            self.assertEqual(
+                runtime.current_extend_kv_rows_for_reuse(
+                    forward_batch,
+                    local_k,
+                    local_rope,
+                ),
+                5056,
+            )
+            self.assertIsNone(
+                runtime.current_extend_kv_rows_for_reuse(
+                    forward_batch,
+                    local_k[:5055],
+                    local_rope,
+                )
+            )
+
     def test_tai_index_mqa_prepare_fast_path_miss_logs_warning(self):
         from sglang.srt.environ import envs
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
@@ -1371,6 +1408,36 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             (0, 0),
         )
 
+    def test_batch_current_slot_spans_follow_prefix_and_extend_pages(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+
+        logical_pages = torch.tensor(
+            [
+                [1, 2, 5, 0],
+                [9, 11, 12, 13],
+            ],
+            dtype=torch.int64,
+        )
+
+        self.assertEqual(
+            runtime.build_batch_current_slot_spans(
+                logical_pages=logical_pages,
+                prefix_lens_cpu=[8, 4],
+                extend_lens_cpu=[2, 7],
+                page_size=4,
+            ),
+            [(2, 3), (5, 7)],
+        )
+        self.assertEqual(
+            runtime.build_batch_current_slot_spans(
+                logical_pages=logical_pages,
+                prefix_lens_cpu=[0, 0],
+                extend_lens_cpu=[8, 4],
+                page_size=4,
+            ),
+            [(0, 2), (4, 5)],
+        )
+
     def test_materialize_batch_prefix_span_and_reuse_current_kv_page_slots(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
         from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -1406,6 +1473,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             prefix_lens_cpu=[8, 4],
             page_size=page_size,
         )
+        current_slot_spans = runtime.build_batch_current_slot_spans(
+            logical_pages=remap_logical_pages,
+            prefix_lens_cpu=[8, 4],
+            extend_lens_cpu=[2, 2],
+            page_size=page_size,
+        )
 
         with patch.object(
             runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce
@@ -1421,6 +1494,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
                     page_size=page_size,
                     prefix_pages=0,
                     prefix_slot_span=prefix_slot_span,
+                    current_slot_spans=current_slot_spans,
                 )
             )
 
@@ -1455,6 +1529,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             prefix_lens_cpu=[8, 4],
             page_size=page_size,
         )
+        current_slot_spans = runtime.build_batch_current_slot_spans(
+            logical_pages=logical_pages,
+            prefix_lens_cpu=[8, 4],
+            extend_lens_cpu=[2, 2],
+            page_size=page_size,
+        )
         current_k = torch.tensor(
             [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]],
             dtype=torch.uint8,
@@ -1476,6 +1556,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
                     index_head_dim=index_head_dim,
                     prefix_pages=0,
                     prefix_slot_span=prefix_slot_span,
+                    current_slot_spans=current_slot_spans,
                     layer_id=2,
                 )
             )
@@ -1588,7 +1669,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         self.assertEqual(dense_locs.tolist(), [4, 8, 12])
         self.assertEqual(list(dense_kv[4:16].flatten().tolist()), list(range(12)))
 
-    def test_materialize_prefix_current_token_kv_uses_ipc_without_all_reduce(self):
+    def test_materialize_prefix_current_token_kv_uses_ipc_and_reduces_current_slots(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
         from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 
@@ -1615,6 +1696,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             ).view(8, 1, 1)
             return True
 
+        range_calls = []
+
+        def record_range_reduce(buffer, cp_size, start_row, end_row, **kwargs):
+            range_calls.append((start_row, end_row, kwargs.get("nvtx_source")))
+            return buffer
+
         with patch.object(
             runtime,
             "_try_tai_ipc_materialize_token_kv_page_slots_into",
@@ -1622,7 +1709,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         ), patch.object(
             runtime,
             "_all_reduce_materialized_buffer_range",
-            side_effect=AssertionError("IPC path must not range all-reduce"),
+            side_effect=record_range_reduce,
         ):
             mixed_kv, mixed_locs = (
                 runtime.materialize_prefix_and_reuse_current_kv_page_slots(
@@ -1641,6 +1728,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         self.assertEqual(mixed_kv[4].item(), 10)
         self.assertEqual(mixed_kv[8].item(), 14)
         self.assertTrue(torch.equal(mixed_kv[12:14], current_kv))
+        self.assertEqual(range_calls, [(12, 16, "mla.partial_current_sync.current")])
 
     def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
@@ -1676,6 +1764,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         )
         prefetcher.handles[1] = handle
         prefetcher.pending_attention_handle = handle
+        range_calls = []
+
+        def record_range_reduce(buffer, cp_size, start_row, end_row, **kwargs):
+            range_calls.append((start_row, end_row, kwargs.get("nvtx_source")))
+            return buffer
 
         with patch.object(
             prefetch.torch.cuda, "current_stream", return_value=current_stream
@@ -1683,6 +1776,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             prefetch,
             "materialize_local_token_kv_page_slots_into",
             side_effect=AssertionError("suffix materialize must not run"),
+        ), patch.object(
+            prefetch,
+            "_all_reduce_materialized_buffer_range",
+            side_effect=record_range_reduce,
         ):
             mixed_kv, mixed_locs = prefetcher.consume_prefix_with_current(
                 layer_id=1,
@@ -1701,6 +1798,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         expected_kv[12:14] = current_kv
         self.assertTrue(torch.equal(mixed_kv, expected_kv))
         self.assertEqual(mixed_locs.tolist(), [[4, 12], [13, 7], [-1, -1]])
+        self.assertEqual(range_calls, [(12, 16, "mla.prefetch_current")])
 
     def test_mla_prefetch_attention_window_defers_pending_event_wait(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch
@@ -1995,7 +2093,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
         )
 
-    def test_materialize_prefix_current_index_uses_ipc_without_all_reduce(self):
+    def test_materialize_prefix_current_index_uses_ipc_and_reduces_current_slots(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
         from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 
@@ -2025,6 +2123,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
             )
             return True
 
+        range_calls = []
+
+        def record_range_reduce(buffer, cp_size, start_row, end_row, **kwargs):
+            range_calls.append((start_row, end_row, kwargs.get("nvtx_source")))
+            return buffer
+
         with patch.object(
             runtime,
             "_try_tai_ipc_materialize_paged_buffer_page_slots_into",
@@ -2032,7 +2136,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         ), patch.object(
             runtime,
             "_all_reduce_materialized_buffer_range",
-            side_effect=AssertionError("IPC path must not range all-reduce"),
+            side_effect=record_range_reduce,
         ):
             dense_page_buffer, dense_pages = (
                 runtime.materialize_prefix_and_reuse_current_index_page_slots(
@@ -2053,6 +2157,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         self.assertEqual(dense_page_buffer[1].tolist(), list(range(page_bytes)))
         self.assertTrue(torch.equal(dense_page_buffer[2, 0:4], current_k[0]))
         self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1]))
+        self.assertEqual(range_calls, [(2, 3, "index.partial_current_sync.current")])
 
     def test_index_prefetch_partial_current_compose_fills_current_page_slots(self):
         from sglang.srt.environ import envs
@@ -2138,7 +2243,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         method_source = source[start:end]
 
         self.assertIn("should_reuse_current_extend_kv(forward_batch)", method_source)
-        self.assertNotIn("is_current_only_extend_batch(forward_batch)", method_source)
+        self.assertNotIn("current_only_compact_unsupported", method_source)
 
     def test_index_current_reuse_prepare_accepts_padded_out_cache_loc(self):
         from pathlib import Path