perf(k3): five cuts on the CP8 128k prefill superstep - #1026
Open
xiaguan wants to merge 6 commits into
Open
Conversation
The CP8 128k prefill anatomy (deep rank 6,130 ms) charged 385 ms to `k3_land`, the f32-partial -> bf16 landing after every dense projection. The TileLang `land_batched` family landed one element per thread with 2-byte stores and reached 2.3 TB/s on the 12288-wide chunk landings (550 us per launch at bucket 16896) — a third of HBM rate. Replace the family with a hand-written CUDA kernel (`csrc/k3/k3_land.cu`): eight columns per thread, two float4 loads per segment and one 16-byte store, grid-stride over rows x n/8, batch a runtime value (no per-bucket instantiation; rows not 16-byte aligned fall back to a scalar path). The arithmetic is the retired kernel's spelling — ascending-segment f32 sum from segment 0, one round-to-nearest-even cast — so at split_k = 1 the landing is the bare cast and the output is bit-identical. The launcher keeps its signature; no executor change. Same-shape A/B (nsys, CP4 64k prefill gate, bucket 16896): nt=12288 n=12288 550 us -> 190 us (1.25 GB, 6.6 TB/s) nt=18432 n=18432 287 us (1.87 GB, 6.5 TB/s) nt=67584 n=33792 532 us (3.4 GB, 6.4 TB/s) A bf16 GEMM epilogue (cublasGemmEx with a bf16 C) was tried first and rejected: cuBLAS picks a different kernel for a bf16 destination, so the accumulation order moves and the 4-layer golden fixture flips at its 3-ULP-margin steps (same deviation distribution against the reference — median 1 / p90 3 / max 12 bf16 ULP on both — but the gate's exact-match contract is per kernel selection). The retired TileLang instantiations (15 spans x 15 buckets) come off the build. Gates (tray03, pruned 224-expert checkpoint): golden_decode 13/13 with every step exact, paged_kv 3/3, spec_verify 6/6, cp_prefill 2/2. Signed-off-by: xiaguan <751080330@qq.com>
The chunked-prefill conv path landed the q/k/v projection to bf16, built every row's 3-slot window with strided device copies (row t's slot j from landed row t-3+j, the carried window for the first rows), ran the batched decode conv over the materialized windows, and copied the last commit row's successor window out as the carry. The CP8 128k anatomy charged that to 182 ms of conv plus 235 ms of cuMemcpy2DAsync per deep rank — the whole non-doorbell gap budget. Replace it with one launch per group (`csrc/k3/k3_conv_silu_chunk.cu`) that reads the f32 partial rows t-3..t directly, takes the positions before the segment from the carried window, and writes the carry itself. Each block owns a column slice and a run of 16 rows: the taps sit in registers and the window slides through them, so every partial element is read once and every output written once. The arithmetic is the batched conv kernel's spelling term for term — `bf16(0 + p)` landing, ascending-tap f32 products, `bf16(ca)` then `sb * (1 / (1 + expf(-sb)))` — under the same -O3 and no fast-math, so the output is bit-identical. The chunk's separate conv-input landing and the `conv_window`/`conv_window_next` slabs go away; the halo path and the decode conv are untouched. Same-context A/B (nsys, CP4 64k prefill gate, device 0): DtoD copies 25,466 / 2,593 ms -> 4,766 / 102 ms conv-input landings -385 ms (3 per KDA layer gone) conv kernel @ bucket 16896 877 us -> 476 us per launch The land kernel also lands onto a zero accumulator now, so a -0 partial lands as +0 exactly as the retired TileLang kernel did. Gates (tray03, pruned 224-expert checkpoint): golden_decode 13/13 every step exact (chunked prefill agrees with the per-token walk), paged_kv 3/3, spec_verify 6/6 (multi-group verify segments), cp_prefill 2/2. Signed-off-by: xiaguan <751080330@qq.com>
Under context-parallel prefill every rank attended its own segment against the whole prefix it had assembled, so rank r paid for r + 1/2 segment-pairs of FMHA while rank 0 paid 1/2: at CP8 128k the deep rank's 2,132 ms of FMHA was ~1,090 ms of pure imbalance (the kernel itself runs at peak), and the lockstep superstep waits for the deepest rank. Stripe the triangle instead (`cp::k3_cp_stripe_kept`): off-diagonal pair (q, k) is attended by the owner q when q - k is odd and by the key holder k when even — the helper already expanded its own keys for its diagonal, so only queries and results move. Each rank publishes its queries once per layer (riding the latent window), attends its kept segments densely and its own segment causally into the f32 log-sum-exp accumulator, then runs (cp_size - 1) / 2 return slots: in slot i it attends the i-th owner it helps — that owner's published queries read in place by the FMHA over the fabric, output + LSE landing in a return slab — and merges what its own slot-i helpers returned, again in place. Two return slabs alternate so slot i+1's FMHA launches inside window i's consume without waiting on window i's readers; every query read is an FMHA launched before a return publish the owner waits on, which orders the owner's next overwrite behind it. Loads go from r + 1/2 to at most (r + 3) / 2 (8 ranks: deep 7.5 -> 4.5, ideal 4). Window kinds gain explicit rank-set bitmasks; the whale slab grows the query and return regions (~1.47 GB per rank at seg_cap 16896); the schedule is a pure function of cp_size and unit-tested for single coverage, mirrored window sets and the leveling bound. Same-context A/B (in-process CP4, pruned checkpoint, tray03): FMHA per rank @64k 1,064 / 1,744 / 2,408 ms -> 1,735 / 1,068 / 1,733 CP4 wall @65536 4,171 / 4,151 ms -> 3,888 / 3,864 ms (-285) deep rank -28% FMHA, exactly the 2.5 / 3.5 the schedule predicts. Gates (tray03): cp_prefill 2/2 (CP4 vs CP1 argmax and boundary tokens equal at 65536/65535, logits inside the noise floor), golden_decode 13/13, paged_kv 3/3, spec_verify 6/6, unit tests for the schedule and the slab layout. Signed-off-by: xiaguan <751080330@qq.com>
The windowed-context merge folded one bf16 FMHA window into the f32 accumulator with a 128-thread block per (q, head) row and one scalar element per thread: 1.6M blocks, 2-byte loads, 1.4 TB/s. In the CP8 128k superstep that was 144 launches x 1.5 ms = 220 ms of a 5,071 ms rank — and with the striped triangle, every returned peer window pays it again over fabric. Now 16 lanes carry a row with one uint4 bf16 load and two float4 accumulator loads/stores each, and a 256-thread block walks 16 rows; every lane of a row recomputes the two weights from the same lse pair and lane 0 alone writes the merged lse. The per-element update keeps the scalar original's expression, so the fold rounds exactly as before. The finalize leaves the accumulator the same way. Gates (pruned 224-expert checkpoint): golden_decode 13/13 exact, paged_kv 3/3, spec_verify 6/6, the kernels crate's windowed walk vs the single causal call, cp_prefill CP4 vs CP1 at 64k rel_l2 within floor. CP4 64k in-process wall 3,888 -> 3,785 ms. Signed-off-by: xiaguan <751080330@qq.com>
add2, mul_sigmoid, situ and o_norm_gate were TileLang batched kernels walking one element per thread with 2-byte loads and stores — 1–2 TB/s on the chunked-prefill rows, 220 ms of the CP8 128k deep rank (add2 91, o_norm_gate 75, situ 36, mul_sigmoid 21) for what is 6 TB/s of traffic. They are hand-written CUDA now, sixteen bytes per thread per operand and shape-agnostic, with the retired kernels' arithmetic kept term for term: TileLang's bfloat16_t is cutlass's, whose `+` is __hadd, `*` is __hmul and cast is cvt.rn, so the bf16x2 forms round identically; the f32 chains keep expf/tanhf/rsqrtf and IEEE division under the same -O3 without fast-math. o_norm_gate's 128-wide xor butterfly (64 and 32 through shared memory, 16..1 by shuffle) is reproduced pair for pair by sixteen lanes of eight columns — lane xor 8/4/2/1 then slots j^4/j^2/j^1 — so the norm scale is the same bits in every lane. eps is a runtime argument (K3_RMS_EPS at the one call site) instead of a compiled constant. Forty TileLang instantiations retire with the family (280 -> 240). Gates (pruned 224-expert checkpoint): golden_decode 13/13 exact — decode runs the same add2/mul_sigmoid/situ — paged_kv 3/3, spec_verify 6/6, cp_prefill matches + sweep, CP4 vs CP1 at 64k rel_l2 unchanged to four digits (8.706e-2 / 6.991e-2). CP4 64k in-process wall 3,785 -> 3,639 ms. Signed-off-by: xiaguan <751080330@qq.com>
Five cuts on the CP whale lane (vectorized landing, in-place conv window, striped FMHA triangle, vectorized lse_merge, vectorized bf16 elementwise family): fleet CP8 128k TTFT 6,387 -> 5,021 ms, 65k 2,810 -> 2,231, 16k 1,150 -> 946. Records the cuBLAS bf16-epilogue rejection (kernel selection changes accumulation order; the golden gate flips at 3-ULP steps) and the open cp_prefill same-process grid-sync flake. Signed-off-by: xiaguan <751080330@qq.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Five independent cuts on the K3 CP whale lane, one commit each, ordered by the nsys anatomy of a CP8 128k superstep (tray14 deep-position rank, 6,130 ms; archive
~/code/bench_results/2026-09-01-k3-cp8-128k-gap-anatomy/, summary indocs/models/k3/cp-lane-design.md). The anatomy overturned two assumptions: the gaps were not host time (host-starved 0 ms — launch-queue backpressure and doorbell waits), and the FMHA kernel was already at peak — ~1,090 ms of its 2,132 ms was deep-position imbalance from the causal triangle.e1fdd881dad55e4bt-3..tin place instead of landing + 3 strided window copies + batched conv + carry copy80f12bb8003b5b5be82f072d+/*/cast reproduced (__hadd/__hmul/cvt.rn), o_norm_gate's 128-wide xor butterfly reproduced pair for pair;epsbecomes a runtime argumentRejected on the way: landing the GEMM in bf16 via cuBLAS epilogue. It is not bit-identical — cuBLAS picks different kernels for bf16 output, the accumulation order changes, and the 4-layer golden gate flips at its 3-ULP-margin steps (
CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTIONmakes no difference). The gate was kept as is; the landing kernel was fixed instead.Evidence
Fleet CP8 (tray03 + tray14, pruned 224-expert checkpoint,
~/bench-tools/k3-ttft-quick.py, same probe before/after, min-of-2/3):1e3b6f23)Superstep wall on the profiled rank 6,130 → 5,071 ms after cut 3 (re-profiled; per-family table in the doc). 24 greedy tokens after the 128k prompt continue the prose verbatim ("Chapter 584: …" following chapter 583). In-process CP4 64k wall 4,171 → 3,639 ms.
Gates per commit (pruned checkpoint, bare host):
golden_decode13/13 step-exact,paged_kv3/3,spec_verify6/6,cp_prefill(CP4 vs CP1 at 16,896/16,895 and 65,536/65,535: argmax and boundary tokens equal, rel_l2 within floor — unchanged to four digits across cuts 4–5, consistent with bit-identical kernels), and the kernels crate'swindowed_walk_matches_single_causal_call.Known flake (pre-existing on
main, not introduced here)cp_prefill's two tests run back to back in one process sometimes hit a DeepGEMMgrid sync timeoutincp4_prefill_ttft_sweep(60 s barrier, arrival count stuck at 151). Reproduced 2/3 on80f12bb8with cut 4 stashed and 1/3 onmain(1e3b6f23) built in a clean worktree; each test passes reliably as its own process, which is how the gates above were run. Root cause not yet found (the counter stalls short of the SM count, so some CTAs never become resident within the 60 s barrier — a same-process teardown/residency interaction, not a numerics issue).Not in this PR
attnres vectorization (~80 ms), doorbell overlap (~144 ms), the residual FMHA imbalance (4.5 vs 4 pairs; helped dense segments over fabric Q run slower than local causal ones) and the mega-MoE wait accounting — the next round, which needs a fresh profile.
🤖 Generated with Claude Code