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

Keep CP shared-KV prefetch warnings actionable

Browse Source 
Expected no-prefetch paths were polluting production logs: no cache prefix, tiny/first-layer windows, and FP8 RAGGED top-k were being reported as fallback warnings. The prefetch contract now treats zero-prefix and first-layer misses as normal skips, while preserving warnings for non-zero misaligned prefixes and real consume misses after the first layer. The same change keeps RAGGED cache-hit prefetch eligible and records the CE/IPM prefetch contract in the plan doc.

Constraint: FP8 sparse prefill uses RAGGED top-k, but CP shared-KV prefix materialization is still page-slot based

Constraint: Layer 0 has no previous attention-window hook that can have prefetched the layer

Rejected: Warn whenever a prefetcher is absent | no-cache and too-short requests are expected synchronous paths and make logs unusable

Confidence: high

Scope-risk: moderate

Directive: Keep CP_SHARED_KV_FALLBACK warnings for unexpected contract failures only; use debug logs for expected skip paths

Tested: Local py_compile for cp_shared_kv_prefetch.py, nsa_indexer.py, nsa_backend.py

Tested: Remote cjy-glm5-new targeted regression: 3 passed, 21 warnings

Tested: Remote cjy-glm5-new full test_cp_shared_kv_runtime.py: 156 passed, 21 warnings, 2 subtests passed

