feat(wo_split=8): opt-in K-parallel W_O GEMV prototype

[Natfii]

Summary

Opt-in CUTE_WO_SPLIT=8 K-parallel W_O GEMV prototype for the production β-coop fused kernel on Qwen3.5-27B-NVFP4 (serve-cute.sh). Default stays wo_split=1 — no behavior change for callers who don't set the env var. Restricted to {1, 2, 4, 8} (the evidenced set).

Builds on PR #7 (audit + harness): the parity audit identified W_O GEMV (R2) at 4 active CTAs as the W_O bottleneck. This PR widens R2 to 32 active CTAs via K-parallel split (each CTA owns 1/8 of one KV-head's K range; gather of 32 partials in slot 0; new pre-W_O barrier for bx>0 W_O CTAs to wait on bx==0 attn producers).

Why this isn't duplicate work

No existing upstream vllm-project/vllm PR covers SM120 K-parallel W_O GEMV for the cute_paged backend (this fork's custom kernel). PR #7 (just merged) shipped the validation harness + parity audit as a separate, smaller-scope artifact; this PR ships the production prototype that consumes that audit's findings.

Performance evidence (scoped to measured region cluster + nsys)

All artifacts under benchmarks/nvllm/traces/wo_k_parallel_audit/2026-05-03-wo-split-8-prod/ — see summary.md for the full breakdown.

Region timing (5-completion synthetic ignore_eos load, R2/R4/R11/R12 only):

Region	wo_split=1	wo_split=8	Δ
R2 phase1_wo_gemv	14121 μs (4 active)	2360 μs (32 active)	−11761 μs (5.99×)
R4 grid_barrier_wait	15211 μs	1753 μs	−13458 μs (8.68×)
R11 phase1_pre_wo_wait	0 μs (mask empty)	250 μs	+250 μs (new consumer wait)
R12 phase1_gather_reduce	73 μs	167 μs	+94 μs (32 vs 4 partials)
Cluster (R2+R4+R11+R12)	29405 μs	4530 μs	−24875 μs (6.49× / −84.6%)

R4 shrinks because R2 finishes faster, so all CTAs hit the grid barrier sooner. Other regions (R0, R1, R3, R5–R10) unchanged.

nsys total-kernel (harness microkernel at production grid shape, 50-launch median):

kernel_cutlass__wo_kernel_body: 13715.248 → 1598.064 μs = −12117 μs (8.58×), stddev 2102 → 5.7 μs (much more stable at wo_split=8). Per feedback_vllm_profiling, vLLM V1 nsys against EngineCore is blocked by CUPTI inheritance — the harness microkernel under nsys is the canonical production-grid measurement; harness reproduces production W_O+gather math bit-exactly (verified below). NCU NOT used (per repo policy).

End-to-end serving (Qwen3.5-27B-NVFP4 PIECEWISE, max-num-seqs=4, GSM8K-50 full-think workload): per-question median 65.7s → 62.3s = −5.2%, 50-question wall 3760s → 3664s = −2.6%. Translation from kernel speedup to e2e is bounded by what fraction of kernel time the W_O cluster represents at this workload.

Quality (parity, not improvement)

GSM8K-50 full-think (seed=42, max_tokens=512, timeout=600s):

Run	Accuracy	Errors / timeouts
baseline wo_split=1	48/50 (96.0%)	0
changed wo_split=8	47/50 (94.0%)	0

−1 question is within ±2% noise; 0 errors both runs (well above the kernel-change ≥30/50 floor). Quality parity, not improvement.

Bit-exact correctness gate

Kernel reproduces reference_split_order(wo_split=N) from docs/research/2026-05-03-w-o-k-parallel-harness/torch_reference.py bit-exactly at both wo_split=1 AND wo_split=8 (max_abs == 0.0 against wo_output[seq, 0, :] post-gather). Methodology: V=constant trick (FP8 V-cache = +1.0 makes Phase 1 attn output deterministic).

Test commands run

# Bit-exact gate
.venv/bin/python /tmp/wo_split_repro.py --wo-split 1 --seed 4242  # PASS, max_abs=0.0
.venv/bin/python /tmp/wo_split_repro.py --wo-split 8 --seed 4242  # PASS, max_abs=0.0

# Region timing (both wo_splits)
bash docs/research/2026-05-02-beta-region-breakdown/run_breakdown.sh
.venv/bin/python docs/research/2026-05-02-beta-region-breakdown/extract_regions.py \
    --buf <region_timings.npy> --slice-ctas 8 --num-k-tiles 8 --num-seqs 1 \
    --tick-source globaltimer --wo-split N --num-kv-heads 4 \
    --out <region_breakdown.csv>

# nsys (harness microkernel at production grid)
nsys profile -t cuda,nvtx -o <out>.nsys-rep \
    .venv/bin/python docs/research/2026-05-03-w-o-k-parallel-harness/run_harness.py \
    --slice-ctas 8 --wo-split N --launches 50

# GSM8K-50 (both wo_splits)
.venv/bin/python scripts/gsm8k_eval_50.py \
    --api http://localhost:8000/v1 --model default \
    --n 50 --seed 42 --max-tokens 512 --timeout 600 \
    --save <out>.json --label task<10|11>_wo_split_<1|8>_<baseline|changed>

Full reproduction in summary.md.

Implementation notes

• serve-cute.sh:118 adds -e CUTE_WO_SPLIT="${CUTE_WO_SPLIT:-1}" and writes the value to /tmp/c2_diag/ENV (the existing sentinel pattern). vllm/nvllm/models/qwen3_5.py:52 reads it at module import via os.environ.setdefault. Without this sentinel pattern, vLLM V1's EngineCore subprocess strips the env var (per feedback_vllm_enginecore_env_strip).
• phase_e_kernel.py: K-parallel W_O block lifted out of the legacy bx==0 && by<num_kv_heads parent into a new gate bx<wo_split && by<num_kv_heads. K-range integer-divide slicing per torch_reference.py:443-446. Slot index by*wo_split+bx. New pre-W_O arrival counter (_phase_e_coop_pre_wo_arrival_count) zeroed per-launch.
• Region-timing taxonomy extends from 11 → 13 regions atomically (host + kernel constants + reducer if/elif chain). R11 in WAIT_NOT_WORK_REGIONS only (avoids first-match-wins misclassification as work). R12 in new DYNAMIC_SINGLE_CTA_REGIONS.
• Cooperative-launch invariant preserved (CLAUDE.md rule 8: atomic-counter spin-wait barriers must run cooperative=True).

Caveats

• nsys evidence is harness microkernel at production grid shape, not full vllm kernel. The harness reproduces production W_O+gather math bit-exactly so the speedup transfers, but other phases (attn QK, MLP) are NOT included in the 8.58× number. Region-timing R2 5.99× is the production-side measurement.
• R12 active CTA count varies (1–3 across runs) depending on concurrent decodes at dump time.
• CUTE_WO_SPLIT=8 requires the sentinel-file workaround to reach EngineCore; bare -e CUTE_WO_SPLIT=8 doesn't propagate.
• Region timing was captured at synthetic ignore_eos load. Per-region speedups at GSM8K may differ (workload sensitivity); the e2e numbers are the GSM8K result.

AI assistance disclosure (per AGENTS.md §1)

This PR was developed with AI assistance (Claude Opus 4.6/4.7). The submitting human reviewed every changed line, ran the test commands listed above, and is responsible for defending the change end-to-end. No code-agent autonomy beyond the bounded plan in /home/natfii/.claude/plans/sorted-crafting-rainbow.md.

🤖 Generated with Claude Code