sglang/gov/work/2026-06-08-fix-bs-1-cp-shared-kv-cross-request-kv-contamination-per-request-page-inverse.toml

#:schema ../schema/work.schema.json

[govctl]
id = "WI-2026-06-08-001"
title = "Fix bs>1 CP shared-KV cross-request KV contamination (per-request page_inverse)"
status = "active"
created = "2026-06-08"
started = "2026-06-08"
refs = ["RFC-0001"]

[content]
description = "Fix bs>1 CP shared-KV cross-request KV contamination. Root cause: value-keyed 1-D page_inverse (build_slot_page_inverse + TAI scatter_page_inverse_kernel) is last-writer-wins, so at bs>1 two requests sharing a physical page alias each other's KV (real corruption when CURRENT_REUSE divergent write lands on the shared dense slot). Fix: per-request inverse stored flat [batch_rows*capacity], indexed req_id*capacity+lp, across torch-reference + TAI CUDA + Triton + prefetch paths. req_id sourced via slot//pages_per_request (build), repeat_interleave (logical_locs), companion-tensor-through-CP-split (current_locs). Validate on B300 with live GLM-5.1-FP8 PD harness (per-request byte-equality vs bs=1 golden)."
notes = ["Review findings to address: (C1) existing test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py ~13 call sites + benchmark use OLD signatures -> must update to new req_id args (do with Phase 5 regression test). (C2 SUSPECTED, verify) build_current_page_mask/build_current_loc_remap (fill_current_kv torch branch ~1658) are value-keyed GLOBAL, not req-id-aware: agent claims at bs>1 page_size=64 + CURRENT_REUSE a request can have a VALID history loc masked to -1 because another request's current page shares the value -> DROPS own KV (different symptom from reported 'got other's KV'; need to verify B's current page CAN collide with A's history page before fixing). (S3 Phase-2) the sparse-page kernel-probe _tai_current_slot_fill_supports_sparse_pages calls kernel with OLD args -> must update when kernel signature changes else TAI disabled forever. (P2 Phase-2 kernel design) do NOT use _broadcast_req_id_to expand+contiguous ([tokens,2048] int64 = 128MB churn/call); design CUDA/Triton kernels to take [tokens,1] req_id and stride it, then delete _broadcast_req_id_to. (S2) prefetcher loc/current_req_id None->zeros defaults are a latent bs>1 footgun (safe now, all callers pass real)."]

[[content.journal]]
date = "2026-06-08"
scope = "cp_shared_kv"
content = "Phase 1 (Python) DONE + opus correctness+perf review. Per-request 2-D page_inverse[batch_rows,capacity] across torch-reference (build_slot_page_inverse, 2 remaps, 2 fill_current, 2 builders, 3 orchestrators), req_id sourcing (build_token_loc_req_id repeat_interleave; get_cp_shared_kv_local_current_req_id companion through same CP valid-split; build_page_table_row_req_id row index), nsa_backend/nsa_indexer/prefetch call sites all wired, 5 files py_compile-clean. PERF: confirmed no new D2H syncs; memory +×batch_rows is per-batch-cached (5-20MB, fine); applied P1 (cache loc_req_id per-batch via get_cp_shared_kv_token_loc_req_id, was rebuilt per-layer)."

[[content.journal]]
date = "2026-06-08"
scope = "tai_kernel"
content = "Phase 2 (TAI kernels) DONE + opus correctness+perf review (no confirmed bugs; ABI 4-way agreement clean, bounds checked, Triton at parity, 1-D legacy degenerates, sparse-probe passes; perf-neutral). Changed: cp_shared_kv_materialize.cu (5 kernels+5 launchers per-request, derive batch_rows/pages_per_request/inner from shapes), _ops.cc (3 schemas +req_id), _cuda.py (4 launchers, +.contiguous() + index row-count guard), Triton .py (4 kernels+wrappers). Compact [num_rows]+row-stride req_id design avoids ~128MB/call broadcast (replaced _broadcast_req_id_to with _row_req_id_for_tai). req_id typed int64. Sparse-page probe updated to [1,cap]+zero req_id (else TAI disabled forever). Reworked tests to per-request (global-optimum, not patch): test_cp_shared_kv_materialize.py (2D reference inverse + req_id remap; bs=2 build/remap; NEW test_per_request_inverse_no_cross_request_contamination = the [[1,5],[9,5]] regression), sweep CUDA/Triton parity (batch_rows+req_id params), benchmark (per-request refs+calls). All py_compile clean."