Not-tested: New ETE run after prefill restart to confirm log volume reduction in production traffic
(cherry picked from commit e08e321e5929fdbb30102ec0b19c6ff0ecac7e7e)
This commit is contained in:
laoyao0822
2026-06-13 00:59:33 +08:00
parent da8bd7f813
commit 4d5c7f32d6
6 changed files with 1509 additions and 582 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md
+118 -20
View File
@@ -430,15 +430,37 @@ L2->L1 load finished on owner rank
### P0. 文档与现状锁定
**文件：**
**状态（2026-06-12）：已完成。**
- 新增：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md`
完成内容：
**验收：**
1. L1 prefetch create 路径不再以 `get_attention_cp_group().pynccl_comm` 作为启用条件。
prefetch 的新合同是 IPC/CE materialize，不再依赖 CP collective communicator。
2. `cp_shared_kv_prefetch.py` 与 `nsa_backend.py` 中 prefetch hot path 已清理旧
`launch_pending_reduce` / `prefix_reduce` / `reduce_enqueue` / `collective_disallowed`
命名，统一改成 `finalize_pending_materialize` / `prefix_materialize` /
`materialize_enqueue` / `materialize_event_missing`。
3. prefetch 文件中不再 import/use `_all_reduce_materialized_buffer_*` 或
`get_attention_cp_group`。如果 IPC/CE materialize 不可用，CP>1 下应显式
fail-fast，而不是 silent fallback 到 all-reduce。
4. 测试已覆盖：`maybe_create()` 不依赖 pynccl，attention-window 不再触发
prefix reduce 旧路径，tiny extend skip 测试也不再 mock pynccl。
- 文档明确 L2->L1 `+2 layer`、L1 prefetch `+1 layer`。
- 文档明确当前 prefetcher 的 bs=1 假设和风险。
- 文档明确 0SM 与现有 SM IPC kernel 的区别。
验收证据：
```text
远端 cjy-glm5-new:
PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
150 passed, 21 warnings, 2 subtests passed
```
静态检查：
```text
rg "pynccl|launch_pending_reduce|prefix_reduce|collective_disallowed|reduce_enqueue|get_attention_cp_group|_all_reduce" \
cp_shared_kv_prefetch.py nsa_backend.py
# no matches
```
### P1. 补 bs>1 prefetch plan 单测
@@ -508,17 +530,17 @@ PYTHONPATH=python python -m pytest -q \
- `current_slot_spans`
- `prefix_page_count`
- `current_page_count`
3. `start_next_layer_prefix()` 只 materialize/reduce `prefix_slot_spans`，
3. `start_next_layer_prefix()` 只 materialize `prefix_slot_spans`，
不再把 batch flattened page table 当成 `[0:prefix_pages)`。
4. `consume_prefix_with_current()` 只 reduce `current_slot_spans`，避免把
row gap / 其他 request prefix 一起 reduce。
4. `consume_prefix_with_current()` 只处理 `current_slot_spans`，避免把
row gap / 其他 request prefix 一起纳入 suffix compose。
5. `consume()` 的 legacy full-materialize suffix 路径也改为使用
`current_slot_spans`，避免 bs>1 bounding suffix。
当前限制：
- MLA prefix spans 仍走现有 materialize + async all-reduce baseline；还没有
接入 spans-list TAI IPC 或 0SM CE。
- MLA prefix spans 已不再走 async all-reduce fallback；当前要求 IPC/CE materialize
生成 completion event，CP>1 下不可用则 fail-fast。
- index prefetcher 仍未改造，继续由 P3 处理。
**文件：**
@@ -531,13 +553,13 @@ PYTHONPATH=python python -m pytest -q \
1. `CpSharedKVMlaPrefetcher.maybe_create()` 不再因为 `batch_size != 1` 直接 skip。
2. create 阶段构造 `prefix_slot_spans/current_slot_spans`。
3. `start_next_layer_prefix()` 只 materialize prefix spans。
4. `consume_prefix_with_current()` 只 fill/reduce current spans。
4. `consume_prefix_with_current()` 只处理 current spans。
5. `consume()` 如果仍存在 legacy suffix path，必须使用 `current_slot_spans`
或 fail-fast，不能回到错误 bounding suffix。
**第一版允许：**
- prefix 使用现有 local materialize + all-reduce 或 SM IPC baseline。
- prefix 使用现有 SM IPC 或 CE materialize baseline。
**第一版不允许：**
@@ -567,8 +589,8 @@ RED 证据：旧代码在 `batch_size=2` 时以
- `current_slot_spans`
- `prefix_page_count`
- `current_page_count`
3. `start_next_layer_prefix()` 只 materialize/reduce index prefix spans。
4. `consume_prefix_with_current()` 只 fill/reduce index current spans。
3. `start_next_layer_prefix()` 只 materialize index prefix spans。
4. `consume_prefix_with_current()` 只处理 index current spans。
5. `consume()` 的 legacy suffix 路径也改为 `current_slot_spans`，
不再使用 batch bounding suffix。
@@ -581,8 +603,8 @@ test_cp_shared_kv_runtime.py
当前限制：
- index prefetch 仍使用现有 materialize + async all-reduce baseline；
spans-list TAI IPC / 0SM CE 留给后续 P5/P6。
- index prefetch 已不再走 async all-reduce fallback；当前要求 IPC/CE materialize
生成 completion event，CP>1 下不可用则 fail-fast。
- active index layer / index skip 的 runtime hook 当前沿用已有
`nsa_backend.py` 调用点；本阶段没有修改 skip 参数语义。
@@ -637,11 +659,11 @@ test_cp_shared_kv_runtime.py
- `MLATokenToKVPool.get_key_buffer_for_prefetch(layer_id, stream)`
- `MLATokenToKVPool.get_index_k_with_scale_buffer_for_prefetch(layer_id, stream)`
- 发现并修正一个 ordering 疏漏：P1-P3 后 prefix materialize 仍在 current stream 上执行，但 `start_next_layer_prefix()` 把 L2->L1 readiness wait 绑定到了 prefetch stream。这样只能保护后续 reduce，不能保护实际读取 L1 raw pages 的 materialize。
- 修正策略：MLA/index 的 `start_next_layer_prefix()` 先取得 `current_stream`，把它传给 prefetch-safe getter，使 L2->L1 ready event 挂到实际读取 KV/index page 的 stream；随后仍保持 `prefetch_stream.wait_stream(current_stream)`，reduce 在 prefetch stream 上异步提交。
- 修正策略：MLA/index 的 `start_next_layer_prefix()` 先取得 `current_stream`，把它传给 prefetch-safe getter，使 L2->L1 ready event 挂到实际读取 KV/index page 的 stream；随后仍保持 stream ordering，由 IPC/CE materialize event 表示 prefix ready。
- 没有新增 collective；没有把 forward stream 改成 CPU 同步等待。
- 单测锁住：
- `test_mla_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream`
- `test_index_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream`
- `test_mla_prefetch_waits_l2_l1_on_materialize_stream`
- `test_index_prefetch_waits_l2_l1_on_materialize_stream`
### P5. TAI SM IPC spans baseline
@@ -995,3 +1017,79 @@ SYMM CE prefix + SM current: gpu_p50=0.521ms cpu_p50=0.546ms
2. 当前 PyTorch symmetric-memory prototype 没有优于 CUDA IPC CE；推测原因包括 symmetric memory backend/registration overhead、统一 window size 带来的空间浪费、以及仍使用 cudaMemcpyBatchAsync per-span submit。
3. 短期最有价值路径是继续用 IPC CE 做 prefix prefetch baseline，同时保留 symmetric window helper 作为后续 NCCL CE/NVSHMEM transport 的地址合同实验入口。
4. 生产接入前必须继续做 span coalescing、跨 layer descriptor 复用，以及更大 bs/不同 cached/extend sweep；不能只凭 symmetric window 概念直接替换现有路径。
---
## 10. FP8/RAGGED 与 index-skip prefetch 修正（2026-06-12）
### 10.1 远端日志现象
最新 prefill 日志中 `index_prefetch` fallback 被刷屏：
```text
[CP_SHARED_KV_FALLBACK][index_prefetch]
current_missing_prefetcher: 53136
not_paged_topk: 1296
```
排查结论：
1. `current_missing_prefetcher` 大多来自无 cache / tiny extend / prefetcher 未创建的正常同步路径，不是运行时错误；consume 侧不应每层每 rank warning。
2. `not_paged_topk` 是错误判定：FP8 sparse prefill 走 `TopkTransformMethod.RAGGED` 是预期行为，但 index/MLA prefix materialize 仍是 page-slot 语义，不能因为 top-k transform 不是 PAGED 就关闭 prefetch。
### 10.2 当前修正
1. `CpSharedKVIndexPrefetcher.maybe_create()` 对 RAGGED top-k 不再 fallback，而是继续创建 index prefetcher。
2. `CpSharedKVMlaPrefetcher.maybe_create()` 对 RAGGED top-k 也继续创建 MLA prefetcher。
3. `nsa_indexer._maybe_materialize_shared_index_buffer()` 不再对缺失 prefetcher 的正常同步路径打 warning；仅在 prefetcher 已存在但 consume miss 时保留 fallback warning。
4. `nsa_backend.py` 的 FP8/RAGGED cache-hit compose 现在先尝试 `mla_prefetcher.consume_prefix_with_current()`，只有 prefetcher 存在但 miss 时才 warning 并回到同步 `materialize_prefix_and_reuse_current_kv_page_slots()`。
5. attention-window 启动 index prefetch 时检查 `token_to_kv_pool.index_layer_to_slot`，index-skip/shared 层没有 index-cache slot 时只跳过 index prefetch，不影响 MLA prefetch。
### 10.3 验证
远端容器 `cjy-glm5-new` 已验证：
```text
python -m py_compile \
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py \
python/sglang/srt/layers/attention/nsa_backend.py
PYTHONPATH=python python -m pytest -q \
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_ragged_mla_uses_page_slot_current_compose_for_cache_hit \
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_mla_prefetch_create_allows_ragged_topk_transform \
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_index_prefetch_create_allows_ragged_topk_transform \
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVTaiMaterializeIntegration::test_attention_window_prefetch_skips_inactive_index_cache_layer
结果：4 passed
```
### 10.4 剩余风险
1. 旧进程日志不会反映本次降噪和 RAGGED prefetch 修正；需要重启 prefill 后再观察 warning 频率。
2. RAGGED 当前层 materialize 的历史代码位置仍在 flashmla_sparse 分支内部；后续如果严格追求“当前层 materialize 完成后再启动下一层 prefetch”，需要把 RAGGED current compose 前移到 attention-window hook 之前。
3. 本次只修正 create/consume 合同和单测；还需要 ETE 确认 FP8 cache-hit 下 index/MLA prefetch 命中率和精度不回退。
### 10.5 预期 skip 不再 warning（2026-06-12）
远端新进程仍出现大量 fallback warning 后，继续按日志聚合定位：
```text
[CP_SHARED_KV_FALLBACK][index_prefetch] reason=current_consume_miss 主要集中在 layer=0
[CP_SHARED_KV_FALLBACK][mla_prefetch] reason=ragged_current_consume_miss 主要集中在 layer=0
[CP_SHARED_KV_FALLBACK][index_prefetch|mla_prefetch] reason=prefix_not_page_aligned prefix_lens=[0]
```
结论：这些都不是异常 fallback：
1. `layer=0` 没有上一层 attention-window 能提前发起本层 prefetch，consume miss 是预期同步路径。
2. `prefix_lens=[0]` 表示 request 本身没有命中 cache，prefix prefetch 不应创建。
3. 真实需要 warning 的场景应收窄为：非零 prefix 不 page-aligned、prefetcher 已创建且 `layer > start_layer` 仍 consume miss、metadata/slot 合同异常等。
修正合同：
1. MLA/index prefetch create 遇到全零 prefix 时只走 debug skip，不再打 `[CP_SHARED_KV_FALLBACK]` warning。
2. index partial-current prefetcher 在 `layer == token_to_kv_pool.start_layer` miss 时不 warning。
3. FP8/RAGGED MLA prefetcher 在 `layer == start_layer` miss 时不 warning。
这类 skip 仍保留 debug 级别路径，方便需要时通过 debug 环境变量观察，但默认生产日志不应被无 cache/tiny/首层预期路径污染。
验证：远端 `cjy-glm5-new` 已通过 targeted regression（3 passed）以及完整 `test_cp_shared_kv_runtime.py`（156 passed, 21 warnings, 2 subtests passed）。
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
+230 -361
View File
@@ -8,8 +8,6 @@ from typing import Any, Optional
import torch
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
_all_reduce_materialized_buffer_async,
_all_reduce_materialized_buffer_range,
build_batch_current_slot_spans,
build_batch_prefix_slot_spans,
cp_shared_kv_debug_enabled,
@@ -34,13 +32,14 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
slot_range_to_token_slice,
_raise_tai_ipc_materialize_required,
_should_fail_fast_tai_ipc_materialize,
_try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into,
_try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into,
_try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into,
_try_tai_ipc_materialize_current_token_kv_page_slot_spans_into,
_try_tai_ipc_materialize_paged_buffer_page_slot_spans_into,
_try_tai_ipc_materialize_token_kv_page_slot_spans_into,
)
from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
from sglang.srt.layers.dp_attention import get_attention_cp_group
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
logger = logging.getLogger(__name__)
@@ -215,32 +214,6 @@ def _materialize_local_paged_buffer_page_slot_spans_into(
)
def _all_reduce_materialized_buffer_ranges_async(
*,
dense_kv_cache: torch.Tensor,
row_slices: list[slice],
cp_size: int,
stream: torch.cuda.Stream,
nvtx_source: str,
nvtx_layer_id: int,
nvtx_cp_rank: int,
) -> torch.cuda.Event | None:
event: torch.cuda.Event | None = None
for rows in row_slices:
event = _all_reduce_materialized_buffer_async(
dense_kv_cache[rows],
cp_size=cp_size,
stream=stream,
nvtx_source=nvtx_source,
nvtx_layer_id=nvtx_layer_id,
nvtx_cp_rank=nvtx_cp_rank,
nvtx_rows=(rows.start, rows.stop),
)
if event is None:
return None
return event
def _record_event_on_stream(stream: torch.cuda.Stream) -> torch.cuda.Event:
event = torch.cuda.Event()
event.record(stream)
@@ -277,7 +250,7 @@ class _PrefetchCpuTiming:
start_max_ms: float = 0.0
get_total_ms: float = 0.0
materialize_total_ms: float = 0.0
reduce_enqueue_total_ms: float = 0.0
materialize_enqueue_total_ms: float = 0.0
consume_count: int = 0
consume_total_ms: float = 0.0
consume_wait_total_ms: float = 0.0
@@ -290,7 +263,7 @@ class _PrefetchCpuTiming:
total_ms: float,
get_ms: float,
materialize_ms: float,
reduce_enqueue_ms: float,
materialize_enqueue_ms: float,
) -> None:
if total_ms < 0.0:
return
@@ -299,7 +272,7 @@ class _PrefetchCpuTiming:
self.start_max_ms = max(self.start_max_ms, total_ms)
self.get_total_ms += max(get_ms, 0.0)
self.materialize_total_ms += max(materialize_ms, 0.0)
self.reduce_enqueue_total_ms += max(reduce_enqueue_ms, 0.0)
self.materialize_enqueue_total_ms += max(materialize_enqueue_ms, 0.0)
def record_consume(
self,
@@ -328,7 +301,7 @@ def _log_prefetch_cpu_start(
total_ms: float,
get_ms: float,
materialize_ms: float,
reduce_enqueue_ms: float,
materialize_enqueue_ms: float,
prefix_pages: int,
total_slots: int,
dense_units: int,
@@ -338,14 +311,14 @@ def _log_prefetch_cpu_start(
if cp_shared_kv_mla_prefetch_should_log_layer(layer_id):
log_fn(
"cpu_timing path=%s stage=start layer=%s total_ms=%.3f "
"get_ms=%.3f materialize_ms=%.3f reduce_enqueue_ms=%.3f "
"get_ms=%.3f materialize_ms=%.3f materialize_enqueue_ms=%.3f "
"prefix_pages=%s total_slots=%s dense_units=%s",
path,
layer_id,
total_ms,
get_ms,
materialize_ms,
reduce_enqueue_ms,
materialize_enqueue_ms,
prefix_pages,
total_slots,
dense_units,
@@ -354,14 +327,14 @@ def _log_prefetch_cpu_start(
starts = timing.start_count
log_fn(
"cpu_summary path=%s starts=%s avg_start_ms=%.3f max_start_ms=%.3f "
"avg_get_ms=%.3f avg_materialize_ms=%.3f avg_reduce_enqueue_ms=%.3f",
"avg_get_ms=%.3f avg_materialize_ms=%.3f avg_materialize_enqueue_ms=%.3f",
path,
starts,
timing.start_total_ms / starts,
timing.start_max_ms,
timing.get_total_ms / starts,
timing.materialize_total_ms / starts,
timing.reduce_enqueue_total_ms / starts,
timing.materialize_enqueue_total_ms / starts,
)
@@ -453,6 +426,7 @@ class CpSharedKVMlaPrefetcher:
slot_sorted_logical_pages_by_row: torch.Tensor | None = None,
slot_sorted_dense_pages_by_row: torch.Tensor | None = None,
dense_num_pages: int,
slot_remap: Any | None = None,
owned_prefix_pages: int = -1,
owned_total_pages: int = -1,
stream: Optional[torch.cuda.Stream] = None,
@@ -465,6 +439,7 @@ class CpSharedKVMlaPrefetcher:
self.slot_sorted_logical_pages_by_row = slot_sorted_logical_pages_by_row
self.slot_sorted_dense_pages_by_row = slot_sorted_dense_pages_by_row
self.dense_num_pages = dense_num_pages
self.slot_remap = slot_remap
self.owned_prefix_pages = owned_prefix_pages
self.owned_total_pages = owned_total_pages
self.total_slots = int(slot_logical_pages.numel())
@@ -519,8 +494,11 @@ class CpSharedKVMlaPrefetcher:
_prefetch_log("create_skip reason=not_in_seq_split")
return None
if not topk_transform_is_paged:
_prefetch_log("create_skip reason=not_paged_topk")
return None
# FP8 sparse prefill uses the RAGGED top-k transform by design.
# MLA prefix prefetch materializes page-slot KV rows before the
# final top-k offset semantics are consumed, so the RAGGED transform
# is still a valid prefetch target.
_prefetch_log("create_continue reason=ragged_topk")
batch_size = int(getattr(forward_batch, "batch_size", 0))
if batch_size <= 0:
_prefetch_log(
@@ -553,11 +531,17 @@ class CpSharedKVMlaPrefetcher:
for prefix_len in prefix_lens
if prefix_len < 0 or prefix_len % page_size != 0
]
if bad_prefix_lens or not any(prefix_len > 0 for prefix_len in prefix_lens):
if not any(prefix_len > 0 for prefix_len in prefix_lens):
_prefetch_log(
"create_skip reason=no_cache_prefix prefix_lens=%s page_size=%s",
prefix_lens,
page_size,
)
return None
if bad_prefix_lens:
_mla_prefetch_fallback_log(
"prefix_not_page_aligned",
"prefix length is zero or not page-aligned. "
"prefix_lens=%s page_size=%s",
"non-zero prefix length is not page-aligned. prefix_lens=%s page_size=%s",
prefix_lens,
page_size,
)
@@ -654,15 +638,6 @@ class CpSharedKVMlaPrefetcher:
)
return None
cp_group = get_attention_cp_group()
if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:
_prefetch_log(
"create_skip reason=missing_pynccl cp_rank=%s cp_size=%s",
layout.cp_rank,
layout.cp_size,
)
return None
prefetch_stream = stream if stream is not None else torch.cuda.Stream()
create_cpu = _cpu_timing_start()
get_cpu = _cpu_timing_start()
@@ -738,6 +713,7 @@ class CpSharedKVMlaPrefetcher:
slot_sorted_logical_pages_by_row=remap.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=remap.slot_sorted_dense_pages_by_row,
dense_num_pages=remap.dense_num_pages,
slot_remap=remap,
owned_prefix_pages=owned_prefix_pages,
owned_total_pages=owned_total_pages,
stream=prefetch_stream,
@@ -779,12 +755,12 @@ class CpSharedKVMlaPrefetcher:
self._log_layer(layer_id, "consume_miss layer=%s", layer_id)
return None
if handle.event is None:
self.launch_pending_reduce()
self.finalize_pending_materialize()
handle = self.handles.get(layer_id)
if handle is None or handle.event is None:
self._log_layer(
layer_id,
"consume_miss reason=prefix_reduce_not_ready layer=%s",
"consume_miss reason=prefix_materialize_not_ready layer=%s",
layer_id,
)
return None
@@ -830,7 +806,11 @@ class CpSharedKVMlaPrefetcher:
spans=suffix_spans,
)
)
if not materialized_suffix_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_kv_cache):
if (
not materialized_suffix_by_ipc
and suffix_spans
and self.layout.cp_size > 1
):
_raise_tai_ipc_materialize_required(
"mla_prefetch_suffix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
@@ -847,19 +827,6 @@ class CpSharedKVMlaPrefetcher:
page_size=self.page_size,
spans=suffix_spans,
)
if not materialized_suffix_by_ipc:
for suffix_rows in _slot_spans_to_token_slices(
self.page_size, suffix_spans
):
_all_reduce_materialized_buffer_range(
dense_kv_cache,
self.layout.cp_size,
suffix_rows.start,
suffix_rows.stop,
nvtx_source="mla.consume_suffix",
nvtx_layer_id=layer_id,
nvtx_cp_rank=self.layout.cp_rank,
)
self._log_layer(
layer_id,
"consume_suffix_done layer=%s suffix_spans=%s",
@@ -929,7 +896,7 @@ class CpSharedKVMlaPrefetcher:
"""Consume the prefetched prefix and append current-layer KV rows.
This is the partial-current-reuse variant of :meth:`consume`: prefix
pages are already materialized/reduced by the prefetch stream, while
pages are already materialized by the prefetch stream, while
current/suffix pages are not copied from the shared pool. Current locs
in ``logical_locs`` are remapped to the appended current KV rows.
"""
@@ -954,12 +921,12 @@ class CpSharedKVMlaPrefetcher:
self._log_layer(layer_id, "consume_prefix_current_miss layer=%s", layer_id)
return None
if handle.event is None:
self.launch_pending_reduce()
self.finalize_pending_materialize()
handle = self.handles.get(layer_id)
if handle is None or handle.event is None:
self._log_layer(
layer_id,
"consume_prefix_current_miss reason=prefix_reduce_not_ready layer=%s",
"consume_prefix_current_miss reason=prefix_materialize_not_ready layer=%s",
layer_id,
)
return None
@@ -1022,7 +989,7 @@ class CpSharedKVMlaPrefetcher:
if (
not current_materialized_by_ipc
and self.current_slot_spans
and _should_fail_fast_tai_ipc_materialize(mixed_kv_cache)
and self.layout.cp_size > 1
):
_raise_tai_ipc_materialize_required(
"mla_prefetch_current_ipc_unavailable",
@@ -1031,19 +998,6 @@ class CpSharedKVMlaPrefetcher:
spans=self.current_slot_spans,
dense_shape=tuple(mixed_kv_cache.shape),
)
if not current_materialized_by_ipc:
for current_rows in _slot_spans_to_token_slices(
self.page_size, self.current_slot_spans
):
_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(
@@ -1106,25 +1060,7 @@ class CpSharedKVMlaPrefetcher:
start_cpu = _cpu_timing_start()
current_stream = torch.cuda.current_stream()
get_cpu = _cpu_timing_start()
try:
kv_cache = _prefetch_pool_get_key_buffer(
token_to_kv_pool=token_to_kv_pool,
layer_id=next_layer_id,
stream=current_stream,
path="mla",
)
get_ms = _cpu_timing_ms(get_cpu)
except Exception:
logger.exception(
"Failed to get next-layer KV cache for CP shared KV MLA prefetch."
)
self.disabled = True
self._log_next_layer(
next_layer_id,
"start_disable reason=get_kv_failed next_layer=%s",
next_layer_id,
)
return
self.stream.wait_stream(current_stream)
try:
prefix_row_spans = _slot_spans_to_token_slices(
@@ -1132,64 +1068,81 @@ class CpSharedKVMlaPrefetcher:
)
prefix_rows = prefix_row_spans[0] if prefix_row_spans else slice(0, 0)
materialize_cpu = _cpu_timing_start()
dense_kv_cache = kv_cache.new_zeros(
(self.dense_num_pages * self.page_size, *kv_cache.shape[1:])
)
self._log_next_layer(
next_layer_id,
"start_prefix_begin next_layer=%s prefix_slot_spans=%s "
"dense_rows=%s",
next_layer_id,
self.prefix_slot_spans,
int(dense_kv_cache.shape[0]),
)
materialized_by_ipc = _try_tai_ipc_materialize_token_kv_page_slot_spans_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
page_size=self.page_size,
spans=self.prefix_slot_spans,
)
if not materialized_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_kv_cache):
_raise_tai_ipc_materialize_required(
"mla_prefetch_prefix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
spans=self.prefix_slot_spans,
dense_shape=tuple(dense_kv_cache.shape),
with torch.cuda.stream(self.stream):
kv_cache = _prefetch_pool_get_key_buffer(
token_to_kv_pool=token_to_kv_pool,
layer_id=next_layer_id,
stream=self.stream,
path="mla",
)
if not materialized_by_ipc:
_materialize_local_token_kv_page_slot_spans_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
page_size=self.page_size,
spans=self.prefix_slot_spans,
get_ms = _cpu_timing_ms(get_cpu)
dense_kv_cache = kv_cache.new_zeros(
(self.dense_num_pages * self.page_size, *kv_cache.shape[1:])
)
materialize_ms = _cpu_timing_ms(materialize_cpu)
reduce_cpu = _cpu_timing_start()
if materialized_by_ipc:
event = _record_event_on_stream(current_stream)
else:
self.stream.wait_stream(current_stream)
with torch.cuda.stream(self.stream):
event = _all_reduce_materialized_buffer_ranges_async(
self._log_next_layer(
next_layer_id,
"start_prefix_begin next_layer=%s prefix_slot_spans=%s "
"dense_rows=%s stream=prefetch",
next_layer_id,
self.prefix_slot_spans,
int(dense_kv_cache.shape[0]),
)
materialized_by_ce = (
_try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
row_slices=prefix_row_spans,
cp_size=self.layout.cp_size,
stream=self.stream,
nvtx_source="mla.prefetch_prefix",
nvtx_layer_id=next_layer_id,
nvtx_cp_rank=self.layout.cp_rank,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
page_size=self.page_size,
spans=self.prefix_slot_spans,
slot_remap=self.slot_remap,
)
reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu)
)
event = _record_event_on_stream(self.stream) if materialized_by_ce else None
materialized_by_ipc = materialized_by_ce
if not materialized_by_ce:
current_stream.wait_stream(self.stream)
materialized_by_ipc = (
_try_tai_ipc_materialize_token_kv_page_slot_spans_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
page_size=self.page_size,
spans=self.prefix_slot_spans,
slot_remap=self.slot_remap,
)
)
if (
not materialized_by_ipc
and self.prefix_slot_spans
and self.layout.cp_size > 1
):
_raise_tai_ipc_materialize_required(
"mla_prefetch_prefix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
spans=self.prefix_slot_spans,
dense_shape=tuple(dense_kv_cache.shape),
)
if not materialized_by_ipc:
_materialize_local_token_kv_page_slot_spans_into(
kv_cache=kv_cache,
dense_kv_cache=dense_kv_cache,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
page_size=self.page_size,
spans=self.prefix_slot_spans,
)
event = _record_event_on_stream(current_stream)
materialize_ms = _cpu_timing_ms(materialize_cpu)
materialize_enqueue_ms = 0.0
if event is None:
self.disabled = True
logger.warning(
"[CP_SHARED_KV_FALLBACK][mla_prefetch] "
"reason=async_reduce_unavailable layer_id=%s cp_rank=%s "
"reason=materialize_event_unavailable layer_id=%s cp_rank=%s "
"cp_size=%s prefix_pages=%s total_slots=%s",
next_layer_id,
self.layout.cp_rank,
@@ -1199,7 +1152,7 @@ class CpSharedKVMlaPrefetcher:
)
self._log_next_layer(
next_layer_id,
"start_disable reason=async_reduce_unavailable next_layer=%s",
"start_disable reason=materialize_event_unavailable next_layer=%s",
next_layer_id,
)
return
@@ -1208,7 +1161,7 @@ class CpSharedKVMlaPrefetcher:
total_ms=total_ms,
get_ms=get_ms,
materialize_ms=materialize_ms,
reduce_enqueue_ms=reduce_enqueue_ms,
materialize_enqueue_ms=materialize_enqueue_ms,
)
_log_prefetch_cpu_start(
log_fn=self._log,
@@ -1219,7 +1172,7 @@ class CpSharedKVMlaPrefetcher:
total_ms=total_ms,
get_ms=get_ms,
materialize_ms=materialize_ms,
reduce_enqueue_ms=reduce_enqueue_ms,
materialize_enqueue_ms=materialize_enqueue_ms,
prefix_pages=self.prefix_pages,
total_slots=self.total_slots,
dense_units=int(dense_kv_cache.shape[0]),
@@ -1246,74 +1199,44 @@ class CpSharedKVMlaPrefetcher:
self._log_next_layer(
next_layer_id,
"start next_layer=%s prefix_pages=%s prefix_rows=%s dense_rows=%s "
"reduce_enqueued=True",
"materialize_enqueued=True",
next_layer_id,
self.prefix_pages,
prefix_rows.stop - prefix_rows.start,
int(dense_kv_cache.shape[0]),
)
def launch_pending_reduce(self) -> None:
def finalize_pending_materialize(self) -> None:
handle = self.pending_attention_handle
if handle is None:
return
if handle.event is not None:
self._log_next_layer(
handle.layer_id,
"start_prefix_reduce_skip reason=already_enqueued next_layer=%s",
"start_prefix_materialize_skip reason=already_enqueued next_layer=%s",
handle.layer_id,
)
return
prefix_row_spans = list(handle.prefix_row_spans or (handle.prefix_rows,))
try:
if _should_fail_fast_tai_ipc_materialize(handle.dense_kv_cache):
_raise_tai_ipc_materialize_required(
"mla_prefetch_deferred_prefix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
dense_shape=tuple(handle.dense_kv_cache.shape),
layer_id=handle.layer_id,
)
current_stream = torch.cuda.current_stream()
self.stream.wait_stream(current_stream)
with torch.cuda.stream(self.stream):
event = _all_reduce_materialized_buffer_ranges_async(
dense_kv_cache=handle.dense_kv_cache,
row_slices=prefix_row_spans,
cp_size=self.layout.cp_size,
stream=self.stream,
nvtx_source="mla.prefetch_prefix",
nvtx_layer_id=handle.layer_id,
nvtx_cp_rank=self.layout.cp_rank,
)
if event is None:
self.disabled = True
self.handles.pop(handle.layer_id, None)
if self.pending_attention_handle is handle:
self.pending_attention_handle = None
self._log_next_layer(
handle.layer_id,
"start_disable reason=async_reduce_unavailable next_layer=%s",
handle.layer_id,
)
return
handle.event = event
self._log_next_layer(
handle.layer_id,
"start_prefix_reduce_enqueued next_layer=%s row_spans=%s",
handle.layer_id,
[(rows.start, rows.stop) for rows in prefix_row_spans],
)
_raise_tai_ipc_materialize_required(
"mla_prefetch_deferred_prefix_materialize_missing",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
dense_shape=tuple(handle.dense_kv_cache.shape),
layer_id=handle.layer_id,
)
except Exception:
logger.exception("Failed to launch CP shared KV MLA prefix prefetch reduce.")
logger.exception(
"CP shared KV MLA prefix prefetch has no IPC/CE completion event."
)
self.disabled = True
self.handles.pop(handle.layer_id, None)
if self.pending_attention_handle is handle:
self.pending_attention_handle = None
self._log_next_layer(
handle.layer_id,
"start_disable reason=reduce_launch_exception next_layer=%s",
"start_disable reason=materialize_event_missing next_layer=%s",
handle.layer_id,
)
return
@@ -1366,6 +1289,7 @@ class CpSharedKVIndexPrefetcher:
slot_sorted_logical_pages_by_row: torch.Tensor | None = None,
slot_sorted_dense_pages_by_row: torch.Tensor | None = None,
dense_num_pages: int,
slot_remap: Any | None = None,
owned_prefix_pages: int = -1,
owned_total_pages: int = -1,
stream: Optional[torch.cuda.Stream] = None,
@@ -1377,6 +1301,7 @@ class CpSharedKVIndexPrefetcher:
self.slot_sorted_logical_pages_by_row = slot_sorted_logical_pages_by_row
self.slot_sorted_dense_pages_by_row = slot_sorted_dense_pages_by_row
self.dense_num_pages = dense_num_pages
self.slot_remap = slot_remap
self.owned_prefix_pages = owned_prefix_pages
self.owned_total_pages = owned_total_pages
self.total_slots = int(slot_logical_pages.numel())
@@ -1448,11 +1373,11 @@ class CpSharedKVIndexPrefetcher:
)
return None
if not topk_transform_is_paged:
_index_prefetch_fallback_log(
"not_paged_topk",
"topk transform is not PAGED.",
)
return None
# FP8 prefill uses the RAGGED top-k transform by design. Index
# prefix materialization is still page-slot based and remains a
# valid prefetch target; the transform method only affects the
# final top-k offset semantics after MQA logits are computed.
_prefetch_log("index_create_continue reason=ragged_topk")
batch_size = int(getattr(forward_batch, "batch_size", 0))
if batch_size <= 0:
_index_prefetch_fallback_log(
@@ -1501,10 +1426,17 @@ class CpSharedKVIndexPrefetcher:
for prefix_len in prefix_lens
if prefix_len < 0 or prefix_len % page_size != 0
]
if bad_prefix_lens or not any(prefix_len > 0 for prefix_len in prefix_lens):
if not any(prefix_len > 0 for prefix_len in prefix_lens):
_prefetch_log(
"index_create_skip reason=no_cache_prefix prefix_lens=%s page_size=%s",
prefix_lens,
page_size,
)
return None
if bad_prefix_lens:
_index_prefetch_fallback_log(
"prefix_not_page_aligned",
"prefix length is zero or not page-aligned. prefix_lens=%s page_size=%s",
"non-zero prefix length is not page-aligned. prefix_lens=%s page_size=%s",
prefix_lens,
page_size,
)
@@ -1605,16 +1537,6 @@ class CpSharedKVIndexPrefetcher:
)
return None
cp_group = get_attention_cp_group()
if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1:
_index_prefetch_fallback_log(
"missing_pynccl",
"pynccl communicator is unavailable. cp_rank=%s cp_size=%s",
layout.cp_rank,
layout.cp_size,
)
return None
prefetch_stream = stream if stream is not None else torch.cuda.Stream()
create_cpu = _cpu_timing_start()
get_cpu = _cpu_timing_start()
@@ -1690,6 +1612,7 @@ class CpSharedKVIndexPrefetcher:
slot_sorted_logical_pages_by_row=remap.slot_sorted_logical_pages_by_row,
slot_sorted_dense_pages_by_row=remap.slot_sorted_dense_pages_by_row,
dense_num_pages=remap.dense_num_pages,
slot_remap=remap,
owned_prefix_pages=owned_prefix_pages,
owned_total_pages=owned_total_pages,
stream=prefetch_stream,
@@ -1730,12 +1653,12 @@ class CpSharedKVIndexPrefetcher:
self._log_layer(layer_id, "index_consume_miss layer=%s", layer_id)
return None
if handle.event is None:
self.launch_pending_reduce()
self.finalize_pending_materialize()
handle = self.handles.get(layer_id)
if handle is None or handle.event is None:
self._log_layer(
layer_id,
"index_consume_miss reason=prefix_reduce_not_ready layer=%s",
"index_consume_miss reason=prefix_materialize_not_ready layer=%s",
layer_id,
)
return None
@@ -1780,7 +1703,11 @@ class CpSharedKVIndexPrefetcher:
spans=suffix_spans,
)
)
if not materialized_suffix_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_page_buffer):
if (
not materialized_suffix_by_ipc
and suffix_spans
and self.layout.cp_size > 1
):
_raise_tai_ipc_materialize_required(
"index_prefetch_suffix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
@@ -1796,17 +1723,6 @@ class CpSharedKVIndexPrefetcher:
layout=self.layout,
spans=suffix_spans,
)
if not materialized_suffix_by_ipc:
for suffix_rows in _slot_spans_to_page_slices(suffix_spans):
_all_reduce_materialized_buffer_range(
dense_page_buffer,
self.layout.cp_size,
suffix_rows.start,
suffix_rows.stop,
nvtx_source="index.consume_suffix",
nvtx_layer_id=layer_id,
nvtx_cp_rank=self.layout.cp_rank,
)
self._log_layer(
layer_id,
"index_consume_suffix_done layer=%s suffix_spans=%s",
@@ -1895,13 +1811,13 @@ class CpSharedKVIndexPrefetcher:
)
return None
if handle.event is None:
self.launch_pending_reduce()
self.finalize_pending_materialize()
handle = self.handles.get(layer_id)
if handle is None or handle.event is None:
self._log_layer(
layer_id,
"index_consume_prefix_current_miss "
"reason=prefix_reduce_not_ready layer=%s",
"reason=prefix_materialize_not_ready layer=%s",
layer_id,
)
return None
@@ -1955,7 +1871,7 @@ class CpSharedKVIndexPrefetcher:
if (
not current_materialized_by_ipc
and self.current_slot_spans
and _should_fail_fast_tai_ipc_materialize(dense_page_buffer)
and self.layout.cp_size > 1
):
_raise_tai_ipc_materialize_required(
"index_prefetch_current_ipc_unavailable",
@@ -1964,17 +1880,6 @@ class CpSharedKVIndexPrefetcher:
spans=self.current_slot_spans,
dense_shape=tuple(dense_page_buffer.shape),
)
if not current_materialized_by_ipc:
for current_pages in _slot_spans_to_page_slices(self.current_slot_spans):
_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(
@@ -2039,87 +1944,83 @@ class CpSharedKVIndexPrefetcher:
start_cpu = _cpu_timing_start()
current_stream = torch.cuda.current_stream()
get_cpu = _cpu_timing_start()
try:
page_buffer = _prefetch_pool_get_index_buffer(
token_to_kv_pool=token_to_kv_pool,
layer_id=next_layer_id,
stream=current_stream,
)
get_ms = _cpu_timing_ms(get_cpu)
except Exception:
logger.exception(
"Failed to get next-layer index buffer for CP shared KV index prefetch."
)
self.disabled = True
self._log_next_layer(
next_layer_id,
"index_start_disable reason=get_index_buffer_failed next_layer=%s",
next_layer_id,
)
return
self.stream.wait_stream(current_stream)
try:
prefix_row_spans = _slot_spans_to_page_slices(self.prefix_slot_spans)
prefix_rows = prefix_row_spans[0] if prefix_row_spans else slice(0, 0)
materialize_cpu = _cpu_timing_start()
dense_page_buffer = page_buffer.new_zeros(
(self.dense_num_pages, *page_buffer.shape[1:])
)
self._log_next_layer(
next_layer_id,
"index_start_prefix_begin next_layer=%s prefix_slot_spans=%s "
"dense_pages=%s",
next_layer_id,
self.prefix_slot_spans,
int(dense_page_buffer.shape[0]),
)
materialized_by_ipc = (
_try_tai_ipc_materialize_paged_buffer_page_slot_spans_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
spans=self.prefix_slot_spans,
with torch.cuda.stream(self.stream):
page_buffer = _prefetch_pool_get_index_buffer(
token_to_kv_pool=token_to_kv_pool,
layer_id=next_layer_id,
stream=self.stream,
)
)
if not materialized_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_page_buffer):
_raise_tai_ipc_materialize_required(
"index_prefetch_prefix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
spans=self.prefix_slot_spans,
dense_shape=tuple(dense_page_buffer.shape),
get_ms = _cpu_timing_ms(get_cpu)
dense_page_buffer = page_buffer.new_zeros(
(self.dense_num_pages, *page_buffer.shape[1:])
)
if not materialized_by_ipc:
_materialize_local_paged_buffer_page_slot_spans_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
spans=self.prefix_slot_spans,
self._log_next_layer(
next_layer_id,
"index_start_prefix_begin next_layer=%s prefix_slot_spans=%s "
"dense_pages=%s stream=prefetch",
next_layer_id,
self.prefix_slot_spans,
int(dense_page_buffer.shape[0]),
)
materialize_ms = _cpu_timing_ms(materialize_cpu)
reduce_cpu = _cpu_timing_start()
if materialized_by_ipc:
event = _record_event_on_stream(current_stream)
else:
self.stream.wait_stream(current_stream)
with torch.cuda.stream(self.stream):
event = _all_reduce_materialized_buffer_ranges_async(
dense_kv_cache=dense_page_buffer,
row_slices=prefix_row_spans,
cp_size=self.layout.cp_size,
stream=self.stream,
nvtx_source="index.prefetch_prefix",
nvtx_layer_id=next_layer_id,
nvtx_cp_rank=self.layout.cp_rank,
materialized_by_ce = (
_try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
spans=self.prefix_slot_spans,
slot_remap=self.slot_remap,
)
reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu)
)
event = _record_event_on_stream(self.stream) if materialized_by_ce else None
materialized_by_ipc = materialized_by_ce
if not materialized_by_ce:
current_stream.wait_stream(self.stream)
materialized_by_ipc = (
_try_tai_ipc_materialize_paged_buffer_page_slot_spans_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
spans=self.prefix_slot_spans,
slot_remap=self.slot_remap,
)
)
if (
not materialized_by_ipc
and self.prefix_slot_spans
and self.layout.cp_size > 1
):
_raise_tai_ipc_materialize_required(
"index_prefetch_prefix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
spans=self.prefix_slot_spans,
dense_shape=tuple(dense_page_buffer.shape),
)
if not materialized_by_ipc:
_materialize_local_paged_buffer_page_slot_spans_into(
page_buffer=page_buffer,
dense_page_buffer=dense_page_buffer,
slot_logical_pages=self.slot_logical_pages,
layout=self.layout,
spans=self.prefix_slot_spans,
)
event = _record_event_on_stream(current_stream)
materialize_ms = _cpu_timing_ms(materialize_cpu)
materialize_enqueue_ms = 0.0
if event is None:
self.disabled = True
_index_prefetch_fallback_log(
"async_reduce_unavailable",
"async reduce unavailable. layer_id=%s cp_rank=%s cp_size=%s "
"materialize_event_unavailable",
"materialize event unavailable. layer_id=%s cp_rank=%s cp_size=%s "
"prefix_pages=%s total_slots=%s",
next_layer_id,
self.layout.cp_rank,
@@ -2129,7 +2030,7 @@ class CpSharedKVIndexPrefetcher:
)
self._log_next_layer(
next_layer_id,
"index_start_disable reason=async_reduce_unavailable next_layer=%s",
"index_start_disable reason=materialize_event_unavailable next_layer=%s",
next_layer_id,
)
return
@@ -2138,7 +2039,7 @@ class CpSharedKVIndexPrefetcher:
total_ms=total_ms,
get_ms=get_ms,
materialize_ms=materialize_ms,
reduce_enqueue_ms=reduce_enqueue_ms,
materialize_enqueue_ms=materialize_enqueue_ms,
)
_log_prefetch_cpu_start(
log_fn=self._log,
@@ -2149,7 +2050,7 @@ class CpSharedKVIndexPrefetcher:
total_ms=total_ms,
get_ms=get_ms,
materialize_ms=materialize_ms,
reduce_enqueue_ms=reduce_enqueue_ms,
materialize_enqueue_ms=materialize_enqueue_ms,
prefix_pages=self.prefix_pages,
total_slots=self.total_slots,
dense_units=int(dense_page_buffer.shape[0]),
@@ -2176,67 +2077,35 @@ class CpSharedKVIndexPrefetcher:
self._log_next_layer(
next_layer_id,
"index_start next_layer=%s prefix_pages=%s dense_pages=%s "
"reduce_enqueued=True",
"materialize_enqueued=True",
next_layer_id,
self.prefix_pages,
int(dense_page_buffer.shape[0]),
)
def launch_pending_reduce(self) -> None:
def finalize_pending_materialize(self) -> None:
handle = self.pending_attention_handle
if handle is None:
return
if handle.event is not None:
self._log_next_layer(
handle.layer_id,
"index_start_prefix_reduce_skip reason=already_enqueued next_layer=%s",
"index_start_prefix_materialize_skip reason=already_enqueued next_layer=%s",
handle.layer_id,
)
return
prefix_row_spans = list(handle.prefix_row_spans or (handle.prefix_rows,))
try:
if _should_fail_fast_tai_ipc_materialize(handle.dense_page_buffer):
_raise_tai_ipc_materialize_required(
"index_prefetch_deferred_prefix_ipc_unavailable",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
dense_shape=tuple(handle.dense_page_buffer.shape),
layer_id=handle.layer_id,
)
current_stream = torch.cuda.current_stream()
self.stream.wait_stream(current_stream)
with torch.cuda.stream(self.stream):
event = _all_reduce_materialized_buffer_ranges_async(
dense_kv_cache=handle.dense_page_buffer,
row_slices=prefix_row_spans,
cp_size=self.layout.cp_size,
stream=self.stream,
nvtx_source="index.prefetch_prefix",
nvtx_layer_id=handle.layer_id,
nvtx_cp_rank=self.layout.cp_rank,
)
if event is None:
self.disabled = True
self.handles.pop(handle.layer_id, None)
if self.pending_attention_handle is handle:
self.pending_attention_handle = None
self._log_next_layer(
handle.layer_id,
"index_start_disable reason=async_reduce_unavailable next_layer=%s",
handle.layer_id,
)
return
handle.event = event
self._log_next_layer(
handle.layer_id,
"index_start_prefix_reduce_enqueued next_layer=%s row_spans=%s",
handle.layer_id,
[(rows.start, rows.stop) for rows in prefix_row_spans],
)
_raise_tai_ipc_materialize_required(
"index_prefetch_deferred_prefix_materialize_missing",
cp_rank=self.layout.cp_rank,
cp_size=self.layout.cp_size,
dense_shape=tuple(handle.dense_page_buffer.shape),
layer_id=handle.layer_id,
)
except Exception:
logger.exception(
"Failed to launch CP shared KV index prefix prefetch reduce."
"CP shared KV index prefix prefetch has no IPC/CE completion event."
)
self.disabled = True
self.handles.pop(handle.layer_id, None)
@@ -2244,7 +2113,7 @@ class CpSharedKVIndexPrefetcher:
self.pending_attention_handle = None
self._log_next_layer(
handle.layer_id,
"index_start_disable reason=reduce_launch_exception next_layer=%s",
"index_start_disable reason=materialize_event_missing next_layer=%s",
handle.layer_id,
)
return
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+247
View File
@@ -2679,6 +2679,253 @@ def _get_or_build_prefix_ipc_slot_descriptors(
return descriptors
def _build_prefix_ipc_slot_descriptors_uncached(
*,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
device: torch.device,
physical_page_capacity: int | None,
) -> _IpcPrefixSlotDescriptors:
flat_slot_logical_pages = _contiguous_for_tai(
slot_logical_pages.reshape(-1).to(device=device)
)
slot_indices = _slot_spans_to_cuda_slot_indices(
spans,
total_slots=int(flat_slot_logical_pages.numel()),
device=device,
)
if slot_indices.numel() == 0:
empty = torch.empty((0,), dtype=torch.long, device=device)
return _IpcPrefixSlotDescriptors(empty, empty, empty, empty)
slot_logical_pages_range = _contiguous_for_tai(
flat_slot_logical_pages.index_select(0, slot_indices)
)
owner_ranks, src_page_indices = build_cp_shared_kv_ipc_page_descriptors(
slot_logical_pages_range,
layout,
physical_page_capacity=physical_page_capacity,
)
return _IpcPrefixSlotDescriptors(
slot_indices=slot_indices.contiguous(),
owner_ranks=owner_ranks,
src_page_indices=src_page_indices,
dense_page_indices=(slot_indices + 1).to(torch.long).contiguous(),
)
def _get_prefix_ipc_slot_descriptors(
*,
slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap | None,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
device: torch.device,
physical_page_capacity: int | None,
cache_kind: str,
) -> _IpcPrefixSlotDescriptors:
if slot_remap is not None:
return _get_or_build_prefix_ipc_slot_descriptors(
slot_remap=slot_remap,
layout=layout,
spans=spans,
device=device,
physical_page_capacity=physical_page_capacity,
cache_kind=cache_kind,
)
return _build_prefix_ipc_slot_descriptors_uncached(
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=device,
physical_page_capacity=physical_page_capacity,
)
def _filter_ipc_prefix_descriptors_for_ce(
descriptors: _IpcPrefixSlotDescriptors,
) -> _IpcPrefixSlotDescriptors:
"""Drop invalid/sentinel slots before CE copy.
The SM IPC materialize kernel can zero-fill invalid slots. The CE path is
intentionally just cudaMemcpyBatchAsync over valid spans; callers allocate
the dense destination with zeros, so invalid slots should be skipped instead
of submitted to the CE primitive.
"""
if descriptors.owner_ranks.numel() == 0:
return descriptors
valid = (
(descriptors.owner_ranks >= 0)
& (descriptors.src_page_indices >= 0)
& (descriptors.dense_page_indices >= 0)
)
if bool(valid.all().item()):
return descriptors
return _IpcPrefixSlotDescriptors(
slot_indices=descriptors.slot_indices[valid].contiguous(),
owner_ranks=descriptors.owner_ranks[valid].contiguous(),
src_page_indices=descriptors.src_page_indices[valid].contiguous(),
dense_page_indices=descriptors.dense_page_indices[valid].contiguous(),
)
def _try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into(
*,
kv_cache: torch.Tensor,
dense_kv_cache: torch.Tensor,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
page_size: int,
spans: list[tuple[int, int]],
slot_remap: SharedTokenKVSlotRemap | None = None,
) -> bool:
if not spans:
return True
if not dense_kv_cache.is_cuda or not dense_kv_cache.is_contiguous():
_log_tai_ipc_materialize_fallback(
"ce_dense_token_tensor_unsupported",
"CP shared KV tai IPC CE token span materialize requires a contiguous "
"CUDA dense tensor; falling back to SM IPC or collective. "
"device=%s contiguous=%s dense_shape=%s spans=%s",
dense_kv_cache.device,
dense_kv_cache.is_contiguous(),
tuple(dense_kv_cache.shape),
spans,
limit=4,
)
return False
ipc_state = _get_or_open_tai_ipc_peer_ptrs(kv_cache, layout)
if ipc_state is None:
return False
kernels, peer_ptrs = ipc_state
ce_kernel = getattr(kernels, "materialize_cuda_ipc_peer_pages_slot_indices_ce", None)
if ce_kernel is None:
_log_tai_ipc_materialize_fallback(
"ce_kernel_missing",
"CP shared KV tai IPC CE token materialize is unavailable; "
"falling back to SM IPC or collective. Upgrade tai-kernel.",
limit=1,
)
return False
try:
descriptors = _get_prefix_ipc_slot_descriptors(
slot_remap=slot_remap,
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=torch.device("cpu"),
physical_page_capacity=kv_cache.shape[0] // page_size,
cache_kind="token-ce",
)
descriptors = _filter_ipc_prefix_descriptors_for_ce(descriptors)
if descriptors.owner_ranks.numel() == 0:
return True
ce_kernel(
peer_ptrs,
dense_kv_cache,
descriptors.owner_ranks,
descriptors.src_page_indices,
descriptors.dense_page_indices,
page_nbytes=_token_kv_page_nbytes(kv_cache, page_size),
)
return True
except Exception as exc:
_log_tai_ipc_materialize_fallback(
"token_span_ce_kernel_failed",
"CP shared KV tai IPC CE token span materialize failed; falling back "
"to SM IPC or collective materialize. cp_rank=%s cp_size=%s spans=%s "
"page_size=%s kv_shape=%s dense_shape=%s error=%s",
layout.cp_rank,
layout.cp_size,
spans,
page_size,
tuple(kv_cache.shape),
tuple(dense_kv_cache.shape),
exc,
)
return False
def _try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into(
*,
page_buffer: torch.Tensor,
dense_page_buffer: torch.Tensor,
slot_logical_pages: torch.Tensor,
layout: CpSharedKVLayout,
spans: list[tuple[int, int]],
slot_remap: SharedPagedBufferSlotRemap | None = None,
) -> bool:
if not spans:
return True
if not dense_page_buffer.is_cuda or not dense_page_buffer.is_contiguous():
_log_tai_ipc_materialize_fallback(
"ce_dense_paged_tensor_unsupported",
"CP shared KV tai IPC CE paged span materialize requires a contiguous "
"CUDA dense tensor; falling back to SM IPC or collective. "
"device=%s contiguous=%s dense_shape=%s spans=%s",
dense_page_buffer.device,
dense_page_buffer.is_contiguous(),
tuple(dense_page_buffer.shape),
spans,
limit=4,
)
return False
ipc_state = _get_or_open_tai_ipc_peer_ptrs(page_buffer, layout)
if ipc_state is None:
return False
kernels, peer_ptrs = ipc_state
ce_kernel = getattr(kernels, "materialize_cuda_ipc_peer_pages_slot_indices_ce", None)
if ce_kernel is None:
_log_tai_ipc_materialize_fallback(
"ce_kernel_missing",
"CP shared KV tai IPC CE paged materialize is unavailable; "
"falling back to SM IPC or collective. Upgrade tai-kernel.",
limit=1,
)
return False
try:
descriptors = _get_prefix_ipc_slot_descriptors(
slot_remap=slot_remap,
slot_logical_pages=slot_logical_pages,
layout=layout,
spans=spans,
device=torch.device("cpu"),
physical_page_capacity=page_buffer.shape[0],
cache_kind="paged-ce",
)
descriptors = _filter_ipc_prefix_descriptors_for_ce(descriptors)
if descriptors.owner_ranks.numel() == 0:
return True
ce_kernel(
peer_ptrs,
dense_page_buffer,
descriptors.owner_ranks,
descriptors.src_page_indices,
descriptors.dense_page_indices,
page_nbytes=_page_nbytes_from_page_tensor(page_buffer),
)
return True
except Exception as exc:
_log_tai_ipc_materialize_fallback(
"paged_span_ce_kernel_failed",
"CP shared KV tai IPC CE paged span materialize failed; falling back "
"to SM IPC or collective materialize. cp_rank=%s cp_size=%s spans=%s "
"page_shape=%s dense_shape=%s error=%s",
layout.cp_rank,
layout.cp_size,
spans,
tuple(page_buffer.shape),
tuple(dense_page_buffer.shape),
exc,
)
return False
def _get_or_build_current_ipc_slot_descriptors(
*,
slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap,
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+8 -37
View File
@@ -20,7 +20,6 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_batch_prefix_slot_spans,
cp_shared_kv_debug_enabled,
cp_shared_kv_debug_log,
cp_shared_kv_mla_prefetch_enabled,
cp_shared_kv_mla_prefetch_log,
cp_shared_kv_mla_prefetch_log_enabled,
cp_shared_kv_mla_prefetch_should_log_layer,
@@ -477,14 +476,6 @@ def _log_cp_shared_kv_index_prefetch_fallback(
)
def _should_log_missing_index_prefetcher() -> bool:
# When MLA/index prefetch is disabled, a missing index prefetcher is the
# expected sync-materialize path, not a fast-path fallback. Still log consume
# misses when a prefetcher exists because those indicate an enabled prefetch
# pipeline failed to provide the requested layer/current buffer.
return cp_shared_kv_mla_prefetch_enabled()
class BaseIndexerMetadata(ABC):
@abstractmethod
def get_seqlens_int32(self) -> torch.Tensor:
@@ -798,25 +789,15 @@ class Indexer(MultiPlatformOp):
)
if prefetched is not None:
return prefetched
_log_cp_shared_kv_index_prefetch_fallback(
"current_consume_miss",
"prefetcher did not provide current index prefix buffer; "
"falling back to sync partial-current index compose. "
"layer=%s cp_rank=%s prefix_lens=%s extend_lens=%s "
"logical_page_table_shape=%s",
layer_id,
layout.cp_rank,
prefix_lens,
extend_lens,
tuple(logical_page_table.shape),
)
else:
if _should_log_missing_index_prefetcher():
if int(layer_id) > int(
getattr(forward_batch.token_to_kv_pool, "start_layer", 0)
):
_log_cp_shared_kv_index_prefetch_fallback(
"current_missing_prefetcher",
"index prefetcher is unavailable; falling back to sync "
"partial-current index compose. layer=%s cp_rank=%s "
"prefix_lens=%s extend_lens=%s logical_page_table_shape=%s",
"current_consume_miss",
"prefetcher did not provide current index prefix buffer; "
"falling back to sync partial-current index compose. "
"layer=%s cp_rank=%s prefix_lens=%s extend_lens=%s "
"logical_page_table_shape=%s",
layer_id,
layout.cp_rank,
prefix_lens,
@@ -901,16 +882,6 @@ class Indexer(MultiPlatformOp):
layout.cp_rank,
tuple(logical_page_table.shape),
)
else:
if _should_log_missing_index_prefetcher():
_log_cp_shared_kv_index_prefetch_fallback(
"missing_prefetcher",
"index prefetcher is unavailable; falling back to sync paged "
"materialize. layer=%s cp_rank=%s logical_page_table_shape=%s",
layer_id,
layout.cp_rank,
tuple(logical_page_table.shape),
)
if cp_shared_kv_debug_enabled():
cp_shared_kv_debug_log(
python/sglang/srt/layers/attention/nsa_backend.py
+63 -41
View File
@@ -76,7 +76,7 @@ from sglang.srt.layers.attention.utils import (
mla_quantize_and_rope_for_fp8,
seqlens_expand_triton,
)
from sglang.srt.layers.dp_attention import get_attention_cp_group, get_attention_tp_size
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils import is_cuda, is_hip
@@ -166,7 +166,7 @@ def _maybe_start_cp_shared_kv_attention_prefetch(
This hook intentionally runs after current-layer index/MLA cache
materialization has finished and immediately before the attention kernel.
Earlier hooks can make the next-layer collective overlap current-layer KV
materialization/reduce instead of attention, which can serialize or contend
materialization instead of attention, which can serialize or contend
with the current layer's required cache work.
"""
@@ -1061,17 +1061,11 @@ class NativeSparseAttnBackend(
prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
try:
pynccl_available = (
getattr(get_attention_cp_group(), "pynccl_comm", None) is not None
)
except Exception:
pynccl_available = False
cp_shared_kv_mla_prefetch_log(
"create_result cp_rank=%s cp_size=%s uses_cp_shared_kv=%s "
"forward_mode=%s topk_transform=%s batch_size=%s has_mla=%s "
"has_index=%s prefix_lens=%s extend_lens=%s seq_lens=%s "
"real_pages=%s page_table_shape=%s pynccl=%s",
"real_pages=%s page_table_shape=%s",
getattr(layout, "cp_rank", None),
getattr(layout, "cp_size", None),
getattr(forward_batch, "uses_cp_shared_kv", None),
@@ -1095,7 +1089,6 @@ class NativeSparseAttnBackend(
tuple(metadata.page_table_1.shape)
if metadata.page_table_1 is not None
else None,
pynccl_available,
)
def _cal_indexer_k_start_end(
@@ -2468,7 +2461,7 @@ class NativeSparseAttnBackend(
forward_batch, "cp_shared_kv_index_prefetcher", None
)
if index_prefetcher is not None:
index_prefetcher.launch_pending_reduce()
index_prefetcher.finalize_pending_materialize()
try:
if nsa_impl == "tilelang":
@@ -2654,40 +2647,69 @@ class NativeSparseAttnBackend(
ragged_current_req_id = ragged_current_req_id[
: int(current_locs_for_reuse.shape[0])
]
slot_remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
kv_cache=kv_cache,
remap_logical_pages=metadata.real_page_table,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefetched_kv = None
mla_prefetcher = getattr(
forward_batch, "cp_shared_kv_mla_prefetcher", None
)
kv_cache, page_table_1_flattened = (
materialize_prefix_and_reuse_current_kv_page_slots(
if mla_prefetcher is not None:
prefetched_kv = (
mla_prefetcher.consume_prefix_with_current(
layer_id=layer.layer_id,
kv_cache=kv_cache,
logical_locs=page_table_1_flattened,
current_kv_cache=current_kv_cache,
current_locs=current_locs_for_reuse,
loc_req_id=logical_locs_row_ids,
current_req_id=ragged_current_req_id,
)
)
if prefetched_kv is not None:
kv_cache, page_table_1_flattened = prefetched_kv
else:
if (
mla_prefetcher is not None
and int(layer.layer_id) > int(getattr(forward_batch.token_to_kv_pool, "start_layer", 0))
):
_log_cp_shared_kv_mla_prefetch_fallback(
"ragged_current_consume_miss",
"MLA prefetch fast path did not provide "
"RAGGED partial-current prefix; using "
"synchronous compose. cp_rank=%s layer=%s "
"prefix_lens=%s extend_lens=%s "
"current_rows=%s page_table_shape=%s",
forward_batch.cp_shared_kv_layout.cp_rank,
layer.layer_id,
[int(x) for x in prefix_lens_cpu],
[int(x) for x in extend_lens_cpu],
int(current_kv_cache.shape[0]),
tuple(page_table_1_flattened.shape),
)
slot_remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
kv_cache=kv_cache,
logical_locs=page_table_1_flattened,
current_kv_cache=current_kv_cache,
current_locs=current_locs_for_reuse,
slot_remap=slot_remap,
remap_logical_pages=metadata.real_page_table,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefix_pages=prefix_pages,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
loc_req_id=logical_locs_row_ids,
current_req_id=ragged_current_req_id,
layer_id=layer.layer_id,
nvtx_source="mla.ragged_partial_current_sync",
current_page_writer_ranks=(
maybe_build_current_page_writer_ranks(
forward_batch=forward_batch,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
layout=forward_batch.cp_shared_kv_layout,
)
),
)
kv_cache, page_table_1_flattened = (
materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=page_table_1_flattened,
current_kv_cache=current_kv_cache,
current_locs=current_locs_for_reuse,
slot_remap=slot_remap,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefix_pages=prefix_pages,
prefix_slot_spans=prefix_slot_spans,
current_slot_spans=current_slot_spans,
loc_req_id=logical_locs_row_ids,
current_req_id=ragged_current_req_id,
layer_id=layer.layer_id,
nvtx_source="mla.ragged_partial_current_sync",
)
)
)
if envs.SGLANG_NSA_DEQUANT_ONLY_TOPK.get():
kv_cache = self._dequant_topk_inplace(
kv_cache,
@@ -2766,7 +2788,7 @@ class NativeSparseAttnBackend(
)
finally:
if mla_prefetcher is not None:
mla_prefetcher.launch_pending_reduce()
mla_prefetcher.finalize_pending_materialize()
mla_prefetcher.wait_attention_window()
index_prefetcher = getattr(
forward_batch, "cp_shared_kv_index_prefetcher", None
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+843 -123
View File
File diff suppressed because it is too large Load Diff