[Model Integration] Adding Kimi K2.7 Code - #192
Open
garv901 wants to merge 5 commits into
Open
Conversation
Implements Kimi-K2.7's text model (HF DeepseekV3ForCausalLM) as a new,
self-contained mstar model package. Golden-verified against vLLM references on a
reduced config and validated end-to-end through the real paged FlashInfer cache
(prefill + decode). Reuses existing mstar abstractions; no shared code modified.
Changes added (mstar/model/kimi_k2_7/):
- config.py: KimiK2Config — MLA latent dims, fine-grained MoE grouping,
deepseek_yarn RoPE + reduced() test config.
- MLA attention (naive/materialized) + deepseek_yarn RoPE (GPT-J interleaved).
- Fine-grained MoE: group-limited sigmoid noaux_tc router + ungated shared
expert (reuses the fused-expert GEMM).
- Decoder layer (dense + MoE), KimiLanguageModel + KimiForCausalLM.
- weight_loader.py: HF loader (name remap + stacked per-expert gate/up/down
rules; MLA loaded by name); INT4 dequant-on-load designed and deferred.
- submodules.py: KimiLLMSubmodule (ARNodeSubmodule prefill/decode lifecycle,
YARN position_ids, CUDA-graph configs); real get_submodule
(meta -> to_empty -> load_weights); registry entry + configs/kimi_k2_7.yaml.
- Tests: dummy-mode modular + GPU integration goldens (components, MoE, MLA,
full forward, real paged attention, weight loading, submodule e2e).
Major issues resolved:
- FlashInfer SM90 prefill static_asserts head_dim_vo in {64,128,256}, so the
naive-MLA head_dim (192 real / 24 reduced) won't JIT-build. Mitigation: pad
q/k/v to the next supported dim and fold a compensating softmax scale
(mscale^2 * sqrt(padded/qk_head_dim)); golden-verified vs SDPA.
- KimiYarnRotaryEmbedding computed inv_freq as an __init__ buffer that
meta->to_empty leaves uninitialized (would silently corrupt YARN on the real
load path). Fix: compute lazily; audit confirms the loaded model has zero
stray buffers.
Builds on the DeepSeek-V3 text backbone to make the reduced/synthetic config both servable end-to-end and correct under tensor parallelism. Golden-verified on GPU; no shared code touched. INT4/fp8, TP8 and the real 1T checkpoint remain deferred (Phase 4). Serve (single-GPU, reduced/synthetic): - kimi_model.py: model_kwargs hook (checkpoint_path / config_variant / tokenizer_mode) lets a serving YAML point the model at a local reduced checkpoint with a UTF-8 byte tokenizer (reduced vocab_size=256). Full-size default behaviour unchanged. - config.py + submodules.py: CUDA-graph prefill capture grid is config-driven (KimiK2Config.prefill_token_buckets / prefill_capture_batch_sizes; reduced() uses a tiny [64]/[1] grid). No env vars. - configs/kimi_k2_7_repro.yaml + test/integration/test_kimi_serve_e2e.py. Verified: live mstar-serve (api_server -> conductor -> worker -> KV_CACHE engine -> decode loop) streams tokens; deterministic in-process gate. Tensor parallelism (tp=2, reduced config): - components/attention.py: MLA head-sharding (rank-local head count; column-parallel q_b/kv_b + row-parallel o_proj; latent down-projs replicated). - components/moe.py: MoE intermediate-sharding (fused_experts reduce_results + all-reduce), mirroring ParallelSparseMoeBlock; router replicated. tp=1 is byte-identical. - configs/kimi_k2_7_tp2.yaml + test/integration/test_kimi_tp.py. Verified: tp=2 == tp=1 (real 2-GPU NCCL + in-process rank-simulation goldens).
- quantization.py: compressed-tensors INT4 parser (dequant-on-load) + Triton fused_moe_kernel_w4a16 (in-kernel dequant of packed int32 experts). KimiSparseMoeBlock keeps experts packed when quantized; bf16 fused-MoE path (shared with Qwen3-Omni) unchanged. - config/kimi_model/weight_loader: k27_code config, nested text_config.quantization_config auto-read, language_model. prefix strip, packed-expert stacked rules; TP8 serve configs (kimi_k2_7_code_tp8[_shm].yaml, SHM = proven load path). - TP robustness: 2h NCCL timeout for large-checkpoint bring-up; worker tp-leader gate on new-token counting. Tests: CPU quant/wiring goldens + GPU quant/kernel goldens.
…n fallback) Port the production Marlin INT4 GEMM (the vLLM/sglang serving path for compressed-tensors W4A16) into mstar for the Kimi routed experts, as a clean, model-agnostic quant module. - utils/marlin/: vendored vLLM Marlin CUDA (device code verbatim + classic-ABI host shims), JIT-compiled to torch.ops._mstar_marlin_C.* via the align.py precedent — no vllm/sgl_kernel runtime dep. Repack + fused_marlin_moe launchers. - model/components/quantization/: reusable FusedMoEQuantizeMethod seam, a generic process_weights_after_loading post-load walker, and MarlinMoEMethod. - Kimi wiring: KimiSparseMoeBlock repacks its packed experts to Marlin layout at load and dispatches through it; quant_kernel=auto picks Marlin on sm80+ with the Triton W4A16 path as fallback (quant_kernel=marlin|triton force either). - Tests: kernel + block goldens vs bf16 and Triton (cosine>0.999, relL2<0.02). Validated e2e: real 1T Kimi-K2.7-Code serves at TP8 with Marlin, coherent output.
…DA-graph capture Fold kv_b_proj into Q/O (w_kc/w_vc) + fuse q_a_proj/kv_a_proj_with_mqa, run MQA over a compressed-latent paged cache (~57× KV shrink). New MlaAbsorbCacheManager with a FlashInfer MLA kernel fast path (real dims/sm90) gated by a probe, SDPA fallback elsewhere; FlashInferMLAWrapper + _create_persistent_wrappers branch make absorbed decode CUDA-graph-capturable. Default via mla_absorb=True (reduced() pins naive as the parity reference). Composes with the Marlin MoE post-load walker. Validated e2e: real 1T serves at TP8 (absorbed kernel + Marlin + capture, coherent output); logits match vLLM up to bf16 near-ties.
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
1 participant
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 does this PR do?
New Kimi-K2.7 model package + registry/config wiring: DeepSeek-V3-
style text backbone, MLA attention, DeepSeek YARN RoPE, fine-grained
MoE router/shared expert, decoder/LM/causal-LM modules, HF weight
loader, reduced test config, and serving configs.
Serving + TP correctness for reduced/synthetic Kimi-K2.7: config-
driven CUDA graph prefill capture buckets, local checkpoint/tokenizer
hooks, mstar serve E2E path, MLA head sharding, MoE intermediate
sharding, and tp=1/tp=2 parity coverage.
Kimi-K2.7-Code INT4 W4A16 path: compressed-tensors INT4 parsing,
dequant-on-load support, Triton W4A16 fused MoE kernel, real-
checkpoint wiring, TP8 configs, packed expert loading, and worker/
distributed robustness updates for large-checkpoint serving.
Marlin routed-expert MoE kernel: vendored/JIT-built Marlin CUDA path
under mstar/utils/marlin, reusable quantization method plumbing, Kimi
MoE post-load repack, quant_kernel=auto|marlin|triton, Marlin default
on sm80+, Triton fallback elsewhere.
Weight-absorbed MLA default path: folds kv_b_proj into Q/O, fuses
q_a_proj + kv_a_proj_with_mqa, stores compressed latent KV cache, adds
MlaAbsorbCacheManager, FlashInfer MLA fast path with SDPA fallback,
and CUDA-graph-capturable absorbed decode. Composes with the Marlin
MoE post-load path.
Todos:
Include support for MoonViT and other optimizations
How was it tested?
Reduced Kimi-K2.7 modular + GPU integration goldens: components, MoE,
MLA, forward pass, paged attention, weight loading, submodule
lifecycle, and serving E2E.
Serving: live mstar-serve path verified through api server ->
conductor -> worker -> KV cache engine -> decode loop with
deterministic reduced-checkpoint behavior.
INT4/W4A16: CPU quant/wiring tests plus GPU quant/kernel goldens for
packed expert loading and in-kernel dequant behavior.
Marlin: kernel and block goldens compared against bf16 and Triton
paths, with cosine/relL2 checks; real 1T Kimi-K2.7-Code served at TP8
with Marlin and coherent output.
MLA absorb: forward/kernel/paged/serve tests added, plus Marlin merge
compatibility coverage; real 1T TP8 serving validated with absorbed
MLA + FlashInfer kernel + Marlin + CUDA graph capture, with logits
matching vLLM up to bf16 near-ties.
Checklist
ruff check .passes