model: Kimi-K3 (kimi_linear + kimi_k3) decode/prefill + K3 MXFP4 kernels + tests - #315
Closed
Andrewxu313 wants to merge 3 commits into
Closed
Andrewxu313 wants to merge 3 commits into
Andrewxu313 wants to merge 3 commits into
Conversation
…els + tests
Model-layer slice on top of the core PR (PR_MERGE_POLICY: models/** + batchgen_kernels/** + tests/**):
- models/moonshotai/kimi_linear/** + kimi_k3/**: whole-model decode graph, MoE graph segments,
DeepEP exchange wiring + planner (--k3-moe-exchange honoring), PSM prewarm, Triton kernels
(situ, row_gather, rmsnorm int64, kda_conv, attn_residual), vocab_parallel, distributed_weight_store,
tp_weight_sharding, moe_tp_reshard, k3/mxfp4_{expert,layout}, tensor_map, tokenizer.
- batchgen_kernels/**: grouped mxfp4/int4 wgmma, marlin_grouped_gemm, kda_fused_decode, qkv_wgmma,
dispatch_scatter_3d, routing gate_sigmoid_topk, JIT registry.
- tests/**: kimi_linear/kimi_k3 unit + gpu + worker + moe (2 stale tests removed).
Excludes GLM-5 (models/glm/glm5/*) — separate effort.
…top-level tests/ (matches vLLM/SGLang; drops from wheel; clears §1.2)
…/ with K3 content (remove in-package duplicates; clears §1.2)
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.
Description
Kimi-K3 model + kernels + tests (
models/moonshotai/kimi_linear+kimi_k3,batchgen_kernels/**,tests/**). Stacked on #314 (k3-core-pr) — review/merge that first.Motivation
Model-layer half of K3 day-0 support: whole-model decode CUDA graph, MoE graph segments, DeepEP-vs-NCCL exchange wiring + planner, PSM prewarm, Triton kernels (situ, row_gather, rmsnorm int64, kda_conv, attn_residual), vocab-parallel embed/lm_head, distributed-weight-store client, K3 MXFP4 layout + tensor map, tokenizer; K3 MXFP4/INT4 WGMMA + marlin grouped GEMM + kda_fused_decode kernels. GLM-5 excluded (rode this branch; goes via the GLM line — the decode-MoE AllReduce change is already on
release/v1.0.10-glm5-cudagraph).Type of Change
model— add/extend model support (models/**+ registration seam + model kernels + tests)File changes
Full list vs the core PR base (
k3-core-pr):Kernel tests were moved out of the runtime package into top-level
tests/kimi_linear/(matches vLLM/SGLang; not shipped in the wheel).Checklist
tests/(runtime tests to run on the H200 pods — dev-machine is CPU-only).MODULE.mdupdated; README News/Roadmap is a separatedocsPR.PR Merge Policy Contract — pre-merge checklist
model); all files within the model allowlist — no scaffolding edits (§2.5/2.6).debug_*/scratch_*/tmp_*scripts (§1.1).kimi_linearkernel tests relocated frombatchgen_kernels/tests/to top-leveltests/kimi_linear/; notest_*.pyinside the runtime package.BATCHGEN_KIMI_ATTENTION_GROUP_SIZE(kimi_initializer.py) +BATCHGEN_SKIP_KV_CALLBACK(wrappers.py) kept for the interim release; documented inbatchgen_design/model_support/kimi_k3/ENV_KNOBS.md(→ server-arg /batchgen_debugnext pass). CI hygiene is report-only.print()/ commented-out code in changed non-test files (§1.4–1.5).MODULE.mdupdated (§2.2).