[Stacked on #957] perf(kimi3): MNNVL CuTe-DSL finalize tail + K3 comm layer extraction - #1062
Draft
dongjiyingdjy wants to merge 43 commits into
Draft
[Stacked on #957] perf(kimi3): MNNVL CuTe-DSL finalize tail + K3 comm layer extraction#1062dongjiyingdjy wants to merge 43 commits into
dongjiyingdjy wants to merge 43 commits into
Conversation
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>
…ather Includes the symmetric cluster-barrier phase fix for the latent tail collective; the stray fused-o_proj wiring from the original development branch is not carried over. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The docs advertise the env fallback but the capability-only gating commit dropped it; restore the operational switch and log engagement so serve logs show which tail path is live. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Capability gating alone decides; the switch existed for the A/B run. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Above 2048 tokens the fused-lane AR keeps the replicated up-projection, whose GEMM dominates the tail. Instead each rank now up-projects only its hidden shard into the shared partial (addmm_ beta-add), and the shared all-reduce stitches the shards: same collective volume, 1/tp of the GEMM. Threshold from the measured tp16 crossover. Real-weight tp16 A/B at 8192 input tokens: median TTFT 685 -> 654 ms (c1), 2546 -> 2425 ms (c4). The MoE tail dispatch is folded into one _moe_tail method (fused decode kernel, column-parallel, fused-lane AR, separate-reduce fallback). Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Stage the routed/shared partials through cached symmetric-memory buffers and reduce with NVLS multimem (ld_reduce) instead of ring NCCL, keeping the column-parallel stitch. ld_reduce measures ~2x under NCCL at these widths and stays flat down to tens of tokens, so the stitch now engages everywhere above the decode range; the NCCL variant remains as the no-multicast fallback. Real-weight tp16 A/B, median TTFT: 8192-token prefill 632 -> 479 ms (c1) and 2300 -> 1749 ms (c4); 1024-token 340 -> 326 ms. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…gence Review findings: the stitch gate lacked a graph-phase exclusion, so decode graphs with bs > 16 captured the multimem path and a later buffer growth would free symmetric memory those graphs still reference; and per-rank eligibility probing meant one failing rank could silently fall back while its peers wait in the barrier. Eligibility is now agreed collectively at init (min over ranks), the stitch is eager-only, replaced buffers are retained for in-flight work, growth doubles, and buffers key on group. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Covers NCCL parity at both sides of every runtime threshold (m=17..4096, both tail widths), buffer growth/doubling with keep-alive retirement, width-cache independence, staging refusals, and the view-invalidation contract. 16 tests, torchrun world size 8/16 (job 8404: 16 passed). Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The graph-phase exclusion is unnecessary now that replaced symmetric buffers are retained: captured graphs keep a valid buffer forever, growth is rank-lockstep, and a capacity miss during capture falls back uniformly. Tiering by tokens keeps the dispatch a single-axis decision. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Tier choice was smeared across forward-side booleans (use_tail, lane enablement, skip_reduce, an inline norm+up-proj block) and an early-return chain re-testing the same conditions. A K3MoETailTier enum plus a pure selector in latent.py now decides once from named rank-uniform axes; every tier is a method receiving raw partials, the shared stitch is one helper, and the capture-time stage miss became an explicit re-selection instead of a hidden fall-through. Selector behavior is pinned by a truth-table test. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Review round 1: the latent-tail init kept the per-rank silent-fallback divergence the multimem path had already fixed (now agreed via MIN reduce, folded into the tier refactor); the multimem gate now refuses attn-DP/CP and world-of-one configs; staging refuses widths that break 16-byte alignment and documents the rank-lockstep capacity contract; and a capture/growth/replay test pins the retired-buffer safety claim that justified tokens-only dispatch. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
…lticast-tail # Conflicts: # python/tokenspeed/runtime/layers/moe/latent.py # python/tokenspeed/runtime/models/kimi_k3.py
GPQA isolates a real accuracy cost to the stitch itself: stitched trees score 173-176 (multimem reduce) and 170 (NCCL reduce) against 180-184 for main, while GSM8K is unaffected — the up-projection output is rounded to bf16 before joining the 16-way reduction, and its magnitude dominates the shared partials it sums with. The multimem tier now reduces both partials in-switch and keeps the replicated norm+up_proj+add3 tail, numerically congruent with the fused-lane path; the NCCL stitch tier is gone. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The stitch removal left dead machinery: the capture-miss re-selection was provably a constant (now a direct fused-lane call, making _moe_tail non-reentrant), the multimem tier threaded a lane that is always None, and the selector kept an unused hidden_size axis. Both AR tiers now share one _projection_tail so their numeric congruence is structural. New guards: MULTIMEM_AR is bounded to the measured 17..8192 window (larger batches take the join's grouped path), the gate refuses width-colliding configs, and the latent-tail gate matches the multimem gate's dp/cp rank-uniformity conditions. Truth table extended; stale comments and docs refreshed. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Buffer-capacity doubling is now clamped at the caller's dispatch ceiling (ragged prefills could otherwise double past 8192 rows of multicast-mapped memory per width); the latent-tail init agrees on the non-collective probe before entering the constructor's rendezvous, closing the one-rank-strands- peers window the MIN flag alone could not cover; the latent-tail distributed test skips via a collectively-agreed probe like the multimem one; and comment/docstring accuracy fixes (the width-collision guard is correctness, not tuning; staging accepts strided input; direction of the 2x claim). Signed-off-by: Nicolas <nperrin.ai@gmail.com>
In full-file order the latent buffer was already at capacity when the graph-replay test staged its growth trigger, so no retirement happened and the retired-buffer replay was silently untested. The test now resets its width symmetrically first and asserts the retirement occurred. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
A blind reviewer (no audit history in its prompt) caught what primed confirmation rounds waved through: _multimem_tail_ok read as the fused tail's gate (now _multimem_ar_ok, matching its tier); the TAIL_FUSION dispatch comment claimed a determinism need the eager tests disprove (the real invariant is tier stability across warmup/capture/replay); the docs scoped the in-switch tier to no-EP although its gate deliberately admits EP; a truth-table test was named for a capture-time miss it does not model; and the max_rows docstring overstated the clamp. Also: an init-time assert pins the fused-tail capacity below the multimem window, and the eligibility gate is spelled via tp_ep_size like its sibling. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The staging buffers are now pre-sized to the dispatch ceiling at init (collective, lockstep), so serving-time growth — a rendezvous inside a forward where one rank's allocation failure strands its peers — can no longer occur, and the symmetric-memory footprint is fixed and visible to memory planning. The vendored-kernel provenance note stops claiming 'unmodified' and enumerates the local deltas a re-sync must preserve; the single-mailbox cleanup-ordering hazard inherited from upstream is documented at the op; tests gain the skinny/dynamic dispatch boundary (m=5,6) and the width-alignment refusal. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The tier refactor moved SEPARATE_REDUCE's reduce/norm/up-projection out of the stream-fork scope, silently serializing work that previously ran concurrently with the shared-expert branch. The routed side now runs back inside the fork (only for this tier — the others need raw partials for their fused collectives), restoring the overlap. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Verified A/B (per-arm module-path and engage-count self-checks, after a harness bug had collapsed earlier cross-tree comparisons into one binary): the fused decode tail holds +12-13% TPOT at c1-c16 and the in-switch reduces hold -20% TTFT on 8192-token prefills against current main, but at decode bucket sizes the two staged reduces lose ~4% TPOT to the single fused-lane AR. The multimem floor moves to 256, above every decode-graph bucket, which also keeps decode numerics on the main-congruent path. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
A blind review round against the 256-token floor: the latent-tail init now dies loudly on a mid-constructor failure instead of falling back locally (a non-uniform failure there strands peers in its collectives, so the silent path was a hang in disguise); the spec-decode caveat on the multimem window is documented (bs*q tokens can re-enter it in-graph — correct but decode-suboptimal; a follow-up should add an is_decode axis); multimem_prealloc's return now reflects buffer readiness; the distributed M_VALUES straddle the real 256/8192 window edges; and stale floor comments, tier shorthand, and the measured -20% TTFT figure are reconciled. Runtime evidence from the round: both suites 26/26 at world 8/16, ~20k serving requests across tier-boundary/abort/endurance storms with zero errors and flat memory. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
A fp64-reference probe shows the fused tail carries a systematic negative output bias absent from the unfused path (~-2e-5/layer floor, -9e-5 at m=1 vs +-1e-5 unbiased), a plausible driver of the GPQA drop when the decode tail dominates: directional error compounds over 61 layers where random error averages out. The upstream-compatible bf16*bf16 norm square contributes a third of it (Jensen: squaring rounded values overestimates the variance, shrinking the output); fp32_internal also lowers the mean relative error from 3.2e-3 to 2.7e-3. The remaining bias is the mailbox's double rounding (store bf16, then residual add) — kernel follow-up. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Brings in the MLA prefill PDL disable (#990), the AMD gfx950 MoE input projections (#878), H200 Marlin SiTU support (#995), and the MLA 0.2.4 pin. Conflicts and their resolution: * layers/moe/latent.py — the K3MoETailTier enum/selector and main's new _marlin_moe_available() landed at the same spot; both kept. * models/kimi_k3.py — #878 inserted methods where this branch initialises the tier-0 tail, which a textual merge strands inside _latent_input_projections after its return; the init block stays in __init__. The SiTU entry takes main's Marlin-aware guard but keeps this branch's contract: the selected tail tier owns the combining reduction, so no internal all-reduce and no skip_reduce parameter. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The model mirrored plan fields under second names (use_trtllm_situ_moe, use_marlin_situ_moe), so every reader had to know both vocabularies and the H200 merge had to guard twice. One name per gate: the plan's own. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The try/except only logged and re-raised, which the scheduler already does for anything that escapes construction. What the reader needs is why the failure must not be caught at all — the constructor rendezvouses, so a rank that dies inside it has already stranded its peers — so that is now the comment. The window assert went with it: the selector tests TAIL_FUSION before the multimem range, so an overlap would resolve there anyway. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Eligibility for the two collective tiers was negotiated in six steps split across the model and kernel layers — a kernel-internal all-reduce for multimem, a collective prealloc, a local tail probe, a model all-reduce to agree it, the tail's own rendezvous, and a second model all-reduce to agree that. Each step was its own way for one rank to strand the others, which is why each carried a warning comment. Now both tiers vote locally, one MIN all-reduce agrees both votes, and only then do the collective allocators run. The staging miss inside the multimem tier is gone with it: its preconditions are exactly what the vote settled, so a miss is a contract violation and says so instead of silently re-deciding the tier that the selector already chose. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The existing tier test only checks which tier the selector returns, and an end-to-end eval only exercises the tiers its own shapes select — GPQA at ebs8 never reaches SEPARATE_REDUCE and the decode path never reaches MULTIMEM_AR, so one of the four had no numerical coverage at all. Each tier is compared against one fp32 reference, and the tiers are compared against each other directly: a request's tier depends on batch size, so two tiers individually inside tolerance but on opposite edges would show up to a user as the same prompt answering differently under load. Runs at world 4 (single node, skips the fused tail, which needs tp 8/16) and at world 8/16 for full coverage: 9 passed at world 8. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Two review findings from the PR. The op cache keyed only on the tensor contract, which every MoE layer shares, so all 61 layers received one mailbox. The gather releases its same-stream dependents before it rewrites sentinels; inside a layer the next collective's completion-scoped wait orders that against the following multicast write, but across layers that edge runs through every kernel in between and nothing guarantees it. A layer landing its shard in storage the previous layer is still clearing corrupts the output or leaves the gather spinning. The key now carries the layer index. Both tiers also voted on local multicast support alone. On a cross-host group without fabric or IMEX that reports available and then hangs inside the symmetric-memory rendezvous rather than falling back, so each now clears fabric_allocation_supported() when the job is wider than one host. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
All four latent-tail kernels hardcoded use_pdl=True at their launch sites, so --disable-pdl (which exports TOKENSPEED_DISABLE_PDL) reached everything in the model except them -- the operator had no kill-switch for exactly the launch mechanism most likely to be implicated when a PDL bug is suspected. The kernel package cannot import runtime modules, so the env var is the channel. Read it once at import into PDL_ENABLED rather than per launch: PDL is a process-wide launch property, and a per-launch lookup would have to become part of every compile cache key. Disabling is always safe -- the in-kernel griddepcontrol waits degenerate to no-ops and plain stream ordering is strictly stronger. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The equivalence suite re-derived norm+projection in torch for every tier, so it agreed with itself no matter what routed_expert_norm, kimi3_latent_projection_add3, kimi3_join_reduce_moe or the fused kernel's epilogue actually computed -- production code was never on the path it was supposed to be covering. The tiers now run through KimiLinearMoE's own _tail_* methods on an instance built by __new__ with just the attributes those methods read, which also adds the join tier (FUSED_LANE_AR), until now untested, and a pure-Python check of the selector boundaries that needs no GPU. Three things fell out of putting real code under the test. The fused tail keyed its mailbox on the layer index, but the NextN draft layer is built with index 0 and would therefore have been handed base layer 0's mailbox -- the exact aliasing the per-layer key exists to prevent. The key is now the module's weight prefix, which is unique across models. The draft's "take the plain reduce path" intent had been inert since those flags moved onto the execution plan: it wrote _fused_moe_ar on the module, which nothing reads, so the draft kept using a lane workspace shaped around the base model. It now updates the execution plan. PDL_ENABLED could not be read at the launch sites after all -- CuTe's JIT resolves names off self, not module globals -- so each kernel binds it in __init__ instead. Fixes a NameError that only surfaced once the fused tail actually compiled, which needs tp>=8. Validated at world 8 (cn04-05): 13 passed, and 13 again under TOKENSPEED_DISABLE_PDL=1. Max relative error against the fp32 reference, by tier and token count: tail_fusion 4.9-5.9e-3, multimem_ar 4.4-6.7e-3, fused_lane_ar 5.9-7.6e-3, separate_reduce 6.2-8.5e-3 -- so the collective's bf16 accumulation in the three unfused tiers costs about 1e-3 against the fused tail's fp32 one, not a correctness gap. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Backs the behavioral half of the previous commit out of this branch and leaves only the note. The draft's "take the plain reduce path" intent has been dead since those flags moved onto the execution plan -- it writes attribute names nothing reads -- but making it effective changes which tail tier the draft selects, and nothing here can run speculative decode to check that. Reinstating it belongs in its own change, validated against a spec run, rather than riding along on a perf branch. The aliasing risk that made this visible is already handled: the tail's mailbox is keyed on the module prefix, so the draft cannot be handed base layer 0's mailbox regardless of which tier it picks. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
The note still described the pre-fix world -- one process-wide mailbox whose sentinel cleanup overlapped the next layer. Mailboxes are per caller now, so the reason the PDL overlap is safe changed: reuse is a full forward apart and no rank can drift that far, because every layer's gather blocks on its peers. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
TOKENSPEED_DISABLE_K3_LATENT_TAIL=1 makes latent_tail_supported report False, so the tier selector sees capacity 0 and every size falls to the ordinary tiers. The tail replaces the entire tail collective with a multicast mailbox -- the branch's most invasive component -- and until now taking it out of the picture meant editing code and rebuilding, which is no way to run an A/B or to respond to a production suspicion. Written as an eligibility failure rather than a later skip so that the whole capability negotiation, allocation and tier map behave exactly as they do on hardware that cannot run the tail at all -- a configuration already covered by the tests. Signed-off-by: Nicolas <nperrin.ai@gmail.com>
Add a decode MoE-tail tier that runs finalize + sum-AR + latent RMSNorm in one collective launch (flashinfer's mnnvl_cutedsl backend, LL protocol), consuming the MoE kernel's deferred-finalize outputs. Decode TPOT 11.31ms -> 10.79ms (-4.5%) on 8xB300 TP8; AIME26 96.67 / GPQA-diamond 93.94 (no accuracy change). Dormant on flashinfer builds without the backend (sentinel import + capability vote); activates with flashinfer > 0.6.17. Extract kimi_k3_comm.py: the communication sites where K3's AttnRes/latent-lane semantics bypass the generic CommManager now route through K3AttnComm/K3MoeTailComm (per-group negotiated states, one MIN-vote instead of one per MoE layer, tier table and thresholds in one module). kimi_k3.py states semantics only; behavior unchanged (same-kernel routing verified by nsys and tier unit tests). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongjiyingdjy <87510204+dongjiyingdjy@users.noreply.github.com>
…utedsl The Kimi-K3 MNNVL CuTe-DSL MoE tail needs the mnnvl_cutedsl backend, merged upstream 2026-08-11 after the 0.6.17 tag and first shipped in nightly-v0.6.18-20260811. Direct wheel-URL pin; switch back to a plain version pin once 0.6.18 is published. Adapt the CuTe-DSL GEMM tuning config to the 0.6.18 autotuner API: per-input initializers moved from DynamicTensorSpec's fifth argument to TuningConfig.tensor_initializers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongjiyingdjy <87510204+dongjiyingdjy@users.noreply.github.com>
… key The packaged MoE tactic tables are keyed by flashinfer version (filename and embedded metadata), so the 0.6.18 nightly missed them all and fell back to heuristic tactics. Clone the three B300 tables under the nightly's key; the tactics themselves are carried over from the 0.6.16 sweep — decode validated unchanged, a fresh sweep on the 0.6.18 cubins should follow before relying on prefill/large-M shapes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: dongjiyingdjy <87510204+dongjiyingdjy@users.noreply.github.com>
Collaborator
|
how about we merge to #957 first |
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
3 participants
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
Fuse the Kimi-K3 decode MoE tail — finalize (top-k weighted expert combine) + sum-AR + latent RMSNorm — into one collective launch via flashinfer's
mnnvl_cutedslbackend (LL protocol), consuming the MoE kernel's deferred-finalize outputs. Extract the K3 communication layer (kimi_k3_comm.py) so all K3 comm sites (attention AR + MoE tail tiers) route through one negotiated module instead of inline branches inkimi_k3.py.perf/k3-latent-moe-multicast-tail; only the top 3 commits belong to this PR. Will rebase onto main once #957 merges.Results (8xB300, TP8/EP8, decode bs=1, 6k ctx)
0.6.18.dev20260811), including one autotuner API adaptation it requires.Commits
perf(kimi3): the fused tail (newK3MoETailTier.MNNVL_CUTEDSL, selector-routed, graph-phase gated, capacity from the workspace) + thekimi_k3_comm.pyextraction (one MIN-vote per process instead of one per MoE layer;kimi_k3.py−429 lines; behavior-neutral moves verified by tier unit tests + nsys).build(deps): pin flashinfer tonightly-v0.6.18-20260811(first build shippingmnnvl_cutedsl; merged upstream after the 0.6.17 tag) + adapt the CuTe-DSL GEMM tuning config to the 0.6.18 autotuner API. The tail is dormant behind a sentinel import on older flashinfer.build(deps): carry the K3 tactic tables forward to the nightly version key (tactics from the 0.6.16 sweep; decode validated unchanged, fresh sweep recommended before relying on prefill/large-M).Review notes
TOKENSPEED_MNNVL_CUTEDSL_MOE_TAIL(default off); PDL follows the globalpdl_enabled().Test plan
test_k3_moe_tail_tier.py(20 cases incl. the new tier's priority/capacity/graph gates)🤖 Generated with Claude Code