K3 model line: decode E2E at ep_size 4 with MegaMoE(situ) transport by xiaguan · Pull Request #867 · pegainfer-project/pegainfer

xiaguan · 2026-08-14T03:17:08Z

What

Brings up the K3 model line end to end: a 93-layer hybrid-attention MoE
(69 KDA linear-attention layers + 24 MLA layers, latent MoE with situ
activation, MXFP4 experts) decoding coherently at --k3-ep-size 4 on
4×GB300, with MegaMoE(situ) as the sole MoE transport — zero host
collectives in the steady state, 43.3 ms/step full-depth.

Every stage landed behind a numerics gate; the EP path is bitwise
identical to single-rank execution, and MegaMoE is bit-parity gated
against the reference kernel.

Commit walk (each stage is reviewable standalone)

Commit	Stage	Gate
`873e6833`	DeepGEMM bump to the openinfer branch (situ activation)	existing glm52 MQA shim re-verified against #377 template refactor
`b49a9f04`	`k3` kernels feature: FP8×FP4 masked grouped GEMM AOT shim + TileLang codegen pipeline	TMA/barrier byte accounting; layout asserts vs host stub
`88ec37ec`	`pegainfer-k3` crate skeleton (config / loader / model line / scheduler)	loader shape audit vs checkpoint index
`53edfa98`	MoE decode chain (local dispatch, activation quant, situ, combine)	per-op parity vs reference
`5c418f8f`	batched TileLang surface, 13 kernel families	per-kernel certification vs reference impl
`a56bd97e`	device model build (loader output → kernel-ready rank model)
`77b10bf0`	batched decode executor — single-rank decode E2E, CUDA graphs per (bucket, parity)	golden gate 38/40 exact; misses sit on the reference's own ≤1-ULP coin-flip steps
`7c19a7a0`	free-running EP4 — fixed-chain MoE oracle, per-rank autonomous schedulers, padding-is-protocol	EP4 vs EP1 bitwise (40 tokens + 163,840 logits)
`e96db007`	docs: full-depth EP4 decode E2E record	93-layer serve, coherent greedy output, idle-padding soak
`33d22f80`	MegaMoE M1: torch-free AOT shim + ep1 swap-in	bit parity vs the reference Python kernel
`e84818eb`	MegaMoE M2: ep4 (peer symm slabs, fixed 92-launch chain)	EP4-mega vs EP1-mega bitwise; traffic invariance bitwise
`789c4d58`	retire the stepwise EP scaffolding (net −1163 LOC)	golden dual-run: mega 39/40 production gate + chain 38/40 anchor (test-only)
`44f1cb7a`	decode step profile doc	see `docs/benchmarks/k3-ep4-decode-profile.md`

(faa93368 is rustfmt drift only.)

Architecture notes for review

Free-running per-rank engines (glm52 pattern): no coordinator; each
rank is an autonomous scheduler partition; the only coupling is the step's
fixed kernel chain; empty batch = padding step is part of the protocol.
MegaMoE is the only production MoE transport — no env flag, no
fallback. The single-rank masked chain survives only as a test-side
numerics anchor behind a doc-hidden constructor.
Peer-reach pitfall now documented in code + docs: cudaMallocAsync slabs
need both cudaDeviceEnablePeerAccess and cudaMemPoolSetAccess
(pool grant before slab allocation).
Launch-count guard (92 launches/step, armed when ranks > 1) replaces the
collective ledger from the retired chain.

Testing

Kernel/chain gates run in CI-shape (cargo test) where no GPU is needed;
GPU gates (golden_decode, ep_mega_oracle, parity gates) documented in
docs/models/k3/bring-up.md with exact commands. Checkpoint paths come
from PEGAINFER_K3_TEST_224 / PEGAINFER_K3_TEST_896 env vars.
Full-depth serve A/B on 4×GB300: mega 43.3 ms/step vs NCCL chain 54.6
(−21%); profile table and levers in the benchmarks doc.

Not in this PR (queued)

Absorbed-MLA latent kernel + paged KV pool (lifts the max_ctx = 128
bring-up cap) — design under review.
Prefill beyond decode-shaped, MTP, EP16 full checkpoint.

