Architecting for Hardware Diversity: Meta’s End-to-End Co-Design for High-Performance AI Training on AMD GPUs

面向硬件多样性的架构：Meta 在 AMD GPU 上高性能 AI 训练的端到端协同设计

Haoyu Zhang 1 , Arshad Rahman 2 , Adnan Aziz 1 , Ajith Prabhudev 1 , Amin Farmahini 1 , Ben Rhoads 1 , Benson Ma 1 , Brian Shi 2 , Cen Zhao 1 , CK Luk 1 , Dingming Wu 1 , Emad El-Haraty 1 , Eric Chen 1 , Frank Li 1 , Gary Huang 1 , Gefei Zuo 1 , Haoran Zhang 1 , Hiwot Kassa 1 , Hongyi Zeng 1 , Jack Taylor 2 , Jeremy Hadidjojo 1 , John Donati 1 , Jonathan Sheffield 1 , Julio Machado Silva 2 , Karthik Ganesan 1 , Kent Russell 2 , KR Kishore 1 , Lei Tian 1 , Liu Ke 1 , Muhammad Osama 2 , Peng Chen 1 , Rakesh Hemnani 1 , Ryan Swann 2 , Saleel Kudchadker 2 , Shankar Yanamandram 1 , Shobhit Kanaujia 1 , Shucai Xiao 2 , Srinath Bayareddy 1 , Srinivas Vaidyanathan 1 , Sudharssun Subramanian 1 , Supadchaya Puangpontip 1 , Taylor Robie 1 , Tianhao Huang 1 , Tyler Hart 1 , Viswesh Sankaran 1 , Weijun Jiang 2 , Wenkai Du 2 , Xiaodong Wang 1 , Xiaohu Guo 2 , Yao Rong 1 , Yinbin Ma 1 , Yue Dong 1 , Yuzhen Huang 1 , Yvonne Liu 1 , Zhan Shu 1 , Dan Rabinovitsj 1 , Ajit Mathews 1 , Chunqiang Tang 1 Affiliation:  1 Meta    2 AMD

Abstract

摘要

Meta deploys a heterogeneous fleet comprising NVIDIA GPUs, AMD GPUs, and proprietary MTIA accelerators to reduce costs and mitigate vendor lock-in. However, the extensive history of NVIDIA-specific optimizations hinders performance portability to other hardware. This paper presents our methods for optimizing alternative accelerators, using AMD Instinct™ MI350X versus NVIDIA B200 for recommendation model training as a case study. We propose three reusable principles. First, hardware-model co-design is essential: rather than directly porting models from NVIDIA GPUs, revising model structures and software recipes allows full exploitation of accelerator-specific strengths (e.g., MI350X’s $1.6\times$ HBM capacity relative to B200). Second, our analytical performance model enables zero-warmup optimal kernel selection across diverse datatypes and matrix shapes in production models, drastically simplifying tuning. Third, increasingly mature agentic workflows have lowered the barrier to porting from the CUDA ecosystem; we extend these workflows for microarchitecture-aware kernel optimizations and custom shim generation. Compared to the directly ported baseline, these methods deliver up to $1.5\times$ improvement in end-to-end training throughput. To our knowledge, this is the first hyperscaler report on productionizing AMD GPUs for large-scale training. Our contributions have been upstreamed to the open-source PyTorch, Triton, and ROCm ecosystems.

Meta 部署由 NVIDIA GPU、AMD GPU 与自研 MTIA 加速器组成的异构机队，以降低成本并缓解供应商锁定。然而，长期围绕 NVIDIA 的优化历史阻碍了性能向其他硬件的可移植性。本文以推荐模型训练上 AMD Instinct™ MI350X 对标 NVIDIA B200 为案例，给出优化替代加速器的方法。我们提出三条可复用原则。第一，硬件—模型协同设计必不可少：与其把模型从 NVIDIA GPU 直接移植，不如修订模型结构与软件配方，以充分挖掘加速器特有优势（例如 MI350X 相对 B200 约 1.6× 的 HBM 容量）。第二，我们的解析性能模型可在生产模型多样的数据类型与矩阵形状上实现零热身最优核选择，大幅简化调优。第三，日益成熟的智能体工作流降低了从 CUDA 生态移植的门槛；我们扩展这些工作流以做微架构感知的核优化与定制 shim 生成。相对直接移植基线，这些方法将端到端训练吞吐最高提升约 1.5×。据我们所知，这是首份关于将 AMD GPU 用于大规模生产训练的超大规模报告。相关贡献已上游至开源 PyTorch、Triton 与 ROCm 生态。

I Introduction

I 引言

Index Terms:

The immense potential of AI applications has catalyzed an unprecedented expansion in the accelerator landscape. While NVIDIA remains the incumbent platform, the AMD Instinct™ [1, 7] GPU has emerged as a viable alternative. Meanwhile, hyperscalers have also developed custom silicon, including Meta MTIA [50, 46, 49, 48], Google TPU [53], Amazon Trainium [60], Microsoft Maia [59], Huawei Ascend [57], Alibaba Hanguang [58], and Baidu Kunlun [61].

索引词：

AI 应用的巨大潜力推动了加速器版图的空前扩张。NVIDIA 仍是主流平台，但 AMD Instinct™ [1, 7] GPU 已成为可行替代。与此同时，超大规模厂商也在自研芯片，包括 Meta MTIA [50, 46, 49, 48]、Google TPU [53]、Amazon Trainium [60]、Microsoft Maia [59]、华为昇腾 [57]、阿里巴巴含光 [58] 与百度昆仑 [61]。

Capitalizing on this hardware proliferation, Meta deploys a heterogeneous fleet comprising NVIDIA GPUs, AMD GPUs, and proprietary MTIA accelerators at scale. This strategic diversification delivers several benefits: (1) enabling cost optimization by preserving the flexibility to adopt the most economical solutions as the market evolves; (2) strengthening supply-chain resilience; and (3) matching model characteristics to the most suitable hardware.

借着硬件供给扩大，Meta 大规模部署由 NVIDIA GPU、AMD GPU 与自研 MTIA 加速器组成的异构机队。这种战略多元化带来几方面收益：（1）保留随市场变化采用最经济方案的灵活性，从而优化成本；（2）增强供应链韧性；（3）把模型特性与最合适的硬件配对。

However, the extensive history of NVIDIA-specific optimizations hinders performance portability to other hardware. For example, fundamental differences in memory hierarchies and execution units (e.g., 32-thread warps vs. 64-thread wavefronts) thwart direct software migration between NVIDIA and AMD GPUs. These challenges are most acute in deep learning recommendation models (DLRMs) [10, 13]. Unlike the relatively uniform large language models (LLMs), DLRMs are characterized by diverse model architectures, fragmented compute kernel sets, irregular matrix dimensions, and complex memory access patterns.

然而，长期围绕 NVIDIA 的优化历史阻碍了性能向其他硬件的可移植性。例如，内存层次与执行单元的根本差异（如 32 线程的 warp 与 64 线程的 wavefront）使 NVIDIA 与 AMD GPU 之间难以直接迁移软件。这些挑战在深度学习推荐模型（DLRM）[10, 13] 上尤为尖锐。与相对规整的大语言模型（LLM）不同，DLRM 的特点是模型结构多样、计算核集合碎片化、矩阵维度不规则，以及复杂的访存模式。

While industry adoption of AMD GPUs has largely focused on inference—evidenced by Meta’s deployment of hundreds of thousands of AMD GPUs—distributed training presents a far greater technical challenge. This paper delivers the first hyperscaler report of enabling AMD GPUs for production-scale DLRM training. We demonstrate that through systematic optimization, AMD MI350X achieves performance comparable to the incumbent NVIDIA B200 GPU across a diverse suite of production DLRMs. We distill our production experience into three reusable principles for new hardware integration:

业界采用 AMD GPU 迄今多集中在推理——Meta 已部署数十万张 AMD GPU 即是例证——而分布式训练的技术难度远更高。本文给出首份超大规模厂商将 AMD GPU 用于生产级 DLRM 训练 的报告。我们证明：经系统优化后，AMD MI350X 在一系列生产 DLRM 上可达到与现役 NVIDIA B200 相当的性能。我们把生产经验提炼为三条可复用的新硬件接入原则：

Hardware-model co-design. Different accelerator microarchitectures offer distinct advantages and trade-offs. Rather than directly porting models from NVIDIA GPUs, hardware-model co-design—revising the model structure and software recipe—is essential to leverage an accelerator’s specific strengths. Specifically, we exploit MI350X’s $1.6\times$ HBM capacity relative to B200. To maximize training efficiency, we expand per-GPU batch sizes and reduce activation checkpointing while simultaneously optimizing data or model parallelism and implementing efficient embedding re-sharding. Beyond training throughput, we utilize the additional HBM to enhance model quality by enabling longer input sequences, larger hidden dimensions, and expanded embedding tables. This holistic approach delivers a $20\%$ training speedup and a $>$0.1% improvement in normalized entropy (NE) [17], demonstrating that architectural synergies can simultaneously drive improvements in both hardware efficiency and model quality.

硬件—模型协同设计。 不同加速器微架构各有优劣。与其把模型从 NVIDIA GPU 直接移植，不如做硬件—模型协同设计——修订模型结构与软件配方——才能吃到加速器特有优势。具体地，我们利用 MI350X 相对 B200 约 $1.6\times$ 的 HBM 容量：为提高训练效率，加大每卡 batch、减少激活检查点，同时优化数据并行或模型并行，并实现高效的 embedding 重分片。除吞吐外，多余 HBM 还用于提升模型质量：更长输入序列、更大隐层维度、更大 embedding 表。这一整体做法带来约 20% 的训练加速，以及归一化熵（NE）[17] 超过 0.1% 的提升，说明架构协同可以同时改进硬件效率与模型质量。

Analytical performance model. Traditionally, selecting high-performance general matrix multiplication (GEMM) kernels relies on offline tuning—a cumbersome process that generalizes poorly across the disparate matrix dimensions common in DLRMs. We propose Origami, an analytical performance model that predicts optimal tile configurations in microseconds, achieving 93% accuracy relative to exhaustively autotuned kernels across diverse matrix shapes. While analytical models are not new, Origami extends prior approaches to multi-chiplet architectures, with accuracy validated at production scale.

解析性能模型。 传统上，高性能通用矩阵乘（GEMM）核的选择依赖离线调优——过程繁琐，且难以泛化到 DLRM 常见、彼此悬殊的矩阵尺寸。我们提出 Origami：一种解析性能模型，可在微秒级预测最优 tile 配置，在多样矩阵形状上相对穷尽自动调优核达到 93% 准确率。解析模型并非新事物，但 Origami 把既有思路扩展到多 chiplet 架构，并在生产规模上验证了精度。

Agentic workflows for kernel optimization. To reduce engineering burden and migration friction, we codify complex kernel-tuning principles into agentic workflows: automating microarchitecture-aware configuration selection, swizzling process IDs, incorporating occupancy hints, and leveraging specialized hardware instructions. Together, they achieve $1.6\times$–$2\times$ speedups across both hand-written and compiler-generated kernels. Furthermore, we leverage agentic workflows to automate the application of solutions that address hardware-specific bottlenecks; e.g., to mitigate the performance impact of atomics on MI350X, we apply a tree-reduction algorithm that utilizes shared memory to deliver up to $4\times$ speedup in custom kernels.

面向核优化的智能体工作流。 为降低工程负担与迁移摩擦，我们把复杂的调核原则编成智能体工作流：自动做微架构感知的配置选择、对进程 ID 做 swizzle、加入 occupancy 提示，并利用专用硬件指令。合在一起，在手写与编译器生成的核上均可获得约 $1.6\times$–$2\times$ 加速。此外，我们用智能体工作流自动套用针对硬件瓶颈的方案；例如为缓解 MI350X 上原子操作的性能影响，采用利用共享内存的树归约，使定制核最高提速约 $4\times$。

Our methods have been validated at scale, supporting production training on tens of thousands of MI350X GPUs at Meta. Across a representative set of production models, MI350X achieves a geometric-mean throughput of 101% relative to our established B200 baseline—a platform highly optimized through extensive deployment. This performance parity validates the viability of our multi-vendor diversification strategy.

上述方法已在规模上得到验证，支撑 Meta 上数万张 MI350X 的生产训练。在一组有代表性的生产模型上，MI350X 相对我们已充分部署、高度优化的 B200 基线，几何均值吞吐约为 101%。这一性能持平，验证了多厂商多元化策略的可行性。

Limitations. While end-to-end performance is comparable, MI350X’s GEMM performance remains an area for improvement, due to both the overhead of emulated TF32 GEMMs in the absence of native hardware support and the relative maturity of NVIDIA cuBLAS over AMD hipBLAS in handling irregular matrix shapes.

局限。 尽管端到端性能相当，MI350X 的 GEMM 仍有提升空间：一是缺少原生 TF32、只能用仿真 TF32 GEMM 带来的开销；二是在不规则矩阵形状上，NVIDIA cuBLAS 相对 AMD hipBLAS 更成熟。

Although we use DLRMs to drive our optimizations, the underlying principles are generalizable to a broader spectrum of ML workloads, including LLMs. Indeed, we have successfully optimized LLM training on clusters of more than 2,000 MI350X GPUs; a comprehensive analysis of LLM-specific performance is reserved for future work.

虽以 DLRM 驱动优化，底层原则可推广到更广的 ML 负载，包括 LLM。事实上，我们已在超过 2,000 张 MI350X 的集群上成功优化 LLM 训练；针对 LLM 的系统分析留待后续工作。

To our knowledge, this is the first hyperscaler report on production-scale training using AMD GPUs. Our purpose is not to rank accelerators, but to share generalizable optimization methods for emerging hardware by synthesizing insights from our large-scale deployment of AMD, NVIDIA, and MTIA platforms. Our contributions have been upstreamed to the open-source PyTorch, Triton, and ROCm ecosystems.

据我们所知，这是首份关于用 AMD GPU 做生产级训练的超大规模报告。目的不是给加速器排名，而是综合我们在 AMD、NVIDIA 与 MTIA 平台上的大规模部署经验，分享面向新兴硬件的可推广优化方法。相关贡献已上游至开源 PyTorch、Triton 与 ROCm 生态。

II Background

II 背景

II-A DLRM Evolution

II-A DLRM 的演进

Fig. 1: Total training compute for deep learning models.

图 1：深度学习模型的总训练算力。

Deep learning recommendation models (DLRMs) [10] have become the canonical architecture for industrial recommendation and personalization systems across major technology companies, such as Meta [12], ByteDance [20], Google [18], and Alibaba [19]. DLRMs often comprise (1) a sparse component that uses embedding tables to map high-cardinality categorical signals—such as user IDs, item IDs, geographic regions, and language preferences—into dense vectors, and (2) a dense component that learns feature interactions to predict user engagement (e.g., click-through or conversion rates). The importance of this model family is reflected both externally—through the MLPerf recommendation benchmark evolution and the adoption of DLRM [11], DLRM-v2 [12], and DLRM-v3 [14, 13]—and internally through the scale of Meta’s production deployment and infrastructure investment dedicated to improving DLRM efficiency.

深度学习推荐模型（DLRM）[10] 已成为工业推荐与个性化系统的典型架构，被 Meta [12]、字节跳动 [20]、Google [18]、阿里巴巴 [19] 等主要科技公司广泛采用。DLRM 通常包含：（1）稀疏部分，用 embedding 表把用户 ID、物品 ID、地理区域、语言偏好等高基数类别信号映射为稠密向量；（2）稠密部分，学习特征交互以预测用户参与（如点击率或转化率）。这一模型族的重要性对外体现于 MLPerf 推荐基准的演进以及 DLRM [11]、DLRM-v2 [12]、DLRM-v3 [14, 13] 的采用；对内则体现于 Meta 为提升 DLRM 效率所投入的生产部署规模与基础设施。

Sparse evolution. Embedding tables grow rapidly in both vocabulary size and dimensionality, forcing distributed sharding and collective-heavy embedding exchange. As training clusters expanded, two practical directions emerged. First, multi-stage pipelining became essential [23]: embedding lookup, all-to-all exchange, and optimizer steps must be organized to reduce exposed latency and avoid synchronization bubbles. Second, sparse parallelism itself evolved beyond single dimension. Recent work on two-dimensional sparse parallelism [27] introduces a hierarchical strategy that adds data-parallel structure on top of model-parallel embedding sharding, reducing all-to-all scope, lowering embedding activation memory growth, and mitigating imbalance/straggler behavior. Together, these trends reflect a broader shift: the sparse subsystem is no longer “just embeddings,” but a distributed system whose parallelism strategy, pipeline structure, and optimizer semantics jointly determine scalability.

稀疏侧演进。 Embedding 表在词表规模与维度上增长迅速，迫使分布式分片，并带来以集合通信为主的 embedding 交换。随着训练集群扩大，出现两个务实方向。其一，多阶段流水线变得关键 [23]：embedding 查找、all-to-all 交换与优化器步骤需妥善组织，以减少暴露延迟、避免同步气泡。其二，稀疏并行本身超越单一维度：关于二维稀疏并行的近期工作 [27] 引入分层策略，在模型并行的 embedding 分片之上叠加数据并行结构，缩小 all-to-all 范围、降低 embedding 激活显存增长，并缓解不均衡/掉队。这些趋势共同表明：稀疏子系统已不再「只是 embedding」，而是一套分布式系统——其并行策略、流水线结构与优化器语义共同决定可扩展性。

Dense evolution. Classical DLRMs pair large embedding tables with relatively shallow dense components (e.g., a few fully connected layers for feature interaction). Recent evolution of DLRMs (Fig. 1) shows that scaling the dense subsystem can yield predictable quality gains analogous to LLM scaling laws [37, 38]. Wukong [15] demonstrates this by stacking factorization-machine layers with a synergistic upscaling strategy, achieving smooth quality improvements over orders of magnitude in dense-side FLOPs. A complementary trend is the adoption of Transformer-like modules to model long user-interaction histories. Meta’s HSTU [14] exemplifies this sequence-centric direction, replacing shallow pooling with attention over hundreds of past events. Together, these advances open a second axis for improving recommendation quality: in addition to scaling embedding tables, practitioners can now scale dense computation (FLOPs) to capture richer feature interactions and temporal patterns.

稠密侧演进。 经典 DLRM 把大 embedding 表与相对较浅的稠密部分配对（例如若干全连接层做特征交互）。DLRM 的近期演进（图 1）表明：扩展稠密子系统可以带来与 LLM 扩展定律 [37, 38] 类似的可预期质量收益。Wukong [15] 通过堆叠因子分解机层并配合协同升维策略证明了这一点，在稠密侧 FLOPs 跨数量级时仍获得平滑的质量提升。另一条互补趋势是采用类 Transformer 模块建模长用户交互历史。Meta 的 HSTU [14] 是这一序列中心方向的代表，用对上百次历史事件的注意力取代浅层池化。合在一起，这些进展为提升推荐质量打开了第二轴：除扩展 embedding 表外，实践者还可扩展稠密计算（FLOPs），以捕捉更丰富的特征交互与时序模式。

