AMD ATOM推理引擎：如何用分页KV缓存和分段编译实现大模型推理性能翻倍

在探讨大模型推理优化时，最引人注目的往往是单一算子的性能提升：GEMM 加速了多少，Attention 快了百分之几，MoE 路由是否已经融合。然而，真正左右线上服务体验的，并非单个 kernel 的峰值指标，而是一条请求从 HTTP 入口进入，经历分词、调度、写入 KV 缓存、进入 GPU 图执行、经过采样，再到流式返回的完整生命周期。

以下图表展示了 DeepSeek R1 0528 模型在 FP8 精度下，不同 GPU 与推理框架在单卡 Token 吞吐量和端到端延迟之间的权衡关系。横轴代表延迟（秒），纵轴表示吞吐量（tok/s/gpu），整体趋势显示吞吐量随延迟升高而增长，凸显了大模型推理中典型的“速度-延迟”性能折衷。其中，AMD MI355X 搭配 ATOM 框架表现最为突出，延迟约 26 秒时吞吐量突破 1100 tok/s/gpu；同系列的 SGLang 版本位列第二；而 NVIDIA B200（TRT）在相近延迟下的吞吐量则显著较低，直观呈现了不同硬件与优化栈在该模型推理性能上的明显差异。更多数据，请访问：https://inferencex.semianalysis.com/

ATOM (AiTer Optimized Model) 是一个轻量级的 vLLM 风格实现，专注于基于 AITER 的集成与优化。
代码仓库：https://github.com/ROCm/ATOM
文档地址：https://rocm.github.io/ATOM/docs
全文约 5000 字，阅读需 22 分钟，播客时长 27 分钟

ATOM（AiTer Optimized Model）正是这样一个项目：它并非包罗万象的训练/推理框架，而是一个轻量级的 vLLM 风格推理引擎。它基于 AMD ROCm 平台上的 AITER 内核，将模型执行、分页 KV 缓存、编译图捕获、多 GPU 并行以及 OpenAI 兼容服务串联成一个可运行、可压测、可剖析的完整系统。其价值不仅在于“能跑”，更在于将 AMD GPU 上的大模型推理优化路径具体化为工程结构。

下表列出了 ATOM 当前支持的主要模型家族，包括 HuggingFace 架构、密集与混合专家（MoE）类型，以及关键备注，方便快速了解兼容范围与注意事项。

模型家族	HF 架构	密集/MoE	备注
Llama[1]	`LlamaForCausalLM`	Dense	Llama 2, Llama 3, Llama 3.1
Qwen3[2]	`Qwen3ForCausalLM`	Dense
Qwen3-MoE[3]	`Qwen3MoeForCausalLM`	MoE	128 专家，top-8 路由
Qwen3-Next[4]	`Qwen3NextForCausalLM`	MoE	混合全注意力 + Gated DeltaNet
DeepSeek V2/V3[5]	`DeepseekV3ForCausalLM`	MoE	MLA 注意力，MTP 推测解码
Mixtral[6]	`MixtralForCausalLM`	MoE	8 专家，top-2 路由
GLM-4-MoE[7]	`Glm4MoeForCausalLM`	MoE
GLM-5[8]	`GlmMoeDsaForCausalLM`	MoE	MLA 注意力，类似 DeepSeek V3.2，详见 recipe[9]
GPT-OSS[10]	`GptOssForCausalLM`	MoE	滑动窗口 + 注意力汇聚
Kimi-K2[11]	通过 `--trust-remote-code`	MoE	详见 recipe[12]
MiMo-V2-Flash[13]	`MiMoV2FlashForCausalLM`	MoE	混合全注意力 + SWA，3 层 MTP，详见 recipe[14]

本文目录