Moves the submodule ecbbe74 -> 8cbafe8: upstream #364 -> #377, the same DG_NO_TORCH patches rebased, plus the merged K3 situ-activation patch for the SM100 FP8xFP4 mega-MoE kernel. The grouped-GEMM headers are identical across the two revisions; all drift is in the MQA logits templates, where the leading 'bool kIsFP4' became a trailing 'typename qk_dtype_t' plus a repositioned 'bool kIsMXSF' — glm52_deepgemm_mqa.cu is adapted with semantics-preserving argument updates (kIsMXSF=false, qk=e4m3). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

…degen New 'k3' feature (independent of 'moe'): - csrc/k3/k3_deepgemm_fp8_fp4_grouped_sm100.cu: AOT instantiation of the DeepGEMM masked grouped GEMM with an FP4 B operand (float_e2m1_unpacksmem_t, gran_k B=32) at K3 expert shapes — W13 N=6144/K=3584, W2 N=3584/K=3072, group counts {56,112,224}, SMs {148,152} — plus a torch-free FP4 K-major TMA descriptor helper and a k3_fp4_sf_prepare kernel (checkpoint u8 e8m0 exponents -> DeepGEMM's packed MN-major i32 SF layout). - Build-time TileLang generation (three tiers: live generation -> pre-generated dir -> NOT_SUPPORTED stubs) for the K3 decode-surface kernels defined in pegainfer-k3/kernels/: batched router top-k and the attention-residual scores/mix pair, E in {224,896}, batch buckets 1..128. Launch geometry recovered from the TVM host stub is asserted against the analytic grid so a TileLang bump fails at build time. - FFI declarations + safe ops wrappers + KERNELS.md rows. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

…skeleton New pegainfer-k3 model line (Kimi K3: 69 KDA + 24 MLA hybrid attention, latent MoE with MXFP4 routed experts, situ activation, no MTP), serving text-only through the multimodal wrapper checkpoint: - config.rs: architecture constants + probe_config_json (glm52 style); the only variable facts are the routed-expert count {224, 896} and the EP topology. - weights.rs + weights/{load,staging,context}.rs: manifest with per-rank name generation and exact coverage validation, name->dtype/shape tensor contract grounded in the real shard headers, EP expert sharding into packed [gate; up] regions at H2D time, MXFP4 payload/scale bytes uploaded raw (byte-isomorphic to the FP4 B operand), A_log padded-lane trim, mmap-lifetime guards and pinned double-buffer staging. - model_line.rs + scheduler/: ModelLine (config probe, --k3-ep-size CLI, CLI-derivable serve plan) and a step-contract Scheduler over a slot-budget registry with the model execution behind a StepExecutor trait; phase 1 launches unwired (requests fail with an explicit message). Server registration behind the 'k3' feature. - docs/models/k3/bring-up.md + index and model-table rows. cargo test -p pegainfer-k3 --lib: 26 passed (checkpoint-gated tests probe both expert counts when the checkpoints are present). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

…ine) Five graph-safe kernels around the masked FP8xFP4 grouped GEMMs: route metadata (race-free per-expert compaction, deterministic token*topk+slot entry order), gather fused into the W13 per-token-group FP8 quant (no expert-major bf16 staging), situ+mul+requant on the masked layout (f32 situ over the bf16 GEMM output, gate|up split-half), UE8M0 scale pack, and weighted combine (topk-slot-order f32 accumulation, no atomics, single bf16 round). GPU-gated chain test runs all seven steps against a dequant reference: rel_l2 ~1.7e-3, scales and slot maps bit-exact, two passes bit-identical; a flipped-nibble negative control fails at rel_l2 ~1.1. First live exercise of the masked FP4 GEMM's TMA/barrier accounting — clean at both shapes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