II-B AMD MI350X vs NVIDIA B200

II-B AMD MI350X 与 NVIDIA B200

Specification	MI350X	MI355X	B200
Architecture	CDNA 4	CDNA 4	Blackwell
Compute Units / Streaming Multiprocessors	256 (8 XCD $\times$ 32)	256 (8 XCD $\times$ 32)	148 (2 chiplets)
Peak FP16/BF16 Dense Performance	2300 TFLOPS	2500 TFLOPS	2250 TFLOPS
HBM Type / Capacity	HBM3e / 288 GB	HBM3e / 288 GB	HBM3e / 180 GB
HBM Bandwidth	Up to 8 TB/s	Up to 8 TB/s	Up to 8 TB/s
LDS per CU or SM	160 KB	160 KB	256 KB
Scale Up BW per GPU (uni-dir)	538 GB/s	538 GB/s	900 GB/s
Cooling	Air-cooled	Direct Liquid Cooled	Air-cooled
Deployed Power Limit	1000 W	1400 W	750 W
Deployment	Meta On-Prem	Public Cloud	Meta On-Prem

Table I: Comparison of AMD Instinct MI350X, MI355X, and NVIDIA Blackwell B200 GPUs. MI355X is a 1400 W design using direct liquid cooling, while Meta uses the 1000 W air-cooled MI350X. Meta’s deployment of B200 operates at 750 W, which is different from the general B200 HGX solution that operates at 1000 W. Meta operates both GPUs below their vendor-specified TDPs on the same air-cooled platform. B200 is capped at a lower power limit to enable air cooling because its power density is higher than that of MI350X.

规格	MI350X	MI355X	B200
架构	CDNA 4	CDNA 4	Blackwell
计算单元 / 流式多处理器	256（8 XCD $\times$ 32）	256（8 XCD $\times$ 32）	148（2 chiplet）
峰值 FP16/BF16 稠密算力	2300 TFLOPS	2500 TFLOPS	2250 TFLOPS
HBM 类型 / 容量	HBM3e / 288 GB	HBM3e / 288 GB	HBM3e / 180 GB
HBM 带宽	最高 8 TB/s	最高 8 TB/s	最高 8 TB/s
每 CU 或 SM 的 LDS	160 KB	160 KB	256 KB
每 GPU 纵向扩展带宽（单向）	538 GB/s	538 GB/s	900 GB/s
冷却	风冷	直接液冷	风冷
部署功耗上限	1000 W	1400 W	750 W
部署场景	Meta 自建机房	公有云	Meta 自建机房

表 I：AMD Instinct MI350X、MI355X 与 NVIDIA Blackwell B200 对比。MI355X 为 1400 W、直接液冷设计；Meta 使用的是 1000 W 风冷 MI350X。Meta 部署的 B200 运行在 750 W，不同于通常 1000 W 的 B200 HGX 方案。两款 GPU 均在同一风冷平台上低于厂商标称 TDP 运行；B200 因功率密度高于 MI350X，被压到更低功耗以支持风冷。

Table I compares the GPU specs. AMD MI350X is built on the CDNA™ 4 architecture and is the first AMD GPU used for production-scale training at Meta. MI350X features eight accelerator complex dies (XCDs) with 32 compute units (CUs) each (a total of 256 CUs), a 288 GB HBM3e memory pool with up to 8 TB/s bandwidth, and 538 GB/s of scale-up bandwidth per GPU. B200 offers 180 GB of HBM capacity, with 900 GB/s of scale-up bandwidth per GPU.

表 I 对比了 GPU 规格。AMD MI350X 基于 CDNA™ 4 架构，是 Meta 首款用于生产级训练的 AMD GPU。MI350X 含八个加速器复合裸片（XCD），各 32 个计算单元（CU），共 256 个 CU；288 GB HBM3e，带宽最高约 8 TB/s；每 GPU 纵向扩展带宽 538 GB/s。B200 的 HBM 容量为 180 GB，每 GPU 纵向扩展带宽 900 GB/s。

B200 and MI350X are same-generation flagships with closely matching theoretical bounds: their near-identical peak compute ($2250$ vs. $2300$ BF16 TFLOPS) and HBM bandwidth ($8$ TB/s) yield largely coincident roofline curves [43], and thus a similar performance envelope for ML workloads. Their strengths are complementary: B200 leads on per-GPU scale-up bandwidth ($900$ vs. $538$ GB/s), an all-to-all NVSwitch fabric (versus AMD’s point-to-point xGMI mesh), shared memory per SM ($256$ vs. $160$ KB), and native TF32 support, whereas MI350X’s principal advantages are its $1.6\times$ HBM capacity ($288$ vs. $180$ GB) and higher deployed power.

B200 与 MI350X 同属一代旗舰，理论上限接近：峰值算力几乎相同（$2250$ 对 $2300$ BF16 TFLOPS），HBM 带宽同为 $8$ TB/s，roofline 曲线大体重合 [43]，因而对 ML 负载的性能包络相近。优势互补：B200 在每 GPU 纵向扩展带宽（$900$ 对 $538$ GB/s）、全互联 NVSwitch 结构（相对 AMD 的点对点 xGMI 网格）、每 SM 共享内存（$256$ 对 $160$ KB）以及原生 TF32 上领先；MI350X 的主要优势则是约 $1.6\times$ 的 HBM 容量（$288$ 对 $180$ GB）以及更高的部署功耗上限。

Fig. 2: Hipification process to enable ML jobs on AMD GPUs.

图 2：使 ML 作业能在 AMD GPU 上运行的 Hipification 流程。

Hipification. Hipification is AMD’s CUDA-to-HIP (Heterogeneous-compute Interface for Portability) translation process that converts NVIDIA CUDA code into HIP equivalents; see Fig. 2. There are two primary approaches: (1) compile-time hipification using AMD’s open-source hipify-clang compiler, which translates CUDA constructs to HIP at build time, and (2) rule-based source-code translation using hipify_torch, which traverses and rewrites source files before compilation. While hipification provides broad functional compatibility and enables most CUDA workloads to run on AMD with minimal manual effort, performance portability remains a challenge. For instance, CUDA uses a warp size of 32, whereas an AMD GPU uses a wavefront size of 64. Kernels that naively hardcode warp sizes will likely perform suboptimally without tuning.

Hipification。 Hipification 是 AMD 将 CUDA 译为 HIP（Heterogeneous-compute Interface for Portability）的过程，把 NVIDIA CUDA 代码转为 HIP 等价物，见图 2。主要有两条路径：（1）编译期 hipification，用 AMD 开源的 hipify-clang 在构建时把 CUDA 构造译为 HIP；（2）基于规则的源码翻译，用 hipify_torch 在编译前遍历并改写源文件。Hipification 提供了广泛的功能兼容，使多数 CUDA 负载只需极少手工即可在 AMD 上跑起来，但性能可移植性仍是难题。例如 CUDA 的 warp 大小为 32，而 AMD GPU 的 wavefront 为 64；若核里天真地写死 warp 大小，不做调优往往表现不佳。

The emergence of agentic AI tools has fundamentally lowered the cost of switching from NVIDIA to AMD. At Meta, we leverage AI agents to address issues in the hipification process. These agents iteratively compile code, interpret error messages, and apply fixes until the build succeeds. When a niche CUDA library lacks a ROCm equivalent, AI agents can even generate custom shims to mimic the missing functionality. This makes AMD GPUs accessible to a broad range of ML researchers, who can now deploy high-performance code without mastering low-level hardware nuances.

智能体式 AI 工具的出现，从根本上降低了从 NVIDIA 迁到 AMD 的成本。在 Meta，我们用 AI 智能体处理 hipification 过程中的问题：它们迭代编译、解读报错并打补丁，直到构建成功。当小众 CUDA 库缺少 ROCm 对应物时，智能体甚至可以生成定制 shim 来模拟缺失功能。这使广大 ML 研究者不必精通底层硬件细节，也能部署高性能代码、用上 AMD GPU。

Meta’s Usage of AMD GPUs. Meta started using the AMD MI300X [51] for inference in 2024. Today, models compatible with AMD GPUs make up over 60% of production model types, weighted by resource footprint. Training DLRMs, on the other hand, presents a greater challenge than inference:

Meta 对 AMD GPU 的使用。 Meta 自 2024 年起将 AMD MI300X [51] 用于推理。如今，按资源占用加权，兼容 AMD GPU 的模型类型已超过生产模型类型的 60%。相对推理，训练 DLRM 的挑战更大：

• • Synchronization and network pressure: Training involves numerous GPUs in tight lockstep with frequent synchronization, heavily straining network efficiency.
• • Kernel and optimizer: Training requires comprehensive backward kernels and advanced optimizer support.
• • Generalization: While inference can be optimized per workload, training must rely on generalized optimizations for out-of-the-box performance—orders of magnitude more training trials run before a model is finalized, making per-job optimization infeasible.

• • 同步与网络压力： 训练涉及大量 GPU 紧密锁步、频繁同步，对网络效率压力极大。
• • 核与优化器： 训练需要完整的反向核与更强的优化器支持。
• • 泛化： 推理可按负载精调；训练则必须依赖开箱即用的通用优化——模型定稿前要跑数量级更多的训练试验，按作业精调不可行。

Addressing these challenges requires systematic optimizations across the software stack, from low-level libraries (Section III) and kernel tuning (Section IV) to high-level model co-design (Section V).

应对这些挑战，需要从底层库（第 III 节）、核调优（第 IV 节）到高层模型协同设计（第 V 节）的全栈系统优化。

III Roofline-Driven Optimization of ROCm Libraries

III 面向 Roofline 的 ROCm 库优化

The ROCm software platform provides optimized libraries, runtimes, and firmware for AMD GPUs. Grounded in roofline models [43] and hardware constraints, we optimize core GEMM kernels and selection algorithms to approach theoretical limits. To optimize RCCL, the collective communication library for distributed training, we derive theoretical bounds from hardware specifications as optimization targets. These roofline-driven principles enable low-level libraries to deliver robust, generalizable performance across diverse AI workloads.

ROCm 软件平台为 AMD GPU 提供优化库、运行时与固件。我们以 roofline 模型 [43] 与硬件约束为依据，优化核心 GEMM 核与选型算法，使之逼近理论极限。对分布式训练的集合通信库 RCCL，我们从硬件规格推导理论界作为优化目标。这些面向 roofline 的原则，使底层库能在多样 AI 负载上给出稳健、可推广的性能。

III-A GEMM Kernels for AMD MI350X

III-A 面向 AMD MI350X 的 GEMM 核

DLRM workloads present a broad range of matrix shapes and sizes, including irregular dimensions that stress standard GEMM kernel selection strategies. To deliver strong performance across this diversity, we integrate performance enhancements into hipBLASLt [5] GEMM kernels through its assembly backend, TensileLite [4]. These enhancements leverage new CDNA 4 [2] hardware capabilities alongside algorithmic refinements that reduce the solution space and enable finer-grained runtime parameter selection via Origami (Section III-B). Beyond exploiting new ISA features, we employ the following algorithmic techniques.

DLRM 负载涵盖广泛的矩阵形状与尺寸，其中包括对标准 GEMM 核选型策略构成压力的不规则维度。为在这种多样性上拿到强性能，我们通过 hipBLASLt [5] 的汇编后端 TensileLite [4]，把性能增强集成进 GEMM 核。这些增强既利用 CDNA 4 [2] 的新硬件能力，也包含缩小解空间、并借助 Origami（第 III-B 节）实现更细粒度运行时参数选择的算法改进。除挖掘新 ISA 特性外，我们还采用下列算法技术。

Stream-K partitions GEMM computations across CUs at finer granularity than traditional data-parallel tiling, improving load balance for tensor sizes that do not evenly tile the device. We also introduce new reduction strategies and runtime dispatch logic, improving robustness across a broad range of compute- and memory-bound tensor sizes. Notably, a single GEMM kernel can now achieve performance previously requiring multiple specialized kernels.

Stream-K 以比传统数据并行分块更细的粒度，把 GEMM 计算划分到各 CU，从而在不能整除设备的张量尺寸上改善负载均衡。我们还引入新的归约策略与运行时分派逻辑，提升在广泛的算力受限与访存受限张量尺寸上的稳健性。值得注意的是，如今单个 GEMM 核即可达到过去需要多个专用核才能达到的性能。

Workgroup mapping via space-filling curves. MI350X GPUs use a multi-XCD architecture in which workgroups are distributed across XCDs, each with a private L2 cache. Scheduling workgroups operating on adjacent or overlapping tiles across different XCDs degrades cache locality and efficiency. To address this, hipBLASLt introduces a mapping strategy based on space-filling curves that preferentially co-locates workgroups processing nearby tiles on the same XCD, improving spatial locality, cache utilization, and power efficiency. Multiple mapping strategies are supported within the GEMM assembly kernels, with runtime-selectable parameters integrated with Origami to select the optimal strategy for a given tensor size, kernel configuration, and target hardware.

基于空间填充曲线的工作组映射。 MI350X 采用多 XCD 架构，工作组分布在各 XCD 上，每个 XCD 有私有 L2 缓存。若把处理相邻或重叠 tile 的工作组调度到不同 XCD，会损害缓存局部性与效率。为此，hipBLASLt 引入基于空间填充曲线的映射策略，优先把处理邻近 tile 的工作组放在同一 XCD，以改善空间局部性、缓存利用率与能效。GEMM 汇编核支持多种映射策略，运行时可选参数与 Origami 集成，按给定张量尺寸、核配置与目标硬件选择最优策略。

III-B Origami: Analytical Kernel Selection

III-B Origami：解析式核选型

Prior kernel selection approaches in hipBLASLt relied on offline autotuning: enumerating kernel configurations on representative tensor shapes and storing the best selections in lookup tables. This approach fails at scale: recommendation models execute thousands of unique GEMM shapes per model, and maintaining exhaustive databases for every shape is impractical.

hipBLASLt 以往的核选型依赖离线自动调优：在代表性张量形状上枚举核配置，并把最优选择存入查找表。这一做法难以规模化：推荐模型每个模型会执行数千种独特的 GEMM 形状，为每种形状维护穷尽数据库不现实。

Origami replaces offline autotuning with an analytical performance model that predicts optimal tile configurations in microseconds. The model can be viewed as inferring the critical path of a GEMM problem across the memory and compute hierarchy. It captures the full tiling hierarchy—from matrix instruction dimensions through local data share (LDS) tiles to implicit cache tiles formed by workgroups sharing the L2 and last level cache (LLC)—and computes cycle counts geometrically at each level.

Origami 用解析性能模型取代离线自动调优，可在微秒级预测最优 tile 配置。该模型可视为在内存与计算层次上推断 GEMM 问题的关键路径。它刻画完整的分块层次——从矩阵指令维度，到本地数据共享（LDS）tile，再到共享 L2 与末级缓存（LLC）的工作组所形成的隐式缓存 tile——并在每一层几何地计算周期数。

Performance Model. We consider a GEMM operation $C=A\times B$, where $A\in\mathbb{R}^{M\times K}$, $B\in\mathbb{R}^{K\times N}$, and $C\in\mathbb{R}^{M\times N}$. To efficiently execute this operation on the GPU, the output matrix $C$ is partitioned into tiles, each computed by a workgroup. A workgroup iterates along the $K$ dimension, loading tiles of $A$ and $B$, performing multiply-accumulate operations, and writing the result back to global memory. The goal of the performance model is to select the tile configuration $(M_{T},N_{T},K_{T})$ that minimizes total execution latency.

性能模型。 考虑 GEMM 运算 $C=A\times B$，其中 $A\in\mathbb{R}^{M\times K}$，$B\in\mathbb{R}^{K\times N}$，$C\in\mathbb{R}^{M\times N}$。为在 GPU 上高效执行，输出矩阵 $C$ 被划分为若干 tile，各由一个工作组计算。工作组沿 $K$ 维迭代，加载 $A$ 与 $B$ 的 tile，做乘加，再写回全局内存。性能模型的目标是选择使总执行延迟最小的 tile 配置 $(M_{T},N_{T},K_{T})$。

For each candidate tile configuration with LDS tile size $(M_{T},N_{T},K_{T})$, Origami computes cycle counts for compute and memory operations, then selects the configuration minimizing total latency. The compute latency for one tile is:

对每个候选 tile 配置及其 LDS tile 尺寸 $(M_{T},N_{T},K_{T})$，Origami 计算计算与访存操作的周期数，再选择使总延迟最小的配置。单 tile 的计算延迟为：

$$L_{\text{compute}}=N_{\text{MI}}\cdot L_{\text{MI}} \tag{1}$$

where $N_{\text{MI}}=\lceil M_{T}/M_{I}\rceil\cdot\lceil N_{T}/N_{I}\rceil\cdot\lceil K_{T}/K_{I}\rceil$ is the number of matrix instructions per tile (macro tile dimensions divided by instruction dimensions) and $L_{\text{MI}}$ is the architecture-specific instruction latency.

其中 $N_{\text{MI}}=\lceil M_{T}/M_{I}\rceil\cdot\lceil N_{T}/N_{I}\rceil\cdot\lceil K_{T}/K_{I}\rceil$ 是每 tile 的矩阵指令数（宏 tile 维度除以指令维度），$L_{\text{MI}}$ 为架构相关的指令延迟。

Memory latency is computed by cascading traffic through the cache hierarchy. The bytes loaded per tile are:

访存延迟通过对缓存层次逐级级联流量来计算。每 tile 加载的字节数为：

$$\text{Bytes}_{\text{tile}}=M_{T}\cdot K_{T}\cdot b_{A}+N_{T}\cdot K_{T}\cdot b_{B} \tag{2}$$

where $b_{A},b_{B}$ are bytes per element. Origami predicts cache hit rates by modeling workgroup sharing within each memory scope. For $N_{\text{CU}}$ CUs sharing a cache, arranged as an $M_{C}\times N_{C}$ grid of output tiles, the predicted hit rate is:

其中 $b_{A},b_{B}$ 为每元素字节数。Origami 通过对各内存范围内的工作组共享建模来预测缓存命中率。对共享同一缓存的 $N_{\text{CU}}$ 个 CU，若排成 $M_{C}\times N_{C}$ 的输出 tile 网格，预测命中率为：

$$H=1-\frac{M_{C}\cdot M_{T}+N_{C}\cdot N_{T}}{M_{C}\cdot N_{C}\cdot(M_{T}+N_{T})} \tag{3}$$

This formula applies to L2 and LLC, with $N_{\text{CU}}$ set to the number of CUs sharing that cache scope. It captures data reuse: $A$ tiles are shared across $N_{C}$ workgroups computing different output columns, while $B$ tiles are shared across $M_{C}$ workgroups computing different rows. Tiles are rounded to the nearest cache line width in the contiguous dimension of the allocation (determined by transpose). The model constrains cache tile dimensions when working sets exceed cache capacity.

该公式适用于 L2 与 LLC，其中 $N_{\text{CU}}$ 取共享该缓存范围的 CU 数。它刻画数据复用：$A$ 的 tile 被计算不同输出列的 $N_{C}$ 个工作组共享，$B$ 的 tile 被计算不同行的 $M_{C}$ 个工作组共享。Tile 在分配的连续维上按最接近的缓存行宽取整（由转置决定）。当工作集超过缓存容量时，模型会约束缓存 tile 维度。