一、快速上手：先把 ATOM 跑起来
二、项目定位：不是“模型库”，而是 ROCm 推理链路样板间
2.1 ATOM 的核心目标
2.2 整体组件关系
三、请求生命周期：一条 prompt 如何穿过 ATOM
3.1 从 API 到 Sequence
3.2 Sequence 是请求在系统里的身份证
四、调度器：prefill-first 与连续 batching 的取舍
4.1 为什么先 prefill
4.2 Decode 阶段：KV 空间不足就抢占
五、分页 KV Cache：把显存切成可复用的街区
5.1 BlockManager 的本质
5.2 Prefix caching：用 xxhash64 识别相同前缀
六、算子层：AITER 是 ATOM 的性能地基
6.1 Linear：按量化类型分派 GEMM
6.2 Attention：把复杂动态逻辑封装成可切分 custom op
七、编译与 CUDA Graph：把 decode 的 CPU 开销压到最低
7.1 四档编译级别
7.2 为什么 prefill eager，decode graph replay
八、MoE、量化与多卡：ATOM 面向大模型的真正压力区
8.1 MoE 路径：路由、专家并行与 MORI 通信
8.2 量化不是单一格式，而是一组执行合同
九、P/D 分离：把 prefill 与 decode 拆到不同 GPU 节点
十、如何评价 ATOM：轻量外壳下的系统工程密度
结语：ATOM 的意义在于把“ROCm 能跑大模型”推进到“ROCm 能服务大模型”

一、快速上手：先把 ATOM 跑起来

ATOM 的官方 README 提供了最简路径：直接使用 nightly Docker 镜像。

环境搭建与快速上手

运行 ATOM 推理引擎的核心前提是配备 AMD GPU、安装 ROCm 软件栈，并准备好 Docker 环境。官方最推荐的部署方式，是直接拉取预构建好的 Docker 镜像，该镜像已集成了 AITER 和 ATOM 所有依赖。根据项目 README.md 的说明，启动方式如下：

docker pull rocm/atom-dev:latest  

docker run -it --network=host   
  --device=/dev/kfd   
  --device=/dev/dri   
  --group-add video   
  --cap-add=SYS_PTRACE   
  --security-opt seccomp=unconfined   
  -v $HOME:/home/$USER   
  -v /mnt:/mnt   
  -v /data:/data   
  --shm-size=16G   
  --ulimit memlock=-1   
  --ulimit stack=67108864   
rocm/atom-dev:latest

若开发者希望从基础 ROCm PyTorch 镜像开始手动构建，README 也提供了另一条路径：先进入 rocm/pytorch:rocm7.0.2_ubuntu24.04_py3.12_pytorch_release_2.8.0 基础镜像，再安装 AITER 与 ATOM。

pip install amd-aiter  
git clone https://github.com/ROCm/ATOM.git && pip install ./ATOM

运行一个最小推理示例，可通过 Python 模块方式直接启动：

python -m atom.examples.simple_inference --model meta-llama/Meta-Llama-3-8B --kv_cache_dtype fp8

对于更关注服务化部署的用户，ATOM 提供了兼容 OpenAI 接口的服务端，可直接暴露 /v1/chat/completions 与 /v1/completions 类接口：

# 单 GPU 场景  
python -m atom.entrypoints.openai_server --model Qwen/Qwen3-0.6B --kv_cache_dtype fp8  

# 多 GPU 场景，启用张量并行  
python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv_cache_dtype fp8 -tp 8  

# 启用 MTP 投机解码  
python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-R1 --kv_cache_dtype fp8 -tp 8   
  --method mtp --num-speculative-tokens 3

关于更完整的架构设计、配置参数、模型支持列表、调度策略、编译流程、多卡部署方案以及性能基准测试，可参考官方项目文档入口：rocm.github.io/ATOM/docs[15]。

二、项目定位：不是“模型库”，而是 ROCm 推理链路样板间

2.1 ATOM 的核心目标

README 对 ATOM 的定义非常清晰：它是一个 轻量级的 vLLM 风格实现，核心在于基于 AITER[16] 进行集成与优化。 换言之，ATOM 的重点并非重新发明 Transformer 架构，而是回答一个更系统性的问题：

如果 AMD 已经拥有一套高性能的 ROCm 内核，如何将它们组织成一个真实可服务的大模型推理引擎？

这一目标决定了 ATOM 的架构关注点非常工程化：

