Kimi K3 support with DSpark speculative decoding

[merceod]

Kimi K3 in M*. The model, the resources it needs, the kernels, and speculative decoding with the DSpark draft. It serves the expert-pruned mgoin/Kimi-K3-pruned75 checkpoint (224 experts, MXFP4) on one 8xH100 node at TP8. The full checkpoint has the same architecture and needs more nodes.

What does this PR do?

Model (mstar/model/kimi_k3/). A plain reference implementation that the tests check against the HF modeling code, plus the serving version. KDA layers run on FlashKDA and fla kernels with their state in a slot pool (RecurrentStatePool, a KDA variant of LinearAttnManager), gated NoPE MLA runs on FlashInfer with absorbed projections, the latent MoE runs on a vendored copy of vLLM's Marlin MXFP4 kernel (JIT built once per machine), and the block attention residuals, norms, SiTU, router and residual adds are fused into a few Triton kernels. Decode runs in CUDA graphs at batch sizes 1 to 64, prefill eagerly in steps of up to 8 prompts.

Resources and engine. Merged input projections behind one GEMM, expert parallelism as an ExpertSharding placement, symmetric-memory all-reduce on by default for single-node TP (Lamport one-shot for small tensors, NVLink multicast above that, NCCL for large prefill messages), TP async scheduling as the model's default, decode rows riding in prefill steps as an option (mixed_prefill_decode), a torch.profiler window over served steps (MSTAR_TORCH_PROFILE), and a few fixes that came up on the way to TP8 (token counting on follower ranks, SIGTERM in the workers, the shared JIT lock, the loader's device mapping).