Traffic reaching each memory level determines hierarchy latency. Bytes cascade through miss rates: $\text{Bytes}_{\text{LLC}}=(1-H_{\text{L2}})\cdot\text{Bytes}_{\text{total}}$ and $\text{Bytes}_{\text{HBM}}=(1-H_{\text{LLC}})\cdot\text{Bytes}_{\text{LLC}}$, where each $H$ is computed using the CU count for that scope. Memory latency in cycles is:

到达每一级内存的流量决定该层延迟。字节数按未命中率级联：$\text{Bytes}_{\text{LLC}}=(1-H_{\text{L2}})\cdot\text{Bytes}_{\text{total}}$，$\text{Bytes}_{\text{HBM}}=(1-H_{\text{LLC}})\cdot\text{Bytes}_{\text{LLC}}$，其中每个 $H$ 用对应范围的 CU 数计算。以周期计的访存延迟为：

$$L_{\text{mem}}=\max\left(\frac{\text{Bytes}_{\text{L2}}}{R_{\text{L2}}},\;\frac{\text{Bytes}_{\text{LLC}}}{R_{\text{LLC}}},\;\frac{\text{Bytes}_{\text{HBM}}}{R_{\text{HBM}}}\right) \tag{4}$$

where $R$ is bytes per compute cycle at each level. Expressing memory bandwidth in bytes per compute cycle keeps the memory and compute latencies in the same unit, so the model need not reconcile the GPU’s separate memory and compute clock frequencies.

其中 $R$ 为每一级每计算周期的字节数。用「每计算周期字节数」表达内存带宽，可使访存与计算延迟同单位，模型不必再协调 GPU 上分离的内存时钟与计算时钟。

Issue Rate Bound. Effective memory throughput can be limited by the rate at which CUs issue loads and stores, not just bandwidth. This creates another parallelism-locality tradeoff on what are usually considered bandwidth-bound problems: smaller tiles increase the number of vector lanes issuing load/store requests but reduce data reuse, while larger tiles improve locality at the cost of issue rate. Origami models achievable throughput as a function of occupied vector lanes to navigate this tradeoff.

发射率上界。 有效内存吞吐不仅受带宽限制，还可能受 CU 发射 load/store 的速率限制。这在通常被视为带宽受限的问题上又带来并行度与局部性的权衡：更小的 tile 增加发射 load/store 的向量通道数，但降低数据复用；更大的 tile 改善局部性，却牺牲发射率。Origami 把可达吞吐建模为占用向量通道数的函数，以把握这一权衡。

Fig. 3: Pipelined GEMM execution. Loads/stores use vector units; compute uses matrix units. The K-loop overlaps load and compute.

图 3：流水线化的 GEMM 执行。Load/store 使用向量单元，计算使用矩阵单元。K 循环使加载与计算重叠。

Latency Model. Fig. 3 illustrates the pipeline structure: the prologue issues the initial load, the steady-state K-loop overlaps loads with compute, and the epilogue completes the final compute and store. Per-tile latency combines these phases:

延迟模型。 图 3 展示流水线结构：序幕发出首次加载，稳态 K 循环使加载与计算重叠，尾声完成最后的计算与写回。每 tile 延迟由这些阶段合成：

$$L_{\text{tile}}=L_{\text{prefetch}}+\bigl\lceil K/K_{T}\bigr\rceil\cdot\max(L_{\text{compute}},L_{\text{mem}})+L_{\text{store}} \tag{5}$$

The number of output tiles $T_{\text{out}}=\lceil M/M_{T}\rceil\cdot\lceil N/N_{T}\rceil$ determines timestep count $\omega=\lceil T_{\text{out}}/N_{\text{CU}}\rceil$. Total latency is $L_{\text{total}}=\omega\cdot L_{\text{tile}}$. Fig. 4 illustrates the tail effect: when tiles do not evenly divide across CUs, the final timestep underutilizes the GPU.

输出 tile 数 $T_{\text{out}}=\lceil M/M_{T}\rceil\cdot\lceil N/N_{T}\rceil$ 决定时间步数 $\omega=\lceil T_{\text{out}}/N_{\text{CU}}\rceil$。总延迟为 $L_{\text{total}}=\omega\cdot L_{\text{tile}}$。图 4 说明尾效应：当 tile 不能在各 CU 间整除时，最后一时间步会欠利用 GPU。

Fig. 4: Wave quantization tail effect. A 4$\times$5 output (20 tiles) on 8 CUs requires 3 timesteps with $20\bmod 8=4$ underutilized CUs in the final timestep. Origami addresses this via two strategies: tile selection to change the number of output tiles, and Stream-K to partition work along the K dimension.

图 4：波量化尾效应。在 8 个 CU 上，$4\times5$ 的输出（20 个 tile）需要 3 个时间步，最后一时间步有 $20\bmod 8=4$ 个 CU 欠利用。Origami 用两种策略应对：通过 tile 选择改变输出 tile 数，以及用 Stream-K 沿 K 维划分工作。

This static analysis is enabled by the regularity of GEMM workloads: uniform tile shapes and predictable access patterns. Origami evaluates all candidate configurations and selects the one minimizing $L_{\text{total}}$. Each candidate configuration evaluates in nanoseconds; the full search across all candidates completes in microseconds.

这种静态分析得益于 GEMM 负载的规整性：tile 形状一致、访存模式可预测。Origami 评估全部候选配置，选择使 $L_{\text{total}}$ 最小者。每个候选的评估在纳秒级完成；对全部候选的完整搜索在微秒级完成。

Grid Size Selection. Beyond tile configuration, Origami selects the optimal Stream-K grid size. Stream-K partitions work in the K dimension across CUs, and Origami balances three considerations: parallelism to maximize CU utilization, locality to enable data reuse across CUs iterating over distinct K slices, and scheduling to load balance across XCDs and minimize the wave quantization problem at each level of the chiplet hierarchy.

网格尺寸选择。 除 tile 配置外，Origami 还选择最优的 Stream-K 网格尺寸。Stream-K 沿 K 维把工作划分到各 CU，Origami 平衡三点：用并行度尽量吃满 CU；用局部性使遍历不同 K 切片的 CU 能复用数据；以及跨 XCD 调度负载均衡，并尽量缓解 chiplet 层次每一级上的波量化问题。

Chiplet-Aware Mapping. MI350X comprises 8 XCDs, each with private L2 cache. Origami extends hierarchical tiling to the chiplet level: a device-scope tile is composed of 8 XCD-scope tiles, each containing workgroups that share L2. XCD-aware workgroup ID swizzling maps spatially adjacent output tiles to the same XCD, improving L2 hit rates. Origami also selects the L2 tile shape dynamically based on the number of CUs per chiplet, arranging the concurrent workgroups into an $M_{C}\times N_{C}$ grid to maximize the cache hit rate predicted by Equation 3.

Chiplet 感知映射。 MI350X 含 8 个 XCD，各有私有 L2。Origami 把分层分块扩展到 chiplet 级：设备范围的 tile 由 8 个 XCD 范围 tile 组成，每个包含共享 L2 的工作组。XCD 感知的工作组 ID swizzle 把空间相邻的输出 tile 映射到同一 XCD，以提高 L2 命中率。Origami 还按每 chiplet 的 CU 数动态选择 L2 tile 形状，把并发工作组排成 $M_{C}\times N_{C}$ 网格，以最大化公式 3 预测的缓存命中率。

Results. Fig. 5 shows Origami performance on production GEMM shapes from recommendation models. The distribution exhibits classical roofline behavior: memory-bound at low arithmetic intensity and compute-bound at high intensity. Compared to exhaustively autotuned GEMM performance, Origami achieves 93% average selection efficiency across datatypes and transpose configurations. Origami is also integrated into tritonBLAS [8], enabling zero-warmup kernel selection for arbitrary matrix shapes at runtime, eliminating the need for offline autotuning entirely.

结果。 图 5 展示 Origami 在来自推荐模型的生产 GEMM 形状上的表现。分布呈经典 roofline 形态：低算术强度时访存受限，高算术强度时算力受限。相对穷尽自动调优的 GEMM 性能，Origami 在各数据类型与转置配置上平均选型效率为 93%。Origami 也已集成进 tritonBLAS [8]，可在运行时对任意矩阵形状做零热身核选型，从而完全省掉离线自动调优。

Generality. Although calibrated for MI350X, Origami’s formulation is general: because GEMM has regular tiling and predictable data movement, its cost reduces to a hardware-parameterized roofline. Retargeting another accelerator is therefore largely a matter of substituting a handful of hardware constants (cache sizes, achievable bandwidths, matrix-unit throughput, and chiplet topology) rather than re-deriving the model from scratch.

通用性。 虽按 MI350X 标定，Origami 的表述是通用的：由于 GEMM 分块规整、数据移动可预测，其代价可归结为带硬件参数的 roofline。因此迁移到另一加速器，只需替换少量硬件常数（缓存大小、可达带宽、矩阵单元吞吐与 chiplet 拓扑），而不必从零重推模型。

Fig. 5: GEMM performance on MI350X for production shapes, shown as percentage of peak achievable FLOPS.

图 5：生产形状在 MI350X 上的 GEMM 性能，以可达峰值 FLOPS 的百分比表示。

III-C GPU Power Excursion Boost

III-C GPU 功耗冲程提升

Fig. 6: GEMM performance ($M,N,K$ dimensions $=8192$) with GPU power excursion across kernel durations and duty ratios on MI350X.

图 6：MI350X 上 GEMM 性能（$M,N,K$ 维均为 8192）随核时长与占空比变化的 GPU 功耗冲程。

The Open Compute Project (OCP) is a collaborative initiative between hyperscale datacenter operators and technology vendors to standardize interfaces and specifications across infrastructure domains, including power delivery. The MI350X accelerator is implemented as a 1000 W open accelerator module (OAM) [3]. While the OAM must comply with the nominal total board power (TBP) over a 1s window, the specification permits transient power excursions over shorter intervals as defined by the OAM’s excursion design power (EDP). These excursions range from 1.1$\times$ TBP over 20 ms up to 2.0$\times$ TBP over 20 $\mu$s.

开放计算项目（OCP）是超大规模数据中心运营商与技术厂商之间的协作倡议，旨在跨基础设施领域标准化接口与规范，包括供电。MI350X 加速器实现为 1000 W 开放加速器模块（OAM）[3]。OAM 须在 1 秒窗口内遵守标称整板功耗（TBP），但规范允许在更短区间内出现由 OAM 冲程设计功耗（EDP）定义的瞬时功耗冲程。冲程范围从 20 ms 上的 1.1$\times$ TBP，到 20 $\mu$s 上的 2.0$\times$ TBP。

Although the OCP specification defines minimum electrical, thermal, and power-delivery capabilities, system designers may provision additional margin. Meta’s AI hardware platform, for example, supports average power levels of up to 1.36$\times$ TBP over a 50 ms window for a 1000 W OAM configuration. By dynamically managing transient power excursions, we leverage the thermal headroom of the MI350X to exceed steady-state power constraints. This allows compute-heavy kernels to surpass the nominal roofline limit, a strategy particularly effective for the hybrid execution profiles of DLRMs, where compute-intensive kernels are interleaved with lower-power operations.

OCP 规范定义的是电气、热与供电能力的下限，系统设计者可以预留额外余量。例如 Meta 的 AI 硬件平台，对 1000 W OAM 配置，可在 50 ms 窗口内支持最高约 1.36$\times$ TBP 的平均功耗。通过动态管理瞬时功耗冲程，我们利用 MI350X 的热余量突破稳态功耗约束，使算力密集核能越过标称 roofline 上限；这对 DLRM 的混合执行形态尤其有效——算力密集核与较低功耗操作交织出现。

To optimize excursion response under these extended capabilities, we developed a synthetic microbenchmark emulating heterogeneous execution phases. Fig. 6 measures performance boost by varying the GEMM duty ratio (percentage of high-power GEMM kernels interleaved with lower-power operations such as collectives or elementwise kernels), relative to steady-state operation at 1000 W with 100% back-to-back GEMMs. Depending on per-kernel duration, we achieve up to 1.33$\times$ performance gain. This improvement results from tuning the control-loop response to shape the boost profile, maximizing performance over the 50 ms window across varying phase periods and duty cycles while strictly respecting EDP limits.

为在这些扩展能力下优化冲程响应，我们开发了模拟异构执行阶段的合成微基准。图 6 通过改变 GEMM 占空比（与集合通信或逐元素核等较低功耗操作交织的高功耗 GEMM 核占比），相对 1000 W、100% 背靠背 GEMM 的稳态运行测量性能提升。视每核时长，最高可获约 1.33$\times$ 性能增益。该改进来自调控制环响应以塑造 boost 曲线，在变化的阶段周期与占空比下最大化 50 ms 窗口内性能，同时严格遵守 EDP 限制。

In end-to-end applications, the benefit from higher excursion limits depends on the frequency and duration of compute-intensive phases, intervening low-power operations, and sensitivity to clock frequency. In Meta’s DLRM production training, this enhancement resulted in end-to-end performance gains of 3–6%.

在端到端应用中，更高冲程上限带来的收益取决于算力密集阶段的频率与时长、其间的低功耗操作，以及对时钟频率的敏感度。在 Meta 的 DLRM 生产训练中，这一增强带来约 3–6% 的端到端性能提升。

III-D Collective Communication with RCCL

III-D 用 RCCL 做集合通信

RCCL provides collective communication primitives for AMD GPUs, supporting PCIe, Infinity Fabric Interconnect [6], and network fabrics (InfiniBand, TCP/IP) across single and multi-node deployments.

RCCL 为 AMD GPU 提供集合通信原语，支持单节点与多节点部署下的 PCIe、Infinity Fabric Interconnect [6]，以及网络结构（InfiniBand、TCP/IP）。

We developed a theoretical performance projection model derived from hardware-level simulations of the memory subsystem, scale-up fabric, and NIC data paths[^1]. Guided by these projections, we co-designed a suite of RCCL optimizations across the RDMA transport, GPU kernel, and protocol layers for the MI350X. These refinements narrowed the communication gap from $2\times$ slower than NCCL to within 10–15% of NCCL performance, enabling production-scale DLRM training. Fig. 7 shows measured AllToAll bandwidth of different message sizes against performance upper bound.

我们基于对内存子系统、纵向扩展互连与 NIC 数据路径的硬件级仿真，建立了理论性能投影模型[^1]。在这些投影指导下，我们针对 MI350X 在 RDMA 传输、GPU 核与协议层共同设计了一系列 RCCL 优化。这些改进把相对 NCCL 约慢 $2\times$ 的通信差距收窄到与 NCCL 相差约 10–15%，使生产级 DLRM 训练成为可能。图 7 展示不同消息大小下测得的 AllToAll 带宽相对性能上界。

Fig. 7: Measured AllToAll bandwidth with different group and message sizes relative to the theoretical upper bounds.

图 7：不同分组与消息大小下测得的 AllToAll 带宽相对理论上界。

Transport and protocol refinements. To resolve a performance conflict between collective types, we adopted collective-specific RDMA queue pair (QP) scaling: AllToAll utilizes a single QP for large-chunk efficiency and 30% higher bus bandwidth, while ring-based collectives (e.g., AllGather) retain four QPs to saturate network bandwidth. We replaced the LL128 low-latency protocol—whose 128-byte stores split across AMD’s 64-byte cache lines—with LL64, a protocol that packs 56-byte payload with 8-byte flags into one cache-line-atomic 64-byte packet, delivering $2\times$ AllGather/ReduceScatter and $1.2\times$–$1.4\times$ AllReduce speedups.

传输与协议改进。 为解决不同集合类型之间的性能冲突，我们采用按集合类型区分的 RDMA 队列对（QP）扩缩：AllToAll 使用单个 QP 以获得大块效率与约 30% 更高的总线带宽；基于环的集合（如 AllGather）保留四个 QP 以吃满网络带宽。我们用 LL64 取代 LL128 低延迟协议——后者 128 字节存储会跨 AMD 的 64 字节缓存行被拆分——LL64 把 56 字节载荷与 8 字节标志打进一个缓存行原子的 64 字节包，带来约 $2\times$ 的 AllGather/ReduceScatter 以及约 $1.2\times$–$1.4\times$ 的 AllReduce 加速。

Algorithmic and routing enhancements. We introduced Direct AllGather to replace multi-hop ring transfers with simultaneous point-to-point sends, yielding a $1.4\times$ speedup for small messages. Our P2P batching design reduced AllToAll latency by 40% at a 128-GPU scale through the aggregation of concurrent communications.

算法与路由增强。 我们引入 Direct AllGather，用同时点对点发送取代多跳环传输，对小消息获得约 $1.4\times$ 加速。P2P 批处理设计通过聚合并发通信，在 128 GPU 规模下把 AllToAll 延迟降低约 40%。

GPU resource and occupancy optimization. A key challenge in production DLRM training is maximizing compute-communication overlap while minimizing contention in shared compute and memory resources. To mitigate CU contention, we transitioned from a thread-block-per-GPU to a warp-per-GPU communication model, capped active channels at 48 (down from 64). By increasing the thread block size from 256 to 512, we achieved a 46% VGPR (vector general-purpose register) reduction and doubled occupancy. These optimizations delivered gains of 63% for ReduceScatter and 48% for AllReduce, while reducing the performance penalty when operating with restricted CU allocations.

GPU 资源与占用率优化。 生产 DLRM 训练的关键挑战是：在最大化计算—通信重叠的同时，尽量减少共享计算与内存资源上的争用。为缓解 CU 争用，我们从「每 GPU 一个线程块」的通信模型转为「每 GPU 一个 warp」，并把活跃通道上限从 64 降到 48。通过把线程块大小从 256 增到 512，VGPR（向量通用寄存器）减少约 46%，占用率翻倍。这些优化使 ReduceScatter 提升约 63%、AllReduce 提升约 48%，并减轻在受限 CU 分配下运行时的性能损失。

III-E Shampoo Optimizer and rocSOLVER

III-E Shampoo 优化器与 rocSOLVER

Shampoo [34, 36] is a second-order optimizer widely deployed due to its faster convergence compared to first-order methods like Adam [35]. Shampoo maintains Kronecker-factored preconditioner matrices, requiring periodic computation of matrix fourth roots to update the preconditioners. This matrix root computation relies on eigen decomposition provided by the rocSOLVER library. For large models, these steps dominate Shampoo’s overhead, making rocSOLVER’s performance critical to overall training throughput.

Shampoo [34, 36] 是广泛部署的二阶优化器，相对 Adam [35] 等一阶方法收敛更快。Shampoo 维护 Kronecker 分解的预条件矩阵，需要周期性地计算矩阵四次方根以更新预条件子。该矩阵根计算依赖 rocSOLVER 库提供的特征分解。对大模型而言，这些步骤主导 Shampoo 开销，使 rocSOLVER 性能对整体训练吞吐至关重要。

