Qwen MoE BringUp

[jtuyls]

Model: Qwen1.5-MoE-A2.7B-Chat — INT4 — 24 layers, 60 experts/layer, top-4 routing, 16 KV heads, head_dim=128

Local fixes to get this to work:

A1. DequantizeLinear block_size decomposition: llvm/torch-mlir#4505

Why: torch-mlir doesn't support block_size attr, so INT4 QDQ models fail.
What: Decompose to extui→uitofp→subf→mulf in the EP. Standalone.
Alternatives:

• torch-mlir (proper fix): Add block_size support to onnx.DequantizeLinear → torch lowering. This is where it belongs — it's a missing ONNX opset feature.

A2. ReduceMean/ReduceSum opset 18+: llvm/torch-mlir#4506

Why: Opset 18 moved axes from attribute to tensor input. torch-mlir expects attribute form.
What: Convert constant tensor input back to attribute in the EP. Standalone.
Alternatives:

• torch-mlir (proper fix): Handle axes-as-tensor in ONNX→torch lowering. Same story — opset version gap in torch-mlir.

A3. Shape refinement + op rewrites

Why: ORT's per-node TypeInfo loses shape precision, producing all-? dims → slow dynamic codegen.
What: Track refined types, rewrite ops to torch dialect with traceable shapes. (~900 lines)
Alternatives:

• ORT (proper fix for part of it): GetOutputTypeInfo should respect add_free_dimension_override_by_name. The EP shouldn't need to re-derive static shapes that ORT already knows.
• torch-mlir: The op rewrites (Reshape→aten.reshape, Unsqueeze→aten.unsqueeze, etc.) are essentially doing ONNX→torch lowering that torch-mlir already does. The EP emits torch.operator "onnx.Reshape" and lets torch-mlir lower it — but the shape info is lost by then. If ORT provided correct types, the generic path would work fine.
• IREE: Could be more aggressive about shape inference on onnx.* ops before lowering, but that's pushing ONNX knowledge into the wrong layer.
• The real root cause is split: ORT gives wrong types + the EP uses a passthrough onnx.* emission that loses shape context torch-mlir could recover.

A4. Expand output type fix

Why: ORT can't infer KV cache broadcast shape through Where(Equal(Concat, -1), ones, Concat) pattern. onnx.Expand output stays [1,1,?,1] instead of [1,16,?,128]. ScatterElements then writes only 1 element per position → nearly all-zero KV cache.
What: Extend ResolveShapeVector to handle Reshape/Where chains; rewrite Expand refinement in PropagateRefinedTypes to resolve shape tensor and compute broadcast output. Depends on A3 infrastructure.
Alternatives:

• ORT (proper fix): Shape inference should trace through Where(Equal(...), ...) patterns to resolve Expand output shapes. The EP is compensating for ORT's shape inference limitations.
• torch-mlir: If the Expand op arrived with correct output type annotation, torch-mlir would lower it correctly. This is an ORT type propagation issue, not a lowering issue.

A5. ScatterElements → tensor.insert_slice for KV cache

Why: ONNX ScatterElements with sequential indices decomposes to hundreds of memcpy dispatches in IREE. The KV cache update pattern (ScatterElements(past_kv, Expand(Reshape(Cast(Add(seq_pos, ...)))), updates)) is a contiguous slice write.
What: Pattern-match the KV cache ScatterElements and emit tensor.insert_slice instead. Depends on A3 (uses refined types for shape resolution).
Alternatives:

• IREE: Recognize contiguous scatter patterns and lower to insert_slice internally. The EP is doing a pattern match that the compiler should do.
• torch-mlir: Better onnx.ScatterElements → torch lowering that preserves the contiguous-slice structure.

A6. Static shape propagation + Cast roundtrip elimination

Why: ORT's node-level TypeInfo returns original model shapes (with dynamic ? dims) even after add_free_dimension_override_by_name("seq_len", 1). So graph inputs have static types in the function signature but ? in every op annotation downstream. Cast(f16→f32)→Cast(f32→f16) roundtrips after Softmax add 24 unnecessary dispatches.
What: Seed refined_types_ from graph input signatures so static shapes propagate everywhere. Add Gather shape refinement and Reshape -1 resolution. Detect Cast roundtrips and alias through them for DCE. Depends on A3.
Alternatives:

• ORT (proper fix for shape propagation): Node-level GetOutputTypeInfo should incorporate free dimension overrides. This is the root cause — the EP is re-deriving information ORT already has in the graph-level type but doesn't expose per-node.
• torch-mlir / IREE (Cast elimination): Cast(A→B)→Cast(B→A) roundtrip canonicalization is a standard compiler optimization. Both torch-mlir and IREE should already handle this. The EP shouldn't need to do DCE prep.
• The Gather refinement and Reshape -1 resolution are shape inference that torch-mlir could do if it received correct input types — again falls out of the ORT type info bug.

A7. MatMul f16→f32 accumulation promotion

Why: IREE selects WMMAR3_F16_16x16x16_F16 (f16 accumulator) for f16 matmuls on gfx1100, which produces NaN for transposed-B operands (the QxK^T pattern in attention). This appears to be a WMMA bug.
What: Wrap every f16 MatMul with Cast(f16→f32) + MatMul(f32) + Cast(f32→f16), forcing IREE to use the f32 path. Depends on A3 (uses ParseVtensorType, BuildVtensorType, ResolveAlias). ~57 lines.
Alternatives:

• IREE (proper fix): Fix WMMA instruction selection to use WMMAR3_F32_16x16x16_F16 (f32 accumulator) for transposed-B cases, or fix the f16 accumulator path itself. This is where the bug lives: [GPU] Fix RDNA3 WMMA f16/bf16 accumulator layout iree#23806
• IREE (alternative): Recognize extf→matmul→truncf as a request for f32 accumulation and select the right WMMA variant. This would make the EP workaround at least intentional rather than a blunt hammer.
• The EP workaround promotes all f16 matmuls, not just the buggy transposed-B ones — unnecessary overhead for matmuls that would be correct with f16 accumulation. A targeted fix in IREE would be more precise.

[jtuyls]

Qwen1.5-MoE-A2.7B-Chat INT4 — IR Guide

Overview

This describes the intermediate representation (IR) generated by the IREE ONNX Runtime EP for the Qwen1.5-MoE-A2.7B-Chat model with batched INT4 QDQ quantization. The IR is torch-dialect MLIR with ONNX operator semantics, compiled by IREE to GPU kernels for AMD gfx1100.

Model: 24 layers, 60 experts/layer, top-4 routing, 16 KV heads, head_dim=128, max_seq=256
Weights: INT4 QDQ (block_size=128), stacked [60, ...] per layer, 16.5 GB total
Performance: 42.9 tok/s decode on RX 7900 XTX (gfx1100)

Flow Diagram

Input tokens [1, seq]
        │
        ▼
   ┌─────────────┐
   │  Embedding   │  Gather from [151936, 2048] → [1, seq, 2048] f16
   └──────┬──────┘
          │
          ▼
   ×24 layers ─────────────────────────────────────────────
   │                                                       │
   │  ┌──────────────────┐                                 │
   │  │   RMSNorm (pre)  │  Cast→Pow→Mean→Add→Sqrt→Div    │
   │  └────────┬─────────┘                                 │
   │           │                                           │
   │  ┌────────┴─────────┐                                 │
   │  │   Self-Attention  │                                │
   │  │  ┌─────────────┐ │                                 │
   │  │  │ QKV Proj     │ │  3× MatMul [2048,2048] + bias  │
   │  │  │ + RoPE       │ │  Split→Neg→Concat + Sin/Cos   │
   │  │  ├─────────────┤ │                                 │
   │  │  │ KV Cache     │ │  2× tensor.insert_slice        │
   │  │  │ Update       │ │  (EP rewrite from ScatterElem) │
   │  │  ├─────────────┤ │                                 │
   │  │  │ Attention    │ │  QK^T→Scale→Mask→Softmax→AV   │
   │  │  ├─────────────┤ │                                 │
   │  │  │ O Proj       │ │  MatMul [2048,2048]            │
   │  │  └─────────────┘ │                                 │
   │  └────────┬─────────┘                                 │
   │           │ + residual                                │
   │  ┌────────┴─────────┐                                 │
   │  │  RMSNorm (post)  │                                 │
   │  └────────┬─────────┘                                 │
   │           │                                           │
   │  ┌────────┴─────────────────────┐                     │
   │  │   MoE Block                  │                     │
   │  │  ┌────────────────────────┐  │                     │
   │  │  │ Router                 │  │                     │
   │  │  │ MatMul[2048,60]→Soft   │  │                     │
   │  │  │ →Cast→TopK(k=4)       │  │ hal.dispatch.extern  │
   │  │  └───────────┬───────────┘  │                     │
   │  │              │               │                     │
   │  │  ┌───────────┴───────────┐  │                     │
   │  │  │ Expert Computation    │  │                     │
   │  │  │ Gather top-4 from     │  │                     │
   │  │  │   [60,...] weights    │  │                     │
   │  │  │ DQL(INT4→f16) ×3     │  │                     │
   │  │  │ gate/up/down MatMul  │  │                     │
   │  │  │ SwiGLU activation    │  │                     │
   │  │  │ Weighted sum         │  │                     │
   │  │  └───────────┬───────────┘  │                     │
   │  │              │               │                     │
   │  │  ┌───────────┴───────────┐  │                     │
   │  │  │ Shared Expert         │  │                     │
   │  │  │ gate/up/down MatMul  │  │  FP16 (not INT4)    │
   │  │  │ SwiGLU + gate sigmoid│  │                     │
   │  │  └───────────┬───────────┘  │                     │
   │  │              │               │                     │
   │  │  Add(expert_out, shared_out)│                     │
   │  └──────────────┬───────────────┘                     │
   │                 │ + residual                          │
   ────────────────────────────────────────────────────────
          │
   ┌──────┴──────┐
   │  Final Norm │  RMSNorm
   └──────┬──────┘
          │
   ┌──────┴──────┐
   │   LM Head   │  MatMul [2048, 151936] → logits
   └─────────────┘