[[content.journal]]
date = "2026-06-08"
scope = "tai_kernel"
content = "B300 KERNEL VALIDATION COMPLETE on b300-036 (g0036). Fresh container sglang-syh-cp from sglang:dev-cu13 (torch 2.11.0+cu130, sgl-deep-gemm 0.1.2, 8x B300), code rsynced to /mnt/beegfs/syh/sglang-stable (system was reinstalled, beegfs rebuilt from scratch). RESULTS: (1) CUDA JIT-compiled clean in 71s (the compile-check I couldn't do locally). (2) test_cp_shared_kv_materialize.py 16/16 PASS incl. the NEW bs>1 contamination regression -- per-request inverse correctly keeps shared page 5 distinct (inverse[0,5]=2 != inverse[1,5]=4) and fill places each request's current write in its OWN dense slot (no cross-request KV). (3) test_..._sweep.py 119/119 PASS = CUDA==Triton parity across batch_rows{1,3,4}, dtypes, page sizes, masks. (4) benchmark correctness-checked + per-request kernels 1.5-9x faster than torch ref (remap 6.6x w/ the i/inner divide, current_fill 9.1x, build_inverse 5.3x) -> the per-request overhead is negligible, matching the perf review. Fixed one brittle assertion in my own new test (scatter_ duplicate-index winner is nondeterministic, not 'last')."

[[content.journal]]
date = "2026-06-08"
scope = "harness"
content = "B300 full P/D harness bring-up. Old sglang container was a live server hogging g0036 GPU (~113GB x5 GPUs) -> caused prefill OOM; user authorized stopping it -> 8 GPUs freed. Relaunched via HOST tmux (not container-internal) running docker exec server scripts with full tee logs to /mnt/beegfs/syh/logs (correct pattern per user). NCCL init PASSED on all 8 ranks (~58s) -- the prior B300 blocker is resolved by DOCA fix + new NCCL env (NET_PLUGIN=spcx, bond0, roce_vf_rail0..7). Both servers loading: prefill g0036 in DeepGEMM warmup (131072 shapes), decode g0037 tp8/dp8. Built run_{prefill,decode,router}_b300.sh + validate_bs_gt1_contamination.py (concurrent shared-prefix requests w/ distinct markers; contamination = request leaks another's marker; compares bsN vs bs1 golden). NOTE: nsa_prefill_backend defaulted to trtllm (the parked server_args.py:1502 flashmla override bug is still active; orthogonal to the contamination fix which is backend-agnostic)."

[[content.journal]]
date = "2026-06-08"
scope = "deepgemm"
content = "CP DeepGEMM warmup fix (user-flagged) — CORRECTED the audit agent's analysis first. Agent claimed ALL GEMMs run at per-rank M (tokens/cp) so divide m_max by cp. VERIFIED WRONG: the MoE grouped GEMM M = sum(num_recv_tokens_per_expert) ~= chunked*topk/ep_size; GLM-5.1 topk=8==ep_size=8 -> MoE M ~= chunked (65536), NOT chunked/cp -- deepep all-to-all re-gathers ALL tokens regardless of CP. A blanket divide-by-cp would UNDER-cover the MoE grouped GEMM -> runtime JIT stall. Implemented the CORRECT per-kernel-type fix in deep_gemm_wrapper/compile_utils.py: _GROUPED_GEMM_KERNEL_TYPES (masked+contig) keep the FULL M grid; non-grouped dense shapes (attn q/k/v/o proj, dense MLP, indexer weights_proj = num_groups==1, all CP-per-rank) use a shrunk _BUILTIN_M_LIST_DENSE = ceil(m_max/cp_div) (floor 2048). Gated by _cp_dense_warmup_divisor: shrink only when nsa prefill CP in-seq-split on AND (no draft OR SGLANG_CP_DRAFT_SHARED_KV) -- because EAGLE DRAFT_EXTEND_V2 is CP only when CP_DRAFT_SHARED_KV set (verified via _is_cp_shared_kv_draft_extend include_v2=True). Result for this launch (chunked 65536, cp8): dense 131072->16384 Ms (8x fewer kernels) while MoE stays full; compiles, M-grid split verified. Assumption: no full-M num_groups==1 hooked GEMM (indexer fp8_mqa/paged_mqa_logits bypass the warmup hook anyway). Validate on next prefill launch (watch per-shape 'with N Ms' log + no first-request dense JIT)."