算子层：将 Linear、Attention、MoE、Norm、Sampling 等重计算路径交由 AITER 处理。
模型层：支持 Llama、Qwen3、DeepSeek、Mixtral、GLM、GPT-OSS、Kimi、MiMo 等多种主流架构。
执行层：通过 Scheduler、BlockManager、ModelRunner 组件实现连续批处理与 KV 缓存管理。
编译层：提供从 0 到 3 共四档编译级别，生产环境默认采用分段式 torch.compile 结合 CUDA graph。
服务层：暴露兼容 OpenAI 的 API，并提供性能基准测试、性能剖析与精度验证工具。
分布式层：支持张量并行（TP）、数据并行（DP）、专家并行（EP），以及基于 MORI 的 MoE all-to-all 通信与 P/D 分离 KV 传输。

因此，ATOM 更像是一座“ROCm 推理系统样板间”：它将 AMD GPU 上的高性能内核、编译策略和线上推理机制整合在一个可验证的框架之中。

2.2 整体组件关系

官方架构文档给出的主链路如下：LLMEngine 面向用户，CoreManager 负责多进程编排，EngineCore 持有调度器与模型执行器，ModelRunner 真正驱动每个 GPU 上的模型前向计算。

我们可以将其类比为一家餐厅的运作流程：

LLMEngine 是前台，负责接收用户请求；
InputOutputProcessor 是点餐员，负责分词、反分词与统计工作；
Scheduler 是后厨调度员，决定优先处理哪些任务；
BlockManager 是仓储管理员，管理 KV 缓存的存储货架；
ModelRunner 是厨师，调用 AITER 内核完成实际计算；
CoreManager 和 ZMQ 类似传菜系统，将请求分发到不同进程。

三、请求生命周期：一条 prompt 如何穿过 ATOM

3.1 从 API 到 Sequence

ATOM 的用户入口定义在 atom/model_engine/llm_engine.py 文件中。

好的，这是根据您的要求对指定文章片段进行的深度重写与降重结果。

LLMEngine.generate() 的执行流程起始于 add_request() 的调用。该方法将用户输入的 prompt 与采样参数传递给 InputOutputProcessor.preprocess()，由其生成内部的 Sequence 对象，随后该对象被转交给 CoreManager 进行后续处理。

# atom/model_engine/llm_engine.py  
def generate(  
    self,  
    prompts: list[str],  
    sampling_params: SamplingParams | list[SamplingParams],  
) -> list[str]:  
    # Reset round-robin counter to ensure consistent DP not core dump  
    self.core_mgr._rr_counter = 0  

    self.add_request(prompts, sampling_params)  
    outputs = {}  
    whilenot self.is_finished() and (  
        self.core_mgr.is_alive() or self.core_mgr.is_rest()  
    ):  
        seqs = self.step()  
        outs = self.io_processor.postprocess(seqs)  
        outputs.update(outs)  

    outputs = [outputs[seq_id] for seq_id in sorted(outputs)]  
    return outputs

这段代码虽然看起来简洁，却精确反映了 ATOM 的核心控制模式：用户线程并不直接驱动模型推理，而是通过一个持续运行的循环，从 engine core 拉取已经处理完成的 Sequence 对象，再执行后续的解码（detokenize）和统计工作。

3.2 Sequence：请求在系统中的唯一标识

Sequence 对象承载了请求的完整生命周期信息，包括 prompt token、生成 token、KV block table、状态、采样参数，以及到达时间、首个 token 生成时间、离开时间等关键性能指标。它既是调度系统的基本操作单元，也是性能数据采集的统计单位。

ATOM 中请求的状态流转大致遵循以下路径：

WAITING -> RUNNING(PREFILL) -> RUNNING(DECODE) -> FINISHED

当 KV cache 空间不足时，正在运行的请求可能被抢占，并被迫回到等待队列，重新经历 prefill 阶段。这种机制看似造成了计算浪费，但在高并发场景下，它赋予了系统在有限的 KV 缓存资源中维持高整体吞吐量的能力。

unsetunset四、调度器：prefill-first 与连续 batching 的取舍unsetunset

4.1 为何优先执行 prefill

大模型推理过程分为两个阶段：prefill 阶段处理完整的输入 prompt，计算量大且输入形状多变；decode 阶段则通常每次只生成一个或少量 token，其操作频繁、对延迟敏感，且非常适合通过 CUDA graph 进行重放优化。

