model: Kimi-K3 (kimi_linear + kimi_k3) decode/prefill + K3 MXFP4 kernels + tests by Andrewxu313 · Pull Request #317 · batchgen-project/batchgen

Andrewxu313 · 2026-09-11T12:19:07Z

Description

Kimi-K3 model support — models/moonshotai/kimi_linear + kimi_k3, the K3 MXFP4 kernels, and tests:

kimi_linear + kimi_k3 model, planner, PSM prewarm, whole-model decode CUDA graph
Triton kernels (situ, row-gather, rmsnorm int64, kda-conv, attn-residual)
vocab-parallel embed/lm_head, TP weight sharding, distributed-weight store config
K3 MXFP4 Marlin / WGMMA grouped-GEMM kernels + sigmoid-topk routing, KDA fused decode
unit / staged-gpu / kernel-parity tests under tests/

Stacked on the core PR #316.

Motivation

Model, kernels, and tests for Kimi-K3, split from core per PR_MERGE_POLICY §2.5 (model = models/** + registration seam + batchgen_kernels/** + tests/**). Builds on the #316 scaffolding.

Type of Change

model — add/extend model support (models/** + registration seam + model kernels only)
kernel — add/optimize a compute kernel (batchgen_kernels/** + in-tree kernel dirs)
core — change scheduling/serving/runtime scaffolding (the only type that may)
fix — narrow bug fix (+ a regression test)
infra — build / CI / packaging / scripts / Docker
docs — documentation only

File changes

File	Δ	Note
`batchgen/models/moonshotai/kimi_k3/MODULE.md`	mod	K3 model
`batchgen/models/moonshotai/kimi_k3/model.py`	mod	K3 model
`batchgen/models/moonshotai/kimi_k3/tokenizer.py`	mod	K3 model
`batchgen/models/moonshotai/kimi_linear/Parallel_Strategy_Manager.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/attn_residual_triton.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/block_residual.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/cuda_graph_segments.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/distributed_weight_store.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/k3/mxfp4_expert.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/k3/mxfp4_layout.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/k3/tensor_map.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/kda_conv_triton.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/kimi_initializer.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/kimi_parameter_server.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/model.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/moe_cuda_graph_segments.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/moe_tp_reshard.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/planner.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/rmsnorm_triton.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/row_gather_triton.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/serving_modules.py`	mod	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/situ_triton.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/tp_weight_sharding.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/vocab_parallel.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/whole_model_cuda_graph_segments.py`	add	kimi_linear model
`batchgen/models/moonshotai/kimi_linear/wrappers.py`	mod	kimi_linear model
`batchgen_kernels/__init__.py`	mod	K3 kernel
`batchgen_kernels/_jit_registry.py`	mod	K3 kernel
`batchgen_kernels/_version.py`	mod	K3 kernel
`batchgen_kernels/attention/dsa/fast_topk_cuda.py`	mod	K3 kernel
`batchgen_kernels/attention/kda_fused_decode.py`	add	K3 kernel
`batchgen_kernels/setup.py`	mod	K3 kernel
`batchgen_kernels/src/attention/kda_fused_decode.cu`	add	K3 kernel
`batchgen_kernels/src/attention/qkv_wgmma.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/dispatch_scatter_3d.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/expert_mxfp4_wgmma.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/fused_int4_wgmma_grouped.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/grouped_int4_wgmma.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/grouped_int4_wgmma_ext.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/grouped_mxfp4_wgmma.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/marlin_grouped_gemm.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/routing/gate_sigmoid_topk.cu`	mod	K3 kernel
`batchgen_kernels/src/moe/routing/routing_extension.cc`	mod	K3 kernel
`batchgen_kernels/src/moe/routing/routing_ops.h`	mod	K3 kernel
`batchgen_kernels/src/moe/single_expert_int4_wgmma.cu`	mod	K3 kernel
`tests/gpu/kimi_linear_tp_weight_parity.py`	add	test
`tests/gpu/mxfp4_resident_ep_world2_parity.py`	add	test
`tests/gpu/run_kimi_k3_kda_gpu.sh`	add	test
`tests/gpu/streamed_sp8_mxfp4_world8_parity.py`	add	test
`tests/gpu/test_gate_sigmoid_topk_k16.py`	add	test
`tests/gpu/test_kimi_k3_attn_residual_triton.py`	add	test
`tests/gpu/test_kimi_k3_dequant_once_prefill.py`	add	test
`tests/gpu/test_kimi_k3_kda_batched_segment_parity.py`	add	test
`tests/gpu/test_kimi_k3_kda_conv_triton.py`	add	test
`tests/gpu/test_kimi_k3_kda_fla_parity.py`	mod	test
`tests/gpu/test_kimi_k3_kda_fused_decode_parity.py`	add	test
`tests/gpu/test_kimi_k3_kda_head_parallel_parity.py`	add	test
`tests/gpu/test_kimi_linear_latent_moe_serving.py`	mod	test
`tests/gpu/test_kimi_linear_mxfp4_latent_moe_serving.py`	add	test
`tests/gpu/verify_k3_mxfp4_expert.py`	mod	test
`tests/integration/paged_kv/test_host_kv_copy_stream_ordering.py`	add	test
`tests/kimi_k3_harness.py`	mod	test
`tests/kimi_k3_oracle_assets/__init__.py`	mod	test
`tests/kimi_linear/test_block_residual_segment_capture.py`	add	test
`tests/kimi_linear/test_conv1d_layout_cpu.py`	mod	test
`tests/kimi_linear/test_conv1d_std.py`	mod	test
`tests/kimi_linear/test_decode_graph_adapter.py`	mod	test
`tests/kimi_linear/test_fused_moe_std.py`	mod	test
`tests/kimi_linear/test_kda_manager_graphready.py`	mod	test
`tests/kimi_linear/test_kda_segment_capture.py`	mod	test
`tests/moe/gpu_parity_mxfp4_marlin.py`	mod	test
`tests/moe/test_mxfp4_marlin_repack.py`	mod	test
`tests/test_batch_scheduler_pool_failure.py`	add	test
`tests/test_batchgen_model_config.py`	mod	test
`tests/test_core_signal_handler_source.py`	add	test
`tests/test_decode_host_kv_stream_order.py`	add	test
`tests/test_dsa_valid_token_kernels.py`	mod	test
`tests/test_glm5_planner.py`	mod	test
`tests/test_gpu_page_table_stability.py`	mod	test
`tests/test_grouped_fp8_blockwise_moe.py`	del	test
`tests/test_jit_extension_name_uniqueness.py`	add	test
`tests/test_kimi_k3_block_residual_prealloc.py`	mod	test
`tests/test_kimi_k3_deepep_ll_graph.py`	add	test
`tests/test_kimi_k3_distributed_weight_config.py`	add	test
`tests/test_kimi_k3_empty_decode_batch_view.py`	add	test
`tests/test_kimi_k3_empty_rank_whole_graph.py`	add	test
`tests/test_kimi_k3_gpu_kv_compaction.py`	add	test
`tests/test_kimi_k3_kda_scratch_reservation.py`	add	test
`tests/test_kimi_k3_kda_segmented.py`	mod	test
`tests/test_kimi_k3_kda_slot_planner.py`	add	test
`tests/test_kimi_k3_model.py`	mod	test
`tests/test_kimi_k3_prefill_offload_tracking.py`	add	test
`tests/test_kimi_k3_prefill_profile_scope.py`	add	test
`tests/test_kimi_k3_resident_prefill_control.py`	add	test
`tests/test_kimi_k3_resident_prefill_handoff.py`	add	test
`tests/test_kimi_k3_row_gather_triton.py`	add	test
`tests/test_kimi_k3_server_startup_init.py`	add	test
`tests/test_kimi_k3_situ_triton.py`	add	test
`tests/test_kimi_k3_tensor_map.py`	mod	test
`tests/test_kimi_k3_tokenizer.py`	mod	test
`tests/test_kimi_k3_vocab_parallel.py`	add	test
`tests/test_kimi_linear_block_residual_serving.py`	mod	test
`tests/test_kimi_linear_ffn_chunk.py`	mod	test
`tests/test_kimi_linear_flashmla_metadata_contract.py`	add	test
`tests/test_kimi_linear_kda_projection_fusion.py`	add	test
`tests/test_kimi_linear_m2a_alog_head_shard.py`	add	test
`tests/test_kimi_linear_m2b_boundary_validator.py`	add	test
`tests/test_kimi_linear_m2b_decode_dp_group.py`	add	test
`tests/test_kimi_linear_m2b_moe_scatter_gather.py`	add	test
`tests/test_kimi_linear_moe_cuda_graph_contract.py`	add	test
`tests/test_kimi_linear_rmsnorm_chunk.py`	add	test
`tests/test_kimi_linear_shared_expert_tp.py`	add	test
`tests/test_pool_output_parsing.py`	add	test
`tests/test_resident_ep_latent_tp_layout.py`	add	test
`tests/test_server_args.py`	add	test
`tests/test_server_worker_readiness.py`	add	test
`tests/test_worker_fatal_propagation.py`	add	test
`tests/test_worker_host_kv_append_completion.py`	add	test
`tests/unit/test_query_book_pool_grow_rebind.py`	del	test
`tests/worker/test_boundary.py`	mod	test
`tests/worker/test_decode.py`	mod	test
`tests/worker/test_kv_manager.py`	mod	test
`tests/worker/test_prefill.py`	mod	test

Checklist

I have read the CONTRIBUTING guide and the PR Merge Policy Contract.
I have updated the tests (if applicable).
I have updated the documentation (if applicable).

…eepEP-LL exchange, resident MoE, --enable-deepep) Core-only slice of the K3 support branch (PR_MERGE_POLICY: scaffolding layer): - core/ C++: Weights_Storage + distributed_weight_daemon (compact per-node host store), HtoD_Engine, GPU_Weight_Buffer, KV_Storage host-paged views, batchgen binding. - batchgen/moe/: deepep_ll.py (DeepEP low-latency EP exchange), marlin_grouped_moe, streamed_sp8_mxfp4, k3_prefill_dequant_once, fused_moe_mxfp4_resident, routing/cuda_routing. - worker/ (decode/prefill/kv_manager/boundary), server/ (server_args --enable-deepep, worker_manager, worker_readiness, batch_scheduler), kv_cache/, continuous_batching, batchgen_worker. - triage: planner/base_planner, sequence, server_worker_main_loop, config, decode_dp_group, kernel_compat, op_builder/core_engine, ckpt_converter, install_deps, docs/troubleshooting.

… + tests Model-only slice of K3 support (PR_MERGE_POLICY: models/** + registration seam + batchgen_kernels/** + tests/**), stacked on the k3-core PR. - kimi_linear + kimi_k3 model, planner, PSM prewarm, whole-model decode CUDA graph, Triton kernels (situ, row-gather, rmsnorm int64, kda-conv, attn-residual int32-wrap fix), vocab-parallel embed/lm_head, tp weight sharding, distributed-weight store config. - K3 MXFP4 Marlin / WGMMA grouped-GEMM kernels + sigmoid-topk routing, KDA fused decode. - tests: unit + staged-gpu + kernel parity; kimi_linear tests live under tests/ (not the runtime package). Checkpoint paths in test docstrings/defaults are placeholders. - decode MoE exchange selected by the core --enable-deepep flag (default off = NCCL).

Andrewxu313 added 2 commits September 11, 2026 13:09

Andrewxu313 mentioned this pull request Sep 11, 2026

core: Kimi-K3 decode/prefill scaffolding (distributed weight store, DeepEP-LL exchange, resident MoE) #316

Merged

9 tasks

Base automatically changed from k3-core to main September 11, 2026 13:36

Andrewxu313 marked this pull request as ready for review September 11, 2026 13:37

Andrewxu313 merged commit 35852fc into main Sep 11, 2026
1 check passed

Andrewxu313 deleted the k3-model branch September 11, 2026 13:37

Andrewxu313 mentioned this pull request Sep 11, 2026

docs: document Kimi-K3 support (README News + server flags) #319

Merged

9 tasks

Navigation Menu

Uh oh!

model: Kimi-K3 (kimi_linear + kimi_k3) decode/prefill + K3 MXFP4 kernels + tests - #317

model: Kimi-K3 (kimi_linear + kimi_k3) decode/prefill + K3 MXFP4 kernels + tests#317
Andrewxu313 merged 2 commits into
mainfrom
k3-model

Andrewxu313 commented Sep 11, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

Andrewxu313 commented Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation

Type of Change

File changes

Checklist

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Andrewxu313 commented Sep 11, 2026 •

edited

Loading