{"slug": "group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient", "title": "Group-GEMM in XLA: Making Mixture-of-Experts Fast and Memory-Efficient", "summary": "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.", "body_md": "# Group-GEMM in XLA: Making Mixture-of-Experts Fast and Memory-Efficient[#](#group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient)\n\n[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).\nBut under the hood, an MoE layer turns into something that GPUs and math\nlibraries have historically been bad at: a large collection of small,\nindependent matrix multiplications, each with a different, data-dependent\nshape.\n\nIn JAX, users can describe this collection of matrix multiplications using an operation called `ragged-dot`, as shown in the following code snippet:\n\n``` php\n# lhs shape: (5, 4) -> 5 total input rows/tokens\nlhs = jnp.array([\n    [1.0, 0.5, 0.0, 1.0],\n    [0.2, 1.0, 0.5, 0.0],\n    [1.0, 1.0, 1.0, 1.0],\n    [0.1, 0.2, 0.3, 0.4],\n    [0.5, 0.5, 0.5, 0.5]\n], dtype=jnp.float32)\n\n# rhs shape: (2, 4, 3) -> 2 expert weight matrices of shape (4, 3)\nrhs = jnp.array([\n    [[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\n    [[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\n], dtype=jnp.float32)\n\n# Group sizes: Expert 0 processes 3 rows, Expert 1 processes 2 rows (3 + 2 = 5 = M)\ngroup_sizes = jnp.array([3, 2], dtype=jnp.int32)\n\n# Compute ragged dot product\noutput = jax.lax.ragged_dot(lhs, rhs, group_sizes)\n```\n\nHence, 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.\nIn 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.\n\nThe naive way to run this operation is to pad everything into a single dense batched\nmatmul — which wastes both compute (on the padding) and memory (on the copies).\nThe better way is a primitive called *Group-GEMM*: a batched matmul in which\neach matrix in the batch can have a different shape, computed in a single\noperation.\n\nThis post walks through how XLA now lowers the `ragged-dot` operation to real\nGroup-GEMM through two GPU backends:\n\n- an AMD-specific backend built on *hipBLASLt* (AMD ROCm), and\n- a new portable code-generation backend built on *Triton* that runs\non both CUDA and ROCm.\n\nBy the end, you’ll understand what MoE looks like to a compiler, what `ragged-dot`\nis, why Group-GEMM is the natural solution, and how each of the two backends\nmakes it work on the GPU.\n\nThe post is organized in five parts:\n\n- **[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.\n- **[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.\n- **[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.\n- **[Part 4 — Performance Evaluation](#part-4-performance-evaluation)** measures both backends against the padded baseline on MI300 and MI350, using MaxText MoE models.\n- **[Part 5 — Conclusion and Next Steps](#part-5-conclusion-and-next-steps)** wraps up and outlines what comes next.\n\n## Part 1 — Why MoE Turns Into a “Ragged” Matmul Problem[#](#part-1-why-moe-turns-into-a-ragged-matmul-problem)\n\n### A Quick Refresher: Dense FFN vs. Mixture-of-Experts[#](#a-quick-refresher-dense-ffn-vs-mixture-of-experts)\n\nIn a standard Transformer, every token passes through the same Feed-Forward\nNetwork (FFN). If you want the model to “know more,” you make that FFN bigger —\nbut then *every* token pays for the extra parameters, and compute grows with\nmodel size.\n\n*Mixture-of-Experts* breaks that link. Instead of one big FFN, an MoE layer\nhas many smaller FFNs called *experts*. A small *router* (a gating network)\nlooks at each token and picks the top-k experts that should process it (often\nk=1 or k=2). Only those experts run for that token.\n\nThe 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.\n\n### The Consequence for a Compiler: Variable-Size Groups[#](#the-consequence-for-a-compiler-variable-size-groups)\n\nHere’s the crucial systems detail. Because the router assigns each token to\nexperts *based on the token’s content*, the number of tokens that land on each\nexpert is different, and it is only known at runtime. One expert might receive\n2,000 tokens this step; its neighbor might receive 50; the split changes with\nevery batch.\n\nIf we gather all the tokens routed to the same expert into a contiguous block, an MoE layer becomes:\n\n- A stack of token rows, partitioned into `G` contiguous groups (one per\nexpert), where group`g` has`m_g` rows and the`m_g` values vary and are\nruntime-dependent.\n- 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]` .\n\nFigure 2 illustrates groups with unequal sizes: each color represents an expert.\n\nThe mental model to keep for the rest of this post:\n\n**An MoE layer is a batch of matmuls whose left-hand-side rows are split into\ncontiguous, variable-size groups.**\n\nThat “variable-size groups” property is what we call *ragged*.\n\n### The Naive Support, and Why It Hurts: Padded-Dot[#](#the-naive-support-and-why-it-hurts-padded-dot)\n\nGPUs are well suited to regular, dense shapes. So the naive solution is to force the ragged problem back into a dense one:\n\n1. Pick a fixed capacity `C` (a maximum number of tokens per expert).\n2. Pad every expert’s token block up to `C` rows.\n3. Now every block has the same shape, so we can run one big **batched GEMM** over`E × [C, k] × [k, n]` .\n\nFigure 3 illustrates this padding process: every expert’s block is padded to the same\n`C` rows, where the grey padding is compute the GPU performs for nothing.\n\nThis works, but it is wasteful in two ways:\n\n- **Wasted compute.** Any expert that received fewer than`C` tokens still does a full`C` -row matmul; the padding rows are pure overhead.\n- **Wasted memory and bandwidth.** The compact ragged input has to be\nexpanded into`E` padded matrices, and the outputs have to be compacted again afterwards. Those scatter/gather and pre/post-processing steps cost\nextra memory and memory traffic.\n\nThis “expand into padded matrices” approach is exactly the XLA **baseline** we’ll compare against later.\n\n### The Ragged-Dot Operation[#](#the-ragged-dot-operation)\n\nRather than having users replace their MoE with a dense problem, JAX lets you describe the ragged\ncomputation directly with `ragged_dot`.\n\nThe basic form:\n\n```\njax.lax.ragged_dot(lhs, rhs, group_sizes, ...)\n```\n\nwith the shapes:\n\n- `lhs` is 2D:`[m, k]` — all the tokens, stacked, in group order.\n- `rhs` is 3D:`[group, k, n]` — one weight matrix per group (per expert).\n- `group_sizes` is 1D:`[group]` — how many rows of`lhs` belong to each group.\n- the output is 2D.\n\nThe `group_sizes` operand is the key: it is a small vector, known only at\nruntime, that partitions the `m` rows of `lhs` into contiguous groups. For\nexample `group_sizes = [m_g0, m_g1, m_g2]` means the first `m_g0` rows are\ngroup 0, the next `m_g1` rows are group 1, and so on.\n\nNote that `jax.lax.ragged_dot` is a JAX-level operation. When the JAX program is\ncompiled, it is converted into the HLO `ragged-dot` operation (`kRaggedDot`),\nwhich is the actual input XLA receives and lowers.\n\nFigure 4 illustrates the operands of a ragged-dot.\n\nJAX also offers a more general form:\n\n```\njax.lax.ragged_dot_general(lhs, rhs, group_sizes, ragged_dimension, ...)\n```\n\nwhich lets you choose which dimension is the ragged one. That choice yields the three variants listed in Table 1 and drawn in Figure 5:\n\n*Table 1: The three `ragged-dot` variants exposed by `jax.lax.ragged_dot_general`.*\n\n| Variant | Ragged dimension | What it means | Group index in output? | \n|---|---|---|---|\n| **Non-contracting** (`m` ) | the “rows” dimension | the classic MoE case | No | \n| **Contracting** (`k` ) | the reduction dimension | groups along the contracted axis | Yes | \n| **Batch** (`b` ) | a batch dimension | effectively a regular batched GEMM | — | \n\nInside XLA these correspond to the `kRaggedDot` variants\n*kRaggedNonContracting*, *kRaggedContracting*, and *kRaggedBatch*. For MoE,\nthe one that matters most is *non-contracting*, where the rows (tokens) are ragged.\n\n### Why Group-GEMM Is *the* Suitable Solution[#](#why-group-gemm-is-the-suitable-solution)\n\nGroup-GEMM is an API for matrix multiplication that generalizes\nthe usual *batched* GEMM: instead of requiring every matrix in the batch to\nshare the same shape, it lets each one differ — in size, transposition, activation, and scaling factors. In other words:\n\n**Group-GEMM = a set of independent GEMMs of different shapes, computed in a single operation.**\n\nThat 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.\nHence, Group-GEMM can take care of “unpacking” the ragged\noperand and computing all of the per-expert products in one kernel launch, with no padding, as shown in Figure 6.\n\nCompared to the padded-dot baseline, this directly removes both costs described\nin [The Naive Support, and Why It Hurts: Padded-Dot](#the-naive-support-and-why-it-hurts-padded-dot):\n\n- No padding FLOPs — each group does exactly the work it needs.\n- No expansion/copies — the compact ragged layout is consumed as-is.\n\nThe rest of this post is about the two backends XLA uses to turn `ragged-dot` into real Group-GEMM for AMD GPUs.\n\n## Part 2 — Backend #1: hipBLASLt Group-GEMM (AMD ROCm)[#](#part-2-backend-1-hipblaslt-group-gemm-amd-rocm)\n\n### Overview[#](#overview)\n\nAMD’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.\n\nXLA already had a hipBLASLt backend for ordinary `dot` operations, so the natural move was to extend that existing\nbackend to also lower `ragged-dot` down to hipBLASLt Group-GEMM.\n\nFigure 7 shows the lowering path of a plain `dot`, and Figure 8 the `ragged-dot` path added alongside it.\n\nThe 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.\n\n```\nstruct GroupedGemmConfig {\n  uint64_t m, n, k;\n  uint64_t batch_count;\n  uint64_t group_count;\n  int64_t lhs_leading_dim_stride;\n  int64_t rhs_leading_dim_stride;\n  int64_t c_leading_dim_stride;\n  int64_t output_leading_dim_stride;\n  blas::Transpose trans_a, trans_b;\n  bool must_swap_operands;\n  xla::complex128 alpha;\n  double beta;\n  blas::DataType type_a, type_b, type_c, type_d;\n  int64_t stride_ragged_dim;\n  int64_t stride_group_dim;\n  int64_t c_stride_ragged_dim;\n  int64_t output_stride_ragged_dim;\n  // PrecisionConfig-level algorithm\n  xla::PrecisionConfig::Algorithm precision_algorithm;\n  int64_t compute_precision;\n  RaggedDotMode ragged_mode;\n  std::optional<blas::ComputationType> compute_type;\n\n  static absl::StatusOr<GroupedGemmConfig> FromProto(\n      const xla::GroupedGemmConfigProto& proto);\n  xla::GroupedGemmConfigProto ToProto() const;\n};\n\nclass GroupedMatmulPlan : public gpu::BlasLt::MatmulPlan {\n   public:\n    friend class BlasLt;\n\n    GroupedMatmulPlan(const BlasLt& blas_lt, const gpu::GroupedGemmConfig& cfg)\n        : blas_lt_(blas_lt), cfg_(cfg) {}\n\n    ~GroupedMatmulPlan() override = default;\n\n    absl::Status ExecuteOnStream(\n        Stream* stream, const gpu::BlasLt::MemoryArgs& args,\n        blas::ProfileResult* profile_result) const override;\n\n    absl::StatusOr<std::vector<MatmulAlgorithm>> GetAlgorithms(\n        size_t max_algorithm_count, size_t max_workspace_size) const override;\n\n    absl::Status SetAlgorithm(const MatmulAlgorithm& algorithm) override {...}\n\n   private:\n    ...\n\n    const BlasLt& blas_lt_;\n    gpu::GroupedGemmConfig cfg_;\n    Epilogue epilogue_ = Epilogue::kDefault;\n    std::unique_ptr<hipblaslt_ext::GroupedGemm> grouped_gemm_;\n    mutable std::optional<hipblasLtMatmulAlgo_t> algorithm_;\n    mutable bool algorithm_dirty_ = false;\n    mutable DeviceAddressBase saved_address_workspace_{};\n    // Saved default activation parameters from hipBLASLt\n    int32_t activation_type_ = 0;\n    int8_t bias_type_ = 0;\n  };  // class GroupedMatmulPlan\n```\n\n(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))\n\n### The hipBLASLt Group-GEMM API[#](#the-hipblaslt-group-gemm-api)\n\nLaunching a Group-GEMM with hipBLASLt follows the fixed sequence of steps shown in Figure 9:\n\nThe important concept here is **user arguments**. hipBLASLt describes each\nindividual 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.\n(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).\n\n### The Compile-Time vs. Runtime Problem[#](#the-compile-time-vs-runtime-problem)\n\nHere 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.\n\nFortunately, `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.\nOnly one condition must be guaranteed: the total problem\nsize 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.\n(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.)\n\nIn XLA, the definition of user arguments is therefore delayed until runtime.\nJust before launching the Group-GEMM, XLA executes small update kernels on the GPU that\ncompute and populate the descriptors with the actual sizes of the groups contained in `group_sizes`.\n\nThese kernels proceed as sketched in Figure 10:\n\nEach 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.\n\n### Autotuning[#](#autotuning)\n\nhipBLASLt’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:\n\nHowever, `ragged-dot` needs one special accommodation. By default, the autotuner profiles candidates on buffers filled with random values. But for\nGroup-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.\n\nTo 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\nsize `M` of the ragged dimension.\n\nFor 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`.\n\nEnforcing `sum(group_sizes) == M` guarantees two things:\n\n1. **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.\n2. **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.\n\n## Part 3 — Backend #2: Triton Ragged-Dot (Portable CUDA + ROCm)[#](#part-3-backend-2-triton-ragged-dot-portable-cuda-rocm)\n\n### Motivation[#](#motivation)\n\nEven though the `hipBLASLt` path provides impressive performance (see [Part 4 — Performance Evaluation](#part-4-performance-evaluation)),\nit is ROCm-specific and not even available on every AMD GPU target.\nOn the other hand, Triton is a generic backend that can be lowered to most GPU targets.\nIn 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.\n\nConsequently, we decided to propose a Triton backend as an alternative to efficiently handle the `ragged-dot` operation in XLA.\nIn 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.\n\n### The Compilation Pipeline[#](#the-compilation-pipeline)\n\nFigure 12 shows the main steps that have been enhanced to enable the lowering of the `ragged-dot` op:\n\n1. **Backend selection.** XLA decides which path serves this`ragged-dot` ,\nfollowing 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.\nNote 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.\n2. **Support validation.** Before committing, XLA checks the`ragged-dot` is\nactually supported: it must be a`ragged-dot` with exactly one ragged\nleft-hand-side dimension, the tiling-propagation feature must be enabled, and the element types must be supported.\n\n1. **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` .\nIt 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\nform a 2-D grid that XLA linearizes into a 1-D stream of program IDs. Whichever\naxis is placed*outer* varies slowly, and the*inner* one varies fastest across\nconsecutive programs. The heuristic exploits this to maximize L2 reuse: if\nthe LHS (activations) is*smaller in bytes* than the RHS (weights), it swaps M\nand N so that N is traversed slowly — keeping the small LHS tile-row resident in\nL2 while streaming across the many N tiles. Concretely it compares`ByteSizeOf(lhs)` against`ByteSizeOf(rhs)` and, when the LHS is smaller, emits\na 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\nraw operand byte-sizes are the wrong thing to compare. The ragged version keeps\nthe*same mechanism* (permute the parallel grid axes before the program-id is\ndelinearized) 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\nhow tall each group’s LHS activation slice really is (the true per-group sizes\nare only known at runtime).\n2. **Autotuning.** The pipeline generates candidate tiling configurations\nspecialized 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) ).\n3. **Emission.** Finally, XLA emits the actual Triton kernel for the selected variant.\n\n### What the Emitted Triton Kernel Does, per Variant[#](#what-the-emitted-triton-kernel-does-per-variant)\n\nNow, 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.\nAs we argued before, the Group-GEMM algorithm is best suited to handle ragged-dot operations.\nWe 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.\n\nThe 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:\n\n- where the raggedness lives,\n- whether `group_sizes` is loaded at all,\n- how the group offset is used.\n\nNote 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:\n\n- 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\n- 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).\n\nTogether these keep each expert’s weights cache-resident while its tiles are being computed.\n\n#### Non-Contracting (the MoE Case)[#](#non-contracting-the-moe-case)\n\n**Shape:** `LHS[M_total, K] × RHS[G, K, N] → out[M_total, N]`. The ragged axis\nis **M** (the stacked tokens). Every output row `m` belongs to exactly one group\n`g = group(m)`, and that group selects *which* weight matrix `RHS[g]` is used.\n\nThe emitted kernel has (conceptually) this signature:\n\n``` python\n@triton.jit\ndef ragged_dot_non_contracting(\n    lhs_ptr,          # [M_total, K]  all tokens, stacked in group order\n    rhs_ptr,          # [G, K, N]     one weight matrix per group (expert)\n    out_ptr,          # [M_total, N]  output, same row layout as lhs\n    group_sizes_ptr,  # [G]           rows per group; sum == M_total\n    M_total, K, N, G,\n    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,\n):\n    pid_m = tl.program_id(0)      # which M-tile this program owns\n    pid_n = tl.program_id(1)      # which N-tile this program owns\n    row_start = pid_m * BLOCK_M   # first LHS row of this tile (in [0, M_total))\n    ...\n```\n\nThe subtle part of the kernel is to “know” which portion of the data the kernel is expected to compute.\nIndeed, 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`:\n\n```\ngroup_sizes         = [ m0, m1,     m2,       … ]\nprefix (exclusive)  = [ 0,  m0,     m0+m1,    … ]   # group start rows\nprefix (inclusive)  = [ m0, m0+m1,  m0+m1+m2, … ]   # group end rows\n```\n\nThe exclusive prefix sum gives the start row of every group; the inclusive\none gives the end row. The program’s group `g` is simply the group whose\n`[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\n`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).\n\nTable 2 below works through a tiny example — `group_sizes = [3, 2, 4]` (so `M_total = 9`) and\n`BLOCK_M = 3`:\n\n*Table 2: Worked example of the group lookup for `group_sizes = [3, 2, 4]` with `BLOCK_M = 3`.*\n\n| `pid_m` | `row_start` | falls in group `g` | `group_start_row` | uses weights | \n|---|---|---|---|---|\n| 0 | 0 | 0 (rows 0–2) | 0 | `RHS[0]` | \n| 1 | 3 | 1 (rows 3–4) | 3 | `RHS[1]` | \n| 2 | 6 | 2 (rows 5–8) | 5 | `RHS[2]` | \n\n(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.)\n\nPer program:\n\n1. Compute the prefix sum over `group_sizes` ; from its own M-tile position\n(`row_start = pid_m * BLOCK_M` ) find which group`g` it falls into (the\ngroup whose`[start, end)` row interval contains`row_start` ) and read that group’s start row`group_start_row = prefix_exclusive[g]` .\n2. Use the group index `g` to select the matching RHS slice`RHS[g, :, :]` , i.e. that expert’s weight matrix.\n3. 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` .\n\n`group_sizes` is loaded — it drives both the group selection and the row\noffset. Figure 14 summarizes this program flow.\n\n#### Contracting[#](#contracting)\n\n**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.\n\nThe emitted kernel has (conceptually) this signature:\n\n``` python\n@triton.jit\ndef ragged_dot_contracting(\n    lhs_ptr,          # [M, K_total]  contracted axis is ragged\n    rhs_ptr,          # [K_total, N]\n    out_ptr,          # [G, M, N]     one output matrix per group\n    group_sizes_ptr,  # [G]           K-slice size per group; sum == K_total\n    M, K_total, N, G,\n    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,\n):\n    pid_g = tl.program_id(0)      # which group this program owns (→ output slab g)\n    pid_m = tl.program_id(1)      # which M-tile\n    pid_n = tl.program_id(2)      # which N-tile\n    ...\n```\n\nThe 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:\n\n- 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.\n- 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\ngroups 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` .\n\nTo illustrate this, Figure 16 shows the *same* pool of blocks\nunder 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.\n\nWhat the program needs from `group_sizes` is where along K its group’s\nslice begins and ends — the ragged K partition is still runtime data. It reads that from the same prefix-sum table:\n\n```\ngroup_sizes         = [ k0,   k1,      k2,        … ]   # K-rows per group\nprefix (exclusive)  = [ 0,    k0,      k0+k1,     … ]   # k_start for each group\nprefix (inclusive)  = [ k0,   k0+k1,   k0+k1+k2,  … ]   # k_end   for each group\n```\n\nFor 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.\n\nPer program:\n\n1. Take its group `g = pid_g` from the grid, and read`k_start` /`k_end` from\nthe prefix sum over`group_sizes` — this time the offset locates the group’s slice along the contracted axis.\n2. Accumulate a partial dot over just that `[k_start, k_end)` slice of the\ncontracted axis.\n3. Write to `out[g, :, :]` — the output position indexed by the group`g` .\n\n`group_sizes` is loaded: it provides the per-group start/end offsets along K.\nFigure 17 summarizes this program flow.\n\n#### Batch[#](#batch)\n\n**Shape:** `LHS[B, M, K] × RHS[B, K, N] → out[B, M, N]` — an ordinary batched\nGEMM. The ragged dimension is a *batch* dimension, but because every batch\nelement is fully computed there are no data-dependent group offsets to work\nout: each program simply owns one `(batch, M-tile, N-tile)` and does a full\ncontraction over K.\n\nThe emitted kernel has (conceptually) this signature:\n\n``` python\n@triton.jit\ndef ragged_dot_batch(\n    lhs_ptr,          # [B, M, K]\n    rhs_ptr,          # [B, K, N]\n    out_ptr,          # [B, M, N]\n    B, M, K, N,\n    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,\n):\n    pid_b = tl.program_id(0)      # which batch element this program owns\n    pid_m = tl.program_id(1)      # which M-tile\n    pid_n = tl.program_id(2)      # which N-tile\n    # no prefix sum: batch index and tile offsets are enough to address the data\n    ...\n```\n\nBecause the batch index `pid_b` comes straight from the grid and every batch\nelement 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.\n\nPer program:\n\n1. Take its batch index `b = pid_b` from the grid (and its`pid_m` ,`pid_n` tile positions).\n2. Accumulate a tiled dot product over K between `LHS[b]` and`RHS[b]` .\n3. Write the `[BLOCK_M, BLOCK_N]` result to`out[b, M-tile, N-tile]` .\n\n`group_sizes` is not loaded: it is used only as a scheduling hint.\nFigure 19 summarizes this program flow.\n\n### Autotuning and the `group_sizes` Seeding Problem[#](#autotuning-and-the-group-sizes-seeding-problem)\n\nThis backend runs into the very same issue as the `hipBLASLt` autotuner in\n[Autotuning](#autotuning).\n\nDuring autotuning, candidate kernels run on synthetic input buffers, which\nare filled with random bytes by default. That’s harmless for a normal GEMM.\nBut `group_sizes` is made of control values that determine loop bounds and\nmemory offsets.\n\nThe fix mirrors the `hipBLASLt` [one](#autotuning): before profiling, the profiler writes\nbalanced, 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.\n\n## Part 4 — Performance Evaluation[#](#part-4-performance-evaluation)\n\nEvery 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):\nthe 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.\n\n### Testbed Characteristics[#](#testbed-characteristics)\n\nThe performance evaluations have been performed on two AMD GPU targets, MI300 (gfx942) and MI350 (gfx950). The software configuration was as follows:\n\n- **Operating System** : Ubuntu 22.04\n- **ROCm version** : 7.2.4\n- **JAX version** : upstream (https://github.com/jax-ml/jax) commit 0eb0f676ba22b15b9dfe29518ab0d40277b7138d\n- **XLA version** : ROCm (https://github.com/ROCm/xla) branch pr/ragged-dot-6-e2e-tests  (upstream commit e5d008bb03d9b133b0881daf3d108adea8e6625b + ragged-dot Triton support commits)\n- **MaxText version** : ROCm (https://github.com/ROCm/maxtext) commit fb25ae8c850fa2b57dd5eb7ebfbe399f08534e69\n\nNote 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.\n\nThe `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`.\nFurthermore, note that this backend is only supported with the Triton experimental tiling feature.\nAs a consequence, the flag `xla_gpu_experimental_enable_tiling_propagation=true` must be set alongside the `xla_gpu_experimental_triton_ragged_dot` flag.\nLast, 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`.\n\n### Use Case: MoE Model Training Using MaxText[#](#use-case-moe-model-training-using-maxtext)\n\n[MaxText](https://github.com/AI-Hypercomputer/maxtext) is a high-performance,\nhighly scalable, open-source LLM library and reference implementation written in\nPython/JAX. Because it is written on top of JAX, its MoE layers ultimately lower to\nJAX primitives — and, on the sparse path, to the `jax.lax.ragged_dot` op that the\nGroup-GEMM backends described above accelerate. This makes MaxText an ideal\nframework in which to evaluate the new `ragged-dot` backends on realistic MoE models.\n\n#### Two MoE Modes: Drop vs. Dropless[#](#two-moe-modes-drop-vs-dropless)\n\nBefore 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.\n\n- **Drop mode.** Each expert is given a*fixed capacity*`expert_capacity` (`expert_capacity = ceil(tokens_per_batch / E) × capacity_factor` ). Tokens routed beyond an\nexpert’s capacity are simply dropped (they skip the FFN and contribute\nnothing), and under-filled experts are padded up to`expert_capacity` . This keeps every\nexpert’s matmul a regular, statically-shaped`[expert_capacity, K] × [K, N]` , which is\nhardware-friendly, but it wastes FLOPs on the padding of under-filled experts\nand loses information on the dropped tokens of over-filled ones.\n- **Dropless mode.** No token is ever dropped and no capacity padding is\nintroduced: every routed (token, expert) pair is computed exactly. This\npreserves the full model quality, but it requires handling the*ragged* ,\ndata-dependent per-expert sizes — either by physically gathering tokens per\nexpert into one ragged matmul (the block-sparse formulation introduced by[MegaBlocks \\[Gale et al., 2022\\]](https://arxiv.org/abs/2211.15841) , which is\nexactly the`ragged-dot` /Group-GEMM path this post is about), or by running a\nfull dense compute over all tokens and masking afterwards.\n\nFigure 21 illustrates the buffers the two modes produce.\n\nThe rest of this section looks at how MaxText realizes these two modes through\nits three concrete compute paths. For a complementary, MaxText-focused deep dive\ninto the dropless MoE path and its optimization, see the AMD ROCm blog\n[*Dropless MoE in MaxText*](https://rocm.blogs.amd.com/software-tools-optimization/maxtext-dropless-moe/README.html).\n\n#### The Three MoE Compute Paths in MaxText[#](#the-three-moe-compute-paths-in-maxtext)\n\nBefore diving into the models, it is worth noticing that `ragged-dot` is only one\nof several ways to run the expert FFNs and combine their outputs.\nIf 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\npaths — two of them implement the dropless mode (the sparse / ragged-dot path\nand the dense mask-and-combine path), while one implements the drop mode (the\ncapacity-based dense dispatch/combine path). The `ragged-dot`/Group-GEMM work\ndescribed so far powers just one of the two dropless paths.\n\nAll three share the same front-end — a router (gating network) picks, for every token, its top-k experts and a weight per selected expert:\n\n```\n# gate_logits: [B, S, E]\ntop_k_weights, top_k_indices = jax.lax.top_k(gate_logits, num_experts_per_tok)  # [B, S, k]\n```\n\nHowever, they differ in the way the selected (token, expert) pairs are turned into matmuls, and in the pre-/post-processing that surrounds those matmuls.\n\nFrom a user’s point of view, two config knobs allow choosing the path:\n\n- `sparse_matmul` —`True` picks the sparse (ragged-dot) path (the dropless,\nblock-sparse formulation of the MoE computation);`False` picks one of the two\ndense paths.\n- `capacity_factor` — inside the dense path,`> 0` enables token dropping\n(fixed per-expert capacity);`<= 0` keeps it dropless.\n\nFigure 22 shows how these two knobs select a path.\n\nTable 3 summarizes the trade-offs before we look at each path in detail:\n\n*Table 3: The three MaxText MoE compute paths and their trade-offs.*\n\n|  | **Path 1 — Sparse (ragged-dot)** | **Path 2 — Dense dropless** | **Path 3 — Dense drop** | \n|---|---|---|---|\n| config | `sparse_matmul = True` | `sparse_matmul = False` ,`capacity_factor ≤ 0` | `sparse_matmul = False` ,`capacity_factor > 0` | \n| expert op | `jax.lax.ragged_dot` | dense per-expert einsum over **all tokens** | dense **capacity-padded** batched einsum | \n| 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 | \n| 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 | \n| compute cost | ≈ `M × K × N` (only routed rows) | ≈ `E × (B·S) × K × N` (**every token × every expert** ) | ≈ `E × C × K × N` (`C` = capacity) | \n| drops tokens? | no | no | yes (overflow beyond `C` ) | \n| padding | none | none (but full dense compute) | per-expert capacity padding | \n\n##### Path 1 — Sparse / Ragged-Dot (Dropless)[#](#path-1-sparse-ragged-dot-dropless)\n\nSelected 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\nso that the expert FFNs become one big ragged matmul.\n\n**Pre-processing — permute.** The router’s `top_k_indices` are used to sort all\n`M = B·S·top_k` token-copies into contiguous per-expert groups, and a `bincount`\nproduces the `group_sizes` vector (real routed counts, summing to `M`):\n\n```\n# permute(): sort token rows by their assigned expert, count per-expert sizes\nsorted_selected_experts = jnp.argsort(jnp.ravel(selected_experts))\nsorted_inputs = _sort_activations(replicated_inputs, sorted_selected_experts)  # [M, K]\ngroup_sizes  = jnp.bincount(jnp.ravel(selected_experts), length=num_experts)   # [E], sums to M\n```\n\n**Compute — Group-GEMM.** The gathered tokens and the per-expert weights feed\ndirectly into `jax.lax.ragged_dot`, which can be later lowered to hipBLASLt / Triton Group-GEMM by XLA.\n\n```\n# gmm(): the expert FFN as a single ragged/group matmul\nout = jax.lax.ragged_dot(lhs=sorted_inputs,        # [M, K]  tokens, grouped by expert\n                         rhs=w,                     # [E, K, N] expert weights\n                         group_sizes=group_sizes,   # [E]\n                         preferred_element_type=dtype)\n```\n\n**Post-processing — un-permute.** The results are scattered back to the original\ntoken order and the top-k contributions are summed with the router weights:\n\n```\n# unpermute(): invert the sort, then weighted-combine the top-k expert outputs\nunsorted = _sort_activations(out, jnp.argsort(sorted_selected_experts))\noutput   = jnp.einsum(\"BKE,BK -> BE\", unsorted.reshape(-1, top_k, N), top_k_weights)\n```\n\nThe whole path is summarized in Figure 23.\n\n**Cost:** only the real routed rows are ever multiplied — compute is\n≈ `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.\nHowever, 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.\n\n##### Path 2 — Dense Dropless (Mask-and-Combine)[#](#path-2-dense-dropless-mask-and-combine)\n\nSelected by `sparse_matmul = False` and `capacity_factor ≤ 0`. In this path, there is no\ngather at all: every token is sent through every expert, and the router\nweights (zero for non-selected experts) mask out the unwanted results at the end.\n\n**Pre-processing — none** (beyond turning top-k into a dense `[B,S,E]` weight\nmatrix, `0` where an expert was not selected):\n\n```\nweights = reshape_and_update_weights(top_k_weights, top_k_indices)  # [B,S,E], 0 off-selection\n```\n\n**Compute — full dense einsums.** Each expert processes the *entire* token batch:\n\n``` php\n# up / gate: every token × every expert\nlayer_w0 = jnp.einsum(\"BSM,EMH -> BSEH\", inputs, w0_kernel)   # [B,S,E,H]\nlayer_w1 = jnp.einsum(\"BSM,EMH -> BSEH\", inputs, w1_kernel)\ninter    = activation(layer_w0) * layer_w1\n# down\ninter    = jnp.einsum(\"BSEH,EHM -> BSEM\", inter, wo_kernel)   # [B,S,E,M]\n```\n\n**Post-processing — weighted combine.** The router weights collapse the expert\naxis, keeping only each token’s selected experts:\n\n``` php\noutput = jnp.einsum(\"BSEM,BSE -> BSM\", inter, weights)        # [B,S,M]\n```\n\nThe whole path is summarized in Figure 24.\n\n**Cost:** dropless and permutation-free, but compute is ≈ `E × (B·S) × K × N` —\nevery token pays for every expert, so it scales with the *full* expert count.\nThis is simple and hardware-friendly (plain dense matmuls, no ragged handling) but\nbecomes prohibitively expensive as `E` grows.\n\n##### Path 3 — Dense Drop (Capacity Dispatch/Combine)[#](#path-3-dense-drop-capacity-dispatch-combine)\n\nSelected by `sparse_matmul = False` and `capacity_factor > 0`. This is the\nclassic fixed-capacity MoE: each expert gets a static slot of\n`C = ceil(tokens_per_batch / E) × capacity_factor` rows; tokens beyond `C` are dropped.\n\n**Pre-processing — dispatch mask (with dropping).** `generate_masks` builds a\none-hot `dispatch_mask` that assigns each kept (token, expert) pair to a capacity\nslot and truncates any overflow:\n\n```\nexpert_capacity = ceil(tokens_per_batch / num_experts) * capacity_factor\n# ... one-hot into [B,S,E,C]; positions past `expert_capacity` are dropped ...\ndispatch_mask, combine_mask = generate_masks(top_k_indices, weights)\n```\n\n**Compute — capacity-padded batched einsum.** Tokens are scattered into the\n`[E, B, C, M]` capacity buffer, run through the experts, and the result stays in\ncapacity layout:\n\n``` php\ndispatch = jnp.einsum(\"BSM,BSEC -> EBCM\", inputs, dispatch_mask)     # scatter to slots\nlayer_w0 = jnp.einsum(\"EBCM,EMH -> EBCH\", dispatch, w0_kernel)       # up/gate\ninter    = activation(layer_w0) * jnp.einsum(\"EBCM,EMH -> EBCH\", dispatch, w1_kernel)\ninter    = jnp.einsum(\"EBCH,EHM -> EBCM\", inter, wo_kernel)          # down\n```\n\n**Post-processing — combine mask.** The `combine_mask` gathers each token’s\noutputs back from its capacity slots and applies the router weights:\n\n``` php\noutput = jnp.einsum(\"EBCM,BSEC -> BSM\", inter, combine_mask)          # gather back to tokens\n```\n\nThe whole path is summarized in Figure 25.\n\n**Cost:** compute is ≈ `E × C × K × N`. Because `C` is fixed, under-filled experts\nwaste FLOPs on padding rows and over-filled experts drop tokens (an accuracy\ntrade-off). The capacity buffer also inflates the activation memory.\n\n##### What About Performance?[#](#what-about-performance)\n\nIn theory, the `ragged-dot`/Group-GEMM path:\n\n- 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\nit wins by a factor that grows with the number of experts — at the cost of the\npermute/un-permute sort.\n- 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\ndropping, while also using less activation memory.\n\nIn other words, Group-GEMM keeps the dropless quality of Path 2 with a\ncompute/memory profile close to the *ideal* (only real work, no padding), trading\nit for a sort and a ragged-aware kernel.\nBut the remaining question is: *how fast is the Group-GEMM kernel itself compared to these dense alternatives on real hardware?*\n\n#### MaxText Models[#](#maxtext-models)\n\nWe evaluated the performance improvement achieved by the new `ragged-dot` backends using MaxText MoE models.\nThe evaluated models are described below, and their configurations are given in\nTable 4 (MI300) and Table 5 (MI350):\n\n**“Routing-heavy” vs. “compute-heavy”:** each MoE layer spends time on (a) the\n*expert GEMM* itself (`ragged_dot`, ≈ M[tokens] × K[emb] × N[moe_mlp]) and\n(b) *routing* — gating top-k selection plus the token permute/sort/scatter that\nbuilds the `num_experts` ragged groups and the un-permute afterwards. A config is\n**routing-heavy** when (b) dominates: this happens with small `base_moe_mlp_dim`\n(N) (each expert GEMM does little arithmetic) and/or many experts (more small\ngroups → more launch + scatter/gather overhead, lower per-group occupancy). The\nopposite corner — large N, few experts — is compute-heavy (big dense expert\nGEMMs dominate, routing negligible).\n\n**mixtral-8x7b** *(MI300 & MI350)* — The baseline dense-attention MoE microbenchmark\n(emb 6144, 24 layers, 32 experts top-4, `dot_product` attention, no MLA). Used to\nprofile the raw ragged-dot expert GEMM without MLA complexity, at 8-way FSDP,\nseq 3072, grad-accum 1. The MI300 and MI350 configurations have identical shapes.\nIt is the primary knob-sweep model for tile sizes and backend (hipBLASLt vs. Triton vs. the default baseline).\n\n**mixtral-8x7b-large** *(MI300 & MI350)* — A deliberately enlarged mixtral variant\n(emb/K 8192, moe_mlp/N 28672, 32 layers, 64 experts top-8) that grows both the\nexpert-GEMM K/N dims and the number of ragged groups. Runs 8-way FSDP at seq 4096,\ngrad-accum 8. It is the compute-heavy corner of the sweep — big, dense expert GEMMs\nwith many groups — used to measure ragged-dot throughput at high arithmetic\nintensity and to check that larger tiles fit within HBM (192 GB HBM3 / 288 GB HBM3e).\n\n**deepseek2-16b** *(MI300 & MI350)* — A scaled-down DeepSeek-V2 test (emb 2048,\n27 layers, 64 experts top-6, 2 shared experts) that keeps the real MLA attention and\n`deepseek` decoder block. It validates the ragged-dot + MLA + YARN-RoPE code path\ncheaply, running autodetected FSDP, grad-accum 4, seq 512. Shared experts add a dense\nside-path alongside the routed experts, exercising the mixed dense/sparse MLP.\n\n**deepseek3-test** *(MI300 & MI350)* — A tiny DeepSeek-V3 CI proxy (emb 512) that\nkeeps the `deepseek` block + MLA attention but shrinks everything else. On MI300 it is\ncut to 16 layers / 32 experts (dense mlp 4096); on MI350 it retains the full 61-layer /\n256-expert topology with the tiny emb 512. Runs autodetected FSDP, grad-accum 4,\nseq 512 — a fast functional check of the ragged-dot + MLA path rather than a perf run.\n\n**deepseek3-671b-reduced** *(MI300 & MI350)* — A perf-representative slice of the\nflagship 671B DeepSeek-V3: it **preserves the real expert-GEMM shapes**\n(K/emb 7168, N/moe_mlp 2048, MLA) but cuts to 4 layers / 32 experts so it fits on 8\nGPUs. Because FSDP can’t hold the full model it is **expert-parallel** (EP=8,\n4 experts/GPU, FSDP off), grad-accum 8, seq 512. This makes the per-expert GEMM timing\nrepresentative of the true 671B while keeping the run tractable.\n\n**gemma4-26b** *(MI300 only)* — Gemma-4 26B MoE (emb 2816, 30 layers, 128 experts\ntop-8, 1 shared) with the `gemma4` decoder block and standard `dot_product` attention.\nSmall `base_moe_mlp_dim` (704) means many narrow expert GEMMs, stressing per-group\nlaunch/routing overhead rather than raw FLOPs. Runs autodetected FSDP, seq 1024,\ngrad-accum 4 — a good test for many-small-experts ragged-dot\nefficiency on the 192 GB HBM3 gfx942 platform.\n\n**qwen3-30b-a3b** *(MI300 & MI350)* — Qwen3 MoE (emb 2048, 48 layers, 128 experts\ntop-8, `qwen3_moe` block, QK-norm) with a very small expert width\n(`base_moe_mlp_dim` 768). Like gemma4 it is dominated by many thin expert GEMMs and\nrouting overhead rather than arithmetic intensity. MI300 and MI350 use identical shapes\n(autodetected FSDP, seq 1024, grad-accum 4). A\nlightweight but routing-heavy ragged-dot workload.\n\n**qwen3-235b-a22b** *(MI350 only)* — Large Qwen3 MoE (emb 4096, **94 layers**,\n128 experts top-8, moe_mlp 1536). The deep 94-layer stack combined with 128 experts\nmakes it the heaviest routing/attention workload in the tables; it runs 8-way FSDP,\nseq 2048, grad-accum 16. Sequence length and layer depth push activation memory, so\nremat and FSDP are essential to fit — a scale/depth stress test for ragged-dot.\n\n**llama4-17b-16e** *(MI350 only)* — Llama-4 17B, 16-expert, **top-1** routing (each\ntoken → a single expert) with wide experts (moe_mlp 8192) and a shared expert.\nThis combination makes it the\ncompute-heavy corner: a few large, dense expert GEMMs where the padded-dot baseline\nalready wastes little, so Group-GEMM has limited headroom to improve on it. Runs\nautodetected FSDP, seq 1024, grad-accum 4.\n\n#### MI300 (MI300X, gfx942, 192 GB HBM3, 8× GPU)[#](#mi300-mi300x-gfx942-192-gb-hbm3-8-gpu)\n\n*Table 4: MaxText MoE model configurations evaluated on MI300.*\n\n| 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) | \n|---|---|---|---|---|---|---|---|---|---|---|---|\n| mixtral-8x7b | 6144 | 20480 | 20480 | 24 | 32 | 4 | – | dot | 3072 | 1 | 8 / 1 / 1 | \n| mixtral-8x7b-large | 8192 | 28672 | 28672 | 32 | 64 | 8 | – | dot | 4096 | 8 | 8 / 1 / 1 | \n| deepseek2-16b | 2048 | 1408 | 5632 | 27 | 64 | 6 | 2 | MLA | 512 | 4 | auto / 1 / 1 | \n| deepseek3-test | 512 | 512 | 4096 | 16 | 32 | 8 | 1 | MLA | 512 | 4 | auto / 1 / 1 | \n| deepseek3-671b-reduced | 7168 | 2048 | 18432 | 4 | 32 | 8 | 1 | MLA | 512 | 8 | 1 / 8 / 1 | \n| gemma4-26b | 2816 | 704 | 2112 | 30 | 128 | 8 | 1 | dot | 1024 | 4 | auto / 1 / 1 | \n| qwen3-30b-a3b | 2048 | 768 | 768 | 48 | 128 | 8 | – | dot | 1024 | 4 | auto / 1 / 1 | \n\n#### MI350-Class (gfx950, 288 GB HBM3e, 8× GPU)[#](#mi350-class-gfx950-288-gb-hbm3e-8-gpu)\n\n*Table 5: MaxText MoE model configurations evaluated on MI350.*\n\n| 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) | \n|---|---|---|---|---|---|---|---|---|---|---|---|\n| mixtral-8x7b | 6144 | 20480 | 20480 | 24 | 32 | 4 | – | dot | 3072 | 1 | 8 / 1 / 1 | \n| mixtral-8x7b-large | 8192 | 28672 | 28672 | 32 | 64 | 8 | – | dot | 4096 | 8 | 8 / 1 / 1 | \n| deepseek2-16b | 2048 | 1408 | 5632 | 27 | 64 | 6 | 2 | MLA | 512 | 4 | auto / 1 / 1 | \n| deepseek3-test | 512 | 512 | 18432 | 61 | 256 | 8 | 1 | MLA | 512 | 4 | auto / 1 / 1 | \n| deepseek3-671b-reduced | 7168 | 2048 | 18432 | 4 | 32 | 8 | 1 | MLA | 512 | 8 | 1 / 8 / 1 | \n| qwen3-30b-a3b | 2048 | 768 | 768 | 48 | 128 | 8 | – | dot | 1024 | 4 | auto / 1 / 1 | \n| qwen3-235b-a22b | 4096 | 1536 | 1536 | 94 | 128 | 8 | – | dot | 2048 | 16 | 8 / 1 / 1 | \n| llama4-17b-16e | 5120 | 8192 | 8192 | 48 | 16 | 1 | 1 | dot | 1024 | 4 | auto / 1 / 1 | \n\n**Legend**:\n\n- `emb_dim (K)` = expert-GEMM contraction dim\n- `moe_mlp_dim (N)` = expert-GEMM output dim\n- `attn` = attention type (MLA = Multi-head Latent Attention — DeepSeek’s low-rank compressed KV attention)\n- `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)\n- `auto` =`-1` (autodetects available GPUs).\n\n### Micro-Benchmark[#](#micro-benchmark)\n\nIn 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.\nThis 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.\nAs 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.\n\n#### From an MoE Layer to 6 Ragged-Dot Classes[#](#from-an-moe-layer-to-6-ragged-dot-classes)\n\nAn 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.\n\nWhen you enumerate every ragged-dot XLA emits across the *forward + backward*\npass, they collapse into exactly 6 structural classes. Two operations belong\nto 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.\n\nIn the following, we use the notation given in Table 6:\n\n*Table 6: Notation used to describe the `ragged-dot` classes.*\n\n| symbol | meaning | \n|---|---|\n| `M` | ragged token dimension = `batch × seq × top_k` (the real routed-token count; no capacity padding on the`ragged_dot` path) | \n| `K` | expert-GEMM contraction / embedding dim (`base_emb_dim` ) | \n| `N` | expert-GEMM output dim (`base_moe_mlp_dim` ) | \n| `E` | number of experts (`num_experts` ) | \n| `group_sizes` | `int32[E]` , per-expert token counts, summing to`M` | \n\nThe `M` (token) axis is the **ragged** axis: it is partitioned into `E` groups, one per expert, and `dtype` is `bfloat16` throughout.\nThe six classes are listed in Table 7:\n\n*Table 7: The six structural `ragged-dot` classes emitted across the forward and backward pass.*\n\n| class | role | lhs | rhs | out | contract (lhs,rhs) | rhs_group | \n|---|---|---|---|---|---|---|\n| `fwd_wi` | fwd gate/up | `[M,K]` | `[E,K,N]` | `[M,N]` | `{1},{1}` | `{0}` | \n| `fwd_wo` | fwd down | `[M,N]` | `[E,N,K]` | `[M,K]` | `{1},{1}` | `{0}` | \n| `dgrad_wi` | dgrad up | `[M,N]` | `[E,K,N]` | `[M,K]` | `{1},{2}` | `{0}` | \n| `dgrad_wo` | dgrad down | `[M,K]` | `[E,N,K]` | `[M,N]` | `{1},{2}` | `{0}` | \n| `wgrad_wi` | wgrad up | `[M,K]` | `[M,N]` | `[E,K,N]` | `{0},{0}` | `{}` (empty) | \n| `wgrad_wo` | wgrad down | `[M,N]` | `[M,K]` | `[E,N,K]` | `{0},{0}` | `{}` (empty) | \n\nAll classes are `bfloat16`, with `lhs_ragged_dimensions = {0}` and\n`group_sizes = int32[E]`.\n\n- 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\ngate/up projection, and`interm[M,N] × w[E,N,K] → out[M,K]` for the down\nprojection.\n- 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.\n- The **wgrad** classes (`wgrad_wi` ,`wgrad_wo` ) are`[M,·] × [M,·] → [E,·,·]` Group-GEMMs that contract the ragged token axis and produce one weight-gradient\nslab per expert (empty`rhs_group_dimensions` ).\n\nFor example, `fwd_wi` — the gate/up projection — is simply:\n\n``` python\nimport jax.numpy as jnp\nfrom jax import lax\n\nlhs = jnp.zeros((M, K), jnp.bfloat16)          # tokens\nrhs = jnp.zeros((E, K, N), jnp.bfloat16)       # expert weights\n                                               # group_sizes: int32[E], sums to M\nout = lax.ragged_dot(lhs, rhs, group_sizes,    # -> [M, N]\n                     preferred_element_type=jnp.bfloat16)\n```\n\nwhile a weight-gradient such as `wgrad_wi` needs the general form with explicit\ndimension numbers:\n\n``` php\n# a[M,K] × b[M,N] -> w[E,K,N]; contract the ragged M axis, group on the output\nwgrad_dn = lax.RaggedDotDimensionNumbers(\n    dot_dimension_numbers=(((0,), (0,)), ((), ())),\n    lhs_ragged_dimensions=(0,),\n    rhs_group_dimensions=(),        # empty -> grouped output on E\n)\nout = lax.ragged_dot_general(lhs, rhs, group_sizes,   # -> [E, K, N]\n                             ragged_dot_dimension_numbers=wgrad_dn,\n                             preferred_element_type=jnp.bfloat16)\n```\n\n#### Fixing `group_sizes` Deterministically[#](#fixing-group-sizes-deterministically)\n\nIn 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:\n\n```\nbase = M // E                      # tokens per expert (floor)\nrem  = M - base * E                # leftover to distribute (0 <= rem < E)\ngs = np.full((E,), base, np.int32)\ngs[:rem] += 1                      # first `rem` experts get one extra\ngroup_sizes = jnp.asarray(gs)      # int32[E], sums to M\n```\n\nThis 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.\n\n#### Results[#](#results)\n\nFor each model variant we measure the runtime of all the 6 classes on each backend (`default` padded-dot baseline, `triton`, and `hip`).\nThe resulting speedups over the baseline are shown in Figure 26.\n\nThe full per-class TFLOP/s results for every measured shape are provided in\nTable 12 of the [Appendix: Micro-Benchmark Full Results](#appendix-micro-benchmark-full-results).\n\n### Performance of MaxText MoE Models[#](#performance-of-maxtext-moe-models)\n\n#### Performance[#](#performance)\n\nNow, 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.\n\n##### MI300[#](#mi300)\n\n*Table 8: MaxText MoE model training-step performance on MI300.*\n\n| Model | Emb Dim | Experts | Exp/Tok | Layers | total_weights | Backend | execution time (s) | TFLOP/s/device | Tokens/s/device | Speedup | \n|---|---|---|---|---|---|---|---|---|---|---|\n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 3.946 | 7.744 | 519.018 | 1.00x | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 1.867 | 16.369 | 1097.136 | 2.11x | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | hip | 1.761 | 17.348 | 1162.700 | 2.24x | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense_matrix | 1.617 | 18.895 | 1266.397 | 2.44x | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense_matrix_dropless | 2.018 | 15.143 | 1014.919 | 1.96x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 5.758 | 10.910 | 355.653 | 1.00x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 2.732 | 22.995 | 749.610 | 2.11x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | hip | 2.373 | 26.470 | 862.896 | 2.43x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense_matrix | 2.261 | 27.780 | 905.610 | 2.55x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense_matrix_dropless | 2.806 | 22.387 | 729.808 | 2.05x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 0.575 | 141.236 | 7126.279 | 1.00x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 0.436 | 186.151 | 9392.514 | 1.32x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | hip | 0.429 | 189.228 | 9547.763 | 1.34x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense_matrix | 0.414 | 196.070 | 9893.003 | 1.39x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense_matrix_dropless | 0.305 | 265.880 | 13415.345 | 1.89x | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | default | 20.531 | 4.751 | 199.507 | 1.00x | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | triton | 3.257 | 29.950 | 1257.783 | 6.30x | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | hip | 3.073 | 31.737 | 1332.846 | 6.68x | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | dense_matrix | 2.914 | 33.474 | 1405.770 | 7.05x | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | dense_matrix_dropless | 4.163 | 23.430 | 983.947 | 4.93x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 20.676 | 3.855 | 198.105 | 1.00x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 5.013 | 15.897 | 817.004 | 4.12x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | hip | 3.505 | 22.740 | 1168.655 | 5.90x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense_matrix | 3.287 | 24.246 | 1246.054 | 6.29x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense_matrix_dropless | 4.896 | 16.278 | 836.600 | 4.22x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 5.044 | 48.057 | 609.025 | 1.00x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 2.817 | 86.045 | 1090.451 | 1.79x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | hip | 2.128 | 113.903 | 1443.498 | 2.37x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix | 2.015 | 120.291 | 1524.461 | 2.50x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 3.405 | 71.187 | 902.164 | 1.48x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | default | 52.107 | 50.128 | 628.859 | 1.00x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | triton | 30.822 | 84.744 | 1063.120 | 1.69x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | hip | 22.940 | 113.863 | 1428.421 | 2.27x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix | 18.438 | 141.662 | 1777.156 | 2.83x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 34.163 | 76.458 | 959.171 | 1.53x | \n\n##### MI350[#](#mi350)\n\n*Table 9: MaxText MoE model training-step performance on MI350.*\n\n| Model | Emb Dim | Experts | Exp/Tok | Layers | total_weights | Backend | execution time (s) | TFLOP/s/device | Tokens/s/device | Speedup | \n|---|---|---|---|---|---|---|---|---|---|---|\n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 2.581 | 11.840 | 793.547 | 1.00x | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 1.433 | 21.316 | 1428.675 | 1.80x | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense_matrix | 1.303 | 23.451 | 1571.765 | 1.98x | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense_matrix_dropless | 1.414 | 21.603 | 1447.888 | 1.83x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 4.208 | 14.929 | 486.677 | 1.00x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 2.220 | 28.302 | 922.617 | 1.90x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense_matrix | 1.778 | 35.344 | 1152.180 | 2.37x | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense_matrix_dropless | 1.984 | 31.662 | 1032.149 | 2.12x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 0.459 | 176.907 | 8926.120 | 1.00x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 0.374 | 216.782 | 10938.067 | 1.23x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense_matrix | 0.305 | 266.230 | 13433.032 | 1.50x | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense_matrix_dropless | 0.188 | 432.514 | 21823.103 | 2.44x | \n| llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | default | 10.249 | 39.299 | 399.631 | 1.00x | \n| llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | triton | 10.157 | 39.655 | 403.258 | 1.01x | \n| llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | dense_matrix | 9.337 | 43.138 | 438.671 | 1.10x | \n| llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | dense_matrix_dropless | 9.344 | 43.108 | 438.372 | 1.10x | \n| qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | default | 22.025 | 14.110 | 1487.731 | 1.00x | \n| qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | triton | 4.082 | 76.132 | 8027.089 | 5.40x | \n| qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | dense_matrix | 3.623 | 85.785 | 9044.923 | 6.08x | \n| qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | dense_matrix_dropless | 5.322 | 58.391 | 6156.577 | 4.14x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 9.696 | 8.220 | 422.434 | 1.00x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 4.012 | 19.863 | 1020.834 | 2.42x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense_matrix | 2.557 | 31.166 | 1601.701 | 3.79x | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense_matrix_dropless | 3.294 | 24.197 | 1243.551 | 2.94x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 2.863 | 84.681 | 1073.164 | 1.00x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 1.917 | 126.458 | 1602.606 | 1.49x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix | 1.432 | 169.318 | 2145.780 | 2.00x | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 2.075 | 116.814 | 1480.395 | 1.38x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | default | 29.188 | 89.489 | 1122.641 | 1.00x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | triton | 21.327 | 122.475 | 1536.450 | 1.37x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix | 13.376 | 195.272 | 2449.693 | 2.18x | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 25.522 | 102.344 | 1283.910 | 1.14x | \n\nThe performance numbers confirm that the new Group-GEMM based backends for `jax.lax.ragged_dot` operations provide a significant speedup for MoE models.\nAs expected, the `hipBLASLt` backend is faster than the Triton one; however, both offer very interesting speedups.\nThat 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:\n\nthe capacity is relatively low (1.25), which means that the padding is relatively small, but potentially a significant number of tokens could be dropped;\n\nthe 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.\n\n#### Memory Savings[#](#memory-savings)\n\nThe memory-usage figures below were extracted directly from the log traces\nemitted by MaxText during each training run. MaxText reports, per run, the\npeak device memory footprint (broken down into total and temporary/scratch\nmemory), which we collected for each backend (`default`, `triton`, `hip`) and\ndiffed against the `default` padded-dot baseline to obtain the savings shown.\nThey are reported in Figure 29 and Table 10 for MI300, and in Figure 30 and Table 11 for MI350.\n\n##### MI300[#](#id1)\n\n*Table 10: MaxText MoE model peak memory footprint on MI300.*\n\n| Model | Emb Dim | Experts | Exp/Tok | Layers | total_weights | Backend | Total Mem (GB) | Temp Mem (GB) | Mem Savings (GB) | Temp Savings (GB) | \n|---|---|---|---|---|---|---|---|---|---|---|\n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 26.7 | 15.7 | 0.0 (0.0%) | 0.0 (0.0%) | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 23.6 | 12.6 | **3.1 (11.6%)** | **3.1 (19.7%)** | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | hip | 23.6 | 12.6 | **3.1 (11.6%)** | **3.1 (19.7%)** | \n| 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%)** | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense_matrix_dropless | 23.9 | 13.0 | **2.8 (10.5%)** | **2.7 (17.2%)** | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 26.3 | 15.2 | 0.0 (0.0%) | 0.0 (0.0%) | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 23.1 | 12.0 | **3.2 (12.2%)** | **3.2 (21.1%)** | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | hip | 23.1 | 12.0 | **3.2 (12.2%)** | **3.2 (21.1%)** | \n| 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%)** | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense_matrix_dropless | 23.6 | 12.6 | **2.7 (10.3%)** | **2.6 (17.1%)** | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 8.6 | 3.4 | 0.0 (0.0%) | 0.0 (0.0%) | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 8.6 | 3.4 | 0.0 (0.0%) | 0.0 (0.0%) | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | hip | 8.6 | 3.4 | 0.0 (0.0%) | 0.0 (0.0%) | \n| 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%) | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense_matrix_dropless | 8.5 | 3.4 | **0.1 (1.2%)** | 0.0 (0.0%) | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | default | 61.0 | 43.3 | 0.0 (0.0%) | 0.0 (0.0%) | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | triton | 43.6 | 26.0 | **17.4 (28.5%)** | **17.3 (40.0%)** | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | hip | 43.6 | 26.0 | **17.4 (28.5%)** | **17.3 (40.0%)** | \n| 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%)** | \n| gemma4-26b | 2816 | 128 | 8 | 30 | 25.233B | dense_matrix_dropless | 41.6 | 23.9 | **19.4 (31.8%)** | **19.4 (44.8%)** | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 57.7 | 36.4 | 0.0 (0.0%) | 0.0 (0.0%) | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 44.5 | 23.2 | **13.2 (22.9%)** | **13.2 (36.3%)** | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | hip | 44.5 | 23.2 | **13.2 (22.9%)** | **13.2 (36.3%)** | \n| 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%)** | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense_matrix_dropless | 45.8 | 24.5 | **11.9 (20.6%)** | **11.9 (32.7%)** | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 120.2 | 54.9 | 0.0 (0.0%) | 0.0 (0.0%) | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 114.4 | 49.1 | **5.8 (4.8%)** | **5.8 (10.6%)** | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | hip | 114.7 | 49.5 | **5.5 (4.6%)** | **5.4 (9.8%)** | \n| 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%)** | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 117.7 | 52.5 | **2.5 (2.1%)** | **2.4 (4.4%)** | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | default | 79.9 | 47.2 | 0.0 (0.0%) | 0.0 (0.0%) | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | triton | 72.7 | 40.1 | **7.2 (9.0%)** | **7.1 (15.0%)** | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | hip | 72.9 | 40.2 | **7.0 (8.8%)** | **7.0 (14.8%)** | \n| 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%)** | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 77.0 | 44.4 | **2.9 (3.6%)** | **2.8 (5.9%)** | \n\n##### MI350[#](#id2)\n\n*Table 11: MaxText MoE model peak memory footprint on MI350.*\n\n| Model | Emb Dim | Experts | Exp/Tok | Layers | total_weights | Backend | Total Mem (GB) | Temp Mem (GB) | Mem Savings (GB) | Temp Savings (GB) | \n|---|---|---|---|---|---|---|---|---|---|---|\n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | default | 26.7 | 15.7 | 0.0 (0.0%) | 0.0 (0.0%) | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | triton | 23.7 | 12.7 | **3.0 (11.2%)** | **3.0 (19.1%)** | \n| 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%)** | \n| deepseek2-16b | 2048 | 64 | 6 | 27 | 15.706B | dense_matrix_dropless | 23.9 | 13.0 | **2.8 (10.5%)** | **2.7 (17.2%)** | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | default | 26.3 | 15.2 | 0.0 (0.0%) | 0.0 (0.0%) | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | triton | 23.2 | 12.2 | **3.1 (11.8%)** | **3.0 (19.7%)** | \n| 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%)** | \n| deepseek3-test | 512 | 256 | 8 | 61 | 15.851B | dense_matrix_dropless | 23.8 | 12.8 | **2.5 (9.5%)** | **2.4 (15.8%)** | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | default | 8.6 | 3.4 | 0.0 (0.0%) | 0.0 (0.0%) | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | triton | 8.6 | 3.4 | 0.0 (0.0%) | 0.0 (0.0%) | \n| 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%) | \n| deepseek3-671b | 7168 | 32 | 8 | 4 | 7.415B | dense_matrix_dropless | 9.4 | 4.2 | -0.8 (-9.3%) | -0.8 (-23.5%) | \n| llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | default | 170.5 | 95.2 | 0.0 (0.0%) | 0.0 (0.0%) | \n| llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | triton | 170.0 | 94.7 | **0.5 (0.3%)** | **0.5 (0.5%)** | \n| 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%) | \n| llama4-17b-16e | 5120 | 16 | 1 | 48 | 107.770B | dense_matrix_dropless | 171.6 | 96.3 | -1.1 (-0.6%) | -1.1 (-1.2%) | \n| qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | default | 74.7 | 66.9 | 0.0 (0.0%) | 0.0 (0.0%) | \n| qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | triton | 24.4 | 16.5 | **50.3 (67.3%)** | **50.4 (75.3%)** | \n| 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%)** | \n| qwen3-235b-a22b | 4096 | 128 | 8 | 4 | 11.196B | dense_matrix_dropless | 28.8 | 21.0 | **45.9 (61.4%)** | **45.9 (68.6%)** | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | default | 57.7 | 36.4 | 0.0 (0.0%) | 0.0 (0.0%) | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | triton | 45.1 | 23.7 | **12.6 (21.8%)** | **12.7 (34.9%)** | \n| 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%)** | \n| qwen3-30b-a3b | 2048 | 128 | 8 | 48 | 30.532B | dense_matrix_dropless | 46.2 | 24.9 | **11.5 (19.9%)** | **11.5 (31.6%)** | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | default | 120.2 | 54.9 | 0.0 (0.0%) | 0.0 (0.0%) | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | triton | 115.2 | 50.0 | **5.0 (4.2%)** | **4.9 (8.9%)** | \n| 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%)** | \n| mixtral-8x7b | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 117.9 | 52.7 | **2.3 (1.9%)** | **2.2 (4.0%)** | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | default | 79.9 | 47.2 | 0.0 (0.0%) | 0.0 (0.0%) | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | triton | 73.7 | 41.1 | **6.2 (7.8%)** | **6.1 (12.9%)** | \n| 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%)** | \n| mixtral-8x7b (large) | 4096 | 8 | 2 | 32 | 46.703B | dense_matrix_dropless | 77.3 | 44.7 | **2.6 (3.3%)** | **2.5 (5.3%)** | \n\nRegarding memory usage, the measured memory footprints confirm that the\nGroup-GEMM backends make better use of memory for `ragged-dot` operations than the default XLA backend by avoiding the storage of unnecessary intermediate data.\nIndeed, 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\n`E × [M, N]` stack of results (and the matching `E × [M, ·]` activations) and then\nmasks each expert’s output down to the rows that actually belong to it. So every\nexpert redundantly processes all `M` tokens, and the intermediate/output tensors\nblow up by a factor of `E`. Thus, the baseline’s wasted memory scales roughly as\n`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.\n\nThe measured numbers line up with this:\n\n- **`qwen3-235b-a22b` (MI350)** —**128 experts** and a large routed-token\nworkload (94 layers, seq 2048, top-8). Its baseline temp memory is**66.9 GB** ,\nof which Group-GEMM removes**50.4 GB (75.3%)** — by far the biggest saving in the\ntables. With 128 experts each redundantly processing all`M` tokens, the`E × [M, N]` stack is enormous, and Group-GEMM never builds it.\n- **`gemma4-26b` (MI300)** —**128 experts** and many routed tokens (top-8),\neven though the expert width is small (`N = 704` ):**17.3 GB (40.0%)** of temp\nmemory saved. The large`E` (×`M` ) dominates the blow-up.\n- **`qwen3-30b-a3b` (MI300 & MI350)** —**128 experts** ,`N = 768` , top-8:**≈13 GB (~35%)** temp saved on both targets — the same many-experts effect.\n- **`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\nredundant stack is comparatively cheap and there is little to reclaim.\n- **`deepseek3-671b-reduced`** —**0 GB (0%)** : expert-parallel with only a few\nexperts per GPU and very few tokens per step, so`E × M` is tiny and there is\nessentially nothing to remove.\n- **`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.\n\nNow, 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:\n\n- 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.\n- in our experiment, the drop capacity is quite low (1.25), which does not incur significant data padding.\n\n## Part 5 — Conclusion and Next Steps[#](#part-5-conclusion-and-next-steps)\n\nUsing Group-GEMM as the backend for XLA `ragged-dot` brings two clear advantages:\nit is **faster** (no padding FLOPs) and it uses **less memory** (no\nexpansion/padding). And XLA now offers two ways to get there:\n\n- **hipBLASLt** — a vendor-library backend on AMD ROCm, extending XLA’s existing hipBLASLt matmul path\n- **Triton** — a portable, autotunable, code-generated backend that runs on both CUDA and ROCm, currently experimental and off by default.\n\n**Next steps:**\n\n- **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.\n- **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.\n\n## Summary[#](#summary)\n\nIn 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.\n\nWe built that argument in five parts, each with a distinct role.\n[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.\n[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).\n[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.\n[Part 5](#part-5-conclusion-and-next-steps) concludes on the Group-GEMM support and proposes a roadmap for future improvements.\n\nAlong the way, the key takeaways were:\n\n- **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).\n- **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.\n\nBy 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.\n\n## Appendix: Micro-Benchmark Full Results[#](#appendix-micro-benchmark-full-results)\n\nTable 12 below reports the full per-class TFLOP/s results for every measured\nshape in the micro-benchmark of [Micro-Benchmark](#micro-benchmark),\nacross the three backends (`default` padded-dot baseline, `triton`, and `hipblas`).\n\n*Table 12: Full per-class micro-benchmark results (TFLOP/s) on MI300.*\n\n| class | M | K | N | E | backend | tflops | status | \n|---|---|---|---|---|---|---|---|\n| fwd_wi | 3840 | 2048 | 1408 | 64 | default | 6.85 | ok | \n| fwd_wi | 3840 | 2048 | 1408 | 64 | triton | 86.44 | ok | \n| fwd_wi | 3840 | 2048 | 1408 | 64 | hipblas | 135.21 | ok | \n| fwd_wo | 3840 | 2048 | 1408 | 64 | default | 9.04 | ok | \n| fwd_wo | 3840 | 2048 | 1408 | 64 | triton | 94.14 | ok | \n| fwd_wo | 3840 | 2048 | 1408 | 64 | hipblas | 120.85 | ok | \n| dgrad_wi | 3840 | 2048 | 1408 | 64 | default | 8.50 | ok | \n| dgrad_wi | 3840 | 2048 | 1408 | 64 | triton | 87.50 | ok | \n| dgrad_wi | 3840 | 2048 | 1408 | 64 | hipblas | 159.97 | ok | \n| wgrad_wi | 3840 | 2048 | 1408 | 64 | default | 8.79 | ok | \n| wgrad_wi | 3840 | 2048 | 1408 | 64 | triton | 94.40 | ok | \n| wgrad_wi | 3840 | 2048 | 1408 | 64 | hipblas | 102.22 | ok | \n| dgrad_wo | 3840 | 2048 | 1408 | 64 | default | 6.49 | ok | \n| dgrad_wo | 3840 | 2048 | 1408 | 64 | triton | 82.86 | ok | \n| dgrad_wo | 3840 | 2048 | 1408 | 64 | hipblas | 119.83 | ok | \n| wgrad_wo | 3840 | 2048 | 1408 | 64 | default | 8.80 | ok | \n| wgrad_wo | 3840 | 2048 | 1408 | 64 | triton | 81.52 | ok | \n| wgrad_wo | 3840 | 2048 | 1408 | 64 | hipblas | 104.17 | ok | \n| fwd_wi | 1920 | 5120 | 1536 | 160 | default | 3.43 | ok | \n| fwd_wi | 1920 | 5120 | 1536 | 160 | triton | 31.05 | ok | \n| fwd_wi | 1920 | 5120 | 1536 | 160 | hipblas | 44.31 | ok | \n| fwd_wo | 1920 | 5120 | 1536 | 160 | default | 4.13 | ok | \n| fwd_wo | 1920 | 5120 | 1536 | 160 | triton | 36.06 | ok | \n| fwd_wo | 1920 | 5120 | 1536 | 160 | hipblas | 33.43 | ok | \n| dgrad_wi | 1920 | 5120 | 1536 | 160 | default | 3.28 | ok | \n| dgrad_wi | 1920 | 5120 | 1536 | 160 | triton | 34.16 | ok | \n| dgrad_wi | 1920 | 5120 | 1536 | 160 | hipblas | 38.35 | ok | \n| wgrad_wi | 1920 | 5120 | 1536 | 160 | default | 3.72 | ok | \n| wgrad_wi | 1920 | 5120 | 1536 | 160 | triton | 4.50 | ok | \n| wgrad_wi | 1920 | 5120 | 1536 | 160 | hipblas | 26.74 | ok | \n| dgrad_wo | 1920 | 5120 | 1536 | 160 | default | 2.95 | ok | \n| dgrad_wo | 1920 | 5120 | 1536 | 160 | triton | 30.68 | ok | \n| dgrad_wo | 1920 | 5120 | 1536 | 160 | hipblas | 45.01 | ok | \n| wgrad_wo | 1920 | 5120 | 1536 | 160 | default | 3.68 | ok | \n| wgrad_wo | 1920 | 5120 | 1536 | 160 | triton | 4.50 | ok | \n| wgrad_wo | 1920 | 5120 | 1536 | 160 | hipblas | 26.80 | ok | \n| fwd_wi | 5120 | 512 | 512 | 256 | default | 2.02 | ok | \n| fwd_wi | 5120 | 512 | 512 | 256 | triton | 32.98 | ok | \n| fwd_wi | 5120 | 512 | 512 | 256 | hipblas | 24.19 | ok | \n| fwd_wo | 5120 | 512 | 512 | 256 | default | 2.02 | ok | \n| fwd_wo | 5120 | 512 | 512 | 256 | triton | 33.16 | ok | \n| fwd_wo | 5120 | 512 | 512 | 256 | hipblas | 26.19 | ok | \n| dgrad_wi | 5120 | 512 | 512 | 256 | default | 1.92 | ok | \n| dgrad_wi | 5120 | 512 | 512 | 256 | triton | 34.07 | ok | \n| dgrad_wi | 5120 | 512 | 512 | 256 | hipblas | 27.99 | ok | \n| wgrad_wi | 5120 | 512 | 512 | 256 | default | 1.52 | ok | \n| wgrad_wi | 5120 | 512 | 512 | 256 | triton | 26.85 | ok | \n| wgrad_wi | 5120 | 512 | 512 | 256 | hipblas | 24.98 | ok | \n| dgrad_wo | 5120 | 512 | 512 | 256 | default | 1.92 | ok | \n| dgrad_wo | 5120 | 512 | 512 | 256 | triton | 33.37 | ok | \n| dgrad_wo | 5120 | 512 | 512 | 256 | hipblas | 26.36 | ok | \n| wgrad_wo | 5120 | 512 | 512 | 256 | default | 1.53 | ok | \n| wgrad_wo | 5120 | 512 | 512 | 256 | triton | 26.85 | ok | \n| wgrad_wo | 5120 | 512 | 512 | 256 | hipblas | 24.53 | ok | \n| fwd_wi | 10240 | 2816 | 704 | 128 | default | 2.14 | ok | \n| fwd_wi | 10240 | 2816 | 704 | 128 | triton | 110.97 | ok | \n| fwd_wi | 10240 | 2816 | 704 | 128 | hipblas | 172.24 | ok | \n| fwd_wo | 10240 | 2816 | 704 | 128 | default | 5.66 | ok | \n| fwd_wo | 10240 | 2816 | 704 | 128 | triton | 128.01 | ok | \n| fwd_wo | 10240 | 2816 | 704 | 128 | hipblas | 137.69 | ok | \n| dgrad_wi | 10240 | 2816 | 704 | 128 | default | 5.47 | ok | \n| dgrad_wi | 10240 | 2816 | 704 | 128 | triton | 128.28 | ok | \n| dgrad_wi | 10240 | 2816 | 704 | 128 | hipblas | 175.65 | ok | \n| wgrad_wi | 10240 | 2816 | 704 | 128 | default | 3.62 | ok | \n| wgrad_wi | 10240 | 2816 | 704 | 128 | triton | 116.71 | ok | \n| wgrad_wi | 10240 | 2816 | 704 | 128 | hipblas | 104.51 | ok | \n| dgrad_wo | 10240 | 2816 | 704 | 128 | default | 2.19 | ok | \n| dgrad_wo | 10240 | 2816 | 704 | 128 | triton | 105.61 | ok | \n| dgrad_wo | 10240 | 2816 | 704 | 128 | hipblas | 176.18 | ok | \n| wgrad_wo | 10240 | 2816 | 704 | 128 | default | 3.67 | ok | \n| wgrad_wo | 10240 | 2816 | 704 | 128 | triton | 96.62 | ok | \n| wgrad_wo | 10240 | 2816 | 704 | 128 | hipblas | 113.53 | ok | \n| fwd_wi | 10240 | 2880 | 2880 | 128 | default | 4.46 | ok | \n| fwd_wi | 10240 | 2880 | 2880 | 128 | triton | 128.32 | ok | \n| fwd_wi | 10240 | 2880 | 2880 | 128 | hipblas | 200.81 | ok | \n| fwd_wo | 10240 | 2880 | 2880 | 128 | default | 4.46 | ok | \n| fwd_wo | 10240 | 2880 | 2880 | 128 | triton | 127.71 | ok | \n| fwd_wo | 10240 | 2880 | 2880 | 128 | hipblas | 199.25 | ok | \n| dgrad_wi | 10240 | 2880 | 2880 | 128 | default | 4.33 | ok | \n| dgrad_wi | 10240 | 2880 | 2880 | 128 | triton | 126.28 | ok | \n| dgrad_wi | 10240 | 2880 | 2880 | 128 | hipblas | 242.96 | ok | \n| wgrad_wi | 10240 | 2880 | 2880 | 128 | default | 4.90 | ok | \n| wgrad_wi | 10240 | 2880 | 2880 | 128 | triton | 130.05 | ok | \n| wgrad_wi | 10240 | 2880 | 2880 | 128 | hipblas | 127.73 | ok | \n| dgrad_wo | 10240 | 2880 | 2880 | 128 | default | 4.32 | ok | \n| dgrad_wo | 10240 | 2880 | 2880 | 128 | triton | 125.88 | ok | \n| dgrad_wo | 10240 | 2880 | 2880 | 128 | hipblas | 242.90 | ok | \n| wgrad_wo | 10240 | 2880 | 2880 | 128 | default | 4.91 | ok | \n| wgrad_wo | 10240 | 2880 | 2880 | 128 | triton | 129.74 | ok | \n| wgrad_wo | 10240 | 2880 | 2880 | 128 | hipblas | 128.08 | ok | \n| fwd_wi | 1280 | 5120 | 8192 | 16 | default | 44.77 | ok | \n| fwd_wi | 1280 | 5120 | 8192 | 16 | triton | 133.88 | ok | \n| fwd_wi | 1280 | 5120 | 8192 | 16 | hipblas | 220.41 | ok | \n| fwd_wo | 1280 | 5120 | 8192 | 16 | default | 35.27 | ok | \n| fwd_wo | 1280 | 5120 | 8192 | 16 | triton | 129.44 | ok | \n| fwd_wo | 1280 | 5120 | 8192 | 16 | hipblas | 224.86 | ok | \n| dgrad_wi | 1280 | 5120 | 8192 | 16 | default | 27.50 | ok | \n| dgrad_wi | 1280 | 5120 | 8192 | 16 | triton | 96.85 | ok | \n| dgrad_wi | 1280 | 5120 | 8192 | 16 | hipblas | 246.03 | ok | \n| wgrad_wi | 1280 | 5120 | 8192 | 16 | default | 40.48 | ok | \n| wgrad_wi | 1280 | 5120 | 8192 | 16 | triton | 113.19 | ok | \n| wgrad_wi | 1280 | 5120 | 8192 | 16 | hipblas | 154.65 | ok | \n| dgrad_wo | 1280 | 5120 | 8192 | 16 | default | 32.66 | ok | \n| dgrad_wo | 1280 | 5120 | 8192 | 16 | triton | 130.98 | ok | \n| dgrad_wo | 1280 | 5120 | 8192 | 16 | hipblas | 232.46 | ok | \n| wgrad_wo | 1280 | 5120 | 8192 | 16 | default | 41.02 | ok | \n| wgrad_wo | 1280 | 5120 | 8192 | 16 | triton | 117.14 | ok | \n| wgrad_wo | 1280 | 5120 | 8192 | 16 | hipblas | 154.67 | ok | \n| fwd_wi | 15360 | 6144 | 20480 | 32 | default | 24.73 | ok | \n| fwd_wi | 15360 | 6144 | 20480 | 32 | triton | 261.23 | ok | \n| fwd_wi | 15360 | 6144 | 20480 | 32 | hipblas | 527.94 | ok | \n| fwd_wo | 15360 | 6144 | 20480 | 32 | default | 17.97 | ok | \n| fwd_wo | 15360 | 6144 | 20480 | 32 | triton | 233.12 | ok | \n| fwd_wo | 15360 | 6144 | 20480 | 32 | hipblas | 471.04 | ok | \n| dgrad_wi | 15360 | 6144 | 20480 | 32 | default | 17.71 | ok | \n| dgrad_wi | 15360 | 6144 | 20480 | 32 | triton | 244.21 | ok | \n| dgrad_wi | 15360 | 6144 | 20480 | 32 | hipblas | 485.43 | ok | \n| wgrad_wi | 15360 | 6144 | 20480 | 32 | default | 24.57 | ok | \n| wgrad_wi | 15360 | 6144 | 20480 | 32 | triton | 284.41 | ok | \n| wgrad_wi | 15360 | 6144 | 20480 | 32 | hipblas | 428.84 | ok | \n| dgrad_wo | 15360 | 6144 | 20480 | 32 | default | 24.01 | ok | \n| dgrad_wo | 15360 | 6144 | 20480 | 32 | triton | 266.45 | ok | \n| dgrad_wo | 15360 | 6144 | 20480 | 32 | hipblas | 470.12 | ok | \n| wgrad_wo | 15360 | 6144 | 20480 | 32 | default | 24.62 | ok | \n| wgrad_wo | 15360 | 6144 | 20480 | 32 | triton | 283.80 | ok | \n| wgrad_wo | 15360 | 6144 | 20480 | 32 | hipblas | 431.66 | ok | \n| fwd_wi | 40960 | 8192 | 28672 | 64 | default | 8.80 | ok | \n| fwd_wi | 40960 | 8192 | 28672 | 64 | triton | 294.11 | ok | \n| fwd_wi | 40960 | 8192 | 28672 | 64 | hipblas | 510.97 | ok | \n| fwd_wo | 40960 | 8192 | 28672 | 64 | default |  | OOM | \n| fwd_wo | 40960 | 8192 | 28672 | 64 | triton | 298.32 | ok | \n| fwd_wo | 40960 | 8192 | 28672 | 64 | hipblas | 507.20 | ok | \n| dgrad_wi | 40960 | 8192 | 28672 | 64 | default |  | OOM | \n| dgrad_wi | 40960 | 8192 | 28672 | 64 | triton | 299.03 | ok | \n| dgrad_wi | 40960 | 8192 | 28672 | 64 | hipblas | 464.65 | ok | \n| wgrad_wi | 40960 | 8192 | 28672 | 64 | default |  | OOM | \n| wgrad_wi | 40960 | 8192 | 28672 | 64 | triton | 78.34 | ok | \n| wgrad_wi | 40960 | 8192 | 28672 | 64 | hipblas | 318.52 | ok | \n| dgrad_wo | 40960 | 8192 | 28672 | 64 | default | 8.74 | ok | \n| dgrad_wo | 40960 | 8192 | 28672 | 64 | triton | 308.75 | ok | \n| dgrad_wo | 40960 | 8192 | 28672 | 64 | hipblas | 497.37 | ok | \n| wgrad_wo | 40960 | 8192 | 28672 | 64 | default |  | OOM | \n| wgrad_wo | 40960 | 8192 | 28672 | 64 | triton | 51.84 | ok | \n| wgrad_wo | 40960 | 8192 | 28672 | 64 | hipblas | 336.43 | ok | \n| fwd_wi | 20480 | 4096 | 1536 | 128 | default | 3.82 | ok | \n| fwd_wi | 20480 | 4096 | 1536 | 128 | triton | 168.07 | ok | \n| fwd_wi | 20480 | 4096 | 1536 | 128 | hipblas | 332.96 | ok | \n| fwd_wo | 20480 | 4096 | 1536 | 128 | default | 5.71 | ok | \n| fwd_wo | 20480 | 4096 | 1536 | 128 | triton | 176.82 | ok | \n| fwd_wo | 20480 | 4096 | 1536 | 128 | hipblas | 323.89 | ok | \n| dgrad_wi | 20480 | 4096 | 1536 | 128 | default | 5.61 | ok | \n| dgrad_wi | 20480 | 4096 | 1536 | 128 | triton | 190.20 | ok | \n| dgrad_wi | 20480 | 4096 | 1536 | 128 | hipblas | 291.06 | ok | \n| wgrad_wi | 20480 | 4096 | 1536 | 128 | default | 4.82 | ok | \n| wgrad_wi | 20480 | 4096 | 1536 | 128 | triton | 175.78 | ok | \n| wgrad_wi | 20480 | 4096 | 1536 | 128 | hipblas | 236.19 | ok | \n| dgrad_wo | 20480 | 4096 | 1536 | 128 | default | 4.43 | ok | \n| dgrad_wo | 20480 | 4096 | 1536 | 128 | triton | 172.91 | ok | \n| dgrad_wo | 20480 | 4096 | 1536 | 128 | hipblas | 324.73 | ok | \n| wgrad_wo | 20480 | 4096 | 1536 | 128 | default | 4.85 | ok | \n| wgrad_wo | 20480 | 4096 | 1536 | 128 | triton | 176.91 | ok | \n| wgrad_wo | 20480 | 4096 | 1536 | 128 | hipblas | 232.07 | ok | \n| fwd_wi | 10240 | 2048 | 768 | 128 | default | 3.33 | ok | \n| fwd_wi | 10240 | 2048 | 768 | 128 | triton | 111.85 | ok | \n| fwd_wi | 10240 | 2048 | 768 | 128 | hipblas | 171.38 | ok | \n| fwd_wo | 10240 | 2048 | 768 | 128 | default | 5.59 | ok | \n| fwd_wo | 10240 | 2048 | 768 | 128 | triton | 116.12 | ok | \n| fwd_wo | 10240 | 2048 | 768 | 128 | hipblas | 135.95 | ok | \n| dgrad_wi | 10240 | 2048 | 768 | 128 | default | 5.39 | ok | \n| dgrad_wi | 10240 | 2048 | 768 | 128 | triton | 115.59 | ok | \n| dgrad_wi | 10240 | 2048 | 768 | 128 | hipblas | 144.88 | ok | \n| wgrad_wi | 10240 | 2048 | 768 | 128 | default | 3.87 | ok | \n| wgrad_wi | 10240 | 2048 | 768 | 128 | triton | 99.24 | ok | \n| wgrad_wi | 10240 | 2048 | 768 | 128 | hipblas | 121.36 | ok | \n| dgrad_wo | 10240 | 2048 | 768 | 128 | default | 3.25 | ok | \n| dgrad_wo | 10240 | 2048 | 768 | 128 | triton | 106.70 | ok | \n| dgrad_wo | 10240 | 2048 | 768 | 128 | hipblas | 171.31 | ok | \n| wgrad_wo | 10240 | 2048 | 768 | 128 | default | 3.87 | ok | \n| wgrad_wo | 10240 | 2048 | 768 | 128 | triton | 97.99 | ok | \n| wgrad_wo | 10240 | 2048 | 768 | 128 | hipblas | 122.84 | ok | \n| fwd_wi | 6400 | 2048 | 512 | 512 | default | 1.11 | ok | \n| fwd_wi | 6400 | 2048 | 512 | 512 | triton | 29.18 | ok | \n| fwd_wi | 6400 | 2048 | 512 | 512 | hipblas | 25.16 | ok | \n| fwd_wo | 6400 | 2048 | 512 | 512 | default | 1.34 | ok | \n| fwd_wo | 6400 | 2048 | 512 | 512 | triton | 30.86 | ok | \n| fwd_wo | 6400 | 2048 | 512 | 512 | hipblas | 16.43 | ok | \n| dgrad_wi | 6400 | 2048 | 512 | 512 | default | 1.27 | ok | \n| dgrad_wi | 6400 | 2048 | 512 | 512 | triton | 27.83 | ok | \n| dgrad_wi | 6400 | 2048 | 512 | 512 | hipblas | 17.26 | ok | \n| wgrad_wi | 6400 | 2048 | 512 | 512 | default | 0.88 | ok | \n| wgrad_wi | 6400 | 2048 | 512 | 512 | triton | 1.82 | ok | \n| wgrad_wi | 6400 | 2048 | 512 | 512 | hipblas | 13.86 | ok | \n| dgrad_wo | 6400 | 2048 | 512 | 512 | default | 1.06 | ok | \n| dgrad_wo | 6400 | 2048 | 512 | 512 | triton | 27.49 | ok | \n| dgrad_wo | 6400 | 2048 | 512 | 512 | hipblas | 30.17 | ok | \n| wgrad_wo | 6400 | 2048 | 512 | 512 | default | 0.87 | ok | \n| wgrad_wo | 6400 | 2048 | 512 | 512 | triton | 1.82 | ok | \n| wgrad_wo | 6400 | 2048 | 512 | 512 | hipblas | 13.80 | ok | \n\n## Disclaimers[#](#disclaimers)\n\nThe 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", "url": "https://wpnews.pro/news/group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient", "canonical_source": "https://rocm.blogs.amd.com/software-tools-optimization/openxla-group-gemm/README.html", "published_at": "2026-10-26 00:00:00+00:00", "updated_at": "2026-10-26 16:15:50.801476+00:00", "lang": "en", "topics": ["machine-learning", "ai-infrastructure", "large-language-models", "ai-tools"], "entities": ["XLA", "JAX", "hipBLASLt", "AMD ROCm", "Triton", "CUDA", "MI300", "MaxText"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient", "markdown": "https://wpnews.pro/news/group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient.md", "text": "https://wpnews.pro/news/group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient.txt", "jsonld": "https://wpnews.pro/news/group-gemm-in-xla-making-mixture-of-experts-fast-and-memory-efficient.jsonld"}}