ATOM 的 Scheduler.schedule() 方法采用了 prefill-first 策略：只要存在可以调度的等待请求，系统就会优先将它们组成一个 prefill batch；只有在没有 prefill 请求需要处理时，才会调度 decode 任务。

核心代码逻辑来自 atom/model_engine/scheduler.py：

# atom/model_engine/scheduler.py  
def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]:  
    """Select the next batch of sequences for a forward pass.  

    Tries prefill first; if no new prefills are ready, falls back to  
    decoding already-running sequences.  
    """  
    scheduled_seqs = {}  
    num_seqs_prefill = 0  
    num_batched_tokens = 0  
    skipped_waiting_requests: deque[Sequence] = deque()  
    num_scheduled_tokens: list[int] = []  
    scheduled_spec_decode_tokens: dict[int, np.ndarray] = {}  

    ifnot self.running andnot self.waiting:  
        returnNone  

    # --- Prefill scheduling ---  
    while self.waiting and num_seqs_prefill < self.max_num_seqs:  
        seq = self.waiting.popleft()  

        num_new_tokens = seq.num_tokens - seq.num_cached_tokens  
        if (  
            num_batched_tokens + num_new_tokens > self.max_num_batched_tokens  
            ornot self.block_manager.can_allocate(seq)  
        ):  
            self.waiting.appendleft(seq)  
            break  

        self.block_manager.allocate(seq)  

        num_seqs_prefill += 1  
        num_new_tokens = seq.num_tokens - seq.num_cached_tokens  
        if self.cache_stats:  
            self.cache_stats.update(seq.num_cached_tokens, seq.num_tokens)  
        num_batched_tokens += num_new_tokens  
        seq.status = SequenceStatus.RUNNING  
        seq.type = SequenceType.PREFILL  
        self.running.append(seq)  
        scheduled_seqs[seq.id] = seq  
        num_scheduled_tokens.append(num_new_tokens)

这段调度逻辑的核心在于对两个关键预算的控制：

4.2 Decode 阶段：KV 空间不足就抢占

当调度器未发现任何 prefill 任务时，便会转入 decode 阶段。在此阶段，它会逐一检查每条正在运行的序列（running sequence）是否仍有空间追加新的 KV block。若空间不足，则会抢占队尾的请求。

# atom/model_engine/scheduler.py  
# --- Decode scheduling ---  
num_seqs_decode = 0  
num_new_tokens = self.mtp_k + 1  
while self.running and num_seqs_decode < self.max_num_seqs:  
seq = self.running.popleft()  
while not self.block_manager.can_append(seq, num_new_tokens):  
if self.running:  
self.preempt(self.running.pop())  
else:  
self.preempt(seq)  
break  
else:  
if seq.spec_token_ids.size > 0:  
scheduled_spec_decode_tokens[seq.id] = seq.spec_token_ids  
num_seqs_decode += 1  
if not getattr(seq, "is_first_decode", False):  
self.block_manager.may_append(seq, num_new_tokens)  
scheduled_seqs[seq.id] = seq  
seq.type = SequenceType.DECODE  
num_scheduled_tokens.append(num_new_tokens)  
seq.is_first_decode = False

这段代码中，num_new_tokens = self.mtp_k + 1 的设计非常巧妙。当启用了 MTP 推测解码（speculative decoding）后，每次 decode 步骤不仅需要为最终的 target token 预留 KV 空间，还必须为 draft token 提前分配好缓存位置。调度器本身并不关心模型内部如何验证 draft token，但它必须确保这些可能被写入的 token 都有对应的缓存位置可用。

五、分页 KV Cache：把显存切成可复用的街区

5.1 BlockManager 的本质

在大模型推理场景中，显存瓶颈往往不是模型权重，而是 KV cache。ATOM 通过 BlockManager 来管理固定大小的 KV block，默认每个 block 的大小为 16 个 token。每个 Sequence 都维护着一张 block_table，记录了其所占用的所有 block。

Block 的数据结构非常精简：

# atom/model_engine/block_manager.py  
class Block:  

def __init__(self, block_id):  
self.block_id = block_id  
self.ref_count = 0  
self.hash = -1  
self.token_ids = []  

def update(self, hash: int, token_ids: list[int]):  
self.hash = hash  
self.token_ids = token_ids  