The de facto algorithm to compute the eigen decomposition of a symmetric input matrix $A$ has three steps: (1) reduce $A$ into tridiagonal form, yielding an orthogonal $U$ and a tridiagonal $T$ such that $A=UTU^{T}$; (2) compute eigenvalues $\Lambda$ and eigenvectors $V$ of $T$; and (3) backtransform eigenvectors $Q=UV$ so that $A=Q\Lambda Q^{T}$. In rocSOLVER, reduction to tridiagonal form is accomplished with block Householder transformations [62] operating on panels of the input matrix; updates to the rest of the matrix are applied once per panel as a pair of GEMM operations. Step (2), based on Cuppen’s Divide-and-Conquer (D.C.) [63], is the most complex step and the focus of our optimizations.

计算对称输入矩阵 $A$ 特征分解的事实标准算法分三步：（1）把 $A$ 化约到三对角形式，得到正交阵 $U$ 与三对角阵 $T$，满足 $A=UTU^{T}$；（2）计算 $T$ 的特征值 $\Lambda$ 与特征向量 $V$；（3）反变换特征向量 $Q=UV$，使 $A=Q\Lambda Q^{T}$。在 rocSOLVER 中，化约到三对角用作用在输入矩阵面板上的分块 Householder 变换 [62] 完成；对其余矩阵的更新按面板以一对 GEMM 施加。第（2）步基于 Cuppen 分治（D.C.）[63]，最复杂，也是我们优化的重点。

Cuppen’s D.C. splits $T$ into sub-matrices that are processed independently and then merged in pairs. Typically, D.C. operates on unreduced sub-blocks of $T$ (i.e., blocks separated by zero off-diagonal elements), but this can lead to load unbalancing: sub-optimal CU allocation often follows because the size of sub-blocks is input-dependent. Our implementation in rocSOLVER lifts this restriction, yielding precise control over CU allocation for improved performance. During a merge step, D.C. identifies deflated eigenvalues (which will not be modified by the merge) and updates all others by solving the secular equation. To improve data locality and reduce data movement, our implementation ensures that eigenvalues (and associated eigenvectors) are properly ordered. Eigenvectors belonging to identical eigenvalues require (possibly costly) extra preprocessing; rocSOLVER includes specialized code to efficiently handle matrices exhibiting several identical eigenvalues in merge steps. Lastly, newly and previously computed eigenvectors are merged using batched GEMMs for optimal performance.

Cuppen 的 D.C. 把 $T$ 拆成可独立处理的子矩阵，再成对合并。通常 D.C. 作用在 $T$ 的未化约子块上（即被零次对角元素隔开的块），但这会导致负载不均：子块大小依赖输入，CU 分配往往不佳。我们在 rocSOLVER 中的实现放开这一限制，从而精确控制 CU 分配以提升性能。合并步骤中，D.C. 识别不会被合并修改的消元特征值，并通过求解世俗方程更新其余特征值。为改善数据局部性、减少数据移动，我们的实现保证特征值（及相应特征向量）有恰当次序。属于相同特征值的特征向量需要（可能昂贵的）额外预处理；rocSOLVER 含专用代码，在合并步骤中高效处理含多个相同特征值的矩阵。最后，新算与已算的特征向量用批处理 GEMM 合并，以获得最优性能。

Fig. 8: Eigen decomposition performance.

图 8：特征分解性能。

Fig. 8 demonstrates that rocSOLVER achieves comparable eigen decomposition performance on AMD MI350X relative to NVIDIA B200. Distributed Shampoo partitions large preconditioner matrices into fixed-size blocks (typically $4096\times 4096$ or $8192\times 8192$) to bound per-GPU computation and enable parallel eigen decomposition across devices. These block sizes represent the practical sweet spot: large enough to preserve preconditioning quality, yet small enough for tractable eigen decomposition. rocSOLVER narrows the gap with cuSOLVER at these critical 4K and 8K dimensions, enabling competitive Shampoo-based training on AMD hardware.

图 8 表明，rocSOLVER 在 AMD MI350X 上的特征分解性能与 NVIDIA B200 相当。分布式 Shampoo 把大预条件矩阵划成固定大小块（通常 $4096\times 4096$ 或 $8192\times 8192$），以限制每 GPU 计算量，并在设备间并行做特征分解。这些块大小是务实甜点：足够大以保持预条件质量，又足够小使特征分解可处理。rocSOLVER 在这些关键的 4K 与 8K 维度上收窄了与 cuSOLVER 的差距，使基于 Shampoo 的训练在 AMD 硬件上具有竞争力。

IV Kernel Tuning and Agentic AI Workflows

IV 核调优与智能体 AI 工作流

In this section, we first identify kernel optimization principles through case studies of key production kernels, and then encode them as reusable skills in an agentic AI workflow that scale to the long tail of model-specific operators.

本节先通过关键生产核的案例研究提炼核优化原则，再把它们编码为智能体 AI 工作流中的可复用技能，以覆盖长尾的模型专用算子。

IV-A Triton Kernel Optimizations

IV-A Triton 核优化

IV-A1 Attention Kernels in DLRM

We focus on two important Triton kernels in Meta’s production DLRM workloads, GDPA and HSTU, and describe the AMD-specific optimization techniques that yield substantial speedups on MI350X.

IV-A1 DLRM 中的注意力核

我们聚焦 Meta 生产 DLRM 负载中两个重要的 Triton 核——GDPA 与 HSTU——并说明在 MI350X 上带来显著加速的 AMD 专用优化技术。

Generalized dot-product attention (GDPA). GDPA [33] extends standard attention [32, 31] by fusing GeLU-activated attention and a feed-forward network into a single operator. Central to Meta’s Kunlun [21] and Interformer [22] models, the fused kernel substantially increases the compute and memory intensity, positioning the kernel as a critical optimization target.

广义点积注意力（GDPA）。 GDPA [33] 扩展标准注意力 [32, 31]，把经 GeLU 激活的注意力与前馈网络融为单一算子。它是 Meta Kunlun [21] 与 Interformer [22] 模型的核心，融合核大幅提高算力与访存强度，因而成为关键优化目标。

Hierarchical sequential transduction unit (HSTU). HSTU [14] is a novel sequential attention architecture designed for large-scale recommendation systems, where the input consists of variable-length user interaction histories with high sparsity. Unlike standard attention, HSTU uses variable-length batching and processes sequences with sparsity ratios as high as 0.95, meaning the effective computation is concentrated in a small fraction of the sequence.

分层序列转导单元（HSTU）。 HSTU [14] 是面向大规模推荐系统的新型序列注意力架构，输入为高稀疏、可变长度的用户交互历史。与标准注意力不同，HSTU 使用可变长度批处理，并处理稀疏比可高达 0.95 的序列，即有效计算集中在序列的一小部分。

(a) Forward latency (ms).

（a）前向延迟（ms）。

(a) Forward latency (ms).

（a）前向延迟（ms）。

IV-A2 Optimization Techniques

We apply five optimization techniques to the Triton kernels, each addressing a different performance bottleneck on the MI350X architecture. While we use GDPA and HSTU as examples, these optimization techniques are generally applicable to a broad set of Triton kernels for ML models.

IV-A2 优化技术

我们对 Triton 核应用五项优化技术，各自针对 MI350X 架构上的不同性能瓶颈。虽以 GDPA 与 HSTU 为例，这些技术通常适用于广泛的 ML 模型 Triton 核。

PID swizzle. MI350X features a multi-XCD (chiplet) architecture where eight XCDs share an L2 cache. The default linear mapping of Triton program IDs to tiles can cause XCDs to access disjoint L2 cache regions, resulting in poor cache utilization. PID swizzle uses tl.swizzle2d to remap the two-dimensional program ID grid so that spatially adjacent tiles are scheduled on the same XCD, improving L2 cache locality.

PID swizzle。 MI350X 为多 XCD（chiplet）架构，八个 XCD 共享 L2 缓存。Triton 程序 ID 到 tile 的默认线性映射会使各 XCD 访问互不相交的 L2 区域，缓存利用率差。PID swizzle 用 tl.swizzle2d 重映射二维程序 ID 网格，使空间相邻的 tile 调度到同一 XCD，从而改善 L2 局部性。

Occupancy hints. The waves_per_eu hint controls the maximum number of wavefronts resident on each execution unit, which directly affects register allocation. By setting waves_per_eu=3 (GDPA) or exposing it as an autotunable parameter with values 0–4 (HSTU), the compiler reduces per-wavefront register usage, allowing more concurrent wavefronts to hide memory latency. This trades per-wavefront performance for higher occupancy, which is beneficial when the kernel is not fully compute-bound.

Occupancy 提示。 waves_per_eu 提示控制每个执行单元上驻留的最大 wavefront 数，直接影响寄存器分配。将 waves_per_eu 设为 3（GDPA），或作为 0–4 可自动调优参数暴露（HSTU），编译器会降低每 wavefront 寄存器用量，使更多并发 wavefront 能掩盖访存延迟。这是用单 wavefront 性能换更高占用率，在核未完全算力受限时有益。

MFMA instruction configuration. AMD CDNA architectures provide matrix fused multiply-add (MFMA) instructions in configurable tile sizes, controlled by the matrix_instr_nonkdim parameter. By setting matrix_instr_nonkdim=16, the system selects $16\times 16$ MFMA tiles, which on CDNA 4 can achieve higher instruction throughput than the default $32\times 32$ tiles for certain matrix shapes.

MFMA 指令配置。 AMD CDNA 架构提供可配置 tile 大小的矩阵融合乘加（MFMA）指令，由 matrix_instr_nonkdim 参数控制。设 matrix_instr_nonkdim=16 时，系统选择 $16\times 16$ 的 MFMA tile；在 CDNA 4 上，对某些矩阵形状可比默认 $32\times 32$ tile 获得更高指令吞吐。

Triton configuration autotuning. Triton kernel configurations (BLOCK_SIZE, num_stages, num_warps, etc.) are critical to performance, and optimal values differ between NVIDIA and AMD GPUs due to architectural differences. AMD’s 64-thread wavefronts, compared to NVIDIA’s 32-thread warps, favor larger num_warps values to achieve sufficient occupancy and hide memory latency. NVIDIA’s Tensor Memory Accelerator (TMA) on H100/B200 enables deep software pipelining (e.g., num_stages $\in[4,7]$) by efficiently managing asynchronous memory transfers. AMD MI350X lacks a TMA equivalent; high stage counts increase LDS pressure without commensurate latency hiding benefits. Empirically, AMD kernels achieve better performance with shallower pipelines (e.g., num_stages $\approx 2$).

Triton 配置自动调优。 Triton 核配置（BLOCK_SIZE、num_stages、num_warps 等）对性能关键，且因架构差异，NVIDIA 与 AMD GPU 上的最优值不同。相对 NVIDIA 的 32 线程 warp，AMD 的 64 线程 wavefront 更倾向更大的 num_warps，以获得足够占用率并掩盖访存延迟。H100/B200 上的 NVIDIA Tensor Memory Accelerator（TMA）通过高效管理异步内存传输，支持深软件流水线（如 num_stages $\in[4,7]$）。AMD MI350X 没有等价 TMA；过高的 stage 数会增加 LDS 压力，却换不来相应的延迟掩盖收益。经验上，AMD 核在更浅流水线（如 num_stages $\approx 2$）上表现更好。

GeLU fast math. The default fast_gelu implementation in Triton uses a tanh-based approximation, which on AMD GPUs decomposes into $\tanh(x)=2\sigma(2x)-1$, requiring a sigmoid evaluation (scaled by two) and additional arithmetic. We replace this with a direct sigmoid-based GeLU approximation using the compiler intrinsics fast_expf and fast_dividef, which map to single hardware instructions with better performance.

GeLU 快速数学。 Triton 默认的 fast_gelu 使用基于 tanh 的近似；在 AMD GPU 上会分解为 $\tanh(x)=2\sigma(2x)-1$，需要一次（乘以二的）sigmoid 求值与额外算术。我们改用基于 sigmoid 的直接 GeLU 近似，并借助编译器内建 fast_expf 与 fast_dividef，它们映射为性能更好的单条硬件指令。

IV-A3 Performance Results

We evaluate these tuning techniques by incrementally applying them to the GDPA and HSTU forward and backward kernels.

IV-A3 性能结果

我们通过逐步把这些调优技术应用到 GDPA 与 HSTU 的前向与反向核上来评估效果。

GDPA. Fig. 9(a) and 9(b) show the incremental forward kernel improvements based on a representative configuration[^2] in GEM models [16]. Starting from the baseline of 1.47 ms latency (58.4 TFLOP/s), the cumulative application of GeLU fast math, occupancy hints, MFMA instruction configuration, and PID swizzle reduces latency to 0.59 ms (145.9 TFLOP/s), a total 60.0% reduction in latency and $2.5\times$ in throughput. Fig. 9(c) and 9(d) show the backward kernel results. The same techniques (except matrix_instr_nonkdim, which regresses the backward pass) reduce latency from 3.67 ms to 2.17 ms, a 40.9% reduction. Compared to B200, the optimized GDPA kernels on MI350X achieve better performance.

GDPA。 图 9(a) 与 9(b) 展示基于 GEM 模型 [16] 中代表性配置[^2] 的前向核逐步改进。从基线延迟 1.47 ms（58.4 TFLOP/s）起步，累计应用 GeLU 快速数学、occupancy 提示、MFMA 指令配置与 PID swizzle 后，延迟降至 0.59 ms（145.9 TFLOP/s），延迟共降 60.0%，吞吐约 $2.5\times$。图 9(c) 与 9(d) 为反向核结果。同样技术（除对反向有回退的 matrix_instr_nonkdim 外）把延迟从 3.67 ms 降到 2.17 ms，降幅 40.9%。相对 B200，MI350X 上优化后的 GDPA 核表现更好。

HSTU. Fig. 10(a) and 10(b) show the forward kernel results for a representative configuration[^3] in production models. Block size tuning, occupancy hints, and PID swizzle together yield a $3.22\times$ speedup, with peak throughput improving from 45.0 to 144.8 TFLOP/s. Fig. 10(c) and 10(d) show the backward kernel results. Out-of-the-box performance is already stronger on MI350X here, leaving less optimization headroom: MFMA parameter tuning delivers a $1.21\times$ speedup, with peak throughput improving from 60.2 to 72.6 TFLOP/s. Averaged across sequence lengths, the optimized MI350X kernels are $1.22\times$ faster than B200 on the forward pass and $1.80\times$ faster on the backward pass, for a combined $1.69\times$ advantage that widens at longer sequences.

HSTU。 图 10(a) 与 10(b) 展示生产模型中代表性配置[^3] 的前向核结果。块大小调优、occupancy 提示与 PID swizzle 合计带来约 $3.22\times$ 加速，峰值吞吐从 45.0 升至 144.8 TFLOP/s。图 10(c) 与 10(d) 为反向核：此处开箱性能在 MI350X 上已较强，优化空间较小；MFMA 参数调优带来约 $1.21\times$ 加速，峰值吞吐从 60.2 升至 72.6 TFLOP/s。按序列长度平均，优化后的 MI350X 核前向比 B200 快约 $1.22\times$、反向快约 $1.80\times$，合计约 $1.69\times$ 优势，且随序列变长而扩大。

IV-B CUDA Kernel Optimizations

IV-B CUDA 核优化

GPU kernels written in CUDA can be compiled for AMD GPUs via HIP translation (Section II-B). While it enables code portability, certain CUDA idioms, particularly those relying on atomic operations, perform poorly on AMD’s multi-chiplet architecture. This section describes the architectural differences that cause these performance gaps and our algorithmic strategies to address them.

用 CUDA 编写的 GPU 核可通过 HIP 翻译编译到 AMD GPU（第 II-B 节）。虽能实现代码可移植，但某些 CUDA 惯用法——尤其依赖原子操作的——在 AMD 多 chiplet 架构上表现不佳。本节说明造成这些性能差距的架构差异，以及我们应对它们的算法策略。

IV-B1 Atomic Operations on Multi-Chiplet GPUs

Unlike NVIDIA GPUs where atomic operations complete in a monolithic L2 cache, MI350X atomics may require cross-chiplet cache coherency (note 8 XCD chiplets in MI350X), incurring round-trip latencies across the Infinity Fabric. When thousands of threads contend on a single address (as is common in reduction-style kernels), these coherency messages serialize execution, leading to poor performance.

IV-B1 多 Chiplet GPU 上的原子操作

与在整块 L2 缓存中完成原子操作的 NVIDIA GPU 不同，MI350X 的原子操作可能需要跨 chiplet 的缓存一致性（MI350X 有 8 个 XCD chiplet），从而在 Infinity Fabric 上产生往返延迟。当数千线程争用同一地址时（归约类核常见），这些一致性消息会使执行串行化，性能很差。

The fundamental issue is that per-thread global atomics scale poorly with thread count on multi-chiplet architectures. On a monolithic GPU, the L2 cache can coalesce nearby atomic operations; on a multi-chiplet GPU, each chiplet must coordinate with the others, and a single contended address becomes a serialization bottleneck. Our measurements show that MI350X can be over $10\times$ slower than B200 on atomic-heavy kernels at high contention rates, while matching or exceeding B200 on atomic-free, bandwidth-bound kernels.

根本问题是：按线程的全局原子操作在多 chiplet 架构上随线程数扩展不佳。在整块 GPU 上，L2 可合并邻近原子操作；在多 chiplet GPU 上，各 chiplet 必须彼此协调，单个争用地址成为串行瓶颈。测量表明：在高争用率的原子密集核上，MI350X 可比 B200 慢超过 $10\times$；而在无原子、带宽受限的核上则可持平或超过 B200。

Algorithmic strategies: hierarchical reduction and caching. We employ two complementary strategies depending on the contention pattern. When many threads contend on a single shared address (e.g., a global counter), we apply hierarchical reduction: threads accumulate results in fast on-chip memory (registers or shared memory), a parallel tree reduction sums them within each block, and a single representative thread issues one global atomic. This converts $N$ global atomics into $\lceil N/\text{block\_size}\rceil$, a reduction factor of up to $1024\times$. When threads target many distinct addresses with temporal locality (e.g., gradient scatter to embedding rows), we apply caching: each thread maintains a small register cache of recently written addresses, merging consecutive updates to the same address locally before issuing a single atomic add on eviction. The tradeoff is additional shared memory usage. For kernels where contention is rare, the overhead may slightly increase latency, but the benefit grows superlinearly with contention rate, making this transformation highly effective for production workloads with moderate-to-high atomic traffic.

算法策略：分层归约与缓存。 我们按争用模式采用两种互补策略。当许多线程争用单个共享地址（如全局计数器）时，用分层归约：线程在快速片上内存（寄存器或共享内存）中累加，块内并行树归约求和，再由一个代表线程发一次全局原子。这把 $N$ 次全局原子变为 $\lceil N/\text{block\_size}\rceil$，最高可降约 $1024\times$。当线程以时间局部性瞄准许多不同地址时（如梯度散射到 embedding 行），用缓存：每线程维护最近写入地址的小寄存器缓存，在本地合并对同一地址的连续更新，驱逐时再发一次原子加。代价是额外共享内存。对争用罕见的核，开销可能略增延迟，但收益随争用率超线性增长，因而对中高原子流量的生产负载很有效。