The generator now emits the batched decode kernel set — 13 families x batch buckets {1,2,4,8,16,32,48,64,96,128}, 610 instantiations (norms, partial landings, elementwise, conv+silu, KDA delta rule, MLA attention with per-slot ctx lengths over a fixed-CAP slot cache, router top-k, attn-res scores/mix, combine) — vendored verbatim from the certified batched kernels (bitwise gates in the source repo's check_batched2.py). TileLang gemv / expert_gemv / packed_expert_gemv are removed end to end: dense projections run on cuBLASLt and routed experts on the masked grouped-GEMM chain, so the bs=1 case is the B=1 bucket of the batched kernels. TMA host-stub machinery is replaced by a hard assert (no body may lower to tensormap creation). Generation parallelized (19 s wall for 610 lowerings); tier-1 build 1m14s, stub tier green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

K3RankModel materializes a rank's loaded tensors into per-layer bundles through a pure, host-testable slot plan: zero-copy adopts for lone projections (checkpoint [out, in] row-major — the executor's cuBLASLt uses OP_T), contiguous row-concat fusions (wbig/wsm/wfu/wsh/wgu), f32 sw taps computed host-side bit-identically, conv weights transposed to (taps, inner) matching the reference spelling (verified bit-pattern against the shard), and a K3ExpertBank that runs k3_fp4_sf_prepare at build (e8m0 -> packed i32 SF in the masked-GEMM layout, source regions dropped immediately; region byte lengths ensure!-tied to the GEMM contract). Layer truncation takes-and-drops by name so the loader's ensure_consumed invariant holds unweakened. GPU-gated 4-layer build on the real checkpoint asserts per-class byte totals (12.42 GiB rank 0 at EP4); 34 host tests cover plan coverage, routing and fuse order. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

- pegainfer-k3/src/executor/: engine-blueprint launch sequence (step.rs), state pools + scratch arena (buffers.rs), banded/offset cuBLASLt wrapper (gemm.rs), per-(bucket,parity) CUDA graphs default-on with eager escape (PEGAINFER_K3_CUDA_GRAPH=0); buckets up to B=128, seat i = row i. - tests/golden_decode.rs + fixture: 4-layer greedy golden replay 38/40 exact, 2 misses on the reference's own <=1-ULP coin-flip steps (structural FP8 activation-quant noise floor, documented); row independence, graph-vs-eager and multi-slot gates bitwise-exact within a bucket. - model_line launch() builds one K3Executor per EP rank; server smoke passes. - weights.rs: layer-limited load planning for truncated bring-up builds (PEGAINFER_K3_LAYERS); kernels/generate.py: per-TU debug-helper renaming to fix duplicate external-linkage symbols when linking multiple TileLang TUs. - docs/models/k3/bring-up.md: executor section + noise-floor finding. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

…e-rank - Free-running per-rank engines (glm52 free-running-dp architecture): the scheduler calls decode() unconditionally (empty batch = padding step), the only cross-rank coupling is the step's fixed collective chain (4 NCCL collectives per MoE layer, count guarded by an ensure! ledger). - EP numerics, bitwise by construction: protocol-max allgather of latents + topk (padding rows constructively zero/-1), expert-windowed route metadata (base=0 delegates bit-identically for single-rank), local masked chain over the global batch, dense entry-major scatter (disjoint support), bf16 allreduce (0+x exact in any order), entry combine with the accumulation spelling copied verbatim from the masked combine. - Discipline from glm52/qwen3 lessons: EP forces eager (no collectives under capture until warmup/two-phase precapture lands), comms minted on the rank's own thread after a condvar-timeout id rendezvous, all ranks load weights before any comm init, EP step errors fail-stop (log + exit), and ensure!(ep_size * max_batch <= masked_cap). - Gates (4-layer, one process each): EP4 rank0 vs EP1 — all 40 greedy tokens and all 163840 final logits bit-identical, idle peers and busy peers both; single-rank golden + moe chain gates untouched and passing. - 8-layer EP4 serve smoke: 4 partitions serve concurrently, idle free-running padding steps hold for 90s with no wedge. Full-depth serve pending free GPUs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

