Day 0 Kimi-K3 Inference Deployment with ATOM on AMD Instinct MI355X GPUs

Jul 27, 2026

Image showing businessman hand working with modern laptop computer in modern office with virtual icon diagram

Kimi-K3 scales the model to 2.8 trillion parameters while introducing Kimi Delta Attention (KDA), Gated Multi-head Latent Attention (MLA), Stable LatentMoE, and Attention Residuals (AttnRes). These architectural choices reduce cache overhead for long-context inference, but the approximately 1.56 TB checkpoint still sets a high HBM capacity bar for single-instance Day 0 deployment.

This post is intended for users who want to evaluate Kimi-K3 as soon as it becomes available. Starting from the model architecture, weight placement, and hardware capacity, it presents a single-instance deployment configuration that we have validated:

8 AMD Instinct™ MI355X GPUs with TP8.

The goal is not to pursue peak performance, but to answer three practical Day 0 questions: why the Kimi-K3 weights fit on these GPUs, how the weights are distributed under TP8, and how to bring up the model quickly with ATOM and run a minimal correctness check. Kernel tuning, communication optimization, long-context acceleration, and end-to-end performance analysis will be covered in separate follow-up posts.

Kimi-K3 Architecture and Weights

Model Architecture

Kimi-K3 is a native multimodal Mixture-of-Experts (MoE) model. Its text backbone contains 93 layers: 69 KDA layers and 24 Gated MLA layers. The layers follow an interleaved KDA x 3 → MLA × 1 pattern, with one additional MLA layer at the end.

Item	Configuration
Total parameters	Approximately 2.78T
Active parameters per token	Approximately 103B
Decoder layers	93
KDA / Gated MLA layers	69 / 24
Dense / MoE layers	1 / 92
Routed experts	896, with top-16 selected per token
Shared experts	2
Hidden / latent hidden size	7168 / 3584
Expert intermediate size	3072
Attention heads	96
Maximum context length	1M tokens
Vision components	MoonViT3d + multimodal projector

In the attention stack, KDA handles most temporal modeling with a fixed-size recurrent state, avoiding a sequence-length-dependent KV cache in every layer. MLA is inserted after every three KDA layers to provide explicit global retrieval, so only the 24 MLA layers maintain per-token latent KV states. AttnRes spans standard residual blocks and learns a weighted combination of historical representations.

In the MoE stack, Stable LatentMoE first projects the 7168-dimensional hidden state down to 3584 dimensions before running the expert computation. Each token activates only 16 of the 896 routed experts. Because different tokens can select different experts, however, the complete expert weights must remain resident across the GPU parallel domain.

HBM Usage Under TP8

ATOM applies the following tensor-parallel placement rules:

Attention heads are sharded across TP ranks.
The large gate/up/down matrices in the Dense MLP and Shared Expert, as well as the w1/w2/w3 matrices in the Routed Expert, use row or column parallelism.
Every rank retains all 896 expert IDs; TP shards each expert's matrices rather than partitioning the expert IDs.
MLA q_a/kv_a, KDA f_a, LatentMoE down/up, Norm, router, and AttnRes score projections are replicated across ranks.
The token embedding and LM head are sharded along the vocabulary dimension. The current text-only service does not load vision_tower or mm_projector.

Weight Distribution

The following table is derived from the safetensors headers of 497,220 tensors in the checkpoint and ATOM's tensor-by-tensor placement rules for ColumnParallelLinear, RowParallelLinear, ReplicatedLinear, VocabParallelEmbedding, and skipped prefixes. The result is not a simple checkpoint-size-divided-by-TP estimate. “Full-checkpoint physical storage” is the size of the complete checkpoint, while the TP8 column shows the original weight data actually loaded on each GPU.

Weight category	Full-checkpoint physical storage	TP8 per GPU	Placement
Dense MLP	1.453 GB	0.182 GB	`gate/up` column parallel; down row parallel
Routed Expert packed values + scales	1446.456 GB	180.807 GB	All expert IDs retained; each expert matrix TP-sharded
Shared Expert	24.310 GB	3.039 GB	`gate/up` column parallel; down row parallel
KDA Attention GEMM	61.214 GB	7.763 GB	Head projections sharded;`f_a` replicated
MLA Attention GEMM	11.145 GB	2.029 GB	Head projections sharded;`q_a/kv_a` replicated
Norm	0.006 GB	0.006 GB	Language-model Norms replicated; Vision/MM Norms skipped
Other weights	16.275 GB	11.231 GB	Embedding sharded; LatentMoE/router replicated; vision weights skipped
Loader-aware total	1560.860 GB	205.057 GB	Includes replicated weights; excludes unloaded vision weights

KV Cache and Runtime-State Distribution

Kimi-K3 supports context lengths of up to 1M tokens. In addition to the model weights, long-context inference requires extra HBM for MLA latent KV states, KDA SSM and convolution states, and the AttnRes prefill chunk. The table below estimates this additional memory for one 1M-token sequence, an FP8 latent KV cache, FP32 KDA SSM states, BF16 KDA convolution states, and an 8K-token AttnRes prefill chunk.

