Shard Kimi projections without padding checkpoint weights

[voipmonitor]

Status

Implemented. This pull request is stacked on #384 and requires local-inference-lab/b12x#215 for unaligned multi-row FP8 output storage.

Behavior

• Preserves the official Kimi-K3 MXFP4 routed-expert shard width instead of padding each TP16 shard from 192 to 256 channels.
• Tensor-parallel shards the merged MLA latent projection and restores logical row order after gather.
• Gathers BF16, FP16, FP32, and FP8 projection payloads through B12X while removing transport-only alignment rows.
• Supports opt-in KDA f_a sharding when the output width is divisible by the tensor-parallel size.
• Declares checkpoint-omitted KDA alignment rows and initializes them to zero before online quantization.

Technical reason

Checkpoint padding increases routed-expert storage and can prevent the official MXFP4 checkpoint from fitting with the required KV cache. Projection sharding reduces replicated dense weight memory without changing logical tensor shapes or executed values.

Compatibility

KDA f_a sharding remains disabled unless requested. TP12 rejects that option because width 128 is not divisible by 12. The replicated path remains the default because TP16 decode measurements favor replication despite its 113.20 MiB per-GPU memory cost.

Validation

• Kimi model, TP-projection, sequence-parallel, KDA, omitted-tail, and native-MoE-shard tests pass.
• The official checkpoint keeps a 192-channel local routed-expert shard at TP16.
• The merged latent projection saves approximately 0.634 GiB per GPU at TP16.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• dev/*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 10f9e7e0-9a15-4ea8-8c8c-4598f2c7f1ef

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}