Per-Layer IR Patterns

1. RMSNorm

RMSNorm normalizes hidden states using root-mean-square, computed in f32 for precision.

Input: hidden [1, seq, 2048] f16 — hidden states from previous layer (or embedding)
Output: normalized [1, seq, 2048] f16 — normalized hidden states, scaled by learned weight
Parameters: weight [2048] f16 — per-channel scale (one per layer)

// Cast f16 → f32 for precision
%cast = onnx.Cast(%input) {to = 1} : [1,?,2048] f16 → [1,?,2048] f32
// x^2
%pow = onnx.Pow(%cast, %two) : [1,?,2048] f32
// mean(x^2, axis=-1, keepdims=True)
%mean = onnx.ReduceMean(%pow, %axes_neg1) {keepdims=1} : [1,?,2048] → [1,?,1] f32
// sqrt(mean + eps)
%norm = onnx.Add(%mean, %eps) → onnx.Sqrt → [1,?,1] f32
// x / norm
%div = onnx.Div(%cast, %norm) : [1,?,2048] f32
// Cast back and apply weight
%out = onnx.Cast(%div) {to = 10} → onnx.Mul(%weight) : [1,?,2048] f16

Two RMSNorm blocks per layer (pre-attention, pre-MoE) = 48 total.

2. Self-Attention (Multi-Head)

QKV Projections

Input: hidden [1, seq, 2048] f16 — normalized hidden states
Output: Q, K, V [1, 16, seq, 128] f16 — query, key, value per head
Parameters: W_q, W_k, W_v [2048, 2048] f16 + bias [2048] f16 (one set per layer)

%q = onnx.MatMul(%hidden, %W_q[2048,2048]) → onnx.Add(%bias_q) : [1,?,2048] f16
%q_4d = onnx.Reshape(%q, [1,seq,16,128]) → onnx.Transpose([0,2,1,3]) : [1,16,?,128] f16
// Same for K and V

Rotary Position Embeddings (RoPE)

Position-dependent rotation applied to Q and K.

Input: Q, K [1, 16, seq, 128] f16 — query/key tensors, seq_position — starting position offset
Output: Q_rot, K_rot [1, 16, seq, 128] f16 — position-encoded query/key
Parameters: freqs [1, 64, 1] f32 — precomputed frequency table

// Position indices: Range(0, seq) + seq_position → [seq] si64
// Frequency computation: Cast(si64→f32) → MatMul(freqs[1,64,1]) → Concat(doubled)
//   → Sin/Cos → Cast(f32→f16) → Reshape([1,1,seq,128])

// rotate_half: Split Q/K into halves, negate second, concat
%q_split0, %q_split1 = onnx.Split(%q_4d, axis=3) : [1,16,?,64] each
%q_rotated = onnx.Mul(%q_4d, %cos) + onnx.Mul(Concat(Neg(%q_split1), %q_split0), %sin)

KV Cache Update (EP rewrite: tensor.insert_slice)

The EP rewrites ScatterElements into efficient slice insertion.

Input: K_rot/V [1, 16, seq, 128] f16 — new key/value, past_key/value [1, 16, 256, 128] f16 — KV cache, seq_position — write offset
Output: updated_key/value [1, 16, 256, 128] f16 — cache with new KV inserted at position

// Bridge from torch to tensor dialect
%kv_data = torch_c.to_builtin_tensor %past_key : tensor<1x16x256x128xf16>
%kv_src_dyn = torch_c.to_builtin_tensor %k_rotated : tensor<?x?x?x?xf16>
// Refine type from KV cache shape (EP derives [1,16,?,128] from data)
%kv_src = tensor.cast %kv_src_dyn : tensor<?x?x?x?xf16> to tensor<1x16x?x128xf16>
// Extract offset and seq length
%offset = tensor.extract %seq_position[] → arith.index_cast : index
%seq_len = tensor.dim %kv_src, %c2 : index
// Insert new KV at position
%updated = tensor.insert_slice %kv_src into %kv_data
    [0, 0, %offset, 0] [1, 16, %seq_len, 128] [1, 1, 1, 1]
    : tensor<1x16x?x128xf16> into tensor<1x16x256x128xf16>

2 insert_slice per layer (key + value) = 48 total. Without this rewrite, ScatterElements decomposes into ~240 memcpy dispatches.

Attention Computation

Input: Q_rot [1, 16, seq, 128] f16 — rotated query, K_cached [1, 16, 256, 128] f16 — full key cache, V_cached [1, 16, 256, 128] f16 — full value cache
Output: attn_out [1, seq, 2048] f16 — attention output projected back to hidden_dim
Parameters: W_o [2048, 2048] f16 — output projection (one per layer)

// QK^T scores
%scores = onnx.MatMul(%q_rotated[1,16,?,128], %k_cached_T[1,16,128,256]) : [1,16,?,256] f16
%scaled = onnx.Mul(%scores, %inv_sqrt_128) : [1,16,?,256] f16

// Causal mask: positions ≤ seq_pos+seq_len are visible
%mask_le = onnx.LessOrEqual(%positions[1,256], %token_positions[?,?])
%mask_lt = onnx.Less(%positions[1,256], %seq_end)
%mask = onnx.And(%mask_le, %mask_lt) : [?,256] i1
%causal = onnx.Where(%mask, %zero, %neg_inf) : [?,256] f16
%masked = onnx.Add(%scaled, onnx.Reshape(%causal, [1,1,?,256]))

// Softmax + Cast roundtrip (f16→f32→f16, for numerical stability)
%attn = onnx.Softmax(%masked, axis=-1) : [1,16,?,256] f16
%attn_f32 = onnx.Cast(%attn, to=f32) → onnx.Cast(to=f16)  // roundtrip

// Attention × Values
%out = onnx.MatMul(%attn, %v_cached[1,16,256,128]) : [1,16,?,128] f16
%proj = onnx.Transpose([0,2,1,3]) → onnx.Reshape([?,?,2048]) → onnx.MatMul(%W_o[2048,2048])

3. MoE Block

Router (TopK via extern GPU kernel)

The EP emits a custom HIP kernel dispatch for fused Softmax→TopK.

Input: hidden [seq, 2048] f16 — flattened hidden states (batch×seq merged)
Output: values [seq, 4] f32 — top-4 routing weights, indices [seq, 4] i64 — top-4 expert indices
Parameters: W_gate [2048, 60] f16 — router weight matrix (one per layer)

// Gate linear: hidden[?,2048] × gate_weight[2048,60] → [?,60] f16
%gate = onnx.MatMul(%hidden_flat, %W_gate) : [?,60] f16

// The EP pattern-matches Softmax→Cast→TopK and replaces with a single
// extern dispatch. The Softmax and Cast ops are NOT emitted — the extern
// kernel handles softmax internally in f32 for stability.

// Bridge gate logits to builtin tensor for extern dispatch
%in = torch_c.to_builtin_tensor %gate : tensor<?x60xf16>
// Dynamic workload: ceil(N / 64) workgroups, 64 threads each (one RDNA wavefront)
%wl = arith.divui(arith.addi(%N, 63), 64)