Fig. 11: Bounds check forward kernel latency vs. violation rate. Tree reduction eliminates atomic contention, enabling MI350X to outperform B200 at high violation rates.

图 11：越界检查前向核延迟随违约率变化。树归约消除原子争用，使 MI350X 在高违约率下优于 B200。

IV-B2 Example: Bounds Check Kernel

The kernel validates that all embedding table indices fall within valid ranges before performing lookups. When operating in WARNING mode, the kernel (1) accumulates a per-thread count of out-of-bounds violations, (2) corrects each invalid index, and (3) increments the global out-of-bounds violation counter via atomic add operations. On MI350X, atomic operations become severe bottlenecks at high violation rates.

IV-B2 示例：越界检查核

该核在查找前校验所有 embedding 表索引是否落在有效范围内。在 WARNING 模式下，核会：（1）按线程累计越界违约次数；（2）修正每个非法索引；（3）通过原子加递增全局越界违约计数器。在 MI350X 上，高违约率时原子操作成为严重瓶颈。

Original implementation. After each thread accumulates its violation count, threads with non-zero counts issue gpuAtomicAdds to global memory. With 1024 threads per block and thousands of blocks processing millions of indices, a typical workload can produce millions of threads that contend on a single address.

原始实现。 每线程累计违约次数后，非零线程向全局内存发 gpuAtomicAdd。每块 1024 线程、数千块处理数百万索引时，典型负载可产生数百万线程争用同一地址。

Optimized implementation. In order to reduce atomic operations, per-thread global adds are replaced with shared-memory parallel tree reduction followed by a single atomic add. Rather than writing to a global counter directly, each thread writes its local violation count to a shared memory buffer allocated at the block level. A 10-step binary reduction tree then sums the per-thread contributions entirely within shared memory, avoiding any cross-XCD traffic. Finally, thread 0 issues a single gpuAtomicAdd with the block-wide sum to the global counter, reducing global atomics from $\sim$ $10^{6}$ to $\sim$ $10^{3}$ per batch for a typical production workload.

优化实现。 为减少原子操作，用共享内存并行树归约加单次原子加取代按线程的全局加。各线程不再直接写全局计数器，而是把本地违约次数写入块级共享内存缓冲；再经 10 步二叉归约树在共享内存内完全求和，避免任何跨 XCD 流量。最后由线程 0 把块级和以一次 gpuAtomicAdd 写到全局计数器，典型生产负载下把全局原子从约 $10^{6}$ 降到约 $10^{3}$ 每次 batch。

Results. Fig. 11 shows the latency on a production embedding configuration (30 tables, 13.8M indices) as violation rate varies from 0% to 20%. With per-thread atomic adds, MI350X forward latency degrades $8.5\times$ from 0% to 20% violations while B200 degrades only $3.3\times$. At 10% violations, tree reduction delivers $4.26\times$ speedup (2.1 ms vs 9.0 ms). At 20% violations, the speedup reaches $4.34\times$, and MI350X with tree reduction is $1.75\times$ faster than B200.

结果。 图 11 展示生产 embedding 配置（30 张表、1380 万索引）上，违约率从 0% 到 20% 时的延迟。使用按线程原子加时，MI350X 前向延迟从 0% 到 20% 违约恶化约 $8.5\times$，而 B200 仅约 $3.3\times$。在 10% 违约时，树归约带来约 $4.26\times$ 加速（2.1 ms 对 9.0 ms）；20% 违约时约 $4.34\times$，且带树归约的 MI350X 比 B200 快约 $1.75\times$。

(a) Forward latency (ms).

（a）前向延迟（ms）。

IV-B3 Example: Grouped Index Select (GIS) Kernel

A core operation in recommendation model training, GIS gathers rows across multiple embedding tables in a single fused kernel. The forward pass is a bandwidth-bound simple gather, while the backward pass scatters gradients back via atomic add, with contention arising when multiple samples select the same row.

IV-B3 示例：分组索引选择（GIS）核

GIS 是推荐模型训练的核心操作，在单个融合核中跨多张 embedding 表收集行。前向是带宽受限的简单 gather；反向通过原子加把梯度散射回去，多样本选中同一行时产生争用。

Forward optimization. AMD’s native 64-thread wavefronts underutilize hardware when embedding dimensions are small. Thus, we introduce 32-thread emulated warps that enable multiple rows to be packed into the same wavefront, eliminating idle lanes and improving memory bandwidth utilization. This optimization delivers $1.16\times$–$1.24\times$ speedup across production configurations (Fig. 12(a)).

前向优化。 当 embedding 维度较小时，AMD 原生 64 线程 wavefront 会欠利用硬件。因此我们引入 32 线程仿真 warp，使多行可打进同一 wavefront，消除空闲通道并改善内存带宽利用率。该优化在生产配置上带来约 $1.16\times$–$1.24\times$ 加速（图 12(a)）。

Backward optimization. In the backward pass, we exploit temporal locality within each thread’s iteration stream. Each thread maintains a 2-slot register cache that maps row indices to accumulated gradient values. When consecutive iterations target the same row (a cache hit), gradients are summed locally in registers with zero global memory traffic. On a miss with both slots occupied, the oldest entry is evicted via a single atomicAdd and replaced. This converts $N$ consecutive updates to the same row into one atomic, with the coalescing factor determined by the locality of the index stream, delivering $1.04\times$–$1.11\times$ backward speedup (Fig. 12(b)).

反向优化。 在反向中，我们利用每线程迭代流内的时间局部性。每线程维护 2 槽寄存器缓存，把行索引映射到累加梯度值。连续迭代命中同一行时，在寄存器本地求和，零全局内存流量；双槽皆满且未命中时，用单次 atomicAdd 驱逐最旧项并替换。这把对同一行的 $N$ 次连续更新变为一次原子，合并因子由索引流局部性决定，反向加速约 $1.04\times$–$1.11\times$（图 12(b)）。

							B200	MI350X
B	T	L	D	Prec	$\alpha$		FWD	BWD	FWD	BWD
							(ms)	(ms)	(ms)	Before	After
65536	5	55	192	fp16	1.0		1.75	14.40	1.48 (0.85$\times$)	23.6 (1.64$\times$)	7.72 (0.54$\times$)
65536	5	55	192	fp16	1.15		0.84	4.61	0.78 (0.93$\times$)	9.92 (2.15$\times$)	3.97 (0.86$\times$)
65536	10	26	256	fp16	1.0		1.86	15.79	1.84 (0.99$\times$)	24.4 (1.55$\times$)	8.06 (0.51$\times$)
65536	10	26	256	fp16	1.15		0.87	4.85	0.91 (1.05$\times$)	11.4 (2.35$\times$)	4.36 (0.90$\times$)
65536	3	55	192	fp32	1.0		1.36	7.56	1.50 (1.10$\times$)	10.6 (1.40$\times$)	5.60 (0.74$\times$)
65536	3	55	192	fp32	1.15		0.53	2.67	0.53 (1.00$\times$)	7.45 (2.79$\times$)	2.69 (1.01$\times$)
65536	6	26	256	fp32	1.0		2.29	8.23	2.03 (0.89$\times$)	10.9 (1.32$\times$)	5.80 (0.70$\times$)
65536	6	26	256	fp32	1.15		0.59	2.98	0.66 (1.12$\times$)	8.95 (3.00$\times$)	3.60 (1.21$\times$)
131072	3	20	256	fp16	1.0		0.84	6.66	0.85 (1.01$\times$)	10.9 (1.64$\times$)	3.69 (0.55$\times$)
131072	3	20	256	fp16	1.15		0.43	2.39	0.46 (1.07$\times$)	8.79 (3.68$\times$)	2.65 (1.11$\times$)
524288	1	20	256	fp16	1.0		1.11	7.02	0.99 (0.89$\times$)	10.5 (1.50$\times$)	4.10 (0.58$\times$)
524288	1	20	256	fp16	1.15		0.54	2.94	0.60 (1.11$\times$)	8.26 (2.81$\times$)	2.83 (0.96$\times$)

Table II: TBE training performance: MI350X forward and backward before/after optimization vs. B200 with representative production configs: B (batch size), T (number of tables), L (lookups per table), D (embedding dimension), Prec (precision), $\alpha$ (Zipf skew factor). Parenthetical values show latency relative to B200 (<1.0$\times$ means MI350X is faster).

							B200	MI350X
B	T	L	D	Prec	$\alpha$		FWD	BWD	FWD	BWD
							(ms)	(ms)	(ms)	优化前	优化后
65536	5	55	192	fp16	1.0		1.75	14.40	1.48 (0.85$\times$)	23.6 (1.64$\times$)	7.72 (0.54$\times$)
65536	5	55	192	fp16	1.15		0.84	4.61	0.78 (0.93$\times$)	9.92 (2.15$\times$)	3.97 (0.86$\times$)
65536	10	26	256	fp16	1.0		1.86	15.79	1.84 (0.99$\times$)	24.4 (1.55$\times$)	8.06 (0.51$\times$)
65536	10	26	256	fp16	1.15		0.87	4.85	0.91 (1.05$\times$)	11.4 (2.35$\times$)	4.36 (0.90$\times$)
65536	3	55	192	fp32	1.0		1.36	7.56	1.50 (1.10$\times$)	10.6 (1.40$\times$)	5.60 (0.74$\times$)
65536	3	55	192	fp32	1.15		0.53	2.67	0.53 (1.00$\times$)	7.45 (2.79$\times$)	2.69 (1.01$\times$)
65536	6	26	256	fp32	1.0		2.29	8.23	2.03 (0.89$\times$)	10.9 (1.32$\times$)	5.80 (0.70$\times$)
65536	6	26	256	fp32	1.15		0.59	2.98	0.66 (1.12$\times$)	8.95 (3.00$\times$)	3.60 (1.21$\times$)
131072	3	20	256	fp16	1.0		0.84	6.66	0.85 (1.01$\times$)	10.9 (1.64$\times$)	3.69 (0.55$\times$)
131072	3	20	256	fp16	1.15		0.43	2.39	0.46 (1.07$\times$)	8.79 (3.68$\times$)	2.65 (1.11$\times$)
524288	1	20	256	fp16	1.0		1.11	7.02	0.99 (0.89$\times$)	10.5 (1.50$\times$)	4.10 (0.58$\times$)
524288	1	20	256	fp16	1.15		0.54	2.94	0.60 (1.11$\times$)	8.26 (2.81$\times$)	2.83 (0.96$\times$)

表 II：TBE 训练性能：MI350X 前向与反向（优化前/后）相对 B200；代表性生产配置：B（batch）、T（表数）、L（每表查找次数）、D（embedding 维）、Prec（精度）、$\alpha$（Zipf 偏斜因子）。括号内为相对 B200 的延迟（$<1.0\times$ 表示 MI350X 更快）。

IV-B4 Table Batched Embeddings (TBE)

Table Batched Embeddings (TBE) [24] is a core primitive in Meta’s recommendation workloads, responsible for looking up and aggregating embeddings from hundreds of tables with billions of rows. TBE fuses multiple embedding lookups into a single kernel launch, amortizing kernel overhead and enabling memory access coalescing across tables. The forward pass gathers embeddings based on sparse indices; the backward pass scatters gradients back to the embedding tables, requiring atomic accumulation for indices that appear multiple times.

IV-B4 表批处理 Embedding（TBE）

表批处理 Embedding（TBE）[24] 是 Meta 推荐负载的核心原语，负责从数百张、共数十亿行的表中查找并聚合 embedding。TBE 把多次 embedding 查找融为单次核启动，摊薄核开销，并实现跨表的访存合并。前向按稀疏索引 gather embedding；反向把梯度散射回 embedding 表，对多次出现的索引需要原子累加。

Backward optimization. We focus on the optimization of the custom backward kernel for MI350X’s wavefront-64 architecture. Key techniques include warp-per-row processing (each wavefront handles one embedding row), hand-tuned vectorized load/store for common embedding dimensions (64, 128, 192, 256), and wavefront-level accumulation that minimizes cross-lane communication.

反向优化。 我们聚焦为 MI350X 的 wavefront-64 架构定制的反向核优化。关键技术包括：按行一个 warp 处理（每个 wavefront 处理一行 embedding）、针对常见 embedding 维（64、128、192、256）手调向量化 load/store，以及最小化跨 lane 通信的 wavefront 级累加。

Results. Table II summarizes TBE training performance across 12 production configurations spanning different batch sizes, table counts, sequence lengths, embedding dimensions, precisions, and index distributions. The backward optimization delivers 1.87–3.32$\times$ speedup on MI350X, with larger gains on FP16 (2.5–3.3$\times$) than FP32 (1.9–2.8$\times$) due to better vectorization. Compared to B200, MI350X with the optimized backward kernel is 1.35–1.96$\times$ faster on uniform distributions ($\alpha$=1.0), where memory bandwidth dominates. On Zipf distributions ($\alpha$=1.15), where skewed access patterns cause cache thrashing, MI350X and B200 are comparable (0.83–1.16$\times$).

结果。 表 II 汇总跨越不同 batch、表数、序列长度、embedding 维、精度与索引分布的 12 组生产配置上的 TBE 训练性能。反向优化在 MI350X 上带来约 1.87–3.32$\times$ 加速；因向量化更好，FP16（约 2.5–3.3$\times$）收益大于 FP32（约 1.9–2.8$\times$）。相对 B200，优化反向核后的 MI350X 在均匀分布（$\alpha$=1.0、内存带宽主导）上快约 1.35–1.96$\times$；在 Zipf 分布（$\alpha$=1.15、偏斜访问导致缓存颠簸）上两者相当（约 0.83–1.16$\times$）。

IV-C Agentic Kernel Optimization Workflow

IV-C 智能体核优化工作流

While GDPA and HSTU are widely shared across models, Meta’s DLRM ecosystem contains a long tail of custom kernels tailored to individual models. Manually tuning each kernel is labor-intensive and does not scale. To address this, we developed an agentic AI workflow that encodes our tuning principles as executable skills, enabling automated kernel performance analysis and optimization. This workflow builds upon Meta’s KernelEvolve framework [42], which provides the underlying agent infrastructure and skill abstraction mechanisms. While KernelEvolve focuses on the general systems architecture for agentic code optimization, our contribution lies in encoding AMD-specific tuning principles as production-ready skills and validating their effectiveness at scale on Meta’s production DLRM training workloads.

GDPA 与 HSTU 虽在多模型间广泛共享，Meta 的 DLRM 生态仍有大量面向个别模型的长尾定制核。手工逐个调核不规模化。为此我们开发了智能体 AI 工作流，把调优原则编码为可执行技能，实现自动化的核性能分析与优化。该工作流建立在 Meta 的 KernelEvolve 框架 [42] 之上，后者提供底层智能体基础设施与技能抽象机制。KernelEvolve 侧重大规模智能体代码优化的通用系统架构；我们的贡献在于把 AMD 专用调优原则编码为生产可用技能，并在 Meta 生产 DLRM 训练负载上规模化验证其有效性。

The agentic workflow operates in three phases:

智能体工作流分三个阶段运行：

1. Benchmark & Profile: Capture representative inputs from production workloads, reproduce the performance in a standalone benchmark, and leverage profiler tools (e.g., AMD rocprofv3 [44], Triton Proton [45]) to identify the main bottlenecks such as low occupancy, poor cache utilization, or instruction stalls.
2. Match & Fix: Map observed bottlenecks to a library of optimization skills, each encoding a tuning principle (e.g., PID swizzling, occupancy hints, atomic avoidance) with code transformation templates and applicability conditions.
3. Generate & Validate: Produce candidate kernel modifications, benchmark them on representative inputs, and verify numerical equivalence (within tolerance, automatically rejecting failures) and performance improvements on a set of production workloads.

1. 基准与剖析： 从生产负载捕获代表性输入，在独立基准中复现性能，并用剖析工具（如 AMD rocprofv3 [44]、Triton Proton [45]）识别主要瓶颈，如低占用率、差的缓存利用或指令停顿。
2. 匹配与修复： 把观察到的瓶颈映射到优化技能库；每项技能编码一条调优原则（如 PID swizzle、occupancy 提示、避开原子操作），并附带代码变换模板与适用条件。
3. 生成与验证： 生成候选核修改，在代表性输入上做基准，并在一组生产负载上验证数值等价（在容差内，自动拒绝失败）与性能改进。

Case Study: Jagged Flash Attention. We illustrate the workflow with the jagged flash attention (JFA) kernel, which represents a challenging optimization target due to its irregular memory access patterns. The JFA kernel, used in our ranking models, exhibited a $2\times$ slowdown on MI350X compared to B200 in its backward pass. Profiling revealed that tl.atomic_add operations dominate execution time, a known bottleneck on AMD’s multi-XCD architecture (see Section IV-B for details). The agent recognized this pattern, noting that the grouped query attention (GQA) mode causes multiple Q heads to race on shared dK/dV gradients. Following the skill’s guidance, the agent rewrote the kernel to accumulate dK/dV in registers per KV-head group and issue a single tl.store at the end, a tree-reduction pattern that eliminates atomic operations. The optimized kernel achieved a $2.3\times$ speedup (32 ms$\to$14 ms), transforming MI350X from $2\times$ slower than B200 to 15–20% faster, and delivering a 25% improvement in end-to-end model training throughput.

案例：Jagged Flash Attention。 我们以 jagged flash attention（JFA）核说明工作流；因其不规则访存模式，它是困难的优化目标。该核用于排序模型，反向在 MI350X 上相对 B200 曾慢约 $2\times$。剖析显示 tl.atomic_add 主导执行时间——这是 AMD 多 XCD 架构上的已知瓶颈（详见第 IV-B 节）。智能体识别出该模式：分组查询注意力（GQA）模式下多个 Q 头在共享的 dK/dV 梯度上竞态。按技能指引，智能体改写核，使每个 KV 头组在寄存器中累加 dK/dV，最后只发一次 tl.store——这是消除原子操作的树归约模式。优化核获得约 $2.3\times$ 加速（32 ms$\to$14 ms），使 MI350X 从比 B200 慢约 $2\times$ 变为快 15–20%，端到端模型训练吞吐提升约 25%。

Applying to dozens of production Triton kernels within Meta’s DLRM ecosystem yields a success rate exceeding 80%. Agent-generated optimizations achieve near-expert performance while reducing development time by an order of magnitude. For a comprehensive evaluation of the effectiveness across diverse kernel types, we refer readers to [42]. By codifying expert knowledge into reusable skills with clear triggering conditions and transformation patterns, the workflow enables systematic exploration of the optimization space while ensuring that anti-patterns for AMD GPUs are automatically detected and addressed.

应用到 Meta DLRM 生态中数十个生产 Triton 核，成功率超过 80%。智能体生成的优化接近专家水平，同时把开发时间降低一个数量级。关于跨多样核类型的全面评估，参见 [42]。通过把专家知识编成带清晰触发条件与变换模式的可复用技能，该工作流能系统探索优化空间，并自动检测、处理 AMD GPU 上的反模式。