Runtime state	Size calculation	Layers	TP8 per GPU	Scaling behavior
MLA Latent KV	1,048,576 × (512 + 64) × 2 bytes/layer	24	14.496 GB	BF16; grows linearly with token count; each rank stores the latent cache
KDA SSM state	(96 ÷ TP) × 128 × 128 × 2 bytes/layer	69	0.054 GB	Fixed per sequence; sharded by head
KDA Conv state	3 × (12288 ÷ TP) × (4 - 1) × 2 bytes/layer	69	0.002 GB	Fixed per sequence; sharded by channel
AttnRes 8K-chunk temporary state	8192 × ⌈93 ÷ 12⌉ × 7168 × 2 bytes	93	0.940 GB	Stores one block residual every 12 layers; scales with the prefill chunk size
Known additional memory beyond weights	—	—	Approximately 15.491 GB	Excludes kernel workspaces and related overhead

Combined HBM Usage

The weight and runtime-state tables above use decimal GB (1 GB = 10^9 bytes). To compare them consistently with the GPU capacities specified in GiB (1 GiB = 2^30 bytes), the values below are converted to GiB.

Deployment	Model weights	1M-token context	Known total	HBM capacity	HBM utilization	HBM remaining
MI355X TP8	190.974 GiB	14.427 GiB	205.401 GiB	288 GiB	71.3%	82.599 GiB

These estimates do not include additional overhead from Grouped GEMM, KDA and attention workspaces, communication buffers, allocator fragmentation, CUDA Graphs, reordered weight copies, or the framework runtime. The deployment provides substantial per-GPU headroom for these unmodeled costs.

AMD Instinct MI355X GPUs

AMD Instinct MI355X GPUs can be used to deploy Kimi-K3. MI355X provides 288 GiB of HBM3E and supports MXFP4, enabling TP8 deployment on an 8-GPU UBB 2.0 platform. Key specifications are summarized below.

Item	AMD Instinct MI355X
GPU architecture	AMD CDNA 4
HBM per GPU	288 GiB HBM3E
Peak memory bandwidth per GPU	8 TB/s
Low-precision compute	10.1 PFLOPS MXFP4
Platform configuration	UBB 2.0 Platform (8 GPUs)

With TP8, this configuration provides sufficient capacity for the model weights and additional HBM headroom for runtime state. Its suitability for single-instance Kimi-K3 deployment and functional validation will depend on hardware availability and communication scale.

Day 0 Kimi-K3 Deployment with ATOM

ATOM provides a Kimi-K3 deployment recipe for AMD Instinct MI355X GPUs, covering the container image, runtime configuration, model service startup.

MI355X: 8 GPUs with TP8

Use the MI355X container image:

		docker pull rocm/atom-dev:rocm7.2.4_ubuntu24.04_py3.12_pytorch2.10.0_20260727_kimi_k3

Inside the container, configure the runtime and start the ATOM OpenAI-compatible server:

		export ATOM_LOADER_USE_THREADPOOL=1
export ATOM_LOADER_THREADPOOL_WORKERS=16
export ATOM_SYNC_AFTER_LOAD=1
export ATOM_DIST_TIMEOUT_SECONDS=3600
export ATOM_USE_TRITON_GEMM=1
export AITER_USE_GROUPED_GEMM=0
export ATOM_USE_TRITON_MOE=0
export AITER_FLYDSL_FORCE=1
export AITER_FORCE_GFX1250=0
export ATOM_USE_UNIFIED_ATTN=1
export ATOM_FORCE_ATTN_TRITON=1

python -m atom.entrypoints.openai_server \
  --model moonshotai/Kimi-K3 \
  --kv_cache_dtype fp8 -tp 8 \
  --trust-remote-code \
  --max-model-len 16384 \
  --max-num-seqs 64 \
  --max-num-batched-tokens 10240 \
  --gpu-memory-utilization 0.93 \
  --block-size 128 \
  --no-enable_prefix_caching

Conclusion

Kimi-K3's 2.78 trillion parameters occupy approximately 1.5609 TB in their mixed MXFP4 and BF16 checkpoint representation. Under ATOM's tensor-by-tensor loader placement rules, MI355X TP8 loads approximately 190.974 GiB of weights per GPU, increasing to approximately 205.401 GiB after adding the known states for one 1M-token sequence with an FP8 latent KV cache.

The configuration can load the model weights and retains approximately 82.6 GiB per GPU after adding the estimated states for a 1M-token context, before accounting for the remaining runtime overhead. For Day 0 evaluation, ATOM has completed full 1,319-sample GSM8K 5-shot validation on MI355X TP8 with a 16K maximum model length, providing a practical path to load the model, start an API service, and verify basic correctness.

This post does not make claims about throughput, time to first token (TTFT), time per output token (TPOT), or kernel efficiency. HBM optimization, TP8 collectives, MXFP4 Grouped MoE, KDA prefill and decode, MLA, and 1M-token context optimization will be addressed in dedicated performance-tuning posts.