- k3_mega_moe_sm100.cu: torch-free AOT shim for the DeepGEMM SM100 FP8xFP4 mega MoE — host replica of the symm-buffer layout and block/pipeline heuristics, gran-8 gate/up weight interleave + UTCCP SF transpose build kernels, gran-32 e4m3 activation quant, and the launch (9 TMA descriptors incl. the hand-built FP4 one, cluster 2, PDL, 152-SM grid). 3 block configs x {situ, swiglu} x ranks=1 instantiated. - Executor: PEGAINFER_K3_MEGA=1 swaps routed_experts to quant_x -> symm write -> one fused launch; expert bank builds exactly one weight form; CUDA graphs stay on (ranks=1 barriers are grid-local, descriptors persistent). mega + ep_size>1 is rejected until the EP form lands. - Parity gate: output bit-identical to the Python-validated kernel for situ and swiglu (negative control trips on a single flipped bit); byte-exact transformed-weight checks. Golden replay under mega: 39/40 exact, 1 step inside the documented noise floor. Perf: 5-8% faster per layer than the masked chain at B in {1,16,128}, eager and graphs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

…e-rank - Shim templated on world size (ranks {1,4}); ranks=4 at 56 experts/rank, 384 max tokens/rank, 209.8 MiB slab. One block config fixed from the protocol max for every rank and every step: sender-side symm writes were verified to depend only on layout constants (never the receiver's block config), and the fixed config makes per-row tiles traffic-independent. - In-process symm plumbing: per-device slabs + pointer-table rendezvous. Peer reach needs BOTH cudaDeviceEnablePeerAccess and cudaMemPoolSetAccess (stream-ordered allocations ignore the former alone; the pool grant must precede the slab allocation) - documented, first failure mode was an illegal address at first launch. - EP mega groups carry no NCCL communicator: the kernel's internal NVLink barriers pair the ranks (collectives/step=0), num_tokens=0 padding steps still serve local experts for peers and hit every barrier, so the free-running contract holds unchanged. EP mega stays eager; ranks=1 keeps graphs. - Gates: EP4-mega vs EP1-mega BITWISE (all 40 tokens + 163840 logits, despite different block configs); traffic invariance at world=4 BITWISE; parity/golden/chain/ep_oracle/lib regressions all green (golden under mega: 39/40 exact). - 93-layer serve: coherent greedy text, 4 concurrent, 60s idle hold clean; 43.3 ms/step single stream vs 54.6 for the NCCL chain on the same box (21% faster; the doc's earlier 72 ms/step chain number did not reproduce and was corrected). - ep_fatal now also writes to stderr (a fatal exit with no logger installed was previously silent). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

…ansport - The NCCL fixed-chain EP path is deleted: routed_experts_ep and its pack/allgather/scatter/allreduce/entry-combine kernels, the comm/id rendezvous phase, the collective ledger, the chain-era ep_size*max_batch cap ensure, and tests/ep_oracle.rs (its bitwise gates live on in ep_mega_oracle). ep.rs is now an honest mega rendezvous; a launch-count guard (92 mega launches/step, armed above one rank) replaces the ledger — its only job is making a lagging rank name itself, since the fused barrier's 60s device timeout is the detector. - The single-rank masked chain retires to test-only numerics anchor: K3ExecutorConfig.moe_transport defaults to Mega, the chain is reachable only through a doc-hidden test constructor (single-rank enforced), its scratch is allocated only when selected, and the PEGAINFER_K3_MEGA env flag is gone (an environment variable must not flip production numerics). golden_decode keeps both transports: mega 39/40 (production gate), chain 38/40 (anchor vs the certified reference). - Net -1163 LOC. Docs: bring-up.md rewritten around mega as the production story with a three-lesson history paragraph; index.md K3 row updated. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

…ottleneck Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 44f1cb7aa0

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-08-14T03:22:01Z

+        &mut self,
+        slot: SlotId,
+        prompt: &[u32],
+        _params: &SamplingParams,


Reject non-greedy requests until sampling is wired

When a request uses temperature > 0 with top_k != 1, the frontend accepts it and passes the sampling parameters here, but prefill_inner discards them and every subsequent step reads argmax_indices. Such requests therefore silently produce greedy output instead of the requested distribution, which can corrupt sampling-based experiments; either retain per-slot parameters and invoke the sampler or reject non-greedy requests explicitly.

AGENTS.md reference: AGENTS.md:L3-L5

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-08-14T03:22:01Z

+        _params: &SamplingParams,
+    ) -> Result<u32> {
+        ensure!(slot < self.max_batch, "K3 slot {slot} is out of range");
+        ensure!(!prompt.is_empty(), "K3 prefill needs at least one token");


Handle empty prompts before entering the EP fatal path

With --k3-ep-size 4, an empty tokenized prompt reaches this check because neither the frontend nor admission_refusal rejects an empty prompt_tokens vector; the resulting error is then converted by StepExecutor::prefill into ep_fatal, which calls process::exit(1). A single empty completion request can therefore terminate the entire EP server even though no distributed step has begun, so this input should be rejected per request before the group-fatal wrapper.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-08-14T03:22:02Z

+                finish_or_retire(active, FinishReason::Stop, emitter);
+                continue;
+            }
+            emitter.push_tokens(&mut active, &[first], &[]);


Reject or populate requested completion logprobs

For a /v1/completions request with logprobs > 0, the stepped bridge records that count in Request.logprobs, but the K3 scheduler never examines it and always pushes an empty logprob slice. StepEmitter converts that to None, so the successful response silently omits the requested probabilities; until logits are converted into token logprobs, admission should fail these requests as unsupported rather than returning incomplete output.

AGENTS.md reference: AGENTS.md:L3-L5

Useful? React with 👍 / 👎.

… in admission Signed-off-by: xiaguan <751080330@qq.com>

…ments Everything removed was verified to have no reachable launch or call site: the GEMV-era combine_land TileLang family, the split-K-8 and routed-situ instantiations (610 -> 430 AOT kernels), two orphaned rank-context helpers, five unused crate dependencies, and the unused K3_SPLIT_K constant. Comments and KERNELS.md rows that still described the retired NCCL collective chain or a future DeepEP path now describe the MegaMoE transport, and the mega kernel surface gets its missing KERNELS.md section. Signed-off-by: xiaguan <751080330@qq.com>

KineticCafe/actions-dco decodes the full commit comparison, whose file entries carry blob_url: null for submodule bumps; its strict decoder rejects the whole payload, so any PR touching a submodule can never pass (v3.1.0 additionally failed on not-yet-linked author emails, upstream issue #216). A DCO check only needs the commit list, so query exactly that. Signed-off-by: xiaguan <751080330@qq.com>

xiaguan and others added 14 commits August 14, 2026 03:17

chore(kernels): rustfmt drift in build.rs and k3_tilelang.rs

a420e6a

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

docs(k3): record full-depth EP4 decode E2E (free-running serve verified)

09cdc39

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

docs(benchmarks): K3 EP4 decode step profile — backbone GEMM is the b…

082e85d

…ottleneck Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: xiaguan <751080330@qq.com>

xiaguan force-pushed the feat/k3-bringup branch from 44f1cb7 to 082e85d Compare August 14, 2026 03:17

chatgpt-codex-connector Bot reviewed Aug 14, 2026

View reviewed changes

xiaguan added 2 commits August 14, 2026 03:27

style(k3): fix clippy gates — let-else in build.rs gencode, then_some…

5da1677

… in admission Signed-off-by: xiaguan <751080330@qq.com>

xiaguan force-pushed the feat/k3-bringup branch from 5f3134b to 3a35cc0 Compare August 14, 2026 03:43

xiaguan force-pushed the feat/k3-bringup branch from e2c276a to 71c88db Compare August 14, 2026 03:53

xiaguan merged commit a5c1884 into main Aug 14, 2026
13 checks passed

xiaguan deleted the feat/k3-bringup branch August 14, 2026 03:57

Navigation Menu

Uh oh!

K3 model line: decode E2E at ep_size 4 with MegaMoE(situ) transport - #867

K3 model line: decode E2E at ep_size 4 with MegaMoE(situ) transport#867
xiaguan merged 17 commits into
mainfrom
feat/k3-bringup

xiaguan commented Aug 14, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Aug 14, 2026

Uh oh!

chatgpt-codex-connector Bot Aug 14, 2026

Uh oh!

chatgpt-codex-connector Bot Aug 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

xiaguan commented Aug 14, 2026

What

Commit walk (each stage is reviewable standalone)

Architecture notes for review

Testing

Not in this PR (queued)

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Aug 14, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant