Group-GEMM in XLA: Making Mixture-of-Experts Fast and Memory-Efficient

XLA now lowers JAX's `ragged-dot` operation to Group-GEMM through two GPU backends: an AMD-specific backend built on hipBLASLt for ROCm and a new portable Triton code-generation backend that runs on both CUDA and ROCm. The work targets mixture-of-experts layers, which decompose into many small matrix multiplications with runtime-determined row counts, and the post evaluates both backends against a padded dense-batched-matmul baseline on AMD MI300 and MI350 GPUs using MaxText MoE models. The Triton backend emits one kernel per `ragged-dot` variant, giving a portable alternative to the vendor library.

Group-GEMM in XLA: Making Mixture-of-Experts Fast and Memory-Efficient group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient Mixture-of-Experts MoE https://arxiv.org/abs/1701.06538 models have become a standard in most recent LLMs DeepSeek, Grok-1, Llama-4, Gemini . But under the hood, an MoE layer turns into something that GPUs and math libraries have historically been bad at: a large collection of small, independent matrix multiplications, each with a different, data-dependent shape. In JAX, users can describe this collection of matrix multiplications using an operation called ragged-dot , as shown in the following code snippet: php lhs shape: 5, 4 - 5 total input rows/tokens lhs = jnp.array 1.0, 0.5, 0.0, 1.0 , 0.2, 1.0, 0.5, 0.0 , 1.0, 1.0, 1.0, 1.0 , 0.1, 0.2, 0.3, 0.4 , 0.5, 0.5, 0.5, 0.5 , dtype=jnp.float32 rhs shape: 2, 4, 3 - 2 expert weight matrices of shape 4, 3 rhs = jnp.array 1.0, 0.0, 0.1 , 0.1, 1.0, 0.0 , 0.0, 0.1, 1.0 , 1.0, 1.0, 0.0 , Expert 0 0.5, 0.5, 0.5 , 0.2, 0.2, 0.2 , 0.3, 0.3, 0.3 , 0.4, 0.4, 0.4 Expert 1 , dtype=jnp.float32 Group sizes: Expert 0 processes 3 rows, Expert 1 processes 2 rows 3 + 2 = 5 = M group sizes = jnp.array 3, 2 , dtype=jnp.int32 Compute ragged dot product output = jax.lax.ragged dot lhs, rhs, group sizes Hence, even though the number of groups i.e., the number of matrix multiplications is known at compile time since Dimension 0 of rhs and group sizes are equal to the number of matrix multiplications , the number of rows i.e., the number of tokens to be extracted from lhs for each group is determined by the contents of group sizes and is not known until runtime. In other words, while the sizes k and n of each m x k x n matrix multiplication are known at compile time, the size m is not known until runtime, which is highly impractical for efficient computation and optimization. The naive way to run this operation is to pad everything into a single dense batched matmul — which wastes both compute on the padding and memory on the copies . The better way is a primitive called Group-GEMM : a batched matmul in which each matrix in the batch can have a different shape, computed in a single operation. This post walks through how XLA now lowers the ragged-dot operation to real Group-GEMM through two GPU backends: - an AMD-specific backend built on hipBLASLt AMD ROCm , and - a new portable code-generation backend built on Triton that runs on both CUDA and ROCm. By the end, you’ll understand what MoE looks like to a compiler, what ragged-dot is, why Group-GEMM is the natural solution, and how each of the two backends makes it work on the GPU. The post is organized in five parts: - Part 1 — Why MoE Turns Into a “Ragged” Matmul Problem part-1-why-moe-turns-into-a-ragged-matmul-problem sets up the problem: how an MoE layer becomes a ragged matmul, why the padded baseline is wasteful, and why Group-GEMM is the right primitive. - Part 2 — Backend 1: hipBLASLt Group-GEMM AMD ROCm part-2-backend-1-hipblaslt-group-gemm-amd-rocm covers the vendor-library backend, and how it bridges XLA’s compile-time planning with runtime group sizes. - Part 3 — Backend 2: Triton Ragged-Dot Portable CUDA + ROCm part-3-backend-2-triton-ragged-dot-portable-cuda-rocm covers the portable code-generation backend, with one emitted kernel per ragged-dot variant. - Part 4 — Performance Evaluation part-4-performance-evaluation measures both backends against the padded baseline on MI300 and MI350, using MaxText MoE models. - Part 5 — Conclusion and Next Steps part-5-conclusion-and-next-steps wraps up and outlines what comes next. Part 1 — Why MoE Turns Into a “Ragged” Matmul Problem part-1-why-moe-turns-into-a-ragged-matmul-problem A Quick Refresher: Dense FFN vs. Mixture-of-Experts a-quick-refresher-dense-ffn-vs-mixture-of-experts In a standard Transformer, every token passes through the same Feed-Forward Network FFN . If you want the model to “know more,” you make that FFN bigger — but then every token pays for the extra parameters, and compute grows with model size. Mixture-of-Experts breaks that link. Instead of one big FFN, an MoE layer has many smaller FFNs called experts . A small router a gating network looks at each token and picks the top-k experts that should process it often k=1 or k=2 . Only those experts run for that token. The result: you can add a lot of parameters many experts while keeping the per-token compute roughly constant, because each token only visits a couple of experts. Figure 1 below summarizes this flow. The Consequence for a Compiler: Variable-Size Groups the-consequence-for-a-compiler-variable-size-groups Here’s the crucial systems detail. Because the router assigns each token to experts based on the token’s content , the number of tokens that land on each expert is different, and it is only known at runtime. One expert might receive 2,000 tokens this step; its neighbor might receive 50; the split changes with every batch. If we gather all the tokens routed to the same expert into a contiguous block, an MoE layer becomes: - A stack of token rows, partitioned into G contiguous groups one per expert , where group g has m g rows and the m g values vary and are runtime-dependent. - Each expert g has its own weight matrix, so we need to compute output g = tokens g × W g , i.e. a matmul of shape m g, k × k, n . Figure 2 illustrates groups with unequal sizes: each color represents an expert. The mental model to keep for the rest of this post: An MoE layer is a batch of matmuls whose left-hand-side rows are split into contiguous, variable-size groups. That “variable-size groups” property is what we call ragged . The Naive Support, and Why It Hurts: Padded-Dot the-naive-support-and-why-it-hurts-padded-dot GPUs are well suited to regular, dense shapes. So the naive solution is to force the ragged problem back into a dense one: 1. Pick a fixed capacity C a maximum number of tokens per expert . 2. Pad every expert’s token block up to C rows. 3. Now every block has the same shape, so we can run one big batched GEMM over E × C, k × k, n . Figure 3 illustrates this padding process: every expert’s block is padded to the same C rows, where the grey padding is compute the GPU performs for nothing. This works, but it is wasteful in two ways: - Wasted compute. Any expert that received fewer than C tokens still does a full C -row matmul; the padding rows are pure overhead. - Wasted memory and bandwidth. The compact ragged input has to be expanded into E padded matrices, and the outputs have to be compacted again afterwards. Those scatter/gather and pre/post-processing steps cost extra memory and memory traffic. This “expand into padded matrices” approach is exactly the XLA baseline we’ll compare against later. The Ragged-Dot Operation the-ragged-dot-operation Rather than having users replace their MoE with a dense problem, JAX lets you describe the ragged computation directly with ragged dot . The basic form: jax.lax.ragged dot lhs, rhs, group sizes, ... with the shapes: - lhs is 2D: m, k — all the tokens, stacked, in group order. - rhs is 3D: group, k, n — one weight matrix per group per expert . - group sizes is 1D: group — how many rows of lhs belong to each group. - the output is 2D. The group sizes operand is the key: it is a small vector, known only at runtime, that partitions the m rows of lhs into contiguous groups. For example group sizes = m g0, m g1, m g2 means the first m g0 rows are group 0, the next m g1 rows are group 1, and so on. Note that jax.lax.ragged dot is a JAX-level operation. When the JAX program is compiled, it is converted into the HLO ragged-dot operation kRaggedDot , which is the actual input XLA receives and lowers. Figure 4 illustrates the operands of a ragged-dot. JAX also offers a more general form: jax.lax.ragged dot general lhs, rhs, group sizes, ragged dimension, ... which lets you choose which dimension is the ragged one. That choice yields the three variants listed in Table 1 and drawn in Figure 5: Table 1: The three ragged-dot variants exposed by jax.lax.ragged dot general . | Variant | Ragged dimension | What it means | Group index in output? | |---|---|---|---| | Non-contracting m | the “rows” dimension | the classic MoE case | No | | Contracting k | the reduction dimension | groups along the contracted axis | Yes | | Batch b | a batch dimension | effectively a regular batched GEMM | — | Inside XLA these correspond to the kRaggedDot variants kRaggedNonContracting , kRaggedContracting , and kRaggedBatch . For MoE, the one that matters most is non-contracting , where the rows tokens are ragged. Why Group-GEMM Is the Suitable Solution why-group-gemm-is-the-suitable-solution Group-GEMM is an API for matrix multiplication that generalizes the usual batched GEMM: instead of requiring every matrix in the batch to share the same shape, it lets each one differ — in size, transposition, activation, and scaling factors. In other words: Group-GEMM = a set of independent GEMMs of different shapes, computed in a single operation. That is a perfect match for ragged-dot , as ragged-dot is nothing more than a set of independent GEMMs of different shapes where the pointer to the input data can be calculated from the base pointer of the operand buffer + an offset derived at runtime from the group-size operands. Hence, Group-GEMM can take care of “unpacking” the ragged operand and computing all of the per-expert products in one kernel launch, with no padding, as shown in Figure 6. Compared to the padded-dot baseline, this directly removes both costs described in The Naive Support, and Why It Hurts: Padded-Dot the-naive-support-and-why-it-hurts-padded-dot : - No padding FLOPs — each group does exactly the work it needs. - No expansion/copies — the compact ragged layout is consumed as-is. The rest of this post is about the two backends XLA uses to turn ragged-dot into real Group-GEMM for AMD GPUs. Part 2 — Backend 1: hipBLASLt Group-GEMM AMD ROCm part-2-backend-1-hipblaslt-group-gemm-amd-rocm Overview overview AMD’s hipBLASLt library has supported Group-GEMM https://rocm.docs.amd.com/projects/hipBLASLt/en/develop/reference/ext-reference.html grouped-gemm since version 1.2.2. XLA already had a hipBLASLt backend for ordinary dot operations, so the natural move was to extend that existing backend to also lower ragged-dot down to hipBLASLt Group-GEMM. Figure 7 shows the lowering path of a plain dot , and Figure 8 the ragged-dot path added alongside it. The implementation reuses the existing matmul backend machinery and simply derives a GroupedMatmulPlan class from MatmulPlan , including the additional parameters required to run the hipBLASLt Group-GEMM. struct GroupedGemmConfig { uint64 t m, n, k; uint64 t batch count; uint64 t group count; int64 t lhs leading dim stride; int64 t rhs leading dim stride; int64 t c leading dim stride; int64 t output leading dim stride; blas::Transpose trans a, trans b; bool must swap operands; xla::complex128 alpha; double beta; blas::DataType type a, type b, type c, type d; int64 t stride ragged dim; int64 t stride group dim; int64 t c stride ragged dim; int64 t output stride ragged dim; // PrecisionConfig-level algorithm xla::PrecisionConfig::Algorithm precision algorithm; int64 t compute precision; RaggedDotMode ragged mode; std::optional<blas::ComputationType compute type; static absl::StatusOr<GroupedGemmConfig FromProto const xla::GroupedGemmConfigProto& proto ; xla::GroupedGemmConfigProto ToProto const; }; class GroupedMatmulPlan : public gpu::BlasLt::MatmulPlan { public: friend class BlasLt; GroupedMatmulPlan const BlasLt& blas lt, const gpu::GroupedGemmConfig& cfg : blas lt blas lt , cfg cfg {} ~GroupedMatmulPlan override = default; absl::Status ExecuteOnStream Stream stream, const gpu::BlasLt::MemoryArgs& args, blas::ProfileResult profile result const override; absl::StatusOr<std::vector<MatmulAlgorithm GetAlgorithms size t max algorithm count, size t max workspace size const override; absl::Status SetAlgorithm const MatmulAlgorithm& algorithm override {...} private: ... const BlasLt& blas lt ; gpu::GroupedGemmConfig cfg ; Epilogue epilogue = Epilogue::kDefault; std::unique ptr<hipblaslt ext::GroupedGemm grouped gemm ; mutable std::optional<hipblasLtMatmulAlgo t algorithm ; mutable bool algorithm dirty = false; mutable DeviceAddressBase saved address workspace {}; // Saved default activation parameters from hipBLASLt int32 t activation type = 0; int8 t bias type = 0; }; // class GroupedMatmulPlan source code available GroupedGemmConfig https://github.com/openxla/xla/blob/2135b8f9071c430ea7afc9940c4ede41071c81c2/xla/stream executor/gpu/gpu blas lt.h L165-L191 and GroupedMatmulPlan https://github.com/openxla/xla/blob/e2451066b70ee17b2b374153a32120df332f5478/xla/stream executor/rocm/hip blas lt.h L154-L194 The hipBLASLt Group-GEMM API the-hipblaslt-group-gemm-api Launching a Group-GEMM with hipBLASLt follows the fixed sequence of steps shown in Figure 9: The important concept here is user arguments . hipBLASLt describes each individual matmul in the group with its own descriptor — a small structure that says how much data to read and from where : the matmul’s sizes, its input pointers, and its strides. There is one descriptor per matmul, and all of them must live in GPU global memory at the moment the Group-GEMM launches. See here https://rocm.docs.amd.com/projects/hipBLASLt/en/develop/reference/ext-reference.html the-userarguments-structure for more details on the user argument structure definition . The Compile-Time vs. Runtime Problem the-compile-time-vs-runtime-problem Here is the central difficulty. XLA is an ahead-of-time compiler: it wants to plan and configure the operation at compile time. But the whole point of ragged-dot is that the group sizes — and therefore each matmul’s dimensions — are only known at runtime. So, we need a way to delay the specification of the matmul’s dimensions to runtime. Fortunately, hipBLASLt provides a neat trick to bridge the gap. The extension allows the operation sizes m, n, k, and batch to be redefined at runtime, in device memory, by modifying the user arguments. Only one condition must be guaranteed: the total problem size defined at compile time must be equal to or greater than the sum of the sizes of all the individual GEMMs that are actually executed. You can see the documentation https://rocm.docs.amd.com/projects/hipBLASLt/en/develop/reference/ext-reference.html grouped-gemm-fixed-mk for more details on this. In XLA, the definition of user arguments is therefore delayed until runtime. Just before launching the Group-GEMM, XLA executes small update kernels on the GPU that compute and populate the descriptors with the actual sizes of the groups contained in group sizes . These kernels proceed as sketched in Figure 10: Each thread is responsible for one matmul’s descriptor. It uses a prefix sum a block-scan over the group sizes to figure out where its group starts in the packed buffers, computes the correct sizes, strides, scales, and data pointers, and writes them to a shared-memory staging buffer. After a synchronization, all threads cooperatively copy the completed set of descriptors into global memory, where hipBLASLt expects to find them at launch. Autotuning autotuning hipBLASLt’s Group-GEMM API can enumerate all the algorithms available for a given problem, and XLA’s autotuner benchmarks them to pick the fastest, as shown in Figure 11: However, ragged-dot needs one special accommodation. By default, the autotuner profiles candidates on buffers filled with random values. But for Group-GEMM, group sizes is not ordinary data: its values are used to compute memory offsets and pointers. Random group sizes would produce invalid offsets and corrupt or crash the memory accesses. To address this problem, the autotuner was enhanced to seed the group sizes buffer with consistent values before profiling. More specifically, the sum of the group sizes values must be equal to the size M of the ragged dimension. For example, for a ragged dimension of size M = 1024 split across G = 4 groups, the autotuner seeds group sizes = 256, 256, 256, 256 256 × 4 = 1024 = M . When M is not divisible by G , the remainder is spread over the first groups e.g. M = 1002 → 251, 251, 250, 250 . In both cases sum group sizes == M . Enforcing sum group sizes == M guarantees two things: 1. In-bounds pointers. Each group’s data pointer is derived from a prefix sum over group sizes group ptr = base ptr + prefix sum group sizes i stride ; keeping the sum equal to M ensures every derived pointer stays inside the operand buffer. 2. A fully-written output. ragged-dot allows sum group sizes < M , but then part of the output buffer is left unwritten — during autotuning this would leave uninitialized garbage values that make result comparison unreliable. Seeding the sum to exactly M ensures the entire output is written, so profiled candidates can be compared safely. Part 3 — Backend 2: Triton Ragged-Dot Portable CUDA + ROCm part-3-backend-2-triton-ragged-dot-portable-cuda-rocm Motivation motivation Even though the hipBLASLt path provides impressive performance see Part 4 — Performance Evaluation part-4-performance-evaluation , it is ROCm-specific and not even available on every AMD GPU target. On the other hand, Triton is a generic backend that can be lowered to most GPU targets. In addition, Triton kernels can more easily be fused together, making it a privileged backend for the XLA compiler, which aims to run on any GPU target. Consequently, we decided to propose a Triton backend as an alternative to efficiently handle the ragged-dot operation in XLA. In a nutshell, the Triton backend consists of lowering the ragged-dot operation into a Group-GEMM Triton kernel, through XLA’s Triton tiling-and-emission pipeline. The Compilation Pipeline the-compilation-pipeline Figure 12 shows the main steps that have been enhanced to enable the lowering of the ragged-dot op: 1. Backend selection. XLA decides which path serves this ragged-dot , following a priority order: cuDNN CUDA → hipBLASLt ROCm → Triton → a default fallback. When the Triton path is chosen, XLA wraps the ragged-dot in a custom Triton fusion and seeds it with default tile sizes that the autotuner can later override. Note that the backend selection order is currently fixed by static rules, but the next step is to let the XLA autotuner choose the best backend for each specific problem’s inputs. 2. Support validation. Before committing, XLA checks the ragged-dot is actually supported: it must be a ragged-dot with exactly one ragged left-hand-side dimension, the tiling-propagation feature must be enabled, and the element types must be supported. 1. Scheduling. XLA decides how the ragged-dot ’s logical dimensions M, N, G, K map onto the GPU grid — which become parallel grid axes one program per tile, run concurrently and which stay sequential a loop inside a program, typically the K reduction .This step includes an M/N-swap heuristic. This heuristic already existed for the ordinary dot and scaled dot but has been extended for ragged-dot . It is therefore useful to briefly recall what it does for a plain GEMM. Plain GEMM. For a regular LHS M,K × RHS K,N → out M,N , the M and N tiles form a 2-D grid that XLA linearizes into a 1-D stream of program IDs. Whichever axis is placed outer varies slowly, and the inner one varies fastest across consecutive programs. The heuristic exploits this to maximize L2 reuse: if the LHS activations is smaller in bytes than the RHS weights , it swaps M and N so that N is traversed slowly — keeping the small LHS tile-row resident in L2 while streaming across the many N tiles. Concretely it compares ByteSizeOf lhs against ByteSizeOf rhs and, when the LHS is smaller, emits a permutation that swaps the M and N grid axes. Ragged-dot. A ragged-dot is not one big matmul but G smaller ones, so the raw operand byte-sizes are the wrong thing to compare. The ragged version keeps the same mechanism permute the parallel grid axes before the program-id is delinearized but changes the decision : it uses the compile-time estimate M avg = M total / G — the average number of rows per group — as a proxy for how tall each group’s LHS activation slice really is the true per-group sizes are only known at runtime . 2. Autotuning. The pipeline generates candidate tiling configurations specialized for ragged-dot and benchmarks them more on the ragged-specific profiling wrinkle in Autotuning and the group sizes Seeding Problem autotuning-and-the-group-sizes-seeding-problem . 3. Emission. Finally, XLA emits the actual Triton kernel for the selected variant. What the Emitted Triton Kernel Does, per Variant what-the-emitted-triton-kernel-does-per-variant Now, let’s take a closer look at the Triton kernels generated to execute the generic ragged-dot operations, which encompass three variants depending on which dimension is the ragged dimension. As we argued before, the Group-GEMM algorithm is best suited to handle ragged-dot operations. We therefore generate different kernels for each variant, which all implement the Group-GEMM algorithm but differ in the way they use the group sizes information to calculate pointer offsets. The tricky point with a Triton kernel is that each kernel defines the program running on one thread block. Then, the GPU launches a grid of programs one per output tile ; so each program needs to figure out which slice of data it owns, perform the matmul for that slice, and write its output tile. So, what differs between the variants is: - where the raggedness lives, - whether group sizes is loaded at all, - how the group offset is used. Note that how these per-tile programs are laid out across the grid is a scheduling decision, not part of the per-program arithmetic above. Two complementary mechanisms shape that layout: - the M/N-swap heuristic from The Compilation Pipeline the-compilation-pipeline , which decides which grid axis varies slowly so that the smaller per-group operand stays resident in L2; and - a tile-id remapping in the emitter specific to the non-contracting/MoE variant , which reorders program IDs so that tiles belonging to the same group run close together, keeping that group’s RHS weight matrix hot in L2 cache. Note that this optimization was replicated for the ROCm AITER https://github.com/ROCm/aiter Group-GEMM implementation . Together these keep each expert’s weights cache-resident while its tiles are being computed. Non-Contracting the MoE Case non-contracting-the-moe-case Shape: LHS M total, K × RHS G, K, N → out M total, N . The ragged axis is M the stacked tokens . Every output row m belongs to exactly one group g = group m , and that group selects which weight matrix RHS g is used. The emitted kernel has conceptually this signature: python @triton.jit def ragged dot non contracting lhs ptr, M total, K all tokens, stacked in group order rhs ptr, G, K, N one weight matrix per group expert out ptr, M total, N output, same row layout as lhs group sizes ptr, G rows per group; sum == M total M total, K, N, G, BLOCK M: tl.constexpr, BLOCK N: tl.constexpr, BLOCK K: tl.constexpr, : pid m = tl.program id 0 which M-tile this program owns pid n = tl.program id 1 which N-tile this program owns row start = pid m BLOCK M first LHS row of this tile in 0, M total ... The subtle part of the kernel is to “know” which portion of the data the kernel is expected to compute. Indeed, the program only knows its M-tile position — i.e. which band of BLOCK M rows of the stacked LHS it is responsible for row start = pid m BLOCK M . What it does not directly know is which expert those rows came from, because the group boundaries depend on the runtime group sizes . To recover the group it runs a prefix sum over group sizes : group sizes = m0, m1, m2, … prefix exclusive = 0, m0, m0+m1, … group start rows prefix inclusive = m0, m0+m1, m0+m1+m2, … group end rows The exclusive prefix sum gives the start row of every group; the inclusive one gives the end row. The program’s group g is simply the group whose start, end interval contains its row start — i.e. the largest g such that prefix exclusive g <= row start . From that same table it reads group start row = prefix exclusive g , which tells it the offset of its group inside the stacked buffer needed to address LHS / out correctly, and to mask the partial tile at the group’s boundary . Table 2 below works through a tiny example — group sizes = 3, 2, 4 so M total = 9 and BLOCK M = 3 : Table 2: Worked example of the group lookup for group sizes = 3, 2, 4 with BLOCK M = 3 . | pid m | row start | falls in group g | group start row | uses weights | |---|---|---|---|---| | 0 | 0 | 0 rows 0–2 | 0 | RHS 0 | | 1 | 3 | 1 rows 3–4 | 3 | RHS 1 | | 2 | 6 | 2 rows 5–8 | 5 | RHS 2 | When a tile straddles a group boundary, the emitter masks the rows outside the group so each row is only ever attributed to a single expert. Per program: 1. Compute the prefix sum over group sizes ; from its own M-tile position row start = pid m BLOCK M find which group g it falls into the group whose start, end row interval contains row start and read that group’s start row group start row = prefix exclusive g . 2. Use the group index g to select the matching RHS slice RHS g, :, : , i.e. that expert’s weight matrix. 3. Accumulate a tiled dot product over the contraction dimension K the usual GEMM inner loop , producing an BLOCK M, BLOCK N output tile, and write it to out . group sizes is loaded — it drives both the group selection and the row offset. Figure 14 summarizes this program flow. Contracting contracting Shape: LHS M, K total × RHS K total, N → out G, M, N . The ragged axis is the K contracted dimension, and — unlike the non-contracting case — the group index shows up in the output: instead of one contraction over all of K, each group g contributes its own result matrix by reducing only over the slice of the contracted axis that belongs to it. The emitted kernel has conceptually this signature: python @triton.jit def ragged dot contracting lhs ptr, M, K total contracted axis is ragged rhs ptr, K total, N out ptr, G, M, N one output matrix per group group sizes ptr, G K-slice size per group; sum == K total M, K total, N, G, BLOCK M: tl.constexpr, BLOCK N: tl.constexpr, BLOCK K: tl.constexpr, : pid g = tl.program id 0 which group this program owns → output slab g pid m = tl.program id 1 which M-tile pid n = tl.program id 2 which N-tile ... The key structural difference from the non-contracting case is that the group dimension does exist in the output, and that changes how a program finds its group: - For the non-contracting kernel, the group index does not appear in the output shape — all groups write into the same M total, N buffer, stacked by row. A program is assigned an M-tile a band of rows and must find out which group those rows belong to by scanning the prefix sum. - For the contracting kernel, the group index is an output dimension out G, M, N . Each group produces its own M, N output, so the groups can simply be distributed over the grid. Consequently, XLA makes G a parallel grid axis , and each program is handed its group directly as pid g . To illustrate this, Figure 16 shows the same pool of blocks under the two kernels. In the non-contracting case a block is placed by its M-tile, N-tile position and only afterwards figures out its group g from the prefix sum — so the group is a runtime label, not a launch coordinate. In the contracting case the G axis is part of the launch grid, so each block is born owning a specific group g = pid g — the groups are literally spread across blocks by the scheduler. What the program needs from group sizes is where along K its group’s slice begins and ends — the ragged K partition is still runtime data. It reads that from the same prefix-sum table: group sizes = k0, k1, k2, … K-rows per group prefix exclusive = 0, k0, k0+k1, … k start for each group prefix inclusive = k0, k0+k1, k0+k1+k2, … k end for each group For its own pid g , the program reads k start = prefix exclusive pid g and k end = prefix inclusive pid g , then contracts over only k start, k end of the K axis. Per program: 1. Take its group g = pid g from the grid, and read k start / k end from the prefix sum over group sizes — this time the offset locates the group’s slice along the contracted axis. 2. Accumulate a partial dot over just that k start, k end slice of the contracted axis. 3. Write to out g, :, : — the output position indexed by the group g . group sizes is loaded: it provides the per-group start/end offsets along K. Figure 17 summarizes this program flow. Batch batch Shape: LHS B, M, K × RHS B, K, N → out B, M, N — an ordinary batched GEMM. The ragged dimension is a batch dimension, but because every batch element is fully computed there are no data-dependent group offsets to work out: each program simply owns one batch, M-tile, N-tile and does a full contraction over K. The emitted kernel has conceptually this signature: python @triton.jit def ragged dot batch lhs ptr, B, M, K rhs ptr, B, K, N out ptr, B, M, N B, M, K, N, BLOCK M: tl.constexpr, BLOCK N: tl.constexpr, BLOCK K: tl.constexpr, : pid b = tl.program id 0 which batch element this program owns pid m = tl.program id 1 which M-tile pid n = tl.program id 2 which N-tile no prefix sum: batch index and tile offsets are enough to address the data ... Because the batch index pid b comes straight from the grid and every batch element has the same, statically-known shape, the program can compute all of its pointers from pid b , pid m , pid n and the compile-time strides alone. There is no group boundary to discover, so group sizes is not even loaded here. It influences scheduling only, but not the arithmetic. Per program: 1. Take its batch index b = pid b from the grid and its pid m , pid n tile positions . 2. Accumulate a tiled dot product over K between LHS b and RHS b . 3. Write the BLOCK M, BLOCK N result to out b, M-tile, N-tile . group sizes is not loaded: it is used only as a scheduling hint. Figure 19 summarizes this program flow. Autotuning and the group sizes Seeding Problem autotuning-and-the-group-sizes-seeding-problem This backend runs into the very same issue as the hipBLASLt autotuner in Autotuning autotuning . During autotuning, candidate kernels run on synthetic input buffers, which are filled with random bytes by default. That’s harmless for a normal GEMM. But group sizes is made of control values that determine loop bounds and memory offsets. The fix mirrors the hipBLASLt one autotuning : before profiling, the profiler writes balanced, valid data into the group sizes buffer — every group gets an equal share of the rows, and the last group absorbs any remainder — guaranteeing the sizes sum exactly to the ragged dimension. Part 4 — Performance Evaluation part-4-performance-evaluation Every comparison is made against the padded-dot baseline from The Naive Support, and Why It Hurts: Padded-Dot the-naive-support-and-why-it-hurts-padded-dot : the ragged-dot is rewritten into regular dot operations by expanding the ragged input of G groups into G padded matrices, with the associated pre- and post-processing of the input and output matrices. This is the “before” picture that Group-GEMM aims to beat, recalled in Figure 20. Testbed Characteristics testbed-characteristics The performance evaluations have been performed on two AMD GPU targets, MI300 gfx942 and MI350 gfx950 . The software configuration was as follows: - Operating System : Ubuntu 22.04 - ROCm version : 7.2.4 - JAX version : upstream https://github.com/jax-ml/jax commit 0eb0f676ba22b15b9dfe29518ab0d40277b7138d - XLA version : ROCm https://github.com/ROCm/xla branch pr/ragged-dot-6-e2e-tests upstream commit e5d008bb03d9b133b0881daf3d108adea8e6625b + ragged-dot Triton support commits - MaxText version : ROCm https://github.com/ROCm/maxtext commit fb25ae8c850fa2b57dd5eb7ebfbe399f08534e69 Note that ROCm version 7.2.4 only offers Group-GEMM support for bf16 on MI300. Support for bf16 on MI350 was not yet available when this blog post was written. The ragged-dot Triton backend is not enabled by default and must be explicitly enabled using the following flag xla gpu experimental triton ragged dot=true . Furthermore, note that this backend is only supported with the Triton experimental tiling feature. As a consequence, the flag xla gpu experimental enable tiling propagation=true must be set alongside the xla gpu experimental triton ragged dot flag. Last, the default Triton autotuning is not yet fine-tuned for ragged-dot , so exhaustive search was used for the evaluation xla gpu exhaustive tiling search=true . Use Case: MoE Model Training Using MaxText use-case-moe-model-training-using-maxtext MaxText https://github.com/AI-Hypercomputer/maxtext is a high-performance, highly scalable, open-source LLM library and reference implementation written in Python/JAX. Because it is written on top of JAX, its MoE layers ultimately lower to JAX primitives — and, on the sparse path, to the jax.lax.ragged dot op that the Group-GEMM backends described above accelerate. This makes MaxText an ideal framework in which to evaluate the new ragged-dot backends on realistic MoE models. Two MoE Modes: Drop vs. Dropless two-moe-modes-drop-vs-dropless Before diving into how MaxText runs the expert FFNs, it helps to recall that there are two ways of handling the per-expert token counts of an MoE layer — and the choice between them is an accuracy/efficiency trade-off. - Drop mode. Each expert is given a fixed capacity expert capacity expert capacity = ceil tokens per batch / E × capacity factor . Tokens routed beyond an expert’s capacity are simply dropped they skip the FFN and contribute nothing , and under-filled experts are padded up to expert capacity . This keeps every expert’s matmul a regular, statically-shaped expert capacity, K × K, N , which is hardware-friendly, but it wastes FLOPs on the padding of under-filled experts and loses information on the dropped tokens of over-filled ones. - Dropless mode. No token is ever dropped and no capacity padding is introduced: every routed token, expert pair is computed exactly. This preserves the full model quality, but it requires handling the ragged , data-dependent per-expert sizes — either by physically gathering tokens per expert into one ragged matmul the block-sparse formulation introduced by MegaBlocks \ Gale et al., 2022\ https://arxiv.org/abs/2211.15841 , which is exactly the ragged-dot /Group-GEMM path this post is about , or by running a full dense compute over all tokens and masking afterwards. Figure 21 illustrates the buffers the two modes produce. The rest of this section looks at how MaxText realizes these two modes through its three concrete compute paths. For a complementary, MaxText-focused deep dive into the dropless MoE path and its optimization, see the AMD ROCm blog Dropless MoE in MaxText https://rocm.blogs.amd.com/software-tools-optimization/maxtext-dropless-moe/README.html . The Three MoE Compute Paths in MaxText the-three-moe-compute-paths-in-maxtext Before diving into the models, it is worth noticing that ragged-dot is only one of several ways to run the expert FFNs and combine their outputs. If you look at the MaxText framework https://github.com/ROCm/maxtext/blob/fb25ae8c850fa2b57dd5eb7ebfbe399f08534e69/src/maxtext/layers/moe.py L3144-L3170 , it exposes three distinct MoE compute paths — two of them implement the dropless mode the sparse / ragged-dot path and the dense mask-and-combine path , while one implements the drop mode the capacity-based dense dispatch/combine path . The ragged-dot /Group-GEMM work described so far powers just one of the two dropless paths. All three share the same front-end — a router gating network picks, for every token, its top-k experts and a weight per selected expert: gate logits: B, S, E top k weights, top k indices = jax.lax.top k gate logits, num experts per tok B, S, k However, they differ in the way the selected token, expert pairs are turned into matmuls, and in the pre-/post-processing that surrounds those matmuls. From a user’s point of view, two config knobs allow choosing the path: - sparse matmul — True picks the sparse ragged-dot path the dropless, block-sparse formulation of the MoE computation ; False picks one of the two dense paths. - capacity factor — inside the dense path, 0 enables token dropping fixed per-expert capacity ; <= 0 keeps it dropless. Figure 22 shows how these two knobs select a path. Table 3 summarizes the trade-offs before we look at each path in detail: Table 3: The three MaxText MoE compute paths and their trade-offs. | | Path 1 — Sparse ragged-dot | Path 2 — Dense dropless | Path 3 — Dense drop | |---|---|---|---| | config | sparse matmul = True | sparse matmul = False , capacity factor ≤ 0 | sparse matmul = False , capacity factor 0 | | expert op | jax.lax.ragged dot | dense per-expert einsum over all tokens | dense capacity-padded batched einsum | | pre-processing | permute sort tokens by expert + bincount group sizes | none broadcast all tokens to all experts | build dispatch mask one-hot into capacity slots , drop overflow | | post-processing | un-permute scatter back + weighted sum | weighted combine BSEM,BSE→BSM , zeros non-selected | combine mask gathers from capacity slots back to tokens | | compute cost | ≈ M × K × N only routed rows | ≈ E × B·S × K × N every token × every expert | ≈ E × C × K × N C = capacity | | drops tokens? | no | no | yes overflow beyond C | | padding | none | none but full dense compute | per-expert capacity padding | Path 1 — Sparse / Ragged-Dot Dropless path-1-sparse-ragged-dot-dropless Selected by sparse matmul = True . This is the path that generates jax.lax.ragged dot ops which are then lowered by XLA. The idea is to physically gather the tokens by expert so that the expert FFNs become one big ragged matmul. Pre-processing — permute. The router’s top k indices are used to sort all M = B·S·top k token-copies into contiguous per-expert groups, and a bincount produces the group sizes vector real routed counts, summing to M : permute : sort token rows by their assigned expert, count per-expert sizes sorted selected experts = jnp.argsort jnp.ravel selected experts sorted inputs = sort activations replicated inputs, sorted selected experts M, K group sizes = jnp.bincount jnp.ravel selected experts , length=num experts E , sums to M Compute — Group-GEMM. The gathered tokens and the per-expert weights feed directly into jax.lax.ragged dot , which can be later lowered to hipBLASLt / Triton Group-GEMM by XLA. gmm : the expert FFN as a single ragged/group matmul out = jax.lax.ragged dot lhs=sorted inputs, M, K tokens, grouped by expert rhs=w, E, K, N expert weights group sizes=group sizes, E preferred element type=dtype Post-processing — un-permute. The results are scattered back to the original token order and the top-k contributions are summed with the router weights: unpermute : invert the sort, then weighted-combine the top-k expert outputs unsorted = sort activations out, jnp.argsort sorted selected experts output = jnp.einsum "BKE,BK - BE", unsorted.reshape -1, top k, N , top k weights The whole path is summarized in Figure 23. Cost: only the real routed rows are ever multiplied — compute is ≈ M × K × N with no padding and no dropping. So, regarding the matrix multiplication itself, this path with Group-GEMM backends provides the most efficient solution. However, we have to pay for the sort/scatter permute/un-permute which is needed to put the data in a form that can be processed by the Group-GEMM backends. Path 2 — Dense Dropless Mask-and-Combine path-2-dense-dropless-mask-and-combine Selected by sparse matmul = False and capacity factor ≤ 0 . In this path, there is no gather at all: every token is sent through every expert, and the router weights zero for non-selected experts mask out the unwanted results at the end. Pre-processing — none beyond turning top-k into a dense B,S,E weight matrix, 0 where an expert was not selected : weights = reshape and update weights top k weights, top k indices B,S,E , 0 off-selection Compute — full dense einsums. Each expert processes the entire token batch: php up / gate: every token × every expert layer w0 = jnp.einsum "BSM,EMH - BSEH", inputs, w0 kernel B,S,E,H layer w1 = jnp.einsum "BSM,EMH - BSEH", inputs, w1 kernel inter = activation layer w0 layer w1 down inter = jnp.einsum "BSEH,EHM - BSEM", inter, wo kernel B,S,E,M Post-processing — weighted combine. The router weights collapse the expert axis, keeping only each token’s selected experts: php output = jnp.einsum "BSEM,BSE - BSM", inter, weights B,S,M The whole path is summarized in Figure 24. Cost: dropless and permutation-free, but compute is ≈ E × B·S × K × N — every token pays for every expert, so it scales with the full expert count. This is simple and hardware-friendly plain dense matmuls, no ragged handling but becomes prohibitively expensive as E grows. Path 3 — Dense Drop Capacity Dispatch/Combine path-3-dense-drop-capacity-dispatch-combine Selected by sparse matmul = False and capacity factor 0 . This is the classic fixed-capacity MoE: each expert gets a static slot of C = ceil tokens per batch / E × capacity factor rows; tokens beyond C are dropped. Pre-processing — dispatch mask with dropping . generate masks builds a one-hot dispatch mask that assigns each kept token, expert pair to a capacity slot and truncates any overflow: expert capacity = ceil tokens per batch / num experts capacity factor ... one-hot into B,S,E,C ; positions past expert capacity are dropped ... dispatch mask, combine mask = generate masks top k indices, weights Compute — capacity-padded batched einsum. Tokens are scattered into the E, B, C, M capacity buffer, run through the experts, and the result stays in capacity layout: php dispatch = jnp.einsum "BSM,BSEC - EBCM", inputs, dispatch mask scatter to slots layer w0 = jnp.einsum "EBCM,EMH - EBCH", dispatch, w0 kernel up/gate inter = activation layer w0 jnp.einsum "EBCM,EMH - EBCH", dispatch, w1 kernel inter = jnp.einsum "EBCH,EHM - EBCM", inter, wo kernel down Post-processing — combine mask. The combine mask gathers each token’s outputs back from its capacity slots and applies the router weights: php output = jnp.einsum "EBCM,BSEC - BSM", inter, combine mask gather back to tokens The whole path is summarized in Figure 25. Cost: compute is ≈ E × C × K × N . Because C is fixed, under-filled experts waste FLOPs on padding rows and over-filled experts drop tokens an accuracy trade-off . The capacity buffer also inflates the activation memory. What About Performance? what-about-performance In theory, the ragged-dot /Group-GEMM path: - vs. Path 2 dense dropless — Group-GEMM computes only the routed rows M × K × N instead of every-token-times-every-expert E × B·S × K × N , so it wins by a factor that grows with the number of experts — at the cost of the permute/un-permute sort. - vs. Path 3 dense drop — Group-GEMM needs no capacity, no padding, and drops no tokens: it consumes the compact M, K layout with real group sizes , avoiding both the wasted-padding FLOPs and the accuracy hit of dropping, while also using less activation memory. In other words, Group-GEMM keeps the dropless quality of Path 2 with a compute/memory profile close to the ideal only real work, no padding , trading it for a sort and a ragged-aware kernel. But the remaining question is: how fast is the Group-GEMM kernel itself compared to these dense alternatives on real hardware? MaxText Models maxtext-models We evaluated the performance improvement achieved by the new ragged-dot backends using MaxText MoE models. The evaluated models are described below, and their configurations are given in Table 4 MI300 and Table 5 MI350 : “Routing-heavy” vs. “compute-heavy”: each MoE layer spends time on a the expert GEMM itself ragged dot , ≈ M tokens × K emb × N moe mlp and b routing — gating top-k selection plus the token permute/sort/scatter that builds the num experts ragged groups and the un-permute afterwards. A config is routing-heavy when b dominates: this happens with small base moe mlp dim N each expert GEMM does little arithmetic and/or many experts more small groups → more launch + scatter/gather overhead, lower per-group occupancy . The opposite corner — large N, few experts — is compute-heavy big dense expert GEMMs dominate, routing negligible . mixtral-8x7b MI300 & MI350 — The baseline dense-attention MoE microbenchmark emb 6144, 24 layers, 32 experts top-4, dot product attention, no MLA . Used to profile the raw ragged-dot expert GEMM without MLA complexity, at 8-way FSDP, seq 3072, grad-accum 1. The MI300 and MI350 configurations have identical shapes. It is the primary knob-sweep model for tile sizes and backend hipBLASLt vs. Triton vs. the default baseline . mixtral-8x7b-large MI300 & MI350 — A deliberately enlarged mixtral variant emb/K 8192, moe mlp/N 28672, 32 layers, 64 experts top-8 that grows both the expert-GEMM K/N dims and the number of ragged groups. Runs 8-way FSDP at seq 4096, grad-accum 8. It is the compute-heavy corner of the sweep — big, dense expert GEMMs with many groups — used to measure ragged-dot throughput at high arithmetic intensity and to check that larger tiles fit within HBM 192 GB HBM3 / 288 GB HBM3e . deepseek2-16b MI300 & MI350 — A scaled-down DeepSeek-V2 test emb 2048, 27 layers, 64 experts top-6, 2 shared experts that keeps the real MLA attention and deepseek decoder block. It validates the ragged-dot + MLA + YARN-RoPE code path cheaply, running autodetected FSDP, grad-accum 4, seq 512. Shared experts add a dense side-path alongside the routed experts, exercising the mixed dense/sparse MLP. deepseek3-test MI300 & MI350 — A tiny DeepSeek-V3 CI proxy emb 512 that keeps the deepseek block + MLA attention but shrinks everything else. On MI300 it is cut to 16 layers / 32 experts dense mlp 4096 ; on MI350 it retains the full 61-layer / 256-expert topology with the tiny emb 512. Runs autodetected FSDP, grad-accum 4, seq 512 — a fast functional check of the ragged-dot + MLA path rather than a perf run. deepseek3-671b-reduced MI300 & MI350 — A perf-representative slice of the flagship 671B DeepSeek-V3: it preserves the real expert-GEMM shapes K/emb 7168, N/moe mlp 2048, MLA but cuts to 4 layers / 32 experts so it fits on 8 GPUs. Because FSDP can’t hold the full model it is expert-parallel EP=8, 4 experts/GPU, FSDP off , grad-accum 8, seq 512. This makes the per-expert GEMM timing representative of the true 671B while keeping the run tractable. gemma4-26b MI300 only — Gemma-4 26B MoE emb 2816, 30 layers, 128 experts top-8, 1 shared with the gemma4 decoder block and standard dot product attention. Small base moe mlp dim 704 means many narrow expert GEMMs, stressing per-group launch/routing overhead rather than raw FLOPs. Runs autodetected FSDP, seq 1024, grad-accum 4 — a good test for many-small-experts ragged-dot efficiency on the 192 GB HBM3 gfx942 platform. qwen3-30b-a3b MI300 & MI350 — Qwen3 MoE emb 2048, 48 layers, 128 experts top-8, qwen3 moe block, QK-norm with a very small expert width base moe mlp dim 768 . Like gemma4 it is dominated by many thin expert GEMMs and routing overhead rather than arithmetic intensity. MI300 and MI350 use identical shapes autodetected FSDP, seq 1024, grad-accum 4 . A lightweight but routing-heavy ragged-dot workload. qwen3-235b-a22b MI350 only — Large Qwen3 MoE emb 4096, 94 layers , 128 experts top-8, moe mlp 1536 . The deep 94-layer stack combined with 128 experts makes it the heaviest routing/attention workload in the tables; it runs 8-way FSDP, seq 2048, grad-accum 16. Sequence length and layer depth push activation memory, so remat and FSDP are essential to fit — a scale/depth stress test for ragged-dot. llama4-17b-16e MI350 only — Llama-4 17B, 16-expert, top-1 routing each token → a single expert with wide experts moe mlp 8192 and a shared expert. This combination makes it the compute-heavy corner: a few large, dense expert GEMMs where the padded-dot baseline already wastes little, so Group-GEMM has limited headroom to improve on it. Runs autodetected FSDP, seq 1024, grad-accum 4. MI300 MI300X, gfx942, 192 GB HBM3, 8× GPU mi300-mi300x-gfx942-192-gb-hbm3-8-gpu Table 4: MaxText MoE model configurations evaluated on MI300. | Model config | emb dim K | moe mlp dim N | mlp dim | layers | experts | top-k | shared | attn | seq len | grad-accum | parallelism FSDP / EP / TP | |---|---|---|---|---|---|---|---|---|---|---|---| | mixtral-8x7b | 6144 | 20480 | 20480 | 24 | 32 | 4 | – | dot | 3072 | 1 | 8 / 1 / 1 | | mixtral-8x7b-large | 8192 | 28672 | 28672 | 32 | 64 | 8 | – | dot | 4096 | 8 | 8 / 1 / 1 | | deepseek2-16b | 2048 | 1408 | 5632 | 27 | 64 | 6 | 2 | MLA | 512 | 4 | auto / 1 / 1 | | deepseek3-test | 512 | 512 | 4096 | 16 | 32 | 8 | 1 | MLA | 512 | 4 | auto / 1 / 1 | | deepseek3-671b-reduced | 7168 | 2048 | 18432 | 4 | 32 | 8 | 1 | MLA | 512 | 8 | 1 / 8 / 1 | | gemma4-26b | 2816 | 704 | 2112 | 30 | 128 | 8 | 1 | dot | 1024 | 4 | auto / 1 / 1 | | qwen3-30b-a3b | 2048 | 768 | 768 | 48 | 128 | 8 | – | dot | 1024 | 4 | auto / 1 / 1 | MI350-Class gfx950, 288 GB HBM3e, 8× GPU mi350-class-gfx950-288-gb-hbm3e-8-gpu Table 5: MaxText MoE model configurations evaluated on MI350. | Model config | emb dim K | moe mlp dim N | mlp dim | layers | experts | top-k | shared | attn | seq len | grad-accum | parallelism FSDP / EP / TP | |---|---|---|---|---|---|---|---|---|---|---|---| | mixtral-8x7b | 6144 | 20480 | 20480 | 24 | 32 | 4 | – | dot | 3072 | 1 | 8 / 1 / 1 | | mixtral-8x7b-large | 8192 | 28672 | 28672 | 32 | 64 | 8 | – | dot | 4096 | 8 | 8 / 1 / 1 | | deepseek2-16b | 2048 | 1408 | 5632 | 27 | 64 | 6 | 2 | MLA | 512 | 4 | auto / 1 / 1 | | deepseek3-test | 512 | 512 | 18432 | 61 | 256 | 8 | 1 | MLA | 512 | 4 | auto / 1 / 1 | | deepseek3-671b-reduced | 7168 | 2048 | 18432 | 4 | 32 | 8 | 1 | MLA | 512 | 8 | 1 / 8 / 1 | | qwen3-30b-a3b | 2048 | 768 | 768 | 48 | 128 | 8 | – | dot | 1024 | 4 | auto / 1 / 1 | | qwen3-235b-a22b | 4096 | 1536 | 1536 | 94 | 128 | 8 | – | dot | 2048 | 16 | 8 / 1 / 1 | | llama4-17b-16e | 5120 | 8192 | 8192 | 48 | 16 | 1 | 1 | dot | 1024 | 4 | auto / 1 / 1 | Legend : - emb dim K = expert-GEMM contraction dim - moe mlp dim N = expert-GEMM output dim - attn = attention type MLA = Multi-head Latent Attention — DeepSeek’s low-rank compressed KV attention - parallelism = ici fsdp parallelism / ici expert parallelism / ici tensor parallelism , i.e. FSDP Fully Sharded Data Parallel — params/grads/optimizer sharded across GPUs / EP Expert Parallel — MoE experts distributed across GPUs / TP Tensor Parallel — individual matmuls split across GPUs - auto = -1 autodetects available GPUs . Micro-Benchmark micro-benchmark In order to evaluate the performance increase provided by the new backends for the ragged-dot operations, we isolated the ragged-dot operations executed in the MaxText models mentioned above and created a synthetic micro-benchmark that “replays” these operations alone. This micro-benchmark is written at the JAX level so it evaluates the whole compilation pipeline , runs the identified ragged-dot ops 50 times after a 10-run warm-up , and calculates the average execution time over these 50 runs. As the hipBLASLt Group-GEMM backend is only available for the bf16 datatype on MI300, we chose to show only numbers for this target in this blog post. From an MoE Layer to 6 Ragged-Dot Classes from-an-moe-layer-to-6-ragged-dot-classes An MoE expert block in MaxText is built from two projections — an up/gate projection wi and a down projection wo — and each of those has to be differentiated, producing a data gradient dgrad and a weight gradient wgrad in the backward pass. When you enumerate every ragged-dot XLA emits across the forward + backward pass, they collapse into exactly 6 structural classes. Two operations belong to the same class if and only if their operands share the same rank number of dimensions and the same contracting / ragged / group dimension numbers, and the output has the same rank and layout. However, the sizes M , K , N , E may differ. In the following, we use the notation given in Table 6: Table 6: Notation used to describe the ragged-dot classes. | symbol | meaning | |---|---| | M | ragged token dimension = batch × seq × top k the real routed-token count; no capacity padding on the ragged dot path | | K | expert-GEMM contraction / embedding dim base emb dim | | N | expert-GEMM output dim base moe mlp dim | | E | number of experts num experts | | group sizes | int32 E , per-expert token counts, summing to M | The M token axis is the ragged axis: it is partitioned into E groups, one per expert, and dtype is bfloat16 throughout. The six classes are listed in Table 7: Table 7: The six structural ragged-dot classes emitted across the forward and backward pass. | class | role | lhs | rhs | out | contract lhs,rhs | rhs group | |---|---|---|---|---|---|---| | fwd wi | fwd gate/up | M,K | E,K,N | M,N | {1},{1} | {0} | | fwd wo | fwd down | M,N | E,N,K | M,K | {1},{1} | {0} | | dgrad wi | dgrad up | M,N | E,K,N | M,K | {1},{2} | {0} | | dgrad wo | dgrad down | M,K | E,N,K | M,N | {1},{2} | {0} | | wgrad wi | wgrad up | M,K | M,N | E,K,N | {0},{0} | {} empty | | wgrad wo | wgrad down | M,N | M,K | E,N,K | {0},{0} | {} empty | All classes are bfloat16 , with lhs ragged dimensions = {0} and group sizes = int32 E . - The forward classes fwd wi , fwd wo are the plain jax.lax.ragged dot form — inputs M,K × w E,K,N → out M,N for the gate/up projection, and interm M,N × w E,N,K → out M,K for the down projection. - The dgrad classes dgrad wi , dgrad wo contract the weight’s output-feature axis rhs contracting = {2} , so they require the explicit ragged dot general form with custom dimension numbers. - The wgrad classes wgrad wi , wgrad wo are M,· × M,· → E,·,· Group-GEMMs that contract the ragged token axis and produce one weight-gradient slab per expert empty rhs group dimensions . For example, fwd wi — the gate/up projection — is simply: python import jax.numpy as jnp from jax import lax lhs = jnp.zeros M, K , jnp.bfloat16 tokens rhs = jnp.zeros E, K, N , jnp.bfloat16 expert weights group sizes: int32 E , sums to M out = lax.ragged dot lhs, rhs, group sizes, - M, N preferred element type=jnp.bfloat16 while a weight-gradient such as wgrad wi needs the general form with explicit dimension numbers: php a M,K × b M,N - w E,K,N ; contract the ragged M axis, group on the output wgrad dn = lax.RaggedDotDimensionNumbers dot dimension numbers= 0, , 0, , , , lhs ragged dimensions= 0, , rhs group dimensions= , empty - grouped output on E out = lax.ragged dot general lhs, rhs, group sizes, - E, K, N ragged dot dimension numbers=wgrad dn, preferred element type=jnp.bfloat16 Fixing group sizes Deterministically fixing-group-sizes-deterministically In real training the per-expert token counts are data-dependent — they come out of the router and change every step. For a reproducible benchmark we instead use a balanced / uniform partition: every expert receives M // E tokens, and the remainder M mod E is spread one extra token per group across the first M mod E experts, so that sum group sizes == M exactly: base = M // E tokens per expert floor rem = M - base E leftover to distribute 0 <= rem < E gs = np.full E, , base, np.int32 gs :rem += 1 first rem experts get one extra group sizes = jnp.asarray gs int32 E , sums to M This gives each expert ≈ M/E rows. Real routing is typically less balanced, which tends to lower the achieved TFLOP/s for the per-group backends. Results results For each model variant we measure the runtime of all the 6 classes on each backend default padded-dot baseline, triton , and hip . The resulting speedups over the baseline are shown in Figure 26. The full per-class TFLOP/s results for every measured shape are provided in Table 12 of the Appendix: Micro-Benchmark Full Results appendix-micro-benchmark-full-results . Performance of MaxText MoE Models performance-of-maxtext-moe-models Performance performance Now, let’s consider the overall performance improvement for the whole model train step. Note that in this section, the shown performance is the median value over training steps 5 to 10. The results are reported in Figure 27 and Table 8 for MI300, and in Figure 28 and Table 9 for MI350. MI300 mi300 Table 8: MaxText MoE model training-step performance on MI300. | Model | Emb Dim | Experts | Exp/Tok | Layers | total weights | Backend | execution time s | TFLOP/s/device | Tokens/s/device | Speedup | |---|---|---|---|---|---|---|---|---|---|---| | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 3.946 | 7.744 | 519.018 | 1.00x | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 1.867 | 16.369 | 1097.136 | 2.11x | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | hip | 1.761 | 17.348 | 1162.700 | 2.24x | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix | 1.617 | 18.895 | 1266.397 | 2.44x | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix dropless | 2.018 | 15.143 | 1014.919 | 1.96x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 5.758 | 10.910 | 355.653 | 1.00x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 2.732 | 22.995 | 749.610 | 2.11x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | hip | 2.373 | 26.470 | 862.896 | 2.43x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix | 2.261 | 27.780 | 905.610 | 2.55x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix dropless | 2.806 | 22.387 | 729.808 | 2.05x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 0.575 | 141.236 | 7126.279 | 1.00x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 0.436 | 186.151 | 9392.514 | 1.32x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | hip | 0.429 | 189.228 | 9547.763 | 1.34x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix | 0.414 | 196.070 | 9893.003 | 1.39x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix dropless | 0.305 | 265.880 | 13415.345 | 1.89x | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | default | 20.531 | 4.751 | 199.507 | 1.00x | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | triton | 3.257 | 29.950 | 1257.783 | 6.30x | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | hip | 3.073 | 31.737 | 1332.846 | 6.68x | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | dense matrix | 2.914 | 33.474 | 1405.770 | 7.05x | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | dense matrix dropless | 4.163 | 23.430 | 983.947 | 4.93x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 20.676 | 3.855 | 198.105 | 1.00x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 5.013 | 15.897 | 817.004 | 4.12x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | hip | 3.505 | 22.740 | 1168.655 | 5.90x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix | 3.287 | 24.246 | 1246.054 | 6.29x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix dropless | 4.896 | 16.278 | 836.600 | 4.22x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 5.044 | 48.057 | 609.025 | 1.00x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 2.817 | 86.045 | 1090.451 | 1.79x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | hip | 2.128 | 113.903 | 1443.498 | 2.37x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix | 2.015 | 120.291 | 1524.461 | 2.50x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 3.405 | 71.187 | 902.164 | 1.48x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | default | 52.107 | 50.128 | 628.859 | 1.00x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | triton | 30.822 | 84.744 | 1063.120 | 1.69x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | hip | 22.940 | 113.863 | 1428.421 | 2.27x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix | 18.438 | 141.662 | 1777.156 | 2.83x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 34.163 | 76.458 | 959.171 | 1.53x | MI350 mi350 Table 9: MaxText MoE model training-step performance on MI350. | Model | Emb Dim | Experts | Exp/Tok | Layers | total weights | Backend | execution time s | TFLOP/s/device | Tokens/s/device | Speedup | |---|---|---|---|---|---|---|---|---|---|---| | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 2.581 | 11.840 | 793.547 | 1.00x | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 1.433 | 21.316 | 1428.675 | 1.80x | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix | 1.303 | 23.451 | 1571.765 | 1.98x | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix dropless | 1.414 | 21.603 | 1447.888 | 1.83x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 4.208 | 14.929 | 486.677 | 1.00x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 2.220 | 28.302 | 922.617 | 1.90x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix | 1.778 | 35.344 | 1152.180 | 2.37x | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix dropless | 1.984 | 31.662 | 1032.149 | 2.12x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 0.459 | 176.907 | 8926.120 | 1.00x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 0.374 | 216.782 | 10938.067 | 1.23x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix | 0.305 | 266.230 | 13433.032 | 1.50x | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix dropless | 0.188 | 432.514 | 21823.103 | 2.44x | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | default | 10.249 | 39.299 | 399.631 | 1.00x | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | triton | 10.157 | 39.655 | 403.258 | 1.01x | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | dense matrix | 9.337 | 43.138 | 438.671 | 1.10x | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | dense matrix dropless | 9.344 | 43.108 | 438.372 | 1.10x | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | default | 22.025 | 14.110 | 1487.731 | 1.00x | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | triton | 4.082 | 76.132 | 8027.089 | 5.40x | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | dense matrix | 3.623 | 85.785 | 9044.923 | 6.08x | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | dense matrix dropless | 5.322 | 58.391 | 6156.577 | 4.14x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 9.696 | 8.220 | 422.434 | 1.00x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 4.012 | 19.863 | 1020.834 | 2.42x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix | 2.557 | 31.166 | 1601.701 | 3.79x | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix dropless | 3.294 | 24.197 | 1243.551 | 2.94x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 2.863 | 84.681 | 1073.164 | 1.00x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 1.917 | 126.458 | 1602.606 | 1.49x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix | 1.432 | 169.318 | 2145.780 | 2.00x | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 2.075 | 116.814 | 1480.395 | 1.38x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | default | 29.188 | 89.489 | 1122.641 | 1.00x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | triton | 21.327 | 122.475 | 1536.450 | 1.37x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix | 13.376 | 195.272 | 2449.693 | 2.18x | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 25.522 | 102.344 | 1283.910 | 1.14x | The performance numbers confirm that the new Group-GEMM based backends for jax.lax.ragged dot operations provide a significant speedup for MoE models. As expected, the hipBLASLt backend is faster than the Triton one; however, both offer very interesting speedups. That said, we also notice that the dense matrix drop case approach outperforms the ragged-dot path for all the tested models. Indeed, as GPUs are optimized for regular computation, the pure matrix multiplication was expected to be faster than Group-GEMM. However, we might have expected that the extra computation due to the padding would result in a larger slowdown, making the ragged-dot approach faster. But this is not the case, for two main reasons: the capacity is relatively low 1.25 , which means that the padding is relatively small, but potentially a significant number of tokens could be dropped; the pre-processing to put the data in shape for the ragged-dot backend is expensive. Consequently, for the relatively small models evaluated, the time overhead introduced by the pre-processing cancels out the benefit of the Group-GEMM performing only useful multiplications. But with larger models the pre-processing overhead could be better balanced against the matmul efficiency, making the ragged-dot Group-GEMM backends attractive. Memory Savings memory-savings The memory-usage figures below were extracted directly from the log traces emitted by MaxText during each training run. MaxText reports, per run, the peak device memory footprint broken down into total and temporary/scratch memory , which we collected for each backend default , triton , hip and diffed against the default padded-dot baseline to obtain the savings shown. They are reported in Figure 29 and Table 10 for MI300, and in Figure 30 and Table 11 for MI350. MI300 id1 Table 10: MaxText MoE model peak memory footprint on MI300. | Model | Emb Dim | Experts | Exp/Tok | Layers | total weights | Backend | Total Mem GB | Temp Mem GB | Mem Savings GB | Temp Savings GB | |---|---|---|---|---|---|---|---|---|---|---| | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 26.7 | 15.7 | 0.0 0.0% | 0.0 0.0% | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 23.6 | 12.6 | 3.1 11.6% | 3.1 19.7% | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | hip | 23.6 | 12.6 | 3.1 11.6% | 3.1 19.7% | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix capacity 1.25 | 23.4 | 12.5 | 3.3 12.4% | 3.2 20.4% | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix dropless | 23.9 | 13.0 | 2.8 10.5% | 2.7 17.2% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 26.3 | 15.2 | 0.0 0.0% | 0.0 0.0% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 23.1 | 12.0 | 3.2 12.2% | 3.2 21.1% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | hip | 23.1 | 12.0 | 3.2 12.2% | 3.2 21.1% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix capacity 1.25 | 23.1 | 12.0 | 3.2 12.2% | 3.2 21.1% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix dropless | 23.6 | 12.6 | 2.7 10.3% | 2.6 17.1% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 8.6 | 3.4 | 0.0 0.0% | 0.0 0.0% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 8.6 | 3.4 | 0.0 0.0% | 0.0 0.0% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | hip | 8.6 | 3.4 | 0.0 0.0% | 0.0 0.0% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix capacity 1.25 | 8.5 | 3.4 | 0.1 1.2% | 0.0 0.0% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix dropless | 8.5 | 3.4 | 0.1 1.2% | 0.0 0.0% | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | default | 61.0 | 43.3 | 0.0 0.0% | 0.0 0.0% | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | triton | 43.6 | 26.0 | 17.4 28.5% | 17.3 40.0% | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | hip | 43.6 | 26.0 | 17.4 28.5% | 17.3 40.0% | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | dense matrix capacity 1.25 | 40.0 | 22.3 | 21.0 34.4% | 21.0 48.5% | | gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | dense matrix dropless | 41.6 | 23.9 | 19.4 31.8% | 19.4 44.8% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 57.7 | 36.4 | 0.0 0.0% | 0.0 0.0% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 44.5 | 23.2 | 13.2 22.9% | 13.2 36.3% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | hip | 44.5 | 23.2 | 13.2 22.9% | 13.2 36.3% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix capacity 1.25 | 44.5 | 23.2 | 13.2 22.9% | 13.2 36.3% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix dropless | 45.8 | 24.5 | 11.9 20.6% | 11.9 32.7% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 120.2 | 54.9 | 0.0 0.0% | 0.0 0.0% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 114.4 | 49.1 | 5.8 4.8% | 5.8 10.6% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | hip | 114.7 | 49.5 | 5.5 4.6% | 5.4 9.8% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix capacity 1.25 | 115.0 | 49.7 | 5.2 4.3% | 5.2 9.5% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 117.7 | 52.5 | 2.5 2.1% | 2.4 4.4% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | default | 79.9 | 47.2 | 0.0 0.0% | 0.0 0.0% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | triton | 72.7 | 40.1 | 7.2 9.0% | 7.1 15.0% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | hip | 72.9 | 40.2 | 7.0 8.8% | 7.0 14.8% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix capacity 1.25 | 73.3 | 40.7 | 6.6 8.3% | 6.5 13.8% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 77.0 | 44.4 | 2.9 3.6% | 2.8 5.9% | MI350 id2 Table 11: MaxText MoE model peak memory footprint on MI350. | Model | Emb Dim | Experts | Exp/Tok | Layers | total weights | Backend | Total Mem GB | Temp Mem GB | Mem Savings GB | Temp Savings GB | |---|---|---|---|---|---|---|---|---|---|---| | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 26.7 | 15.7 | 0.0 0.0% | 0.0 0.0% | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 23.7 | 12.7 | 3.0 11.2% | 3.0 19.1% | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix capacity 1.25 | 23.5 | 12.5 | 3.2 12.0% | 3.2 20.4% | | deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense matrix dropless | 23.9 | 13.0 | 2.8 10.5% | 2.7 17.2% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 26.3 | 15.2 | 0.0 0.0% | 0.0 0.0% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 23.2 | 12.2 | 3.1 11.8% | 3.0 19.7% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix capacity 1.25 | 23.2 | 12.1 | 3.1 11.8% | 3.1 20.4% | | deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense matrix dropless | 23.8 | 12.8 | 2.5 9.5% | 2.4 15.8% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 8.6 | 3.4 | 0.0 0.0% | 0.0 0.0% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 8.6 | 3.4 | 0.0 0.0% | 0.0 0.0% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix capacity 1.25 | 9.7 | 4.5 | -1.1 -12.8% | -1.1 -32.4% | | deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense matrix dropless | 9.4 | 4.2 | -0.8 -9.3% | -0.8 -23.5% | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | default | 170.5 | 95.2 | 0.0 0.0% | 0.0 0.0% | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | triton | 170.0 | 94.7 | 0.5 0.3% | 0.5 0.5% | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | dense matrix capacity 1.25 | 171.6 | 96.3 | -1.1 -0.6% | -1.1 -1.2% | | llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | dense matrix dropless | 171.6 | 96.3 | -1.1 -0.6% | -1.1 -1.2% | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | default | 74.7 | 66.9 | 0.0 0.0% | 0.0 0.0% | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | triton | 24.4 | 16.5 | 50.3 67.3% | 50.4 75.3% | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | dense matrix capacity 1.25 | 23.2 | 15.4 | 51.5 68.9% | 51.5 77.0% | | qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | dense matrix dropless | 28.8 | 21.0 | 45.9 61.4% | 45.9 68.6% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 57.7 | 36.4 | 0.0 0.0% | 0.0 0.0% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 45.1 | 23.7 | 12.6 21.8% | 12.7 34.9% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix capacity 1.25 | 44.8 | 23.5 | 12.9 22.4% | 12.9 35.4% | | qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense matrix dropless | 46.2 | 24.9 | 11.5 19.9% | 11.5 31.6% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 120.2 | 54.9 | 0.0 0.0% | 0.0 0.0% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 115.2 | 50.0 | 5.0 4.2% | 4.9 8.9% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix capacity 1.25 | 115.2 | 49.9 | 5.0 4.2% | 5.0 9.1% | | mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 117.9 | 52.7 | 2.3 1.9% | 2.2 4.0% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | default | 79.9 | 47.2 | 0.0 0.0% | 0.0 0.0% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | triton | 73.7 | 41.1 | 6.2 7.8% | 6.1 12.9% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix capacity 1.25 | 73.6 | 41.0 | 6.3 7.9% | 6.2 13.1% | | mixtral-8x7b large | 4096 | 8 | 2 | 32 | 46.703B | dense matrix dropless | 77.3 | 44.7 | 2.6 3.3% | 2.5 5.3% | Regarding memory usage, the measured memory footprints confirm that the Group-GEMM backends make better use of memory for ragged-dot operations than the default XLA backend by avoiding the storage of unnecessary intermediate data. Indeed, the default backend rewrites the single ragged-dot lhs M,K , rhs E,K,N → out M,N into E independent dot products, each over the full M rows — i.e. it computes an E × M, N stack of results and the matching E × M, · activations and then masks each expert’s output down to the rows that actually belong to it. So every expert redundantly processes all M tokens, and the intermediate/output tensors blow up by a factor of E . Thus, the baseline’s wasted memory scales roughly as E × M × N , whereas Group-GEMM keeps a single compact M, N output. The savings are therefore largest when the product E × M × N is large — many experts, a large routed-token dimension M , and/or a wide output N — and small when any of those is small. The measured numbers line up with this: - qwen3-235b-a22b MI350 — 128 experts and a large routed-token workload 94 layers, seq 2048, top-8 . Its baseline temp memory is 66.9 GB , of which Group-GEMM removes 50.4 GB 75.3% — by far the biggest saving in the tables. With 128 experts each redundantly processing all M tokens, the E × M, N stack is enormous, and Group-GEMM never builds it. - gemma4-26b MI300 — 128 experts and many routed tokens top-8 , even though the expert width is small N = 704 : 17.3 GB 40.0% of temp memory saved. The large E × M dominates the blow-up. - qwen3-30b-a3b MI300 & MI350 — 128 experts , N = 768 , top-8: ≈13 GB ~35% temp saved on both targets — the same many-experts effect. - mixtral-8x7b only 8 experts — despite large M / N , only ~5 GB ~10% temp saved: with just 8 experts the E × factor is small, so the redundant stack is comparatively cheap and there is little to reclaim. - deepseek3-671b-reduced — 0 GB 0% : expert-parallel with only a few experts per GPU and very few tokens per step, so E × M is tiny and there is essentially nothing to remove. - llama4-17b-16e MI350 — 0.5 GB 0.5% : only 16 experts and top-1 routing small routed-token count , combined with wide dense GEMMs — a small E × M , so the baseline’s redundant stack is modest. Now, if we compare the memory usage of ragged-dot Group-GEMM against dense-matrix approaches, we notice that in most cases the memory usage is quite similar. This can be explained by: - the fact that MaxText does not allocate extra matrices to process the data. It only uses the activation matrix as it is, using masking and data movements to compute the experts’ FFNs. - in our experiment, the drop capacity is quite low 1.25 , which does not incur significant data padding. Part 5 — Conclusion and Next Steps part-5-conclusion-and-next-steps Using Group-GEMM as the backend for XLA ragged-dot brings two clear advantages: it is faster no padding FLOPs and it uses less memory no expansion/padding . And XLA now offers two ways to get there: - hipBLASLt — a vendor-library backend on AMD ROCm, extending XLA’s existing hipBLASLt matmul path - Triton — a portable, autotunable, code-generated backend that runs on both CUDA and ROCm, currently experimental and off by default. Next steps: - Autotuner-driven backend selection. As mentioned earlier, the current selection of ragged-dot backends is based on static rules. Given that performance on different targets and the scope for fusion and optimization can vary from one operation to another, a next step for ragged-dot backends would be to delegate the selection of the backend to the XLA autotuner. - Faster data pre-/post-processing around ragged-dot. As highlighted in The Three MoE Compute Paths in MaxText the-three-moe-compute-paths-in-maxtext , the sparse ragged-dot path pays for a permute sort tokens by expert + build group sizes before the Group-GEMM and an un-permute scatter back + weighted combine after it. This pre-processing can dominate the layer, capping the end-to-end speedup even when the Group-GEMM kernel itself is fast. Improving this pre-/post-processing is therefore a key next step to reduce the routing overhead and fully unlock the performance of the XLA Group-GEMM backend. Summary summary In this blog post, we showed how XLA turns the jax.lax.ragged dot operation into real Group-GEMM on GPUs, avoiding the wasted compute and memory of the padded-dot baseline. We built that argument in five parts, each with a distinct role. Part 1 part-1-why-moe-turns-into-a-ragged-matmul-problem framed the problem, moving from the MoE layer a model author writes to the ragged matmul a compiler actually sees, and identified Group-GEMM as the primitive that removes the padded baseline’s two costs. Part 2 part-2-backend-1-hipblaslt-group-gemm-amd-rocm and Part 3 part-3-backend-2-triton-ragged-dot-portable-cuda-rocm then gave two implementations of that idea — one driving a vendor library hipBLASLt , the other generating code Triton . Part 4 part-4-performance-evaluation supplied the evidence, and is also where the honest limits surfaced: the kernels are dramatically faster in isolation, yet the end-to-end gain is capped by the permute/un-permute routing overhead that surrounds them. Part 5 part-5-conclusion-and-next-steps concludes on the Group-GEMM support and proposes a roadmap for future improvements. Along the way, the key takeaways were: - Two GPU backends. We delivered 1 a hipBLASLt backend on AMD ROCm that extends XLA’s existing matmul path and bridges the compile-time/runtime gap by filling the Group-GEMM descriptors from the runtime group sizes ; and 2 a portable Triton code-generation backend that runs on both CUDA and ROCm and emits a specialized kernel for each ragged-dot variant non-contracting, contracting, and batch . - Performance and memory results. On real MaxText MoE models, the Group-GEMM backends deliver significant end-to-end speedups up to ~6.7× on MI300 and ~5.4× on MI350 over the padded-dot baseline, with hipBLASLt generally ahead of Triton. They also substantially reduce activation memory — the savings scale with E × M × N , reaching ~75% temp-memory savings for many-expert models such as qwen3-235b-a22b . We also saw that, for the relatively small models evaluated, the permute/un-permute pre-processing overhead can offset the raw matmul efficiency, which is why improving that routing overhead is a key next step. By reading this post you learned what MoE looks like to a compiler, why ragged-dot /Group-GEMM is the right abstraction, how each backend is implemented inside XLA, and how to enable and evaluate them on AMD GPUs. Appendix: Micro-Benchmark Full Results appendix-micro-benchmark-full-results Table 12 below reports the full per-class TFLOP/s results for every measured shape in the micro-benchmark of Micro-Benchmark micro-benchmark , across the three backends default padded-dot baseline, triton , and hipblas . Table 12: Full per-class micro-benchmark results TFLOP/s on MI300. | class | M | K | N | E | backend | tflops | status | |---|---|---|---|---|---|---|---| | fwd wi | 3840 | 2048 | 1408 | 64 | default | 6.85 | ok | | fwd wi | 3840 | 2048 | 1408 | 64 | triton | 86.44 | ok | | fwd wi | 3840 | 2048 | 1408 | 64 | hipblas | 135.21 | ok | | fwd wo | 3840 | 2048 | 1408 | 64 | default | 9.04 | ok | | fwd wo | 3840 | 2048 | 1408 | 64 | triton | 94.14 | ok | | fwd wo | 3840 | 2048 | 1408 | 64 | hipblas | 120.85 | ok | | dgrad wi | 3840 | 2048 | 1408 | 64 | default | 8.50 | ok | | dgrad wi | 3840 | 2048 | 1408 | 64 | triton | 87.50 | ok | | dgrad wi | 3840 | 2048 | 1408 | 64 | hipblas | 159.97 | ok | | wgrad wi | 3840 | 2048 | 1408 | 64 | default | 8.79 | ok | | wgrad wi | 3840 | 2048 | 1408 | 64 | triton | 94.40 | ok | | wgrad wi | 3840 | 2048 | 1408 | 64 | hipblas | 102.22 | ok | | dgrad wo | 3840 | 2048 | 1408 | 64 | default | 6.49 | ok | | dgrad wo | 3840 | 2048 | 1408 | 64 | triton | 82.86 | ok | | dgrad wo | 3840 | 2048 | 1408 | 64 | hipblas | 119.83 | ok | | wgrad wo | 3840 | 2048 | 1408 | 64 | default | 8.80 | ok | | wgrad wo | 3840 | 2048 | 1408 | 64 | triton | 81.52 | ok | | wgrad wo | 3840 | 2048 | 1408 | 64 | hipblas | 104.17 | ok | | fwd wi | 1920 | 5120 | 1536 | 160 | default | 3.43 | ok | | fwd wi | 1920 | 5120 | 1536 | 160 | triton | 31.05 | ok | | fwd wi | 1920 | 5120 | 1536 | 160 | hipblas | 44.31 | ok | | fwd wo | 1920 | 5120 | 1536 | 160 | default | 4.13 | ok | | fwd wo | 1920 | 5120 | 1536 | 160 | triton | 36.06 | ok | | fwd wo | 1920 | 5120 | 1536 | 160 | hipblas | 33.43 | ok | | dgrad wi | 1920 | 5120 | 1536 | 160 | default | 3.28 | ok | | dgrad wi | 1920 | 5120 | 1536 | 160 | triton | 34.16 | ok | | dgrad wi | 1920 | 5120 | 1536 | 160 | hipblas | 38.35 | ok | | wgrad wi | 1920 | 5120 | 1536 | 160 | default | 3.72 | ok | | wgrad wi | 1920 | 5120 | 1536 | 160 | triton | 4.50 | ok | | wgrad wi | 1920 | 5120 | 1536 | 160 | hipblas | 26.74 | ok | | dgrad wo | 1920 | 5120 | 1536 | 160 | default | 2.95 | ok | | dgrad wo | 1920 | 5120 | 1536 | 160 | triton | 30.68 | ok | | dgrad wo | 1920 | 5120 | 1536 | 160 | hipblas | 45.01 | ok | | wgrad wo | 1920 | 5120 | 1536 | 160 | default | 3.68 | ok | | wgrad wo | 1920 | 5120 | 1536 | 160 | triton | 4.50 | ok | | wgrad wo | 1920 | 5120 | 1536 | 160 | hipblas | 26.80 | ok | | fwd wi | 5120 | 512 | 512 | 256 | default | 2.02 | ok | | fwd wi | 5120 | 512 | 512 | 256 | triton | 32.98 | ok | | fwd wi | 5120 | 512 | 512 | 256 | hipblas | 24.19 | ok | | fwd wo | 5120 | 512 | 512 | 256 | default | 2.02 | ok | | fwd wo | 5120 | 512 | 512 | 256 | triton | 33.16 | ok | | fwd wo | 5120 | 512 | 512 | 256 | hipblas | 26.19 | ok | | dgrad wi | 5120 | 512 | 512 | 256 | default | 1.92 | ok | | dgrad wi | 5120 | 512 | 512 | 256 | triton | 34.07 | ok | | dgrad wi | 5120 | 512 | 512 | 256 | hipblas | 27.99 | ok | | wgrad wi | 5120 | 512 | 512 | 256 | default | 1.52 | ok | | wgrad wi | 5120 | 512 | 512 | 256 | triton | 26.85 | ok | | wgrad wi | 5120 | 512 | 512 | 256 | hipblas | 24.98 | ok | | dgrad wo | 5120 | 512 | 512 | 256 | default | 1.92 | ok | | dgrad wo | 5120 | 512 | 512 | 256 | triton | 33.37 | ok | | dgrad wo | 5120 | 512 | 512 | 256 | hipblas | 26.36 | ok | | wgrad wo | 5120 | 512 | 512 | 256 | default | 1.53 | ok | | wgrad wo | 5120 | 512 | 512 | 256 | triton | 26.85 | ok | | wgrad wo | 5120 | 512 | 512 | 256 | hipblas | 24.53 | ok | | fwd wi | 10240 | 2816 | 704 | 128 | default | 2.14 | ok | | fwd wi | 10240 | 2816 | 704 | 128 | triton | 110.97 | ok | | fwd wi | 10240 | 2816 | 704 | 128 | hipblas | 172.24 | ok | | fwd wo | 10240 | 2816 | 704 | 128 | default | 5.66 | ok | | fwd wo | 10240 | 2816 | 704 | 128 | triton | 128.01 | ok | | fwd wo | 10240 | 2816 | 704 | 128 | hipblas | 137.69 | ok | | dgrad wi | 10240 | 2816 | 704 | 128 | default | 5.47 | ok | | dgrad wi | 10240 | 2816 | 704 | 128 | triton | 128.28 | ok | | dgrad wi | 10240 | 2816 | 704 | 128 | hipblas | 175.65 | ok | | wgrad wi | 10240 | 2816 | 704 | 128 | default | 3.62 | ok | | wgrad wi | 10240 | 2816 | 704 | 128 | triton | 116.71 | ok | | wgrad wi | 10240 | 2816 | 704 | 128 | hipblas | 104.51 | ok | | dgrad wo | 10240 | 2816 | 704 | 128 | default | 2.19 | ok | | dgrad wo | 10240 | 2816 | 704 | 128 | triton | 105.61 | ok | | dgrad wo | 10240 | 2816 | 704 | 128 | hipblas | 176.18 | ok | | wgrad wo | 10240 | 2816 | 704 | 128 | default | 3.67 | ok | | wgrad wo | 10240 | 2816 | 704 | 128 | triton | 96.62 | ok | | wgrad wo | 10240 | 2816 | 704 | 128 | hipblas | 113.53 | ok | | fwd wi | 10240 | 2880 | 2880 | 128 | default | 4.46 | ok | | fwd wi | 10240 | 2880 | 2880 | 128 | triton | 128.32 | ok | | fwd wi | 10240 | 2880 | 2880 | 128 | hipblas | 200.81 | ok | | fwd wo | 10240 | 2880 | 2880 | 128 | default | 4.46 | ok | | fwd wo | 10240 | 2880 | 2880 | 128 | triton | 127.71 | ok | | fwd wo | 10240 | 2880 | 2880 | 128 | hipblas | 199.25 | ok | | dgrad wi | 10240 | 2880 | 2880 | 128 | default | 4.33 | ok | | dgrad wi | 10240 | 2880 | 2880 | 128 | triton | 126.28 | ok | | dgrad wi | 10240 | 2880 | 2880 | 128 | hipblas | 242.96 | ok | | wgrad wi | 10240 | 2880 | 2880 | 128 | default | 4.90 | ok | | wgrad wi | 10240 | 2880 | 2880 | 128 | triton | 130.05 | ok | | wgrad wi | 10240 | 2880 | 2880 | 128 | hipblas | 127.73 | ok | | dgrad wo | 10240 | 2880 | 2880 | 128 | default | 4.32 | ok | | dgrad wo | 10240 | 2880 | 2880 | 128 | triton | 125.88 | ok | | dgrad wo | 10240 | 2880 | 2880 | 128 | hipblas | 242.90 | ok | | wgrad wo | 10240 | 2880 | 2880 | 128 | default | 4.91 | ok | | wgrad wo | 10240 | 2880 | 2880 | 128 | triton | 129.74 | ok | | wgrad wo | 10240 | 2880 | 2880 | 128 | hipblas | 128.08 | ok | | fwd wi | 1280 | 5120 | 8192 | 16 | default | 44.77 | ok | | fwd wi | 1280 | 5120 | 8192 | 16 | triton | 133.88 | ok | | fwd wi | 1280 | 5120 | 8192 | 16 | hipblas | 220.41 | ok | | fwd wo | 1280 | 5120 | 8192 | 16 | default | 35.27 | ok | | fwd wo | 1280 | 5120 | 8192 | 16 | triton | 129.44 | ok | | fwd wo | 1280 | 5120 | 8192 | 16 | hipblas | 224.86 | ok | | dgrad wi | 1280 | 5120 | 8192 | 16 | default | 27.50 | ok | | dgrad wi | 1280 | 5120 | 8192 | 16 | triton | 96.85 | ok | | dgrad wi | 1280 | 5120 | 8192 | 16 | hipblas | 246.03 | ok | | wgrad wi | 1280 | 5120 | 8192 | 16 | default | 40.48 | ok | | wgrad wi | 1280 | 5120 | 8192 | 16 | triton | 113.19 | ok | | wgrad wi | 1280 | 5120 | 8192 | 16 | hipblas | 154.65 | ok | | dgrad wo | 1280 | 5120 | 8192 | 16 | default | 32.66 | ok | | dgrad wo | 1280 | 5120 | 8192 | 16 | triton | 130.98 | ok | | dgrad wo | 1280 | 5120 | 8192 | 16 | hipblas | 232.46 | ok | | wgrad wo | 1280 | 5120 | 8192 | 16 | default | 41.02 | ok | | wgrad wo | 1280 | 5120 | 8192 | 16 | triton | 117.14 | ok | | wgrad wo | 1280 | 5120 | 8192 | 16 | hipblas | 154.67 | ok | | fwd wi | 15360 | 6144 | 20480 | 32 | default | 24.73 | ok | | fwd wi | 15360 | 6144 | 20480 | 32 | triton | 261.23 | ok | | fwd wi | 15360 | 6144 | 20480 | 32 | hipblas | 527.94 | ok | | fwd wo | 15360 | 6144 | 20480 | 32 | default | 17.97 | ok | | fwd wo | 15360 | 6144 | 20480 | 32 | triton | 233.12 | ok | | fwd wo | 15360 | 6144 | 20480 | 32 | hipblas | 471.04 | ok | | dgrad wi | 15360 | 6144 | 20480 | 32 | default | 17.71 | ok | | dgrad wi | 15360 | 6144 | 20480 | 32 | triton | 244.21 | ok | | dgrad wi | 15360 | 6144 | 20480 | 32 | hipblas | 485.43 | ok | | wgrad wi | 15360 | 6144 | 20480 | 32 | default | 24.57 | ok | | wgrad wi | 15360 | 6144 | 20480 | 32 | triton | 284.41 | ok | | wgrad wi | 15360 | 6144 | 20480 | 32 | hipblas | 428.84 | ok | | dgrad wo | 15360 | 6144 | 20480 | 32 | default | 24.01 | ok | | dgrad wo | 15360 | 6144 | 20480 | 32 | triton | 266.45 | ok | | dgrad wo | 15360 | 6144 | 20480 | 32 | hipblas | 470.12 | ok | | wgrad wo | 15360 | 6144 | 20480 | 32 | default | 24.62 | ok | | wgrad wo | 15360 | 6144 | 20480 | 32 | triton | 283.80 | ok | | wgrad wo | 15360 | 6144 | 20480 | 32 | hipblas | 431.66 | ok | | fwd wi | 40960 | 8192 | 28672 | 64 | default | 8.80 | ok | | fwd wi | 40960 | 8192 | 28672 | 64 | triton | 294.11 | ok | | fwd wi | 40960 | 8192 | 28672 | 64 | hipblas | 510.97 | ok | | fwd wo | 40960 | 8192 | 28672 | 64 | default | | OOM | | fwd wo | 40960 | 8192 | 28672 | 64 | triton | 298.32 | ok | | fwd wo | 40960 | 8192 | 28672 | 64 | hipblas | 507.20 | ok | | dgrad wi | 40960 | 8192 | 28672 | 64 | default | | OOM | | dgrad wi | 40960 | 8192 | 28672 | 64 | triton | 299.03 | ok | | dgrad wi | 40960 | 8192 | 28672 | 64 | hipblas | 464.65 | ok | | wgrad wi | 40960 | 8192 | 28672 | 64 | default | | OOM | | wgrad wi | 40960 | 8192 | 28672 | 64 | triton | 78.34 | ok | | wgrad wi | 40960 | 8192 | 28672 | 64 | hipblas | 318.52 | ok | | dgrad wo | 40960 | 8192 | 28672 | 64 | default | 8.74 | ok | | dgrad wo | 40960 | 8192 | 28672 | 64 | triton | 308.75 | ok | | dgrad wo | 40960 | 8192 | 28672 | 64 | hipblas | 497.37 | ok | | wgrad wo | 40960 | 8192 | 28672 | 64 | default | | OOM | | wgrad wo | 40960 | 8192 | 28672 | 64 | triton | 51.84 | ok | | wgrad wo | 40960 | 8192 | 28672 | 64 | hipblas | 336.43 | ok | | fwd wi | 20480 | 4096 | 1536 | 128 | default | 3.82 | ok | | fwd wi | 20480 | 4096 | 1536 | 128 | triton | 168.07 | ok | | fwd wi | 20480 | 4096 | 1536 | 128 | hipblas | 332.96 | ok | | fwd wo | 20480 | 4096 | 1536 | 128 | default | 5.71 | ok | | fwd wo | 20480 | 4096 | 1536 | 128 | triton | 176.82 | ok | | fwd wo | 20480 | 4096 | 1536 | 128 | hipblas | 323.89 | ok | | dgrad wi | 20480 | 4096 | 1536 | 128 | default | 5.61 | ok | | dgrad wi | 20480 | 4096 | 1536 | 128 | triton | 190.20 | ok | | dgrad wi | 20480 | 4096 | 1536 | 128 | hipblas | 291.06 | ok | | wgrad wi | 20480 | 4096 | 1536 | 128 | default | 4.82 | ok | | wgrad wi | 20480 | 4096 | 1536 | 128 | triton | 175.78 | ok | | wgrad wi | 20480 | 4096 | 1536 | 128 | hipblas | 236.19 | ok | | dgrad wo | 20480 | 4096 | 1536 | 128 | default | 4.43 | ok | | dgrad wo | 20480 | 4096 | 1536 | 128 | triton | 172.91 | ok | | dgrad wo | 20480 | 4096 | 1536 | 128 | hipblas | 324.73 | ok | | wgrad wo | 20480 | 4096 | 1536 | 128 | default | 4.85 | ok | | wgrad wo | 20480 | 4096 | 1536 | 128 | triton | 176.91 | ok | | wgrad wo | 20480 | 4096 | 1536 | 128 | hipblas | 232.07 | ok | | fwd wi | 10240 | 2048 | 768 | 128 | default | 3.33 | ok | | fwd wi | 10240 | 2048 | 768 | 128 | triton | 111.85 | ok | | fwd wi | 10240 | 2048 | 768 | 128 | hipblas | 171.38 | ok | | fwd wo | 10240 | 2048 | 768 | 128 | default | 5.59 | ok | | fwd wo | 10240 | 2048 | 768 | 128 | triton | 116.12 | ok | | fwd wo | 10240 | 2048 | 768 | 128 | hipblas | 135.95 | ok | | dgrad wi | 10240 | 2048 | 768 | 128 | default | 5.39 | ok | | dgrad wi | 10240 | 2048 | 768 | 128 | triton | 115.59 | ok | | dgrad wi | 10240 | 2048 | 768 | 128 | hipblas | 144.88 | ok | | wgrad wi | 10240 | 2048 | 768 | 128 | default | 3.87 | ok | | wgrad wi | 10240 | 2048 | 768 | 128 | triton | 99.24 | ok | | wgrad wi | 10240 | 2048 | 768 | 128 | hipblas | 121.36 | ok | | dgrad wo | 10240 | 2048 | 768 | 128 | default | 3.25 | ok | | dgrad wo | 10240 | 2048 | 768 | 128 | triton | 106.70 | ok | | dgrad wo | 10240 | 2048 | 768 | 128 | hipblas | 171.31 | ok | | wgrad wo | 10240 | 2048 | 768 | 128 | default | 3.87 | ok | | wgrad wo | 10240 | 2048 | 768 | 128 | triton | 97.99 | ok | | wgrad wo | 10240 | 2048 | 768 | 128 | hipblas | 122.84 | ok | | fwd wi | 6400 | 2048 | 512 | 512 | default | 1.11 | ok | | fwd wi | 6400 | 2048 | 512 | 512 | triton | 29.18 | ok | | fwd wi | 6400 | 2048 | 512 | 512 | hipblas | 25.16 | ok | | fwd wo | 6400 | 2048 | 512 | 512 | default | 1.34 | ok | | fwd wo | 6400 | 2048 | 512 | 512 | triton | 30.86 | ok | | fwd wo | 6400 | 2048 | 512 | 512 | hipblas | 16.43 | ok | | dgrad wi | 6400 | 2048 | 512 | 512 | default | 1.27 | ok | | dgrad wi | 6400 | 2048 | 512 | 512 | triton | 27.83 | ok | | dgrad wi | 6400 | 2048 | 512 | 512 | hipblas | 17.26 | ok | | wgrad wi | 6400 | 2048 | 512 | 512 | default | 0.88 | ok | | wgrad wi | 6400 | 2048 | 512 | 512 | triton | 1.82 | ok | | wgrad wi | 6400 | 2048 | 512 | 512 | hipblas | 13.86 | ok | | dgrad wo | 6400 | 2048 | 512 | 512 | default | 1.06 | ok | | dgrad wo | 6400 | 2048 | 512 | 512 | triton | 27.49 | ok | | dgrad wo | 6400 | 2048 | 512 | 512 | hipblas | 30.17 | ok | | wgrad wo | 6400 | 2048 | 512 | 512 | default | 0.87 | ok | | wgrad wo | 6400 | 2048 | 512 | 512 | triton | 1.82 | ok | | wgrad wo | 6400 | 2048 | 512 | 512 | hipblas | 13.80 | ok | Disclaimers disclaimers The information presented in this document is for informational purposes only and may contain technical inaccuracies, omissions, and typographical errors. The information contained herein is subject to change and may be rendered inaccurate for many reasons, including but not limited to product and roadmap changes, component and motherboard version changes, new model and/or product releases, product differences between differing manufacturers, software changes, BIOS flashes, firmware upgrades, or the like. Any computer system has risks of security vulnerabilities that cannot be completely prevented or mitigated. AMD assumes no obligation to update or otherwise correct or revise this information. However, AMD reserves the right to revise this information and to make changes from time to time to the content hereof without obligation of AMD to notify any person of such revisions or changes. THIS INFORMATION IS PROVIDED ‘AS IS.” AMD MAKES NO REPRESENTATIONS OR WARRANTIES WITH RESPECT TO THE CONTENTS HEREOF AND ASSUMES NO RESPONSIBILITY FOR ANY INACCURACIES, ERRORS, OR OMISSIONS THAT MAY APPEAR IN THIS INFORMATION. AMD SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR ANY PARTICULAR PURPOSE. IN NO EVENT WILL AMD BE LIABLE TO ANY PERSON FOR ANY RELIANCE, DIRECT, INDIRECT, SPECIAL, OR OTHER CONSEQUENTIAL DAMAGES ARISING FROM THE USE OF ANY INFORMATION CONTAINED HEREIN, EVEN IF AMD IS EXPRESSLY ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. AMD, the AMD Arrow logo, insert all other AMD trademarks used in the material here per AMD Trademarks and combinations thereof are trademarks of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies. Insert any third party trademark attribution here per AMD’s Third Party Trademark List. © Insert year written Advanced Micro Devices, Inc. All rights reserved