// Extern TopK kernel — pre-compiled HIP kernel object (moe_topk_gfx1100.co)
%values, %indices = hal.dispatch.extern "moe_topk"[%wl](
    %N_i32, %num_experts_i32, %k_i32, %in)
    : (i32, i32, i32, tensor<?x60xf16>{%N}) → (tensor<?x4xf32>{%N}, tensor<?x4xi64>{%N})
  count(%device, %workload) → (index, index, index) {
    hal.return %workload, 1, 1   // 1D dispatch: one workgroup per ceil(N/64)
  }
  layout(#hal.pipeline.layout<
    constants = 3,               // push constants: N, num_experts, top_k
    bindings = [
      storage_buffer<ReadOnly>,  // binding0: gate_logits f16[N × 60]
      storage_buffer,            // binding1: top_k_values f32[N × 4]
      storage_buffer             // binding2: top_k_indices i64[N × 4]
    ]>)
  objects({ #executable_target = [#hal.executable.object<{path = "moe_topk_gfx1100.co"}>] })
  attributes {workgroup_size = [64, 1, 1]}

Kernel internals (moe_topk_kernel.cu): One thread per token, 64 threads per workgroup.
Each thread independently processes one token's 60-expert routing:

1. Softmax (in-kernel, f32): Read 60 f16 logits → convert to f32 → subtract max → exp → normalize. Uses a local float probs[64] array (registers/stack for ≤64 experts).
2. Top-4 selection (selection sort): O(k×E) = O(4×60) = 240 comparisons per token. Iterates k=4 times, each time scanning all 60 experts to find the next-best unselected one. Simple and fast for small k and E.
3. Write outputs: 4 float values (routing weights) + 4 int64 indices (expert IDs).

The kernel avoids IREE's sort-based TopK decomposition which creates ~216 sort dispatches for 60-element sort + top-4 extraction.

24 extern dispatches (one per layer). The extern kernel works around two IREE sort codegen bugs:

1. Sort shared memory overflow (topk_sort_bug_minimal.mlir): The ONNX model casts routing probabilities from f16→f32 before TopK (inherited from PyTorch's .float() for numerical safety, though f16 sorting gives identical results). When IREE fuses this Cast into the sort dispatch, the producer's output goes to workgroup memory. With dynamic dims, IREE allocates for worst-case size (up to 983KB), exceeding gfx1100's 64KB shared memory limit.

2. ArrayRef assertion crash (topk_reshape_crash.mlir): When TopK's [1,?,4] output is reshaped with shape=[-1], IREE's dispatch creation crashes during dimension collapsing.

Both bugs only trigger in the full model context, not standalone. The f16→f32 Cast before TopK is unnecessary for correctness (standalone f16 TopK produces identical top-4 indices), but removing it at the model level still hits bug #2. Both need IREE fixes to eliminate the extern kernel.

Expert Computation (INT4 QDQ, gather-then-compute)

Despite being called "batched," only the top-4 experts are computed per token.

Input: hidden [seq, 2048] f16 — hidden states, indices [seq, 4] i64 — which experts to use, values [seq, 4] f32 — routing weights
Output: moe_out [seq, 2048] f16 — weighted expert output
Parameters per layer: gate/up/down_weights [60, *, *] ui8 + scales [60, *, *] f16 + zeros [60, *, *] ui8 — INT4 QDQ packed weights for all 60 experts (only top-4 gathered at runtime)

// 1. Gather top-4 expert weights from stacked [60,...] tensors
//    indices_flat = [seq*4] (for decode: [4], for prefill seq=8: [32])
%indices_flat = onnx.Reshape(%topk_indices[seq,4], [-1]) : [seq*4] si64
%down_w = onnx.Gather(%all_down_weights[60,2048,1408], %indices_flat, axis=0) : [seq*4,2048,1408] ui8
%down_s = onnx.Gather(%all_down_scales[60,2048,11], %indices_flat, axis=0)   : [seq*4,2048,11] f16
%down_z = onnx.Gather(%all_down_zeros[60,2048,11], %indices_flat, axis=0)    : [seq*4,2048,11] ui8
// Same for gate and up projections
// (IR shows [?,..] because seq is dynamic; concrete shapes shown for decode seq=1)

// 2. Dequantize INT4 → f16 (block_size=128)
//    Lowered by torch-mlir: reshape → uitofp → sub(zp) → mul(scale) → reshape
%down_f16 = onnx.DequantizeLinear(%down_w, %down_s, %down_z)
    {block_size = 128, axis = 2} : [seq*4,2048,1408] f16
%down_transposed = onnx.Transpose(%down_f16, perm=[0,2,1]) : [seq*4,1408,2048] f16

// 3. Replicate hidden states for top-4 experts
//    For decode: [1,2048] → [1,1,2048] → [1,4,2048] → [4,1,2048]
%hidden_tiled = onnx.Unsqueeze(%hidden, dim=1)    // [seq,1,2048]
    → onnx.Tile([1,4,1])                           // [seq,4,2048]
    → onnx.Flatten(axis=2) → onnx.Reshape([seq*4,1,2048])

// 4. Expert MatMuls (batch dim = seq × top_k)
//    For decode (seq=1): [4,1,2048] × [4,2048,1408] → [4,1,1408]
//    For prefill (seq=8): [32,1,2048] × [32,2048,1408] → [32,1,1408]
%up_out = onnx.MatMul(%hidden_tiled, %up_transposed)     : [seq*4,1,1408] f16
%gate_out = onnx.MatMul(%hidden_tiled, %gate_transposed)  : [seq*4,1,1408] f16

// 5. SwiGLU activation
%act = onnx.Mul(%gate_out, onnx.Sigmoid(%gate_out)) → onnx.Mul(%up_out)

// 6. Down projection
%expert_out = onnx.MatMul(%act, %down_transposed) : [seq*4,1,2048] f16

// 7. Weighted sum with routing values
//    routing weights [seq,4] → [seq,1,4], expert outputs [seq*4,2048] → [seq,4,2048]
//    MatMul: [seq,1,4] × [seq,4,2048] → [seq,1,2048]  (weighted combination)
%weights_reshaped = onnx.Reshape(%topk_values, [seq,1,4])
%expert_reshaped = onnx.Reshape(%expert_out, [seq,4,2048])
%moe_out = onnx.MatMul(%weights_reshaped, %expert_reshaped) : [seq,1,2048] f16

Per layer: 3 Gather ops (from [60,...]), 3 DequantizeLinear (block_size=128), 3 expert MatMuls + 1 weighted-sum MatMul.

Shared Expert (FP16, no quantization)

Each layer has a shared expert that processes all tokens (not gated by TopK).

Input: hidden [seq, 2048] f16 — same hidden states as routed experts
Output: shared_out [seq, 2048] f16 — shared expert output, gated by a learned sigmoid
Parameters: W_shared_gate [2048, 1] f16, W_shared_gate_proj/up/down [2048, 5632] / [5632, 2048] f16 — FP16 weights (not quantized, larger intermediate dim 5632 vs 1408 for routed experts)

// Shared expert gate (controls contribution)
%shared_gate = onnx.MatMul(%hidden, %W_shared_gate[2048,1]) → onnx.Sigmoid : [?,1] f16

// Shared expert FFN (SwiGLU, same as regular experts but FP16)
%shared_up = onnx.MatMul(%hidden, %W_shared_up[2048,5632])    : [?,5632] f16
%shared_gate_proj = onnx.MatMul(%hidden, %W_shared_gate_proj[2048,5632]) : [?,5632] f16
%shared_act = onnx.Mul(%shared_gate_proj, onnx.Sigmoid(%shared_gate_proj)) → onnx.Mul(%shared_up)
%shared_down = onnx.MatMul(%shared_act, %W_shared_down[5632,2048]) : [?,2048] f16

// Gated combination
%shared_out = onnx.Mul(%shared_gate, %shared_down) : [?,2048] f16

// Final: expert output + shared expert output + residual
%layer_out = onnx.Add(%moe_out, %shared_out) → onnx.Reshape → onnx.Add(%residual)

Op Counts (Full Model, 24 Layers)

Category	Op	Count	Notes
Reshape/Shape	onnx.Reshape	391	Shape manipulation for attention heads, experts
	onnx.Shape	73	Dynamic dim extraction
	onnx.Unsqueeze	97	Broadcasting
	onnx.Concat	223	Shape tensor construction
Linear	onnx.MatMul	362	QKV, O, gate, expert, shared expert projections
Gather	onnx.Gather	290	Expert weight selection + shape indexing
Elementwise	onnx.Mul	289	RoPE, SwiGLU, scaling
	onnx.Add	267	Bias, residual, mask
	onnx.Sigmoid	72	SwiGLU activation (3/layer)
Cast	onnx.Cast	197	f16↔f32 for RMSNorm, attention, routing
Transpose	onnx.Transpose	193	Head reshape, expert weight transpose
Quantization	onnx.DequantizeLinear	72	INT4→f16 (3/layer, block_size=128)
Attention	onnx.Softmax	48	Attention + routing softmax (2/layer)
	onnx.Split	48	RoPE rotate_half
	onnx.Neg	48	RoPE rotate_half
Normalization	onnx.ReduceMean	49	RMSNorm (2/layer + final)
	onnx.Pow	49	RMSNorm x^2
	onnx.Sqrt/Div	49 each	RMSNorm normalization
Expert dispatch	onnx.Tile	24	Replicate hidden for top-k experts
EP rewrites	hal.dispatch.extern	24	TopK extern GPU kernel
	tensor.insert_slice	48	KV cache update

Total: ~5253 lines of MLIR, ~3400 ONNX ops + 72 EP-specific ops.

EP-Specific Rewrites

Two patterns are handled by the EP rather than generic ONNX→torch-mlir lowering:

1. TopK Extern Kernel (hal.dispatch.extern "moe_topk")

• Why: Two IREE bugs: (1) sort shared memory overflow when Cast(f16→f32) fuses into sort dispatch with dynamic dims, (2) ArrayRef crash when TopK output is reshaped with -1. The f16→f32 Cast before TopK is inherited from PyTorch but unnecessary — f16 sorting gives identical results.
• What: Replaces Softmax→Cast→TopK with a pre-compiled HIP kernel
• Impact: Avoids ~216 sort dispatches per TopK (24 per model)
• Reproducers: topk_sort_bug_minimal.mlir (bug 1), topk_reshape_crash.mlir (bug 2)

2. ScatterElements → tensor.insert_slice

• Why: IREE decomposes onnx.ScatterElements into hundreds of memcpy dispatches
• What: Pattern-matches KV cache scatter (ScatterElements(past_kv, Expand(Range+offset), updates)) and emits tensor.insert_slice directly
• Impact: 2 ops per layer vs ~240 dispatches. The tensor.cast refines the dynamic update type from the static KV cache shape.

Optimization Opportunities

Near-term (EP/torch-mlir level)

1. Remove TopK extern kernel: Fix two IREE bugs: (a) sort shared memory overflow when the f16→f32 Cast fuses into the sort dispatch — fix by tiling the producer per-workgroup instead of pre-materializing the full intermediate in shared memory, (b) ArrayRef crash in dispatch creation when TopK output is reshaped with -1. Also consider removing the unnecessary f16→f32 Cast at the model level (f16 sorting gives identical results). Standalone reproducers: topk_sort_bug_minimal.mlir, topk_reshape_crash.mlir.

2. Remove ScatterElements rewrite: Either fix IREE's ScatterElements decomposition to recognize contiguous scatter patterns, or accept the ~240 dispatch overhead.

3. Eliminate Cast→Cast roundtrips: The f16→f32→f16 Cast pair after each Softmax (24 per model) should be folded by torch-mlir/IREE canonicalization. Currently produces 48 unnecessary dispatches.

4. RoPE fusion: The RoPE pattern (Split→Neg→Concat + Sin/Cos multiply) generates many small dispatches. Could be fused into a single kernel.

5. Pre-transpose expert weights at export: The current flow does DequantizeLinear → Transpose → MatMul for each expert projection. The Transpose (perm=[0,2,1]) converts [4, 2048, 1408] to [4, 1408, 2048] for the MatMul. If the weights were stored pre-transposed in the ONNX model, the 72 Transpose ops (3 per layer) would be eliminated. This saves memory bandwidth (no materialized transpose) and dispatch overhead. Straightforward change in export_qwen_moe.py — transpose the weight tensors before saving.

Medium-term (IREE compiler level)

6. DQL+MatMul fusion: IREE's FuseDequantizationMatmul pass should fuse the DequantizeLinear(block_size=128) → Transpose → MatMul pattern into a single kernel. Currently the DQL decomposition (uitofp→subf→mulf) is a separate dispatch.

7. Dynamic shape specialization: Many tensors have shapes like [?,16,?,128] where ? is batch*seq. For decode (seq=1), these could be specialized at compile time via seq_len_divisor hints.

8. Attention kernel: Replace the decomposed QK^T → Scale → Mask → Softmax → AV with IREE's fused attention kernel (iree_linalg_ext.attention). Currently disabled due to known bugs on gfx1100.

Long-term (architecture level)

9. Fused MoE kernel: Replace the current Gather→DQL→Tile→MatMul×3→SwiGLU→WeightedSum decomposition with a fused kernel. The current flow has 6+ dispatch launches per layer and materializes intermediate tensors (Tile replicates hidden 4×, each expert projection is a separate buffer).

   How others do it: The standard approach (vLLM, ORT, Megablocks) is sort-permute-grouped-GEMM-unpermute:

   1. Sort tokens by expert: Flatten topk_ids[N,4] → sort by expert → sorted_token_ids (tokens grouped contiguously per expert) + expert_ids (which expert each group belongs to)
   2. Grouped GEMM: A single kernel launch processes all experts. Each thread block handles one tile of one expert's GEMM, using expert_ids to select the weight matrix and sorted_token_ids as a gather index into activations. No physical permutation needed — vLLM's Triton kernel uses indirection:

      offs_token = tl.load(sorted_token_ids_ptr + pid_m * BLOCK_M + tl.arange(0, BLOCK_M))
      a = tl.load(a_ptr + offs_token[:, None] * stride_am + ...)  # gather via indirection
      b = tl.load(b_ptr + expert_id * stride_be + ...)            # select expert weight

   3. Unpermute + weighted reduce: Each original token's k expert outputs are gathered and combined with routing weights.

   ONNX Runtime has com.microsoft.MoE (and QMoE for INT4/INT8) contrib ops with a CUTLASS grouped-GEMM backend that follows this exact pattern.

   For IREE, the most practical path is hal.dispatch.extern with a hand-written HIP kernel (similar to the TopK extern). The sort preprocessing can stay as standard IREE ops. Longer-term, a compiler pass could pattern-match the MoE subgraph and emit a grouped-GEMM schedule using iree_linalg_ext.custom_op or a new iree_linalg_ext.moe op.

   Expected impact: Eliminates the Tile (4× hidden state replication), 3 separate MatMul dispatches, separate DQL dispatches, and the weighted-sum MatMul — down to 1-2 kernel launches per MoE layer.

10. Static KV cache: Re-export with PyTorch's StaticCache so the cache is a fixed-size pre-allocated buffer. The insert_slice write pattern remains (new KV must still be written at the current position), but the cache tensor doesn't need to be passed in/out as function arguments — it becomes a mutable global, eliminating the 48-in/48-out ping-pong overhead per decode step.

11. Separate prefill/decode models: Currently, a single ONNX model serves both prefill and decode, with add_free_dimension_override_by_name("seq_len", 1) specializing for decode. This requires two separate ORT sessions (different seq_len configs) compiled independently. Exporting separate ONNX models with static seq_len=1 for decode would eliminate dimension overrides and shape inference complexity, but the current approach already works since ORT provides correct types natively.

[jtuyls]

Qwen1.5-MoE-A2.7B-Chat INT4 — IR Guide (Update)

What Changed

Two major optimizations since the previous IR guide:

1. Fused MoE extern kernels — gate+up+SwiGLU and down+reduce replaced with hal.dispatch.extern HIP kernels
2. Stateful KV cache — com.iree:LoadGlobal/StoreGlobal custom ops replace KV function I/O with util.global mutable

Performance Summary

Configuration	Decode	Throughput	vs Original
Original (IREE compiles everything)	26.4 ms	37.8 tok/s	baseline
Original + stateful KV	17.8 ms	56.2 tok/s	+49%
Fused MoE extern kernels	22.6 ms	44.0 tok/s	+16%
Fused MoE + stateful KV	15.3 ms	65.4 tok/s	+73%

Hardware: RX 7900 XTX (gfx1100), decode seq_len=1, max_seq=256, INT4 QDQ weights.

GPU Time Breakdown (rocprof, per inference)

Both configurations use stateful KV cache. "Fused" uses extern HIP kernels for MoE experts; "No Fused" lets IREE compile the gather+DQL+matmul.

Category	Fused MoE + Stateful	No Fused + Stateful	Delta
Shared expert	5.4ms (72 disp)	6.5ms (121 disp)	+1.1ms
MoE expert	1.2ms extern (48 disp)	2.4ms IREE (122 disp)	+1.2ms
LM head	1.4ms (1 disp)	2.8ms (2 disp)	+1.4ms
Attention	2.7ms (168 disp)	2.7ms (168 disp)	same
TopK	1.4ms (24 disp)	1.3ms (24 disp)	same
KV scatter	0.1ms	0.1ms	same
Copies	0.4ms	0.4ms	same
RMSNorm / elementwise	0.2ms (152 disp)	0.3ms (179 disp)	same
Wall clock	15.3ms	17.8ms	-2.5ms

Notes:

• Shared expert dominates (36% fused, 39% no-fused) — FP16 with large intermediate dim (5632 vs 1408 for routed experts), runs on every token
• Fused MoE externs achieve 347-351 GB/s (36-37% peak BW) with INT4 warp-per-row matvec; gate+up+SwiGLU: 34.7µs/call, down+reduce: 13.8µs/call
• TopK regressed to 59µs/call in the full model (was 4µs standalone) — not yet investigated
• LM head doubles in the no-fused model (2 dispatches vs 1) — a graph structure artifact where IREE doesn't fuse through Reshape in the larger graph
• Copies are only ~4 tiny H2D transfers for input_ids + seq_position; KV copies eliminated entirely by stateful cache

New EP Features

Fused MoE Extern Kernels

The gather→DQL→matmul→SwiGLU→weighted-sum decomposition (6+ dispatches/layer) is replaced by two hal.dispatch.extern kernels per layer:

Kernel 1: moe_gate_up_swiglu — Fused gate + up projection with SwiGLU activation

• Input: hidden [1, 2048] f16 (broadcast to 4 experts), gate/up weights [60, 1408, 2048] uint8 + scales/zeros, topk_indices [4] i64
• Output: swiglu [4, 1408] f16
• Design: warp-per-row matvec, each warp computes one element of 1408-dim intermediate
• 34.7 µs/call, 351 GB/s (37% peak BW), 61 VGPRs

Kernel 2: moe_grouped_gemm_int4 — Fused down projection + weighted reduce

• Input: swiglu [4, 1408] f16, down weights [60, 2048, 1408] uint8 + scales/zeros, topk_indices [4] i64, topk_weights [4] f32
• Output: final_output [1, 2048] f16 (already weighted-summed)
• Design: warp-per-row, each warp iterates over 4 token-expert pairs, accumulates weighted dot products
• 13.8 µs/call, 347 GB/s (36% peak BW), 46 VGPRs

Both kernels use vectorized 4-byte loads (uint32 for INT4 weights, 2×f16 for activations) and warp shuffle reduction. The INT4 dequantization (block_size=128) is fused into the dot product loop.

Graph surgery (fused_moe_graph_surgery.py) replaces all expert computation nodes per layer with these two ExternDispatch nodes. 1296 nodes removed, 144 added across 24 layers.

// Kernel 1: gate+up+SwiGLU
%swiglu = hal.dispatch.extern "moe_gate_up_swiglu"[176](
    %N_i32, %K_i32, %E_i32, %D_i32, %I_i32,
    %hidden, %gate_w, %gate_s, %gate_z, %up_w, %up_s, %up_z, %flat_indices)
  : (...) → tensor<4x1408xf16>
  layout(#hal.pipeline.layout<constants = 5, bindings = [...]>)
  objects({#executable_target = [#hal.executable.object<{path = "moe_gate_up_swiglu_gfx1100.co"}>]})

// Kernel 2: down+reduce
%output = hal.dispatch.extern "moe_grouped_gemm_int4"[256](
    %N_out_i32, %N_in_i32, %E_i32, %D_i32, %G_i32,
    %swiglu, %down_w, %down_s, %down_z, %flat_indices, %topk_weights)
  : (...) → tensor<1x2048xf16>

Stateful KV Cache (com.iree:LoadGlobal / com.iree:StoreGlobal)

KV cache state is moved inside the IREE VM using util.global mutable. The ONNX graph's 48 KV inputs and 48 KV outputs are replaced by custom ops that load/store from a single combined global buffer.

Why: The ONNX EP's KV I/O overhead was ~10ms/inference:

• 5.6ms: IREE copies all 137 input buffers into transient allocations (baked into compiled VMFB by stream.async.transfer)
• 4.4ms: 48 bind_ortvalue_input calls at 92µs each for KV rebinding per call

No EP-level or IREE compiler fix could eliminate these — IREE's stream execution model always copies external inputs into packed transient memory.

How: Graph surgery (convert_kv_to_stateful()) replaces KV I/O:

Before: func(input_ids, seq_pos, past_key_0, ..., past_value_23) → (logits, present_key_0, ..., present_value_23)
After:  func(input_ids, seq_pos) → (logits)    [KV lives in util.global]

1. Each past_key_N graph input → com.iree:LoadGlobal node (no inputs, produces KV tensor)
2. Each present_key_N graph output → com.iree:StoreGlobal node (consumes updated KV, no outputs)
3. Graph inputs reduced from 50 to 2, outputs from 49 to 1

Combined KV buffer: All 48 KV tensors are packed into a single util.global mutable @kv_cache : tensor<2x24x16x256x128xf16> (dim0=K/V, dim1=layer). This avoids an IREE buffer aliasing bug where per-tensor globals cause read-after-write hazards — IREE aliases the insert_slice destination to the global's storage, corrupting attention reads within the same function call. With a single combined buffer, each layer operates on different sub-regions via extract_slice/insert_slice, preventing aliasing.

MLIR emission (in mlir_gen.cc):

// Module-level global declaration
util.global private mutable @kv_cache = dense<0.0> : tensor<2x24x16x256x128xf16>

// LoadGlobal → extract_slice from combined buffer
func.func @main(%input_ids: ..., %seq_position: ...) -> (...) {
    %buf = util.global.load @kv_cache : tensor<2x24x16x256x128xf16>
    // Extract layer 0 key: [kv_type=0, layer=0, ...]
    %past_key_0 = tensor.extract_slice %buf[0, 0, 0, 0, 0] [1, 1, 16, 256, 128] [1,1,1,1,1]
        : tensor<2x24x16x256x128xf16> to tensor<1x16x256x128xf16>
    ...
    // After computation: insert_slice updated KV back
    %updated_buf = util.global.load @kv_cache : tensor<2x24x16x256x128xf16>
    %new_buf = tensor.insert_slice %present_key_0 into %updated_buf[0, 0, 0, 0, 0] ...
    util.global.store %new_buf, @kv_cache : tensor<2x24x16x256x128xf16>
}

EP registration: LoadGlobalOp and StoreGlobalOp are registered as stub custom ops in the com.iree domain (same as ExternDispatch). They are never executed by ORT — the EP compiles them to MLIR during graph compilation.

Runtime: No changes to ComputeImpl. The IREE function now takes 2 inputs and returns 1 output. The 1:1 sequential mapping works as-is. KV state persists across iree_runtime_call_invoke calls via the util.global mutable in the VM module state.

Updated Flow Diagram

Input tokens [1, 1]  (decode: single token)
        │
        ▼
   ┌─────────────┐
   │  Embedding   │  Gather from [151936, 2048] → [1, 1, 2048] f16
   └──────┬──────┘
          │
          ▼
   ×24 layers ─────────────────────────────────────────────
   │                                                       │
   │  ┌──────────────────┐                                 │
   │  │   RMSNorm (pre)  │  Cast→Pow→Mean→Add→Sqrt→Div    │
   │  └────────┬─────────┘                                 │
   │           │                                           │
   │  ┌────────┴─────────┐                                 │
   │  │   Self-Attention  │                                │
   │  │  QKV proj + RoPE │  3× MatMul + bias + rotate     │
   │  │  KV cache update │  2× extract_slice from global   │
   │  │                  │  2× tensor.insert_slice          │
   │  │                  │  2× insert_slice back to global  │
   │  │  Attention comp  │  QK^T→Scale→Mask→Softmax→AV    │
   │  │  O proj          │  MatMul [2048,2048]             │
   │  └────────┬─────────┘                                 │
   │           │ + residual                                │
   │  ┌────────┴─────────┐                                 │
   │  │  RMSNorm (post)  │                                 │
   │  └────────┬─────────┘                                 │
   │           │                                           │
   │  ┌────────┴─────────────────────┐                     │
   │  │   MoE Block                  │                     │
   │  │  Router: MatMul[2048,60]     │                     │
   │  │    → hal.dispatch.extern     │ "moe_topk"          │
   │  │      (fused softmax+topk)    │                     │
   │  │                              │                     │
   │  │  Experts (fused):            │                     │
   │  │    → hal.dispatch.extern     │ "moe_gate_up_swiglu"│
   │  │      (INT4 DQL+matvec+SwiGLU)│                     │
   │  │    → hal.dispatch.extern     │ "moe_grouped_gemm"  │
   │  │      (INT4 DQL+matvec+reduce)│                     │
   │  │                              │                     │
   │  │  Shared expert:              │                     │
   │  │    3× MatMul (FP16, d=5632) │                     │
   │  │    SwiGLU + gate sigmoid    │                     │
   │  │                              │                     │
   │  │  Add(expert_out, shared_out)│                     │
   │  └──────────────┬───────────────┘                     │
   │                 │ + residual                          │
   ────────────────────────────────────────────────────────
          │
   ┌──────┴──────┐
   │  Final Norm │  RMSNorm
   └──────┬──────┘
          │
   ┌──────┴──────┐
   │   LM Head   │  MatMul [2048, 151936] → logits
   └─────────────┘

Remaining Optimization Opportunities

Opportunity	Expected Impact	Effort
Shared expert BW efficiency (32% of peak, 3.1x gap)	5.4ms → ~2ms	Medium-High
Proper prefill model (batch all tokens)	Prefill: 2.1s → ~50ms	Medium
TopK regression (59µs vs 4µs standalone)	-1.3ms/decode	Low — investigate

Shared Expert Analysis

The shared expert is the #1 bottleneck at 5.4ms (36% of decode). Each of the 72 dispatches (3 matmuls × 24 layers) is a matvec: [1, 2048] × [2048, 5632] or [1, 5632] × [5632, 2048], reading 23 MB of FP16 weights per call.

Metric	Value
Measured	5.4ms (72 dispatches, 75µs/call)
Speed of light (960 GB/s)	1.7ms (24µs/call)
BW efficiency	32% (3.1x gap)

Paths to improve:

• Check IREE matvec codegen — 32% efficiency suggests suboptimal tiling/vectorization for these shapes. IREE's generated kernel may not saturate memory bandwidth for the [1,N]×[N,M] matvec pattern.
• INT4 quantize shared expert weights — halves data read (23→11.5 MB/call), cuts SoL to 0.9ms. Currently FP16 because the original model doesn't quantize shared experts. Requires re-export with quantized shared expert weights.
• Extern matvec kernel — same warp-per-row design as the fused MoE kernels (which achieve 36-37% BW). Would need 3 separate kernels (gate, up, down) or a fused gate+up+SwiGLU+down kernel. The fused MoE kernels already proved this design works at ~350 GB/s for similar shapes.

[jtuyls]

Update: 11.6ms / 86 tok/s decode on RX 7900 XTX

Note: New system, so numbers could slightly differ.

Replicated and extended the fused MoE pipeline on a new RX 7900 XTX system. Starting from the FP16 sparse ONNX model, the full pipeline now includes:

1. Routed expert INT4 ExternDispatch — moe_gate_up_swiglu + moe_down_proj_reduce (same as before)
2. Shared expert INT4 ExternDispatch (new) — shared_gate_silu + shared_up_proj + shared_down_proj with GPTQ INT4 weights downloaded from HuggingFace.
3. Stateful KV cache — LoadGlobal/StoreGlobal with combined tensor<2x24x16x256x128xf16> buffer
4. TopK extern dispatch — fused softmax + top-k selection (wave32)
5. O2 optimization — IREE handles RMSNorm and attention fusion natively (extern RMSNorm was slower than IREE-compiled at O2)

Decode breakdown comparison

Category	This update	Disp	Previous (REPRO.md)	Disp	Delta
Shared expert	2.8ms	108	5.4ms (FP16)	72	-2.6ms
MoE expert	2.7ms	72	1.2ms	48	+1.5ms
LM head	2.2ms	1	1.4ms	1	+0.8ms
Attention	3.8ms	179	2.7ms	168	+1.1ms
TopK	2.0ms	36	1.4ms	24	+0.6ms
Copies/elem/other	0.7ms	380	1.0ms	~170
Wall clock	11.6ms	776	15.3ms	~650	-3.7ms
Throughput	86 tok/s		65 tok/s		+32%

Shared expert BW utilization (peak 960 GB/s)

Kernel	Time	Data	BW	% peak
shared_gate_silu (INT4, 2D WG)	24us	6.0 MB	251 GB/s	26%
shared_up_proj (INT4, 2D WG)	24us	6.0 MB	251 GB/s	26%
shared_down_proj (INT4, 2D WG)	27us	6.0 MB	221 GB/s	23%
moe_gate_up_swiglu (INT4)	48us	6.0 MB	125 GB/s	13%
moe_down_proj_reduce (INT4)	28us	6.0 MB	213 GB/s	22%

Key findings

• Shared expert INT4 is 2x faster than FP16: 2.8ms vs 5.4ms. The GPTQ INT4 kernels with 2D workgroups read 4x less data with reasonable (23-26%) BW utilization.
• IREE native TopK is much worse than extern: disabling the extern TopK dispatch and letting IREE compile Softmax+TopK at O2 increased wall time from 11.6ms to 16.8ms (IREE uses sort-based decomposition generating hundreds of dispatches).
• RMSNorm extern is slower than IREE at O2: extern kernel at 17us/call vs IREE-compiled ~2us. Removed extern RMSNorm in favor of IREE native.
• Dispatch overhead dominates: 776 dispatches with 3.2ms overhead (4.1us/dispatch) = 28% of wall time. Reducing dispatch count is the main remaining opportunity.

Remaining gaps

• Matvec BW utilization: should be 50-70%.
• MoE routed expert kernels (+1.5ms): same kernel code as before but 72 vs 48 dispatches and slightly higher per-kernel time.
• Attention (+1.1ms): 179 vs 168 dispatches at O2.
• Dispatch overhead (3.2ms): 776 total dispatches. Needs to go down.

Note:

Also bringing up Qwen3-30B-A3B Q4 and Qwen3.5-35B-A3B Q4 for better competitive analysis.

[jtuyls]

Qwen3-30B-A3B & Qwen3.5-35B-A3BMoE Models — GPTQ INT4

Model	Status	Decode	VRAM
Qwen3-30B-A3B	Working, correct numerics	18.5 tok/s	~19 GB
Qwen3.5-35B-A3B	Eager correct, GPU blocked	—	~21 GB

Hardware: RX 7900 XTX (gfx1100, 96 CUs, 24 GB VRAM, 960 GB/s)

---

Model Architectures

	Qwen3-30B-A3B	Qwen3.5-35B-A3B
Layers	48 (all full attention)	40 (30 GatedDeltaNet + 10 full SDPA)
Hidden / Head dim	2048 / 128	2048 / 256
Q / KV heads	32 / 4 (GQA 8:1)	16 / 2 (GQA 8:1)
Experts	128 routed, top-8	256 routed, top-8, +1 shared
Expert FFN dim	768	512
Vocab	151,936	248,320
Total / Active params	30.5B / 3.3B	35B / 3B
RoPE	Full	Partial (25% dims)
State	Paged KV (48 layers)	Paged KV (10) + recurrent (30) + conv (30)

---

Architecture Diagram — Qwen3-30B-A3B

Input tokens [1, seq] (int64)
 │
 ▼
┌─────────────────────────────────────────────────────────────┐
│ Embedding     Gather from [151936, 2048] FP16               │  IREE-compiled
│               → [1, seq, 2048] FP16                         │
└──────────────────────────────┬──────────────────────────────┘
                               │
           ╔═══════════════════╧════════════════════╗
           ║          ×48 Transformer Layers         ║
           ╠════════════════════════════════════════╣
           ║                                         ║
           ║  ┌───────────────────────────────────┐  ║
           ║  │ RMSNorm (pre-attention)           │  ║  IREE-compiled
           ║  │   w * (x / rms(x)), FP16→FP32→FP16│  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ Q/K/V Projections (FP16 matmul)   │  ║  IREE-compiled
           ║  │   Q: [seq, 2048] → [seq, 4096] FP16│  ║  INT4→FP16 at export
           ║  │   K: [seq, 2048] → [seq, 512] FP16 │  ║
           ║  │   V: [seq, 2048] → [seq, 512] FP16 │  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ RoPE (full, 128 dims) + Q/K Norm  │  ║  IREE-compiled, FP16
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ Paged Attention (2-phase)         │  ║
           ║  │   Phase 1: KV Write to page pool  │  ║  ★ Extern HIP kernel
           ║  │     [256 pages, 4 KV heads,       │  ║    FP16 KV cache
           ║  │      16 tok/page, 128 dim]         │  ║
           ║  │   Phase 2: GQA (32Q:4KV)          │  ║  ★ Extern HIP kernel
           ║  │     FP16 in, FP32 accum            │  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ O Projection + Residual           │  ║  IREE-compiled
           ║  │   [seq, 4096] → [seq, 2048] FP16  │  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ RMSNorm (pre-MoE)                 │  ║  IREE-compiled
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ MoE Dispatch (GPTQ INT4)          │  ║
           ║  │   Router: →[128] softmax + top-8  │  ║  Inside HIP kernel
           ║  │   8 Expert FFNs (fused):           │  ║  ★ Extern HIP kernel
           ║  │     Gate/Up: INT4→FP32 → [768]    │  ║    1024 threads = 8×128
           ║  │     SwiGLU (FP32)                  │  ║
           ║  │     Down: → [2048] FP32→FP16      │  ║
           ║  │   Weighted sum + Residual          │  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ╚══════════════════╧═════════════════════╝
                               │
 ┌─────────────────────────────┴──────────────────────────────┐
 │ Final RMSNorm → LM Head [2048]→[151936] FP16               │  IREE-compiled
 └─────────────────────────────┬──────────────────────────────┘
                               ▼
                     logits [1, seq, 151936] FP16 → argmax

Architecture Diagram — Qwen3.5-35B-A3B

Input tokens [1, seq] (int64)
 │
 ▼
┌─────────────────────────────────────────────────────────────┐
│ Embedding     [248320, 2048] FP16 → [1, seq, 2048] FP16    │  IREE-compiled
└──────────────────────────────┬──────────────────────────────┘
                               │
           ╔═══════════════════╧════════════════════╗
           ║   ×40 Layers (30 GDN + 10 full SDPA)   ║
           ╠════════════════════════════════════════╣
           ║                                         ║
           ║  RMSNorm: (1+w) * (x/rms(x)), FP32     ║  IREE-compiled
           ║                  │                       ║
           ║          ┌───────┴───────┐               ║
           ║      GDN(30)          Full SDPA(10)      ║
           ║          │                │              ║
           ║  ┌───────┴──────┐ ┌──────┴─────────┐   ║
           ║  │GatedDeltaNet │ │ Q(→8192)+gate  │   ║
           ║  │ QKV proj     │ │ K(→512) V(→512)│   ║  IREE FP16 matmul
           ║  │ Conv1d state │ │ Partial RoPE   │   ║
           ║  │ Delta rule   │ │ GQA 16Q:2KV    │   ║
           ║  │ Gated norm   │ │ head_dim=256   │   ║
           ║  │ Out proj     │ │ sigmoid gate   │   ║
           ║  │              │ │                │   ║
           ║  │ FP16 in/out  │ │                │   ║
           ║  │ FP32 recur.  │ │                │   ║
           ║  │★Extern HIP   │ │★Extern HIP     │   ║
           ║  └──────┬───────┘ └───────┬────────┘   ║
           ║         └────────┬────────┘              ║
           ║                  │                       ║
           ║  + Residual + RMSNorm (pre-MoE)          ║  IREE-compiled
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ MoE (GPTQ INT4, 256 experts)      │  ║
           ║  │   Router →[256] top-8              │  ║  ★ Extern HIP kernel
           ║  │   Gate/Up →[512] SwiGLU Down→[2048]│  ║    1024 threads
           ║  │                                    │  ║
           ║  │ + Shared Expert (FP16)             │  ║  ★ Extern HIP kernel
           ║  │   silu(h@W_g)*(h@W_u)@W_d         │  ║    256 threads
           ║  │   * sigmoid(h @ gate) (IREE)       │  ║
           ║  │ + Residual                        │  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ╚══════════════════╧═════════════════════╝
                               │
 ┌─────────────────────────────┴──────────────────────────────┐
 │ Final RMSNorm → LM Head [2048]→[248320] FP16               │  IREE-compiled
 └─────────────────────────────┬──────────────────────────────┘
                               ▼
                     logits [1, seq, 248320] FP16 → argmax

---

What Runs Where

Component	Execution	Weight fmt	Compute dtype	30B	35B
Embedding	IREE	FP16	FP16	×1	×1
RMSNorm	IREE	FP16	FP16→FP32→FP16	×96	×80
Q/K/V/O projections	IREE	INT4→FP16 at export	FP16 matmul	×48	×10
Paged Attention	Extern HIP	—	FP16 KV, FP32 accum	×48	×10
MoE dispatch	Extern HIP	INT4 packed	INT4→FP32 on-the-fly	×48	×40
Shared expert	Extern HIP	FP16	FP16 in, FP32 accum	—	×40
GatedDeltaNet	Extern HIP	FP16	FP16 in, FP32 state	—	×30
LM head	IREE	FP16	FP16 matmul	×1	×1

---

Decode Performance (Qwen3-30B, RX 7900 XTX)

54 ms/token, 18.5 tok/s decode, 285 tok/s prefill, 18.2 GB VRAM

Decode Breakdown (rocprofv3)

Component	Per-layer	×48 layers	%
MoE dispatch	1.3 ms	62 ms	88%
Q/K/V/O projections	0.1 ms	5 ms	7%
Paged attention	3 us	0.14 ms	<1%
LM head	—	1.4 ms	2%
Overhead	—	~10 ms	3%

Kernel Trace (per-dispatch us, rocprofv3)

Kernel	Count	us
moe_dispatch_gptq	48	1,314
Q proj (4096×2048)	48	59
O proj (2048×4096)	48	42
K/V proj (512×2048)	96	7
Paged attn (WMMA)	48	~3
LM head (151936×2048)	1	1,380

---

Performance Comparison (RX 7900 XTX)

Engine	Quantization	Decode tok/s	Source
IREE (this, v4)	GPTQ INT4	18.5	this work
llama.cpp (Vulkan)	Q4_0	95-105	HF llama-bench
llama.cpp (ROCm)	Q4_0	76-78	HF llama-bench
llama.cpp (ROCm)	Q8_0	~45	HF llama-bench

Note: llama.cpp numbers are Qwen3.5-35B-A3B (~15% less active compute). Adjusted equivalent: ~85-90 Vulkan, ~65-70 ROCm.

Gap: 5-6x vs llama.cpp, almost entirely MoE (88% of decode, 1 CU out of 96). Multi-CU v5 → ~47 tok/s, + WMMA → ~60 tok/s.

---

Bugs/Issues Found

1. GPTQ qzeros offset: auto_gptq stores zero=7 (+1 offset), Qwen3.5 symmetric stores zero=8 (+0 offset). Mismatched offset produces wrong output. Fix: auto-detection by median zero value.

2. IREE LDS codegen (gfx1100): Dynamic batch_matmul ?x2048 @ 2048xN triggers >64 KB LDS in full model context. Workaround: iree-opt --linalg-generalize-named-ops --linalg-fold-unit-extent-dims.

3. Python iree.runtime extern dispatch bug: FunctionInvoker delivers zeros/stale data to extern dispatch input bindings. C++ iree-run-module works correctly on the same VMFB. Blocks Qwen3.5 GDN. Repro: iree-workspace/repro/extern_dispatch_binding0/

4. Tied output unsupported for extern dispatch: flow.dispatch -> (%buf{%size}) doesn't preserve kernel writes. Workaround: separate in/out bindings with full-buffer copy in kernel.

5. INT4 attention divergence: GPTQ INT4 attention weights cause numerical divergence after ~24 layers via Q/K RMSNorm noise amplification. Fix: hybrid FP16 attention + INT4 experts.

---

Next Steps

Opportunity	Expected impact	Complexity
Multi-CU MoE (split dispatch)	62→10 ms MoE → ~47 tok/s	Medium
+ WMMA matmul in MoE	2-3x throughput → ~60 tok/s	Medium
Fix Python runtime bug	Unblock Qwen3.5 GPU	Low (IREE PR)
Grouped GEMM (IREE native MoE)	Eliminate custom kernel	High
Reduce dispatch overhead	10→3 ms	Medium

Related issues/PRs filed:

• RDNA3/4 f32/f16/bf16 Matvec (M=1) Performance iree#23507
• [rocm] Extern dispatch (hal.executable.source) receives stale data for function argument inputs on binding 0 iree#24092
• Fix ExpandDestinationForallOp rejecting non-store users iree#24073
• [Stream] Prevent CSE from aliasing dispatch-allocated buffers iree#23665

[jtuyls]

Qwen3-30B-A3B & Qwen3.5-35B-A3B MoE Models — GPTQ INT4

Model Status

Model	Status	Decode	Prefill	VRAM
Qwen3-30B-A3B (extern dispatch)	Working, correct numerics	75 tok/s	1866 tok/s	~19 GB
Qwen3-30B-A3B (IREE native SDPA)	Working, correct numerics	44.5 tok/s	35-43 tok/s (sequential)	~19 GB
Qwen3.5-35B-A3B	Working, correct numerics	4.3 tok/s (not looked at much)	—	~21 GB

Hardware: RX 7900 XTX (gfx1100, 96 CUs, 24 GB VRAM, 960 GB/s)

Architecture Diagram — Qwen3-30B-A3B (Extern Dispatch, 75 tok/s)

Input tokens [1, seq] (int64)
 │
 ▼
┌─────────────────────────────────────────────────────────────┐
│ Embedding     Gather from [151936, 2048] FP16               │  IREE-compiled
│               → [1, seq, 2048] FP16                         │
└──────────────────────────────┬──────────────────────────────┘
                               │
           ╔═══════════════════╧════════════════════╗
           ║          ×48 Transformer Layers         ║
           ╠════════════════════════════════════════╣
           ║                                         ║
           ║  ┌───────────────────────────────────┐  ║
           ║  │ RMSNorm (pre-attention)           │  ║  IREE-compiled
           ║  │   w * (x / rms(x)), FP16→FP32→FP16│  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ Q/K/V Projections (FP16 matvec)   │  ║  IREE-compiled
           ║  │   Q: [1, 2048] → [1, 4096] FP16   │  ║  INT4→FP16 at export
           ║  │   K: [1, 2048] → [1, 512] FP16    │  ║  matvec_like tiling
           ║  │   V: [1, 2048] → [1, 512] FP16    │  ║
           ║  └───────────────┬───────────���───────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ RoPE (full, 128 dims) + Q/K Norm  │  ║  IREE-compiled, FP16
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ Paged Attention (3-phase decode)   │  ║
           ║  │   Phase 1: KV Write to page pool  │  ║  ★ Extern HIP kernel
           ║  │   Phase 2: Split-KV attention     │  ║  ★ Flash Decode kernel
           ║  │     (256-tok chunks, 96 CUs)      │  ║    GQA handled in kernel
           ║  │   Phase 3: Reduce (log-sum-exp)   │  ║  ★ Extern HIP kernel
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ O Projection + Residual           │  ║  IREE-compiled
           ║  │   [1, 4096] → [1, 2048] FP16      │  ║  matvec_like tiling
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ RMSNorm (pre-MoE)                 │  ║  IREE-compiled
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ MoE Split Dispatch (GPTQ INT4)    │  ║
           ║  │   Routing: softmax + top-8        │  ║  ★ Extern HIP (multi-CU)
           ║  │   Gate+Up: WMMA grouped GEMM      │  ║  ★ Extern HIP (multi-CU)
           ║  │   SwiGLU (FP32)                   │  ║
           ║  │   Down: → [2048] + reduce         │  ║  ★ Extern HIP (multi-CU)
           ║  │   Weighted sum + Residual          │  ║
           ║  └───────────────┬───────────────────┘  ��
           ║                  │                       ║
           ╚══════════════════╧═════════════════════╝
                               │
 ┌─────────────────────────────┴──────────────────────────────┐
 │ Final RMSNorm → LM Head [2048]→[151936] FP16               │  IREE-compiled
 └─────────────────────────────┬──────────────────────────────┘
                               ▼
                     logits [1, 1, 151936] FP16 → argmax

Architecture Diagram — Qwen3-30B-A3B (IREE Native SDPA, 44.5 tok/s)

Input tokens [1, seq] (int64)
 │
 ▼
┌─────────────────────────────────────────────────────────────┐
│ Embedding     [151936, 2048] FP16 → [1, 1, 2048] FP16      │  IREE-compiled
└──────────────────────────────┬──────────────────────────────┘
                               │
           ╔═══════════════════╧════════════════════╗
           ║          ×48 Transformer Layers         ║
           ╠════════════════════════════════════════╣
           ║                                         ║
           ║  RMSNorm → Q/K/V proj → Q/K Norm       ║  IREE-compiled
           ║  → RoPE (full, 128 dims)                ║  batch_matmul tiling
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ KV Cache Write (slice_scatter)    │  ║  IREE-compiled
           ║  │   flat [25165824] FP16 buffer      │  ║  Device-resident
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ GQA Expand (4→32 heads)           │  ║  Separate dispatches
           ║  │   K: broadcast [4,512,128]→[32,.] │  ║  at O0/O1 (~1ms total)
           ║  │   V: broadcast+transpose          │  ║  Fuses at O2 (blocked)
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ IREE Native Attention             │  ║
           ║  │   F.scaled_dot_product_attention   │  ║  IREE-compiled
           ║  │   → iree_linalg_ext.attention      │  ║  → online_attention
           ║  │   ~11 µs/layer (0.5 ms total)      │  ║  (flash-style tiling)
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ║  O proj + Residual + RMSNorm             ║  IREE-compiled
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ MoE Split Dispatch (same kernels) │  ║  ★ Extern HIP
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ╚══════════════════╧═════════════════════╝
                               │
                     LM Head → logits → argmax

Architecture Diagram — Qwen3.5-35B-A3B (4.3 tok/s)

Input tokens [1, seq] (int64)
 │
 ▼
┌─────────────────────────────────────────────────────────────┐
│ Embedding     [248320, 2048] FP16 → [1, seq, 2048] FP16    │  IREE-compiled
└──────────────────────────────┬──────────────────────────────┘
                               │
           ╔═══════════════════╧════════════════════╗
           ║   ×40 Layers (30 GDN + 10 full SDPA)   ║
           ╠════════════════════════════════════════╣
           ║                                         ║
           ║  RMSNorm: (1+w) * (x/rms(x)), FP32     ║  IREE-compiled
           ║                  │                       ║
           ║          ┌───────┴───────┐               ║
           ║      GDN(30)          Full SDPA(10)      ║
           ║          │                │              ║
           ║  ┌───────┴──────┐ ┌──────┴─────────┐   ║
           ║  │GatedDeltaNet │ │ Q(→8192)+gate  │   ║
           ║  │ QKV proj     │ │ K(→512) V(→512)│   ║  IREE FP16 matmul
           ║  │ Conv1d state │ │ Partial RoPE   │   ║
           ║  │ Delta rule   │ │ GQA 16Q:2KV    │   ║
           ║  │ Gated norm   │ │ head_dim=256   │   ║
           ║  │ Out proj     │ │ sigmoid gate   │   ║
           ║  │              │ │                │   ║
           ║  │ FP16 in/out  │ │                │   ║
           ║  │ FP32 recur.  │ │                │   ║
           ║  │★Extern HIP   │ │★Extern HIP     │   ║
           ║  └──────┬───────┘ └───────┬────────┘   ║
           ║         └────────┬────────┘              ║
           ║                  │                       ║
           ║  + Residual + RMSNorm (pre-MoE)          ║  IREE-compiled
           ║                  │                       ║
           ║  ┌───────────────┴───────────────────┐  ║
           ║  │ MoE (GPTQ INT4, 256 experts)      │  ║
           ║  │   Router →[256] top-8              │  ║  ★ Extern HIP kernel
           ║  │   Gate/Up →[512] SwiGLU Down→[2048]│  ║    1024 threads
           ║  │                                    │  ║
           ║  │ + Shared Expert (FP16)             │  ║  ★ Extern HIP kernel
           ║  │   silu(h@W_g)*(h@W_u)@W_d         │  ║    256 threads
           ║  │   * sigmoid(h @ gate) (IREE)       │  ║
           ║  │ + Residual                        │  ║
           ║  └───────────────┬───────────────────┘  ║
           ║                  │                       ║
           ╚══════════════════╧═════════════════════╝
                               │
 ┌─────────────────────────────┴──────────────────────────────┐
 │ Final RMSNorm → LM Head [2048]→[248320] FP16               │  IREE-compiled
 └─────────────────────────────┬──────────────────────────────┘
                               ▼
                     logits [1, seq, 248320] FP16 → argmax

Decode Performance (Qwen3-30B, RX 7900 XTX)

Extern Dispatch: 75 tok/s, 14 ms/tok

Component	Per-layer us	x48 layers ms	%
MoE split dispatch	~105	5.0	36%
Attention (flash decode)	~20	2.1	15%
Q proj (matvec_like 4096x2048)	31.7	1.5	11%
O proj (matvec_like 2048x4096)	23.9	1.1	8%
K/V proj (512x2048)	4.5	0.4	3%
Router (128x2048)	3.7	0.2	1%
Other (norms, RoPE, KV write)	—	2.7	19%
LM head (151936x2048)	—	0.8	6%
Total		~14

Native SDPA: 44.5 tok/s, 22.5 ms/tok

Component	Per-layer us	x48 layers ms	%	vs Extern
MoE split dispatch	~124	5.9	26%	1.2x
Q proj (batch_matmul 1x1x4096x2048)	45.9	2.2	10%	1.5x
O proj (batch_matmul 1x1x2048x4096)	40.4	1.9	9%	1.7x
GQA K broadcast	3.7	0.18	1%	new
GQA V broadcast+transpose	17.2	0.83	4%	new
Attention (iree_linalg_ext.attention)	11.3	0.54	2%	0.5x faster
K/V proj (512x2048)	6.3	0.6	3%	1.4x
K/V cache scatter	5.5	0.5	2%	similar
Other (norms, RoPE, elementwise)	—	4.4	20%	8.4x
LM head	—	1.4	6%	similar
Total		~22.5		1.6x

Performance Comparison (RX 7900 XTX)

Engine	Quantization	Decode tok/s	Notes
IREE extern dispatch	GPTQ INT4	75	Split MoE extern dispatch + flash decode
IREE native SDPA	GPTQ INT4	44.5	IREE codegen attention + split MoE extern dispatch
llama.cpp (Vulkan)	Q4_0	95-105	HF llama-bench
llama.cpp (ROCm), measured now	Q4_0	97	HF llama-bench

Next Steps

• Fix RoPE dispatch O2 crash: [Codegen][ROCM] negative shared memory size (-8192 bytes) for fused RoPE dispatch on gfx1100 iree#24149
• Improve RMSNorm fusion, will remove 192 dispatches: RMSNorm reduction doesn't fuse with its normalize+scale consumer iree#24148
• Improve RDNA3 matvec perf: RDNA3/4 f32/f16/bf16 Matvec (M=1) Performance iree#23507
• [rocm] Extern dispatch (hal.executable.source) receives stale data for function argument inputs on binding 0 iree#24092
• [Stream] Prevent CSE from aliasing dispatch-allocated buffers iree#23665
• Ensure i4 dequant is fused into all matmuls
• Use Stateful KV cache for SDPA version
• 35B split dispatch MoE: 4.3 → ~15+ tok/s
• 35B GDN state copy elimination: Reduce 3.6 GB/token wasted traffic