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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user