Route CP shared MLA store through TAI fused kernels without runtime spam

The shared-KV prefill path now optionally calls tai_kernel.nsa_prefill.fused_store_mla_kv before falling back to logical_locs_to_physical plus set_mla_kv_buffer. The fast path supports packed FP8 and BF16/FP16 direct KV buffers, while debug mode and kernel failures still preserve the existing fallback behavior. Success logging was removed after path verification because per-layer/per-rank logs are too noisy in normal server runs.

Constraint: Runtime must remain safe when tai-kernel is absent or debug checks are enabled
Rejected: Keep success logs permanently | floods prefill logs once every rank/layer starts using the fast path
Confidence: high
Scope-risk: moderate
Directive: Keep fallback warnings; do not re-add per-layer success logs outside explicit debug instrumentation
Tested: g0034 container python -m py_compile python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
Tested: g0034 container PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (40 passed)
Not-tested: Full multi-node PD server throughput after log removal
This commit is contained in:
laoyao0822
2026-05-06 00:54:47 +08:00
parent 49eaf9ffde
commit 5e5ac5e2e7
5 changed files with 364 additions and 3 deletions

View File

@@ -474,6 +474,50 @@ SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 fast path
- fallback reason 可解释;
- profiler 中 `materialize.token.local_copy` / `materialize.paged.local_copy` 下降。
### P0 fused MLA persistent store 接入
在 materialize 之前shared KV direct-write 的 MLA persistent KV 写入还有一段
独立开销:
```text
logical_locs_to_physical(...)
quantize_k_cache_separate(...)
set_mla_kv_buffer_triton(...)
```
对应的 tai-kernel fast path 已通过 SGLang 实验开关接入:
```text
SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=0/1
```
默认关闭。开启后只影响 prefill CP shared-KV 的 MLA local persistent write
```text
forward_mla.py
_maybe_write_cp_shared_local_mla_kv(...)
-> try_tai_fused_mla_store(...)
-> tai_kernel.nsa_prefill.fused_store_mla_kv(
k_nope, k_rope, raw_kv_buffer, logical_locs,
page_size, cp_size)
-> fallback logical_locs_to_physical + token_to_kv_pool.set_mla_kv_buffer
```
接入约束:
1. `fused_store_mla_kv` 根据 raw KV buffer dtype 分发:
- `uint8 [capacity, 1, 656]`packed NSA FP8 路径,
fused logical->physical + FP8 quant + scale/rope store
- `bf16/fp16 [capacity, 1, 576]`non-FP8 direct-store 路径,
fused logical->physical + MLA KV direct write。
2. tai-kernel import 失败、shape/dtype 不支持、kernel runtime 失败时自动 fallback。
3. `SGLANG_DEBUG_CP_SHARED_KV=1` 时强制保留原 PyTorch path便于继续暴露
shared KV correctness 问题。
4. fused kernel 消费 shared logical token locs调用前仍要求 direct-write gate
已确认 `local_out_cache_loc` 属于当前 CP owner。
5. 该路径不替代后续 attention 前的 shared KV materialize只减少 persistent
write 阶段的 remap/quant/store kernel 碎片。
## Benchmark plan
新增 benchmark 只测本地 materialize不包含 CP all-reduce