Problem Description
ATOM Plugin with SGLang Consistently results in 20+ GB extra being allocated for weights when compared to SGLang operating without the plugin.
This happens for many models so it seems to be framework related to the ATOM plugin:
Why the plugin needs ~20 GB more
The keyword is bpreshuffle (block‑pre‑shuffle). AITER's high‑performance a8w8 block‑scale GEMM and fused‑MoE kernels don't consume weights in the checkpoint's native layout — they require the weight tensors pre‑swizzled and padded to the kernel's MFMA tile boundaries, with their block‑scale factor tensors alongside. So when the ATOM plugin builds each linear/MoE layer, it allocates those reordered + tile‑padded weight buffers (dominated by DeepSeek's 256‑expert MoE, where padding every expert to tile alignment adds up). Vanilla SGLang keeps the plain fp8 weights and uses generic kernels, so it doesn't pay that ~20 GB. It's essentially the memory price of AITER's tuned kernels.
Operating System
22.04.5 LTS (Jammy Jellyfish)
CPU
AMD EPYC 9575F 64-Core Processor
GPU
AMD Instinct MI355X
ROCm Version
7.2.4
ROCm Component
No response
Steps to Reproduce
Reproducing the SGLang (no plugin) vs ATOM‑plugin K3 memory comparison
Prerequisites
8× AMD Instinct MI355X (gfx950), ROCm 7.2.
~1.6 TB disk for the Kimi‑K3 checkpoint, and Hugging Face access to moonshotai/Kimi-K3.
Docker with GPU passthrough.
0. Download the weights once (shared by both runs)
export MODEL_DIR=/path/to/Kimi-K3 # pick any host dir with ~1.6 TB free
huggingface-cli download moonshotai/Kimi-K3 --local-dir "$MODEL_DIR"
Common docker flags used below:
DOCKER_FLAGS="--rm --device=/dev/kfd --device=/dev/dri --group-add video
--shm-size 16G --network host --security-opt seccomp=unconfined
--security-opt apparmor=unconfined --cap-add=SYS_PTRACE
-v $MODEL_DIR:/model_weights"
The four memory‑accounting lines we read are printed by SGLang during startup — you do not need to send any request. Launch, wait until you see Memory pool end, then Ctrl‑C.
A. WITHOUT the ATOM plugin (native SGLang + AITER)
Image: lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727 (SGLang's day‑0 K3 build).
docker pull lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727
docker run $DOCKER_FLAGS --entrypoint bash
lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727 -lc '
export SGLANG_USE_AITER=1 SGLANG_AITER_K3_OPT=1
AITER_FLYDSL_FORCE=1 AITER_SITUV2_A8W4=1
python3 -m sglang.launch_server --model-path /model_weights
--trust-remote-code --tp-size 8
--attention-backend triton --dtype bfloat16
--mem-fraction-static 0.85 --page-size 1
--context-length 16384 --disable-radix-cache
--host 127.0.0.1 --port 30000 2>&1 | tee /tmp/k3_native.log'
B. WITH the ATOM plugin
Image: rocm/atom-dev:sglang-latest (must include ATOM commit b687428 "[ATOM SGL][model] Kimi K3" or newer). The image does not bundle the KDA deps, so install them first (per ATOM's recipe).
docker pull rocm/atom-dev:sglang-latest
docker run $DOCKER_FLAGS --entrypoint bash
rocm/atom-dev:sglang-latest -lc '
pip install "fla-core==0.5.1" "flash-linear-attention==0.5.1"
export SGLANG_PLUGINS=atom_sglang
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE=atom.plugin.sglang.models
SGLANG_USE_AITER=1
ATOM_LOADER_NUM_THREADS=16 ATOM_SYNC_AFTER_LOAD=1 ATOM_DIST_TIMEOUT_SECONDS=3600
ATOM_USE_TRITON_GEMM=1 AITER_USE_GROUPED_GEMM=0 ATOM_USE_TRITON_MOE=0
AITER_FLYDSL_FORCE=1 AITER_FORCE_GFX1250=0
ATOM_MLA_PAGE_SIZE=1 ATOM_USE_UNIFIED_ATTN=1 ATOM_FORCE_ATTN_TRITON=1
AITER_LOG_LEVEL=WARNING
python3 -m sglang.launch_server --model-path /model_weights
--trust-remote-code --tp-size 8
--attention-backend aiter --kv-cache-dtype fp8_e4m3
--page-size 128 --mem-fraction-static 0.93
--context-length 16384 --disable-radix-cache
--host 127.0.0.1 --port 30000 2>&1 | tee /tmp/k3_atom.log'
Both use --context-length 16384 so both load cleanly (ATOM OOMs only when a 64k request is actually run). The per‑token KV density and the weight size are what we compare, and both are independent of context length, so this reproduces the finding without needing 64k.
Extract the numbers (from each *.log)
grep -E 'Load weight end|KV Cache is allocated|Mamba Cache is allocated|Memory pool end'
/tmp/k3_native.log | grep 'TP0'
grep -E 'Load weight end|KV Cache is allocated|Mamba Cache is allocated|Memory pool end'
/tmp/k3_atom.log | grep 'TP0'
Read off, per GPU:
Weights (GB) — Load weight end … mem usage= GB
KV pool — KV Cache is allocated … #tokens: … size: GB
KDA/mamba state — Mamba Cache is allocated … conv_state/ssm_state
Free after pools — Memory pool end. avail mem= GB
Per‑token KV = (KV size GB × 1e9) / #tokens.
Expected reference (what we measured, TP8)
Notes for the reader
Both configurations are SGLang + AITER; the only difference is the ATOM plugin env (SGLANG_PLUGINS=atom_sglang + SGLANG_EXTERNAL_MODEL_PACKAGE=...). So any delta is ATOM‑plugin‑specific.
The two recipes differ (native = triton/bf16‑KV/page‑1; ATOM = aiter/fp8‑KV/page‑128) because those are each stack's own documented K3 recipe. Note the KV gap persists against fp8's 2× advantage — i.e. it's structural (the ATOM pool allocates the 576‑wide MLA latent across all num_attention_heads/TP logical heads plus an unused paired‑V buffer), not a dtype artifact.
Numbers are per GPU; all TPk ranks match within ~1 GB.
(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support
rocminfo --support output
Additional Information
No response
Problem Description
ATOM Plugin with SGLang Consistently results in 20+ GB extra being allocated for weights when compared to SGLang operating without the plugin.
This happens for many models so it seems to be framework related to the ATOM plugin:
Why the plugin needs ~20 GB more
The keyword is bpreshuffle (block‑pre‑shuffle). AITER's high‑performance a8w8 block‑scale GEMM and fused‑MoE kernels don't consume weights in the checkpoint's native layout — they require the weight tensors pre‑swizzled and padded to the kernel's MFMA tile boundaries, with their block‑scale factor tensors alongside. So when the ATOM plugin builds each linear/MoE layer, it allocates those reordered + tile‑padded weight buffers (dominated by DeepSeek's 256‑expert MoE, where padding every expert to tile alignment adds up). Vanilla SGLang keeps the plain fp8 weights and uses generic kernels, so it doesn't pay that ~20 GB. It's essentially the memory price of AITER's tuned kernels.
Operating System
22.04.5 LTS (Jammy Jellyfish)
CPU
AMD EPYC 9575F 64-Core Processor
GPU
AMD Instinct MI355X
ROCm Version
7.2.4
ROCm Component
No response
Steps to Reproduce
Reproducing the SGLang (no plugin) vs ATOM‑plugin K3 memory comparison
Prerequisites
8× AMD Instinct MI355X (gfx950), ROCm 7.2.
~1.6 TB disk for the Kimi‑K3 checkpoint, and Hugging Face access to moonshotai/Kimi-K3.
Docker with GPU passthrough.
0. Download the weights once (shared by both runs)
export MODEL_DIR=/path/to/Kimi-K3 # pick any host dir with ~1.6 TB free
huggingface-cli download moonshotai/Kimi-K3 --local-dir "$MODEL_DIR"
Common docker flags used below:
DOCKER_FLAGS="--rm --device=/dev/kfd --device=/dev/dri --group-add video
--shm-size 16G --network host --security-opt seccomp=unconfined
--security-opt apparmor=unconfined --cap-add=SYS_PTRACE
-v $MODEL_DIR:/model_weights"
The four memory‑accounting lines we read are printed by SGLang during startup — you do not need to send any request. Launch, wait until you see Memory pool end, then Ctrl‑C.
A. WITHOUT the ATOM plugin (native SGLang + AITER)
Image: lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727 (SGLang's day‑0 K3 build).
docker pull lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727
docker run $DOCKER_FLAGS --entrypoint bash
lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727 -lc '
export SGLANG_USE_AITER=1 SGLANG_AITER_K3_OPT=1
AITER_FLYDSL_FORCE=1 AITER_SITUV2_A8W4=1
python3 -m sglang.launch_server --model-path /model_weights
--trust-remote-code --tp-size 8
--attention-backend triton --dtype bfloat16
--mem-fraction-static 0.85 --page-size 1
--context-length 16384 --disable-radix-cache
--host 127.0.0.1 --port 30000 2>&1 | tee /tmp/k3_native.log'
B. WITH the ATOM plugin
Image: rocm/atom-dev:sglang-latest (must include ATOM commit b687428 "[ATOM SGL][model] Kimi K3" or newer). The image does not bundle the KDA deps, so install them first (per ATOM's recipe).
docker pull rocm/atom-dev:sglang-latest
docker run $DOCKER_FLAGS --entrypoint bash
rocm/atom-dev:sglang-latest -lc '
pip install "fla-core==0.5.1" "flash-linear-attention==0.5.1"
export SGLANG_PLUGINS=atom_sglang
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE=atom.plugin.sglang.models
SGLANG_USE_AITER=1
ATOM_LOADER_NUM_THREADS=16 ATOM_SYNC_AFTER_LOAD=1 ATOM_DIST_TIMEOUT_SECONDS=3600
ATOM_USE_TRITON_GEMM=1 AITER_USE_GROUPED_GEMM=0 ATOM_USE_TRITON_MOE=0
AITER_FLYDSL_FORCE=1 AITER_FORCE_GFX1250=0
ATOM_MLA_PAGE_SIZE=1 ATOM_USE_UNIFIED_ATTN=1 ATOM_FORCE_ATTN_TRITON=1
AITER_LOG_LEVEL=WARNING
python3 -m sglang.launch_server --model-path /model_weights
--trust-remote-code --tp-size 8
--attention-backend aiter --kv-cache-dtype fp8_e4m3
--page-size 128 --mem-fraction-static 0.93
--context-length 16384 --disable-radix-cache
--host 127.0.0.1 --port 30000 2>&1 | tee /tmp/k3_atom.log'
Both use --context-length 16384 so both load cleanly (ATOM OOMs only when a 64k request is actually run). The per‑token KV density and the weight size are what we compare, and both are independent of context length, so this reproduces the finding without needing 64k.
Extract the numbers (from each *.log)
grep -E 'Load weight end|KV Cache is allocated|Mamba Cache is allocated|Memory pool end'
/tmp/k3_native.log | grep 'TP0'
grep -E 'Load weight end|KV Cache is allocated|Mamba Cache is allocated|Memory pool end'
/tmp/k3_atom.log | grep 'TP0'
Read off, per GPU:
Weights (GB) — Load weight end … mem usage= GB
KV pool — KV Cache is allocated … #tokens: … size: GB
KDA/mamba state — Mamba Cache is allocated … conv_state/ssm_state
Free after pools — Memory pool end. avail mem= GB
Per‑token KV = (KV size GB × 1e9) / #tokens.
Expected reference (what we measured, TP8)
Notes for the reader
Both configurations are SGLang + AITER; the only difference is the ATOM plugin env (SGLANG_PLUGINS=atom_sglang + SGLANG_EXTERNAL_MODEL_PACKAGE=...). So any delta is ATOM‑plugin‑specific.
The two recipes differ (native = triton/bf16‑KV/page‑1; ATOM = aiter/fp8‑KV/page‑128) because those are each stack's own documented K3 recipe. Note the KV gap persists against fp8's 2× advantage — i.e. it's structural (the ATOM pool allocates the 576‑wide MLA latent across all num_attention_heads/TP logical heads plus an unused paired‑V buffer), not a dtype artifact.
Numbers are per GPU; all TPk ranks match within ~1 GB.
(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support
rocminfo --support output
Additional Information
No response