[WIP] perf(k3): optimize decode - #834
Draft
nperrin-fr wants to merge 24 commits into
Draft
Conversation
…code The single-CTA packed-key sigmoid-bias top-16 kernel was gated to gfx950; NVIDIA decode routed through the grouped multi-pass path (7.0us vs 3.7us per MoE layer on B300 at the K3 decode shape). Enable it for NVIDIA and keep it selected under the deployed defaults: - pass the CDNA-only waves_per_eu launch hint only on AMD - fuse the static expert-location dispatch map into the kernel (the server's default static dispatch previously forced every K3 decode step onto the grouped path even at ep_size=1) - support bf16 route-weight output so the sidecar adapter's cast stays a no-op; selection and normalization stay in FP32 bs1 decode TPOT on 8x B300 TP8: 10.71ms -> 10.53ms (93.4 -> 95.0 tok/s). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Replace the fused-AR + replicated up-projection tail with three fused stages inside the decode CUDA graphs: one kernel for the latent all-reduce + RMSNorm + shared-expert reduce-scatter, a per-rank-sharded up-projection (1/tp of the replicated weight traffic) whose epilogue multicast-stores the shard into every rank's mailbox over NVLS, and a barrier-free Lamport gather. Buffers come from stock PyTorch symmetric memory; the CuTe-DSL kernels are vendored under thirdparty/cute_dsl/ (see the package header for provenance and license). Engages only inside the CUDA-graph phase on plain TP8/TP16 (no EP), SM100-family, bf16, hidden 7168 / latent 3584, and at most 16 tokens; eager forwards and larger batches keep the existing fused-AR tail. TOKENSPEED_K3_MULTICAST_TAIL=0 disables it. Standalone (8x B300, graph replay): 8.15us vs ~18.6us per layer for the current lamport-AR + replicated-gemv tail. bs1 decode TPOT: 10.53ms -> 9.96ms (95.0 -> 100.4 tok/s), stable P99 == median. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Drop the enablement environment variable: the tail already has a full eligibility chain (platform/TP/shape checks, a symmetric-memory multicast probe, and a per-layer try/except fallback to the fused-AR tail), so the switch added no selection power — optional knobs just rot unconfigured. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Vendor the warp-specialized SM100 attention-residual forward kernel (online-softmax mix over the block snapshots + prefix in one launch, TMEM V-row caching, PDL chaining, optional in-place delta accumulate and fused out-norm; provenance and license in the file headers) and register it as the CUDA solution for the attn_res family. The split partial/combine arrangement stays as the portable fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Add an mnnvl-structured one-shot kernel to the vendored AR-fusion fork (single multicast store through the switch, local-only lamport poll, deterministic rank-ordered reduction, self-rotating device flags) and instantiate the existing FusedOp epilogues over it, covering the plain, residual+RMSNorm and AttnRes-combine patterns. The workspace comes from torch symmetric memory; arming is a collective capability probe with the IPC lamport workspace as automatic fallback, both in the kernel-package manager and the runtime backend's plain all-reduce. The latent-lane norm pattern stays on the IPC path (measured slower on the wide lane). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…ernels One launch now quantizes the NoPE query and commits the latent row into the fp8 FlatKV cache (byte-exact vs the previous pair, with the same NaN sanitization). The q/kv dual RMSNorm moves to a two-CTA Triton kernel, and the SiTU/sigmoid-mul/gated-RMSNorm glue kernels get launch shapes tuned for decode (fewer, lighter CTAs that co-reside with aux-stream work). Derived offline from Inductor codegen of the same modules; no runtime compile dependency. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…code Every Triton kernel on the K3 decode path gains an enable_pdl flag with conservative griddepcontrol fencing (wait before the first predecessor-output load, launch_dependents after the last store), bit-identical output either way; enablement stays with the caller and defaults off pending an uncontended A/B. Also check in the load-shaped regression tests distilled from the bs64 stability investigation: an eval-shaped multi-rank chain test for the deferred-delta arrangement and a sidecar contract test that pins the deferred-finalize out-of-bounds write (xfail on the affected sidecar build; the reland gate for both suspended features). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…itive The router GEMV at one token underfills the dense kernel (~1.8 TB/s on 896 rows); route it through the row-per-CTA streaming GEMV instead (4.0us vs 6.2us uncontended), shifting the main-stream front left while the top-k stays hidden under the latent down-projection on the aux stream. Also add a single-read merged [gate|latent] front kernel with an auto strategy: measured slower at decode (the top-k serializes behind the merged sweep) so it stays off at bs1 and is available for mid-batch prefill fronts; design source credited in comments. bs1 decode: 9.91 -> 9.77ms TPOT; GSM8K 0.978 / GPQA 0.9242 at bs64, zero faults. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Plumbs the flag from the public try-wrapper through the registry adapter to the megafuse launcher (default off). The runtime call site stays unwired: the serving A/B measured noise-to-worse, consistent with every PDL experiment on this decode graph. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…ather The multicast tail's final gather now takes an optional addend and folds the layer's residual stream into the copy (fp32 add, bf16 round — bit identical to the eager add, asserted with torch.equal in the distributed test), removing one elementwise launch per MoE tail layer. bs1 decode: 9.76 -> 9.62ms TPOT; bs64 eval clean (GSM8K 0.975; GPQA within its temp=1.0 sampling band for a bit-identical change). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The NVLS one-shot kernel's .cu was absent from the trtllm_comm build group, so fresh builds silently lacked the symbol and the capability probe fell back to the IPC workspace. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The lane pattern was routed back to the IPC workspace on measurement; the test still asserted mnnvl parity and failed on the supports() guard. It now asserts the routing decision itself. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The rejected GEMM+AR experiment's runtime hooks rode along in the residual-fusion commit via a concurrent edit and referenced a module that no longer exists, breaking model import. Restore the model to the tail-fusion state. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Rewrite the single-launch KDA decode (conv+SiLU, f_b gate GEMV, QK L2-norm, delta-rule recurrence against the paged fp32 state) in CuTe DSL: warp-per-row-group state streaming with a cp.async prefetch front, cluster-synchronized conv-window commit, and PDL with producer- independent loads issued before the dependency wait. Ties the previous kernel at single-token decode and wins from four tokens up (+24% at 16); recurrence-phase design derived from an Apache-2.0 packed-decode reference (header note), the fusion math ported from the in-tree kernel. The Triton megafusion it replaces is removed outright: its in-place conv-window shift raced across the column-split programs once the grid exceeded one co-resident wave, double-shifting windows at large decode batches (reproduced deterministically at 64). The verify variant writes per-position windows to its own scratch and is unaffected. Platforms without CuTe DSL fall back to the unfused conv/recurrence chain. bs1 decode holds 9.62ms; bs64 eval clean with GSM8K 0.9788 / GPQA-diamond 0.9293 (both above the campaign gates). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The fused single-read projection+activation kernel was gfx950-only; enable it for NVIDIA at one token with a B300-tuned launch (one output row per CTA, 1024-wide K tile — the tile must be a power of two that divides the hidden size, or the unmasked fast path reads past K). Replaces the cublasLt GEMM + separate SiTU launch on the shared-expert fork: 8.0us vs 9.2 standalone, bs1 decode 9.62 -> 9.52ms TPOT (105.0 tok/s); bs64 eval clean, GSM8K 0.975 / GPQA-diamond 0.9394. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011oPfcso15ywRoX9NTbbycp Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Fold rmsnorm(o) * w * sigmoid(gate) into the CuTe DSL KDA fused decode kernel's epilogue, normalizing the fp32 accumulator directly (one fewer bf16 round-trip than the separate kernel). NV=1 owns the full V row; NV=2 exchanges sum-of-squares partials over distributed shared memory on the existing single cluster-barrier phase, with the state writeback deferred past the barrier to keep the store stream drained. Handoff is attempt-and-verify via KdaGatedNormRequest: the model stashes weight/gate/eps, the wrapper fuses only when shapes/dtypes/strides match and marks the stash consumed; otherwise the host-side norm still runs. Armed only in the CUDA-graph phase: the eager 1-token-prefill path faults when epilogue fusion arms there (same dark corner as the latent tail; root cause tracked separately). Decode: 9.52 -> 9.47 ms median TPOT at bs1 (4k/500, TP8). Quality: GSM8K 0.9795, GPQA-diamond 0.9293 at bs64. 9 new kernel tests cover both NV paths, strided gates, fallback, and graph capture. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…rm epilogue The NV=2 fused-norm epilogue issued its st.shared::cluster partial-sum exchange without a completed cluster-barrier phase, violating the CUDA requirement that the target CTA has begun execution before its shared memory is accessed, and its peer partial landed in a smem word still read by the phase-2 normalizer (cross-CTA WAR). Under irregular SM scheduling (a 1-token prefill interleaving with graph replay - which every prefix-cached bench warmup produces) this surfaced as rare cudaErrorIllegalAddress plus sporadic per-head output-scale corruption. Give the peer partial a dedicated smem slot and restore the early cluster arrive with an epilogue-side wait before the DSM store, keeping the deferred state writeback ordering. compute-sanitizer: 294 hazards + invalid peer writes -> 0. Stability: 14 bench runs under concurrent 1-token-prefill injection and double-bench load, zero faults (unfixed build faulted within 4 runs); TPOT unchanged at 9.47 ms; GSM8K 0.9750, GPQA-diamond 0.9141 at bs64. Also sync the capture tests' side-stream warmups (missing wait_stream ordering) and add a sanitizer-gated regression test that fails on the unfixed kernel. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Graph-replay A/B on B300 shows the fused gated-norm epilogue wins only while launch overhead dominates: +0.03 to +0.09 ms/step for B<=16 but -0.08 at B=32, -0.19 at B=64, -0.46 at B=128 (x69 layers) - at larger batches the epilogue's per-row reductions serialize inside the kernel while the standalone norm amortizes its launch cost. Decline the stash for t > 16 so large-batch decode keeps the separate norm kernel; the decision remains shape-deterministic and CUDA-graph stable. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…lective Two of the three cluster-barrier uses inside the collective's per-token loop arrived on every CTA but waited only on CTA0's first warp. Arrive/ wait counts must stay symmetric on every CTA: with the loop wrapping at m >= 2, the skewed phase counters let a later iteration's pre-exchange wait match an already-completed phase, so the st.shared::cluster store into the sibling CTA ran without the peer-entered guarantee - an intermittent illegal memory access. The window opens exactly when a prefill admits a second sequence and the next step replays the bs>=2 decode graph, which matches every observed crash (including the earlier bs64-eval faults previously attributed elsewhere by elimination); pure bs1 replay never wraps the loop, which is why bs1-only soaks always passed. Make all three waits unconditional. Cost is ~0.01 ms/step at bs1 (noise level). Soak: 5x bs1 + 3x forced-bs2 benches and a full bs64 eval, zero faults; GSM8K 0.9757, GPQA-diamond 0.9192; 7/7 distributed tail tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…ecode Restructure the CuTe DSL fused decode's state handling around batch size. At NV=1 (batch*heads >= 96) a single thread stages the full 64KB [K,V] state slab with one 1D bulk-TMA copy behind an mbarrier (re-armed every launch, so graph replays stay self-contained), and pass 2 recomputes the decayed row from shared memory instead of holding a 64-register strip that spilled. At NV>1 the register-resident pass 2 remains (faster at small batch). The f_a gate operand loads as bf16x4 vectors (the front's top long-scoreboard stall), and the gated-norm epilogue now fuses at every batch size - the NV=1 epilogue is block-local, which removes the large-batch crossover that motivated the old t<=16 gate. Recurrence-phase design notes and provenance are documented in the module docstring. Auto-selected; misaligned pool pitches fall back to the layout-agnostic NV>1 path. Standalone (graph replay, us/layer): B1 5.82->5.51, B8 6.32->6.13, B16 9.18->7.91, B64 27.9->19.0. e2e decode: 9.47-9.48 -> 9.35 ms median TPOT at bs1 TP8 (x6), forced-bs2 11.86 (x3). Quality: three greedy GSM8K-128 evals at bs16 match the baseline failure distribution (0.9844-0.9922); full bs64 eval runs as the final gate. 34 kernel tests including envelope-strided pools and large-batch fused-norm parity. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Collapse branch-added multi-line comment blocks to single-line invariants, drop the narrative doc addition, and remove test files whose gated features were reverted or quarantined (their history and repro recipes live in the campaign records). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
nperrin-fr
force-pushed
the
dev/decode-k3
branch
from
July 29, 2026 04:21
a547f00 to
81b1edf
Compare
Both landed during the campaign but ended up with no runtime callers: the online-softmax v2 attn-res family lost its e2e A/B and carries an unresolved race (quarantined from the build since; the sources and history remain on the backup branch), and the merged-front GEMV primitive was superseded by the rowcta router path. Drop the kernels, their wrappers, registry entries, and tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The v2 removal dropped the trailing side-effect imports along with the v2 entry point, so no attn_res backend registered and every K3 decode raised NoKernelFoundError at startup. Import-only smoke checks and test collection both pass without them; only a live server exercises the registry. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The guard coerced a falsy value to -1, so a request asking for prompt logprobs from position 0 - the natural way to ask for all of them - was accepted and answered with no logprobs at all, while 1 and above were correctly rejected. Compare against None explicitly and handle the batch (list) form. Also correct the FlatKV note about recycled pages: the MLA decode kernel masks the K side but feeds V-tail rows into the PV MMA verbatim, so the writer's sanitize and zeroed pages are what keep a previous tenant's bytes out of the output, not a kernel-side fixup that does not exist. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013y4dRxfUyRVFZWhf5ZSkJE Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Collaborator
Author
|
this is one is probably too big and I'll break it up later. |
nperrin-fr
added a commit
that referenced
this pull request
Jul 29, 2026
At decode batch sizes the NoPE fp8 path was launch-latency bound: one kernel assembled the fp8 query, a second committed the latent row into the fp8 FlatKV cache. mla_nope_query_kv_fp8 does both in one launch, byte-exact against the previous pair and with the same NaN sanitization. The pool gains mla_fp8_commit_view, which runs set_mla_kv_buffer's structural location checks and hands out the raw per-layer fp8 view only when a direct dtype-cast store matches the write path's semantics. It returns None otherwise, and the caller keeps the two-kernel path -- so non-FlatKV pools, non-fp8 caches and non-unit-stride views are unaffected. Split out of #834 (dev/decode-k3). The original commit also carried a dual-RMSNorm rewrite and launch-shape tuning for the SiTU/sigmoid-mul glue kernels; those are independent and follow separately. Tests: 10 kernel cases (fp8 parity vs the unfused pair, NaN sanitization, PDL on/off) and 5 pool cases (view gating and location checks). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WKty3aLeTKb2vhRu4kvRgc Signed-off-by: Nicolas <nperrin.ai@gmail.com>
nperrin-fr
added a commit
that referenced
this pull request
Jul 29, 2026
At decode batch sizes the NoPE fp8 path was launch-latency bound: one kernel assembled the fp8 query, a second committed the latent row into the KV cache. mla_nope_query_kv_fp8 does both in one launch, byte-exact against the previous pair and with the same NaN sanitization. The destination is the existing get_key_buffer(layer_id): both MLA pools expose an absolute-token-indexed latent buffer there, which is exactly what the kernel indexes with out_cache_loc. No new pool API, and no capability probing at the call site -- the surrounding branch already establishes MLA kernel backend, fp8 data_type and k_scale == 1.0, and the kernel asserts the rest (row width, unit last stride) while its store casts to the destination dtype. Split out of #834 (dev/decode-k3). The original commit also carried a dual-RMSNorm rewrite and launch-shape tuning for the SiTU/sigmoid-mul glue kernels; those are independent and follow separately. Tests: 10 kernel cases (fp8 parity vs the unfused pair, NaN sanitization, PDL on/off). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WKty3aLeTKb2vhRu4kvRgc Signed-off-by: Nicolas <nperrin.ai@gmail.com>
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.
Summary
Test Plan