Preserve FP8 CP shared-KV page contracts
NSA FP8 CP shared-KV reuse must operate on packed page-slot rows, not bf16 compact rows. The change keeps current-only and partial-current reuse inside the page-aligned materialization contract, fails fast for non-page-aligned CP split inputs, and prevents FP8 FlashMLA-KV prefill from reaching incompatible in-seq CP metadata. Constraint: NSA FP8 persistent MLA KV rows are packed 656-byte records and CP shared KV cache management is page-granular.\nConstraint: FlashMLA-KV prefill metadata is not CP-local after NSA in-seq splitting.\nRejected: Silently splice bf16 current rows into FP8 materialized cache | corrupts the packed cache layout.\nRejected: Keep FP8 FlashMLA-KV auto prefill under NSA CP | reaches num_splits shape errors after q-row splitting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable FP8 FlashMLA-KV prefill for NSA in-seq CP until metadata is rebuilt after CP splitting or made CP-local.\nTested: Local git diff --check and py_compile for touched SGLang files.\nTested: Remote g0034 related unit sweep recorded in docs: test_nsa_cp_utils.py, test_cp_shared_kv_layout.py, test_cp_shared_kv_runtime.py, test_cp_hicache_metadata.py passed.\nNot-tested: Full FP8 ETE startup and performance run after this commit.
This commit is contained in:
laoyao0822
@@ -4290,3 +4290,307 @@ Risk / not yet covered:
|
||||
capability, this can be tightened to fail-fast for the production GLM5 config.
|
||||
- Remote CUDA verification is still required after syncing both `tai-kernel` and
|
||||
`sglang-dev`; local CPU import is blocked by missing optional dependencies.
|
||||
|
||||
### C103 — 2026-06-01 FP8 MLA current reuse must splice packed cache rows, not bf16 compact rows
|
||||
|
||||
Problem found:
|
||||
|
||||
- NSA FP8 KV cache stores MLA rows in the packed persistent layout:
|
||||
`[512 FP8 nope bytes][4 FP32 scales][64 BF16 rope values]`, i.e. `656` bytes
|
||||
per token row.
|
||||
- The CP shared-KV partial-current MLA path previously built fresh current rows
|
||||
with `cat([k_nope, k_rope])`, i.e. the direct bf16/fp16 `576`-wide layout.
|
||||
That is only valid for non-FP8 cache. In FP8 mode it would splice rows whose
|
||||
dtype/row-width do not match the materialized prefix buffer.
|
||||
|
||||
Implementation direction:
|
||||
|
||||
- `tai-kernel` now provides `pack_quant_mla_kv(k_nope, k_rope) -> uint8[N,1,656]`
|
||||
by reusing the fused quantized MLA store kernel with identity logical locs.
|
||||
This is correctness-first and keeps the quantization/store contract in TAI
|
||||
rather than rebuilding it with torch ops in SGLang.
|
||||
- SGLang detects packed FP8 MLA KV cache by row width `656` plus `uint8` or
|
||||
`float8_e4m3*` dtype, then packs current rows before current-only or
|
||||
partial-current reuse.
|
||||
- `fp8_e5m2` remains fail-fast for NSA CP shared-KV FP8 reuse; current model
|
||||
config only supports `fp8_e4m3` for NSA.
|
||||
- No silent bf16 fallback is allowed for FP8 current reuse. Missing TAI pack
|
||||
capability raises `[CP_SHARED_KV_FAIL_FAST][fp8_current_pack]`.
|
||||
|
||||
IPC coverage:
|
||||
|
||||
- The descriptor-driven IPC materialize kernel is byte-count based. It already
|
||||
covers FP8 MLA pages when `page_nbytes = page_size * 656`, so no separate IPC
|
||||
copy kernel is required for packed FP8 rows.
|
||||
- A CUDA test was added to lock the FP8 packed MLA page size through the IPC
|
||||
slot-dense materialize path.
|
||||
|
||||
Verification status:
|
||||
|
||||
- Local and remote `py_compile` passed for touched SGLang and TAI files.
|
||||
- Local source-level backend hook test passed. Local dynamic helper tests remain
|
||||
blocked by missing local `orjson`; the remote container is authoritative for
|
||||
those tests.
|
||||
- Remote `g0034` SGLang targeted tests passed for FP8 current-pack success,
|
||||
fail-fast-on-missing-TAI, and backend current-reuse hook detection.
|
||||
- Remote `g0034` TAI CUDA tests passed for `pack_quant_mla_kv` identity layout
|
||||
and FP8 packed MLA page IPC slot-dense materialize.
|
||||
|
||||
### C104 — 2026-06-01 current-only MLA reuse must still return page-slot KV buffers
|
||||
|
||||
Problem found:
|
||||
|
||||
- FP8 current-only CP shared-KV reuse packed the freshly computed current rows
|
||||
into a compact tensor shaped like `[valid_tokens, 1, 656]` and assigned it
|
||||
directly to `kv_cache`.
|
||||
- `flashmla_kv` consumes MLA KV cache as pages and reshapes it with
|
||||
`view(-1, page_size, 1, kv_cache_dim)`. A short current-only request such as
|
||||
4 tokens therefore produced `4 * 656 = 2624` elements, which cannot be viewed
|
||||
as full 64-token pages.
|
||||
- This violated the page-aligned cache contract. FP8 does not get an exception:
|
||||
even packed FP8 MLA rows must be stored and exposed through page-sized slots.
|
||||
|
||||
Implementation direction:
|
||||
|
||||
- The current-only branch now uses the same page-slot compose helper as
|
||||
partial-current reuse, with `prefix_pages=0`.
|
||||
- `prefix_pages=0` means there are no historical prefix pages to materialize or
|
||||
reduce. It does **not** mean the output buffer has zero pages. The current
|
||||
suffix still occupies `ceil(valid_tokens / page_size)` physical page slots.
|
||||
- The compact current tensor remains only an input to the slot-fill step. The
|
||||
attention-facing `kv_cache` is the composed page buffer, and returned locs are
|
||||
remapped into that page buffer with non-current tail slack masked out.
|
||||
|
||||
Verification status:
|
||||
|
||||
- Local source-level regression test passed for the backend current-only branch:
|
||||
it now requires `materialize_prefix_and_reuse_current_kv_page_slots`,
|
||||
`prefix_pages=0`, and forbids `kv_cache = current_kv_cache`.
|
||||
- Local `py_compile` passed for `nsa_backend.py` and `cp_shared_kv_runtime.py`.
|
||||
- Remote `g0034` container `py_compile` passed for `nsa_backend.py` and
|
||||
`cp_shared_kv_runtime.py`.
|
||||
- Remote `g0034` targeted runtime tests passed for current-only page-slot
|
||||
layout, padded-current reuse gating, FP8 current-pack success, and
|
||||
fail-fast-on-missing-TAI.
|
||||
- Local dynamic runtime helper tests are still blocked by missing optional local
|
||||
dependency `orjson`; the remote container remains authoritative for those
|
||||
tests.
|
||||
|
||||
### C105 — 2026-06-01 current-only tiny suffixes must also skip NSA in-seq CP split
|
||||
|
||||
Problem found:
|
||||
|
||||
- The short-page CP-split guard only handled radix-hit suffixes with a positive
|
||||
page-aligned prefix. Current-only FP8 warmup/first-token paths with
|
||||
`prefix_len=0` and `extend_len=64` still entered NSA in-seq CP split.
|
||||
- With `cp_size=8` and `page_size=64`, that is only one physical page. The
|
||||
page-aligned in-seq split produces mostly zero-token CP lanes. The downstream
|
||||
FlashMLA KV path then sees inconsistent q-batch / metadata cardinality and the
|
||||
kernel rejects `num_splits` with `num_splits must have shape (b+1)`.
|
||||
|
||||
Implementation direction:
|
||||
|
||||
- The page-unit guard now treats `prefix_len=0` as a valid page-aligned prefix
|
||||
for the purpose of deciding whether CP split is safe.
|
||||
- If `ceil(extend_len / page_size) < cp_size`, CP split is skipped for CP shared
|
||||
KV regardless of whether the request is current-only or a radix hit.
|
||||
- This preserves the page cache contract: we avoid splitting a single page over
|
||||
eight owner lanes instead of falling back to token-balanced sub-page splits.
|
||||
|
||||
Verification status:
|
||||
|
||||
- Added a unit regression that `extend_len=64`, `prefix_len=0`, `page_size=64`,
|
||||
`cp_size=8` does not use NSA in-seq CP split.
|
||||
- Remote `g0034` container `py_compile` passed for `nsa/utils.py`.
|
||||
- Remote `g0034` targeted tests passed for the new current-only skip guard and
|
||||
adjacent short-page / one-page-per-rank CP-split behavior.
|
||||
- Remote `g0034` combined targeted tests passed for the current-only page-slot
|
||||
FP8 reuse checks plus the new CP-split guard.
|
||||
|
||||
### C106 — 2026-06-01 effective NSA impl must drive metadata and top-k layout
|
||||
|
||||
Finding:
|
||||
|
||||
- `forward_extend()` selects the effective NSA implementation from forward mode:
|
||||
target-verify and draft-extend use `nsa_decode_impl`, while normal extend uses
|
||||
`nsa_prefill_impl`.
|
||||
- Two metadata/layout decisions still used global fields instead of that effective
|
||||
implementation:
|
||||
- `flashmla_metadata` construction in normal `init_forward_metadata()` was gated
|
||||
only by `nsa_decode_impl == "flashmla_kv"`.
|
||||
- `get_topk_transform_method()` selected ragged-vs-paged from
|
||||
`nsa_prefill_impl` only.
|
||||
|
||||
Impact:
|
||||
|
||||
- The default Hopper FP8 setup (`prefill=flashmla_auto -> flashmla_kv`,
|
||||
`decode=flashmla_kv`) did not expose this, but explicit backend combinations
|
||||
could mismatch: e.g. decode `flashmla_kv` with prefill `flashmla_sparse` could
|
||||
make target/draft use FlashMLA-KV attention with ragged top-k metadata.
|
||||
- This is the same class of inconsistency as the recent FP8 failures: attention
|
||||
kernel selection, top-k/index layout, and FlashMLA scheduler metadata must be
|
||||
derived from one effective implementation for the current forward mode.
|
||||
|
||||
Correction:
|
||||
|
||||
- Added `_effective_nsa_impl_for_forward_batch()` and routed normal metadata
|
||||
creation, forward-layer top-k transform selection, and indexer metadata through
|
||||
it.
|
||||
- CUDA-graph decode metadata remains decode-backend gated because that path is
|
||||
decode/verify/draft graph state, not normal prefill metadata.
|
||||
|
||||
Verification:
|
||||
|
||||
- Remote g0034 container `py_compile` passed for `nsa_backend.py` and the runtime
|
||||
test file.
|
||||
- Remote targeted tests passed for effective top-k transform, FlashMLA metadata
|
||||
source gating, and current-only page-slot reuse.
|
||||
|
||||
### C107 — 2026-06-01 non-page-aligned CP shared-KV prefixes must fail before CP split
|
||||
|
||||
Finding:
|
||||
|
||||
- Production radix/HiCache prefix matching is supposed to floor CP shared-KV
|
||||
cache hits to page boundaries before scheduling. Evidence:
|
||||
- base radix match floors keys to page size;
|
||||
- HiRadix CP paths floor exact valid-tail hits and backed partial-tail splits;
|
||||
- scheduler runtime checks require `cache_protected_len % page_size == 0`.
|
||||
- The CP split builder still has a legacy token-balanced fallback when called
|
||||
with a non-page-aligned `extend_prefix_len`. If production ever reaches that
|
||||
branch, CP split would reintroduce sub-page split semantics instead of exposing
|
||||
the upstream page-flooring bug.
|
||||
|
||||
Correction:
|
||||
|
||||
- `can_cp_split()` now fails fast for single-request CP shared-KV in-seq split
|
||||
when `extend_prefix_lens_cpu[0]` is not page-aligned.
|
||||
- This does not change the normal current-only path (`prefix_len=0`) or valid
|
||||
page-aligned radix-hit path. It only makes an invalid production state explicit.
|
||||
|
||||
Verification status:
|
||||
|
||||
- Added unit coverage that a non-page-aligned CP shared-KV prefix raises
|
||||
`[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix]` before CP split.
|
||||
- Remote g0034 targeted test passed for non-page-aligned prefix fail-fast plus adjacent current-only and page-aligned CP-split cases.
|
||||
|
||||
### C108 — 2026-06-01 current-only index reuse tests must model the real CPU-length contract
|
||||
|
||||
Finding:
|
||||
|
||||
- A wider remote unit run exposed a stale test stub for NSA index current-only
|
||||
reuse. The production gate is `is_current_only_extend_batch()`, which requires
|
||||
`forward_mode.is_extend_without_speculative()`, zero prefix lengths, and
|
||||
`seq_lens_cpu == extend_seq_lens_cpu`.
|
||||
- The test passed `current_index_kv` but omitted those CPU metadata fields, so
|
||||
the production helper correctly treated it as not current-only and attempted
|
||||
shared index materialization.
|
||||
|
||||
Correction:
|
||||
|
||||
- The test stub now carries the same CPU-length contract as production current-only
|
||||
batches. This keeps the assertion meaningful: current-only index reuse skips
|
||||
shared index materialization only when the batch metadata proves that all
|
||||
visible rows are current rows.
|
||||
|
||||
Verification:
|
||||
|
||||
- Remote g0034 container two-file unit suite passed:
|
||||
`test_nsa_cp_utils.py` + `test_cp_shared_kv_runtime.py`
|
||||
(`124 passed, 5 warnings, 2 subtests passed`).
|
||||
- Remote g0034 container page-aligned/cache runtime suite passed:
|
||||
`test_nsa_cp_utils.py` + `test_cp_shared_kv_layout.py` +
|
||||
`test_cp_shared_kv_runtime.py`
|
||||
(`149 passed, 5 warnings, 2 subtests passed`).
|
||||
|
||||
### C109 — 2026-06-01 FP8 HiCache target/draft budget split uses packed row bytes
|
||||
|
||||
Finding:
|
||||
|
||||
- CP HiCache shares `--hicache-size` across target and EAGLE/NextN draft when a
|
||||
draft KV pool is attached.
|
||||
- The split is byte-based but chooses a target token capacity first:
|
||||
`target_pages = budget / ((target_bytes_per_token + draft_bytes_per_token) * page_size)`.
|
||||
The remaining bytes go to draft, and draft capacity must be at least target
|
||||
capacity.
|
||||
- For NSA FP8, `NSATokenToKVPool.kv_cache_dim` is the packed persistent row width
|
||||
(`656` uint8 bytes/token/layer), not the BF16 compact MLA width. HiCache host
|
||||
pool construction passes the device `kv_cache_dim` as `override_kv_cache_dim`,
|
||||
so `NSATokenToKVPoolHost.get_size_per_token()` and
|
||||
`_estimate_hicache_size_per_token()` agree.
|
||||
- GLM-5 EAGLE draft has one executable draft layer. Therefore under FP8 the
|
||||
draft host allocation can be small in bytes (for example around 1/78 of the
|
||||
target KV bytes) while still having equal or slightly larger token capacity.
|
||||
A small draft GB log is not by itself a capacity bug.
|
||||
|
||||
Correction:
|
||||
|
||||
- No production code change was needed for the current target/draft HiCache size
|
||||
split. Added regression tests to pin the FP8 packed-row byte contract and the
|
||||
one-layer draft capacity implication.
|
||||
|
||||
Verification status:
|
||||
|
||||
- Remote g0034 container targeted FP8 HiCache budget tests passed
|
||||
(`4 passed, 3 warnings`).
|
||||
- Remote g0034 container full HiCache metadata unit file passed
|
||||
(`102 passed, 5 warnings`).
|
||||
|
||||
### C110 — 2026-06-01 FP8 FlashMLA-KV prefill is incompatible with NSA in-seq CP split metadata
|
||||
|
||||
Finding:
|
||||
|
||||
- Remote g0034 prefill log `sglang_cp_hicache_20260531_183930.log` died on all
|
||||
CP ranks with `RuntimeError: num_splits must have shape (b+1)` inside
|
||||
`flash_mla_with_kvcache()`.
|
||||
- Startup was Hopper FP8 KV cache with `nsa_prefill_backend=flashmla_auto`,
|
||||
`nsa_decode_backend=flashmla_kv`, `--enable-nsa-prefill-context-parallel`, and
|
||||
`--nsa-prefill-cp-mode in-seq-split`.
|
||||
- The failing request was a normal target prefill/extend with about 40k input
|
||||
tokens and no cache prefix. Auto dispatch selected `flashmla_kv` for FP8
|
||||
prefill on Hopper.
|
||||
- In in-seq CP split, the model input hidden states are split by
|
||||
`prepare_input_dp_with_cp_dsa()`/`cp_split_and_rebuild_data()`, but normal NSA
|
||||
metadata still builds FlashMLA-KV scheduler metadata from the unsplit full
|
||||
`seqlens_expanded`. `_forward_flashmla_kv()` then reshapes local CP q rows as
|
||||
`q=[local_q, 1, heads, dim]` while passing `num_splits` sized for the full q
|
||||
row count, violating the FlashMLA-KV `(b+1)` contract.
|
||||
- BF16 did not hit this because the default prefill implementation is
|
||||
`flashmla_sparse`, which does not use FlashMLA-KV `num_splits` metadata.
|
||||
|
||||
Correction direction:
|
||||
|
||||
- Treat FP8 FlashMLA-KV as decode/verify/draft-only for NSA prefill CP until the
|
||||
in-seq CP metadata is made FlashMLA-KV-aware.
|
||||
- Auto FP8 prefill under NSA CP should dispatch to `flashmla_sparse` instead of
|
||||
`flashmla_kv`; explicit invalid `flashmla_kv` prefill under active in-seq CP
|
||||
should fail loudly rather than reaching the kernel shape error.
|
||||
|
||||
Verification target:
|
||||
|
||||
- Add unit coverage for FP8 `flashmla_auto` + NSA prefill CP selecting
|
||||
`flashmla_sparse`.
|
||||
- Sync to g0034 and rerun the remote unit slice before ETE restart.
|
||||
|
||||
C110 implementation note:
|
||||
|
||||
- `flashmla_auto` now selects `flashmla_sparse` for FP8 normal prefill whenever
|
||||
NSA prefill CP is enabled and the forward mode is context-parallel extend.
|
||||
This avoids constructing full-q FlashMLA-KV scheduler metadata before the
|
||||
model's in-seq CP split narrows q rows.
|
||||
- If an explicit configuration still routes an actual CP-split FP8 prefill layer
|
||||
into `flashmla_kv`, `forward_extend()` raises
|
||||
`[NSA_FP8_FLASHMLA_KV_CP_PREFILL_UNSUPPORTED]` before calling the kernel.
|
||||
This is intentional: the alternative is the opaque kernel error
|
||||
`num_splits must have shape (b+1)`.
|
||||
|
||||
C110 verification update:
|
||||
|
||||
- Remote g0034 container `py_compile` passed for `nsa_backend.py` and
|
||||
`test_cp_shared_kv_runtime.py`.
|
||||
- Remote targeted tests passed for FP8 auto prefill CP dispatch, explicit
|
||||
FlashMLA-KV CP prefill fail-fast, effective top-k transform selection, and
|
||||
FlashMLA metadata gating (`4 passed, 5 warnings`).
|
||||
- Remote related unit sweep passed:
|
||||
`test_nsa_cp_utils.py`, `test_cp_shared_kv_layout.py`,
|
||||
`test_cp_shared_kv_runtime.py`, and `test_cp_hicache_metadata.py`
|
||||
(`253 passed, 5 warnings, 2 subtests passed`).
|
||||
|
||||
Reference in New Issue
Block a user