Generality. Our workflow builds on Meta’s KernelEvolve [42], which encodes tuning expertise as reusable skills across NVIDIA, AMD, and MTIA. The skill abstraction transfers across ISAs. The atomics-to-tree-reduction skill above has a direct analog on Meta’s MTIA accelerators [48, 49], whose PE array also handles contended atomics inefficiently: the same global-atomic reduction is rewritten as a per-PE local reduction followed by a global reduction over the on-chip reduction-engine network, yielding ${\sim}10\times$ (up to $65\times$) speedups on reduce-style kernels. Because the workflow needs only a portable tile-level kernel DSL and a compiler, which Google TPU (XLA/Pallas) [52] and AWS Trainium (Neuron/NKI) [60] also provide, these skills extend to a new accelerator as an incremental port rather than a from-scratch rebuild.

通用性。 我们的工作流建立在 Meta KernelEvolve [42] 之上，后者跨 NVIDIA、AMD 与 MTIA 把调优专长编码为可复用技能。技能抽象可跨 ISA 迁移。上述「原子到树归约」技能在 Meta 的 MTIA 加速器 [48, 49] 上有直接对应——其 PE 阵列同样不擅长争用原子：同样的全局原子归约被改写为每 PE 本地归约，再经片上归约引擎网络做全局归约，在归约类核上约 ${\sim}10\times$（最高 $65\times$）加速。由于工作流只需可移植的 tile 级核 DSL 与编译器——Google TPU（XLA/Pallas）[52] 与 AWS Trainium（Neuron/NKI）[60] 也具备——这些技能迁到新加速器是增量移植，而非从零重建。

(a) ResNet-50.

（a）ResNet-50。

IV-D PyTorch Inductor Kernel Optimizations

IV-D PyTorch Inductor 核优化

PyTorch Inductor is the default backend for compiler kernel fusion: transforming multiple eager operations into a single generated kernel launch to eliminate intermediate memory traffic, reducing bandwidth requirements by $2\times$–$5\times$ compared to unfused eager execution.

PyTorch Inductor 是编译器核融合的默认后端：把多个 eager 操作变成单次生成的核启动，以消除中间内存流量，相对未融合的 eager 执行可把带宽需求降低约 $2\times$–$5\times$。

Inductor generates four main categories of kernels: pointwise (triton_poi) for elementwise operations (accounts for 40–60% of the total kernel count), reduction (triton_red) for aggregations, persistent reduction (triton_per) for reductions fitting in registers or shared memory, and for-each (triton_for) for batched tensor operations.

Inductor 主要生成四类核：逐点（triton_poi）做逐元素运算（约占核总数的 40–60%）；归约（triton_red）做聚合；持久归约（triton_per）做可放入寄存器或共享内存的归约；以及 for-each（triton_for）做批处理张量运算。

Inductor’s performance benefits are largely rooted in NVIDIA-specific optimizations refined over years of development and iterative tuning. Porting to alternative accelerators, such as AMD GPUs, necessitates more than mechanical adaptation; it requires re-evaluating core heuristics, expanding autotuning search spaces, and developing bespoke micro-architecture aware optimizations.

Inductor 的性能收益植根于经多年开发与迭代调优的 NVIDIA 专用优化。迁到 AMD GPU 等替代加速器，需要的不只是机械适配：还要重审核心启发式、扩大自动调优搜索空间，并开发微架构感知的定制优化。

IV-D1 Optimization Techniques

We apply AMD-specific optimizations leveraging architectural insights from Section IV-A:

IV-D1 优化技术

我们应用基于第 IV-A 节架构洞见的 AMD 专用优化：

Software pipelining. We replace Python range() with Triton’s tl.range() using num_stages=2, enabling better memory-compute overlap. This reduces latency by 11.7% for ResNet-50 reduction kernels.

软件流水线。 用 Triton’s tl.range()（num_stages=2）取代 Python range()，以更好重叠访存与计算。这使 ResNet-50 归约核延迟降低约 11.7%。

Native Triton intrinsics are used to replace library function calls where available: tl.rsqrt() instead of libdevice.rsqrt(), and tl.minimum / tl.maximum instead of helper functions.

原生 Triton 内建在可用处取代库函数调用：用 tl.rsqrt() 取代 libdevice.rsqrt()，用 tl.minimum / tl.maximum 取代辅助函数。

Register configuration. With 256 vector general purpose registers (VGPRs) per thread, we raise Inductor’s register spill threshold from 16 to 32 for larger tiles and better instruction-level parallelism.

寄存器配置。 每线程有 256 个向量通用寄存器（VGPR），我们把 Inductor 的寄存器溢出阈值从 16 提到 32，以支持更大 tile 与更好的指令级并行。

Expanded autotuning. We extend Inductor’s autotuning configurations with AMD-specific tile sizes and occupancy hints. For example, we add non-square tiles ($64\times 32$, $128\times 16$, $128\times 32$, $32\times 512$) for 2D pointwise kernels.

扩大自动调优。 我们为 Inductor 的自动调优配置扩展 AMD 专用 tile 尺寸与 occupancy 提示。例如，为二维逐点核增加非方 tile（$64\times 32$、$128\times 16$、$128\times 32$、$32\times 512$）。

IV-D2 Model-Level Results

We evaluate these optimizations on standard vision and language models. Fig. 13 shows per-category kernel latencies for three models trained on MI350X before/after optimization versus B200. For ResNet-50 [39], reduction kernels improve 11.7%, reaching $1.16\times$ of B200 latency; other Triton inductor kernels outperform B200. For DenseNet-121 [40], pointwise kernels improve 9.2%; reduction improves 6.9%, reaching B200 parity; persistent reduction improves 3.5%. Overall kernel time is reduced by 4.2%. For NanoGPT [41], persistent reduction improves 4.8%; pointwise improves 3.8%. Overall kernel time is reduced by 3.4%.

IV-D2 模型级结果

我们在标准视觉与语言模型上评估这些优化。图 13 展示三个模型在 MI350X 上优化前/后相对 B200 的分类别核延迟。对 ResNet-50 [39]，归约核改进 11.7%，达到 B200 延迟的约 $1.16\times$；其他 Triton Inductor 核优于 B200。对 DenseNet-121 [40]，逐点核改进 9.2%；归约改进 6.9%，达到与 B200 持平；持久归约改进 3.5%。总核时间降低 4.2%。对 NanoGPT [41]，持久归约改进 4.8%；逐点改进 3.8%。总核时间降低 3.4%。

In summary, AMD-specific optimizations deliver 6.9–11.7% improvement for reduction, 3.5–4.8% for persistent reduction, and up to 9.2% for pointwise. While MI350X matches or exceeds B200 across several categories for CNNs, it is still 30% slower on transformer Inductor kernels, indicating optimization headroom remains.

总之，AMD 专用优化使归约改进约 6.9–11.7%，持久归约约 3.5–4.8%，逐点最高约 9.2%。在 CNN 的若干类别上 MI350X 已持平或超过 B200，但在 Transformer 的 Inductor 核上仍约慢 30%，说明仍有优化空间。

V Model-System Co-Design

V 模型—系统协同设计

AMD MI350X’s large HBM capacity creates opportunities for model-system co-design, where training configurations and model architectures are jointly optimized to exploit the additional memory resources. We explore two complementary directions: improving training throughput (examples/second) and improving model quality.

AMD MI350X 的大 HBM 容量为模型—系统协同设计创造了空间：联合优化训练配置与模型架构，以吃到额外内存资源。我们探索两个互补方向：提升训练吞吐（样本/秒）与提升模型质量。

V-A Strategies for Training Efficiency

V-A 训练效率策略

Increased batch size. Larger HBM capacity enables training with larger batch sizes, improving GPU utilization by amortizing kernel launch overhead and better saturating compute units. For example, a friend recommendation model achieves 26% speedup at $2\times$ batch size, 46% at $4\times$, and 51% at $6\times$ (Fig. 15(a)). For a social ranking model migrated to MI350X, doubling the batch size yields 12–15% throughput improvement over B200 baseline with the same batch size (B200 encounters out-of-memory errors at doubled batch size).

加大 batch。 更大的 HBM 容量支持更大 batch 训练，通过摊薄核启动开销、更好吃满计算单元来提高 GPU 利用率。例如，某好友推荐模型在 $2\times$ batch 时加速 26%，$4\times$ 时 46%，$6\times$ 时 51%（图 15(a)）。某迁到 MI350X 的社交排序模型，把 batch 加倍后相对同 batch 的 B200 基线吞吐提升 12–15%（B200 在加倍 batch 时会 OOM）。

Removal of activation checkpointing. Activation checkpointing (AC) trades memory for computation by recomputing activations during the backward pass instead of storing them. With larger HBM, AC can be disabled for memory-constrained layers, eliminating redundant recomputation. The activation memory budget controls how aggressively PyTorch checkpoints activations: a budget of 1.0 saves all activations (maximum memory, no recomputation), while lower values trigger selective recomputation to reduce memory at the cost of throughput. Fig. 14 shows the tradeoff between activation memory budget and training throughput on a 1.3B-parameter GPT model. MI350X’s larger HBM enables it to run without AC entirely, achieving 96.4 examples/second. In contrast, B200 must use aggressive AC (budget $\leq$ 0.35) just to fit, achieving 71.5 examples/second, a $1.35\times$ throughput disadvantage solely attributable to memory constraints.

去掉激活检查点。 激活检查点（AC）用重算换内存：反向时重算激活而非存储。有了更大 HBM，可对内存受限层关闭 AC，消除多余重算。激活内存预算控制 PyTorch 检查点的激进程度：预算 1.0 保存全部激活（内存最大、无重算）；更低值触发选择性重算以省内存、牺牲吞吐。图 14 展示 13 亿参数 GPT 模型上激活内存预算与训练吞吐的权衡。MI350X 更大的 HBM 使其可完全不用 AC，达到 96.4 样本/秒；而 B200 必须用激进 AC（预算 $\leq$ 0.35）才能装下，仅 71.5 样本/秒——仅因内存约束就有约 $1.35\times$ 的吞吐劣势。

Fig. 14: Training throughput vs. HBM tradeoff on a GPT-2 1.3B model with activation checkpointing.

图 14：带激活检查点的 GPT-2 1.3B 模型上，训练吞吐与 HBM 的权衡。

Dense parallelism strategy tuning. Additional HBM allows for relaxed sharding strategies, mitigating communication overhead. For FSDP-based training [25], moving from ZeRO-3 to ZeRO-2 or ZeRO-1 [26] reduces AllGather and ReduceScatter communication volume. Switching to HSDP (Hybrid Sharded Data Parallelism) further hides communication latency by overlapping with computation. For a large Ads foundation model at 1K GPU scale, HSDP provides 6% speedup at the cost of 750 MB additional memory per GPU, a tradeoff feasible only with HBM capacity headroom.

稠密并行策略调优。 额外 HBM 允许放宽分片策略，减轻通信开销。对基于 FSDP 的训练 [25]，从 ZeRO-3 转到 ZeRO-2 或 ZeRO-1 [26] 可减少 AllGather 与 ReduceScatter 通信量。再切到 HSDP（Hybrid Sharded Data Parallelism）可通过与计算重叠进一步掩盖通信延迟。对 1K GPU 规模的大型广告基础模型，HSDP 以每 GPU 额外约 750 MB 内存换来约 6% 加速——这一权衡只有在有 HBM 余量时才可行。

Sparse embedding placement optimization. Recommendation models use large embedding tables that often reside in host memory due to GPU memory constraints. With larger HBM, embedding tables can be moved from host to HBM, eliminating PCIe transfer bottlenecks. Additionally, row-wise sharding can be replaced with table-wise or column-wise sharding, reducing AlltoAll communication and CPU/GPU synchronization overhead. Experiments on a feed ranking model show $>$10% speedup from embedding placement optimization, with the expectations for larger gains when combined with 2D embedding sharding strategies.

稀疏 embedding 放置优化。 推荐模型使用大 embedding 表，常因 GPU 内存限制放在主机内存。有了更大 HBM，可将表从主机迁到 HBM，消除 PCIe 传输瓶颈。此外，可用按表或按列分片取代按行分片，减少 AlltoAll 通信与 CPU/GPU 同步开销。某信息流排序模型实验显示，embedding 放置优化带来超过 10% 加速；若再结合二维 embedding 分片策略，预期收益更大。

Combined results. Fig. 15(b) shows the combined effect of co-design strategies on a feed ranking model. After porting to MI350X, we increased the batch size, removed AC, and optimized embedding placement, resulting in a $1.26\times$ speedup. In conjunction with kernel optimizations, the model realizes a total speedup of $1.5\times$ compared to the initial baseline. These results demonstrate that combining multiple strategies delivers substantial efficiency improvements.

组合结果。 图 15(b) 展示协同设计策略在信息流排序模型上的组合效果。迁到 MI350X 后，我们加大 batch、去掉 AC 并优化 embedding 放置，获得约 $1.26\times$ 加速；再叠加核优化，相对初始基线总加速约 $1.5\times$。说明多策略组合可带来显著效率提升。

V-B Strategies for Model Quality

V-B 模型质量策略

Model architecture upscaling. The MI350X’s expanded HBM capacity accommodates larger model architectures that would otherwise require additional GPU nodes. Increased memory footprint directly translates into superior model quality through several key avenues: (1) Longer input sequences: transformer-based recommendation models can leverage longer sequences to capture richer user context. (2) Expanded embedding tables: larger hash sizes reduce collision rates, thereby preserving higher-fidelity information within the learned representations. (3) Increased dense complexity: support for more dense layers and larger hidden dimensions enables the extraction of more expressive features.

模型结构升维。 MI350X 扩大的 HBM 容量可容纳否则需要更多 GPU 节点的更大模型结构。更大内存足迹通过几条路径直接转化为更好的模型质量：（1）更长输入序列：基于 Transformer 的推荐模型可用更长序列捕捉更丰富的用户上下文；（2）更大 embedding 表：更大哈希规模降低冲突率，从而在学到的表示中保留更高保真信息；（3）更高稠密复杂度：支持更多稠密层与更大隐层维度，以抽取更具表达力的特征。

Results. For a feed ranking model, increasing the number of attention layers from 8 to 16 yields 0.08% normalized entropy (NE) [17] improvement; further increasing to 20 layers yields 0.11% NE improvement.[^4] These architectural changes fit within AMD MI350X’s memory capacity, whereas B200 would require cluster expansion.

结果。 对某信息流排序模型，把注意力层从 8 增到 16 带来约 0.08% 的归一化熵（NE）[17] 提升；再增到 20 层带来约 0.11% NE 提升。[^4] 这些结构改动可装进 AMD MI350X 的内存容量，而 B200 则需要扩展集群。

Generality. This co-design principle, matching the model to the accelerator’s memory hierarchy, is not AMD-specific. It transfers directly to Meta’s MTIA training accelerator (MTIA-300 [49]), whose larger HBM relative to the incumbent GPU is likewise exploited for larger batches and fewer trainers. The principle is direction-agnostic: on the SRAM-bound prior inference generation (MTIA-200 [48]) it instead drives the model to shrink its activation footprint to fit on chip, e.g., decomposing batched matmuls to remove redundant in-batch broadcasts cuts activation memory by ${\sim}25\%$, admitting $10$–$50\times$ larger embedding tables.

通用性。 这一「把模型与加速器内存层次配对」的协同设计原则并非 AMD 专属。它直接迁移到 Meta 的 MTIA 训练加速器（MTIA-300 [49]）：相对现役 GPU 更大的 HBM 同样被用于更大 batch 与更少 trainer。原则与方向无关：在 SRAM 受限的上一代推理芯片（MTIA-200 [48]）上，它反而推动模型收缩激活足迹以装入片上——例如分解批处理矩阵乘、去掉多余的 batch 内广播，可削减约 ${\sim}25\%$ 激活内存，从而容纳约 $10$–$50\times$ 更大的 embedding 表。

(a) A model with increased batch sizes.

（a）加大 batch 后的某模型。

VI Evaluation

VI 评估

The previous sections presented the efficacy of specific ML stack optimizations. This section assesses their collective impact on production recommendation models, illustrating how the cross-layer optimizations jointly maximize end-to-end system-level performance.

前几节展示了具体 ML 栈优化的效果。本节评估它们对生产推荐模型的整体影响，说明跨层优化如何共同最大化端到端系统级性能。

VI-A Production Results

VI-A 生产结果

Table III summarizes the job scale and measured throughput of ten representative production workloads from multiple product groups on MI350X and B200, using matched batch sizes. GPU specifications follow the details in Section II-B; readers should note the differences between our deployment and standard public offerings. Evaluated models range from single-node (8 GPUs) to multi-node (128 GPUs), representing the diversity of production training scales.[^5] Here, Params is the total parameter count (sparse embeddings plus dense layers), and Seq. len is the longest input sequence a model ingests. The geometric mean performance ratio of 101.2% demonstrates that MI350X matches B200 throughput at equal GPU count and batch size across the model suite.

表 III 汇总来自多个产品组的十个代表性生产负载在 MI350X 与 B200 上、匹配 batch 下的作业规模与测得吞吐。GPU 规格见第 II-B 节；读者应注意我们的部署与标准公开产品的差异。评估模型从单节点（8 GPU）到多节点（128 GPU），代表生产训练规模的多样性。[^5] 此处 Params 为总参数量（稀疏 embedding 加稠密层），Seq. len 为模型摄入的最长输入序列。几何均值性能比 101.2% 表明：在等 GPU 数与等 batch 下，MI350X 在该模型套件上与 B200 吞吐持平。

We use an equal number of B200 GPUs to give a straightforward comparison between same-generation hardware from both vendors (Section II-B). As a widely adopted platform, Blackwell serves as a familiar anchor point for readers to evaluate our results. The B200 baseline is not a strawman: as the long-established incumbent, it has received equal or greater optimization than the MI350X: years of model adoption, CUDA library tuning, and kernel development. The B200 keeps an edge on performance-per-watt, but power is a small share of capital-dominated total cost of ownership (TCO), so the MI350X’s lower TCO drives provisioning.

我们使用同等数量的 B200 GPU，以便在两家同代硬件之间做直白比较（第 II-B 节）。作为广泛采用的平台，Blackwell 是读者评估结果的熟悉锚点。B200 基线并非稻草人：作为长期现役平台，它获得的优化不少于甚至多于 MI350X——多年模型适配、CUDA 库调优与核开发。B200 在每瓦性能上仍占优，但功耗在资本主导的总拥有成本（TCO）中占比很小，因此推动扩容决策的是 MI350X 更低的 TCO。

The production models also achieve higher per-GPU training throughput by leveraging larger HBM and batch sizes on MI350X, with gains in Model-D (98% to 102%), Model-E (94% to 110%), and Model-F (92% to 95%). For the kernel latency analysis in the next section, we intentionally use only the same batch sizes. While this underutilizes the MI350X, it provides consistent shapes needed to compare kernel-level performance between platforms.

生产模型还通过在 MI350X 上利用更大 HBM 与更大 batch，获得更高的每 GPU 训练吞吐，例如 Model-D（98%→102%）、Model-E（94%→110%）、Model-F（92%→95%）。下一节的核延迟分析中，我们有意只使用相同 batch。这会欠利用 MI350X，但能提供比较平台间核级性能所需的一致形状。

