# AMD upstreams Group-GEMM for MoE models to OpenXLA with ROCm and portable Triton backends

> **Key takeaway:** AMD has contributed a production Group-GEMM implementation for Mixture-of-Experts inference to OpenXLA, combining a hipBLASLt backend for ROCm GPUs with a portable Triton backend that works on both CUDA and ROCm.

**Published:** 2026-10-26T18:03:31.841302+00:00 
**Source:** AMD ROCm Official Blog 
**Category:** ai-local-edge 
**Canonical URL:** https://fosswire.org/news/amd-upstreams-group-gemm-for-moe-models-to-openxla-with-rocm.html 

## Summary
AMD engineers have published a detailed technical walkthrough of new Group-GEMM support in OpenXLA that eliminates the compute and memory waste of padded matrix multiplications in Mixture-of-Experts (MoE) large language models. The implementation provides two GPU backends: a vendor-optimized path through hipBLASLt for AMD ROCm, and a portable code-generation backend using Triton that targets both CUDA and ROCm. MoE architectures from DeepSeek, Grok-1, Llama-4, and Gemini route tokens to variable numbers of experts per batch, creating ragged-dot operations where each expert processes a different, runtime-dependent number of tokens. The previous approach padded all expert token blocks to a fixed capacity, burning GPU cycles and memory on empty work. The new XLA lowering converts JAX's ragged-dot primitive directly to Group-GEMM, a batched matrix multiplication where each batch element can have a distinct shape.

The hipBLASLt backend handles the runtime variability by separating XLA's compile-time planning from the actual group sizes discovered at execution, while the Triton backend emits specialized kernels per ragged-dot variant for cross-platform deployment. Performance evaluation on AMD MI300 and MI350 accelerators using MaxText MoE models demonstrates the gains over the padded baseline. The work is published under AMD's ROCm blog with full technical detail, representing a concrete upstream contribution to the open-source ML compiler stack rather than a proprietary optimization.

## XLA lowering architecture and backend design
- **Ragged-dot primitive:** JAX exposes MoE variable-size expert routing through `jax.lax.ragged_dot(lhs, rhs, group_sizes)`, where `group_sizes` determines per-expert token counts only known at runtime, while expert weight dimensions are compile-time constants.
- **hipBLASLt backend (ROCm):** The AMD-specific path bridges XLA's ahead-of-time compilation with runtime-discovered group sizes by targeting hipBLASLt's Group-GEMM APIs, keeping the vendor library responsible for kernel dispatch and memory coalescing on MI300/MI350.
- **Triton backend (portable CUDA+ROCm):** A code-generation alternative emits one specialized Triton kernel per ragged-dot variant, enabling the same MoE workloads to run on NVIDIA and AMD hardware without vendor-library dependencies.
- **Memory and compute efficiency:** Both backends avoid the padded-dot baseline's dual waste: compute cycles spent on zero-valued padding rows and memory bandwidth consumed by padding copies and redundant data movement.
- **Evaluation methodology:** Benchmarking uses MaxText MoE models on AMD MI300 and MI350 accelerators, measuring against the padded baseline to quantify real-world speedups for production-scale LLM inference.

## Why it matters
This contribution strengthens OpenXLA as a vendor-neutral compiler for production AI workloads, giving JAX users a performant, open-source path for MoE inference that is not locked to NVIDIA's CUDA stack. For the ROCm ecosystem specifically, hipBLASLt Group-GEMM support closes a critical gap with cuBLASLt for emerging MoE architectures, while the Triton backend advances portable GPU code generation that benefits both AMD and NVIDIA hardware owners. Linux-based ML deployments, from cloud instances to local inference servers, gain a memory-efficient primitive that directly reduces operational costs for serving models like DeepSeek and Llama-4 variants.