feat: add kimi-k3 (kimi-linear) serving support - #308
Merged
Conversation
Contributor
Author
|
Validation recorded internally; CI is the public record. |
Andrewxu313
force-pushed
the
tairanxu/k3-prefill-core
branch
from
August 12, 2026 21:45
1dd89e2 to
8d961b7
Compare
Andrewxu313
force-pushed
the
tairanxu/k3-prefill-model
branch
from
August 12, 2026 21:46
4db6b16 to
0dac884
Compare
Model directory (config, initializer, parameter server, planner, model graph, serving modules, tokenizer, wrappers, parallel-strategy manager), the registration seams (get_initializer / get_parallel_strategy_manager), and the model kernels: the causal-conv1d CUDA extension and the fused BF16 MoE triton kernel. Core-side wiring (worker, kv_cache, registries) lands in the companion core PR.
INSTALL.md gains the NUMA dev-header prerequisite (numa.h is required by the core_engine JIT build) and a pointer to the new docs/troubleshooting.md, which collects the commonly-met build, runtime, and model bring-up problems.
The conv1d extension (setup.py:281) placed _C_causal_conv1d.so under batchgen_kernels/conv1d/ in the wheel, but the subpackage was missing from packages/package_dir, so conv1d/__init__.py (causal_conv1d_fwd / causal_conv1d_update wrappers) was not installed.
M1 bring-up fix batch, model side: F1-F4, the FA3 (flash_attn_interface) import, meta-parameter replacement via _replace_param, the attention_mask kwarg drop for worker parity, and KDA state-pool slots read from the engine config instead of the environment.
Planner half of the M1 fix batch: F6 planner overrides. The worker-side half (F5, decode heartbeat) is in the companion core PR.
P-6(c) chained-serving test caught serving-vs-model-oracle drift up to 6.6e-2 (gate 2e-2): model.py and the fla reference apply no activation between the f/z low-rank projections. With silu removed all 24 checks pass at <=5.9e-3 (verified via in-process monkeypatch diag).
First smoke run: 6-7 empty decode ranks crashed at KimiMoEGate scores.view(0, -1) (model.py:201) via moe_forward_serving:523. Empty ranks now build empty routing and still drive the streamed-expert ring in lockstep. F3 covered the KDA attention wrapper; the MoE gate is the FFN-side counterpart.
load_tokenizer() constructs tokenizers with no model path, so KimiLinearTokenizer fell back to _K25_ASSETS_DIR — the server rendered every prompt with the kimi_k25 chat template (different media/tool/think handling) instead of the checkpoint's own template. Vendor the checkpoint's chat_template.jinja + tokenizer_config.json as kimi_linear/assets/ and default there; k25 assets remain only the tiktoken.model fallback (byte-identical vocab).
Template and special-token config now fail fast when missing (never substitute another model's); the byte-identical tiktoken.model fallback remains but warns. First application of the fail-fast principle; repo-wide audit tracked separately.
configure_decoding builds shards once (empty routed_expert copy task -> no decode H2D streamer); decode dispatch on planner decode_moe_mode (default resident_ep, streamed fallback kept launchable); prefill untouched. tests/kimi_linear: fused_moe + conv1d suites at the standard BF16 gate (1e-5 + 1.6e-2*|ref|, outlier <1e-4, BF16-TC reference) per D2.
Model half of M5.1: KimiLinearKDAWrapper pools become views of KDAStateGPUManager memory (single fixed-address allocation per pool — the capture requirement); KDASlotManager kept as a stateless facade so PSM/worker call sites are untouched. Adds graph-readiness tests + single-layer KDA decode capture smoke. The manager itself lands in the companion core PR.
The KDA state manager requires an indexed device (cuda:N), not the bare "cuda" string; the M5.1 tests now pass one.
One graph per layer covering the attention span (input_ln + attn +
residual + post_ln); MoE runs eager between replays because resident-EP
performs NCCL collectives every MoE layer. Layer 0 is dense
(first_k_dense_replace=1, no collectives) so its whole layer folds into
one graph. Buckets {1,2,4,8,16} captured lazily; static page-table /
cache-seqlen / slot buffers refreshed in place (KDA slots bound to the
manager's persistent buffer, one staging call per step for all 20 KDA
layers). Padded and warmup rows use a reserved KDA scratch slot, never
-1 (fla indexes OOB before the pool base). Mode via planner
decode_graph_mode + batchgen_debug.kimi_decode_graph_mode
(graph|eager|compare); default stays eager until the M5.5 gates.
Eager-fallback paths warn with reason and re-warn at decade counts so a
persistent fallback cannot look like healthy operation.
The graph path pads to the bucket and parks padded rows on the reserved scratch slot; the unpadded eager reference never writes there, so a whole-pool compare fails on scratch garbage. Compare live slots, and add an explicit per-slot offender check so a real divergence cannot hide behind the exclusion.
Gating installation on the planner mode meant batchgen_debug could not switch graph/compare on a live server — every experiment needed a full cold restart, against the project batch-flag debug policy (glm5_moe_mode precedent). Eager now installs the adapter as a pass-through (buckets capture lazily, so no cost); new off value installs nothing at all.
The MLA decode spans bake `_k_cache[physical_layer]` into the captured
graph. That slice's base is `data_ptr + layer * num_pages * page_stride`,
so it is a function of the FULL cache geometry, not just the base pointer.
The worker re-creates the KV manager per batch job and sizes it from that
rank's free HBM, so num_pages varies between jobs. A re-init that reuses
the same base address with a different page count left data_ptr unchanged
while every slice above physical layer 0 moved -- by an amount
proportional to the layer index. `_signature` tracked only data_ptr, so
the graphs were never dropped and replayed against relocated K-cache
slices.
Symptom: MLA spans diverged from eager while all KDA spans stayed bitwise
identical (KDA pools are separate fixed tensors), with the error growing
monotonically with physical KV-layer index. Per rank the dirty MLA layers were always 0/7 or 7/7, never a mix, and
re-capture healed it -- both signatures of baked-in state rather than
per-step buffer staleness. Graph-mode accuracy fell well below the eager
control and recovered after a forced re-capture of the same build.
Fix: carry k_cache.shape and k_cache.stride() in the signature so a
geometry change drops and re-captures. This adds no steady-state
re-captures -- the KV manager is re-created per batch job, not per step.
Also corrects the comment in the MLA span that asserted the false
invariant ('a fixed address for the life of the manager'), and adds
regression checks that the signature carries shape/stride plus a
non-vacuity check that the per-layer slice offset really does move with
num_pages.
The old gate never reached the failing condition: it ran only inside a single batch job, so the KV manager never changed shape under a live graph, and the bug passed every span bitwise while materially degrading MMLU accuracy. This regrows the fake K cache mid-run and re-runs the full schedule, asserting the adapter drops, re-captures, and still matches eager logits. The old tensor is deliberately kept alive across the swap, so a graph that failed to drop would replay against still-valid but stale memory -- wrong logits and no crash, which is exactly how it presented in production.
Growing the cache allocates at a new address, which even the old pointer-only signature detects -- so the end-to-end logits assertion passed with the fix reverted, making it decorative. Verified by reverting the fix on the node: the signature checks failed but 'logits still match eager' passed. Production drifts DOWNWARD (page counts shrink between jobs), and freeing then re-allocating smaller is what makes the caching allocator return the same address. Pointer-reused + shape-changed is the only combination that reproduces the bug. The test now reports whether reuse actually happened rather than assuming it, so a run that fails to reproduce it says so instead of silently claiming coverage.
The shrink went below the highest assigned page (assign_pages hands out 0..15 for 8 seqs x 2 pages), so the eager pass hit IndexError: index 14 out of bounds for size 13. Harness bug, not a product bug. NUM_PAGES 16 -> 24 and the shrink target is now derived from the pages actually assigned (used+2) with an assert that headroom exists, so the test cannot silently stop shrinking if the slot config changes.
Adds the mechanical three-way check (checkpoint index <-> model state_dict <-> module_shapes/name map) that would have found the K3 blocker offline in seconds, plus the K3 tensor map, MXFP4 layout helpers and the parameter-server/initializer wiring that makes the checkpoint loadable. The reconciler reports four buckets rather than a boolean: a model parameter with no source (which stays at its init value -- the GLM-5 Q/K RMSNorm incident), a module_shapes entry with no host source (allocated but never written, so the slot keeps zeros then the previous module's weights), a checkpoint tensor with no destination (the K3 blocker), and shape/dtype mismatch. It runs offline from the index plus a meta-device model, and also as a startup assert. Quantized weight_packed/weight_scale pairs are treated as one logical parameter, and unmapped tensors must be declared ignorable with a reason -- an unlisted one is an error. Two review findings changed the design rather than being papered over. The gate originally defaulted to index-only mode, which cannot see a wrong shape; the default is now shard-header mode, verified on the real released checkpoint (a transposed o_proj gives ok=True in index mode and ok=False with 24 tensor mismatches in header mode). And the offline suite was blind to transposes, so the fixture was rebuilt with per-tensor shape and dtype from an independent template table -- mutation-proved by transposing o_proj, which turns six previously-green tests red. Also fixes a real 48B bug found en route: o_norm.weight is BF16[128], not F32, confirmed from the shard headers. Reconciler placed under batchgen/models/ so it falls inside MODEL_ALLOW_RE and the whole change is one model PR rather than a core/model split. 77 passed, 2 skipped, CPU only.
K3 has no jinja chat template and no chat_template key -- its template is Python, apply_chat_template at tokenization_kimi.py:357. That third case is easy to miss and is exactly how a model ends up borrowing a sibling's assets, which is the 2026-07-31 incident. Vendors the eight metadata files from HF into models/moonshotai/kimi_k3/assets/, md5-verified byte-identical to the served checkpoint, and builds KimiK3Tokenizer against them. Resolution is by explicit path and raises rather than falling back; kimi_linear no longer claims the kimi_k3 registration. The design changed on a finding neither review caught. Both had accepted that scanning rendered text for structural markers was sufficient to prove the segment path and the string path agree. Fuzzing showed 89 of 600 MARKER-FREE renders diverge: encoding_k3.py:93-99 emits an attribute as four adjacent text segments, so the flat string lets a BPE merge cross a boundary the reference never crossed. An argument key of ' spaced ' is enough, and no marker scan can see it. The scan was deleted and replaced by re-encoding the rendered string with the same encode() the worker calls and requiring equality with the reference ids -- exact, about ten lines, and it also catches marker-in-dict-key and split-across-parts. 400/400 realistic conversations pass at a small fraction of one encode's cost. encode() is narrowed to the four structural markers rather than allowed_special="all", licensed by a guard that probes the renderer across seven configs and confirms it emits only those four. "What does [EOS] mean in a tokenizer?" now encodes identically to HF user content instead of injecting stop token 163585 into the prompt body. Both eos ids are configured, since generation_config.json (163586) and tokenizer_config.json (163585) disagree. Registry routing is deliberately NOT included -- batchgen/config/ is outside MODEL_ALLOW_RE and needs its own core PR.
kimi_k3/config.py: strict parser (text_config nesting, 1-BASED layer lists,
key allowlists at all three levels, MTP -> NotImplementedError). Raises on
unknown keys and violated invariants rather than defaulting.
kimi_k3/model.py: the M2 decoder -- KDA with the full-rank gate and the
A_log F32[128] pad the checkpoint actually ships (the reference's [96]
allocation cannot strict-load its own weights), q-LoRA NoPE-MLA with the
bf16 sigmoid output gate, LatentMoE + SiTU, and the memory-lean Block
Attention Residual mixer. Hard-fail perimeter throughout; no fallbacks.
kimi_k3/kda_reference.py: fla's own torch KDA core, vendored, because the
HF oracle has NO pure-torch KDA -- it delegates the whole recurrence to
triton. Parity oracle only, never a serving path.
THE LOAD-BEARING FIX (found by adversarial review, not by the build): the
fla version story was wrong in a way that would have shipped silently wrong
numerics. The vendored torch core was labelled fla 0.4.2 while being 0.5.2
byte-for-byte, and the GPU test then PINNED the env to 0.4.2 -- whose
chunk_kda has no `use_beta_sigmoid_in_kernel` parameter at all. It swallows
the kwarg via **kwargs and sigmoids nothing, so the oracle's call form feeds
unbounded beta logits into the delta rule: finite, plausible, wrong.
The guard is a SIGNATURE probe, not a version string -- a version string is
precisely what misled us here. `_import_chunk_kda` refuses any fla that does
not NAME the parameter, so a wrong fla fails loudly at import instead of
quietly computing something else. Provenance headers corrected to 0.5.2.
Two detectors were decorative and are now proven to bite:
* T10 (the lean mixer must not materialize the full tensor) counted fp32
ELEMENTS, so a variant materializing (T,nb+1,H) in BF16 passed it. Now
bounded in BYTES across every dtype: verified to catch that evader at
151.0 MB against a 39.6 MB bound while the honest form uses 37.7 MB.
* The GPU beta test asserted only that raw and pre-sigmoided beta differ
-- true whether or not the kernel sigmoids. Replaced with the
discriminating identity kernel(raw, flag=True) == kernel(sigmoid(raw),
flag=False), which holds iff sigmoid is applied exactly once.
* The mutation runner treated ANY nonzero pytest exit as a catch, so a
renamed test would silently retire its detector (this project has lost
two that way). It now requires collected>=1 and failed>=1; verified that
a vacated detector reports SURVIVOR.
Also: media_placeholder_token_id no longer defaults (guarding a guessed id
is a silent fallback inside the hard-fail perimeter); the layer-partition
check no longer skips when one list is empty; the harness docstring no
longer claims conv/gated-norm are cross-validated on CPU (both transcribe
the same reading of fla, so a shared misread cancels -- only GPU Part B
discriminates); dead if/else branch in _run_attn collapsed.
Verified: 51/51 CPU tests, 40/40 mutations red, clean run green.
GPU KDA parity staged and NOT yet run -- the KDA recurrence
interior, conv and gated-norm remain CPU-cancelling until it does.
GPU Part A failed at fail_frac 5.6e-2 against the project BF16 gate. It is not a formula error, and the evidence is that switching the inputs bf16 -> fp32 drops fail_frac 240x (5.6e-2 -> 2.3e-4) with the formula untouched -- the signature of rounding. In fp32 every remaining failure sits at |ref| < 0.1*RMS, i.e. positions whose true value is a cancellation residue, where a relative gate measures nothing. So the instrument was wrong, not the kernel. fla's own test suite compares this exact pair -- its triton chunk_kda against its naive torch recurrence -- with a SCALE-RELATIVE RMS ratio, RMS(err)/RMS(ref) < 0.005 (fla/utils/_testing.py::get_err_ratio, fla/tests/ops/test_kda.py). Adopting the library authors' bar rather than inventing one, we measure: syn bf16 0.004497 | fp32 0.001827 real bf16 0.004480 | fp32 0.001829 (the true 96 x 128 geometry) oddT bf16 0.004477 | fp32 0.001827 all under the 0.005 bar. Note bf16 sits at ~90% of it -- that is thin margin and worth watching if the chunk size or geometry changes. This is a whole-tensor budget, so it is not a weaker test of what actually matters: a wrong gate branch, a missing sigmoid or a dropped l2norm blows past 0.005 immediately. The per-element BF16 gate stays in force everywhere magnitudes are O(1) and it discriminates -- module outputs, logits. Also drops the FLA_PIN constant and corrects the file header and run script, which still described the retired 0.4.2 pin. NOT addressed here and still failing: Part D (lean mixer transient just past its memory bound) and Part E (full model, err_ratio far past the bar in both dtypes). E does NOT improve in fp32, so it is structural rather than rounding, and the harness's claim that it forces the oracle's MLA to eager is false -- transformers logs "Ignoring the provided attention implementation eager / Using flash_attention_2 backend instead". Under investigation; M2 is not signed off until both close.
The old limit was a round number, and the true peak sat slightly over it -- a failure that said nothing about correctness. Peak = output + N live fp32 chunk working sets, with N measured just under 4 (and the peak scaling linearly with chunk_size, which confirms the model). The bound is now derived from that with N = 4. More importantly, a magnitude bound cannot tell "chunked" from "materialized at a small T". Added the discriminating check: the NON-OUTPUT transient must not grow when T doubles. That is the actual claim -- the (T, nb+1, H) tensor is never built -- and reverting to the reference form trips it at once.
Part E compared our 25-layer stack against the oracle across TWO independent
kernel implementations of the same two ops, and asserted argmax equality.
That assertion is unsatisfiable for any correct port, and the bisect shows
why (all measured in fp32, so it is not a bf16 story):
* Our CausalConv1dSilu and KimiGatedRMSNormSigmoid differ from fla's triton
ShortConvolution / FusedRMSNormGated by ~1e-7 -- and both are EQUIDISTANT
from an fp64 ground truth (6.5e-8 vs 7.2e-8). Neither is wrong.
* chunk_kda has a perturbation-response FLOOR: output difference saturates
near 1e-5 however small the input difference (gain 113x at 1e-7, 31x at
1e-6, 11x at 1e-5 -- shrinking gain = noise floor, not ill-conditioning;
an fp64 reference shows gain ~1.0).
* The top-16-of-64 sigmoid router is discontinuous: 11-20 tokens per 1024
sit within 1e-4 of the rank-16/17 boundary. One flipped token accounted
for 7.466e-3 of layer 1's 7.472e-3 MoE divergence; positive feedback ran
it to 955/1024 tokens by layer 24.
Intrinsically, given identical input, EVERY non-KDA module is bit-exact:
MLA 0.000e+00 at all 7 layers, MoE 0.000e+00 at all 24, dense MLP 0.000e+00,
router index sets 0/1024 mismatched. Only KDA differs (1.8e-05..8.2e-05), and
only because of the conv/gated-norm seam feeding it.
So test_E now SHARES fla's kernels into our stack (weights shared, not
copied) and the assertion is TIGHTENED from a bf16 gate to bit-equality --
measured err_ratio 0.000000, top-1 100.00%, max_abs 0.0, in both fp32 and
bf16. The pure-torch modules keep their coverage in test_B_*_micro_parity,
which is where a kernel-vs-kernel comparison belongs.
test_E_kernel_seam_amplification pins both amplifiers so a future failure
cannot be explained away by loosening a tolerance, and _share_fla_kernels
hard-fails if it matches nothing (a rename must not silently revert the test
to comparing across kernel families).
RETRACTED: the flash_attention_2 warning I flagged earlier is stale output
from KimiLinearModel.__init__ (modeling_kimi_linear.py:1110-1117). The
harness reset DOES take effect -- config, model.config and every
self_attn.config are the same object -- and at forward time the eager path
provably runs. FA2 explains none of the divergence.
Part B compares our KDA module against the oracle's and hit the identical seam as Part E (fail_frac 3.3e-3 / 1.1e-3): both stacks run the same chunk_kda, but ours feeds it our pure-torch conv/gated-norm output and the oracle feeds it fla's, and the kernel's ~1e-5 perturbation floor magnifies that 1e-7 difference. Sharing the kernels leaves ONLY our wiring under test -- projection order, the A_log slice, dt_bias layout, beta at the call site, the o_norm gate source -- and the assertion tightens from a bf16 gate to bit-equality. The kernels keep their own coverage in test_B_conv_micro_parity / test_B_gated_norm_micro_parity. Also fixes my own broken test from the previous commit: it read cfg.n_routed_experts / num_experts_per_tok, which do not exist on KimiK3Config (they are num_experts / num_experts_per_token) -- an AttributeError, so the router half of the amplification claim was never actually being measured. Now reads the real fields, and asserts the minimum boundary gap rather than a token count, so the check is robust to the seed while still pinning the claim that the top-k boundary is approached far below the perturbation floor.
Production BatchGen prefill is PACKED — (1, total_tokens, H) + cu_seqlens —
not the padded (B, T) batch this decoder was bit-exact on. Running the padded
code on a packed batch does not merely drift, it produces unrelated numbers.
Measured with K3-SYN-25 (seqlens [37, 53, 128]):
* eager MLA with the plain triangular mask lets sequence 2 attend to all of
sequence 1: err_ratio 7.74e-1 at the MLA module. With a block-diagonal
causal mask the packed MLA is BIT-EXACT to the per-sequence oracle.
* CausalConv1dSilu run densely over the packed axis replaces the left
zero-pad with the previous sequence's tail, corrupting exactly the first
W-1 = 3 tokens of every non-first sequence (max abs 1.598e-1).
* chunk_kda needs cu_seqlens AND a boundary-correct conv; with both, the
packed KDA module is bit-exact to the per-sequence oracle.
All three channels are wired: the mask (MLA), per-segment convs, cu_seqlens
into chunk_kda. _apply_attn_res_lean is token-parallel and is deliberately
untouched.
In _build_block_causal_mask, torch.bucketize(..., right=True) is load-bearing:
right=False assigns each sequence's FIRST token to the previous segment and
left the fp32 whole-model error at 7.4e-2 — small enough to be mistaken for
drift. Commented at the call site.
cu_seqlens=None keeps the dense path bit-for-bit: verified by running the
pre-change and post-change modules side by side in one process on identical
seeded weights — torch.equal on the logits across syn25/skew10 x bf16/fp32 x
(4,128)/(1,257)/(2,33), max_abs 0.0 in all 12 cases. On GPU the unchanged
test_E_full_model_gpu still asserts bitwise identity to the oracle. The
oracle-parity claim rests on that.
Backend selector fla_chunk -> fla_triton: "chunk" in fla's chunk_kda names the
chunkwise-PARALLEL algorithm (64-token internal chunks), NOT chunked prefill.
In a codebase whose headline feature is unchunked prefill that name invites
precisely the wrong conclusion.
New GPU Part F, each with a MANDATORY control that must fail:
* KDA and MLA sub-modules at UNEQUAL lengths vs the oracle per sequence:
torch.equal. Controls: the plain triangular mask (7.74e-1).
* Whole model at EQUAL lengths (3x37, 2x53) vs the oracle's dense (N, T)
batch: torch.equal in bf16 AND fp32. Control cu_seqlens=None: ~1.0.
Equal lengths are load-bearing, not convenient — they make the packed
total equal the dense batch's flattened M, so every GEMM shape matches and
the comparison isolates packing.
* Unequal lengths at FULL depth cannot be gated to parity, and the test says
so with the measurement instead of a loosened bar: the packed run differs
from a per-sequence reference by err_ratio 1.62e-1 bf16 / 4.29e-1 fp32,
because packing changes every GEMM's M (a bare Linear already moves 4.2e-7
in fp32), chunk_kda's documented perturbation floor lifts that ~170x, and
the top-16-of-64 router turns it into a flipped expert at layer 2 (layers 0
and 1 are bit-identical; layer-1 router flips 0 of 218 tokens). Removing
the router discontinuity entirely still leaves 5.0e-3, so the ceiling is
structural. That test asserts what is attainable: decoder layer 0 (KDA +
dense MLP, no router) bit-identical, and a >4x separation from the control.
_share_fla_kernels no longer hardcodes cu_seqlens=None when rebinding the conv,
which would have silently reintroduced the leak Part F hunts.
added 19 commits
August 13, 2026 05:54
M-PR-6. The 48B keeps attention, KDA and shared experts resident and streams
only the routed experts. K3 cannot: 69 KDA x 846.67 MiB + 24 MLA x 442.88 MiB
+ 92 shared x 252.00 MiB = 90.07 GiB — nearly the GPU's whole HBM, before
the skeleton, the KDA state pools or a single activation.
Three things were missing for the other three rings, all of them wiring:
* _build_weight_copy_task declared attn/kda_attn/shared_expert and left them
empty, so the producer had nothing to drain. Now populated layer-major
ascending -- the order the consumer requests them in, which is what the
ring requires (an out-of-order request finds no slot and dies on
get_weights' 2 s throw rather than hanging).
* kda_attn has no key in base_planner's default num_prefill_module_buffer,
and GPU_Weight_Buffer::Init() iterates that map -- an absent key is zero
slots and an unbreakable producer stall. Declared model-side, in the Kimi
planner, so no core change is needed.
* the three _load_* methods permanently .to(device)'d the weights. Under
streaming they empty the params instead (clearing meta so model.to() is
safe) and the wrappers are built persistent=False. The shared expert gets
an ExpertWrapperBase with expert_idx=-1, whose module key is already the
"shared_expert_{L}" the parameter server serves.
Gated behind stream_all_modules, default OFF, so the validated 48B path is
untouched; the planner turns it on for K3 only. configure_decoding raises
while it is set -- the worker only starts a decode H2D streamer for routed
experts, so decode would stall the other three rings.
Separately: num_experts. K3's config.json says num_experts: 896 and ships no
n_routed_experts. Both `getattr(cfg, "n_routed_experts", 256) or 256` sites
silently returned 256, and the PSM's copy was worse -- it read a ModelConfig,
which has no such field at all, so it returned 256 for every model on every
path. Replaced with require_num_routed_experts(), which reads either spelling
and raises rather than defaulting.
KDA state pools were sized 256 slots at K3 dims, larger than the whole
GPU. K3 gets 4 slots (428.6 MiB each), a deliberate cap on concurrent
sequences rather than a silent 256.
The M3 rehearsal in PREFILL_PLAN is 'force the 48B to stream all four rings and check its logits do not move' — on a model whose right answer is already known. That was unreachable: the ring depths were nested inside the K3-only branch, so there was no way to turn streaming on for the 48B. Split them: _configure_streamed_rings() runs whenever stream_all_modules is set, _adjust_config_for_k3() keeps only what is genuinely K3-shaped (the KDA state-pool slot count). stream_all_modules is now a constructor argument defaulting to is_k3, so the rehearsal is KimiLinearPlanner(is_k3=False, stream_all_modules=True) and nothing about the 48B default changes.
moe_forward_serving dispatched the 7168 hidden straight into K3's routed experts, whose w1/w3 are 3584->3072: routed_expert_down_proj, routed_expert_norm and routed_expert_up_proj were never called, so on K3 the streamed prefill MoE was not merely inaccurate, it could not run. Op order now follows the eager reference (kimi_k3/model.py, bit-exact to the HF oracle): router on the PRE-down hidden; down_proj ONCE per token before dispatch; experts in the latent; FP32 combine; norm once post-combine; up_proj; shared expert on the identity path in hidden space. Gated on config.routed_expert_hidden_size, so the 48B keeps the hidden-space path unchanged. No silent fallbacks: a kimi_k3 config that reaches the non-latent branch (or has latent_moe_use_norm off) raises, and resident-EP decode — which routes in the hidden space and has no latent seam — refuses a LatentMoE config instead of computing the wrong thing. tests/gpu/test_kimi_linear_latent_moe_serving.py pins the serving forward against the M2 eager block on identical weights (2 synthetic configs, bf16 and fp32), the router index sets, the once-per-token projections, the 0-token drive-every-expert invariant, and both hard-fails.
k3_skeleton_declaration emits 'language_model.'-prefixed names (that is how the C++ parameter server keys skeleton_state_dict_), while _load_model_skeleton looked up the bare model.named_parameters() name. On the real 93-layer K3 config that is 1026 declared skeleton tensors and 0 hits -> RuntimeError before a single weight is copied. k3_skeleton_key() already existed as the intended bridge and had no caller outside a test; give it its one caller. The translation happens in exactly one place, _skeleton_ckpt_key, and in the model-name -> ckpt-name direction: that direction is a total function, whereas stripping on ingest is not (K3's skeleton_state_dict_ also holds the 168 unprefixed vision_tower./mm_projector. tensors). Identity for the 48B, whose checkpoint has no prefix. Still exactly one lookup per parameter: a missing skeleton tensor hard-fails with both names in the message, and there is no bare-name second probe. Measured, meta build, no weights: K3 (93L/896E) found=1026 missing=0 (was found=0 missing=1026); Kimi-Linear-48B found=112 missing=0, unchanged, and every 48B lookup key is byte-identical to the parameter name.
The engine serves w{1,2,3}.weight_packed + .weight_scale; KimiBlockSparseMLP
declares w{1,2,3}.weight. apply_weights (models/wrappers/base.py:167-170)
skips a served name that is not a module parameter and has no else, so every
K3 routed expert was computing on torch.empty(0) — 247,296 unserved expert
params at 896 experts/layer, silently.
K3MXFP4Expert declares the six served names verbatim (name sets equal, nothing
skippable); KimiK3MXFP4ExpertWrapper validates the ring slot before every use
and asserts apply_weights matched everything. Compute goes through the marlin
MXFP4 path (SiTU epilogue compiled in, 9/9 GPU parity ladder), not the WGMMA
one, whose epilogue is gpt-oss SwiGLU and cannot express K3's activation.
attn_res_block_size=12 is in the real config.json, but nothing in the serving modules or the PSM knew about block_residual. K3 REPLACES the classic residual body with depth-attention over block boundaries, so the gap is not a missing feature — it silently changes every layer's input, and on the prefill path it does not even get that far: the worker passes no block_residual, so layer 0's boundary append hits torch.cat([None, ...]). Ground truth is the M2 eager model (kimi_k3/model.py, bit-exact to the HF oracle). The layer body already matched it; what was missing was the two seams the SERVING caller does not provide. - block_residual.py (new, torch-only so a test can load it without batchgen/fla): the lean depth mixer, the between-layer carrier, the serving layer forward, and the output-stage pre-hook. - The mixer is the M2 memory-lean form, verified bitwise identical to the unchunked one it replaces (fp32 and bf16, nb in 0/1/3/9). It matters here: its fp32 transient is O(chunk*(nb+1)*H) instead of O(T*(nb+1)*H), and in serving T is a whole packed micro-batch. - The PSM installs both seams for a config with attn_res_block_size and nothing at all otherwise, so the 48B path is byte-for-byte unchanged. Decode already worked: it enters through KimiLinearModel.forward, which threads block_residual explicitly and applies the output mix itself. The injected forward keeps that convention working untouched and the hook no-ops for it, so the mix is applied exactly once on either path — which the test pins by comparing the two paths for bitwise equality. The carrier exists only because the worker's prepack-prefill loop is core-scope. Every hand-off is checked: layer 0 re-seeds (block_residual is intra-forward scratch), any other layer must find its predecessor's tensor, and the output hook refuses a stack that did not run to the end. Wrong ordering raises instead of quietly producing a different model.
--meta proves the name-set arithmetic in BOTH directions on a real-config meta build; --gpu proves the device repack is bit-identical to the oracle-gated CPU one (so it inherits that validation instead of asserting its own), that the wrapped expert matches an oracle-dequant reference under the project gate, and that every slot mutation hard-fails.
The Phase-A adapter's patched layer forward returns a 1-tuple and its captured span runs the classic residual body, so a replayed K3 layer neither produces nor consumes block_residual. The adapter is installed even in "eager" mode precisely so batchgen_debug can flip modes on a live server — which left K3 one flag away from decoding without its depth residuals, with no error and no log line. So do not install it at all for a config with attn_res_block_size, and raise if a graph mode was actually requested. Nothing changes for the 48B, which has no attention residuals and keeps the live-switch path.
The first run failed at fail_frac 1.15e-01. Measured cause: the all-fp32 reference disagrees with the TRUE bf16 answer by err_ratio 3.40e-3 / fail_frac 1.15e-01 — IDENTICAL to the kernel's apparent error, i.e. the instrument was being measured, not the kernel. kernel-validation.md already says a BF16 GEMM's ground truth must be a BF16 tensor-core matmul. Against a reference at the kernel's own precision the chain lands at err_ratio ~4e-4. Gate is the RMS-relative one (assert_kernel_err_ratio, bar 5e-3): a K=3072 reduction output concentrates its per-element failures at |ref| < 0.1*RMS, which the harness already documents as the rounding signature that makes the per-element gate the wrong instrument. Two mutation arms (gate/up swap, SiLU-for-SiTU) keep that bar honest.
Four tracks landed the missing K3 serving pieces concurrently (2c2d6f0, bdb32a3, d61a6b1/2b70478d, d80f9c0/14dc674c). This pass reconciles them and applies the adversarial-review findings that survived re-measurement. WHY EACH PIECE WAS NEEDED (verified on the real 93-layer config) * LatentMoE in the serving MoE (2c2d6f0). K3's routed experts live in a 3584 latent, not the 7168 hidden. Without the down/up seam the serving forward could not run K3 at all: `RuntimeError: mat1 and mat2 shapes cannot be multiplied (23x448 and 224x192)`. * Skeleton lookup by CHECKPOINT name (bdb32a3). The C++ parameter server keys skeleton_state_dict_ by the checkpoint name and every K3 text tensor carries `language_model.`. Before: found=0 missing=1026. After: found=1026 missing=0 — verified here against the REAL converted checkpoint (96 shards, 497,220 tensors), not only against the declaration. * MXFP4 expert modules (d61a6b1). The checkpoint ships no `.weight` for any routed expert. apply_weights silently skips a served name that is not a module parameter, so BF16 experts were not "unquantized", they were empty. Before: 247,296 unserved expert params (82,432 experts x 3). After: 0. * Block Attention Residuals wiring (d80f9c0). The worker's prepack-prefill loop passes no block_residual and calls model.model.norm directly, so layer 0 died on torch.cat([None, ...]) and the output depth mix never ran. WHAT THIS COMMIT CHANGES 1. Parallel_Strategy_Manager: `k3.tensor_map` is imported inside the K3 branch, not at module level — `k3/__init__.py` promises nothing is exported eagerly, and the module-level form dragged k3.tensor_map -> k3.mxfp4_layout -> weight_reconciler into every 48B run that imports the PSM. The skeleton error message now reads K3_CKPT_PREFIX instead of a duplicated literal; with a mutated prefix the old message still printed "language_model." while the lookup used something else. 2. k3/mxfp4_expert.is_mxfp4_quantized: a `kimi_k3` config with no quantization_config now WARNS instead of selecting BF16 experts in silence. It does not raise: validate_k3_config already refuses that config on all three real entry points (kimi_initializer:205, build_k3_state_dict_name_map :392, load_k3_config:685), and an unquantized K3 is a supported synthetic shape — the M2 eager ground truth builds BF16 experts for exactly the K3-SYN-25 / K3-SKEW-10 harness configs. 3. k3/mxfp4_expert banner: the per-forward marlin repack now carries its measured cost qualitatively — flat in token count, a several-fold single-expert slowdown, and a serious per-layer cost at prefill occupancy. The old "collapses to two .view()s" end state was self-contradictory: marlin_grouped_moe.py:358 hard-fails a non-bf16 scale, so uint8 marlin-order scales keep a per-expert expansion every forward, and bf16 scales cost +5.88% bytes per expert across 82,432 experts. 4. serving_modules: the two LatentMoE seam errors are no longer conflated. Norm per expert / after up_proj IS a different function (err_ratio 1.486e-01 / 8.180e-01). Down-proj per (token, expert) is NOT — the projection is a bias-free nn.Linear, so it is idempotent on a duplicated row and measures 0.000e+00 (bf16) / 2.259e-07 (fp32). Only the call-count hook test pins it. `_require_k3_latent_moe`'s docstring no longer claims to catch a config that never went through from_hf_dict; its only trigger is model_type. 5. block_residual / model.py (both families): the chunked depth mixer is NOT bit-identical to the unchunked reference, and two source comments said it was. MEASURED H=512 fp32: torch.equal True at T in {13,1024,2048,8192}, False at T in {1025,4097} with max_abs 2.4e-7 (nb=3) / 2.6e-6 (nb=9) — the ragged final chunk picks a different ATen/cuBLAS strategy. The gate is the 1e-6 tolerance, not bit equality. The memory note also claimed the peak extra allocation is T-independent; re-measurement disproved that (several fp32 chunk tensors are live at once, plus an O(T*H) output buffer, so the peak grows with T), and torch.inference_mode is load-bearing — with grad enabled the same call allocates an order of magnitude more at large T. 6. Parallel_Strategy_Manager._init_decode_graph: the CUDA-graph refusal was justified by a silent-corruption hazard that does not exist. A replayed 1-tuple raises `ValueError: not enough values to unpack (expected 2, got 1)` at model.py:880. The guard stays — that error names neither CUDA graphs nor block residuals and lands mid-decode — but it is now sold on the real failure mode. VERIFICATION (meta build, no server, no multi-GPU) K3, real config, stream_all_modules=ON: unserved/unlanded = 0/0 for all four module types; skeleton 1026/1026 against the declaration AND against the real converted checkpoint's promoted key set; copy-task element-for-element equal to build_k3_state_dict_name_map's ordering on all four rings. 48B, real config, stream_all_modules=OFF: unchanged, digests identical to ea135c8. CPU: tests/test_kimi_k3_model.py 51 passed; mutation runner 40/40 red, clean run green.
model.py gains `_apply_output_attn_res`, a method wrapping the output depth-mix the eager forward already did inline. The worker (companion core PR) calls that instead of importing a model-private function, so core stays model-agnostic and there is exactly one implementation of the mix — the eager and serving paths cannot drift.
Kimi-K3 ships q/k/v_conv1d.weight as fp32 while the KDA projections are bf16. The causal_conv1d CUDA kernel enforces weight dtype == input dtype (causal_conv1d.cu:239/337), so K3 prefill died with 'RuntimeError: weight type must equal input type' on the first KDA layer. Cast weight (and bias, if present) to the activation dtype at the call sites. No-op when they already match, e.g. Kimi-Linear 48B.
PREFILL_MEMORY_AUDIT fix #2. Model-scope only (models/** + batchgen_kernels/** + tests/**); no core files touched. causal_conv1d_fwd's varlen path stages a channel-major (dim, total) copy, runs the kernel in place on it, and returns x_cm.t() — a NON-contiguous view. fla's @input_guard on ChunkKDAFunction.forward then does arg.contiguous() on q, k and v, allocating a third full copy of each while the staging bases stay pinned by the caller's locals. At K3 prefill scale that is three full (total, dim) tensors of pure layout duplication live inside chunk_kda, on every one of the 69 KDA layers. New overwrite_x=True (varlen only) transposes the result back into x's own storage and returns x, so the staging buffer dies at return and exactly one (total, dim) tensor survives. kda_prefill_serving passes it for q/k/v, whose projection outputs are dead after the conv. Bit-exact. The op the flag adds, x.copy_(x_cm.t()), is the same same-dtype same-device element copy that fla's .contiguous() was performing — same source bytes, same destination values, no arithmetic and no reassociation. The kernel call itself is unchanged (identical args, identical strides), so conv outputs and pooled final states are bit-identical; the CPU check asserts torch.equal on both. Default stays False: no existing caller changes behaviour. It removes three of the sixteen (S, dim) tensors live at the layer-0 OOM point and lifts the KDA-bound S_max accordingly; peak during the three conv calls themselves is unchanged, and the copy count is unchanged (two transposing copies either way) — no extra bandwidth. Verification (no GPU used): tests/kimi_linear/test_conv1d_layout_cpu.py — CPU, stubbed extension, 11/11 pass: returns the caller's storage, token-major contiguous, values and pool writes bit-identical to the default path, and .contiguous() on the rearranged (1, L, H, D) tensor now returns self instead of allocating. test_conv1d_std.py gains case 4 for the same properties on the real kernel.
Prefill fix 1 of PREFILL_MEMORY_AUDIT.md. KimiMLP.forward now runs the FFN body in token tiles of 8192 rows and writes into a preallocated output, so SituAndMul's five concurrent (tokens, intermediate) fp32 tensors plus the bf16 [gate, up] cat — 24 bytes per (token, intermediate) element — stop scaling with S. The layer-0 dense MLP and MoE shared-expert FFN transients drop by more than an order of magnitude at long-context S and are now flat in S: 24*8192*I of tile working set plus the (S,H) output. This removes the FFN-bound S_max wall; the next binding term is layer-0 KDA. BIT-EXACT, not approximately: `_ffn` is the pre-change body verbatim, applied per tile, so no sum is reassociated and no op order changes; every output row depends only on its own input row. Tiles are EVEN rather than fixed-width + ragged remainder, because a few-row remainder is a degenerate GEMM shape that can move the BLAS onto a different reduction order (MEASURED: a 1-row fp32 F.linear does not reproduce the corresponding row of the full GEMM on CPU). <= one tile keeps the original single call, so decode and CUDA-graph capture are untouched.
K3's depth-residual accumulator grew by `torch.cat` at each of the 8 block boundaries (layers 0,12,...,84). A cat allocates the (S,nb+1,H) result while the (S,nb,H) input is still live, so at long-context S the last boundary held the old and new accumulators simultaneously and the eight appends churned several times the buffer's size in allocation. BlockResidualBuffer now allocates one (S, num_boundaries, H) tensor per pass and each boundary writes column nb in place: peak flat at one buffer, one allocation, no old+new co-live moment. `PREFILL_MEMORY_AUDIT.md` §4/§7 fix 3. What is threaded is still a plain tensor, and still the NARROWED view buf[:, :nb+1] — never the whole buffer. That is the trap in this change: every consumer reads block_residual.shape[1] as "boundaries so far" (the shape[1] > 0 gate in _forward_attn_residual, the block_residual[start:end] read in apply_attn_res, the worker's bres= log), and a caller handed the full (S,8,H) buffer would mix 8 all-zero keys into layer 0's depth-attention and silently compute a different model. NUMERICS: bit-exact, not "within tolerance". The append is a same-dtype copy_ of exactly the bytes cat copied, into exactly the same logical column. The only consumer of the view immediately re-cats a token slice of it into a fresh contiguous fp32 tensor, so the wider row stride never reaches a reduction: apply_attn_res's `v` comes out byte-identical AND identically laid out either way. Nothing is reassociated, no op order changes, no reduction is chunked. Semantics preserved verbatim: prefix_sum still RESETS by assignment at a boundary, block_residual is still intra-forward scratch re-zeroed every forward (seed() allocates a fresh zeroed buffer — a recycled one's untouched columns would hold another request's activations), and the depth-mix output still feeds only the norms, never the accumulator. Gated by tests/test_kimi_k3_block_residual_prealloc.py, CPU-only (torch alone; fla and the batchgen package are stubbed when absent). Three drives of the same 93-layer / 8-boundary stack must be torch.equal per layer and at the output stage: the untouched M2 eager ground truth (kimi_k3/model.py, still cat-based), the patched kimi_linear body with the buffer unseeded (cat fallback), and the patched body with it seeded. 12 passed. Mutating the append to return the whole buffer, or to skip a column, is caught.
PREFILL_MEMORY_AUDIT fix #5, on top of fixes #1-#4. Model-scope only (models/** + tests/**); no core file touched. chunk_kda holds ~17 tensors of (S, 96*128) at once, so one KDA layer's working set is 417,792 B/token and scales with the whole packed batch, on all 69 KDA layers, on every rank (attention is replicated under EP-8). kda_prefill_serving now drives it in chunk-aligned token segments (KDA_PREFILL_SEGMENT_TOKENS = 16384), reading each sequence's fp32 recurrent state out of the pool at the top of a segment and writing the segment's final state back at the bottom. 12 of the 17 tensors then scale with T instead of S; 6 (q/k/v after the conv, f, z, the stitched output) stay resident. WHY IT IS BIT-EXACT, and what is deliberately NOT segmented. Only chunk_kda is segmented. The projections, the three convs and o_norm/o_proj still see the whole packed range. - No GEMM shape may change. test_F_full_model_packed_equal_lengths measured that re-running an nn.Linear at a different M moves it by ~3e-7 because cuBLAS picks its kernel from the problem shape, and K3's top-16 router amplifies that without bound. So the projections stay at M = total. - Segmenting the conv would therefore save nothing: with overwrite_x its output IS the projection buffer, both full-size either way. It would only add the width-1 = 3 token carry across every cut — a known trap for zero bytes. Not done. - Every segment cut is a sequence boundary or k*64 tokens from the start of the sequence it falls in, so each segment's per-sequence chunk grid is the restriction of the unsegmented grid. Chunk-local work (the gate cumsum, the WY transform, intra-chunk Aqk/Akk, the ragged tail) is untouched, and the only thing crossing a cut is the inter-chunk state — fp32 out of chunk_delta_h.py L314, fp32 into L137 (`b_h1 += load(h0)`), fp32 in the pool. No reduction is re-cut and no sum is re-associated. - Segment size cannot select a different kernel either: every autotune key on the KDA forward path (chunk_delta_h, kda/gate, kda/chunk_intra, kda/wy_fast, gla/chunk) lists head/dim/BT constants and never T, and the FlashKDA backend verifier keys off flags, not shapes. Recomputed at S=131,072: 12 of the 17 working-set tensors now scale with the segment size T instead of S, the per-layer peak drops by more than half at the default T=16,384, and the KDA-bound S_max rises well past the target context; the binding term becomes block_residual at nb=8. Verification (no GPU used, no server launched): tests/test_kimi_k3_kda_segmented.py — CPU, 46 pass. torch.equal on both the output and the pool's final states, over 5 packed layouts x {fp32, bf16}, driving the production _kda_segment_plan/_kda_chunk_segments against a transcription of fla's own naive_chunk_kda unrolled to one chunk at a time. The answer is invariant to the segment size for every aligned size. Controls: cutting mid-chunk (_KDA_CHUNK_SIZE forced to 1) is NOT bit-exact; a shared pool slot is NOT bit-exact; the plan is asserted to actually cut inside a sequence and to actually span a sequence boundary. Mutation-checked offline: dropping the cu_seqlens rebasing (1.0e-1), the slot slice (2.3e-1) or the state hand-off (3.7e-1) all fail the gate. NOT verified here: that fla's Triton kernels are token-count invariant on real hardware. That is argued from the 0.5.2 source above and needs the integration run to confirm.
Lands fixes 1-5 of batchgen_design/model_support/kimi_k3/PREFILL_MEMORY_AUDIT.md coherently on top of 482374d / 666af61 / b507844 / bcbf409 / fe2c180, and applies the adversarial-review findings that survived re-derivation. NO CORE FILE IS TOUCHED BY THIS COMMIT. bcbf409 (batchgen_worker.py: seed the preallocated block_residual, del inputs_embeds) is the one core-scope change in this stack and must be split into its own core PR. WHAT THE FIVE FIXES DO, AND WHY EACH IS BIT-EXACT 1 KimiMLP token tiling (T=8192). The FFN is elementwise in the token axis: every output row depends only on its own input row, and `_ffn` is the pre-change body verbatim. No sum is reassociated, no fp32 island is demoted. Tiles are EVEN, never fixed-width + ragged remainder, because a few-row tail is a degenerate GEMM M (MEASURED on CPU fp32: a 1-row F.linear does not reproduce the corresponding row of the full GEMM). The layer-0 dense-MLP and shared-expert transients drop by an order of magnitude and are now flat in S. 2 conv1d overwrite_x. Adds one same-dtype, same-device, non-overlapping Tensor.copy_ - a pure permutation of memory locations, and the same copy fla's @input_guard was already making. The kernel call is byte-identical. 3 Preallocated block_residual. torch.cat -> copy_ into a column of a buffer that already exists; append returns the NARROWED buf[:, :nb+1] view so shape[1] still counts boundaries. The only data consumer re-cats a token slice into a fresh contiguous fp32 tensor, so the stride difference dies before any reduction. Nearly halves the L84-boundary transient. 4 del inputs_embeds (in bcbf409). Frees nothing that is still read. 5 KDA segmentation (T=16,384). Every cut is a sequence boundary or k*64 tokens from its sequence start, so each segment's chunk grid is the RESTRICTION of the unsegmented grid; only the inter-chunk recurrent state crosses, in fp32 both ways, through the fp32 pool. No autotune key and no Triton specialization on the KDA path contains T. No GEMM shape changes: the projections, the convs and o_norm/o_proj still see the whole range. Cuts the per-KDA-layer working set by more than half at the default T. REVIEW FINDINGS APPLIED R1 (ffn) the tiled peak was understated by one output tile: Python rebinds `y` only on the next iteration's assignment, so the previous tile is co-live with the peak. Added `del y` (one output tile back) and test_previous_tile_is_freed_before_the_next, which is weakref-based and kills the mutation directly. R1 (ffn) the CUDA parity case only covered the shared expert's 6144 intermediate. Widened to layer 0's 33792 - a different cuBLAS regime (N=67,584), and the width fix 1 exists for. Still skips off-GPU; it is the last open assumption in fix 1 and must be run on the node. R1 (ffn) documented that the fp32 arm of the sweep is load-bearing: it is the only arm with the resolution to catch a degenerate tile (MEASURED: the ragged-tail mutant reddens 4 fp32 cases and 0 bf16 cases). R2 (conv) the whole-sweep saving does NOT survive fix 5, which the same integration requires. fla copies whatever slice it is handed, and it is handed one segment: the real saving is 3 x (segment, dim) bf16 per layer, an 8x overstatement corrected in the code comment and in the audit. The fix is still correct and still free - nothing else makes the segment slices contiguous - but it is much smaller than first claimed. R2 (conv) deleting the `overwrite_x requires contiguous x` assert was not caught. Added a check that it fires AND that it fires before the kernel mutates conv_states, plus the segment-slice contiguity property above. The file is now also a pytest entry point instead of `python -m` only. R3 (block_residual) BLOCKING, and fixed: BlockResidualBuffer._buf is process-wide class state and nothing in production ever released it. The tensor it replaced was a plain local that died with the prefill frame, so this was a regression: the whole block-residual buffer stayed pinned across configure_decoding() and the resident-EP build. BlockResidualCarrier.reset() now drops the buffer too (one switch for "the pass is over"), which covers take() at the end of a carried pass and both PSM phase switches; KimiLinearModel.forward drops it after the output mix on the explicit path. Three weakref tests. R4 (kda) the zero-length-sequence guard only inspected the first and last segment, so an INTERIOR zero-length sequence was silently dropped while the guard claimed otherwise. Replaced with a strictly-increasing precondition on cu_seqlens plus a real coverage check on the plan. R4 (kda) _KDA_CHUNK_SIZE was unfalsifiable: every test derived its expectation from it, so setting it to 32 left all 46 green while producing mid-chunk cuts. The cut-alignment assertion now uses a literal 64, and test_chunk_size_constant_matches_flas_own pins the constant against fla 0.5.2's own default. REVIEW FINDINGS REJECTED - R1's claim that test_kimi_linear_ffn_chunk is weak because the 4x-too-wide tile mutant survives. A wider tile is bit-exact; it costs memory, not correctness, and the suite's job here is exactness. Not worth a test. - R2's implication that the conv fix should be dropped now that it saves a per-segment saving rather than a whole-sweep one. It is a precondition for fix 5 collecting anything: without it every segment slice is non-contiguous and fla copies it. Kept, with the arithmetic corrected. - The suggestion to also chunk apply_attn_res or the MLA path. Out of scope, and apply_attn_res is already chunked at 1024 and is gated at 1e-6, not bit-exactness; touching it would need its own exactness argument. TESTS (all real counts, CPU, py3.11 / torch 2.8.0 / transformers 4.57.6) tests/test_kimi_k3_model.py 51 passed tests/mutation_check_kimi_k3.py 40/40 RED, clean GREEN tests/test_kimi_linear_ffn_chunk.py 56 passed, 5 skipped (CUDA) tests/test_kimi_k3_block_residual_prealloc.py 15 passed tests/test_kimi_k3_kda_segmented.py 50 passed tests/kimi_linear/test_conv1d_layout_cpu.py 1 passed (16 checks) Every change above was mutation-checked: each new assertion was shown to go red when the thing it guards is removed. PREDICTED PEAK AT S=131,072, ALL FIVE FIXES The binding layer is now MLA at nb=8, not KDA and not the layer-0 MLP; the predicted 128K peak fits the per-rank PyTorch budget with margin. Note the audit's post-fix-5 S_max claim is wrong: it prices only block_residual and omits MLA's 275,749 B/token, which none of these five fixes touches; the honest ceiling is materially lower.
tests/test_kimi_linear_block_residual_serving.py covers this wiring against the M2 ground truth but cannot run without a GPU — its import chain reaches batchgen/__init__, which JIT-builds a CUDA op, and it fails at collection with CUDA_VISIBLE_DEVICES="" too. So the property the review flagged as blocking is pinned on CPU instead: drive the real decoder_layer_forward_block_residual over all 93 layers exactly as the worker's prepack loop does, fire the real model.norm pre-hook, and assert by weakref that nothing still references the (num_tokens, 8, hidden) buffer afterwards — the full-context block-residual buffer. Also asserts a second pass reproduces the first, so the release is not trading a leak for stale scratch. Mutation-checked: removing the BlockResidualBuffer.reset() from BlockResidualCarrier.reset() reddens this and the two unit-level release tests. tests/test_kimi_k3_block_residual_prealloc.py 16 passed.
Andrewxu313
force-pushed
the
tairanxu/k3-prefill-model
branch
from
August 12, 2026 21:54
0dac884 to
711117f
Compare
Andrewxu313
marked this pull request as ready for review
August 12, 2026 21:54
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
2 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.
What
Adds serving support for Moonshot's Kimi-K3 and the kimi-linear model family: the
kimi_linearmodel directory (packed varlen prefill with a block-causal mask, KDAprefill segmentation, block attention residual carry, latent-space MoE with MXFP4
experts, streamed attn/kda/shared/routed module rings, planner overrides, bucketed
decode-graph adapter with a K-cache-geometry-aware capture signature), the
kimi_k3model directory (strict config, tokenizer with vendored metadata and pinned chat
rendering, weight-name reconciler and MXFP4 tensor map, eager reference decoder), the
registration seams, the packaged conv1d kernel wrapper, and the CPU parity / mutation
/ staged GPU test suites. Decode CUDA graphs are refused for block-attention-residual
models by design. Stacked on
tairanxu/k3-prefill-coreand requirestairanxu/k3-prefill-kernel; this PR merges last, after both.Why
Kimi-K3 combines linear KDA attention, interleaved MLA, block attention residuals, and
MXFP4 MoE — none of which the existing model directories cover. The implementation is
gated by a CPU parity suite against the vendored reference implementation (bit-exact
where attainable, tolerance-gated where a kernel seam makes bit-exactness unattainable,
with the reasons pinned in the tests), a mutation-detection runner over the model
modules, and staged GPU parity tests; validated for long-context prefill and
single-token completion through the serving path.
Type of Change
modelkernelcorefixinfradocsFile changes
batchgen/get_initializer.pybatchgen/get_parallel_strategy_manager.pybatchgen/models/moonshotai/kimi_k3/MODULE.mdbatchgen/models/moonshotai/kimi_k3/__init__.pybatchgen/models/moonshotai/kimi_k3/assets/__init__.pybatchgen/models/moonshotai/kimi_k3/assets/config.jsonbatchgen/models/moonshotai/kimi_k3/assets/configuration_kimi_k3.pybatchgen/models/moonshotai/kimi_k3/assets/encoding_k3.pybatchgen/models/moonshotai/kimi_k3/assets/generation_config.jsonbatchgen/models/moonshotai/kimi_k3/assets/kimi_k3_processor.pybatchgen/models/moonshotai/kimi_k3/assets/tiktoken.modelbatchgen/models/moonshotai/kimi_k3/assets/tokenization_kimi.pybatchgen/models/moonshotai/kimi_k3/assets/tokenizer_config.jsonbatchgen/models/moonshotai/kimi_k3/config.pybatchgen/models/moonshotai/kimi_k3/kda_reference.pybatchgen/models/moonshotai/kimi_k3/model.pybatchgen/models/moonshotai/kimi_k3/tokenizer.pybatchgen/models/moonshotai/kimi_linear/Parallel_Strategy_Manager.pybatchgen/models/moonshotai/kimi_linear/__init__.pybatchgen/models/moonshotai/kimi_linear/assets/chat_template.jinjabatchgen/models/moonshotai/kimi_linear/assets/tokenizer_config.jsonbatchgen/models/moonshotai/kimi_linear/block_residual.pybatchgen/models/moonshotai/kimi_linear/config.pybatchgen/models/moonshotai/kimi_linear/cuda_graph_segments.pybatchgen/models/moonshotai/kimi_linear/k3/__init__.pybatchgen/models/moonshotai/kimi_linear/k3/mxfp4_expert.pybatchgen/models/moonshotai/kimi_linear/k3/mxfp4_layout.pybatchgen/models/moonshotai/kimi_linear/k3/tensor_map.pybatchgen/models/moonshotai/kimi_linear/kimi_initializer.pybatchgen/models/moonshotai/kimi_linear/kimi_parameter_server.pybatchgen/models/moonshotai/kimi_linear/model.pybatchgen/models/moonshotai/kimi_linear/planner.pybatchgen/models/moonshotai/kimi_linear/serving_modules.pybatchgen/models/moonshotai/kimi_linear/tokenizer.pybatchgen/models/moonshotai/kimi_linear/wrappers.pybatchgen/models/weight_reconciler.pybatchgen/models/wrappers/attention.pybatchgen_kernels/_jit_registry.pybatchgen_kernels/conv1d/__init__.pybatchgen_kernels/setup.pybatchgen_kernels/src/conv1d/causal_conv1d.cubatchgen_kernels/triton/fused_moe_bf16.pydocs/INSTALL.mddocs/troubleshooting.mdtests/gpu/__init__.pytests/gpu/run_kimi_k3_kda.shtests/gpu/test_kimi_k3_kda_fla_parity.pytests/gpu/test_kimi_linear_latent_moe_serving.pytests/gpu/verify_k3_mxfp4_expert.pytests/kimi_k3_harness.pytests/kimi_k3_oracle_assets/__init__.pytests/kimi_k3_oracle_assets/configuration_kimi_k3.pytests/kimi_k3_oracle_assets/fla_cpu_shim.pytests/kimi_k3_oracle_assets/modeling_kimi_linear.pytests/kimi_linear/test_conv1d_layout_cpu.pytests/kimi_linear/test_conv1d_std.pytests/kimi_linear/test_decode_graph_adapter.pytests/kimi_linear/test_fused_moe_std.pytests/kimi_linear/test_kda_manager_graphready.pytests/kimi_linear/test_kda_segment_capture.pytests/mutation_check_kimi_k3.pytests/test_kimi_k3_block_residual_prealloc.pytests/test_kimi_k3_kda_segmented.pytests/test_kimi_k3_model.pytests/test_kimi_k3_tensor_map.pytests/test_kimi_k3_tokenizer.pytests/test_kimi_linear_block_residual_serving.pytests/test_kimi_linear_ffn_chunk.pytests/test_weight_reconciler.pyChecklist
model/kernelPR does not touch the scheduling/scaffolding layer (§2.5–§2.6).debug_*/scratch_*/tmp_*scripts,BATCHGEN_*env-guards, strayprint(), committed artifacts, orCo-Authored-By(§1/§4). (bench_*benchmarks are fine.)tests/; touched modules'MODULE.mdupdated if the public API changed (§2).bash .github/workflows/scripts/check-pr-hygiene.sh origin/mainlocally; the CI hygiene check is green.🤖 Generated with Claude Code