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.6x 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.5x 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.
Figures & tables
Fig. 1 : Total training compute for deep learning models.
Specification
MI350X
MI355X
B200
Architecture
CDNA 4
CDNA 4
Blackwell
Compute Units / Streaming Multiprocessors
256 (8 XCD × 32)
256 (8 XCD × 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
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.
Fig. 2 : Hipification process to enable ML jobs on AMD GPUs.
Fig. 3 : Pipelined GEMM execution. Loads/stores use vector units; compute uses matrix units. The K-loop overlaps load and compute.
Fig. 4 : Wave quantization tail effect. A 4 × 5 output (20 tiles) on 8 CUs requires 3 timesteps with 20mod8=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.
Fig. 5 : GEMM performance on MI350X for production shapes, shown as percentage of peak achievable FLOPS.
Fig. 6 : GEMM performance ( M,N,K dimensions =8192 ) with GPU power excursion across kernel durations and duty ratios on MI350X.
Fig. 7 : Measured AllToAll bandwidth with different group and message sizes relative to the theoretical upper bounds.
Fig. 8 : Eigen decomposition performance.
Fig. 9 : GDPA kernel performance with incremental optimization techniques applied on AMD MI350X, compared to optimal configuration on B200. Lower latency and higher throughput are better.
Fig. 10 : HSTU kernel performance with incremental optimization techniques applied on AMD MI350X, compared to optimal configuration on B200. Lower latency and higher throughput are better.
Fig. 11 : Bounds check forward kernel latency vs. violation rate. Tree reduction eliminates atomic contention, enabling MI350X to outperform B200 at high violation rates.
Fig. 12 : GIS kernel performance. For the bandwidth-bound forward kernel, MI350X exceeds B200. For backward, multi-row packing helps on mixed-dimension configs but cannot fully overcome atomic contention on uniform, high-collision workloads.
B200
MI350X
B
T
L
D
Prec
α
FWD
BWD
FWD
BWD
(ms)
(ms)
(ms)
Before
After
65536
5
55
192
fp16
1.0
1.75
14.40
1.48 (0.85 × )
23.6 (1.64 × )
7.72 (0.54 × )
65536
5
55
192
fp16
1.15
0.84
4.61
0.78 (0.93 × )
9.92 (2.15 × )
3.97 (0.86 × )
65536
10
26
256
fp16
1.0
1.86
15.79
1.84 (0.99 × )
24.4 (1.55 × )
8.06 (0.51 × )
65536
10
26
256
fp16
1.15
0.87
4.85
0.91 (1.05 × )
11.4 (2.35 × )
4.36 (0.90 × )
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), α (Zipf skew factor). Parenthetical values show latency relative to B200 (<1.0 × means MI350X is faster).
Fig. 13 : Triton Inductor kernel latencies by category for three representative models, showing MI350X before/after AMD optimizations compared to B200. Lower latency is better.
Fig. 14 : Training throughput vs. HBM tradeoff on a GPT-2 1.3B model with activation checkpointing.
Fig. 15 : (a) Batch size scaling yields up to 51% throughput improvement. (b) Combined co-design strategies (batch size increase, AC removal, embedding placement optimization) yield 26% speedup.
Model
GPUs
Seq. len
Params
B200 ex/s
MI350X ex/s
MI350X/B200
A
8
512
∼100 M
921,714
754,936
82%
B
8
256
∼100 B
1,592,926
1,810,036
114%
C
64
200
∼200 B
728,001
884,216
121%
D
64
1000
∼300 B
774,324
755,221
98%
E
64
200
∼300 B
283,641
267,635
94%
F
128
1000
∼3.6 T
330,481
304,781
92%
Table III : End-to-end training throughput (examples/second) on MI350X vs B200 across ten production models. Higher is better.
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.
Making deep learning recommendation model (DLRM) training and inference fast and efficient is important. However, this presents three key system challenges - model architecture diversity, kernel primitive diversity, and hardware generation and architecture heterogeneity. This paper presents KernelEvolve-an agentic kernel coding framework-to tackle heterogeneity at-scale for DLRM. KernelEvolve is designed to take kernel specifications as input and automate the process of kernel generation and optimization for recommendation model across heterogeneous hardware architectures. KernelEvolve does so by operating at multiple programming abstractions, from Triton and CuTe DSL to low-level hardware agnostic languages, spanning the full hardware-software optimization stack. The kernel optimization process is described as graph-based search with selection policy, universal operator, fitness function, and termination rule, dynamically adapts to runtime execution context through retrieval-augmented prompt synthesis. We designed, implemented, and deployed KernelEvolve to optimize a wide variety of production recommendation models across generations of NVIDIA and AMD GPUs, as well as Meta's AI accelerators. We validate KernelEvolve on the publicly-available KernelBench suite, achieving 100% pass rate on all 250 problems across three difficulty levels, and 160 PyTorch ATen operators across three heterogeneous hardware platforms, demonstrating 100% correctness. KernelEvolve reduces development time from weeks to hours and achieves substantial performance improvements over PyTorch baselines across diverse production use cases and for heterogeneous AI systems at-scale. Beyond performance efficiency improvements, KernelEvolve significantly mitigates the programmability barrier for new AI hardware by enabling automated kernel generation for in-house developed AI hardware.
We introduce AMDKernelVault, an open HIP and Triton kernel corpus and training framework for recent AMD CDNA GPUs. Existing LLM-based kernel agents are largely CUDA/NVIDIA-centric and often depend on repeated frontier-LLM calls for generation, reflection, and optimization. To address this gap, we develop HIPKernelGen and TritonKernelGen, agent-driven pipelines that transform PyTorch references into HIP or Triton kernels, compile and validate candidates under ROCm, and latency-profile them on AMD hardware. The corpus contains 62,153 execution-verified HIP kernel samples, 2,377 production-grounded ROCm Libraries QA entries, and 39,893 Triton kernels. We further train Qwen3-8B with supervised fine-tuning and execution-aware reinforcement learning as a demonstration of the corpus's utility. Under fixed evaluation budgets, it achieves the highest correctness among the compared models on PyTorch-to-HIP (34.0% Pass@1), TritonBench-G (33.2% Corr@3), and ROCmBench (41.94% Corr@3), but does not uniformly lead compilation or speed metrics. The corpus and documentation are available at https://huggingface.co/datasets/amd/AIG-Datasets, and the associated training and kernel-generation code is available at https://github.com/AMD-AGI/hip_kernel_llm_lab.
Training and deployed inference often cross export, conversion, and platform-specific runtime boundaries. Meganeura asks whether one compact native compiler can span both phases on consumer GPUs. Its typed static graph, automatic differentiation, optimizer, checkpoint, memory planner, and runtime lower specialized programs through Vulkan and Metal. We compare five matched workloads with PyTorch on NVIDIA and AMD discrete GPUs, an AMD APU, Apple silicon, and an Intel iGPU. The protocol separates strict f32 from validated fast paths and gates forward and backward independently. Forty-eight of 50 device-workload-mode cells pass both gates; the other two share one unresolved backward-reference disagreement on a newly supported APU. In strict f32, Meganeura wins 12 of 20 GPU-referenced minimal-latency cells and has a median valid training gap of 1.8x. On the discrete AMD GPU, four of five inference workloads are within 1.10x of compiled ROCm PyTorch and three training workloads are faster. Under accelerated contracts, the worst training gap is 4.6x. Compilation takes 0.1-2.4 seconds versus 6-96 seconds for torch.compile on supported GPU paths; the stripped binary is 13 MiB. Dispatch profiles localize the largest gaps to convolution derivatives and attention backward. A physical Android XR case study transfers a Meganeura-trained decoder into an Adreno/OpenXR application sharing the graphics queue. The results show that general consumer graphics APIs can support a compact shared train-to-deploy stack at useful, sometimes vendor-competitive performance. The measured gaps point to kernel coverage, scheduling, and arithmetic policy rather than an identified API limitation.