def reset(self):  
self.ref_count = 1  
self.hash = -1  
self.token_ids = []

我们可以把 KV cache 想象成一家酒店：请求到来时分配房间，请求结束后退房；如果两位客人拥有完全相同的前缀 prompt，它们甚至可以共享前面的几间房间。这里的 ref_count 就是记录共享房间的住户数量。

5.2 Prefix caching：用 xxhash64 识别相同前缀

ATOM 的 prefix caching 机制采用了 xxhash64，并且使用了链式哈希（chained hash）策略：当前 block 的哈希值会包含前一个 block 的哈希值。这样一来，即使两个 block 的 token 内容完全一致，只要它们所处的上下文不同（即前文不同），最终的哈希值也会不同，从而避免了错误的上下文共享。

# atom/model_engine/block_manager.py  
@classmethod  
def compute_hash(cls, token_ids: list[int], prefix: int = -1):  
h = xxhash.xxh64()  
if prefix != -1:  
h.update(prefix.to_bytes(8, "little"))  
h.update(np.array(token_ids).tobytes())  
return h.intdigest()

这绝非简单的“对 token 列表做哈希”。链式哈希确保了每个 KV block 的语义与其所处的上下文紧密绑定。对于多轮对话、批量相似 prompt 以及 RAG 模板复用等场景，这种机制能够显著减少重复的 prefill 计算，从而大幅提升吞吐量。

六、算子层：AITER 是 ATOM 的性能地基

6.1 Linear：按量化类型分派 GEMM

ATOM 的 atom/model_ops/linear.py 是理解其算子策略的核心文件。该文件根据不同的量化格式，将计算任务分派给 AITER 中对应的 GEMM 实现：

无量化：tgemm.mm
per-tensor FP8：带 scale 的 tgemm.mm
per-token INT8：gemm_a8w8
per-token FP8：gemm_a8w8_bpreshuffle
per-1×128 FP8：gemm_a8w8_blockscale_bpreshuffle
per-1×32 MXFP4：gemm_a4w4

核心调度逻辑如下：

# atom/model_ops/linear.py  
@mark_trace  
def forward(  
self, x: torch.Tensor, x_scale: Optional[torch.Tensor] = None, otype=dtypes.bf16  
) -> torch.Tensor:  
if self.quant_type.value == QuantType.No.value:  
y = tgemm.mm(  
x,  
self.weight,  
self.bias,  
otype=otype,  
)  
else:  
if x_scale is None:  
quant_func = self.quant_func  
if self.quant_type.value == QuantType.per_1x128.value:  
quant_func = functools_partial(  
self.quant_func, transpose_scale=True  
)  
if self.quant_type.value != QuantType.per_1x32.value:  
x, x_scale = quant_func(  
x,  
quant_dtype=self.params_dtype,  
scale=getattr(self, "input_scale", None),  
)  
if self.quant_type.value == QuantType.per_Tensor.value:  
y = tgemm.mm(  
x,  
self.weight,  
self.bias,  
otype=otype,  
scale_a=x_scale,  
scale_b=self.weight_scale,  
)  
elif self.quant_type.value == QuantType.per_Token.value:  
if self.params_dtype == dtypes.i8:  
y = gemm_a8w8(  
x,  
self.weight,  
x_scale,  
self.weight_scale,  
self.bias,  
dtype=otype,  
)  
else:  
y = gemm_a8w8_bpreshuffle(  
x,  
self.weight,  
x_scale,  
self.weight_scale,  
dtype=otype,  
)  
elif self.quant_type.value == QuantType.per_1x128.value:  
y = gemm_a8w8_blockscale_preshuffle_impl(  
x,  
self.weight,  
x_scale,  
self.weight_scale,  
dtype=otype,  
prefix=self.prefix,  
)  
elif self.quant_type.value == QuantType.per_1x32.value:  
y = gemm_a4w4_quant(  
x,  
x_scale,  
self.weight,  
otype,  
self.weight_scale.data,  
self.params_dtype,  
getattr(self, "input_scale", None),  
self.output_size,  
)  
if self.tp_dim == 1 and self.tp_size > 1 and self.reduce_results:  
y = get_tp_group().all_reduce(y, ca_fp8_quant=False)  
return y