Speculative decoding. model_kwargs.speculative_tokens turns each decode step into a verify of k drafted tokens inside the same declare / admit / plan / forward / commit protocol and the same captured graphs. Verification is greedy or sampled (FlashInfer's chain sampling), the KV trims rejected tails, and the KDA state is never rolled back. The pool keeps a checkpoint plus the pending prefix and each step runs prefix and block in one Triton launch. speculative_draft loads the Inferact/Kimi-K3-DSpark draft and speculative_schedule picks the block length per capture bucket ({16: 7, 32: 4, 64: 0} for pruned75).

Prefill. The KDA prefill path used to sync with the host twice per layer through fla's chunk-table cache. That is gone. The Marlin prefill slice is 8192 tokens, the routed latent is reduce-scattered at prefill sizes with a sharded RMSNorm, and slices of 6144 tokens or more run the experts as bf16 grouped GEMMs on weights dequantized once per layer from the Marlin tiles (MSTAR_MOE_BF16_TOKENS, warmed at setup).

Numbers

pruned75 at TP8 on 8xH100, 1024 tokens in and 512 out, 8C requests per level, output tokens per second. vLLM 0.28 plain with prefix caching off on the same prompts, each system on its own node of the same kind. Measured on the branch before main was merged in.

C	vLLM plain	M* plain	ratio	M* + DSpark, schedule {16: 7, 32: 4, 64: 0}	ratio	TTFT mean, vLLM / M* plain (ms)
1	55.6	59.3	1.07	127.6	2.30	227 / 206
2	102.3	108.2	1.06	163.2	1.60	287 / 305
4	184.9	191.2	1.03	334.8	1.81	491 / 530
8	321.3	323.7	1.01	465.5	1.45	948 / 1155
16	527.6	537.7	1.02	691.8	1.31	1350 / 1469
32	821.7	840.6	1.02	937.0	1.14	2093 / 2240
64	936.0	1188.4	1.27	1118.4	1.19	11971 / 4113

Steady decode (inter-token latency after the first token) is under vLLM's at every level up to C 32, 16.5 vs 17.6 ms at C 1 and 33.8 vs 34.9 at C 32. The C 8 TTFT mean carries one 3 s request that paid the bf16 path's first compile (the p50 is 1039 ms). The warm-up at setup that removes that landed after this sweep and has not been re-swept on the big model. The C 1 and C 2 draft rows have only 8 and 16 requests and the acceptance on this pruned checkpoint moves with the text, so read them as noisy.

The vLLM column is plain vLLM 0.28 without speculation, which is its best configuration on H100. We did measure vLLM with the same DSpark draft and it came out far slower than plain vLLM, 33 to 58 tok/s at C 1 to 16, because vLLM's DSpark needs the FlashInfer MLA backend that is Blackwell-only, falls back to Triton MLA on H100 and loses full CUDA graphs (about 176 ms per step). So the "M* plain" row is the like-for-like comparison, and the draft rows show what speculation adds on our side against the fastest vLLM setup available on this hardware, not against vLLM's speculative path. On Blackwell, where vLLM's DSpark runs at full speed, the fair comparison would be both systems with the draft, and we have not measured that yet.

How to run it

mstar-serve --config configs/kimi_k3_pruned75_tp8.yaml --port 8200 --tensor-comm-protocol SHM

configs/kimi_k3_pruned75_tp8_dspark_sched.yaml adds the draft with the block schedule. The tiny random-weight configs (kimi_k3_tiny*.yaml, TP1 and TP2 variants) exercise every path on one or two GPUs.

How was it tested?

• CPU tests. test/kimi_k3 (parity with the HF modeling code on the tiny checkpoint, dense vs paged, TP sharding, quantized experts, tokenizer) and the modular tests for the recurrent pool, the verify path, acceptance, the sampler and the worker.
• GPU tests on H100. Bit-exactness and reference tests for every fused kernel, the Marlin, FlashInfer and MXFP4 backends, the paged paths on the pool, the verify path against the torch reference, the Marlin layout inverse.
• Served gates on 1 and 2 GPUs, run on every merged step. Greedy texts identical to the dense reference with speculation on and off (tiny fp32 TP1, tiny TP2), 16/16 smokes on the real-dims TP2 checkpoint with the plain, draft and schedule servers, and every CUDA-graph bucket captured (a failed capture falls back to eager, so the gate counts "Failed to capture" lines).
• The sweeps above on 8xH100, plus profiler windows of both engines' decode and prefill steps.

main is merged in (one conflict, the model registry).

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

[NSagan271]

Some blocking issues, especially some correctness issues in the speculative decoding resource path and the fact that a combination of inline chunks and tensors read through transport can result in out-of-order outputs.

[NSagan271]

This class now includes three different implementations of all-reduce; it works but is not the most modular and readable. I would recommend moving some of the logic (the all-reduce itself, as well as logic for when each method is applicable) to a new abstract Communicator class with subclasses for SymmMem, NCCL, and Lamport (or for NCCL and SymmMem, with the Lamport logic in SymmMem, which may be technically more correct).

[NSagan271]

@merceod how do you plan to compose this with the BatchedModelOutput I added in #258? I do like that the BatchedModelOutput makes exposes this behavior to the model implementor (e.g., they won't see a performance dip if they actually return something non-contiguous; they also can return a single tensor instead of splitting by rid), but it is also nice that the StorageSpan objects are created automatically from whatever any model returns.

[NSagan271]

The fallback is not cuda graph compatible, so I'd prefer to fail loudly here if use_cuda_graph is True, instead of waiting for an error during cuda graph capture itself.

[NSagan271]

The CPU page pool currently allocates its tensor based on the default KVLayout.NHD. It should be changed to take the dimensions of the real KV cache (after the first two dimensions, which are config.num_layers, and max_cpu_pages).

[NSagan271]

Nit: can use an enum here

[NSagan271]

Every new shape allocates a new buffer; this is fine for cuda graph mode because the number of shapes is finite, but will keep on allocating new buffer for eager mode.

[NSagan271]

Would _symm_multimem == False ever be a correctness issue; e.g., for BAGEL, where the MoT has two different output projections?

[NSagan271]

This could still cause a hang if a server dies while holding the lock and is restarted before 20 minutes. Claude's recommendation is (take it with a grain of salt):

Put our own fcntl.flock around load(). Use a sidecar file, e.g. <build_dir>/mstar.flock.
The kernel drops a flock when the holder dies, including on SIGKILL. On NFS, Linux turns flock into a POSIX lock that the server drops when the dead client's lease runs out. So the lock can't be left behind, even across nodes.
While holding it, any torch lock file you find must be stale, because every mstar loader goes through the wrapper. You can delete it unconditionally, and both the /proc scan and the 20-minute check go away.
One helper would cover align.py and marlin/init.py.
Caveat: on an NFS mount with nolock, flock only works within a node. That is no worse than today's /proc scan.

[NSagan271]

If there are ever inline tensors queued after larger chunks, the inline tensors can get sent to the client while the previous chunks are still being read, resulting in an out-of-order output. This is not relevant to Kimi because all output chunks are small and can be inlined, but would be relevant for an audio model that has chunk-to-chunk variability.

[NSagan271]

Does the arch=compute_90,code=sm_90" block the marlin kernels from building on Blackwell? If so marlin_supports in mstar/model/kimi_k3/components/language_model.py should check for sm capability before selecting marlin kernels.