VI-B Kernel Category Analysis

VI-B 核类别分析

To understand the sources of performance differences, we analyze GPU traces collected during training to categorize kernel latency by type: GEMM, Triton Inductor, CUDA embedding, elementwise, collectives, and hand-written Triton kernels. Fig. 16 shows the per-model breakdown of kernel latency ratio for each category.

为理解性能差异来源，我们分析训练期间采集的 GPU 轨迹，按类型归类核延迟：GEMM、Triton Inductor、CUDA embedding、逐元素、集合通信，以及手写 Triton 核。图 16 展示各模型按类别的核延迟比分解。

Model	GPUs	Seq. len	Params	B200 ex/s	MI350X ex/s	MI350X/B200
A	8	512	${\sim}100$M	921,714	754,936	82%
B	8	256	${\sim}100$B	1,592,926	1,810,036	114%
C	64	200	${\sim}200$B	728,001	884,216	121%
D	64	1000	${\sim}300$B	774,324	755,221	98%
E	64	200	${\sim}300$B	283,641	267,635	94%
F	128	1000	${\sim}3.6$T	330,481	304,781	92%
G	64	200	${\sim}30$B	544,537	499,066	92%
H	64	24,600	${\sim}500$B	120,223	115,301	96%
I	80	4096–8192	${\sim}1.9$T	546,072	558,127	102%
J	24	700	${\sim}20$B	185,728	242,325	130%
Geomean	—	—	—	—	101.2%

Table III: End-to-end training throughput (examples/second) on MI350X vs B200 across ten production models. Higher is better.

模型	GPU 数	Seq. len	Params	B200 ex/s	MI350X ex/s	MI350X/B200
A	8	512	${\sim}100$M	921,714	754,936	82%
B	8	256	${\sim}100$B	1,592,926	1,810,036	114%
C	64	200	${\sim}200$B	728,001	884,216	121%
D	64	1000	${\sim}300$B	774,324	755,221	98%
E	64	200	${\sim}300$B	283,641	267,635	94%
F	128	1000	${\sim}3.6$T	330,481	304,781	92%
G	64	200	${\sim}30$B	544,537	499,066	92%
H	64	24,600	${\sim}500$B	120,223	115,301	96%
I	80	4096–8192	${\sim}1.9$T	546,072	558,127	102%
J	24	700	${\sim}20$B	185,728	242,325	130%
几何均值	—	—	—	—	101.2%

表 III：十个生产模型上 MI350X 相对 B200 的端到端训练吞吐（样本/秒）。越高越好。

Fig. 16: Kernel latency ratio (MI350X / B200) by category. For each model, the sum of kernel latencies on B200 is normalized to 100%. This analysis does not consider multi-stream overlap.

图 16：按类别的核延迟比（MI350X / B200）。对每个模型，B200 上核延迟之和归一化为 100%。本分析未考虑多流重叠。

Where MI350X excels. Hand-written Triton kernels show consistent MI350X advantages, delivering a 6% to 52% reduction in execution latency through the optimizations discussed in Section IV-A. Model-H benefits particularly from faster HSTU kernels, achieving $2.1\times$ speedup for the Triton category. Embedding operations show mixed results but generally favor MI350X ($0.61\times$–$1.04\times$), leveraging the CUDA kernel optimizations detailed in Section IV-B.

MI350X 占优之处。 手写 Triton 核 持续显示 MI350X 优势：经第 IV-A 节所述优化，执行延迟降低约 6% 到 52%。Model-H 尤其受益于更快的 HSTU 核，Triton 类别达到约 $2.1\times$ 加速。Embedding 操作 结果不一，但总体偏向 MI350X（约 $0.61\times$–$1.04\times$），得益于第 IV-B 节详述的 CUDA 核优化。

Overall, collective performance is comparable between the two GPU types, with MI350X outperforming B200 in multi-host, small-message AllToAll collectives, specifically benefiting Models C and J. Note that overlapped kernel execution presents a challenge for latency analysis: due to the multi-stream nature on GPUs, concurrent compute and communication kernels complicate the measurement of their individual contributions to end-to-end latency. For instance, collectives in Models E and F take longer to execute but are masked by compute kernels in the B200 traces, resulting in no net impact on throughput. Consequently, total kernel latency may overstate the true impact of collectives, leading to discrepancies between Fig. 16 and measured throughput in Table III.

总体而言，集合通信性能 在两类 GPU 上相当；MI350X 在多主机、小消息 AllToAll 集合上优于 B200，尤其惠及 Model C 与 J。需注意：重叠的核执行使延迟分析变难——GPU 上的多流特性下，并发的计算与通信核使各自对端到端延迟的贡献难以单独度量。例如 Model E 与 F 中集合执行更久，但在 B200 轨迹里被计算核掩盖，对吞吐无净影响。因此总核延迟可能高估集合通信的真实影响，导致图 16 与表 III 测得吞吐之间出现偏差。

Where B200 leads. GEMM operations consistently favor B200, with MI350X showing $1.3\times$–$1.7\times$ higher latency. MI350X lacks native hardware support for TF32, instead emulating them with three BF16 GEMMs at a performance penalty. Model-A’s overall 18% performance deficit is largely attributable to TF32 GEMM operations. While this is a notable limitation, modern models are increasingly moving toward BF16 or lower precision. We expect continued model modernization and improvements to emulated TF32 performance to further narrow the overall gap. Additionally, cuBLAS outperforms hipBLAS due to its relative maturity, especially in the case of irregular matrix dimensions. Triton Inductor kernels show model-dependent results: some models achieve parity ($0.99\times$–$1.08\times$) while others show MI350X incurring $1.22\times$ higher latency than B200, reflecting ongoing efforts to refine Inductor and eliminate heuristics specifically tailored for NVIDIA GPUs (Section IV-D).

B200 占优之处。 GEMM 操作 持续偏向 B200，MI350X 延迟约高 $1.3\times$–$1.7\times$。MI350X 缺少原生 TF32 硬件支持，只能用三次 BF16 GEMM 仿真，带性能代价。Model-A 整体约 18% 的性能逆差主要归因于 TF32 GEMM。虽是明显短板，现代模型正越来越多地转向 BF16 或更低精度。我们预期持续的模型现代化与仿真 TF32 性能改进会进一步收窄整体差距。此外，因相对更成熟，cuBLAS 在不规则矩阵维度上优于 hipBLAS。Triton Inductor 核 结果随模型而异：有的持平（约 $0.99\times$–$1.08\times$），有的 MI350X 延迟比 B200 高约 $1.22\times$，反映仍在完善 Inductor、消除专为 NVIDIA GPU 定制的启发式（第 IV-D 节）。

VI-C Operational Considerations

VI-C 运维考量

Integrating a second GPU vendor into a previously NVIDIA-only fleet raised operational challenges in power, reliability, and observability. In addition to performance, operational maturity is just as critical to adopting a new vendor, and much of the tooling and fleet-management machinery we built generalizes to future accelerators.

把第二家 GPU 厂商接入此前仅有 NVIDIA 的机队，在功耗、可靠性与可观测性上带来运维挑战。除性能外，运维成熟度对采用新厂商同样关键；我们构建的大量工具与机队管理机制也可推广到未来加速器。

Reliability. We track MI350X on the same metrics as our NVIDIA fleet (training-failure rate, unplanned unavailability, and hardware-swap rate), and it is steadily converging toward NVIDIA-fleet levels: the training-failure rate has more than halved from its peak as we resolved new at-scale failure modes, namely HBM RAS, multi-chiplet memory-access faults, and silent data corruption (SDC). SDC, the bit-flips that corrupt results without crashing, grows with cluster size and is the hardest to tame, so we address it end-to-end: a model-level checksum-and-replay check flags non-deterministic steps (distinguishing on-chip corruption from communication non-determinism); an all-pairs vote across a host’s GPUs pinpoints the culprit device (the one that disagrees with all its peers) and auto-quarantines the host; and the same check runs during host qualification and burn-in to screen SDC-prone hosts before they enter the training pool, catching over $2\times$ the SDC of the vendor’s field-health tool. These mitigations preserve correctness: across the tens of models we train on both vendors, results reach numerical parity (training-loss difference within $0.3\%$ and NE within ${\sim}10^{-4}$ at multi-thousand-GPU scale).

可靠性。 我们用与 NVIDIA 机队相同的指标跟踪 MI350X（训练失败率、计划外不可用、硬件更换率），并正稳步向 NVIDIA 机队水平收敛：随着我们解决新的规模化故障模式——即 HBM RAS、多 chiplet 访存故障与静默数据损坏（SDC）——训练失败率已从峰值降一半以上。SDC 是不崩溃却腐蚀结果的位翻转，随集群变大而加剧，也最难治理，因此我们端到端应对：模型级校验和—重放检查标记非确定性步骤（区分片上损坏与通信非确定性）；主机内 GPU 全对投票定位罪魁设备（与所有对等者不一致的那张）并自动隔离主机；同样检查在主机验收与烧机中运行，在进入训练池前筛出易 SDC 主机，捕获量超过厂商现场健康工具的 $2\times$。这些缓解措施保住正确性：在我们跨两家厂商训练的数十个模型上，结果达到数值持平（数千 GPU 规模下训练损失差在 $0.3\%$ 内，NE 在 ${\sim}10^{-4}$ 内）。

Observability. AMD telemetry initially lagged NVIDIA’s mature stack, so we closed the gap: we integrated the efficiency metrics our NVIDIA fleet relies on (TFLOPS, SM/CU and memory utilization, throughput) and a hardware-agnostic trace-analysis interface (compute/communication breakdown, imbalance, overlap, and timeline analyses) that also enables side-by-side cross-vendor trace comparison (the basis for the apples-to-apples measurements here) and extends to MTIA.

可观测性。 AMD 遥测起初落后于 NVIDIA 的成熟栈，我们补齐差距：集成 NVIDIA 机队依赖的效率指标（TFLOPS、SM/CU 与内存利用率、吞吐），以及与硬件无关的轨迹分析接口（计算/通信分解、不均衡、重叠与时间线分析），并支持跨厂商并排轨迹比较（本文苹果对苹果测量的基础），还可扩展到 MTIA。

Power management. Operating the MI350X against its excursion envelope rather than its steady-state limit is a fleet-level capability: it depends on platform provisioning beyond the OCP minimum and on tuning the control-loop response to production phase patterns. Section III-C describes the mechanism and the resulting 3–6% end-to-end gain.

电源管理。 按冲程包络而非稳态限运行 MI350X 是机队级能力：它依赖超出 OCP 下限的平台供给，以及按生产阶段模式调控制环响应。第 III-C 节描述了机制及由此带来的约 3–6% 端到端收益。

VII Related Work

VII 相关工作

Hardware Accelerators. Custom hardware accelerators exploit the structure of neural network computations. Meta’s MTIA [46, 47, 50] an in-house family of custom AI chips, designed to provide superior cost efficiency over commercial GPUs for the company’s large-scale DLRM and GenAI workloads. Google’s TPU [52] employs systolic arrays for efficient matrix multiplication. Recent accelerators target emerging workloads: spiking neural networks [54], heterogeneous precision [55], and processing-in-memory architectures that reduce the memory bandwidth bottleneck [56].

硬件加速器。 定制硬件加速器利用神经网络计算的结构。Meta 的 MTIA [46, 47, 50] 是自研 AI 芯片家族，旨在为公司大规模 DLRM 与 GenAI 负载提供优于商用 GPU 的成本效率。Google 的 TPU [52] 用脉动阵列做高效矩阵乘。近期加速器面向新兴负载：脉冲神经网络 [54]、异构精度 [55]，以及减轻内存带宽瓶颈的存内计算架构 [56]。

Distributed Training Infrastructure. The LLM pretraining study on AMD MI300X GPUs [9] presents the transformer sharding strategy and the practical guidance for fault-tolerant distributed training. Scaling DLRM training requires purpose-built infrastructure. Zion [30] is a scale-up DLRM training platform that co-locates CPUs and accelerators to handle the memory capacity demands of embedding tables while leveraging GPU compute for dense operations. TorchRec [23] provides a PyTorch domain library with optimized distributed primitives, including automatic sharding planners that partition embedding tables across devices.

分布式训练基础设施。 关于 AMD MI300X GPU 上 LLM 预训练的研究 [9] 给出 Transformer 分片策略与容错分布式训练的实践指引。扩展 DLRM 训练需要专用基础设施。Zion [30] 是纵向扩展的 DLRM 训练平台，把 CPU 与加速器共置，以应对 embedding 表的内存容量需求，并用 GPU 算力做稠密运算。TorchRec [23] 提供带优化分布式原语的 PyTorch 领域库，包括把 embedding 表划分到各设备的自动分片规划器。

Embedding Table Sharding and Placement. DLRMs present unique systems challenges due to their massive embedding tables, which can exceed hundreds of gigabytes and dominate both memory capacity and bandwidth requirements. Intelligent partitioning of embedding tables across heterogeneous memory hierarchies is critical for training throughput. RecShard [28] formulates embedding placement as a mixed-integer linear program, leveraging per-feature access statistics—including categorical value frequency distributions and pooling factors—to achieve over 2.5$\times$ throughput improvement compared to size-based sharding. The key insight is that not all embedding rows are accessed equally; exploiting this skew enables fine-grained placement across GPU HBM, CPU DRAM, and SSDs. 2D sparse parallelism [27] combines data and model parallelism for embeddings, enabling training to scale to thousands of GPUs.

Embedding 表分片与放置。 DLRM 因其巨大 embedding 表而带来独特系统挑战——表可达数百 GB，主导内存容量与带宽需求。跨异构内存层次智能划分 embedding 表对训练吞吐至关重要。RecShard [28] 把 embedding 放置表述为混合整数线性规划，利用每特征访问统计——包括类别值频率分布与池化因子——相对按大小分片获得超过 2.5$\times$ 的吞吐提升。关键洞见是：并非所有 embedding 行被均等访问；利用这种偏斜可在 GPU HBM、CPU DRAM 与 SSD 间做细粒度放置。二维稀疏并行 [27] 对 embedding 结合数据并行与模型并行，使训练可扩展到数千 GPU。

Kernel-Level Optimizations. Efficient GPU kernels are essential for embedding operations. FBGEMM [24] provides high-performance implementations of table batched embedding lookups, fused optimizers, and quantized operations optimized for recommendation workloads. These kernels exploit memory access coalescing, warp-level parallelism, and mixed-precision computation to maximize GPU utilization. For inference, mixed-precision embeddings [29] achieve 200$\times$ compression by learning per-feature precision assignments, grouping features by frequency to reduce the search space. KernelEvolve [42] provides an agentic framework for kernel optimization; we build specialized workflows that incorporates expert-tuning heuristics and architectural insights tailored for training workloads, simplifying cross-platform GPU kernel porting and optimization.

核级优化。 高效 GPU 核对 embedding 操作至关重要。FBGEMM [24] 提供面向推荐负载优化的表批处理 embedding 查找、融合优化器与量化运算的高性能实现。这些核利用访存合并、warp 级并行与混合精度计算以最大化 GPU 利用率。在推理侧，混合精度 embedding [29] 通过学习每特征精度分配、按频率分组特征以缩小搜索空间，达到约 200$\times$ 压缩。KernelEvolve [42] 提供核优化的智能体框架；我们构建的专用工作流纳入面向训练负载的专家调优启发式与架构洞见，简化跨平台 GPU 核移植与优化。

VIII Conclusion

VIII 结论

We present the first comprehensive study demonstrating that AMD MI350X GPUs achieve performance parity with NVIDIA B200 for production DLRM training through cross-stack optimizations spanning core ROCm libraries, kernel tuning, and model co-design. Our work addresses architectural differences between vendors—memory capacity, atomic performance, instruction sets—through principled algorithmic adaptations, hardware-aware autotuning, and HBM-conscious model co-design, providing a transferable framework for hardware-agnostic training at scale. Our optimizations have been upstreamed and open-sourced, establishing a foundation for multi-vendor ML infrastructure.

我们给出首份全面研究，表明通过跨越核心 ROCm 库、核调优与模型协同设计的跨栈优化，AMD MI350X GPU 可在生产 DLRM 训练上与 NVIDIA B200 性能持平。我们通过有原则的算法适配、硬件感知自动调优与关注 HBM 的模型协同设计，应对厂商间架构差异——内存容量、原子性能、指令集——为规模化、与硬件无关的训练提供可迁移框架。相关优化已上游并开源，为多厂商 ML 基础设施打下基础。

References

参考文献

[1] A. Smith and V. K. Alla (2025) AMD Instinct MI300X: A Generative AI Accelerator and Platform Architecture . In IEEE/ACM MICRO , Cited by: §I .

[2] AMD (2026) AMD Instinct MI350X GPUs (Website) External Links: Link Cited by: §I .

[3] Y. J. Song, A. Tulloch, H. Reddy, C. Tang, and V. Thakkar (2026) Four MTIA Chips in Two Years: Scaling AI Experiences for Billions (Website) External Links: Link Cited by: §I , §VII .

[4] A. Firoozshahian, J. Coburn, R. Levenstein, R. Nattoji, A. Kamath, O. Wu, G. Grewal, H. Aepala, B. Jakka, B. Dreyer, A. Hutchin, U. Diril, K. Nair, E. K. Aredestani, M. Schatz, Y. Hao, R. Komuravelli, K. Ho, S. Abu Asal, J. Shajrawi, K. Quinn, N. Sreedhara, P. Kansal, W. Wei, D. Jayaraman, L. Cheng, P. Chopda, E. Wang, A. Bikumandla, A. Karthik Sengottuvel, K. Thottempudi, A. Narasimha, B. Dodds, C. Gao, J. Zhang, M. Al-Sanabani, A. Zehtabioskuie, J. Fix, H. Yu, R. Li, K. Gondkar, J. Montgomery, M. Tsai, S. Dwarakapuram, S. Desai, N. Avidan, P. Ramani, K. Narayanan, A. Mathews, S. Gopal, M. Naumov, V. Rao, K. Noru, H. Reddy, P. Venkatapuram, and A. Bjorlin (2023) MTIA: First Generation Silicon Targeting Meta’s Recommendation Systems . In ACM/IEEE ISCA , Cited by: §I , §VII .

[5] MTIA Team (2026) MTIA 300: Meta’s First Training Chip Featuring Built-in NICs and Collective Offloading Engines . In ACM/IEEE ISCA , Cited by: §I , §IV-C , §V-B .