此处蕴含两个关键设计思路：

第一，ATOM 并未将量化视为“模型加载时一次性完成的转换操作”，而是在每一层的前向传播（layer forward）过程中，始终保留 scale、weight layout 以及 kernel dispatch 的清晰路径。

第二，Tensor Parallel 的相关逻辑同样被整合进 Linear wrapper 内部进行处理。例如，针对 row parallel 输出所需的 all-reduce 操作，都在此处统一完成，从而使模型文件本身能够保持相对简洁的结构。

6.2 Attention：将复杂动态逻辑封装为可切分的自定义算子

在推理系统中，Attention 机制无疑是编译器最难“优雅处理”的环节。这主要是因为其内部高度依赖动态的 sequence length、block table、KV cache 以及 prefill/decode 分支。面对这一挑战，ATOM 并未试图强行让 Dynamo/Inductor 理解所有细节，而是采取了一种更务实的策略：将 Attention 注册为 custom op（自定义算子）和 splitting op（分割算子）。

# atom/model_ops/base_attention.py  
# Dynamo 将不会尝试检查 prefill 或 decode 的任何内部操作
# 这样一来，尽管 Attention 操作本身非常复杂，
# 我们依然能够将模型的计算图作为一个完整的图来捕获
@mark_spliting_op(is_custom=True, gen_fake=fake_, mutates_args=[])
def unified_attention_with_output_base(
q: torch.Tensor,
q_scale: Optional[torch.Tensor],
k: torch.Tensor,
v: torch.Tensor,
positions: torch.Tensor,
layer_name: str,
use_mla: bool,
qkv: torch.Tensor,
) -> torch.Tensor:
atom_config = get_current_atom_config()
self = atom_config.compilation_config.static_forward_context[layer_name]
if use_mla:
return self.impl.forward(
layer=self,
query=q,
k_nope=k,
k_rope=v,
positions=positions,
q_scale=q_scale,
)
else:
return self.impl.forward(
layer=self,
query=q,
key=k,
value=v,
position=positions,
q_scale=q_scale,
qkv=qkv,
)

这段代码清晰揭示了 ATOM 的编译哲学：
并非所有组件都需交由编译器优化；该采用黑盒策略的地方就黑盒，该划分图结构的地方就划分。

在 Attention 内部，系统会根据 MHA/MLA、prefill/decode、Triton/ASM 及 KV cache dtype 等不同条件，调用相应的 AITER 内核。对于 DeepSeek 这类采用 MLA 的模型，ATOM 还特别支持压缩 KV 表示以及专用的 MLA decode/prefill kernel。

unsetunset七、编译与 CUDA Graph：将 decode 阶段的 CPU 开销降至极限unsetunset

7.1 四档编译级别

ATOM 官方文档定义了从 0 到 3 的四个 compilation level：

Level 0：纯 eager 模式；
Level 1：直接使用 Dynamo，后端为 eager；
Level 2：Dynamo + Inductor 组合；
Level 3：分段编译（Piecewise compilation）+ CUDA graph，生产环境默认选择。

Level 3 的核心思路是：在 Attention 等 splitting op 处将模型图拆分开来。对于能够编译的子图，交给 Inductor 处理；而对于复杂动态的 Attention，则保留给专用 kernel。随后，对 decode batch 进行 CUDA graph 捕获，在运行时直接重放。

这个过程好比修建高速公路：并非要把每一条乡间小道都改造成高速路，而是识别出那些车流量稳定、形状固定的主干路段，铺设最平整的路面；对于那些容易拥堵、入口复杂的路口，则保留专门的调度机制。

7.2 为何 prefill 用 eager，decode 用 graph replay

根据官方编译文档的说明：prefill 阶段序列长度变化剧烈，通常采用 eager 模式；而 decode 阶段的 batch 形状更为稳定，非常适合预先捕获 CUDA graph。ModelRunner.run_model() 中的逻辑可以概括如下：

# docs/compilation_cudagraph_guide.md
def run_model(self, input_ids):
forward_context = get_forward_context()
context = forward_context.context
bs = context.batch_size
is_prefill = context.is_prefill
positions = context.positions

