Add Kimi-K2-Instruct-0905 contrib model (1T MoE on trn2.48xlarge)

[jimburtoft]

Summary

• Adds a full NxDI contrib implementation of Kimi-K2-Instruct-0905 (moonshotai/Kimi-K2-Instruct-0905), a 1 trillion parameter MoE model based on the DeepSeek-V3 architecture
• Validated on trn2.48xlarge with blockwise FP8 quantization
• Recommended config: LNC=2, TP=32, EP=2 (64 cores) — 76% faster than LNC=1/TP=64
• All 4 integration tests pass: smoke, generation, coherence, and TPOT performance

Model Details

Parameter	Value
Total parameters	~1,000B
Active parameters/token	~32B
Architecture	DeepSeek-V3 variant with MLA attention
Experts	384 routed (8 active) + 1 shared per MoE layer
Quantization	Blockwise FP8 (e4m3, 128x128 blocks)
Attention	Multi-Latent Attention (MLA) with KV LoRA rank=512

Benchmark Results (BS=1, seq_len=1024)

LNC=2 vs LNC=1 Comparison

Config	TP	EP	Cores	TPOT	tok/s	Speedup
LNC=2 (recommended)	32	2	64	165.5 ms	6.0	+76%
LNC=1	64	2	128	297 ms	3.4	baseline

LNC=2 gives 2x HBM bandwidth per logical core. Since MoE decode is purely bandwidth-bound (192 expert weight loads per step), this translates directly to throughput improvement.

Full Benchmark Sweep (LNC=2)

Output Tokens	TTFT P50 (ms)	TPOT P50 (ms)	tok/s	E2E P50 (ms)
16	1,420.4	166.38	6.0	3,916.1
32	1,419.8	165.58	6.0	6,553.0
64	1,419.7	165.56	6.0	11,849.8
128	1,419.8	165.48	6.0	22,435.8
256	1,419.9	165.42	6.0	43,604.1
512	1,420.0	165.47	6.0	85,974.4

Files Added

contrib/models/Kimi-K2-Instruct-0905/
├── README.md                          # Full model info, architecture, benchmarks, usage
├── src/
│   ├── __init__.py                    # Exports NeuronKimiK2ForCausalLM, KimiK2InferenceConfig
│   └── modeling_kimi_k2.py            # Complete model implementation (1548 lines)
└── test/
    ├── __init__.py
    ├── integration/
    │   ├── __init__.py
    │   └── test_model.py              # 4 integration tests
    └── unit/
        └── __init__.py

Key Implementation Notes

• Streaming checkpoint loader: Processes 62 safetensor shards one at a time to avoid OOM on 2TB host RAM
• Custom FP8 handling: Blockwise symmetric quantization with 128x128 blocks, non-expert modules dequanted to BF16
• MLA weight absorption: Avoids KV decompression during decode for efficient inference
• CPU greedy sampling (no_on_device_sampling): Required due to vocabulary size (163840) exceeding on-device sampling limits
• Router with bias: Uses RouterTopK(bias=True) to load e_score_correction_bias as linear_router.bias
• Selective loading threshold: Requires patching to 0.0 in neuronx_distributed/modules/moe/model_utils.py

Known Limitations

• EOS logit elevation: The <|im_end|> token (163586) has an elevated logit in context encoding outputs. Mitigated with min_tokens_before_eos parameter. Documented in README.
• Batch size: Higher batch sizes provide no throughput improvement due to MoE bandwidth bottleneck (TPOT scales linearly with BS)
• Instance requirement: Requires trn2.48xlarge — model is too large for smaller instances

Testing

# On trn2.48xlarge with NEURON_LOGICAL_NC_CONFIG=2:
NEURON_LOGICAL_NC_CONFIG=2 LOCAL_WORLD_SIZE=64 pytest test_model.py -v --capture=tee-sys

[FThompsonAWS]

Try setting tkg_batch_size to the same value as max_batch_size to compile TKG with the correct shapes.