feat(core): core-side wiring for kimi-linear (k3) serving - #306
Merged
Conversation
Contributor
Author
|
Validation recorded internally; CI is the public record. |
Worker carry-loop, KDA state GPU manager and kimi-linear KV coordinator under kv_cache/, host-KV profile entry, model/tokenizer registry rows, and the worker_manager seam. Core half of the kimi-linear bring-up; the model-side stack (models/moonshotai/kimi_linear/**) lands in the companion model PR.
Worker half of the kimi-linear M1 fix batch: F5 worker-side sequencing and the decode heartbeat. The planner-override half (F6) rides in the companion model PR.
_report_completion read prompt/decoded_length from rank-0's local replica, stale for non-rank-0-owned sequences on the outer-loop drain path (only the completion bit is all-reduced there) -> short usage.completion_tokens now, routinely wrong at 128-seq L2. Sync metadata before reporting. Also reset seq._buffer_slot after free_slot so a re-entered report cannot free a slot already owned by another sequence.
KDAStateGPUManager owns allocation, slot accounting, and F4 zero-on-alloc; conv pools migrate to the causal_conv1d.cu contract ((L, slots, dim, W-1), contiguous per-layer 3-D views), and each pool is a single fixed-address allocation — the CUDA-graph capture requirement. The wrapper-side views and the graph-readiness tests ride in the companion model PR.
uuid[:8] is unique for random-hex ids but /v1/batches reuses structured
custom_ids (mmlu-{run}-{idx}) whose first 8 chars are identical across a
whole run — the warning named 'mmlu-030' for every sequence, hiding which
one actually looped and costing real debugging time. The remaining 58
uuid[:8] log sites are queued for the pre-PR sweep.
Registry: kimi-k3 / moonshotai/Kimi-K3 now resolve to KimiK3Tokenizer, and _import_tokenizers warns per-module on import failure instead of one module's error aborting the loop. kimi_linear no longer owns the kimi_k3 name. Scheduler: prompt construction (chat rendering, tool formatting) can raise; previously the exception reached _run's bare logger.exception with the batch already marked IN_PROGRESS, so no terminal status was ever set and the client polled forever. The batch now fails LOUDLY with the exception in the error field. Deliberately still batch-granular -- per-request rejection at admission is a named follow-up; this change only makes the failure visible and terminal, per the no-silent-fallback rule. Core PR per PR_MERGE_POLICY (batchgen/config/ and batchgen/server/ are outside MODEL_ALLOW_RE).
Three guards on the host-to-device weight copy, replacing an error path that constructed a std::runtime_error and discarded it before falling through to weights_copy_complete: - dst.find() instead of operator[], which default-inserts an undefined torch::Tensor on every miss into a buffer map reused across ring slots. - A host tensor with no GPU slot now throws. Continuing dropped it silently and the consumer read whatever the slot last held -- the GLM-5 Q/K RMSNorm incident class (glm5_initializer.py:141-146). - src/dst byte-size equality before the copy. blocking_copy_ writes src_byte_size bytes with no bound check; a short slot is overrun into its neighbour, a long one keeps a stale tail, both silent. Latent today (48B name maps agree, no occurrences in current logs), but K3 lines up 497,220 checkpoint tensors against a freshly written module_shapes -- exactly the condition these guards exist for. Per the 2026-08-04 ledger this lands as its own fix PR before the M3 force-stream rehearsal, whose negative tests (byte-size mismatch, omitted tensor) must fail loudly against THIS code and hang/corrupt without it. Compile verification is staged for the GPU machine; this translation unit builds only there.
The repack reinterprets the uint8 buffer as packed uint4b8 and converts the scale tensor VALUE to bfloat16 -- an E8M0 exponent byte of 121 would become 121.0 instead of 2**(121-127). Both corruptions are silent, and convert_checkpoint.py:144 passes marlin=True by default, so converting the K3 checkpoint today would quietly produce garbage marlin weights. Scoped to tensors the repack actually matches (uint8 scale = E8M0 = MXFP4): an unscoped check would reject checkpoints the repack never touches. The MXFP4-aware repack is the in-flight kernel work unit; until it lands, K3 conversion must run --no-marlin.
Worker half of the block-residual carry (the model half, kimi_linear model.py's _apply_output_attn_res, rides in the companion model PR). K3 replaces the classic residual body with Block Attention Residuals: a depth-mix across block boundaries whose state is carried layer to layer. The prepacked prefill loop did `hidden_states = layer_outputs[0]`, which never supplied that state and discarded the second return value. The model would have loaded, streamed and produced logits while every layer saw a zero-width residual — wrong text, no error, nothing to grep for. That is the exact silent-wrong-output class this project bans, and it is why the first real 8K run had not been attempted. The loop now mirrors KimiLinearModel.forward, the eager reference: initialise the (N, 0, hidden) residual, thread it through layers that ask for it, and apply the output depth-mix BEFORE the final norm — mix-then-norm is the reference order and swapping it changes every output. Models without attn_res_block_size keep the exact previous code path, so kimi-linear-48B, GLM-5 and K2.5 are untouched. The worker calls the model's _apply_output_attn_res method 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.
K3 was aliased onto _KIMI_LINEAR_MLA_PROFILE, which declares num_layers=27 because that is the 48B's depth. K3 has 93 engine layers with 24 MLA layers at engine indices 3,7,...,87,91,92, and wrappers.py::_offload_prepacked_kv indexes the pool by the ENGINE layer index rather than by a dense MLA counter. Every MLA layer at index >= 27 therefore wrote past the end of the pool -- silent host-memory corruption starting at layer 28 of the very first prefill, with no error and nothing to grep for. The geometry is otherwise identical (compressed_kv_dim 576, MQA single head, bf16), so this is purely the layer count. K3 now resolves to its own kimi_k3_mla key; the kimi-linear aliases are unchanged, so the 48B keeps the 27-layer profile it wants. Found by reading the load path while diagnosing an unrelated OOM, before it could execute -- the first real K3 prefill had not reached layer 28.
Prefill computes the first generated token (`_select_tokens` on the last-token logits) and appends it to output_tokens. For a prefill-only model that token is then thrown away: a max_tokens=1 request still enters decode (PREFILL_PLAN C4), decode raises, and the client receives the decode error string instead of the token. Full 93-layer Kimi-K3 prefills have now completed successfully (all 82,432 experts streamed) without anyone being able to see what the model produced. Rank 0, ids only -- the worker has no tokenizer, so decode them client-side. The proper fix is to return the prefill token when max_tokens=1 instead of entering decode; that is C4 and a larger change.
Two prefill-memory fixes from batchgen_design/model_support/kimi_k3/PREFILL_MEMORY_AUDIT.md section 7. Depends on the model-side prealloc commit in the companion model PR, which adds KimiLinearModel._new_block_residual. fix 3 — the prepack-prefill loop seeded block_residual itself with new_zeros(S, 0, H), so every block boundary reallocated it by cat, and the old and new buffers were co-live at K3's last boundary. It now asks the model for a zero-column view of a buffer preallocated for all 8 boundaries. The existing `block_residual = None` immediately above is load-bearing: it drops the previous micro-batch's view before the next buffer is allocated. fix 4 — `hidden_states = inputs_embeds.unsqueeze(0)` is a VIEW, so hidden_states alone already keeps the embedding storage alive for exactly as long as layer 0 needs it. Keeping inputs_embeds bound as well pinned that storage for the whole 93-layer forward, dead weight across layers 1-92. One `del`.
Contract
--------
Prefill already samples the first token and appends it through the normal
decode write path: the writeback loop at the end of prefill()/
prefill_prepacked() writes it into query_book[local_idx].decoded_tokens at
seq.decoded_length and bumps decoded_length. But _update_batch_status(
prefill_uuids, PREFILLED) then handed EVERY prefilled sequence to the decode
phase unconditionally, so a max_tokens=1 request had to make a full decode
round-trip to be told it was already finished.
_finish_prefill_completed_sequences() now runs right after that status
update and completes the sequences whose decode budget is already satisfied
(decoded_length >= max_decode_length). They are reported through the same
path decode uses -- incremental writer, _gather_completed_tokens,
KV release, PREFILLED -> COMPLETED, _report_completion -- so the HTTP
response carries the decoded token text with the existing length-capped
finish_reason. No parallel token-recording path is introduced and the
handoff of output_tokens to decode for the surviving sequences is untouched,
so nothing is double-appended.
Only the length budget is tested. That is decode's own first completion test
(CompletionHandler.is_sequence_completed / _check_and_handle_completions) and
the test that wins in get_finish_reason, so the reported finish_reason is
unchanged. Sequences that stop for any other reason (EOS, context limit) stay
PREFILLED and reach decode exactly as before.
Skipping decode
---------------
Sequences still needing tokens remain PREFILLED, so the decode while-loop
condition (has_prefilled() or has_in_decode() or has_on_hold()) is false only
when NOTHING remains -- no configure_decoding, no decode-model load. The skip
therefore falls out of replicated batch state rather than a separate branch.
Rank alignment: decoded_length is advanced only on the owning rank, so the
completed set is derived AFTER _sync_sequence_metadata() replicates it to
every rank. Every rank computes the identical set from identical batch-global
state; a rank-divergent skip would deadlock the next collective. That sync is
needed anyway, because _report_completion reads decoded_length on rank 0 for
sequences owned elsewhere (same reason the decode path syncs before
reporting).
Behaviour change
----------------
max_tokens > 1: unchanged on every model, including the 48B path -- still
enters decode (and under Kimi-K3 stream_all_modules still raises the M-PR-6
prefill-only error, which is correct until decode exists).
max_tokens == 1: now returns the prefill token WITHOUT the decode round-trip.
This is a deliberate, documented improvement, not just a K3 fix -- the 48B
path used to load the decode model and configure decoding purely to discover
the sequence was already complete. Same token, same finish_reason ("length"),
one less model load. For a prefill-only model it is the difference between
the answer and decode's NotImplementedError string reaching the client.
Verified on CPU: the branch logic was extracted from this source with ast and
driven with a stubbed 3-rank mixed batch (max_tokens 1,1,5,3) -- identical
completed set on every rank, correct COMPLETED/PREFILLED split, token
recorded exactly once, identical collective sequence on every rank, decode
skipped only when all sequences finish. py_compile + tests/test_kimi_k3_model.py
(51 passed).
Nothing model-specific is introduced. The problem ----------- QueryBookBufferPool allocated torch.zeros((num_sequences, model_context_length)) per worker. With a large --max-pool-size and K3's 1,048,576-token context that is 8 B x rows x context of int64 zeros for input_ids, per worker, times every rank on the node -- and generate_persistent() passed model_context_length as max_decoding_length too, so decoded_tokens was a second buffer of the same size. That OOM-killed the node twice and made --max-pool-size 1 load-bearing. get_input_ids_view(slot, seq_extended_size) always slices to the real length; the full width was never used. Fix 1: share it. The tokenized global batch is IDENTICAL on every rank (_tokenize_global_batch all-gathers the results to all of them), and every rank fills a row for every sequence, so the ranks were holding world_size byte-identical copies. input_ids_buffer is now ONE multiprocessing shared_memory segment per node: the node's rank%NUM_GPUS_PER_NODE==0 creates it, dist.barrier, everyone else attaches (allocate_node_shared_int64). Not the parent-allocated allocate_host_kv_cache pattern, because the parent cannot size this buffer -- the widths come from the tokenized requests, which only exist inside the worker loop after dist is up. decoded_tokens_buffer stays PRIVATE per rank: only the owning rank writes a sequence's decoded tokens, so those copies legitimately differ. It is, however, now sized by need as well -- the pool-mode path set it to model_context_length, the second full-width buffer, not what the spec's "already width-bounded" assumed. Fix 2: size by need. Width = the widest seq_extended_size the batch will actually ask for (prompt + THAT request's decode budget), capped at the model context length. Never the context length, never --max-pool-size, which keeps its row-count meaning: allocate_slot() still hard-fails past it. In pool mode the allocation is deferred out of generate_persistent() into the first admission, because that is the first moment either width is known. Overflow is never silent. A later admission needing more grows the pool with a WARNING naming both old and new sizes, copies the live rows across (adopt()) and rebinds every seq.input_ids / query_book view; the superseded segment stays mapped (untracked views must not be unmapped under) with its name unlinked. get_input_ids_view() raises QueryBookPoolCapacityError rather than silently returning a SHORT view if anything slips past that, and allocate_slot()'s exhaustion error is now the same named type. The slot trap ------------- A shared buffer is only safe if slot -> row is globally consistent. It is, almost: _tokenize_global_batch Phase 3 and _tokenize_admitted_sequences allocate a slot for EVERY sequence on EVERY rank in the same order, and _report_completion frees on every rank. The one exception was host-KV migration -- the source rank called free_slot() and set _buffer_slot = -1 while the destination reused that same slot index and every other rank kept it. So the source could hand row S to a new admission while everyone else still read S as the migrated sequence, and the -1 left behind made an eviction re-entry write into row -1, i.e. the LAST row of the pool. Chosen fix: derive nothing new -- just stop the divergence. The source-side free is removed (the slot IS still that sequence's, globally), and the two places that could reintroduce a rank-local slot now hard-fail with QueryBookPoolCapacityError instead of quietly allocating: migration-receive with _buffer_slot < 0, and re-entry with _buffer_slot < 0. Per-rank disjoint row ranges were rejected: every rank needs a row for every sequence, so disjoint ranges would save nothing. Measured (CPU, 4 processes, real code lifted from this file with ast) --------------------------------------------------------------------- - one segment: each process sees every other process's disjoint row writes; a private-pool control run sees none of them - per-process private growth for a 64 MiB buffer, 4 procs: shared 0.2/0.2/0.2/0.1 MiB (total 0.7) vs private 64.2 x4 (total 256.8) - identical buffer digest in every process - QueryBookPoolCapacityError fires on both width overflow and row overflow, naming both sizes - grow: adopt() preserves every live row, carries slot bookkeeping, the wider view then works py_compile + tests/test_kimi_k3_model.py (51 passed).
Fixes a regression introduced by the preceding C4 commit.
Symptom
-------
Every /v1/inference request whose sequences complete at prefill (max_tokens=1,
the C4 path) came back with "Results unexpectedly empty after inference"
(server_worker_main_loop.py:552). The batch endpoint was unaffected.
Root cause
----------
generate() gathers legacy results at the very end by iterating
self._local_to_uuid_map. _finish_prefill_completed_sequences() completes C4
sequences during PREFILL and calls _report_completion(), which calls
release_local_query_slot() (query_book.py:51) -- and that pops
uuid_to_local_map, local_to_uuid_map AND the query_book entry. By the time the
gather runs, a C4-completed sequence is gone from every structure the gather
reads, so local_results is empty on every rank and rank 0 returns {}.
The batch path never saw this because C4 already feeds it through
_submit_completed_to_incremental_writer() before the pop.
Fix
---
Capture the decoded text while the slot still exists -- immediately before the
_report_completion loop -- into self._prefill_completed_results
(global_idx -> str), and have the legacy gather seed local_results from it
before walking the still-live map. Ordering is the whole fix: capture must
precede the pop.
The capture is gated on `self._response_queue is None`, i.e. legacy mode only.
Pool/batch mode is fed by _report_completion and the incremental writer, so
accumulating there as well would both be redundant and grow without bound in a
persistent server that never drains the store. Nothing is reported twice on
either path: the store feeds only generate()'s return value, which pool mode
discards.
The store is cleared in process_new_batch() alongside the new SequenceBatch, so
it never carries across batches.
Verified on CPU (real method lifted from this source with ast, stubbed so that
_report_completion actually pops -- a capture-after-pop ordering mistake fails
the test):
- legacy (no response queue): gather is NON-empty, contains both C4 texts plus
the survivor gathered from the live map
- pool (response queue present): store stays empty, each sequence reported
exactly once, incremental writer called exactly once
- no C4 sequences: store untouched, gather unchanged
py_compile + tests/test_kimi_k3_model.py.
Implements the prefill-metrics proposal, with one deliberate departure
documented below.
Why
---
prefill_s -- the most-quoted BatchGen performance number -- is scraped out of a
tqdm progress bar. That string is presentation, not an API: it breaks on any
desc=/bar_format/width change, it is emitted with \r so one physical log line
can carry several bars, it reports a rate rather than an elapsed time, and tqdm
prints the final bar twice with the copies disagreeing in the last digit.
prefill_prepacked now emits one JSON line after the forward loop:
[METRICS] {"phase":"prefill","prefill_s":...,"sequences":...,
"tokens_total":...,"seq_len_min":...,"seq_len_max":...,
"micro_batches":...,"max_tokens_per_micro_batch":...,
"world_size":...,"rank":...,"first_sampled_token_ids":[...]}
The report tooling already prefers this over the tqdm scrape and drops
the "came from a progress bar" warning when it is present. Additive-only: the
presence of "phase" is the whole handshake, so older readers are unaffected.
The timing window starts immediately before the `with torch.inference_mode():`
that wraps the loop, so it is the pure forward pass. configure_prefill is NOT
folded in -- it is already reported separately as `Config completed`.
Departure from the proposal: emitted by EVERY prefilling rank, not rank 0
-------------------------------------------------------------------------
The proposal specifies "rank 0 only" and fixes "rank": 0. That is wrong, and
a 131,069-token run proved it: the batch was owned by rank 2 and the run
was left unmeasurable, reporting no prefill record at all.
Two independent reasons rank-0 gating cannot work:
1. prefill_prepacked is called only under `if local_prefill_indices:`, so a
rank that owns none of the batch never reaches the emit point at all -- the
line is simply absent, not zero. The same is true of the tqdm bar, which is
additionally disable=(self.rank != 0), so nothing is emitted anywhere.
2. The proposal's snippet reads output_tokens[0] unguarded, which is an
IndexError on a rank that ran zero micro-batches.
So the line carries the emitting rank's real id and its LOCAL sequence/token
counts. The wall time for the batch is the MAX of prefill_s over the emitting
ranks; aggregating that is the parser's job.
PARSER DEPENDENCY: the report tooling currently takes prefill_s from the first
matching phase=="prefill" line. Until it aggregates over ranks, a multi-rank
prefill will record an arbitrary rank's forward time rather than the max. This
is still strictly better than today, where such a run yields no number at all.
Verified on CPU (dict literal lifted from this source with ast and evaluated
against stubs, so the contract is checked against shipping code): all 11
contract fields present and correctly typed, no undeclared fields, rank
reflects the emitting rank, compact single-line JSON that re-parses, and both
crashes the proposal's snippet would have hit (empty output_tokens, empty
seq_lengths_list) return safe values instead.
py_compile + tests/test_kimi_k3_model.py (51 passed).
Follow-up to the C4 gather fix, which introduced self._prefill_completed_results in __init__. Hot reload (/v1/reload) rebinds methods on a LIVE worker instance and never re-runs __init__, so a reloaded process would get the new code without the new state and raise AttributeError on the first C4 completion. This is not hypothetical: _validate_reload (server_worker_main_loop.py:68) diffs the two __init__ sources, lists attributes the running worker lacks, and logs "These will cause AttributeError if accessed." It reports the hazard and deliberately does not repair it. Both read sites now tolerate the attribute being absent and create it on first use. __init__ still declares it, so a cold start is unchanged. Verified: the C4 result-store checks still pass on both paths (legacy gather non-empty with C4 texts; pool mode accumulates nothing and reports once). py_compile.
Implements the ruling that /v1/inference is DEPRECATED and all
inference goes through the batch API.
The hazard
----------
The legacy path carried no request-id routing. On a pool-mode server
worker_manager.infer() put its payload on the shared worker request queue and
then blocked on response_queue.get(). The worker admission loop recognises only
None and {"type": "admit"}, so the legacy payload matched no branch and was
SILENTLY DROPPED -- while the caller sat on the shared response queue and took
the next completion belonging to somebody else's batch. That is not theory: it
corrupted a real production batch, which was left parked in in_progress with
its completion consumed by a /v1/inference request.
Rejected at the earliest seam
-----------------------------
The HTTP handler now raises 410 immediately, before any queue interaction, with
a JSON detail naming the deprecation and pointing at /v1/batches. The route is
kept rather than deleted so a caller gets the explanation instead of a 404 it
would read as a typo, and it takes NO request body so an unparseable legacy
payload still gets the deprecation notice rather than a 422 about its schema.
The old body is deleted, not left unreachable below the raise.
BatchGenHttpClient.submit_inference raises the same named error without touching
the network. The symbol is kept deliberately: deleting it yields an
AttributeError, which tells a caller only that something vanished.
Defense in depth, and what it deliberately does NOT catch
--------------------------------------------------------
Both worker admission polls -- _poll_admissions and the idle-wait block inside
generate() -- now raise LegacyInferenceDeprecated on a dict carrying "prompts",
the shape worker_manager.infer builds. Reaching either line means some producer
other than the now-closed HTTP route is putting legacy payloads on the queue, so
failing loudly is right; parking is exactly the bug.
This is deliberately NOT a catch-all on unrecognised messages.
{"command": "reload"} also reaches _poll_admissions and is dropped there today,
so a catch-all would convert every mid-batch /v1/reload into a dead server. The
guard keys on the legacy shape only, and a test asserts the reload message still
passes through untouched.
Relationship to the C4 gather fix
---------------------------------
The C4 legacy-gather fix STANDS and is not reverted. Verifying this change showed its premise
was too narrow: the legacy result gather at the end of generate() is not
exclusive to /v1/inference. worker.infer() has a second caller,
batch_scheduler.py:313, which is the /v1/batches path taken whenever
max_pool_size == 0. So the gather remains live for legacy-mode batches -- a
batch-API mode, squarely inside the ruling -- while the deprecated HTTP entry is
rejected at the door. Reverting it would have re-broken max_tokens=1 batches on
a --max-pool-size 0 server. (It would also have conflicted: the hot-reload
commit rewrote both hunks the gather fix added.)
The shared seam
---------------
batchgen/deprecation.py holds the message and exception. The three consumers
cannot import one another: batchgen_client is imported by batchgen/__init__ and
must carry no heavy deps, http_server pulls FastAPI/pydantic, batchgen_worker
pulls torch. A module that imports nothing is the only seam all three can share,
and a test asserts it stays that way.
Verified on CPU (py3.11):
- tests/test_kimi_k3_model.py: 51 passed
- py_compile on all four files
- 38 checks over the SHIPPING source: the handler is checked structurally with
ast (http_server will not import on the dev box -- JIT core_engine needs
ninja) and asserts the body is one raise, 410, the shared constants, no body
param, and no reference to worker/infer/request_queue/response_queue/put/get;
_poll_admissions is LIFTED from source with ast and executed against stubs,
confirming a legacy payload raises and admits nothing, while admit, reload,
None and an empty queue all keep their prior behaviour. Mutation-checked:
stubbing the guard out makes the legacy-payload assertion fail, so the test
is not vacuous.
NOT verified live. The handler is not hot-reloadable: _reload_worker_module
reloads batchgen_worker plus four dependency modules, and http_server is not
among them; the routes are closures registered in create_app() in the parent
process. The curl check rides the next scheduled restart rather than forcing one.
…anch The copy-live-rows / re-point-views path in _ensure_buffer_pool has never run with a live sequence: wave admission serialises batches so every observed grow rebound 0. This CPU test forces a grow while a mid-decode sequence occupies a row and asserts the row survives byte-for-byte, the slot mapping and query_book entry rebind onto the grown buffer, and a second admission does not clobber the first. It exercises the real batchgen_worker.py source (AST-extracted, since the module import pulls the JIT core_engine) and carries a mutation check that neuters the row-copy to prove the assertions are load-bearing.
Andrewxu313
force-pushed
the
tairanxu/k3-prefill-core
branch
from
August 12, 2026 21:45
1dd89e2 to
8d961b7
Compare
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
Core-side scaffolding required to serve the kimi-linear model family (Kimi-K3): a
manager-authoritative KDA state GPU pool and a kimi-linear KV coordinator under
kv_cache/, a per-model host-KV profile so kimi-k3 sizes its host region from its ownengine-layer count, worker support for carrying block attention residuals through the
packed prefill loop, prefill-only completion for
max_tokens=1requests (result storekept visible to the legacy gather path and across hot reloads), a structured
[METRICS]prefill record emitted per prefilling rank, a shared per-node input-ids pool sized by
need, hard-fail semantics in the HtoD copy loop (missing slot / byte-size mismatch now
raise), tokenizer routing for kimi-k3 with terminal prompt-construction failures, a
ckpt-converter guard refusing MXFP4 tensors on the uniform-INT4 marlin repack path, and
admission-side rejection of
/v1/inferencewhile the batch pool is active. CompanionPRs:
tairanxu/k3-prefill-kernel(MoE kernels) andtairanxu/k3-prefill-model(modelstack); this PR merges first.
Why
The kimi-linear family (linear KDA attention + interleaved MLA + block attention
residuals + MXFP4 MoE) is the first model whose serving path needs per-layer recurrent
state pools, residual carry across the prefill loop, and prefill-only completion in the
core worker. Each change is the minimal core seam the model PR builds on, split out per
the merge-policy contract. Validated for long-context prefill and for decode parity on
the existing supported models (no behavior change with the feature paths off).
Type of Change
modelkernelcorefixinfradocsFile changes
batchgen/batchgen_client.pybatchgen/batchgen_worker.pybatchgen/ckpt_converter/ckpt_converter.pybatchgen/config/model_registry.pybatchgen/config/tokenizer_registry.pybatchgen/deprecation.pybatchgen/kv_cache/host_kv_mananger_config.pybatchgen/kv_cache/kda_state_gpu_manager.pybatchgen/kv_cache/kimi_linear_kv_coordinator.pybatchgen/server/batch_scheduler.pybatchgen/server/http_server.py/v1/inferenceat the door while the pool servesbatchgen/server/worker_manager.pycore/HtoD_Engine/HtoD_Engine.cuChecklist
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