if is_prefill or self.enforce_eager or bs > self.graph_bs[-1]:
# Eager 路径：处理 prefill、强制 eager 模式或超规格 batch
hidden_states = self.model(input_ids, positions)
else:
# Graph replay 路径：处理在捕获范围内的 decode batch
graph_bs = context.graph_bs
max_q_len = forward_context.attn_metadata.max_seqlen_q
graph_key = (graph_bs, max_q_len)
self.graphs[graph_key].replay()
num_tokens = context.batch_size * max_q_len
hidden_states = self.forward_vars["outputs"][:num_tokens]

return self.model.compute_logits(hidden_states), hidden_states

这段伪代码揭示了一个非常务实的核心原则：
prefill 阶段追求高吞吐与灵活性，而 decode 阶段则致力于实现稳定且低延迟的响应。

在在线服务场景中，每生成一个 token 都需要执行一次完整的模型推理。如果每次推理都由 Python 层重新发射大量 kernel，那么 CPU 侧的 launch overhead 会被显著放大。CUDA graph replay 的核心价值，就在于将这一调度开销彻底压制下来。

unsetunset八、MoE、量化与多卡：ATOM 直面大模型真实压力unsetunset

8.1 MoE 路径：路由、专家并行与 MORI 通信

ATOM 原生支持 DeepSeek、Mixtral、Qwen3-MoE、GLM 以及 GPT-OSS 等 MoE 架构模型。MoE 的难点远不止专家 GEMM 计算，还涉及 token routing、top-k 选择、跨专家通信以及最终的 combine 操作。根据项目文档，FusedMoE 会根据量化配置自动选择不同的执行路径，例如 FP8、MXFP4 或无量化模式，并且能够结合 MORI 实现 expert parallel all-to-all 通信。

在大模型推理中，MoE 就像一座大型分拣中心：每个 token 根据 router 的评分被分发到少数几个专家，专家完成计算后再汇总结果。如果分拣（routing）、运输（通信）和计算（GEMM）三个阶段无法流水化，即使理论 FLOPs 再高，也难以转化为实际的吞吐量。ATOM 引入 MORI all-to-all 与 Two-Batch Overlap 机制的根本目的，就是将通信从“等待阻塞项”转化为可被计算隐藏的流水线阶段。

8.2 量化并非单一格式，而是一套执行契约

README 明确指出 ATOM 支持 FP8、MXFP4、INT8、INT4 等多种量化格式，并能从 HuggingFace 配置中自动检测。从源码层面看，量化格式会直接影响：

参数的数据类型（dtype）；
scale tensor 的形状（shape）；
权重是否需要预重排（preshuffle）；
激活值是否执行动态量化（dynamic quant）；
GEMM kernel 的选择；
张量并行（TP）分片的加载方式；
MoE 专家 kernel 的执行路径。

这意味着，ATOM 中的量化绝不仅仅是“把权重压小”那么简单，而是一套贯穿模型加载、内存布局、计算调度和分布式通信的完整执行契约。不同的量化格式对应着截然不同的 kernel 和 layout，错误的 layout 甚至可能完全抵消量化带来的性能收益。

unsetunset九、P/D 分离：将 prefill 与 decode 拆分至不同 GPU 节点unsetunset

ATOM 在 2026/03 的 README 更新中，重点强调了其对 Prefill/Decode disaggregation 的支持。其 atom/kv_transfer/disaggregation/README.md 文档详细说明了机制：prefill 节点负责计算 KV cache，随后通过 MORI-IO 的 RDMA 能力将 KV 传输给 decode 节点，decode 节点无需执行 prefill 阶段，直接开始生成 token。

简化后的流程如下：

Client -> Proxy
-> Prefill Node: 计算 KV cache
-> Proxy: 接收 block metadata
-> Decode Node: RDMA 读取 KV 并生成 token
-> Proxy -> Client

这一设计背后的系统动机非常清晰：prefill 和 decode 的资源需求画像截然不同。prefill 更像是大规模矩阵计算，而 decode 则是低延迟的小步迭代。将它们拆分到不同节点，可以分别针对吞吐量和延迟进行独立优化，同时也便于在服务集群中灵活地进行扩缩容。

