[Model Integration] Adding Kimi K2.7 Code - #192
Conversation
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
Looks good to me! A few (very) small changes/updates and it will be ready to merge:
-
Rebase onto current main. The mergeable_state is dirty, but it's one trivial conflict i.e. KVCacheConfig gained flashinfer_backend on main (#184) while this PR adds softmax_scale/mla_ckv_dim (keep both fields). Everything was verified on a local merge preview with exactly that resolution.
-
Label the 1T claims' provenance in the description: "real 1T Kimi-K2.7-Code served at TP8 ... logits matching vLLM up to bf16 near-ties" is author-side evidence reviewers can't cheaply reproduce (8 GPUs + ~600GB)
-
The ruff checkbox is unticked so please run it
Also, @NSagan271 please do a pass and let Garv know if you have any suggestions.
| init_method=init_method, | ||
| world_size=self.num_workers, | ||
| rank=self.global_rank, | ||
| timeout=timedelta(hours=2), |
There was a problem hiding this comment.
Non-critical but please address: this 2h timeout (here and in new_group below) applies to EVERY multi-GPU model, not just 1T Kimi loads. So a genuinely hung collective now sits 2 hours before NCCL aborts instead of the old default. I suggest an env var or config field defaulting to the previous behavior, with the Kimi TP8 configs opting into 2h.
583a16d to
9df9253
Compare
NSagan271
left a comment
NSagan271
left a comment
There was a problem hiding this comment.
Overall, clean model port. My biggest comment is about abstracting out quantization configuration / data (e.g., scales) associated with quantized parameters sooner rather than later.
| runs at capture time only and is not replayed. | ||
| """ | ||
| wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | None = None | ||
| wrapper: FlashInferPrefillWrapper | FlashInferDecodeWrapper | FlashInferMLAWrapper | None = None |
There was a problem hiding this comment.
The number of wrapper types is growing (I also added one in #203, though it doesn't get included here). Maybe we want a base FlashInferWrapper class (can just have a dummy init, if there's nothing actually shared between the wrappers; otherwise, includes any base shared code) instead of doing a type union here.
| }) | ||
| q_start += sl | ||
|
|
||
| ps.mla = { |
There was a problem hiding this comment.
Nit (readability / type checking): consider using a dataclass for mla and requests (just an opinion though---don't need to implement)
| cfg = self.kv_cache_config | ||
|
|
||
| # Only the FlashInfer MLA path is capturable; absorbed SDPA stays eager. | ||
| use_mla_kernel = ( |
There was a problem hiding this comment.
I might be mistaken, but I think if cfg.attention_backend == "mla_absorb" but _mla_kernel_available returns False, e.g., because of sm_major != 9:, this function falls through to instantiating either a prefill or decode wrapper, which I believe is incorrect behavior.
I don't see anything (either here or in Kimi's submodules.py) enforcing that cuda graphs not be captured when _mla_kernel_available returns False, which the "absorbed SDPA stays eager" comment implies. If my assessment is correct, then the fix is to cancel cuda graph capture for a cell where cfg.attention_backend == "mla_absorb but the MLA kernel is not available.
There was a problem hiding this comment.
Yes @garv901 Naomi is right.
I forced mla_absorb=True on reduced dims (temporary worktree edit, reverted). The server boots, every capture fails with a logged warning, and then every request 500s with AttributeError: 'FlashInferDecodeWrapper' object has no attribute 'set_latent' because the wrong wrapper leaks into the plan path. Any absorb deployment on non-sm90 hardware would be broken.
There was a problem hiding this comment.
Yeah I saw that too, sorry for the oversight on this. I am working on the changes and will try to push the changes by tomorrow
| page_size, num_kv_heads, head_dim, | ||
| dtype=kv_cache_type, device=device, | ||
| ).contiguous() | ||
| if cfg.attention_backend == "mla_absorb": |
There was a problem hiding this comment.
Nit/future (unsure): maybe this logic (specifically, computing the kv cache dimensions, not instantiating the tensor) should be pushed to the KV cache config?
| ) -> torch.Tensor: | ||
| if self.mla_absorb: | ||
| return self._forward_absorbed(hidden_states, cache_handle, position_ids) | ||
| num_tokens = hidden_states.shape[0] |
There was a problem hiding this comment.
Nit: worth the comment that the code below is a fallback for reduced-capability testing (like the comment in config.py)
| checkpoint_path = kwargs.get("checkpoint_path") | ||
| self.model_path_hf = checkpoint_path or model_path_hf | ||
| self._config_variant = kwargs.get("config_variant", "full") | ||
| if self._config_variant == "reduced": |
There was a problem hiding this comment.
Other models override config keys with those on huggingface; I think it's worth it to do so here to or leave a comment in the code.
| topk_ids: torch.Tensor, | ||
| activation: str = "silu", | ||
| reduce_results: bool = True, | ||
| w1_scale: torch.Tensor | None = None, |
There was a problem hiding this comment.
This introduces several different arguments related to different quantization methods. It's manageable/readable for just W4A16 and the Kimi variant, but we will conceivably want to add many more in the future (INT8, W4A4, FP8, etc.).
I think this is the right time to abstract out some of the quantization-related data. e.g. adding a QuantizationType Enum instead of having the logic quantized = w1_scale is not None; having, e.g.,
class QuantizationData(ABC):
@property
@abstractmethod
def quant_type(self) -> QuantizationType:
pass
@dataclass
class W4A16Data(QuantizationData):
w1_scale: torch.Tensor
w2_scale: torch.Tensor
pack_factor: int
group_size: int = 16
@property
def quant_type(self) -> QuantizationType:
return QuantizationType.W4A16
...
for instance, potentially with a version of the class CompressedTensorsQuantConfig class pulled out as a common component.
|
Overall this looks good to me and ready to merge. Just address @NSagan271 comment on mstar/engine/cuda_graph_runner.py line 336 above and also a couple of minor points (none of them urgent):
Once you are done with the cuda graph runner comment, it should be ready to merge (lmk). |
Addresses PR #192 review; configs de-hardcoded and synthetic ones moved. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
5dfca30 to
7f3bbf1
Compare
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.
Reduce comment density across the Kimi-K2.7 port and its supporting engine/utils changes to match the surrounding mstar style. Explanatory prose that duplicated the code, reference-file line citations, and narrative docstrings are condensed to short summaries; the non-obvious rationale (meta/to_empty buffer handling, MLA softmax scale, cache layout) is kept as one- or two-line notes. Comments only -- no behavioral change. Kimi CPU tests (27) unchanged and ruff reports no new findings.
The weight-absorbed MLA work added a third ATTENTION_BACKENDS entry
("mla_absorb") but left test_registry_names asserting the exact
pre-existing set, so the test has been failing since that change.
Add the new backend to the expected set.
The 2h NCCL timeout applied to every multi-GPU model, so a genuinely hung collective sat 2h before aborting. Unset now means PyTorch's default; the Kimi TP8 configs opt in via dist_timeout_s, overridable per process with MSTAR_DIST_TIMEOUT_S.
7f3bbf1 to
58ede84
Compare
3 participants
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