Protect CP shared-KV cache-hit correctness under batched FP8 reuse
Cache-hit GSM8K regressions only appeared after the second pass reused request-specific suffix pages, so this change adds fail-fast transfer validation, masks stale rectangular page-table tails, and extends CUDA/unit coverage across FP8 CP shared-KV write, load, top-k, and materialization paths. The temporary ledger records eliminated hypotheses to prevent re-debugging the same L2 and persistent-cache paths.\n\nConstraint: CP shared KV stores physical pages but scheduler-visible semantics must remain valid-token/page-bounded.\nConstraint: bs>1 FP8 prefill must preserve existing CP shared-KV fast paths without silent fallback.\nRejected: Blame raw HiCache L2 load without tests | L2 KV and index backup/load/materialize roundtrips pass on remote CUDA.\nRejected: Disable current/partial reuse broadly | hides the cache-hit contract regression and costs performance.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not weaken CP shared-KV fail-fast or rectangular-tail masking without rerunning second-pass cache-hit accuracy tests.\nTested: remote CUDA pytest for fused FP8 MLA store, fused persistent index store, L2-loaded FP8 KV materialize, L2-loaded index materialize, ragged top-k offset, TAI batched index MQA prepare.\nTested: local py_compile for touched test files and git diff --check.\nNot-tested: full second-pass GSM8K accuracy after these diagnostic tests; root cause remains under investigation.
This commit is contained in:
laoyao0822
@@ -2032,3 +2032,178 @@ Current next validation requirement:
|
||||
|
||||
- Treat this as an active post-fix hang/regression, not an old-process artifact.
|
||||
- If this hang needs deeper diagnosis, run with an uncapped or higher CP bs>1 timing log limit so the final scheduler/bootstrap boundary is visible; the current log reaches watchdog without a direct failing traceback.
|
||||
|
||||
### 2026-06-07 repeat GSM8K post-fix death: current bounded finding
|
||||
|
||||
New evidence from `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_164817.log`:
|
||||
|
||||
- The process did not die with a Python traceback, explicit OOM, CP fail-fast, or `Scheduler hit an exception`.
|
||||
- Last normal prefill batch was `17:14:54` with `bs=5`, `new_token=640`, `cached_token=3200`, `queue_req=47`, `inflight_req=9` on all CP ranks.
|
||||
- The final batch's node/rid write-through backup reached `submit_write_cp_layer final ack`, `attached prepared CP backup`, and `writing_check released 5 write locks` on all CP ranks.
|
||||
- The final watchdog pending list still contained the final-batch rids (`3732391e...`, `17f4f067...`, `55a9e51c...`, `2827e296...` etc.), so the `POST /generate 200 OK` lines at `17:14:54` are not proof that the final batch completed; they may be earlier inflight completions.
|
||||
- `GET /health` was healthy through `17:14:18`; the watchdog started at `17:15:17` and reported `last_heartbeat=17:14:54`.
|
||||
- The timing diagnostics were ineffective at the actual failure point because the run had a malformed env assignment: `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT` parsed as `"-1SGLANG_CP_SHARED_KV_BS_GT1_DEBUG=1"` and fell back to default `256`, so common scheduler boundary events hit their per-event cap around `16:59`.
|
||||
|
||||
Current inference:
|
||||
|
||||
- The strongest current suspicion is a scheduler-side hang after the final batch, before any later detokenizer response. The likely boundaries are still:
|
||||
1. `process_disagg_prefill_inflight_queue()` polling inflight KV transfer;
|
||||
2. `PrefillBootstrapQueue.pop_bootstrapped()` / subsequent scheduler admission for the next batch;
|
||||
3. `stream_output()` / ZMQ send if the detokenizer side stops consuming.
|
||||
- Existing logs cannot distinguish these because the final boundary markers were capped before the failure.
|
||||
- Do not keep re-reading this log expecting a hidden traceback; the current log does not contain one.
|
||||
|
||||
Diagnostic change added locally:
|
||||
|
||||
- Added an env-gated CP shared-KV poll queue consensus guard behind the existing `SGLANG_CP_SHARED_KV_BS_GT1_DEBUG=1`.
|
||||
- The guard validates `(queue_len, rid_hash)` across attn-TP and attn-CP CPU groups before the old variable-length poll tensor all-reduce. If queue length or rid order diverges, it raises `[CP_SHARED_KV_FAIL_FAST][poll_queue]` instead of risking a silent collective hang.
|
||||
- This intentionally adds extra scalar collectives only when the existing `SGLANG_CP_SHARED_KV_BS_GT1_DEBUG=1` is enabled; normal production and timing-only paths do not pay this cost.
|
||||
- Regression test added: `test_poll_consensus_debug_fails_before_shape_mismatch_hang`.
|
||||
|
||||
Next reproduction requirements:
|
||||
|
||||
- Use correctly separated env assignments:
|
||||
- `SGLANG_CP_SHARED_KV_BS_GT1_TIMING=1`
|
||||
- `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT=0`
|
||||
- `SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS=0`
|
||||
- optional: `SGLANG_CP_SHARED_KV_BS_GT1_DEBUG=1` for one diagnostic run only.
|
||||
- If the next run reports `[CP_SHARED_KV_FAIL_FAST][poll_queue]`, root cause is queue length/order divergence across ranks.
|
||||
- If it still watchdogs without poll fail-fast, the next suspect is a call-count mismatch between ranks or a non-poll blocking site (`stream_output`, admission/eviction planning, or receive loop); use the uncapped boundary markers to locate the last `*_start` without matching `*_done`.
|
||||
|
||||
## 2026-06-07 cache-hit GSM8K accuracy regression
|
||||
|
||||
New user-provided evidence on `g0034:17100` with full GSM8K (`1319` questions, parallel=64, temp=0):
|
||||
|
||||
```text
|
||||
cold/first run: Accuracy: 0.955 Invalid: 0.000 Latency: 207.174 s Output throughput: 668.366 token/s
|
||||
second run: Accuracy: 0.687 Invalid: 0.001 Latency: 211.220 s Output throughput: 660.263 token/s
|
||||
```
|
||||
|
||||
Interpretation:
|
||||
|
||||
- The cold path can still reach the expected GLM-5.1 GSM8K accuracy, so the base model, decode EAGLE, FP8 decode, and non-hit prefill path are not globally broken.
|
||||
- The immediate second run is a cache-hit path and drops accuracy sharply. This shifts the top suspect away from generic EAGLE/FP8 decode and toward cache-hit semantics: radix/HiCache match length, page-tail valid length, CP shared-KV load/materialize, or cached index/MLA reuse under page-aligned storage.
|
||||
- Throughput remains similar between runs, so this is correctness corruption rather than a service-death or obvious slow-path-only issue.
|
||||
|
||||
Current top hypothesis to verify next:
|
||||
|
||||
- Some path exposes page-rounded/padded KV as valid cached prefix, or uses cached prefix metadata without flooring to the last valid page boundary. A later request can then skip recomputing tail tokens or consume stale/padded page data, which would be invisible on the cold run and show up only after cache hits.
|
||||
|
||||
Immediate code audit targets:
|
||||
|
||||
1. `hiradix_cache.py` CP page-floor/prune/match logic: cached match length returned to scheduler must be valid-token/page-floor length, never physical padded extent.
|
||||
2. `radix_cache.py` / `schedule_batch.py`: `prefix_indices`, `host_hit_length`, `cache_protected_len`, and `extend_input_len` must not include padded tail tokens.
|
||||
3. `nsa/utils.py` and current reuse paths: batch compute padding and physical page padding must not alter request logical lengths used for top-k/MLA attention.
|
||||
4. HiCache load path: host backed CP node should obey the same page-minimum contract; if a node is backed only to page floor, cache hit must not round up to the next page.
|
||||
|
||||
## 2026-06-07 继续定位：cache-hit 回退集中到 slot-remap 的矩形 page table 合同
|
||||
|
||||
新增证据：
|
||||
|
||||
- 两次 GSM8K 全量同进程复跑：cold run `Accuracy=0.955`，immediate cache-hit run `Accuracy=0.687`。cold 正确说明模型、decode 基本路径、FP8 非命中路径不是全局损坏；错误集中在 cache-hit 前缀被 prefill 重新消费后。
|
||||
- 远端日志 `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_180649.log` 第二轮出现稳定 cache-hit batch：`#new-seq: 5`, `#new-token: 320/384`, `#cached-token: 3520/3712`，约等于每请求 704+ cached tokens。
|
||||
- `hiradix_cache` 初查没有看到把 non-page-aligned hit 向上 round-up 的证据；当前仍是 page-floor 方向。因此下一主线不是 radix round-up，而是 cache-hit prefix materialize/compose。
|
||||
- `nsa_backend.py` 使用 `metadata.real_page_table` 构造 token/index slot-remap；`real_page_table` 来自矩形 `req_to_token[:, :max_seq_len]` 按 page stride 取样。bs>1 且请求长度不同/slot 复用时，矩形尾部可能包含请求 valid page 之外的旧 loc。
|
||||
- `cp_shared_kv_runtime.build_slot_page_inverse()` 对 duplicate logical page 使用 `scatter_`，同一 logical page 会映射到最后一次出现的 slot。若 valid prefix/current page 也出现在某个请求的矩形无效尾部，page_inverse 可能把合法 topk loc remap 到一个未被 prefix/current spans materialize/reduce 的 stale slot，cache-hit attention 会读到 0/错误 KV。该路径只在 cache-hit prefix materialize 时高概率触发，符合 cold 正常、hit 掉点的现象。
|
||||
|
||||
待验证/修复方向：
|
||||
|
||||
1. 给 `metadata.real_page_table`/index logical page table 增加 batch-aware valid-page mask：每个 request 只保留 `ceil(seq_len/page_size)` 或 `ceil((prefix_len+extend_len)/page_size)` 范围内的 page，范围外置 0，不参与 slot inverse。
|
||||
2. 增加单测：构造 bs>1 矩形 logical_pages，尾部 stale slot 重复了前面 valid logical page，验证 slot remap 不能把 valid loc 映射到 stale slot。
|
||||
3. 修复应优先在 slot-remap 输入处清洗 invalid rectangular tail，而不是禁止 cache-hit/current reuse。
|
||||
|
||||
## 2026-06-07 修复：在 NSA metadata 边界清理 rectangular page-table tail
|
||||
|
||||
- 变更点：`NSAMetadata.real_page_table` 构建后立即调用 `mask_batch_logical_pages_to_valid_lengths()`，按每个 request 的有效 `seq_len` 保留 `ceil(seq_len/page_size)` 个 page，后续矩形尾部统一置 0。
|
||||
- 选择该位置的原因：`real_page_table` 同时喂给 MLA KV materialize、index materialize、prefetcher 与 slot-remap cache。集中清理比在每个 consume/materialize 调用点单独清理更不容易遗漏。
|
||||
- 保留的合同：最后一个有效 partial page 仍保留；不会把 token 级 valid length 截成 compact token cache，只清理 page-table 的 invalid row tail。
|
||||
- 覆盖的风险：当 bs>1 不同 seq_len 混排、request slot 复用导致矩形尾部出现旧 logical page，`build_slot_page_inverse`/TAI inverse 的 duplicate page 写入不再能把有效 page remap 到 stale slot。
|
||||
- 验证：远端容器 `py_compile` 通过；`test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_valid_page_mask_prevents_stale_rectangular_tail_remap` 通过；完整 `test_cp_shared_kv_runtime.py` 通过 `114 passed, 5 warnings, 2 subtests passed`。
|
||||
- 未验证：最新 ETE cache-hit GSM8K 二次全量精度仍需用户重启新代码后复测。
|
||||
|
||||
### 2026-06-07 澄清：cuda graph replay 不是当前 prefill ETE 主线
|
||||
|
||||
- 用户确认当前 prefill 不走 cuda graph。
|
||||
- 因此 `init_forward_metadata_replay_cuda_graph()` 中同类 rectangular-tail 复用风险不作为本次 GSM8K cache-hit 掉点的 active root cause。
|
||||
- 当前主线收敛到非 cuda-graph `init_forward_metadata()` 构建的 `real_page_table`，以及其下游 MLA/index materialize、prefetch、partial/current reuse。
|
||||
|
||||
### 2026-06-07 继续审计：`forward_indexer()` 的 direct req_to_token page-table 路径
|
||||
- 发现 `nsa_indexer.forward_indexer()` 仍直接从 `forward_batch.req_to_token_pool.req_to_token[:, ::page_size]` 构造 `block_tables`，再传入 `_maybe_materialize_shared_index_buffer()`。
|
||||
- 这条路径若在 CP shared-KV + cache-hit 下启用，会绕过 `nsa_backend.init_forward_metadata()` 中对 `real_page_table` 的 logical-tail mask，存在同类 stale rectangular-tail 风险。
|
||||
- 但当前 CUDA/Hopper prefill 的 `forward_cuda()` 在 `_is_cuda or _is_hip` 分支内走 `_get_topk_paged/_get_topk_in_seq_cp_pair/_get_topk_ragged`，只有非 CUDA/HIP 分支才调用 `forward_indexer()`。
|
||||
- 结论：这是潜在维护风险，不是当前 prefill ETE 的 active root cause；当前继续审计 CUDA active path。
|
||||
|
||||
### 2026-06-07 继续审计：CUDA graph replay / MTP precompute 的 direct req_to_token 路径
|
||||
- 发现 `nsa_backend.init_forward_metadata_replay_cuda_graph()` 以及 `nsa_backend_mtp_precompute.py` 的 decode/target-verify/draft precompute 分支也直接从 `req_to_token` 构造 token/page table，并在 `real_page_size > 1` 时 transform 成 `real_page_table`。
|
||||
- 这些路径当前未调用 `mask_batch_logical_pages_to_valid_lengths()`，因此在 decode/target-verify/draft CUDA graph replay 或 MTP precompute 场景中存在同类 stale rectangular-tail 风险。
|
||||
- 用户已确认当前 prefill 不走 cuda graph；因此它们不是本次 prefill GSM8K cache-hit 掉点 active root cause。
|
||||
- 后续若要覆盖 decode graph / draft graph，应在 precompute/replay 边界同样按 logical `seq_lens` mask real-page-table，且需要专门测试，不能把它混入当前 prefill 修复判断。
|
||||
|
||||
### 2026-06-07 再确认：page 物理粒度不等于 token 可见长度向上取整
|
||||
|
||||
本轮代码复核结论：
|
||||
|
||||
- `ReqToTokenPool.alloc()` 复用 slot 时只从 `free_slots` 取回 `req_pool_idx`，`free()` 也只把 slot 放回队列并清空 `req.req_pool_idx`；`req_to_token` 行内容不会被清零。因此矩形 batch 读取到 `max_seq_len` 时，短请求 valid token 之后的 tail 可以携带旧请求残留。
|
||||
- `write_cache_indices()` 只写 `[0, prefix_len)` 与 `[prefix_len, seq_len)`，不会清理 `[seq_len, max_seq_len)`。这进一步确认 stale tail 是真实输入状态，不是单测构造出来的假风险。
|
||||
- CP HiCache 的合同是：radix/scheduler 可见长度使用 `logical_len`/`valid_len`，物理 host/device reservation 使用 `padded_len`，且 `padded_len` 必须是 `page_size` 倍数。也就是说 cache 存储和释放以 page 为最小物理单位，但命中长度和 split 仍不能把 padding token 暴露为有效 token。
|
||||
- CP HiCache split 强制 `split_len % page_size == 0`；exact non-page-aligned hit 会 floor 到上一页，stale tail child 会被 prune；chunked prefill backup 也只 backup 到 `floor_to_page_len()`。这些都支持同一个合同：可以牺牲一个 sub-page tail 来换 page 管理一致性，但不能把 page padding 当成有效语义 token。
|
||||
- 因此 `real_page_table` 的 tail mask 是正确方向：保留 `ceil(valid_seq_len / page_size)` 个 page，包括最后一个 partial page；只把每行 valid page 数之后的矩形尾部置 0，避免 stale duplicate page 参与 `page_inverse.scatter_()`。这不会丢失以 page 存储的合法 cache，只会阻断无效矩形尾部覆盖 slot inverse。
|
||||
|
||||
继续审计到的类似风险：
|
||||
|
||||
1. `metadata.page_table_1` 仍保留 token 粒度矩形表，未整体 mask。当前 CUDA prefill 的 CP shared-KV materialize/slot-remap 已经改用 masked `metadata.real_page_table`，但 `topk_transform()` 的 paged fused path 仍以 `page_table_1` 为输入。现有 debug validator 只检查输出是否在整个 page table 的 min/max 范围内，不检查“输出 loc 是否属于该 request 的 valid prefix”。如果 cache-hit 复测仍掉点，下一步应加一个 env/debug 级 per-row membership validator，而不是继续禁用 current/partial reuse。
|
||||
2. `nsa_indexer.forward_indexer()` 的非 CUDA/HIP 路径会直接从 `req_to_token[:, ::page_size]` 构造 `block_tables`，绕过 metadata 边界 mask。当前 Hopper CUDA prefill 不走这条路径，归类为维护风险。
|
||||
3. CUDA graph replay / MTP precompute 路径仍存在 direct `req_to_token` page-table 构造。用户确认 prefill 无 cuda graph，所以不是当前 active root cause；后续若要覆盖 decode/draft graph，应单独补 mask 和测试。
|
||||
4. MHA FP8 的 `page_table_1_flattened` 是按每行 `indexer_seq_lens_cpu` 做 `page_table[i, :kv_len]` 的 valid slice concat，不读矩形 tail；当前看不是同类 stale-tail 风险。
|
||||
5. prefill disaggregation transfer 当前读 `req_to_token[req_pool_idx, start:end]` / `[:seq_len]` 这类 valid token slice；这类路径本身不消费矩形 tail。仍需保留已有 page-count fail-fast，但它不是本轮 cache-hit 精度掉点的主嫌。
|
||||
|
||||
后续判据：
|
||||
|
||||
- 如果最新 ETE 二次 GSM8K 精度恢复，说明主要 root cause 是 `real_page_table` slot-remap tail。
|
||||
- 如果仍掉点，优先打开/补强 topk per-request membership validator，检查 `topk_indices` 是否引用了每个 request valid range 外或错误 request 的 loc；不要先做大范围 fallback/禁用。
|
||||
|
||||
### 2026-06-07 二次 GSM8K 仍掉点：风险从“所有 cache-hit”收敛到 non-page-aligned CP HiCache 命中长度/内容
|
||||
|
||||
新远端证据来自 `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_193039.log`：
|
||||
|
||||
- 用户同进程连续两轮 GSM8K：第一轮 `Accuracy=0.953`，第二轮 `Accuracy=0.692`，吞吐基本相同。
|
||||
- 第一轮并不是完全无 cache：从 `19:35:08` 起已经出现大量 `#new-seq: 5, #cached-token: 3200`，即约每请求 640 cached tokens；第一轮仍正确。
|
||||
- 第二轮主要变成 `#new-token: 320/384, #cached-token: 3520/3584/3648...`，即命中更长、更贴近每个 GSM8K 问题自身的缓存。
|
||||
- 第一轮开始处 CP HiCache 写入了 `node_id=4 logical_len=698`；这是 non-page-aligned node，物理 page 覆盖会到 704 tokens。后续第二轮的典型 `cached-token=3520` 正好对应 5 个请求各 704 tokens。
|
||||
|
||||
当前推断：
|
||||
|
||||
- “cache-hit 本身全部坏”不是最强解释，因为第一轮已有 640-token 共享前缀命中且精度正常。
|
||||
- 更强嫌疑是：某个 CP HiCache/radix/load-back/scheduler-visible 前缀长度路径，把 `logical_len=698` 的 node 以物理 `padded_len=704` 暴露给 scheduler/attention，或者 load-back/backup 的最后 partial page 内容/索引与 logical valid token 不一致。
|
||||
- 这会让第二轮复用 per-question cache 时跳过 6 个应计算 token，或消费 page padding/stale KV；冷跑和只命中 640 共享前缀时不明显，复跑命中 question-specific non-page-aligned node 时掉点。
|
||||
|
||||
下一步审计目标：
|
||||
|
||||
1. `hiradix_cache.match_prefix()` / `_match_prefix_helper()` / `load_back()`：确认返回给 scheduler 的 `prefix_indices` 长度来自 logical valid length，而不是 `padded_len`。
|
||||
2. `schedule_batch.py`：确认 `cache_protected_len`、`extend_input_len`、日志里的 `#cached-token` 是否可能为满足 page 对齐被向上取整。
|
||||
3. CP HiCache backup/load path：确认 host/device node 的 `value` 只包含 logical indices；物理 page reservation/owned positions 不应改变 node visible length。
|
||||
4. 若代码确认没有 length exposure，再回到 top-k membership validator；但当前日志优先级应先查 698→704 这条链。
|
||||
|
||||
### 2026-06-07 复测后收敛：第二轮掉点不是所有 cache-hit 坏，而是 first-run suffix page 持久化复用可疑
|
||||
|
||||
最新证据来自 `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_193039.log`，对应用户同进程两轮 GSM8K：第一轮 `Accuracy=0.953`，第二轮 `Accuracy=0.692`。
|
||||
|
||||
日志事实：
|
||||
|
||||
- 该日志内没有 `CP_SHARED_KV_FAIL_FAST`、`CP_SHARED_KV_FALLBACK`、`RuntimeError`、`Traceback`、`Health check failed`、`HiCache-load`、`CacheCtrl-load`、`load_back`。
|
||||
- 第一轮已经有稳定 cache-hit：CP0 典型 `#new-seq: 5, #new-token: 640/704/576, #cached-token: 3200`，约等于每请求 640 cached tokens；第一轮仍正确。
|
||||
- 第二轮主要变成 `#new-seq: 5, #new-token: 320/384, #cached-token: 3520/3584/...`，即每请求多复用了约 64 tokens 的 question-specific suffix page。
|
||||
- 第一轮 CP0 写入大量 non-page-aligned suffix node，例如 `logical_len=95/81/140/66/86`，且 `owned_positions=64`；bs=5 时 CP0 常见 `valid_local_tokens=320`，即每个请求写入首个 owner page。
|
||||
- 第二轮 CP0 `local_out_cache_loc` 多为小 suffix：如 `split_tokens=165 out_cache_tokens=168 local_tokens=165`、`split_tokens=63 out_cache_tokens=64 local_tokens=63`，说明当前 extend/write 仍按 valid rows 截掉 physical padding。
|
||||
|
||||
修正当前 root-cause 优先级：
|
||||
|
||||
- 仅用“radix/scheduler 把 698 向上暴露成 704”解释不充分。代码中 CP exact hit 仍有 page-floor 方向，且日志也能解释为第二轮命中共享 640 page 后，再命中第一轮写入的每题 suffix 首个 64-token page。
|
||||
- 因此当前更强嫌疑是 bs>1 下持久化 direct write 的 row-order/physical-loc contract：第一轮正常使用 current K/V 参与计算并得到正确答案，但同一轮写入 device/HiCache 的 MLA KV 或 NSA index K 在第二轮作为 prefix 被复用时内容或行序错误。
|
||||
- 需要优先审计 `get_cp_shared_kv_local_out_cache_loc()`、`select_cp_local_valid_rows_for_cache_write()`、MLA direct store、index direct store 之间的行序一致性，以及 prefix materialize 读取这些 page 时是否按相同 owner-lane order 解释。
|
||||
|
||||
下一步最小验证：
|
||||
|
||||
1. 构造/补强单测覆盖 bs=5、page_size=64、prefix=640、extend=[95,81,140,66,86] 这类 GSM8K 形态，验证 local valid row selector 输出顺序与 local physical out loc 顺序逐 request/page 对齐。
|
||||
2. 若单测无法暴露，加入 env-gated、限频的 runtime validator：在 first-run suffix direct write 后读取刚写入的 index/MLA buffer，与 current local rows 做 checksum/row-id 对比。优先验证 index K，因为它有较清晰的 `local_key -> physical_out_loc` 写入路径；MLA FP8 packed store 需要额外确认量化/布局。
|
||||
3. 不再通过大范围禁用 current/partial reuse 判断根因；否则会掩盖“第一轮 direct write 与第二轮 persistent reuse 不等价”的核心问题。
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
# CP shared-KV bs>1 GSM8K cache-hit 掉点临时排查记录
|
||||
|
||||
> 目的：上下文压缩后先读本文件，避免重复扫描远端大日志与重复审计同一批路径。
|
||||
|
||||
## 当前复现
|
||||
|
||||
- 远端日志：`g0034:/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_193039.log`
|
||||
- 用户连续两次全量 GSM8K，同一 prefill 进程：
|
||||
- 第 1 次：`Accuracy=0.953`, `Invalid=0.000`, `Latency=211.445s`, `Output throughput=654.008 token/s`
|
||||
- 第 2 次：`Accuracy=0.692`, `Invalid=0.002`, `Latency=214.264s`, `Output throughput=654.376 token/s`
|
||||
|
||||
## 日志已确认事实
|
||||
|
||||
- 无显式错误：`CP_SHARED_KV_FAIL_FAST=0`, `CP_SHARED_KV_FALLBACK=0`, `RuntimeError=0`, `Traceback=0`, `Health check failed=0`。
|
||||
- 无明显 L2 load-back：`HiCache-load=0`, `CacheCtrl-load=0`, `load_back=0`。
|
||||
- 第一轮并非全冷：CP0 典型 batch 已有 `#new-seq=5, #new-token=640/704/576, #cached-token=3200`，约每请求 640 cached tokens，仍然精度正常。
|
||||
- 第二轮典型变为 `#new-seq=5, #new-token=320/384, #cached-token=3520/3584/...`，约每请求比第一轮多复用 64 tokens 的 question-specific suffix page。
|
||||
- 第一轮 CP0 写入大量 non-page-aligned suffix node：如 `logical_len=95/81/140/66/86`, `owned_positions=64`。bs=5 时 CP0 常见 `valid_local_tokens=320`，即每请求首个 owner page。
|
||||
- 第二轮 CP0 `local_out_cache_loc` 仍按 valid rows 计算，例如 `split_tokens=63 out_cache_tokens=64 local_tokens=63 valid_local_tokens=63`，说明当前 write 侧没有直接把 padding token 当 valid row 写。
|
||||
|
||||
## 已排除 / 降低优先级
|
||||
|
||||
- 不是所有 cache-hit 都坏：第一轮已有共享 640-token cache-hit，精度仍正常。
|
||||
- 不是明显服务错误或 fallback 慢路径：日志无 fail-fast/fallback/crash。
|
||||
- 不是当前 prefill CUDA graph replay：用户确认 prefill 不走 cuda graph。
|
||||
- 单纯 “radix 把 698 向上暴露成 704” 不是最强解释：代码仍有 CP exact hit page-floor 方向，日志也可解释为第二轮额外命中了第一轮写入的每题 suffix 首 page。
|
||||
|
||||
## 当前主嫌
|
||||
|
||||
bs>1 下 first-run suffix page 的持久化 direct write 与 second-run prefix reuse 不等价：
|
||||
|
||||
1. 第一轮当前请求使用 current K/V 或 current index 参与计算，答案正确。
|
||||
2. 同一轮把 suffix 首 page 写入 device/HiCache 持久 cache。
|
||||
3. 第二轮把这个 page 当 prefix 复用后精度掉点。
|
||||
|
||||
因此优先审计：
|
||||
|
||||
- `get_cp_shared_kv_local_out_cache_loc()`
|
||||
- `select_cp_local_valid_rows_for_cache_write()`
|
||||
- MLA direct store (`_maybe_write_cp_shared_local_mla_kv`)
|
||||
- NSA index direct store (`_store_cp_shared_local_index_k_cache`)
|
||||
- prefix materialize/slot-remap 读取持久 page 时是否按相同 owner-lane order 解释。
|
||||
|
||||
## 下一步验证方向
|
||||
|
||||
1. 先补/跑 CPU 单测：bs=5, page_size=64, prefix=640, extend=[95,81,140,66,86]，验证 local valid row selector 输出顺序与 local physical out loc 顺序逐 request/page 对齐。
|
||||
2. 如果单测没暴露，加入 env-gated 限频 runtime validator：first-run suffix direct write 后读回刚写的 index/MLA buffer，与 current local rows 做 checksum/row-id 对比。
|
||||
3. 不先大范围禁用 current/partial reuse；这会掩盖 direct write 与 persistent reuse 不等价的问题。
|
||||
|
||||
## 2026-06-07 增量发现：CPU 侧 bs=5 row/loc 顺序探针
|
||||
|
||||
远端容器已跑过一个不依赖 CUDA 的顺序探针：
|
||||
|
||||
- 脚本：`/sgl-workspace/sglang-tai/tmp_cp_bs5_order_probe.py`
|
||||
- 参数：`page_size=64`, `cp_size=8`, `prefix=[640]*5`, `extend=[95,81,140,66,86]`
|
||||
- 结论：Python planner / valid-row selector / local out loc 在该 GSM8K-like 形状下顺序一致。
|
||||
|
||||
输出摘要：
|
||||
|
||||
```text
|
||||
rank0: compute=320 valid_rows=320 locs=320 per_req=[64,64,64,64,64]
|
||||
rank1: compute=320 valid_rows=136 locs=136 per_req=[31,17,64,2,22]
|
||||
rank2: compute=320 valid_rows=12 locs=12 per_req=[0,0,12,0,0]
|
||||
rank3-7: valid_rows=0 locs=0
|
||||
```
|
||||
|
||||
含义：
|
||||
|
||||
- rank0 写每个 request 的 suffix 首 page，符合第二轮多命中约 64 tokens/request 的现象。
|
||||
- rank1/rank2 写剩余 valid tail；dummy compute-padding rows 没进入 persistent write rows。
|
||||
- 这降低了 “Python split/selector 把 valid rows 和 out loc 排错序” 的优先级。
|
||||
- 不能排除：TAI fused store、index fused store、或 persistent prefix materialize 对同一批 loc 的解释不一致。
|
||||
|
||||
若需复跑：
|
||||
|
||||
```bash
|
||||
ssh g0034 "docker exec sglang-glm5-dev-2 bash -lc 'cd /sgl-workspace/sglang-tai && PYTHONPATH=python python /sgl-workspace/sglang-tai/tmp_cp_bs5_order_probe.py'"
|
||||
```
|
||||
|
||||
## 2026-06-07 增量发现：direct write 代码路径事实
|
||||
|
||||
MLA direct write：
|
||||
|
||||
- 文件：`python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
|
||||
- 函数：`_maybe_write_cp_shared_local_mla_kv`
|
||||
- 流程：
|
||||
1. `get_cp_shared_kv_local_out_cache_loc(forward_batch)` 取得 local logical locs。
|
||||
2. `select_cp_local_valid_rows_for_cache_write(...)` 从 compute-padded local rows 中只取 valid rows。
|
||||
3. 优先走 `try_tai_fused_mla_store(...)`，把 logical locs 交给 TAI fused store。
|
||||
4. fallback 才通过 `get_cp_shared_kv_local_physical_out_cache_loc()` 转 physical loc 后调用 pool setter。
|
||||
|
||||
Index direct write：
|
||||
|
||||
- 文件：`python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
|
||||
- 函数：`_store_cp_shared_local_index_k_cache`
|
||||
- 流程：
|
||||
1. 取 local logical locs。
|
||||
2. 用同一个 valid-row selector 取 `local_key`。
|
||||
3. 转 physical locs。
|
||||
4. `_store_index_k_cache(... out_loc_override=physical_out_loc)`。
|
||||
|
||||
当前判断：
|
||||
|
||||
- Python 侧 valid row / local loc 长度与顺序在探针形状下是自洽的。
|
||||
- MLA TAI fused store 比 index store 更可疑：它接收 logical locs 并在 kernel 内自行解释 CP shared layout；index store 当前传的是 physical locs。
|
||||
- FP8 packed MLA buffer shape 需要重点确认。远端曾出现过 `kv_cache_dim=656` 相关问题，TAI fused MLA store 必须按真实 packed layout 写入，而不能按旧 bf16/compact 假设写。
|
||||
|
||||
## 2026-06-07 增量发现：日志噪音与非根因
|
||||
|
||||
最新日志里的 `Invalid` 计数主要来自环境变量拼接错误：
|
||||
|
||||
```text
|
||||
Invalid value for SGLANG_CP_SHARED_KV_BS_GT1_TIMING_LIMIT:
|
||||
"-1SGLANG_CP_SHARED_KV_BS_GT1_DEBUG=1" is not a valid integer value, using default "256"
|
||||
```
|
||||
|
||||
该问题会影响 timing limit/debug 配置，但不是 GSM8K 第二轮掉点的直接根因：同一进程第一轮准确率正常，第二轮 cache 命中更高后才掉点。
|
||||
|
||||
## 当前最小下一步
|
||||
|
||||
优先做一个 CUDA/remote 级验证，不继续靠日志猜：
|
||||
|
||||
1. 在 `tai-kernel` 给 `fused_store_mla_kv` 补 GSM8K-like bs=5 非连续 logical loc + FP8 packed shape 的单测。
|
||||
2. 对比：
|
||||
- TAI fused MLA store 写入 raw layer buffer；
|
||||
- reference path：`CpSharedKVLayout.logical_locs_to_physical()` + 现有 pool setter / torch reference。
|
||||
3. 如果 MLA fused store 通过，再对 index fused store 做同形状验证。
|
||||
|
||||
该验证能直接回答：第一轮写入的 suffix page 是否已经在 persistent KV cache 中损坏。
|
||||
|
||||
## 2026-06-07 新假设：第二轮掉点可能来自 L2->L1 load，而不是写入
|
||||
|
||||
用户指出一个关键区分：
|
||||
|
||||
- 第一轮的 cache-hit 可能主要来自同 prefix 的 radix/device cache，仍然正确；
|
||||
- 第二轮多命中的 question-specific suffix page 可能来自第一轮 backup 到 L2/HiCache 后再 load 回 L1；
|
||||
- 如果 L2 host layout、page_first_direct load kernel、或 load-back page-slot remap 错，现象就是第二轮掉点且不一定有显式 fallback/error。
|
||||
|
||||
这解释了为什么日志里“第一轮已有 cached-token 但正确”不能完全排除 cache 路径问题：第一轮命中的 cache 类型可能与第二轮额外命中的 cache 类型不同。
|
||||
|
||||
需要把验证拆成两类：
|
||||
|
||||
1. **写入验证**：direct write 后立刻从 L1 persistent buffer 读回，确认 first-run suffix page 写入无损。
|
||||
2. **L2 roundtrip 验证**：L1 -> L2 backup -> 释放/覆盖 L1 -> L2 -> L1 load 后再读回，确认 host page_first_direct/direct backend 的 layout 与 L1 一致。
|
||||
|
||||
当前日志曾统计 `HiCache-load=0/CacheCtrl-load=0/load_back=0`，但这只能说明没有这些字符串；不能证明没有发生 HiCache L2 load。下一步需要按实际代码路径找 load/restore 日志点或补 env-gated validator。
|
||||
|
||||
### L2 roundtrip 单测结果
|
||||
|
||||
已新增并在远端容器验证：
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages -s
|
||||
```
|
||||
|
||||
结果：
|
||||
|
||||
```text
|
||||
1 passed, 3 warnings in 10.05s
|
||||
```
|
||||
|
||||
该测试覆盖：
|
||||
|
||||
- `NSATokenToKVPool(dtype=torch.float8_e4m3fn, kv_cache_dim=656)`
|
||||
- `NSATokenToKVPoolHost(layout=page_first_direct, io_backend=direct)`
|
||||
- 多层 KV bytes：L1 -> L2 backup -> L1 清零 -> L2 -> L1 load
|
||||
- NSA indexer bytes 同步 roundtrip
|
||||
- 非连续 page 映射：`src_pages=[1,2,5,8]`, `host_pages=[3,4,7,12]`, `dst_pages=[10,11,13,15]`
|
||||
|
||||
含义：
|
||||
|
||||
- 简单的 TAI `page_first_direct/direct` byte-copy kernel 不是当前最高优先级根因。
|
||||
- 仍未排除控制路径错误：HiCache metadata 中 host_indices/device_indices 选择、node split/page-floor 后 stale tail、或 prefix materialize 对 loaded pages 的解释仍可能出错。
|
||||
|
||||
## 2026-06-07 C: bs=5 valid-row FP8 L2 roundtrip test
|
||||
|
||||
新增并远端验证：
|
||||
|
||||
- `test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_fp8_page_first_direct_bs5_valid_rows_survive_l2_roundtrip`
|
||||
- 形状：`cp_size=8,page_size=64,bs=5,prefix=[640]*5,extend=[95,81,140,66,86]`
|
||||
- 覆盖：compute padding、只写 valid rows、FP8 packed MLA KV、NSA indexer page row、`page_first_direct/direct`、L1→L2 backup、L2→L1 load。
|
||||
- 远端命令：
|
||||
`PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_fp8_page_first_direct_bs5_valid_rows_survive_l2_roundtrip -s`
|
||||
- 结果：`1 passed, 3 warnings in 9.97s`。
|
||||
|
||||
结论：
|
||||
|
||||
- 真实形状下，valid-row direct-write 之后再做 padded owned-row L2 backup/load，KV 与 index valid rows 均可逐字节保持。
|
||||
- 因此“TAI direct L2 raw copy 坏”以及“bs=5 valid-row/padded backup 基本顺序坏”目前不是最高优先级。
|
||||
- 剩余更可疑方向：
|
||||
1. 线上第二轮 GSM8K 是否真的触发了 `load_back`；此前日志没有 `[HiCache-load]`，需要用最新日志再确认。
|
||||
2. radix/HiCache node metadata 在真实 split/page-floor/evict 后的 `visible_device_indices` 是否与 page table/current-reuse consume 侧一致。
|
||||
3. cache-hit consume 侧（page table、current/partial reuse、index topk offset、MLA materialize）是否在第二轮复用时使用了错误 page/order，而不是 L2 copy 本身损坏。
|
||||
|
||||
## 2026-06-07 D: latest log still does not show L2 load_back
|
||||
|
||||
复查 `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260606_193039.log`：
|
||||
|
||||
- `[HiCache-load]`: 0
|
||||
- `load_back CP`: 0
|
||||
- `load_cp`: 0
|
||||
- `load_to_device`: 0
|
||||
- `CP_SHARED_KV_FAIL_FAST`: 0
|
||||
- `CP_SHARED_KV_FALLBACK`: 0
|
||||
- `Prefill batch`: 4384
|
||||
- `cached-token`: 4384
|
||||
- `partial`: 2048
|
||||
- `index_topk`: 5007
|
||||
|
||||
代码事实：`hiradix_cache.load_back()` 的 CP path 会打 `[HiCache-load] load_back CP...` info 日志。当前日志完全没有该前缀。
|
||||
|
||||
结论：这份 GSM8K 二轮掉点证据目前不支持“第二轮主要来自 L2 load-back 拷贝损坏”。更像是 device/radix cache-hit consume 侧，或 current/partial reuse/index/MLA materialize/page table 的复用语义问题。L2 仍不能完全排除，但优先级下降；如要继续验证，需要在新运行中显式制造 L1 evict 后再复测同一 prompts。
|
||||
|
||||
## 2026-06-07 E: actual TAI fused MLA store bs=5 test
|
||||
|
||||
新增并远端验证：
|
||||
|
||||
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_fused_mla_store_matches_fallback_for_bs5_compute_padding_rows`
|
||||
- 覆盖：`bs=5,prefix=[640]*5,extend=[95,81,140,66,86]`，compute padding 后通过 `get_cp_shared_kv_local_out_cache_loc()` 和 `select_cp_local_valid_rows_for_cache_write()` 得到真实 local valid rows，再调用实际 `runtime.try_tai_fused_mla_store()` 写 FP8 MLA KV。
|
||||
- 参考：`quantize_k_cache_separate()` + `set_mla_kv_buffer_triton()` fallback。
|
||||
- 远端结果：`1 passed, 3 warnings in 8.76s`。
|
||||
|
||||
结论：第一轮 prefill 的 TAI fused MLA direct-write 在该 bs=5/tiny/FP8 形状下与 fallback 字节一致。它不是当前最高优先级根因。
|
||||
|
||||
当前更可疑方向收敛为 cache-hit consume 侧：
|
||||
|
||||
1. page table / visible device indices 在 page-floored cache hit 下是否按 request 边界传递。
|
||||
2. index topk 的 current/partial reuse 在 second-run cache-hit 下是否使用了错误 offset/layout。
|
||||
3. MLA/index materialize shared cache 时是否用错 logical page order 或 dense output slot。
|
||||
|
||||
## 2026-06-07 F: CP=8 bs=5 index partial-current compose rank-merge 单测
|
||||
|
||||
新增并远端验证：
|
||||
|
||||
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_cp8_index_partial_current_compose_matches_rank_merged_reference_for_bs5`
|
||||
- 形状：`cp_size=8,page_size=64,bs=5,prefix=[640]*5,extend=[95,81,140,66,86]`
|
||||
- 前缀页按相同 logical page 在 5 个 request 中重复，模拟同 prefix cache hit；suffix/current 页按 owner lane 分配，模拟 CP shared KV direct-write/current-reuse。
|
||||
- 测试方式：每个 CP rank 只 materialize 自己拥有的 page，patch collective 为 identity，然后在测试中手工把 8 个 rank 的输出相加，等价模拟 all-reduce 后的 dense page buffer。
|
||||
- 远端结果：`1 passed`。
|
||||
|
||||
结论：当前 CPU/reference index partial-current compose 在 CP=8、bs=5、重复 prefix + owner-lane current suffix 的情况下，rank-merge 后等价完整 dense reference。该路径暂未暴露 index page materialize/order 错误。
|
||||
|
||||
## 2026-06-07 G: CP=8 KV partial-current remote-current visibility 单测
|
||||
|
||||
新增并远端验证：
|
||||
|
||||
- `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_cp8_kv_partial_current_keeps_remote_current_valid_locs_after_reduce`
|
||||
- 目的：确认本 rank 在 current slot all-reduce 后不会把远端 owner 的 valid current loc 错误 mask 成 `-1`。
|
||||
- 远端结果：`1 passed`。
|
||||
|
||||
结论：`build_current_page_mask()` 当前只 mask 本 rank `current_locs` 所在 page；远端 owner current page 的 valid loc 会保留，符合 all-reduce 后所有 rank 可见 current page 的语义。仍需注意：远端 owner current page 的 tail slack 也不会由本 rank mask，是否安全依赖 index/topk 侧正确使用 valid sequence length，不把 tail slack 选入 `logical_locs`。
|
||||
|
||||
## 2026-06-07 H: FP8/RAGGED MLA materialization domain mismatch suspect
|
||||
|
||||
代码审计发现一个比 L2 copy 更具体的风险：
|
||||
|
||||
- `get_topk_transform_method()` 在 FP8 + `flashmla_sparse` prefill 下返回 `RAGGED`。
|
||||
- RAGGED topk 的 `page_table_1/topk_indices` 是 ragged flattened request/KV 坐标，用于索引 `page_table_1_flattened`，不是 token_to_kv_pool 的 raw logical loc。
|
||||
- 但当前 CP shared-KV MLA generic materialize/current-reuse block 在进入 `flashmla_sparse` attention 分支前，无条件把 `page_table_1` 当 `logical_locs` 传给 `materialize_prefix_and_reuse_current_kv_page_slots()` / `materialize_shared_token_kv_buffer()`。
|
||||
- 这在 cache-hit/prefix sharing 下会把“ragged 坐标”误当 “KV logical loc”，非常符合：第一轮部分 cache-hit 仍正常、第二轮更高 prefix/suffix cache-hit 后精度掉点。
|
||||
|
||||
下一步按 TDD 补 source/behavior guard 单测：FP8 RAGGED 路径必须绕过 generic CP MLA materialize，交给后面的 RAGGED `page_table_1_flattened` materialize/dequantize 路径处理；后续再做基于 flattened locs 的 partial-current 优化。
|
||||
|
||||
## 2026-06-07 I: RAGGED generic MLA materialize guard 已存在
|
||||
|
||||
继续审计后确认当前本地代码已经有 guard：
|
||||
|
||||
```python
|
||||
if (
|
||||
forward_batch.uses_cp_shared_kv
|
||||
and topk_transform_method == TopkTransformMethod.PAGED
|
||||
):
|
||||
```
|
||||
|
||||
因此 FP8/`flashmla_sparse` 的 RAGGED topk 不会进入 generic CP MLA materialize/current-reuse block，避免了把 ragged flattened topk 坐标误当 raw KV logical loc。已补回归测试锁住这个 guard。
|
||||
|
||||
结论：H 是合理风险，但在当前版本已被代码规避；它不是这次第二轮 GSM8K 掉点的剩余根因。继续排查应转向 RAGGED 分支内部的 `page_table_1_flattened` materialize/dequantize、radix visible loc、或 cache-hit 后的 req_to_token/page_table 内容。
|
||||
|
||||
## 2026-06-07 J: 用户提出的 L2 区分需要更高层单测
|
||||
|
||||
用户指出：第一轮的 cache-hit 可能主要是同 prefix/device radix cache，第二轮额外 cache-hit 才可能来自 L2/HiCache load。因此仅有 raw L1↔L2 byte-copy roundtrip 还不够，需要补一个更贴近真实 consume 的测试：
|
||||
|
||||
1. first-run 形状写入 FP8 packed MLA KV + NSA index valid rows；
|
||||
2. 备份 full padded owned pages 到 page_first_direct/direct host；
|
||||
3. 清空/污染 L1 后从 host load 到新的 L1 pages；
|
||||
4. 用 CP shared-KV `materialize_shared_token_kv_buffer()` 按 RAGGED `page_table_1_flattened` 风格 logical locs 读取 loaded pages；
|
||||
5. 对比 materialize 后 dense loc 对应 bytes 与 first-run 写入 bytes。
|
||||
|
||||
目的：直接验证“L2 load 本身正确，但 load 后被 RAGGED/MLA materialize 或 logical→physical remap 错读”的可能性。若该测试通过，L2 load 作为 GSM8K 二轮掉点根因的优先级继续降低，后续应重点看真实运行中的 page_table/topk offset/current reuse 控制流。
|
||||
|
||||
### L2-loaded suffix + RAGGED materialize 单测结果
|
||||
|
||||
新增并在远端验证：
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_fp8_l2_loaded_bs5_suffix_materializes_from_ragged_logical_locs -s
|
||||
```
|
||||
|
||||
结果：`1 passed, 3 warnings in 9.82s`。
|
||||
|
||||
覆盖内容：
|
||||
|
||||
- GSM8K-like `bs=5,prefix=[640]*5,extend=[95,81,140,66,86]`；
|
||||
- CP=8 owner-lane suffix valid rows；
|
||||
- FP8 packed MLA KV (`kv_cache_dim=656`)；
|
||||
- `page_first_direct/direct` L1→L2 backup，再 L2→L1 load 到新的 logical pages；
|
||||
- 之后按 RAGGED `page_table_1_flattened` 风格的 token logical locs 调 `materialize_shared_token_kv_buffer()`；
|
||||
- 手工合并 8 个 rank 的 local materialize 输出，验证 dense loc 对应 bytes 与 first-run 写入 bytes 一致。
|
||||
|
||||
结论：在该单测覆盖的控制面里，L2-loaded suffix page 被 RAGGED MLA materialize 读取是逐字节正确的。当前 GSM8K 二轮掉点更不像 raw L2 copy 或 basic RAGGED logical-loc materialize 错；下一步应看真实运行中的 `page_table_1_flattened/topk_indices_offset` 是否在 bs>1/cache-hit 后按 request 正确构造，以及 decode transfer/FP8 dtype 协议是否一致。
|
||||
|
||||
## 2026-06-07 K: RAGGED fused topk offset 语义验证
|
||||
|
||||
新增并在远端验证：
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/layers/test_nsa_topk_transform.py::TestNSATopkTransform::test_ragged_topk_transform_offsets_are_request_relative_with_row_starts -s
|
||||
```
|
||||
|
||||
结果：`1 passed, 5 warnings in 9.77s`。
|
||||
|
||||
验证点：`fast_topk_transform_ragged_fused` 在传入 `row_starts` 时，会把选中的 compact column 转成：
|
||||
|
||||
```text
|
||||
request_offset + (selected_compact_col - row_start)
|
||||
```
|
||||
|
||||
而不是：
|
||||
|
||||
```text
|
||||
request_offset + selected_compact_col
|
||||
```
|
||||
|
||||
结论：当前 `_get_topk_ragged_with_cp()` 给 batch compact path 传 `topk_indices_offset_override=request_kv_base` 的方向是对的；不需要改成 `request_kv_base - k_base`。这降低了“RAGGED topk offset 基本语义错误”作为二轮 GSM8K 掉点根因的优先级。
|
||||
|
||||
## 2026-06-07 L: TAI batched index MQA prepare GSM8K-like 形状验证
|
||||
|
||||
新增并在远端验证：
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_tai_batched_index_mqa_prepare_matches_getk_gets_reference_gsm8k_bs5 -s
|
||||
```
|
||||
|
||||
结果：`1 passed, 3 warnings in 10.00s`。
|
||||
|
||||
覆盖内容：
|
||||
|
||||
- `bs=5`, 每请求 `seq_len=704`, `q_start=640`, `q_len=64`；
|
||||
- `page_size=64`, `index_head_dim=128`；
|
||||
- 实际调用 `try_tai_prepare_cp_mqa_index_batch()`；
|
||||
- 与 reference 路径逐项对比：`index_buf_accessor.GetK/GetS`、`ks`、`ke_offset`。
|
||||
|
||||
结论：至少在 GSM8K-like cache-hit/tiny-extend 形状下，TAI batched index prepare kernel 的 K/S compact 与 range descriptor 与 Python reference 一致。当前二轮掉点不应优先归因于该 kernel 的基本 gather/descriptor 错误。
|
||||
|
||||
## 2026-06-07 M: L2 与 persistent index cache 需要区分验证
|
||||
|
||||
用户指出第一轮 GSM8K 里的 cache hit 可能主要来自相同 prefix 的 L1/device radix cache，而第二轮新增的更高命中率可能来自 L2/HiCache load；因此不能只用第一轮正常来排除 L2/load 路径。
|
||||
|
||||
当前已覆盖并通过的单测降低了以下路径优先级：raw page_first_direct L1↔L2 copy、L2-loaded FP8 MLA KV 经 RAGGED logical loc materialize、RAGGED topk offset、TAI batched index MQA prepare。下一步需要补更贴近第二轮 consume 的 persistent NSA index cache 写入/读取单测，尤其是 `fused_store_index_k_cache` 在 bs=5/tiny/FP8/physical loc direct-write 下是否与 fallback `act_quant + SetKAndS` 等价。
|
||||
|
||||
原因：第一轮 topk 可以使用 freshly computed `current_index_kv/local_key`，即使 persistent index cache 写坏也不一定立即影响；第二轮高 cache-hit 时会更依赖持久化 index cache，这与“第一轮 0.95、第二轮 0.69”的症状吻合。
|
||||
|
||||
### M 追加验证结果：persistent index fused store 基本路径通过
|
||||
|
||||
新增并在远端验证：
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_index_fused_store_persistent_pages_survive_bs5_materialize -s
|
||||
```
|
||||
|
||||
结果：`1 passed, 3 warnings in 10.22s`。
|
||||
|
||||
覆盖内容：bs=5/tiny extend、CP=8、compute padding、`get_cp_shared_kv_local_out_cache_loc()`、`select_cp_local_valid_rows_for_cache_write()`、`logical_locs_to_physical()`、`fused_store_index_k_cache()`、`materialize_shared_paged_buffer()` rank merge、`GetK/GetS` consume。
|
||||
|
||||
结论：不经过 L2 的 persistent NSA index fused-store + paged materialize 基本路径未复现行顺序/page offset 错误。单测仍保留，因为它能锁住第一轮写 persistent index cache 的核心合同。
|
||||
|
||||
## 2026-06-07 N: L2-loaded persistent index cache 单测
|
||||
|
||||
新增并在远端验证：
|
||||
|
||||
```bash
|
||||
PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py::TestNSAHiCacheTransfer::test_fp8_l2_loaded_index_pages_materialize_after_fused_store_bs5 -s
|
||||
```
|
||||
|
||||
结果：`1 passed, 3 warnings in 10.33s`。
|
||||
|
||||
覆盖内容：
|
||||
|
||||
- 第一轮形状用 `fused_store_index_k_cache()` 写入 persistent index pages；
|
||||
- `page_first_direct/direct` 备份到 L2 host；
|
||||
- 污染目标 L1 pages 后从 L2 load 到新 L1 pages；
|
||||
- 对 loaded index pages 走 `materialize_shared_paged_buffer()` + rank merge；
|
||||
- 通过 `GetK/GetS` 按 request page table 读取并反量化，对比原始 valid key rows。
|
||||
|
||||
结论：在单测覆盖的 GSM8K-like bs=5/tiny/FP8 形状中，persistent index cache 经 fused store → L2 backup/load → paged materialize → GetK/GetS 后仍可还原 valid rows。用户提出的“第二轮来自 L2 cache load 的错误”目前没有被 KV 或 index 的直接单测支持；下一步应把排查焦点上移到真实运行中的 radix/req_to_token/page_table 可见范围、chunked/full prompt cache-hit 后 valid length 与 page table 的组合、以及 prefill→decode transfer 元数据。
|
||||
@@ -524,6 +524,8 @@ class PrefillBootstrapQueue:
|
||||
[req.disagg_kv_sender for req in self.queue],
|
||||
self.scheduler.attn_cp_cpu_group,
|
||||
self.scheduler.attn_tp_cpu_group,
|
||||
debug_label="bootstrap",
|
||||
debug_ids=[req.rid for req in self.queue],
|
||||
)
|
||||
_cp_shared_kv_bs_gt1_prefill_timing(
|
||||
"bootstrap_poll_done",
|
||||
@@ -937,6 +939,8 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
[req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
|
||||
self.attn_cp_cpu_group,
|
||||
self.attn_tp_cpu_group,
|
||||
debug_label="inflight",
|
||||
debug_ids=[req.rid for req in self.disagg_prefill_inflight_queue],
|
||||
)
|
||||
_cp_shared_kv_bs_gt1_prefill_timing(
|
||||
"inflight_poll_done",
|
||||
@@ -1050,6 +1054,8 @@ class SchedulerDisaggregationPrefillMixin:
|
||||
[req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
|
||||
self.attn_cp_cpu_group,
|
||||
self.attn_tp_cpu_group,
|
||||
debug_label="get_transferred_rids",
|
||||
debug_ids=[req.rid for req in self.disagg_prefill_inflight_queue],
|
||||
)
|
||||
|
||||
transferred_rids: List[str] = []
|
||||
|
||||
@@ -63,11 +63,85 @@ def poll_and_all_reduce(pollers, gloo_group: dist.ProcessGroup):
|
||||
return tensor_to_reduce.tolist()
|
||||
|
||||
|
||||
def _cp_shared_kv_poll_debug_enabled() -> bool:
|
||||
return envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get()
|
||||
|
||||
|
||||
def _poll_queue_debug_hash(debug_ids: Optional[list[str]]) -> int:
|
||||
if not debug_ids:
|
||||
return 0
|
||||
|
||||
# Stable bounded FNV-1a style digest. Keep it in signed-int64 range because
|
||||
# the debug consensus uses torch.int64 CPU collectives.
|
||||
value = 1469598103934665603
|
||||
mask = (1 << 63) - 1
|
||||
for item in debug_ids:
|
||||
for byte in str(item).encode("utf-8", errors="replace"):
|
||||
value ^= byte
|
||||
value = (value * 1099511628211) & mask
|
||||
value ^= 0xFF
|
||||
value = (value * 1099511628211) & mask
|
||||
return int(value)
|
||||
|
||||
|
||||
def _validate_poll_queue_consensus(
|
||||
*,
|
||||
label: str,
|
||||
scope: str,
|
||||
local_len: int,
|
||||
debug_ids: Optional[list[str]],
|
||||
group: dist.ProcessGroup,
|
||||
) -> None:
|
||||
local_hash = _poll_queue_debug_hash(debug_ids)
|
||||
signature = torch.tensor([local_len, local_hash], dtype=torch.int64, device="cpu")
|
||||
min_signature = signature.clone()
|
||||
max_signature = signature.clone()
|
||||
|
||||
dist.all_reduce(min_signature, op=dist.ReduceOp.MIN, group=group)
|
||||
dist.all_reduce(max_signature, op=dist.ReduceOp.MAX, group=group)
|
||||
|
||||
if torch.equal(min_signature, max_signature):
|
||||
return
|
||||
|
||||
ids_head = list(debug_ids[:8]) if debug_ids is not None else None
|
||||
message = (
|
||||
"[CP_SHARED_KV_FAIL_FAST][poll_queue] "
|
||||
f"label={label} scope={scope} local_len={local_len} "
|
||||
f"min_len={int(min_signature[0].item())} "
|
||||
f"max_len={int(max_signature[0].item())} "
|
||||
f"local_hash={local_hash} min_hash={int(min_signature[1].item())} "
|
||||
f"max_hash={int(max_signature[1].item())} ids_head={ids_head}"
|
||||
)
|
||||
logger.error(message)
|
||||
raise RuntimeError(message)
|
||||
|
||||
|
||||
def poll_and_all_reduce_attn_cp_tp_group(
|
||||
pollers,
|
||||
attn_cp_cpu_group: dist.ProcessGroup,
|
||||
attn_tp_cpu_group: dist.ProcessGroup,
|
||||
*,
|
||||
debug_label: Optional[str] = None,
|
||||
debug_ids: Optional[list[str]] = None,
|
||||
):
|
||||
if _cp_shared_kv_poll_debug_enabled():
|
||||
label = debug_label or "unknown"
|
||||
local_len = len(pollers)
|
||||
_validate_poll_queue_consensus(
|
||||
label=label,
|
||||
scope="attn_tp",
|
||||
local_len=local_len,
|
||||
debug_ids=debug_ids,
|
||||
group=attn_tp_cpu_group,
|
||||
)
|
||||
_validate_poll_queue_consensus(
|
||||
label=label,
|
||||
scope="attn_cp",
|
||||
local_len=local_len,
|
||||
debug_ids=debug_ids,
|
||||
group=attn_cp_cpu_group,
|
||||
)
|
||||
|
||||
# First sync across attn-tp ranks so all TP participants for a given (dp, cp)
|
||||
# shard observe the same status transitions.
|
||||
polls = poll_and_all_reduce(pollers, attn_tp_cpu_group)
|
||||
|
||||
@@ -228,6 +228,128 @@ def _slot_remap_cache_key(
|
||||
)
|
||||
|
||||
|
||||
def _to_cpu_int_list(values: Any, *, name: str) -> list[int]:
|
||||
if values is None:
|
||||
raise ValueError(f"{name} must not be None")
|
||||
if isinstance(values, torch.Tensor):
|
||||
if values.dim() == 0:
|
||||
result = [int(values.item())]
|
||||
else:
|
||||
result = [int(x) for x in values.detach().cpu().reshape(-1).tolist()]
|
||||
else:
|
||||
result = [int(x) for x in values]
|
||||
for idx, value in enumerate(result):
|
||||
if value < 0:
|
||||
raise ValueError(f"{name} contains negative value: idx={idx} value={value}")
|
||||
return result
|
||||
|
||||
|
||||
def _normalize_seq_lens_for_logical_page_rows(
|
||||
*,
|
||||
seq_lens_cpu: Any,
|
||||
rows: int,
|
||||
repeat_lens_cpu: Any | None = None,
|
||||
) -> list[int]:
|
||||
seq_lens = _to_cpu_int_list(seq_lens_cpu, name="seq_lens_cpu")
|
||||
if rows < 0:
|
||||
raise ValueError(f"rows must be non-negative, got {rows}")
|
||||
if rows == 0:
|
||||
if seq_lens:
|
||||
raise ValueError(
|
||||
"logical page table has no rows but seq_lens_cpu is non-empty: "
|
||||
f"seq_lens={seq_lens}"
|
||||
)
|
||||
return []
|
||||
if len(seq_lens) == rows:
|
||||
return seq_lens
|
||||
|
||||
repeat_lens = (
|
||||
_to_cpu_int_list(repeat_lens_cpu, name="repeat_lens_cpu")
|
||||
if repeat_lens_cpu is not None
|
||||
else None
|
||||
)
|
||||
if (
|
||||
repeat_lens is not None
|
||||
and len(repeat_lens) == len(seq_lens)
|
||||
and sum(repeat_lens) == rows
|
||||
):
|
||||
expanded: list[int] = []
|
||||
for seq_len, repeat in zip(seq_lens, repeat_lens, strict=True):
|
||||
expanded.extend([seq_len] * int(repeat))
|
||||
return expanded
|
||||
|
||||
if len(seq_lens) > 0 and rows % len(seq_lens) == 0:
|
||||
repeat = rows // len(seq_lens)
|
||||
expanded = []
|
||||
for seq_len in seq_lens:
|
||||
expanded.extend([seq_len] * repeat)
|
||||
return expanded
|
||||
|
||||
raise ValueError(
|
||||
"Cannot align seq_lens_cpu with logical page rows: "
|
||||
f"rows={rows} seq_lens={seq_lens} repeat_lens={repeat_lens}"
|
||||
)
|
||||
|
||||
|
||||
def mask_batch_logical_pages_to_valid_lengths(
|
||||
logical_pages: torch.Tensor,
|
||||
*,
|
||||
seq_lens_cpu: Any,
|
||||
page_size: int,
|
||||
repeat_lens_cpu: Any | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Zero rectangular page-table tails beyond each request's valid pages.
|
||||
|
||||
CP shared-KV slot remap gives every flattened page-table slot a deterministic
|
||||
dense page id. Therefore stale values in the rectangular tail are not just
|
||||
unused padding: duplicate stale pages can overwrite the logical-page inverse
|
||||
and redirect valid top-k/cache-hit locations to the wrong dense slot. Mask
|
||||
those tails at the metadata boundary while keeping page-level padding inside
|
||||
the final valid page visible.
|
||||
"""
|
||||
|
||||
if page_size <= 0:
|
||||
raise ValueError(f"page_size must be positive, got {page_size}")
|
||||
if logical_pages.dim() == 0:
|
||||
raise ValueError("logical_pages must have at least one dimension")
|
||||
|
||||
if logical_pages.dim() == 1:
|
||||
rows = 1
|
||||
pages_per_request = int(logical_pages.numel())
|
||||
else:
|
||||
rows = int(logical_pages.shape[0])
|
||||
pages_per_request = int(logical_pages.reshape(rows, -1).shape[1])
|
||||
|
||||
seq_lens = _normalize_seq_lens_for_logical_page_rows(
|
||||
seq_lens_cpu=seq_lens_cpu,
|
||||
rows=rows,
|
||||
repeat_lens_cpu=repeat_lens_cpu,
|
||||
)
|
||||
if rows == 0 or pages_per_request == 0:
|
||||
return logical_pages.clone()
|
||||
|
||||
valid_pages_cpu = [
|
||||
(int(seq_len) + int(page_size) - 1) // int(page_size)
|
||||
for seq_len in seq_lens
|
||||
]
|
||||
valid_pages = torch.tensor(
|
||||
valid_pages_cpu,
|
||||
device=logical_pages.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
columns = torch.arange(
|
||||
pages_per_request,
|
||||
device=logical_pages.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
keep_mask = columns.unsqueeze(0) < valid_pages.unsqueeze(1)
|
||||
|
||||
masked = logical_pages.clone()
|
||||
masked_2d = masked.reshape(rows, pages_per_request)
|
||||
masked_2d.masked_fill_(~keep_mask, 0)
|
||||
return masked
|
||||
|
||||
|
||||
def _log_slot_remap_cache_not_reused(
|
||||
*,
|
||||
kind: str,
|
||||
|
||||
@@ -31,6 +31,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
get_or_build_shared_token_kv_slot_remap,
|
||||
is_current_only_extend_batch,
|
||||
is_packed_fp8_mla_kv_cache,
|
||||
mask_batch_logical_pages_to_valid_lengths,
|
||||
materialize_prefix_and_reuse_current_kv_page_slots,
|
||||
materialize_shared_token_kv_buffer,
|
||||
pack_current_mla_kv_for_reuse,
|
||||
@@ -932,6 +933,21 @@ class NativeSparseAttnBackend(
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
paged_mqa_schedule_metadata = None
|
||||
|
||||
real_page_table = self._transform_table_1_to_real(page_table)
|
||||
real_page_valid_seq_lens_cpu = indexer_seq_lens_cpu
|
||||
if draft_token_num:
|
||||
real_page_valid_seq_lens_cpu = (
|
||||
indexer_seq_lens_cpu + draft_token_num
|
||||
if isinstance(indexer_seq_lens_cpu, torch.Tensor)
|
||||
else [int(x) + int(draft_token_num) for x in indexer_seq_lens_cpu]
|
||||
)
|
||||
real_page_table = mask_batch_logical_pages_to_valid_lengths(
|
||||
real_page_table,
|
||||
seq_lens_cpu=real_page_valid_seq_lens_cpu,
|
||||
page_size=self.real_page_size,
|
||||
repeat_lens_cpu=extend_seq_lens_cpu,
|
||||
)
|
||||
|
||||
metadata = NSAMetadata(
|
||||
page_size=self.real_page_size,
|
||||
cache_seqlens_int32=cache_seqlens_int32,
|
||||
@@ -957,7 +973,7 @@ class NativeSparseAttnBackend(
|
||||
nsa_cu_seqlens_k=nsa_cu_seqlens_k,
|
||||
nsa_seqlens_expanded=seqlens_expanded,
|
||||
nsa_extend_seq_lens_list=extend_seq_lens_cpu,
|
||||
real_page_table=self._transform_table_1_to_real(page_table),
|
||||
real_page_table=real_page_table,
|
||||
nsa_max_seqlen_q=1,
|
||||
topk_indices_offset=topk_indices_offset,
|
||||
indexer_k_start_end=indexer_k_start_end,
|
||||
|
||||
@@ -4,8 +4,12 @@ import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.srt.disaggregation.base import KVPoll
|
||||
from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue
|
||||
from sglang.srt.disaggregation.utils import poll_and_all_reduce_attn_cp_tp_group
|
||||
from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
@@ -25,6 +29,9 @@ class FakeSender:
|
||||
if self.should_fail:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
def poll(self):
|
||||
return KVPoll.Success
|
||||
|
||||
|
||||
class TestPrefillBootstrapQueue(CustomTestCase):
|
||||
def _make_req(self, rid, bootstrap_room, origin_input_ids, sender):
|
||||
@@ -124,6 +131,70 @@ class TestPrefillBootstrapQueue(CustomTestCase):
|
||||
self.assertEqual(skipped.disagg_kv_sender.init_calls, [])
|
||||
self.assertEqual(checked.disagg_kv_sender.init_calls, [(3, 0)])
|
||||
|
||||
def test_poll_consensus_debug_fails_before_shape_mismatch_hang(self):
|
||||
reduce_calls = []
|
||||
|
||||
def fake_all_reduce(tensor, op, group):
|
||||
reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group))
|
||||
# The new debug guard uses an int64 [queue_len, queue_hash] scalar
|
||||
# vector before building the uint8 per-request poll tensor. Simulate
|
||||
# another rank having a different queue length and assert that we
|
||||
# fail fast before reaching the old variable-length uint8 all_reduce.
|
||||
if tensor.dtype == torch.int64 and int(tensor.numel()) == 2:
|
||||
if op == dist.ReduceOp.MIN:
|
||||
tensor[0] = 1
|
||||
elif op == dist.ReduceOp.MAX:
|
||||
tensor[0] = 2
|
||||
return
|
||||
raise AssertionError(
|
||||
"poll tensor all_reduce should not run after queue mismatch"
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get",
|
||||
return_value=True,
|
||||
),
|
||||
patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce),
|
||||
):
|
||||
with self.assertRaisesRegex(RuntimeError, "poll_queue.*inflight"):
|
||||
poll_and_all_reduce_attn_cp_tp_group(
|
||||
[FakeSender(), FakeSender()],
|
||||
MagicMock(name="cp_group"),
|
||||
MagicMock(name="tp_group"),
|
||||
debug_label="inflight",
|
||||
debug_ids=["rid-a", "rid-b"],
|
||||
)
|
||||
|
||||
self.assertTrue(reduce_calls)
|
||||
|
||||
def test_poll_consensus_debug_disabled_preserves_old_collective_shape(self):
|
||||
reduce_calls = []
|
||||
|
||||
def fake_all_reduce(tensor, op, group):
|
||||
reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get",
|
||||
return_value=False,
|
||||
),
|
||||
patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce),
|
||||
):
|
||||
polls = poll_and_all_reduce_attn_cp_tp_group(
|
||||
[FakeSender(), FakeSender()],
|
||||
MagicMock(name="cp_group"),
|
||||
MagicMock(name="tp_group"),
|
||||
debug_label="inflight",
|
||||
debug_ids=["rid-a", "rid-b"],
|
||||
)
|
||||
|
||||
self.assertEqual(polls, [KVPoll.Success, KVPoll.Success])
|
||||
self.assertEqual(
|
||||
[(dtype, size) for dtype, size, _op, _group in reduce_calls],
|
||||
[(torch.uint8, 2), (torch.uint8, 2)],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -14,6 +14,37 @@ register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestNSATopkTransform(unittest.TestCase):
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
|
||||
def test_ragged_topk_transform_offsets_are_request_relative_with_row_starts(self):
|
||||
from sgl_kernel import fast_topk_transform_ragged_fused
|
||||
|
||||
# CP bs>1 compacts each request segment into a temporary K buffer.
|
||||
# The selected compact column must be translated back to the normal
|
||||
# request-concatenated ragged layout as:
|
||||
# request_base + (selected_compact_col - row_start)
|
||||
# not request_base + selected_compact_col.
|
||||
columns = 5000
|
||||
logits = torch.zeros((2, columns), device="cuda", dtype=torch.float32)
|
||||
logits[0, 2999] = 10.0
|
||||
logits[1, 1000 + 2999] = 10.0
|
||||
lengths = torch.tensor([3000, 3000], device="cuda", dtype=torch.int32)
|
||||
row_starts = torch.tensor([0, 1000], device="cuda", dtype=torch.int32)
|
||||
request_offsets = torch.tensor(
|
||||
[0, 100000], device="cuda", dtype=torch.int32
|
||||
)
|
||||
|
||||
out = fast_topk_transform_ragged_fused(
|
||||
score=logits,
|
||||
lengths=lengths,
|
||||
topk_indices_offset=request_offsets,
|
||||
topk=2048,
|
||||
row_starts=row_starts,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self.assertEqual(int(out[0, 0].item()), 2999)
|
||||
self.assertEqual(int(out[1, 0].item()), 100000 + 2999)
|
||||
|
||||
def test_paged_topk_transform_raises_when_fused_output_is_not_from_page_table(self):
|
||||
page_table = torch.tensor(
|
||||
[
|
||||
|
||||
@@ -1316,6 +1316,31 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
self.assertIsNone(result)
|
||||
|
||||
def test_fp8_ragged_mla_defers_cp_materialize_to_flattened_path(self):
|
||||
from pathlib import Path
|
||||
|
||||
source = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "python/sglang/srt/layers/attention/nsa_backend.py"
|
||||
).read_text()
|
||||
paged_gate = (
|
||||
"if (\n"
|
||||
" forward_batch.uses_cp_shared_kv\n"
|
||||
" and topk_transform_method == TopkTransformMethod.PAGED\n"
|
||||
" ):"
|
||||
)
|
||||
self.assertIn(
|
||||
paged_gate,
|
||||
source,
|
||||
"RAGGED topk indices are flattened request/KV coordinates, not raw KV logical locs; "
|
||||
"generic CP MLA materialize must stay gated to PAGED topk.",
|
||||
)
|
||||
ragged_start = source.index(" if topk_transform_method == TopkTransformMethod.RAGGED:")
|
||||
ragged_end = source.index(" attn_output = self._forward_flashmla_sparse", ragged_start)
|
||||
ragged_source = source[ragged_start:ragged_end]
|
||||
self.assertIn("page_table_1_flattened", ragged_source)
|
||||
self.assertIn("materialize_shared_token_kv_buffer", ragged_source)
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
|
||||
def test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel(
|
||||
self,
|
||||
@@ -1474,6 +1499,47 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
[(0, 2), (4, 5)],
|
||||
)
|
||||
|
||||
def test_valid_page_mask_prevents_stale_rectangular_tail_remap(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
|
||||
|
||||
page_size = 4
|
||||
layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0)
|
||||
kv_cache = torch.zeros((64, 1, 1), dtype=torch.float32)
|
||||
# Row 0 is valid for 3 pages. The 4th column simulates stale req_to_token
|
||||
# rectangular tail that duplicates logical page 1. Without masking,
|
||||
# build_slot_page_inverse scatter_ maps page 1 to that stale slot.
|
||||
raw_logical_pages = torch.tensor(
|
||||
[
|
||||
[1, 2, 5, 1],
|
||||
[9, 10, 0, 0],
|
||||
],
|
||||
dtype=torch.int64,
|
||||
)
|
||||
|
||||
valid_logical_pages = runtime.mask_batch_logical_pages_to_valid_lengths(
|
||||
raw_logical_pages,
|
||||
seq_lens_cpu=[12, 8],
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
self.assertEqual(valid_logical_pages.tolist(), [[1, 2, 5, 0], [9, 10, 0, 0]])
|
||||
|
||||
slot_remap = runtime.build_shared_token_kv_slot_remap(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=torch.tensor([4], dtype=torch.int64),
|
||||
remap_logical_pages=valid_logical_pages,
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
dense_loc = runtime.remap_logical_locs_to_slot_dense_locs(
|
||||
torch.tensor([4], dtype=torch.int64),
|
||||
page_inverse=slot_remap.page_inverse,
|
||||
page_size=page_size,
|
||||
)
|
||||
self.assertEqual(dense_loc.tolist(), [4])
|
||||
|
||||
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
|
||||
@@ -3157,6 +3223,747 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
rtol=0,
|
||||
)
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
|
||||
def test_fp8_fused_mla_store_matches_fallback_for_bs5_compute_padding_rows(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.layers.attention.nsa.quant_k_cache import (
|
||||
quantize_k_cache_separate,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.utils import (
|
||||
NSAContextParallelMetadata,
|
||||
build_batch_page_aligned_in_seq_split_plan,
|
||||
get_cp_shared_kv_local_out_cache_loc,
|
||||
select_cp_local_valid_rows_for_cache_write,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
|
||||
build_in_seq_page_compute_owners,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton
|
||||
|
||||
try:
|
||||
runtime._load_tai_fused_mla_store_kernel()
|
||||
except Exception as exc:
|
||||
self.skipTest(f"TAI fused MLA store kernel unavailable: {exc}")
|
||||
|
||||
page_size = 64
|
||||
cp_size = 8
|
||||
prefix_lens = [640, 640, 640, 640, 640]
|
||||
extend_lens = [95, 81, 140, 66, 86]
|
||||
|
||||
def build_valid_locs():
|
||||
chunks = []
|
||||
logical_page_cursor = 0
|
||||
for req_id, (extend_len, prefix_len) in enumerate(
|
||||
zip(extend_lens, prefix_lens)
|
||||
):
|
||||
owners = build_in_seq_page_compute_owners(
|
||||
extend_len=int(extend_len),
|
||||
extend_prefix_len=int(prefix_len),
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
)
|
||||
self.assertIsNotNone(owners)
|
||||
remaining = int(extend_len)
|
||||
req_chunks = []
|
||||
for owner in owners:
|
||||
logical_page = (
|
||||
int(owner)
|
||||
+ 1
|
||||
+ cp_size * (req_id * (len(owners) + 4) + logical_page_cursor)
|
||||
)
|
||||
page_locs = torch.arange(
|
||||
logical_page * page_size,
|
||||
(logical_page + 1) * page_size,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
valid_rows = min(page_size, max(remaining, 0))
|
||||
if valid_rows > 0:
|
||||
req_chunks.append(page_locs[:valid_rows])
|
||||
remaining -= valid_rows
|
||||
logical_page_cursor += 1
|
||||
self.assertEqual(remaining, 0)
|
||||
chunks.append(torch.cat(req_chunks, dim=0))
|
||||
return torch.cat(chunks, dim=0)
|
||||
|
||||
out_cache_loc = build_valid_locs().to(device="cuda")
|
||||
|
||||
for rank in range(cp_size):
|
||||
layout = CpSharedKVLayout(
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=rank,
|
||||
)
|
||||
plan = build_batch_page_aligned_in_seq_split_plan(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=rank,
|
||||
)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=layout,
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=len(extend_lens),
|
||||
batch_plan=plan,
|
||||
page_aligned=True,
|
||||
page_size=page_size,
|
||||
extend_prefix_len=prefix_lens[0],
|
||||
),
|
||||
out_cache_loc=out_cache_loc.clone(),
|
||||
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
||||
return_value=False,
|
||||
):
|
||||
logical_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
||||
if logical_locs.numel() == 0:
|
||||
continue
|
||||
|
||||
local_compute_rows = sum(
|
||||
int(x) for x in plan.request_compute_rank_local_tokens
|
||||
)
|
||||
torch.manual_seed(20260607 + rank)
|
||||
k_nope = (
|
||||
torch.randn(
|
||||
(local_compute_rows, 1, 512),
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 2.0
|
||||
) + 0.25
|
||||
latent_cache = torch.randn(
|
||||
(local_compute_rows, 576), device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
k_rope = latent_cache[:, 512:].unsqueeze(1)
|
||||
self.assertFalse(k_rope.is_contiguous())
|
||||
valid_k_nope = select_cp_local_valid_rows_for_cache_write(
|
||||
forward_batch, k_nope
|
||||
)
|
||||
valid_k_rope = select_cp_local_valid_rows_for_cache_write(
|
||||
forward_batch, k_rope
|
||||
)
|
||||
self.assertEqual(int(valid_k_nope.shape[0]), int(logical_locs.numel()))
|
||||
|
||||
physical_locs = layout.logical_locs_to_physical(logical_locs)
|
||||
capacity_tokens = int(physical_locs.max().item()) + page_size + 1
|
||||
expected = torch.zeros(
|
||||
(capacity_tokens, 1, 656), dtype=torch.uint8, device="cuda"
|
||||
)
|
||||
nope_part, rope_part = quantize_k_cache_separate(
|
||||
valid_k_nope, valid_k_rope
|
||||
)
|
||||
set_mla_kv_buffer_triton(expected, physical_locs, nope_part, rope_part)
|
||||
|
||||
class FakePool:
|
||||
def __init__(self):
|
||||
self.nsa_kv_cache_store_fp8 = True
|
||||
self.page_size = page_size
|
||||
self.start_layer = 0
|
||||
self.kv_buffer = [
|
||||
torch.zeros(
|
||||
(capacity_tokens, 1, 656),
|
||||
dtype=torch.uint8,
|
||||
device="cuda",
|
||||
)
|
||||
]
|
||||
|
||||
pool = FakePool()
|
||||
layer = SimpleNamespace(layer_id=0)
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=True
|
||||
), patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=False):
|
||||
used = runtime.try_tai_fused_mla_store(
|
||||
token_to_kv_pool=pool,
|
||||
layer=layer,
|
||||
layout=layout,
|
||||
logical_locs=logical_locs,
|
||||
k_nope=valid_k_nope,
|
||||
k_rope=valid_k_rope,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
self.assertTrue(used)
|
||||
torch.testing.assert_close(pool.kv_buffer[0], expected, atol=0, rtol=0)
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
|
||||
def test_fp8_index_fused_store_persistent_pages_survive_bs5_materialize(self):
|
||||
from sglang.jit_kernel.fused_store_index_cache import (
|
||||
can_use_nsa_fused_store,
|
||||
fused_store_index_k_cache,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa import (
|
||||
cp_shared_kv_runtime as runtime,
|
||||
index_buf_accessor,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
|
||||
from sglang.srt.layers.attention.nsa.utils import (
|
||||
NSAContextParallelMetadata,
|
||||
build_batch_page_aligned_in_seq_split_plan,
|
||||
get_cp_shared_kv_local_out_cache_loc,
|
||||
select_cp_local_valid_rows_for_cache_write,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
|
||||
build_in_seq_page_compute_owners,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
if not hasattr(torch, "float8_e4m3fn"):
|
||||
self.skipTest("torch.float8_e4m3fn unavailable")
|
||||
if not can_use_nsa_fused_store(torch.bfloat16, torch.int64, 64):
|
||||
self.skipTest("NSA fused index store JIT unavailable")
|
||||
|
||||
page_size = 64
|
||||
cp_size = 8
|
||||
batch_size = 5
|
||||
index_head_dim = 128
|
||||
page_bytes = page_size * (index_head_dim + 4)
|
||||
prefix_lens = [640, 640, 640, 640, 640]
|
||||
extend_lens = [95, 81, 140, 66, 86]
|
||||
|
||||
request_pages: list[list[int]] = []
|
||||
request_valid_locs: list[torch.Tensor] = []
|
||||
all_valid_locs: list[torch.Tensor] = []
|
||||
next_generation = 32
|
||||
for req_id, (prefix_len, extend_len) in enumerate(
|
||||
zip(prefix_lens, extend_lens)
|
||||
):
|
||||
owners = build_in_seq_page_compute_owners(
|
||||
extend_len=extend_len,
|
||||
extend_prefix_len=prefix_len,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
)
|
||||
self.assertIsNotNone(owners)
|
||||
req_pages = []
|
||||
req_locs = []
|
||||
remaining = int(extend_len)
|
||||
for owner in owners:
|
||||
logical_page = int(owner) + 1 + cp_size * next_generation
|
||||
next_generation += 1
|
||||
req_pages.append(logical_page)
|
||||
valid_rows = min(page_size, remaining)
|
||||
page_locs = torch.arange(
|
||||
logical_page * page_size,
|
||||
logical_page * page_size + valid_rows,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
req_locs.append(page_locs)
|
||||
remaining -= valid_rows
|
||||
self.assertEqual(remaining, 0, f"request {req_id} loc construction")
|
||||
request_pages.append(req_pages)
|
||||
request_valid_locs.append(torch.cat(req_locs, dim=0))
|
||||
all_valid_locs.append(request_valid_locs[-1])
|
||||
|
||||
max_request_pages = max(len(pages) for pages in request_pages)
|
||||
logical_pages = torch.zeros(
|
||||
(batch_size, max_request_pages), dtype=torch.int64, device="cuda"
|
||||
)
|
||||
for req_id, pages in enumerate(request_pages):
|
||||
logical_pages[req_id, : len(pages)] = torch.tensor(
|
||||
pages, dtype=torch.int64, device="cuda"
|
||||
)
|
||||
out_cache_loc = torch.cat(all_valid_locs, dim=0).to(device="cuda")
|
||||
max_logical_page = max(max(pages) for pages in request_pages)
|
||||
physical_page_capacity = (max_logical_page - 1) // cp_size + 3
|
||||
|
||||
rank_dense_buffers = []
|
||||
expected_by_loc: dict[int, torch.Tensor] = {}
|
||||
for rank in range(cp_size):
|
||||
layout = CpSharedKVLayout(
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=rank,
|
||||
)
|
||||
plan = build_batch_page_aligned_in_seq_split_plan(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=rank,
|
||||
)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=layout,
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=batch_size,
|
||||
batch_plan=plan,
|
||||
page_aligned=True,
|
||||
page_size=page_size,
|
||||
extend_prefix_len=prefix_lens[0],
|
||||
),
|
||||
out_cache_loc=out_cache_loc.clone(),
|
||||
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
||||
return_value=False,
|
||||
):
|
||||
logical_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
||||
local_rows = sum(int(x) for x in plan.request_compute_rank_local_tokens)
|
||||
torch.manual_seed(20260607 + rank)
|
||||
local_key = (
|
||||
torch.randn(
|
||||
(local_rows, index_head_dim),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
* 2.0
|
||||
+ 0.125
|
||||
).contiguous()
|
||||
valid_key = select_cp_local_valid_rows_for_cache_write(
|
||||
forward_batch,
|
||||
local_key,
|
||||
).contiguous()
|
||||
self.assertEqual(int(valid_key.shape[0]), int(logical_locs.numel()))
|
||||
|
||||
page_buffer = torch.zeros(
|
||||
(physical_page_capacity, page_bytes),
|
||||
dtype=torch.uint8,
|
||||
device="cuda",
|
||||
)
|
||||
if logical_locs.numel() > 0:
|
||||
physical_locs = layout.logical_locs_to_physical(logical_locs)
|
||||
fused_store_index_k_cache(
|
||||
valid_key,
|
||||
page_buffer,
|
||||
physical_locs.contiguous(),
|
||||
page_size=page_size,
|
||||
)
|
||||
for loc, row in zip(logical_locs.cpu().tolist(), valid_key):
|
||||
expected_by_loc[int(loc)] = row.detach()
|
||||
|
||||
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
|
||||
page_buffer,
|
||||
logical_pages,
|
||||
layout,
|
||||
)
|
||||
with patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
|
||||
), patch.object(
|
||||
runtime, "_try_tai_materialize_shared_pages", return_value=None
|
||||
), patch.object(
|
||||
runtime, "_try_tai_ipc_materialize_paged_buffer_page_slots_into",
|
||||
return_value=False,
|
||||
):
|
||||
dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
|
||||
page_buffer,
|
||||
logical_pages,
|
||||
layout,
|
||||
slot_remap=slot_remap,
|
||||
)
|
||||
rank_dense_buffers.append(dense_page_buffer.to(torch.int16))
|
||||
|
||||
merged_dense_buffer = (
|
||||
torch.stack(rank_dense_buffers, dim=0).sum(dim=0).to(torch.uint8)
|
||||
)
|
||||
|
||||
class FakePool:
|
||||
page_size = 64
|
||||
index_head_dim = 128
|
||||
quant_block_size = 128
|
||||
device = "cuda"
|
||||
|
||||
# Cross-check fallback and fused quantization on the same valid rows so
|
||||
# this test fails on row/order/page corruption, not on an unrelated
|
||||
# scale-format policy difference.
|
||||
for req_id, (extend_len, req_locs) in enumerate(
|
||||
zip(extend_lens, request_valid_locs)
|
||||
):
|
||||
req_page_count = len(request_pages[req_id])
|
||||
page_indices = dense_pages[req_id, :req_page_count].contiguous()
|
||||
k_u8 = index_buf_accessor.GetK.execute(
|
||||
FakePool,
|
||||
merged_dense_buffer,
|
||||
seq_len=int(extend_len),
|
||||
page_indices=page_indices,
|
||||
)
|
||||
scale_u8 = index_buf_accessor.GetS.execute(
|
||||
FakePool,
|
||||
merged_dense_buffer,
|
||||
seq_len=int(extend_len),
|
||||
page_indices=page_indices,
|
||||
)
|
||||
stored_deq = (
|
||||
k_u8.contiguous().view(torch.float8_e4m3fn).float()
|
||||
* scale_u8.contiguous().view(torch.float32).view(-1, 1)
|
||||
)
|
||||
expected_key = torch.stack(
|
||||
[expected_by_loc[int(loc)] for loc in req_locs.tolist()],
|
||||
dim=0,
|
||||
).to(device="cuda")
|
||||
ref_fp8, ref_scale = act_quant(
|
||||
expected_key.contiguous(),
|
||||
block_size=128,
|
||||
scale_fmt="ue8m0",
|
||||
)
|
||||
ref_deq = ref_fp8.float() * ref_scale.float().view(-1, 1)
|
||||
|
||||
torch.testing.assert_close(
|
||||
stored_deq,
|
||||
expected_key.float(),
|
||||
rtol=0.20,
|
||||
atol=0.65,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
stored_deq,
|
||||
ref_deq,
|
||||
rtol=0.35,
|
||||
atol=0.90,
|
||||
)
|
||||
|
||||
|
||||
def test_cp8_index_partial_current_compose_matches_rank_merged_reference_for_bs5(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
|
||||
build_in_seq_page_compute_owners,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
page_size = 64
|
||||
cp_size = 8
|
||||
batch_size = 5
|
||||
prefix_len = 640
|
||||
prefix_pages = prefix_len // page_size
|
||||
extend_lens = [95, 81, 140, 66, 86]
|
||||
index_head_dim = 8
|
||||
scale_bytes = 4
|
||||
page_bytes = page_size * index_head_dim + page_size * scale_bytes
|
||||
pages_per_request = prefix_pages + 3
|
||||
|
||||
shared_prefix_pages = list(range(1, prefix_pages + 1))
|
||||
next_owner_generation = 32
|
||||
logical_rows = []
|
||||
current_locs_by_rank = [[] for _ in range(cp_size)]
|
||||
current_k_by_rank = [[] for _ in range(cp_size)]
|
||||
current_scale_by_rank = [[] for _ in range(cp_size)]
|
||||
expected = torch.zeros(
|
||||
(batch_size * pages_per_request + 1, page_bytes), dtype=torch.uint8
|
||||
)
|
||||
scale_offset = page_size * index_head_dim
|
||||
|
||||
def fill_index_row(buffer, dense_page, offset, k_row, scale_row):
|
||||
buffer[dense_page, offset * index_head_dim : (offset + 1) * index_head_dim] = k_row
|
||||
scale = scale_row.reshape(1, -1).view(torch.uint8).reshape(-1)
|
||||
start = scale_offset + offset * scale_bytes
|
||||
buffer[dense_page, start : start + scale_bytes] = scale
|
||||
|
||||
current_pages_by_req = []
|
||||
for req_id, extend_len in enumerate(extend_lens):
|
||||
owners = build_in_seq_page_compute_owners(
|
||||
extend_len=extend_len,
|
||||
extend_prefix_len=prefix_len,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
)
|
||||
current_pages = []
|
||||
for owner in owners:
|
||||
logical_page = int(owner) + 1 + cp_size * next_owner_generation
|
||||
next_owner_generation += 1
|
||||
current_pages.append(logical_page)
|
||||
current_pages_by_req.append(current_pages)
|
||||
logical_rows.append(
|
||||
shared_prefix_pages
|
||||
+ current_pages
|
||||
+ [0] * (pages_per_request - prefix_pages - len(current_pages))
|
||||
)
|
||||
|
||||
logical_pages = torch.tensor(logical_rows, dtype=torch.int64)
|
||||
max_logical_page = int(logical_pages.max().item())
|
||||
physical_capacity = (max_logical_page - 1) // cp_size + 2
|
||||
|
||||
prefix_page_bytes = {}
|
||||
for logical_page in shared_prefix_pages:
|
||||
row = ((torch.arange(page_bytes, dtype=torch.int64) + logical_page * 17) % 251).to(torch.uint8)
|
||||
prefix_page_bytes[logical_page] = row
|
||||
|
||||
for req_id, current_pages in enumerate(current_pages_by_req):
|
||||
remaining = int(extend_lens[req_id])
|
||||
for page_offset, logical_page in enumerate(current_pages):
|
||||
dense_page = req_id * pages_per_request + prefix_pages + page_offset + 1
|
||||
valid_rows = min(page_size, remaining)
|
||||
owner = int((logical_page - 1) % cp_size)
|
||||
for offset in range(valid_rows):
|
||||
logical_loc = logical_page * page_size + offset
|
||||
k_row = ((torch.arange(index_head_dim, dtype=torch.int64) + req_id * 31 + page_offset * 7 + offset) % 253).to(torch.uint8)
|
||||
scale_row = torch.tensor(
|
||||
[[req_id * 1000.0 + page_offset * 100.0 + offset + 0.5]],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
current_locs_by_rank[owner].append(logical_loc)
|
||||
current_k_by_rank[owner].append(k_row)
|
||||
current_scale_by_rank[owner].append(scale_row.reshape(1))
|
||||
fill_index_row(expected, dense_page, offset, k_row, scale_row)
|
||||
remaining -= valid_rows
|
||||
self.assertEqual(remaining, 0)
|
||||
|
||||
for req_id in range(batch_size):
|
||||
for prefix_slot, logical_page in enumerate(shared_prefix_pages):
|
||||
dense_page = req_id * pages_per_request + prefix_slot + 1
|
||||
expected[dense_page] = prefix_page_bytes[logical_page]
|
||||
|
||||
prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
|
||||
logical_pages=logical_pages,
|
||||
prefix_lens_cpu=[prefix_len] * batch_size,
|
||||
page_size=page_size,
|
||||
)
|
||||
current_slot_spans = runtime.build_batch_current_slot_spans(
|
||||
logical_pages=logical_pages,
|
||||
prefix_lens_cpu=[prefix_len] * batch_size,
|
||||
extend_lens_cpu=extend_lens,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
rank_outputs = []
|
||||
for rank in range(cp_size):
|
||||
layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)
|
||||
page_buffer = torch.zeros((physical_capacity, page_bytes), dtype=torch.uint8)
|
||||
for logical_page, row in prefix_page_bytes.items():
|
||||
if int((logical_page - 1) % cp_size) == rank:
|
||||
physical_page = int((logical_page - 1) // cp_size) + 1
|
||||
page_buffer[physical_page] = row
|
||||
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
|
||||
page_buffer,
|
||||
logical_pages,
|
||||
layout,
|
||||
)
|
||||
current_locs = torch.tensor(current_locs_by_rank[rank], dtype=torch.int64)
|
||||
if current_k_by_rank[rank]:
|
||||
current_k = torch.stack(current_k_by_rank[rank], dim=0)
|
||||
current_scale = torch.stack(current_scale_by_rank[rank], dim=0)
|
||||
else:
|
||||
current_k = torch.empty((0, index_head_dim), dtype=torch.uint8)
|
||||
current_scale = torch.empty((0, 1), dtype=torch.float32)
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce):
|
||||
dense_page_buffer, dense_pages = runtime.materialize_prefix_and_reuse_current_index_page_slots(
|
||||
page_buffer=page_buffer,
|
||||
current_index_k=current_k,
|
||||
current_index_scale=current_scale,
|
||||
current_locs=current_locs,
|
||||
slot_remap=slot_remap,
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
index_head_dim=index_head_dim,
|
||||
prefix_pages=0,
|
||||
prefix_slot_spans=prefix_slot_spans,
|
||||
current_slot_spans=current_slot_spans,
|
||||
layer_id=0,
|
||||
)
|
||||
expected_dense_pages = logical_pages.clone()
|
||||
flat_positive = expected_dense_pages.reshape(-1) > 0
|
||||
expected_dense_pages.reshape(-1)[flat_positive] = torch.arange(
|
||||
1,
|
||||
int(expected_dense_pages.numel()) + 1,
|
||||
dtype=expected_dense_pages.dtype,
|
||||
)[flat_positive]
|
||||
self.assertEqual(dense_pages.tolist(), expected_dense_pages.tolist())
|
||||
rank_outputs.append(dense_page_buffer.to(torch.int16))
|
||||
|
||||
merged = torch.stack(rank_outputs, dim=0).sum(dim=0).to(torch.uint8)
|
||||
torch.testing.assert_close(merged, expected, atol=0, rtol=0)
|
||||
|
||||
|
||||
def test_cp8_kv_partial_current_keeps_remote_current_valid_locs_after_reduce(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
|
||||
|
||||
page_size = 4
|
||||
cp_size = 8
|
||||
# Two current pages: logical page 33 is rank0-owned, logical page 42 is
|
||||
# rank1-owned. After the current slot all-reduce both valid pages are
|
||||
# present on every rank, so rank0 must not mask the rank1-owned valid loc.
|
||||
logical_pages = torch.tensor([[1, 33, 42]], dtype=torch.int64)
|
||||
remote_current_loc = 42 * page_size + 1
|
||||
local_current_locs = torch.tensor(
|
||||
[33 * page_size + 0, 33 * page_size + 1], dtype=torch.int64
|
||||
)
|
||||
logical_locs = torch.tensor(
|
||||
[1 * page_size + 0, 33 * page_size + 1, remote_current_loc],
|
||||
dtype=torch.int64,
|
||||
)
|
||||
layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=0)
|
||||
physical_capacity = 8
|
||||
kv_cache = torch.zeros((physical_capacity * page_size, 2), dtype=torch.float32)
|
||||
kv_cache[layout.logical_locs_to_physical(torch.tensor([page_size], dtype=torch.int64))[0]] = torch.tensor([7.0, 8.0])
|
||||
slot_remap = runtime.build_shared_token_kv_slot_remap(
|
||||
kv_cache,
|
||||
logical_locs,
|
||||
logical_pages,
|
||||
layout,
|
||||
page_size,
|
||||
)
|
||||
current_kv_cache = torch.tensor([[10.0, 11.0], [12.0, 13.0]])
|
||||
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce):
|
||||
_mixed_cache, mixed_locs = runtime.materialize_prefix_and_reuse_current_kv_page_slots(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
current_kv_cache=current_kv_cache,
|
||||
current_locs=local_current_locs,
|
||||
slot_remap=slot_remap,
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
prefix_pages=1,
|
||||
current_slot_spans=[(1, 3)],
|
||||
layer_id=0,
|
||||
)
|
||||
|
||||
self.assertGreaterEqual(
|
||||
int(mixed_locs[-1].item()),
|
||||
0,
|
||||
"remote-rank valid current loc must remain visible after current slot all-reduce",
|
||||
)
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
|
||||
def test_tai_batched_index_mqa_prepare_matches_getk_gets_reference_gsm8k_bs5(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.layers.attention.nsa import (
|
||||
cp_shared_kv_runtime as runtime,
|
||||
index_buf_accessor,
|
||||
)
|
||||
|
||||
page_size = 64
|
||||
index_head_dim = 128
|
||||
scale_bytes = 4
|
||||
row_bytes = index_head_dim + scale_bytes
|
||||
page_bytes = page_size * row_bytes
|
||||
batch_size = 5
|
||||
seq_lens = [704, 704, 704, 704, 704]
|
||||
q_starts = [640, 640, 640, 640, 640]
|
||||
q_lens = [64, 64, 64, 64, 64]
|
||||
pages_per_request = 11
|
||||
num_pages = batch_size * pages_per_request + 1
|
||||
|
||||
index_buffer = torch.empty(
|
||||
(num_pages, page_bytes), dtype=torch.uint8, device="cuda"
|
||||
)
|
||||
token_offsets = torch.arange(
|
||||
page_size, dtype=torch.float32, device="cuda"
|
||||
).view(page_size, 1)
|
||||
byte_offsets = torch.arange(
|
||||
index_head_dim, dtype=torch.int64, device="cuda"
|
||||
).view(1, index_head_dim)
|
||||
for page_id in range(num_pages):
|
||||
page_view = index_buffer[page_id].view(page_size, row_bytes)
|
||||
k_bytes = (
|
||||
byte_offsets
|
||||
+ page_id * 17
|
||||
+ torch.arange(page_size, dtype=torch.int64, device="cuda").view(
|
||||
page_size, 1
|
||||
)
|
||||
).remainder_(251).to(torch.uint8)
|
||||
scales = token_offsets + float(page_id * 1000) + 0.25
|
||||
page_view[:, :index_head_dim] = k_bytes
|
||||
page_view[:, index_head_dim : index_head_dim + scale_bytes] = (
|
||||
scales.contiguous().view(torch.uint8).view(page_size, scale_bytes)
|
||||
)
|
||||
|
||||
block_tables = torch.tensor(
|
||||
[
|
||||
[1 + req_id * pages_per_request + page for page in range(pages_per_request)]
|
||||
for req_id in range(batch_size)
|
||||
],
|
||||
dtype=torch.int64,
|
||||
device="cuda",
|
||||
)
|
||||
batch_indices = torch.arange(batch_size, dtype=torch.int32, device="cuda")
|
||||
kv_lens = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
|
||||
q_starts_tensor = torch.tensor(q_starts, dtype=torch.int32, device="cuda")
|
||||
q_lens_tensor = torch.tensor(q_lens, dtype=torch.int32, device="cuda")
|
||||
k_bases = torch.tensor(
|
||||
[sum(seq_lens[:i]) for i in range(batch_size)],
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
q_bases = torch.tensor(
|
||||
[sum(q_lens[:i]) for i in range(batch_size)],
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_tai_index_mqa_prepare_enabled", return_value=True
|
||||
):
|
||||
prepared = runtime.try_tai_prepare_cp_mqa_index_batch(
|
||||
index_buffer=index_buffer,
|
||||
block_tables=block_tables,
|
||||
batch_indices=batch_indices,
|
||||
kv_lens=kv_lens,
|
||||
q_starts=q_starts_tensor,
|
||||
q_lens=q_lens_tensor,
|
||||
k_bases=k_bases,
|
||||
q_bases=q_bases,
|
||||
total_kv_len=sum(seq_lens),
|
||||
total_q_count=sum(q_lens),
|
||||
max_kv_len=max(seq_lens),
|
||||
max_q_len=max(q_lens),
|
||||
page_size=page_size,
|
||||
index_head_dim=index_head_dim,
|
||||
)
|
||||
if prepared is None:
|
||||
self.skipTest("TAI batched index MQA prepare kernel unavailable")
|
||||
|
||||
k_fp8_u8, k_scale, ks, ke_offset = prepared
|
||||
fake_pool = SimpleNamespace(
|
||||
page_size=page_size,
|
||||
index_head_dim=index_head_dim,
|
||||
quant_block_size=index_head_dim,
|
||||
device="cuda",
|
||||
)
|
||||
ref_k = []
|
||||
ref_scale = []
|
||||
ref_ks = []
|
||||
ref_ke_offset = []
|
||||
for req_id, (seq_len, q_start, q_len) in enumerate(
|
||||
zip(seq_lens, q_starts, q_lens)
|
||||
):
|
||||
ref_k.append(
|
||||
index_buf_accessor.GetK.execute(
|
||||
fake_pool,
|
||||
index_buffer,
|
||||
seq_len=seq_len,
|
||||
page_indices=block_tables[req_id],
|
||||
)
|
||||
)
|
||||
ref_scale.append(
|
||||
index_buf_accessor.GetS.execute(
|
||||
fake_pool,
|
||||
index_buffer,
|
||||
seq_len=seq_len,
|
||||
page_indices=block_tables[req_id],
|
||||
)
|
||||
.view(torch.float32)
|
||||
.squeeze(-1)
|
||||
)
|
||||
ref_ks.append(
|
||||
torch.full(
|
||||
(q_len,),
|
||||
int(k_bases[req_id].item()),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
)
|
||||
ref_ke_offset.append(
|
||||
torch.arange(
|
||||
q_start + 1,
|
||||
q_start + q_len + 1,
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
)
|
||||
ref_k = torch.cat(ref_k, dim=0)
|
||||
ref_scale = torch.cat(ref_scale, dim=0)
|
||||
ref_ks = torch.cat(ref_ks, dim=0)
|
||||
ref_ke_offset = torch.cat(ref_ke_offset, dim=0)
|
||||
|
||||
torch.testing.assert_close(k_fp8_u8, ref_k, atol=0, rtol=0)
|
||||
torch.testing.assert_close(k_scale, ref_scale, atol=0, rtol=0)
|
||||
torch.testing.assert_close(ks, ref_ks, atol=0, rtol=0)
|
||||
torch.testing.assert_close(ke_offset, ref_ke_offset, atol=0, rtol=0)
|
||||
|
||||
def test_token_range_materialize_uses_tai_kernel_when_enabled(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
|
||||
|
||||
@@ -1,8 +1,19 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa.utils import (
|
||||
NSAContextParallelMetadata,
|
||||
build_batch_page_aligned_in_seq_split_plan,
|
||||
get_cp_shared_kv_local_out_cache_loc,
|
||||
select_cp_local_valid_rows_for_cache_write,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
|
||||
build_in_seq_page_compute_owners,
|
||||
)
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
|
||||
from sglang.srt.mem_cache.memory_pool_host import (
|
||||
ALLOC_MEMORY_FUNCS,
|
||||
@@ -133,6 +144,815 @@ class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
].cpu()
|
||||
self.assertTrue(torch.equal(got_kv, expected_kv))
|
||||
|
||||
@staticmethod
|
||||
def _build_owned_logical_out_locs(
|
||||
*,
|
||||
extend_lens,
|
||||
prefix_lens,
|
||||
page_size: int,
|
||||
cp_size: int,
|
||||
dst_page_stride: int = 0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
valid_chunks = []
|
||||
padded_chunks = []
|
||||
logical_page_cursor = 0
|
||||
for req_id, (extend_len, prefix_len) in enumerate(
|
||||
zip(extend_lens, prefix_lens)
|
||||
):
|
||||
owners = build_in_seq_page_compute_owners(
|
||||
extend_len=int(extend_len),
|
||||
extend_prefix_len=int(prefix_len),
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
)
|
||||
if owners is None:
|
||||
raise AssertionError("test shape must support owner planning")
|
||||
remaining = int(extend_len)
|
||||
req_valid_chunks = []
|
||||
req_padded_chunks = []
|
||||
for page_idx, owner in enumerate(owners):
|
||||
logical_page = (
|
||||
int(owner)
|
||||
+ 1
|
||||
+ cp_size
|
||||
* (
|
||||
dst_page_stride
|
||||
+ req_id * (len(owners) + 4)
|
||||
+ logical_page_cursor
|
||||
)
|
||||
)
|
||||
page_locs = torch.arange(
|
||||
logical_page * page_size,
|
||||
(logical_page + 1) * page_size,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
valid_rows = min(page_size, max(remaining, 0))
|
||||
if valid_rows > 0:
|
||||
req_valid_chunks.append(page_locs[:valid_rows])
|
||||
remaining -= valid_rows
|
||||
req_padded_chunks.append(page_locs)
|
||||
logical_page_cursor += 1
|
||||
if remaining != 0:
|
||||
raise AssertionError(
|
||||
f"failed to cover request extend_len={extend_len}"
|
||||
)
|
||||
valid_chunks.append(torch.cat(req_valid_chunks, dim=0))
|
||||
padded_chunks.append(torch.cat(req_padded_chunks, dim=0))
|
||||
return torch.cat(valid_chunks, dim=0), torch.cat(padded_chunks, dim=0)
|
||||
|
||||
@staticmethod
|
||||
def _index_row_bytes(row_values: torch.Tensor, page_size: int) -> torch.Tensor:
|
||||
# index_k_with_scale_buffer stores one page as
|
||||
# [token0(index+scale), token1(index+scale), ...].
|
||||
return row_values.view(row_values.shape[0], page_size, -1)
|
||||
|
||||
def test_fp8_page_first_direct_bs5_valid_rows_survive_l2_roundtrip(self):
|
||||
page_size = 64
|
||||
cp_size = 8
|
||||
layer_num = 2
|
||||
# GSM8K-like second-pass shape: same long page-aligned prompt prefix,
|
||||
# short per-request suffixes, and compute padding on every request.
|
||||
prefix_lens = [640, 640, 640, 640, 640]
|
||||
extend_lens = [95, 81, 140, 66, 86]
|
||||
src_valid_locs, src_padded_locs = self._build_owned_logical_out_locs(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
)
|
||||
dst_valid_locs, dst_padded_locs = self._build_owned_logical_out_locs(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
dst_page_stride=128,
|
||||
)
|
||||
max_physical_page = (
|
||||
int(torch.div(dst_padded_locs.max(), page_size, rounding_mode="floor"))
|
||||
// cp_size
|
||||
+ 4
|
||||
)
|
||||
size = max_physical_page * page_size
|
||||
|
||||
device_pool = NSATokenToKVPool(
|
||||
size=size,
|
||||
page_size=page_size,
|
||||
kv_lora_rank=512,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
qk_rope_head_dim=64,
|
||||
layer_num=layer_num,
|
||||
device="cuda",
|
||||
enable_memory_saver=False,
|
||||
kv_cache_dim=656,
|
||||
index_head_dim=128,
|
||||
)
|
||||
host_pool = NSATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=page_size,
|
||||
layout="page_first_direct",
|
||||
pin_memory=True,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
for cp_rank in range(cp_size):
|
||||
layout = CpSharedKVLayout(
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=cp_rank,
|
||||
)
|
||||
plan = build_batch_page_aligned_in_seq_split_plan(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=cp_rank,
|
||||
)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=layout,
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=len(extend_lens),
|
||||
batch_plan=plan,
|
||||
page_aligned=True,
|
||||
page_size=page_size,
|
||||
extend_prefix_len=prefix_lens[0],
|
||||
),
|
||||
out_cache_loc=src_valid_locs.clone(),
|
||||
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
||||
return_value=False,
|
||||
):
|
||||
local_valid_logical_locs = get_cp_shared_kv_local_out_cache_loc(
|
||||
forward_batch
|
||||
)
|
||||
local_valid_rows = int(local_valid_logical_locs.numel())
|
||||
if local_valid_rows == 0:
|
||||
continue
|
||||
|
||||
local_compute_rows = sum(
|
||||
int(x) for x in plan.request_compute_rank_local_tokens
|
||||
)
|
||||
kv_rows = (
|
||||
torch.arange(
|
||||
local_compute_rows * 656,
|
||||
device="cuda",
|
||||
dtype=torch.int64,
|
||||
)
|
||||
.view(local_compute_rows, 1, 656)
|
||||
.add_(1000 * cp_rank)
|
||||
.remainder_(251)
|
||||
.to(torch.uint8)
|
||||
)
|
||||
index_row_stride = device_pool.index_k_with_scale_buffer[0].shape[1] // page_size
|
||||
index_rows = (
|
||||
torch.arange(
|
||||
local_compute_rows * index_row_stride,
|
||||
device="cuda",
|
||||
dtype=torch.int64,
|
||||
)
|
||||
.view(local_compute_rows, index_row_stride)
|
||||
.add_(2000 * cp_rank)
|
||||
.remainder_(253)
|
||||
.to(torch.uint8)
|
||||
)
|
||||
valid_kv_rows = select_cp_local_valid_rows_for_cache_write(
|
||||
forward_batch, kv_rows
|
||||
)
|
||||
valid_index_rows = select_cp_local_valid_rows_for_cache_write(
|
||||
forward_batch, index_rows
|
||||
)
|
||||
self.assertEqual(int(valid_kv_rows.shape[0]), local_valid_rows)
|
||||
self.assertEqual(int(valid_index_rows.shape[0]), local_valid_rows)
|
||||
|
||||
local_valid_physical_locs = layout.logical_locs_to_physical(
|
||||
local_valid_logical_locs
|
||||
).to(device="cuda")
|
||||
src_owned_padded_locs = src_padded_locs[
|
||||
layout.owned_by_this_rank(src_padded_locs)
|
||||
]
|
||||
dst_owned_padded_locs = dst_padded_locs[
|
||||
layout.owned_by_this_rank(dst_padded_locs)
|
||||
]
|
||||
src_physical_padded_locs = layout.logical_locs_to_physical(
|
||||
src_owned_padded_locs
|
||||
)
|
||||
dst_physical_padded_locs = layout.logical_locs_to_physical(
|
||||
dst_owned_padded_locs
|
||||
)
|
||||
self.assertEqual(
|
||||
int(src_physical_padded_locs.numel()),
|
||||
int(dst_physical_padded_locs.numel()),
|
||||
)
|
||||
host_indices = torch.arange(
|
||||
cp_rank * 4096,
|
||||
cp_rank * 4096 + int(src_physical_padded_locs.numel()),
|
||||
dtype=torch.int64,
|
||||
)
|
||||
# The real reservation backs up full owned pages. Poison both src
|
||||
# and dst first so the assertion proves valid rows were written and
|
||||
# copied, not accidentally preserved.
|
||||
for layer_id in range(layer_num):
|
||||
device_pool.kv_buffer[layer_id].fill_(17 + cp_rank)
|
||||
device_pool.index_k_with_scale_buffer[layer_id].fill_(23 + cp_rank)
|
||||
device_pool.kv_buffer[layer_id][
|
||||
local_valid_physical_locs
|
||||
] = valid_kv_rows
|
||||
page_ids = torch.div(
|
||||
local_valid_physical_locs, page_size, rounding_mode="floor"
|
||||
).to(torch.long)
|
||||
page_offsets = torch.remainder(
|
||||
local_valid_physical_locs, page_size
|
||||
).to(torch.long)
|
||||
index_view = self._index_row_bytes(
|
||||
device_pool.index_k_with_scale_buffer[layer_id], page_size
|
||||
)
|
||||
index_view[page_ids, page_offsets] = valid_index_rows
|
||||
|
||||
expected_kv = [
|
||||
device_pool.kv_buffer[layer_id][local_valid_physical_locs]
|
||||
.detach()
|
||||
.clone()
|
||||
for layer_id in range(layer_num)
|
||||
]
|
||||
expected_index = []
|
||||
page_ids = torch.div(
|
||||
local_valid_physical_locs, page_size, rounding_mode="floor"
|
||||
).to(torch.long)
|
||||
page_offsets = torch.remainder(local_valid_physical_locs, page_size).to(
|
||||
torch.long
|
||||
)
|
||||
for layer_id in range(layer_num):
|
||||
expected_index.append(
|
||||
self._index_row_bytes(
|
||||
device_pool.index_k_with_scale_buffer[layer_id], page_size
|
||||
)[page_ids, page_offsets]
|
||||
.detach()
|
||||
.clone()
|
||||
)
|
||||
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
src_physical_padded_locs,
|
||||
io_backend="direct",
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(layer_num):
|
||||
device_pool.kv_buffer[layer_id][
|
||||
dst_physical_padded_locs.to(device="cuda")
|
||||
].fill_(91 + cp_rank)
|
||||
dst_pages = torch.div(
|
||||
dst_physical_padded_locs, page_size, rounding_mode="floor"
|
||||
).to(device="cuda", dtype=torch.long)
|
||||
device_pool.index_k_with_scale_buffer[layer_id][dst_pages].fill_(
|
||||
97 + cp_rank
|
||||
)
|
||||
|
||||
host_pool.begin_load_to_device_op(
|
||||
host_indices,
|
||||
dst_physical_padded_locs,
|
||||
io_backend="direct",
|
||||
)
|
||||
try:
|
||||
for layer_id in range(layer_num):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
dst_physical_padded_locs,
|
||||
layer_id=layer_id,
|
||||
io_backend="direct",
|
||||
)
|
||||
finally:
|
||||
host_pool.end_load_to_device_op()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
dst_valid_logical_locs_rank = dst_valid_locs[
|
||||
layout.owned_by_this_rank(dst_valid_locs)
|
||||
]
|
||||
dst_valid_physical_locs = layout.logical_locs_to_physical(
|
||||
dst_valid_logical_locs_rank
|
||||
).to(device="cuda")
|
||||
self.assertEqual(
|
||||
int(dst_valid_physical_locs.numel()), int(local_valid_rows)
|
||||
)
|
||||
for layer_id in range(layer_num):
|
||||
got_kv = device_pool.kv_buffer[layer_id][dst_valid_physical_locs]
|
||||
self.assertTrue(
|
||||
torch.equal(got_kv, expected_kv[layer_id]),
|
||||
f"KV valid-row L2 roundtrip mismatch cp_rank={cp_rank} layer={layer_id}",
|
||||
)
|
||||
got_pages = torch.div(
|
||||
dst_valid_physical_locs, page_size, rounding_mode="floor"
|
||||
).to(torch.long)
|
||||
got_offsets = torch.remainder(dst_valid_physical_locs, page_size).to(
|
||||
torch.long
|
||||
)
|
||||
got_index = self._index_row_bytes(
|
||||
device_pool.index_k_with_scale_buffer[layer_id], page_size
|
||||
)[got_pages, got_offsets]
|
||||
self.assertTrue(
|
||||
torch.equal(got_index, expected_index[layer_id]),
|
||||
f"index valid-row L2 roundtrip mismatch cp_rank={cp_rank} layer={layer_id}",
|
||||
)
|
||||
|
||||
def test_fp8_l2_loaded_bs5_suffix_materializes_from_ragged_logical_locs(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
page_size = 64
|
||||
cp_size = 8
|
||||
layer_num = 1
|
||||
prefix_lens = [640, 640, 640, 640, 640]
|
||||
extend_lens = [95, 81, 140, 66, 86]
|
||||
src_valid_locs, src_padded_locs = self._build_owned_logical_out_locs(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
)
|
||||
dst_valid_locs, dst_padded_locs = self._build_owned_logical_out_locs(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
dst_page_stride=128,
|
||||
)
|
||||
max_physical_page = (
|
||||
int(torch.div(dst_padded_locs.max(), page_size, rounding_mode="floor"))
|
||||
// cp_size
|
||||
+ 4
|
||||
)
|
||||
size = max_physical_page * page_size
|
||||
expected_rows_by_dst_loc: dict[int, torch.Tensor] = {}
|
||||
dense_outputs = []
|
||||
dense_locs_reference = None
|
||||
device_pool = NSATokenToKVPool(
|
||||
size=size,
|
||||
page_size=page_size,
|
||||
kv_lora_rank=512,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
qk_rope_head_dim=64,
|
||||
layer_num=layer_num,
|
||||
device="cuda",
|
||||
enable_memory_saver=False,
|
||||
kv_cache_dim=656,
|
||||
index_head_dim=128,
|
||||
)
|
||||
host_pool = NSATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=page_size,
|
||||
layout="page_first_direct",
|
||||
pin_memory=True,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
for cp_rank in range(cp_size):
|
||||
layout = CpSharedKVLayout(
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=cp_rank,
|
||||
)
|
||||
plan = build_batch_page_aligned_in_seq_split_plan(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=cp_rank,
|
||||
)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=layout,
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=len(extend_lens),
|
||||
batch_plan=plan,
|
||||
page_aligned=True,
|
||||
page_size=page_size,
|
||||
extend_prefix_len=prefix_lens[0],
|
||||
),
|
||||
out_cache_loc=src_valid_locs.clone(),
|
||||
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
||||
return_value=False,
|
||||
):
|
||||
local_src_valid_locs = get_cp_shared_kv_local_out_cache_loc(
|
||||
forward_batch
|
||||
)
|
||||
local_rows = int(local_src_valid_locs.numel())
|
||||
|
||||
local_dst_valid_locs = dst_valid_locs[
|
||||
layout.owned_by_this_rank(dst_valid_locs)
|
||||
]
|
||||
self.assertEqual(int(local_dst_valid_locs.numel()), local_rows)
|
||||
src_owned_padded_locs = src_padded_locs[
|
||||
layout.owned_by_this_rank(src_padded_locs)
|
||||
]
|
||||
dst_owned_padded_locs = dst_padded_locs[
|
||||
layout.owned_by_this_rank(dst_padded_locs)
|
||||
]
|
||||
src_physical_padded_locs = layout.logical_locs_to_physical(
|
||||
src_owned_padded_locs
|
||||
)
|
||||
dst_physical_padded_locs = layout.logical_locs_to_physical(
|
||||
dst_owned_padded_locs
|
||||
)
|
||||
host_indices = torch.arange(
|
||||
cp_rank * 4096,
|
||||
cp_rank * 4096 + int(src_physical_padded_locs.numel()),
|
||||
dtype=torch.int64,
|
||||
)
|
||||
|
||||
for layer_id in range(layer_num):
|
||||
device_pool.kv_buffer[layer_id].fill_(11 + cp_rank)
|
||||
if local_rows > 0:
|
||||
local_src_physical_locs = layout.logical_locs_to_physical(
|
||||
local_src_valid_locs
|
||||
).to(device="cuda")
|
||||
row_bytes = (
|
||||
torch.arange(
|
||||
local_rows * 656,
|
||||
device="cuda",
|
||||
dtype=torch.int64,
|
||||
)
|
||||
.view(local_rows, 1, 656)
|
||||
.add_(cp_rank * 37)
|
||||
.remainder_(251)
|
||||
.to(torch.uint8)
|
||||
)
|
||||
device_pool.kv_buffer[0][local_src_physical_locs] = row_bytes
|
||||
for dst_loc, row in zip(
|
||||
local_dst_valid_locs.tolist(),
|
||||
row_bytes.detach().cpu(),
|
||||
strict=True,
|
||||
):
|
||||
expected_rows_by_dst_loc[int(dst_loc)] = row
|
||||
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
src_physical_padded_locs,
|
||||
io_backend="direct",
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
for layer_id in range(layer_num):
|
||||
device_pool.kv_buffer[layer_id][
|
||||
dst_physical_padded_locs.to(device="cuda")
|
||||
].fill_(99 + cp_rank)
|
||||
|
||||
host_pool.begin_load_to_device_op(
|
||||
host_indices,
|
||||
dst_physical_padded_locs,
|
||||
io_backend="direct",
|
||||
)
|
||||
try:
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
dst_physical_padded_locs,
|
||||
layer_id=0,
|
||||
io_backend="direct",
|
||||
)
|
||||
finally:
|
||||
host_pool.end_load_to_device_op()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
with patch.object(
|
||||
runtime,
|
||||
"_all_reduce_materialized_buffer",
|
||||
side_effect=lambda buffer, *args, **kwargs: buffer,
|
||||
):
|
||||
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=device_pool.kv_buffer[0],
|
||||
logical_locs=dst_valid_locs.to(device="cuda"),
|
||||
layout=layout,
|
||||
page_size=page_size,
|
||||
nvtx_source="test.ragged_l2_loaded",
|
||||
nvtx_layer_id=0,
|
||||
)
|
||||
dense_outputs.append(dense_kv.detach().cpu().to(torch.int16))
|
||||
dense_locs = dense_locs.detach().cpu()
|
||||
if dense_locs_reference is None:
|
||||
dense_locs_reference = dense_locs
|
||||
else:
|
||||
self.assertTrue(torch.equal(dense_locs_reference, dense_locs))
|
||||
|
||||
self.assertEqual(set(expected_rows_by_dst_loc), set(dst_valid_locs.tolist()))
|
||||
assert dense_locs_reference is not None
|
||||
merged_dense = torch.stack(dense_outputs, dim=0).sum(dim=0).to(torch.uint8)
|
||||
got = merged_dense[dense_locs_reference]
|
||||
expected = torch.stack(
|
||||
[expected_rows_by_dst_loc[int(loc)] for loc in dst_valid_locs.tolist()],
|
||||
dim=0,
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(got, expected),
|
||||
"L2-loaded FP8 suffix pages must remain readable through RAGGED logical-loc materialize",
|
||||
)
|
||||
|
||||
def test_fp8_l2_loaded_index_pages_materialize_after_fused_store_bs5(self):
|
||||
from sglang.jit_kernel.fused_store_index_cache import (
|
||||
can_use_nsa_fused_store,
|
||||
fused_store_index_k_cache,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa import (
|
||||
cp_shared_kv_runtime as runtime,
|
||||
index_buf_accessor,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
|
||||
|
||||
if not hasattr(torch, "float8_e4m3fn"):
|
||||
self.skipTest("torch.float8_e4m3fn unavailable")
|
||||
if not can_use_nsa_fused_store(torch.bfloat16, torch.int64, 64):
|
||||
self.skipTest("NSA fused index store JIT unavailable")
|
||||
|
||||
page_size = 64
|
||||
cp_size = 8
|
||||
layer_num = 1
|
||||
prefix_lens = [640, 640, 640, 640, 640]
|
||||
extend_lens = [95, 81, 140, 66, 86]
|
||||
src_valid_locs, src_padded_locs = self._build_owned_logical_out_locs(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
)
|
||||
dst_valid_locs, dst_padded_locs = self._build_owned_logical_out_locs(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
dst_page_stride=128,
|
||||
)
|
||||
max_physical_page = (
|
||||
int(torch.div(dst_padded_locs.max(), page_size, rounding_mode="floor"))
|
||||
// cp_size
|
||||
+ 4
|
||||
)
|
||||
size = max_physical_page * page_size
|
||||
device_pool = NSATokenToKVPool(
|
||||
size=size,
|
||||
page_size=page_size,
|
||||
kv_lora_rank=512,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
qk_rope_head_dim=64,
|
||||
layer_num=layer_num,
|
||||
device="cuda",
|
||||
enable_memory_saver=False,
|
||||
kv_cache_dim=656,
|
||||
index_head_dim=128,
|
||||
)
|
||||
host_pool = NSATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=page_size,
|
||||
layout="page_first_direct",
|
||||
pin_memory=True,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
dst_pages_by_req = []
|
||||
cursor = 0
|
||||
for extend_len in extend_lens:
|
||||
num_pages = (int(extend_len) + page_size - 1) // page_size
|
||||
req_padded = dst_padded_locs[cursor : cursor + num_pages * page_size]
|
||||
dst_pages_by_req.append(
|
||||
torch.div(
|
||||
req_padded[::page_size], page_size, rounding_mode="floor"
|
||||
).tolist()
|
||||
)
|
||||
cursor += num_pages * page_size
|
||||
max_pages = max(len(pages) for pages in dst_pages_by_req)
|
||||
logical_pages = torch.zeros(
|
||||
(len(extend_lens), max_pages), dtype=torch.int64, device="cuda"
|
||||
)
|
||||
for req_id, pages in enumerate(dst_pages_by_req):
|
||||
logical_pages[req_id, : len(pages)] = torch.tensor(
|
||||
pages, dtype=torch.int64, device="cuda"
|
||||
)
|
||||
|
||||
expected_rows_by_dst_loc: dict[int, torch.Tensor] = {}
|
||||
dense_outputs = []
|
||||
dense_pages_reference = None
|
||||
for cp_rank in range(cp_size):
|
||||
layout = CpSharedKVLayout(
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=cp_rank,
|
||||
)
|
||||
plan = build_batch_page_aligned_in_seq_split_plan(
|
||||
extend_lens=extend_lens,
|
||||
prefix_lens=prefix_lens,
|
||||
page_size=page_size,
|
||||
cp_size=cp_size,
|
||||
cp_rank=cp_rank,
|
||||
)
|
||||
forward_batch = SimpleNamespace(
|
||||
uses_cp_shared_kv=True,
|
||||
cp_shared_kv_layout=layout,
|
||||
nsa_cp_metadata=NSAContextParallelMetadata(
|
||||
batch_size=len(extend_lens),
|
||||
batch_plan=plan,
|
||||
page_aligned=True,
|
||||
page_size=page_size,
|
||||
extend_prefix_len=prefix_lens[0],
|
||||
),
|
||||
out_cache_loc=src_valid_locs.clone(),
|
||||
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
||||
return_value=False,
|
||||
):
|
||||
local_src_valid_locs = get_cp_shared_kv_local_out_cache_loc(
|
||||
forward_batch
|
||||
)
|
||||
local_rows = int(local_src_valid_locs.numel())
|
||||
local_dst_valid_locs = dst_valid_locs[
|
||||
layout.owned_by_this_rank(dst_valid_locs)
|
||||
]
|
||||
self.assertEqual(int(local_dst_valid_locs.numel()), local_rows)
|
||||
|
||||
src_owned_padded_locs = src_padded_locs[
|
||||
layout.owned_by_this_rank(src_padded_locs)
|
||||
]
|
||||
dst_owned_padded_locs = dst_padded_locs[
|
||||
layout.owned_by_this_rank(dst_padded_locs)
|
||||
]
|
||||
src_physical_padded_locs = layout.logical_locs_to_physical(
|
||||
src_owned_padded_locs
|
||||
)
|
||||
dst_physical_padded_locs = layout.logical_locs_to_physical(
|
||||
dst_owned_padded_locs
|
||||
)
|
||||
host_indices = torch.arange(
|
||||
cp_rank * 4096,
|
||||
cp_rank * 4096 + int(src_physical_padded_locs.numel()),
|
||||
dtype=torch.int64,
|
||||
)
|
||||
|
||||
device_pool.index_k_with_scale_buffer[0].fill_(0)
|
||||
if local_rows > 0:
|
||||
local_compute_rows = sum(
|
||||
int(x) for x in plan.request_compute_rank_local_tokens
|
||||
)
|
||||
torch.manual_seed(20260607 + cp_rank)
|
||||
local_key = (
|
||||
torch.randn(
|
||||
(local_compute_rows, 128),
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 2.0
|
||||
+ 0.125
|
||||
).contiguous()
|
||||
valid_key = select_cp_local_valid_rows_for_cache_write(
|
||||
forward_batch,
|
||||
local_key,
|
||||
).contiguous()
|
||||
self.assertEqual(int(valid_key.shape[0]), local_rows)
|
||||
src_physical_valid_locs = layout.logical_locs_to_physical(
|
||||
local_src_valid_locs
|
||||
).to(device="cuda")
|
||||
fused_store_index_k_cache(
|
||||
valid_key,
|
||||
device_pool.index_k_with_scale_buffer[0],
|
||||
src_physical_valid_locs.contiguous(),
|
||||
page_size=page_size,
|
||||
)
|
||||
for dst_loc, row in zip(
|
||||
local_dst_valid_locs.tolist(),
|
||||
valid_key.detach().cpu(),
|
||||
strict=True,
|
||||
):
|
||||
expected_rows_by_dst_loc[int(dst_loc)] = row
|
||||
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
src_physical_padded_locs,
|
||||
io_backend="direct",
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
device_pool.index_k_with_scale_buffer[0][
|
||||
torch.div(
|
||||
dst_physical_padded_locs.to(device="cuda"),
|
||||
page_size,
|
||||
rounding_mode="floor",
|
||||
).to(torch.long)
|
||||
].fill_(77 + cp_rank)
|
||||
host_pool.begin_load_to_device_op(
|
||||
host_indices,
|
||||
dst_physical_padded_locs,
|
||||
io_backend="direct",
|
||||
)
|
||||
try:
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
dst_physical_padded_locs,
|
||||
layer_id=0,
|
||||
io_backend="direct",
|
||||
)
|
||||
finally:
|
||||
host_pool.end_load_to_device_op()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
slot_remap = runtime.build_shared_paged_buffer_slot_remap(
|
||||
device_pool.index_k_with_scale_buffer[0],
|
||||
logical_pages,
|
||||
layout,
|
||||
)
|
||||
with patch.object(
|
||||
runtime,
|
||||
"_all_reduce_materialized_buffer",
|
||||
side_effect=lambda buffer, *args, **kwargs: buffer,
|
||||
), patch.object(
|
||||
runtime, "_try_tai_materialize_shared_pages", return_value=None
|
||||
), patch.object(
|
||||
runtime,
|
||||
"_try_tai_ipc_materialize_paged_buffer_page_slots_into",
|
||||
return_value=False,
|
||||
):
|
||||
dense_index, dense_pages = runtime.materialize_shared_paged_buffer(
|
||||
device_pool.index_k_with_scale_buffer[0],
|
||||
logical_pages,
|
||||
layout,
|
||||
slot_remap=slot_remap,
|
||||
nvtx_source="test.index_l2_loaded",
|
||||
nvtx_layer_id=0,
|
||||
)
|
||||
dense_outputs.append(dense_index.detach().cpu().to(torch.int16))
|
||||
dense_pages = dense_pages.detach().cpu()
|
||||
if dense_pages_reference is None:
|
||||
dense_pages_reference = dense_pages
|
||||
else:
|
||||
self.assertTrue(torch.equal(dense_pages_reference, dense_pages))
|
||||
|
||||
self.assertEqual(set(expected_rows_by_dst_loc), set(dst_valid_locs.tolist()))
|
||||
assert dense_pages_reference is not None
|
||||
merged_dense = torch.stack(dense_outputs, dim=0).sum(dim=0).to(torch.uint8)
|
||||
|
||||
class FakePool:
|
||||
page_size = 64
|
||||
index_head_dim = 128
|
||||
quant_block_size = 128
|
||||
device = "cuda"
|
||||
|
||||
for req_id, (extend_len, req_locs) in enumerate(
|
||||
zip(extend_lens, torch.split(dst_valid_locs, extend_lens))
|
||||
):
|
||||
page_count = len(dst_pages_by_req[req_id])
|
||||
page_indices = dense_pages_reference[req_id, :page_count].to(
|
||||
device="cuda"
|
||||
)
|
||||
dense_gpu = merged_dense.to(device="cuda")
|
||||
got_k = index_buf_accessor.GetK.execute(
|
||||
FakePool,
|
||||
dense_gpu,
|
||||
seq_len=int(extend_len),
|
||||
page_indices=page_indices.contiguous(),
|
||||
)
|
||||
got_s = index_buf_accessor.GetS.execute(
|
||||
FakePool,
|
||||
dense_gpu,
|
||||
seq_len=int(extend_len),
|
||||
page_indices=page_indices.contiguous(),
|
||||
)
|
||||
got_deq = (
|
||||
got_k.contiguous().view(torch.float8_e4m3fn).float()
|
||||
* got_s.contiguous().view(torch.float32).view(-1, 1)
|
||||
)
|
||||
expected_key = torch.stack(
|
||||
[expected_rows_by_dst_loc[int(loc)] for loc in req_locs.tolist()],
|
||||
dim=0,
|
||||
).to(device="cuda")
|
||||
ref_fp8, ref_scale = act_quant(
|
||||
expected_key.contiguous(),
|
||||
block_size=128,
|
||||
scale_fmt="ue8m0",
|
||||
)
|
||||
ref_deq = ref_fp8.float() * ref_scale.float().view(-1, 1)
|
||||
torch.testing.assert_close(
|
||||
got_deq,
|
||||
expected_key.float(),
|
||||
rtol=0.20,
|
||||
atol=0.65,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
got_deq,
|
||||
ref_deq,
|
||||
rtol=0.35,
|
||||
atol=0.90,
|
||||
)
|
||||
|
||||
def test_device_to_host_indexer_kernel_layer_first(self):
|
||||
self._run_device_to_host_indexer_copy(
|
||||
io_backend="kernel", layout="layer_first"
|
||||
@@ -148,6 +968,127 @@ class TestNSAHiCacheTransfer(CustomTestCase):
|
||||
io_backend="direct", layout="page_first_direct"
|
||||
)
|
||||
|
||||
def test_fp8_page_first_direct_roundtrip_preserves_kv_and_indexer_pages(self):
|
||||
page_size = 64
|
||||
layer_num = 3
|
||||
size = page_size * 20
|
||||
|
||||
device_pool = NSATokenToKVPool(
|
||||
size=size,
|
||||
page_size=page_size,
|
||||
kv_lora_rank=512,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
qk_rope_head_dim=64,
|
||||
layer_num=layer_num,
|
||||
device="cuda",
|
||||
enable_memory_saver=False,
|
||||
kv_cache_dim=656,
|
||||
index_head_dim=128,
|
||||
)
|
||||
host_pool = NSATokenToKVPoolHost(
|
||||
device_pool=device_pool,
|
||||
host_to_device_ratio=2.0,
|
||||
host_size=0,
|
||||
page_size=page_size,
|
||||
layout="page_first_direct",
|
||||
pin_memory=True,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
for layer_id in range(layer_num):
|
||||
kv_buf = device_pool.kv_buffer[layer_id]
|
||||
kv_data = torch.arange(
|
||||
kv_buf.numel(), device="cuda", dtype=torch.int64
|
||||
).view_as(kv_buf)
|
||||
kv_buf.copy_(((kv_data + 17 * layer_id) % 251).to(torch.uint8))
|
||||
|
||||
index_buf = device_pool.index_k_with_scale_buffer[layer_id]
|
||||
index_data = torch.arange(
|
||||
index_buf.numel(), device="cuda", dtype=torch.int64
|
||||
).view_as(index_buf)
|
||||
index_buf.copy_(((index_data + 29 * layer_id) % 253).to(torch.uint8))
|
||||
|
||||
src_pages = torch.tensor([1, 2, 5, 8], dtype=torch.int64)
|
||||
host_pages = torch.tensor([3, 4, 7, 12], dtype=torch.int64)
|
||||
dst_pages = torch.tensor([10, 11, 13, 15], dtype=torch.int64)
|
||||
device_indices = self._token_indices_for_pages(
|
||||
src_pages, page_size, device="cpu"
|
||||
)
|
||||
host_indices = self._token_indices_for_pages(
|
||||
host_pages, page_size, device="cpu"
|
||||
)
|
||||
load_device_indices = self._token_indices_for_pages(
|
||||
dst_pages, page_size, device="cpu"
|
||||
)
|
||||
|
||||
expected_kv = []
|
||||
expected_index = []
|
||||
for layer_id in range(layer_num):
|
||||
expected_kv.append(
|
||||
[
|
||||
device_pool.kv_buffer[layer_id][
|
||||
int(page) * page_size : (int(page) + 1) * page_size
|
||||
]
|
||||
.detach()
|
||||
.clone()
|
||||
for page in src_pages.tolist()
|
||||
]
|
||||
)
|
||||
expected_index.append(
|
||||
[
|
||||
device_pool.index_k_with_scale_buffer[layer_id][int(page)]
|
||||
.detach()
|
||||
.clone()
|
||||
for page in src_pages.tolist()
|
||||
]
|
||||
)
|
||||
|
||||
host_pool.backup_from_device_all_layer(
|
||||
device_pool, host_indices, device_indices, io_backend="direct"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(layer_num):
|
||||
for page in dst_pages.tolist():
|
||||
start = int(page) * page_size
|
||||
device_pool.kv_buffer[layer_id][start : start + page_size].fill_(0)
|
||||
device_pool.index_k_with_scale_buffer[layer_id][int(page)].fill_(0)
|
||||
|
||||
host_pool.begin_load_to_device_op(
|
||||
host_indices, load_device_indices, io_backend="direct"
|
||||
)
|
||||
try:
|
||||
for layer_id in range(layer_num):
|
||||
host_pool.load_to_device_per_layer(
|
||||
device_pool,
|
||||
host_indices,
|
||||
load_device_indices,
|
||||
layer_id=layer_id,
|
||||
io_backend="direct",
|
||||
)
|
||||
finally:
|
||||
host_pool.end_load_to_device_op()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
for layer_id in range(layer_num):
|
||||
for page_idx, dst_page in enumerate(dst_pages.tolist()):
|
||||
dst_start = int(dst_page) * page_size
|
||||
got_kv = device_pool.kv_buffer[layer_id][
|
||||
dst_start : dst_start + page_size
|
||||
]
|
||||
self.assertTrue(
|
||||
torch.equal(got_kv, expected_kv[layer_id][page_idx]),
|
||||
f"KV roundtrip mismatch layer={layer_id} dst_page={dst_page}",
|
||||
)
|
||||
|
||||
got_index = device_pool.index_k_with_scale_buffer[layer_id][
|
||||
int(dst_page)
|
||||
]
|
||||
self.assertTrue(
|
||||
torch.equal(got_index, expected_index[layer_id][page_idx]),
|
||||
f"index roundtrip mismatch layer={layer_id} dst_page={dst_page}",
|
||||
)
|
||||
|
||||
|
||||
class TestPageFirstDirectAllLayerBackupRoute(CustomTestCase):
|
||||
def test_mla_page_first_direct_all_layer_backup_uses_tai_per_layer_route(self):
|
||||
|
||||
Reference in New Issue
Block a user