unsetunset十、如何评价 ATOM：轻量外壳下的系统工程密度unsetunset

ATOM 的代码中 99% 以上是 Python，但这绝不意味着它“只是一个 Python 框架”。恰恰相反，它的 Python 层承担的是高性能推理系统中最难统一的部分：调度、状态管理、缓存、编译边界、分布式通信、服务接口以及性能验证。

其核心价值可以总结为以下四点：

将 AITER kernel 系统化
单个 kernel 再快，如果没有模型 wrapper、调度器和服务入口，也很难落地。ATOM 将 AITER 放入了完整的推理链路中。
将 vLLM-like 架构移植到 ROCm 语境
Scheduler、paged KV、continuous batching、OpenAI API 都是成熟推理系统的关键构件，ATOM 将它们与 AMD GPU 的优化特性相结合。
务实处理编译器与手写 kernel 的边界
Attention 等动态复杂路径通过 custom op 切出，Linear、MLP 等稳定路径则交给编译器和 CUDA graph。这是工程上更稳健的选择。
前置大模型服务的复杂场景
MoE、MTP speculative decoding、FP8/MXFP4、TP/DP/EP、P/D 分离以及 benchmark dashboard，都表明 ATOM 面向的不是 demo，而是真实的线上推理压力场景。

当然，它也有自身的局限性。ATOM 高度依赖 ROCm、AMD GPU 与 AITER；首次运行时可能存在编译耗时；部分高级能力如 P/D 分离、MORI、EP 通信依赖于特定的硬件和网络环境。对于普通单卡用户而言，最佳入门方式是先运行 README 中的 simple inference 与 OpenAI server；对于系统工程读者来说，更值得深入阅读的是 atom/model_engine/、atom/model_ops/、docs/compilation_cudagraph_guide.md 与 docs/scheduling_kv_cache_guide.md。

unsetunset结语：ATOM 的意义在于将“ROCm 能跑大模型”推进到“ROCm 能服务大模型”unsetunset

如果只看项目名称，ATOM 像是一个模型优化仓库；但如果读完代码，你会发现它真正关心的是端到端推理服务的工程闭环。

它将 AITER 的底层 kernel、分页 KV cache 的显存管理、prefill-first 的动态调度、piecewise compile 的图优化、CUDA graph 的低延迟重放、多 GPU 并行以及 OpenAI-compatible server 完美地拼接在一起。

这类项目的意义不在于“替代所有框架”，而在于清晰地展示了一条路线：AMD ROCm 生态中的大模型推理，不应停留在单算子 benchmark，而应走向可部署、可观测、可压测、可扩展的系统形态。ATOM 正是在这条路上，将抽象的硬件能力翻译为具体工程结构的一个关键样本。

参考资料[1]

Llama: https://huggingface.co/meta-llama

[2]

Qwen3: https://huggingface.co/Qwen

[3]

Qwen3-MoE: https://huggingface.co/Qwen

[4]

Qwen3-Next: https://huggingface.co/Qwen

[5]

DeepSeek V2/V3: https://huggingface.co/deepseek-ai

[6]

Mixtral: https://huggingface.co/mistralai/Mixtral-8x7B-v0.1

[7]

GLM-4-MoE: https://huggingface.co/THUDM

[8]

GLM-5: https://huggingface.co/zai-org/GLM-5-FP8

[9]

recipe: https://github.com/ROCm/ATOM/blob/main/recipes/Kimi-K2-Thinking.md

[10]

GPT-OSS: https://huggingface.co/openai

参考文献

Kimi-K2 模型： https://huggingface.co/moonshotai/Kimi-K2-Thinking
对应推理配置（recipe）： https://github.com/ROCm/ATOM/blob/main/recipes/Kimi-K2-Thinking.md
MiMo-V2-Flash 模型： https://huggingface.co/XiaomiMiMo/MiMo-V2-Flash
混合全注意力 + 滑动窗口注意力，3层 MTP 架构配置： https://github.com/ROCm/ATOM/blob/main/recipes/MiMo-V2-Flash.md
ATOM 官方文档： https://rocm.github.io/ATOM/docs
AITER 算子库： https://github.com/ROCm/aiter

延伸阅读

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：https://www.itsolotime.com/archives/33112