[[content.journal]]
date = "2026-06-08"
scope = "harness"
content = "B300 server bring-up round 3 fixes. (1) PREFILL CRASH root-caused: first forward (EAGLE target-extend) crashed in nsa_backend._forward_trtllm -> mla_quantize_and_rope_for_fp8 -> flashinfer rope 'k_rope_in.size(0)==nnz (512 vs 64)' (512=64 q_tokens x 8 cp_size). The k_rope carried all CP ranks' tokens but the per-rank rope kernel expects 64 -> a CP shape bug in the trtllm NSA path. Crash is in code I did NOT change and runs BEFORE the materialize -> contamination fix exonerated. Root cause = wrong backend: nsa_prefill_backend forced to trtllm by the parked server_args.py:1500-1503 bug (dp_size==1 && major>=10 set trtllm UNCONDITIONALLY, only branch missing the user_set guard). APPLIED the parked fix: guard with if not user_set_prefill/decode (whole-repo correct; the 'no DP attn' limit doesn't apply under CP) + added --nsa-prefill-backend flashmla_auto to run_prefill_b300.sh. Confirmed: prefill now nsa_prefill_backend=flashmla_auto. (2) DECODE: per user, removed --disable-cuda-graph, added --cuda-graph-max-bs 64 (dp8 per-rank); confirmed disable_cuda_graph=False cuda_graph_max_bs=64 nsa_decode_backend=flashmla_kv dp_size=8 dp-attention. (3) tmux: relaunched both as HOST tmux running 'docker exec -it' foreground (operable, attach via tmux attach -t prefill/decode; no orphan). Kill pattern: pkill -9 -f sglang (catches sglang::scheduler_DP* workers) + pkill -9 -f run_<role>, verify GPU=0 before relaunch (orphan/OOM was relaunch-then-kill race). Both relaunching."

[[content.journal]]
date = "2026-06-08"
content = "B300 live validation of the contamination fix surfaced TWO pre-existing prefill crashes (both fixed, prefill now UP at Uvicorn :50001, 0 crashes): (1) CP+EAGLE compute-padding: rebuild_cp_kv_cache all-gathers current k to global 512 under current-reuse, but the compute-padding cache-write branch (unlike the non-padding branch) never re-localized -> selector got 512 vs expected 64. Fix: new cp_localize_current_kv_to_compute_rows (utils.py) re-localizes via the batch-plan compute split before select_cp_local_valid_rows_for_cache_write; called in nsa_backend forward_extend. Validated live: 'valid_rows_select local_rows=64 expected_compute_rows=64' from a k_rows=512 forward. Whole-repo: index path correct by construction (separate local_key pre-all-gather); forward_mla:589 pre-rebuild. (2) flashmla fp8 dtype 'kv must be bf16': my server_args fp8 guard defaulted nsa_decode_backend=trtllm, and calculate_mla_kv_cache_dim only packs fp8 when NEITHER backend is trtllm -> unpacked fp8 -> flashmla_sparse skipped dequant. Fix: keep prefill/decode backend pair in same family (server_args.py). 3 files uncommitted: nsa/utils.py, nsa_backend.py, server_args.py."

[[content.acceptance_criteria]]
text = "Per-request flat page_inverse [batch_rows*capacity] in torch-reference (cp_shared_kv_runtime.py): build + 2 remaps + 2 fill_current + 2 dataclasses + 3 orchestrators thread req_id"
status = "done"
category = "added"

[[content.acceptance_criteria]]
text = "TAI CUDA + .cc schemas + _cuda.py launchers updated (4 live kernels gain req_id / pages_per_request); JIT rebuild clean on B300"
status = "done"
category = "added"

[[content.acceptance_criteria]]
text = "Triton parity + dead build_rank_major fail-fast guard + prefetch path fixed"
status = "pending"
category = "added"

[[content.acceptance_criteria]]
text = "bs>1 shared-prefix+current-reuse regression test + strengthened BS_GT1_DEBUG per-request assert"
status = "pending"
category = "added"

[[content.acceptance_criteria]]
text = "B300 live validation: per-request byte-equality vs bs=1 golden, TAI path, no contamination"
status = "pending"
category = "added"