[6] J. Coburn, C. Tang, S. A. Asal, N. Agrawal, R. Chinta, H. Dixit, B. Dodds, S. Dwarakapuram, A. Firoozshahian, C. Gao, K. Gondkar, T. Graf, J. Hu, J. Huang, S. Hughes, A. Hutchin, B. Jakka, G. J. Chen, I. Kalyanaraman, A. Kamath, P. Kansal, E. Kazi, R. Levenstein, M. Maddury, A. Mastro, S. Medaiyese, P. Modi, J. Montgomery, S. Nadathur, A. Nagpal, A. Narasimha, M. Naumov, E. Ozer, J. Park, P. Ramani, H. Reddy, D. Reiss, D. Roy, S. Sekar, A. Sharma, P. Shetty, A. Sukumaran-Rajam, E. Tal, M. Tsai, S. Varshini, R. Wareing, O. Wu, X. Xie, J. Yang, H. Yu, T. Zargar, Z. Zeng, F. Zhang, A. Matthews, X. Jiao, J. Zhang, E. Menage, T. E. Stokke, and M. Sourouri (2025) Meta’s Second Generation AI Chip: Model-Chip Co-Design and Productionization Experiences . In ACM/IEEE ISCA , Cited by: §I , §IV-C , §V-B .

[7] N. P. Jouppi, D. H. Yoon, M. Ashcraft, M. Gottscho, T. B. Jablin, G. Kurian, J. Laudon, S. Li, P. Ma, X. Ma, T. Norrie, N. Patil, S. Prasad, and C. Young (2021) Ten Lessons From Three Generations Shaped Google’s TPUv4i . In ACM/IEEE ISCA , Cited by: §I .

[8] X. Fu, Z. Zhang, H. Fan, G. Huang, M. El-Shabani, R. Huang, R. Solanki, F. Wu, R. Diamant, and Y. Wang (2024) Distributed Training of Large Language Models on AWS Trainium . In ACM SoCC , Cited by: §I , §IV-C .

[9] R. Borkar, A. Wall, P. Pulavarthi, and Y. Yu (2026) Azure Maia for the era of AI: From silicon to software to systems (Website) External Links: Link Cited by: §I .

[10] H. Liao, J. Tu, J. Xia, H. Liu, X. Zhou, H. Yuan, and Y. Hu (2021) Ascend: A Scalable and Unified Architecture for Ubiquitous Deep Neural Network Computing . In IEEE HPCA , Cited by: §I .

[11] Y. Jiao, L. Han, and X. Long (2020) Hanguang 800 NPU—The Ultimate AI Inference Solution for Data Centers . In IEEE HotChips , Cited by: §I .

[12] J. Ouyang, M. Noh, Y. Wang, W. Qi, Y. Ma, C. Gu, S. Kim, K. Hong, W. Bae, Z. Zhao, J. Wang, P. Wu, X. Gong, J. Shi, H. Zhu, and X. Du (2020) Baidu Kunlun: An AI Processor for Diversified Workloads . In IEEE HotChips , Cited by: §I .

[13] M. Naumov, D. Mudigere, H. M. Shi, J. Huang, N. Sundaraman, J. Park, X. Wang, U. Gupta, C. Wu, A. G. Azzolini, D. Dzhulgakov, A. Mallevich, I. Cherniavskii, Y. Lu, R. Krishnamoorthi, A. Yu, V. Kondratenko, S. Pereira, X. Chen, W. Chen, V. Rao, B. Jia, L. Xiong, and M. Smelyanskiy (2019) Deep Learning Recommendation Model for Personalization and Recommendation Systems . arXiv abs/1906.00091 . External Links: Link Cited by: §I , §II-A .

[14] MLCommons (2026) DLRMv3: Generative recommendation benchmark in MLPerf Inference (Website) External Links: Link Cited by: §I , §II-A .

[15] X. He, J. Pan, O. Jin, T. Xu, B. Liu, T. Akkaya, Y. Moon, L. Saleski, J. Yang, W. Zheng, and J. Quiñonero Candela (2014) Practical Lessons from Predicting Clicks on Ads at Facebook . In Proceedings of the Eighth International Workshop on Data Mining for Online Advertising , Cited by: §I , §V-B .

[16] D. Mudigere, Y. Hao, J. Huang, Z. Jia, A. Tulloch, S. Sridharan, X. Liu, M. Ozdal, J. Nie, J. Park, L. Luo, J. (. Yang, L. Gao, D. Ivchenko, A. Basant, Y. Hu, J. Yang, E. K. Ardestani, X. Wang, R. Komuravelli, C. Chu, S. Yilmaz, H. Li, J. Qian, Z. Feng, Y. Ma, J. Yang, E. Wen, H. Li, L. Yang, C. Sun, W. Zhao, D. Melts, K. Dhulipala, K. Kishore, T. Graf, A. Eisenman, K. K. Matam, A. Gangidi, G. J. Chen, M. Krishnan, A. Nayak, K. Nair, B. Muthiah, M. khorashadi, P. Bhattacharya, P. Lapukhov, M. Naumov, A. Mathews, L. Qiao, M. Smelyanskiy, B. Jia, and V. Rao (2022) Software-Hardware Co-Design for Fast and Scalable Training of Deep Learning Recommendation Models . In ACM/IEEE ISCA , Cited by: §II-A .

[17] Z. Liu, L. Zou, X. Zou, C. Wang, B. Zhang, D. Tang, B. Zhu, Y. Zhu, P. Wu, K. Wang, and Y. Cheng (2022) Monolith: Real Time Recommendation System With Collisionless Embedding Table . In ACM RecSys , Cited by: §II-A .

[18] H. Cheng, L. Koc, J. Harmsen, T. Shaked, T. Chandra, H. Aradhye, G. Anderson, G. Corrado, W. Chai, M. Ispir, R. Anil, Z. Haque, L. Hong, V. Jain, X. Liu, and H. Shah (2016) Wide & Deep Learning for Recommender Systems . In Proceedings of the 1st Workshop on Deep Learning for Recommender Systems (DLRS) , Cited by: §II-A .

[19] G. Zhou, X. Zhu, C. Song, Y. Fan, H. Zhu, X. Ma, Y. Yan, J. Jin, H. Li, and K. Gai (2018) Deep Interest Network for Click-Through Rate Prediction . In ACM KDD , Cited by: §II-A .

[20] U. Gupta, C. Wu, X. Wang, M. Naumov, B. Reagen, D. Brooks, B. Cottel, K. Hazelwood, M. Hempstead, B. Jia, H. S. Lee, A. Malevich, D. Mudigere, M. Smelyanskiy, L. Xiong, and X. Zhang (2020) The Architectural Implications of Facebook’s DNN-Based Personalized Recommendation . In IEEE HPCA , Cited by: §II-A .

[21] J. Zhai, L. Liao, X. Liu, Y. Wang, R. Li, X. Cao, L. Gao, Z. Gong, F. Gu, J. He, Y. Lu, and Y. Shi (2024) Actions Speak Louder than Words: Trillion-Parameter Sequential Transducers for Generative Recommendations . In ICML , Cited by: §II-A , §II-A , §IV-A1 .

[22] D. Ivchenko, D. Van Der Staay, C. Taylor, X. Liu, W. Feng, R. Kindi, A. Sudarshan, and S. Sefati (2022) TorchRec: A PyTorch Domain Library for Recommendation Systems . In ACM RecSys , Cited by: §II-A , §VII .

[23] X. Zhang, Q. Zhu, L. Xu, Z. Huda, W. Zhou, J. Fang, D. van der Staay, Y. Hu, J. Nie, J. Yang, and C. Yang (2025) Two-dimensional Sparse Parallelism for Large Scale Deep Learning Recommendation Model Training . arXiv abs/2508.03854 . External Links: Link Cited by: §II-A , §VII .

[24] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei (2020) Scaling Laws for Neural Language Models . arXiv abs/2001.08361 . External Links: Link Cited by: §II-A .

[25] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre (2022) Training Compute-Optimal Large Language Models . arXiv abs/2203.15556 . External Links: Link Cited by: §II-A .

[26] B. Zhang, L. Luo, Y. Chen, J. Nie, X. Liu, S. Li, Y. Zhao, Y. Hao, Y. Yao, E. D. Wen, J. Park, M. Naumov, and W. Chen (2024) Wukong: Towards a Scaling Law for Large-Scale Recommendation . In ICML , Cited by: §II-A .

[27] S. Williams, A. Waterman, and D. Patterson (2009) Roofline: An Insightful Visual Performance Model for Multicore Architectures . Communications of the ACM 52 ( 4 ), pp. 65–76 . Cited by: §II-B , §III .

[28] D. Rabinovitsj and O. Baldonado (2026) Meta’s open AI hardware vision (Website) External Links: Link Cited by: §II-B .

[29] AMD (2026) hipBLASLt: General Matrix Multiplication Library for AMD GPUs (Website) External Links: Link Cited by: §III-A .

[30] D. E. Tanner (2018) Tensile: Auto-Tuning GEMM GPU Assembly for All Problem Sizes . In 2018 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW) , Cited by: §III-A .

[31] AMD (2025) Introducing AMD CDNA™ 4 Architecture . White Paper Advanced Micro Devices, Inc. . External Links: Link Cited by: §III-A .

[32] R. Swann, M. Osama, X. Guo, B. Nelson, L. Zhang, A. Brown, Y. Ong, A. Yazdani, S. Siddens, G. Dasika, and A. Underwood (2025) tritonBLAS: Triton-based Analytical Approach for GEMM Kernel Parameter Selection . arXiv abs/2512.04226 . External Links: Link Cited by: §III-B .

[33] O. C. Project (2026) OAI OAM Base Specification, Revision 2.0 (Website) External Links: Link Cited by: §III-C .

[34] J. Kolla, P. Alizadeh, and G. Lee (2026) Understanding RCCL Bandwidth and xGMI Performance on AMD Instinct™ MI300X (Website) External Links: Link Cited by: §III-D .

[35] V. Gupta, T. Koren, and Y. Singer (2018) Shampoo: Preconditioned Stochastic Tensor Optimization . In ICML , Cited by: §III-E .

[36] H. M. Shi, T. Lee, S. Iwasaki, J. Gallego-Posada, Z. Li, K. Rangadurai, D. Mudigere, and M. Rabbat (2023) A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale . arXiv abs/2309.06497 . External Links: Link Cited by: §III-E .

[37] D. P. Kingma and J. Ba (2015) Adam: A Method for Stochastic Optimization . In ICLR , Cited by: §III-E .

[38] R. S. Martin, C. Reinsch, and J. H. Wilkinson (1968) Householder’s Tridiagonalization of a Symmetric Matrix . Numerische Mathematik 11 ( 3 ), pp. 181–195 . Cited by: §III-E .

[39] J. J. M. Cuppen (1980) A Divide and Conquer Method for the Symmetric Eigenproblem . Numerische Mathematik 36 , pp. 177–195 . Cited by: §III-E .

[40] J. (. Xu, C. Chen, S. Yang, M. Hoehnerbach, X. (. Liu, T. Zadouri, D. (. Shankar, J. Zhou, H. Yu, M. Ren, H. Xu, C. Yang, J. Nie, H. Zhang, H. Li, M. Shu, M. Sultan, M. Leung, J. Bocharov, and T. Dao (2026) Generalized Dot-Product Attention: Tackling Real-World Challenges in GPU Training Kernels (Website) External Links: Link Cited by: §IV-A1 .

[41] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention Is All You Need . In NeurIPS , Cited by: §IV-A1 .

[42] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré (2022) FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness . In NeurIPS , Cited by: §IV-A1 .

[43] B. Hou, X. Liu, X. Liu, J. Xu, Y. Badr, M. Hang, S. Chanpuriya, J. Zhou, Y. Yang, H. Xu, Q. Suo, L. Chen, Y. Hu, J. Zhang, H. Xiong, Y. Huang, C. Chen, Y. Dong, Y. Yang, S. Chang, X. Gan, W. Chen, S. Kolay, D. Liu, J. Nie, C. Yang, E. Wen, J. Yang, and H. Li (2026) Kunlun: Establishing Scaling Laws for Massive-Scale Recommendation Systems through Unified Architecture Design . arXiv abs/2602.10016 . External Links: Link Cited by: §IV-A1 .

[44] Z. Zeng, X. Liu, M. Hang, X. Liu, Q. Zhou, C. Yang, Y. Liu, Y. Ruan, L. Chen, Y. Chen, Y. Hao, J. Xu, J. Nie, X. Liu, B. Zhang, W. Wen, S. Yuan, H. Yin, X. Zhang, K. Wang, W. Chen, Y. Han, H. Li, C. Yang, B. Long, P. S. Yu, H. Tong, and J. Yang (2025) InterFormer: Effective Heterogeneous Interaction Learning for Click-Through Rate Prediction . arXiv abs/2411.09852 . External Links: Link Cited by: §IV-A1 .

[45] Meta Engineering (2026) Meta’s Generative Ads Model (GEM): The Central Brain Accelerating Ads Recommendation AI Innovation (Website) External Links: Link Cited by: §IV-A3 .

[46] D. Khudia, J. Huang, P. Basu, S. Deng, H. Liu, J. Park, and M. Smelyanskiy (2021) FBGEMM: Enabling High-Performance Low-Precision Deep Learning Inference . arXiv abs/2101.05615 . External Links: Link Cited by: §IV-B4 , §VII .

[47] G. Liao, H. Qin, Y. Wang, A. Golden, M. Kuchnik, Y. Yetim, J. J. Ang, C. Fu, Y. He, S. Hsia, Z. Jiang, D. Li, U. Pashkevich, V. Puvvada, F. Shi, M. Steiner, R. Xiao, N. Yan, X. Yu, Z. Fang, R. Levenstein, K. Ho, H. Zhu, A. Hammond, R. Li, A. Mathews, K. Gondkar, A. Zainul-Abedin, K. Singh, H. Yu, W. Chi, B. Huang, S. Zhang, N. Weller, Z. Marine, W. Cook, C. Wu, and G. Liu (2026) KernelEvolve: Scaling Agentic Kernel Coding for Heterogeneous AI Accelerators at Meta . In ACM/IEEE ISCA , Cited by: §IV-C , §IV-C , §IV-C , §VII .

[48] AMD (2026) ROCprofiler-SDK Documentation (Website) External Links: Link Cited by: item 1 .

[49] Triton Developers (2024) Proton: A Lightweight Profiler for Triton (Website) Note: Part of the Triton language and compiler project External Links: Link Cited by: item 1 .

[50] N. P. Jouppi, C. Young, N. Patil, D. Patterson, G. Agrawal, R. Bajwa, S. Bates, S. Bhatia, N. Boden, A. Borchers, R. Boyle, P. Cantin, C. Chao, C. Clark, J. Coriell, M. Daley, M. Dau, J. Dean, B. Gelb, T. V. Ghaemmaghami, R. Gottipati, W. Gulland, R. Hagmann, C. R. Ho, D. Hogberg, J. Hu, R. Hundt, D. Hurt, J. Ibarz, A. Jaffey, A. Jaworski, A. Kaplan, H. Khaitan, A. Koch, N. Kumar, S. Lacy, J. Laudon, J. Law, D. Le, C. Leary, Z. Liu, K. Lucke, A. Lundin, G. MacKean, A. Maggiore, M. Mahony, K. Miller, R. Nagarajan, R. Narayanaswami, R. Ni, K. Nix, T. Norrie, M. Omernick, N. Penukonda, A. Phelps, J. Ross, M. Ross, A. Salek, E. Samadiani, C. Severn, G. Sizikov, M. Snelham, J. Souter, D. Steinberg, A. Swing, M. Tan, G. Thorson, B. Tian, H. Toma, E. Tuttle, V. Vasudevan, R. Walter, W. Wang, E. Wilcox, and D. H. Yoon (2017) In-Datacenter Performance Analysis of a Tensor Processing Unit . In ACM/IEEE ISCA , Cited by: §IV-C , §VII .

[51] K. He, X. Zhang, S. Ren, and J. Sun (2016) Deep Residual Learning for Image Recognition . In CVPR , Cited by: §IV-D2 .

[52] G. Huang, Z. Liu, L. van der Maaten, and K. Q. Weinberger (2017) Densely Connected Convolutional Networks . In CVPR , Cited by: §IV-D2 .

[53] A. Karpathy (2026) nanoGPT: The simplest, fastest repository for training/finetuning medium-sized GPTs (Website) External Links: Link Cited by: §IV-D2 .

[54] Y. Zhao, A. Gu, R. Varma, L. Luo, C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, A. Desmaison, C. Balioglu, P. Damania, B. Nguyen, G. Chauhan, Y. Hao, A. Mathews, and S. Li (2023) PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel . Proceedings of the VLDB Endowment 16 ( 12 ), pp. 3848–3860 . Cited by: §V-A .

[55] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He (2020) ZeRO: Memory Optimizations Toward Training Trillion Parameter Models . In SC , Cited by: §V-A .

[56] Meta Platforms, Inc. (2024) Meta AI Acceleration in the Next-Gen Meta MTIA for Recommendation Inference . In IEEE HotChips , Cited by: §VII .

[57] B. Xu, R. Boone, and P. Li (2026) SpikeX: Exploring Accelerator Architecture and Network-Hardware Co-Optimization for Sparse Spiking Neural Networks . arXiv abs/2505.12292 . External Links: Link Cited by: §VII .

[58] C. Zhou, P. Savarese, V. Richard, Z. Hassman, X. Yuan, M. Maire, M. DiBrino, and Y. Li (2024) SySMOL: Co-designing Algorithms and Hardware for Neural Networks with Heterogeneous Precisions . arXiv abs/2311.14114 . External Links: Link Cited by: §VII .

[59] H. Kwon, P. Chatarasi, M. Pellauer, A. Parashar, V. Sarber, and T. Krishna (2019) Understanding Reuse, Performance, and Hardware Cost of DNN Dataflows: A Data-Centric Approach . arXiv abs/1805.02566 . External Links: Link Cited by: §VII .

[60] Q. Anthony, Y. Tokpanov, S. Szot, S. Rajagopal, P. Medepalli, A. Golubeva, V. Shyam, R. Washbourne, R. Iyer, A. Chaurasia, T. Figliolia, X. Yang, A. Sarje, D. Thorstensen, A. Pearson, Z. Grossbart, J. van Patten, E. Barsoum, Z. Gu, Y. Fu, and B. Millidge (2025) Training Foundation Models on a Full-Stack AMD Platform: Compute, Networking, and System Design . arXiv abs/2511.17127 . External Links: Link Cited by: §VII .

[61] M. Naumov, J. Kim, D. Mudigere, S. Sridharan, X. Wang, W. Zhao, S. Yilmaz, C. Kim, H. Yuen, M. Ozdal, K. Nair, I. Gao, B. Su, J. Yang, and M. Smelyanskiy (2020) Deep Learning Training in Facebook Data Centers: Design of Scale-up and Scale-out Systems . arXiv abs/2003.09518 . External Links: Link Cited by: §VII .

[62] G. Sethi, B. Acun, N. Agrawal, C. Kober, Z. Zhang, D. Mudigere, K. Nair, M. Smelyanskiy, C. Wu, and K. Hazelwood (2022) RecShard: Statistical Feature-Based Memory Optimization for Industry-Scale Neural Recommendation . In MLSys , Cited by: §VII .

[63] S. Li, Z. Hu, X. Tang, H. Wang, S. Xu, W. Luo, Y. Li, X. He, and R. Li (2024) Mixed-Precision Embeddings for Large-Scale Recommendation Models . arXiv abs/2409.20305 . External Links: Link Cited by: §VII .