diff --git a/iree/turbine/kernel/ops/wave_ops.py b/iree/turbine/kernel/ops/wave_ops.py
index 65231330f..1a3a4d12c 100644
--- a/iree/turbine/kernel/ops/wave_ops.py
+++ b/iree/turbine/kernel/ops/wave_ops.py
@@ -42,23 +42,20 @@
 
 def allocate(
     shape: tuple[IndexExpr], dtype: DataType, address_space: IndexSymbol
-) -> "Memory":
-    ...
+) -> "Memory": ...
 
 
 def self_index(
     dim: IndexExpr,
     dtype: DataType,
     elements_per_thread: Optional[IndexExpr | int] = None,
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
 def extract(
     register: "Register",
     offsets: tuple[IndexExpr],
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
 def extract_slice(
@@ -66,12 +63,10 @@ def extract_slice(
     offsets: tuple[IndexExpr],
     sizes: tuple[IndexExpr],
     strides: tuple[IndexExpr],
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
-def shared_memory_barrier():
-    ...
+def shared_memory_barrier(): ...
 
 
 def read(
@@ -79,28 +74,25 @@ def read(
     elements_per_thread: Optional[IndexExpr | int] = None,
     mapping: Optional[IndexMapping] = None,
     mapping_dynamic_vals: "Register" | tuple["Register", ...] = (),
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
 def conditional(
     condition: "Register" | IndexExpr,
-) -> Callable[[Callable[[], None]], None]:
-    ...
+) -> Callable[[Callable[[], None]], None]: ...
 
 
 def iterate(
     axis: IndexExpr, init_args: Sequence["Register"]
-) -> Callable[[Callable[[AccT], AccT]], AccT]:
-    ...
+) -> Callable[[Callable[[AccT], AccT]], AccT]: ...
 
 
-def register(shape: tuple[IndexExpr, ...], dtype: DataType, value: float) -> "Register":
-    ...
+def register(
+    shape: tuple[IndexExpr, ...], dtype: DataType, value: float
+) -> "Register": ...
 
 
-def mma(lhs: "Register", rhs: "Register", acc: "Register") -> "Register":
-    ...
+def mma(lhs: "Register", rhs: "Register", acc: "Register") -> "Register": ...
 
 
 def write(
@@ -109,32 +101,27 @@ def write(
     elements_per_thread: Optional[IndexExpr | int] = None,
     mapping: Optional[IndexMapping] = None,
     mapping_dynamic_vals: "Register" | tuple["Register", ...] = (),
-):
-    ...
+): ...
 
 
-def apply_expr(value: "Register" | Sequence["Register"], expr: Callable) -> "Register":
-    ...
+def apply_expr(
+    value: "Register" | Sequence["Register"], expr: Callable
+) -> "Register": ...
 
 
-def set_symbol(symbol: IndexExpr, value: "Register"):
-    ...
+def set_symbol(symbol: IndexExpr, value: "Register"): ...
 
 
-def exp2(src: "Register") -> "Register":
-    ...
+def exp2(src: "Register") -> "Register": ...
 
 
-def log2(src: "Register") -> "Register":
-    ...
+def log2(src: "Register") -> "Register": ...
 
 
-def reciprocal(src: "Register") -> "Register":
-    ...
+def reciprocal(src: "Register") -> "Register": ...
 
 
-def abs(src: "Register") -> "Register":
-    ...
+def abs(src: "Register") -> "Register": ...
 
 
 def tanh_approx(src: "Register") -> "Register":
@@ -145,88 +132,73 @@ def tanh(src: "Register") -> "Register":
     ...
 
 
-def roundeven(src: "Register") -> "Register":
-    ...
+def roundeven(src: "Register") -> "Register": ...
 
 
-def maximum(lhs: "Register", rhs: "Register") -> "Register":
-    ...
+def maximum(lhs: "Register", rhs: "Register") -> "Register": ...
 
 
-def minimum(lhs: "Register", rhs: "Register") -> "Register":
-    ...
+def minimum(lhs: "Register", rhs: "Register") -> "Register": ...
 
 
 def broadcast(
     arg: "Register", target_shape: Optional[Sequence[IndexExpr | int]] = None
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
 def sum(
     src: "Register",
     acc: Optional["Register"] = None,
     dim: Optional[IndexExpr | int] = None,
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
 def max(
     src: "Register",
     acc: Optional["Register"] = None,
     dim: Optional[IndexExpr | int] = None,
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
 def min(
     src: "Register",
     acc: Optional["Register"] = None,
     dim: Optional[IndexExpr | int] = None,
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
-def shuffle(src: "Register", offset: int, width: int) -> "Register":
-    ...
+def shuffle(src: "Register", offset: int, width: int) -> "Register": ...
 
 
-def gt(lhs: "Register", rhs: "Register") -> "Register":
-    ...
+def gt(lhs: "Register", rhs: "Register") -> "Register": ...
 
 
-def ge(lhs: "Register", rhs: "Register") -> "Register":
-    ...
+def ge(lhs: "Register", rhs: "Register") -> "Register": ...
 
 
-def lt(lhs: "Register", rhs: "Register") -> "Register":
-    ...
+def lt(lhs: "Register", rhs: "Register") -> "Register": ...
 
 
-def le(lhs: "Register", rhs: "Register") -> "Register":
-    ...
+def le(lhs: "Register", rhs: "Register") -> "Register": ...
 
 
-def eq(lhs: "Register", rhs: "Register") -> "Register":
-    ...
+def eq(lhs: "Register", rhs: "Register") -> "Register": ...
 
 
-def cast(src: "Register", dtype: DataType) -> "Register":
-    ...
+def cast(src: "Register", dtype: DataType) -> "Register": ...
 
 
-def permute(src: "Register", target_shape: Sequence[IndexExpr]) -> "Register":
-    ...
+def permute(src: "Register", target_shape: Sequence[IndexExpr]) -> "Register": ...
 
 
 def reshape(
     inputs: Sequence["Register"], target_vector_shape: dict[IndexSymbol, int]
-) -> "Register":
-    ...
+) -> "Register": ...
 
 
-def select(cond: "Register", if_true: "Register", if_false: "Register") -> "Register":
-    ...
+def select(
+    cond: "Register", if_true: "Register", if_false: "Register"
+) -> "Register": ...
 
 
 def define_op(op_name: str) -> Callable[[T], T]:
@@ -430,7 +402,10 @@ def custom_string(self, value_map: dict[str, str]) -> str:
         if hasattr(self.fx_node, "index") and self.fx_node.index:
             vars_list.append(f"index={self.fx_node.index}")
         vars_str = ", ".join(vars_list)
-        return f"{self.tkw_op_name}({vars_str}) type({self.fx_node.type})"
+        return f"""{self.tkw_op_name}({vars_str})
+        type({self.fx_node.type})
+        indexing_dims({self.indexing_dims if hasattr(self, 'indexing_dims') else None})
+        vector_shapes({self.vector_shapes if hasattr(self, 'vector_shapes') else None})"""
 
     def add_to_graph(self, region_graph: RegionGraph, type: Any = None) -> fx.Node:
         arg_list = tuple([value for _, value in vars(self).items()])
@@ -1226,7 +1201,10 @@ def custom_string(self, value_map: dict[str, str]) -> str:
         custom_str += f"lhs={self.lhs} (index = {self.lhs_index}), "
         custom_str += f"rhs={self.rhs} (index = {self.rhs_index}), "
         custom_str += f"acc={self.acc} (index = {self.acc_index}))"
-        custom_str += f" type({self.fx_node.type})"
+        custom_str += f"\n\ttype({self.fx_node.type})"
+        custom_str += f"\n\tindex({self.index})"
+        custom_str += f"\n\tindexing_dims({self.indexing_dims if hasattr(self, 'indexing_dims') else None})"
+        custom_str += f"\n\tvector_shapes({self.vector_shapes if hasattr(self, 'vector_shapes') else None})"
         return custom_str
 
     def align_index(self, constraints: list["Constraint"]) -> None:
@@ -1529,20 +1507,19 @@ def outputs(self, graph: fx.Graph) -> list[fx.Node]:
 
     @property
     def index(self) -> list[dict[IndexSymbol, IndexSequence]]:
+        def get_index(node: fx.Node) -> dict[IndexSymbol, IndexSequence]:
+            custom = get_custom(node)
+            if isinstance(custom, MMA):
+                return custom.acc_index
+            return custom.index
+
         for node in self.get_root_graph().subgraphs[self.subgraph_name].nodes:
             if isinstance(output := get_custom(node), Output):
                 return_vals = output.return_vals[0]
                 return (
-                    [
-                        (
-                            get_custom(val).acc_index
-                            if isinstance(get_custom(val), MMA)
-                            else val.index
-                        )
-                        for val in return_vals
-                    ]
+                    [get_index(val) for val in return_vals]
                     if isinstance(return_vals, (Sequence))
-                    else return_vals.index
+                    else get_index(return_vals)
                 )
 
     @index.setter
@@ -1759,6 +1736,8 @@ def index(self) -> dict[IndexSymbol, IndexSequence]:
             return None
         if not isinstance(custom, Iterate):
             return custom_index
+        if not isinstance(custom_index, Sequence):
+            return custom_index
         assert isinstance(custom_index, Sequence) and self.res_idx < len(
             custom.indexing_dims
         ), f"Invalid {custom_index=} with {self.res_idx=} and {custom.indexing_dims=}\n{custom}"
diff --git a/iree/turbine/kernel/wave/analysis/index_sequence_analysis.py b/iree/turbine/kernel/wave/analysis/index_sequence_analysis.py
index cb8d367aa..96ab4570c 100644
--- a/iree/turbine/kernel/wave/analysis/index_sequence_analysis.py
+++ b/iree/turbine/kernel/wave/analysis/index_sequence_analysis.py
@@ -9,6 +9,7 @@
     BinaryPyOp,
     Broadcast,
     CustomOp,
+    GetResult,
     IterArg,
     MMA,
     NestedRegionOp,
@@ -37,9 +38,7 @@
     get_workgroup_constraints,
     partial,
 )
-from ..utils.mma_utils import (
-    get_mma_dimensional_mapping,
-)
+from ..utils.mma_utils import get_mma_dimensional_mapping
 from ..utils.graph_utils import (
     get_inputs,
     get_users,
@@ -462,6 +461,11 @@ def propagate_indices(
     visited: set[CustomOp],
     symbolic_constraints: list[SymbolicAlias],
 ):
+    def get_index(custom: CustomOp):
+        if isinstance(custom, MMA):
+            return custom.acc.index
+        return custom.index
+
     """
     Propagate the index and vector shapes through the graph
     starting with priveleged nodes (like MMA, Read, Write).
diff --git a/iree/turbine/kernel/wave/analysis/partition_strided_operators.py b/iree/turbine/kernel/wave/analysis/partition_strided_operators.py
index 17b76f4f7..750ac0c76 100644
--- a/iree/turbine/kernel/wave/analysis/partition_strided_operators.py
+++ b/iree/turbine/kernel/wave/analysis/partition_strided_operators.py
@@ -39,53 +39,59 @@
 logger = get_logger("turbine.wave.partition_strided_operators")
 
 
-def get_vector_shape(
-    vector_shapes: dict[IndexSymbol, int],
+def get_concrete_shape(
+    vector_shapes_map: dict[IndexSymbol, int],
     symbolic_shape: list[IndexSymbol],
 ) -> list[int]:
-    vector_shapes = [max(vector_shapes[dim], 1) for dim in symbolic_shape]
-    return vector_shapes
+    concrete_shape = [max(vector_shapes_map[dim], 1) for dim in symbolic_shape]
+    return concrete_shape
 
 
-def _get_symbolic_shape_and_vector_shapes(
+def _get_read_symbolic_shape(
     custom: CustomOp,
 ):
-    register_shape = custom.register_type.symbolic_shape
-    vector_shapes = custom.vector_shapes
+    return custom.memory_type.symbolic_shape
+
+
+def _get_write_symbolic_shape(
+    custom: CustomOp,
+):
+    vector_shapes_map = custom.vector_shapes
     memory_shape = custom.memory_type.symbolic_shape
+    register_shape = custom.register_type.symbolic_shape
     # Check to see if the memory shape does not match with the vector shapes.
-    if not set(memory_shape).issubset(set(vector_shapes.keys())):
-        return register_shape, vector_shapes
+    if not set(memory_shape).issubset(set(vector_shapes_map.keys())):
+        return register_shape
     # Pick the shape with the most dimensions.
     if len(memory_shape) > len(register_shape):
-        return memory_shape, vector_shapes
-    return register_shape, vector_shapes
+        return memory_shape
+    return register_shape
 
 
 def partition_strided_operators(trace: CapturedTrace, constraints: list[Constraint]):
     """
     This function analyzes the index sequence of operators in the graph
-    that are writes on 2d tensors. If the operator has an access pattern where
+    that are reads or writes on 2d tensors. If the operator has an access pattern where
     the strides are greater than one on a single dimension, this function splits the
     operands into individual elements and constructs a write for
     each individual element.
     """
 
+    op_types = Read | Write
+    op_index_fn = {Read: lambda x: x.index, Write: lambda x: x.register_index}
+
     def has_strided_access(node: fx.Node) -> bool:
         """
-        Checks for writes on 2d tensors with strided access on a single dimension that
+        Checks for reads or writes on 2d tensors with strided access on a single dimension that
         read more than a single element.
         """
+
         custom = get_custom(node)
-        if isinstance(custom, Write):
-            strides = [
-                simplify_index(custom.register_index[dim]).stride
-                for dim in custom.register_index
-            ]
-            elements_per_thread = [
-                simplify_index(custom.register_index[dim]).size
-                for dim in custom.register_index
-            ]
+        if isinstance(custom, op_types):
+            op_type = Read if isinstance(custom, Read) else Write
+            index = op_index_fn[op_type](custom)
+            strides = [simplify_index(index[dim]).stride for dim in index]
+            elements_per_thread = [simplify_index(index[dim]).size for dim in index]
             strides = [x for x, y in zip(strides, elements_per_thread) if y > 1]
             num_strided_accesses = sum(1 for stride in strides if stride > 1)
             if num_strided_accesses > 1:
@@ -98,25 +104,33 @@ def has_strided_access(node: fx.Node) -> bool:
     strided_operators = trace.walk(has_strided_access)
     for operator in strided_operators:
         custom = get_custom(operator)
+        op_type = Read if isinstance(custom, Read) else Write
+        index = op_index_fn[op_type](custom)
+
         simplified_index = {
-            dim: simplify_index(custom.register_index.get(dim, custom.index[dim]))
+            dim: simplify_index(index.get(dim, custom.index[dim]))
             for dim in custom.index
         }
 
-        symbolic_shape, vector_shapes = _get_symbolic_shape_and_vector_shapes(custom)
+        vector_shapes_map = custom.vector_shapes
+        symbolic_shape = (
+            _get_read_symbolic_shape(custom)
+            if op_type == Read
+            else _get_write_symbolic_shape(custom)
+        )
 
-        shape = get_vector_shape(vector_shapes, symbolic_shape)
+        concrete_shape = get_concrete_shape(vector_shapes_map, symbolic_shape)
         elements_per_thread = subs_idxc(custom.elements_per_thread)
         max_stride_dim, max_stride = max(
             [(dim, seq.stride) for dim, seq in simplified_index.items()],
             key=lambda item: item[1],
         )
 
-        # Compute offsets we will aplly to each index element for each partitioned
+        # Compute offsets we will apply to each index element for each partitioned
         # write.
         offsets = np.array(
             [
-                np.unravel_index(int(i * max_stride), shape)
+                np.unravel_index(int(i * max_stride), concrete_shape)
                 for i in range(elements_per_thread)
             ]
         )
@@ -149,7 +163,7 @@ def check_contiguous_index():
                         (j if i == fastest_mem_dim_idx else 0)
                         for j in range(elements_per_thread)
                     ]
-                    for i in range(len(shape))
+                    for i in range(len(concrete_shape))
                 ]
             ).T
             if not np.array_equal(offsets, expected_offsets):
@@ -175,27 +189,61 @@ def check_contiguous_index():
             for i in range(elements_per_thread):
                 # Non-contiguous access patterns can have varying offsets. We
                 # handle that here.
-                extract = ExtractSlice(custom.register_, [i], [1], [1]).add_to_graph(
+                offset = offsets[i]
+                if op_type == Write:
+                    extract = ExtractSlice(
+                        custom.register_, [i], [1], [1]
+                    ).add_to_graph(custom.graph)
+
+                    write = Write(
+                        extract,
+                        custom.memory,
+                        mapping=custom.mapping,
+                        elements_per_thread=1,
+                    ).add_to_graph(custom.graph)
+                    write.index = {
+                        dim: IndexSequence(
+                            simplified_index[dim].start.subs({GPR_NUM: 0}) + offset[j],
+                            1,
+                            1,
+                        )
+                        for j, dim in enumerate(symbolic_shape)
+                    }
+                    ops_to_combine.append(write)
+                else:
+                    read = Read(
+                        custom.memory,
+                        elements_per_thread=1,
+                        mapping=custom.mapping,
+                        _write_dependency=custom._write_dependency,
+                    ).add_to_graph(custom.graph)
+                    read.index = {
+                        dim: IndexSequence(
+                            simplified_index[dim].start.subs({GPR_NUM: 0}) + offset[j],
+                            1,
+                            1,
+                        )
+                        for j, dim in enumerate(symbolic_shape)
+                    }
+                    ops_to_combine.append(read)
+
+        # Update users of original op.
+        if isinstance(custom, Write):
+            # Useful to handle write/read dependency
+            custom.replace_all_uses_with(ops_to_combine)
+        elif isinstance(custom, (Read, SelfIndex)):
+            with custom.graph.inserting_before(operator):
+                reshape = Reshape(ops_to_combine, custom.vector_shapes).add_to_graph(
                     custom.graph
                 )
+            reshape.expanded_dims = custom.expanded_dims
+            reshape.vector_shapes = custom.vector_shapes
+
+            # Save the original index on the reshape op so later we can
+            # detect if op was part of `gpr_offset` partition.
+            reshape.index = custom.index
+            custom.replace_all_uses_with(reshape)
 
-                offset = offsets[i]
-                write = Write(
-                    extract,
-                    custom.memory,
-                    mapping=custom.mapping,
-                    elements_per_thread=1,
-                ).add_to_graph(custom.graph)
-                write.index = {
-                    dim: IndexSequence(
-                        simplified_index[dim].start.subs({GPR_NUM: 0}) + offset[j], 1, 1
-                    )
-                    for j, dim in enumerate(symbolic_shape)
-                }
-                ops_to_combine.append(write)
-
-        # Useful to handle write/read dependency
-        custom.replace_all_uses_with(ops_to_combine)
         custom.graph.erase_node(operator)
 
 
diff --git a/iree/turbine/kernel/wave/compile.py b/iree/turbine/kernel/wave/compile.py
index 33bf2182c..15401079f 100644
--- a/iree/turbine/kernel/wave/compile.py
+++ b/iree/turbine/kernel/wave/compile.py
@@ -120,6 +120,7 @@ def wave_compile(options: WaveCompileOptions, kernel: "LaunchableWave") -> WaveK
     if options.compile_to_mlir:
         return WaveKernel(options, None, asm)
 
+    print(asm)
     compiled_wave_vmfb = compile_to_vmfb(asm, options)
     if options.create_vmfb_file:
         _write_file(options.create_vmfb_file, "wb", compiled_wave_vmfb)
diff --git a/iree/turbine/kernel/wave/constraints.py b/iree/turbine/kernel/wave/constraints.py
index faa211f77..96ac78db0 100644
--- a/iree/turbine/kernel/wave/constraints.py
+++ b/iree/turbine/kernel/wave/constraints.py
@@ -63,6 +63,7 @@ class MMAOperand(Enum):
     M = 0
     N = 1
     K = 2
+    B = 3
 
 
 @dataclass
@@ -180,6 +181,7 @@ def mma_index_offset(self, mma_type: Optional[MMAType]):
                     ),  # M
                     lane % 16,  # N
                     4 * floor(lane / 16),  # K
+                    0,  # B
                 ]
             case MMAType.F32_32x32x8_F16 | MMAType.I32_32x32x8_I8:
                 offset = [
@@ -194,6 +196,7 @@ def mma_index_offset(self, mma_type: Optional[MMAType]):
                     ),  # M
                     lane % 32,  # N
                     4 * floor(lane / 32),  # K
+                    0,  # B
                 ]
             case (
                 MMAType.F32_16x16x32_F8
@@ -207,6 +210,7 @@ def mma_index_offset(self, mma_type: Optional[MMAType]):
                     ),  # M
                     lane % 16,  # N
                     8 * floor(lane / 16),  # K
+                    0,  # B
                 ]
                 if mma_type == MMAType.F32_16x16x32_K4_F8:
                     offset = [
@@ -217,6 +221,7 @@ def mma_index_offset(self, mma_type: Optional[MMAType]):
                         (16 * floor(GPR_NUM / 4))
                         + 4 * floor(lane / 16)
                         + (GPR_NUM % 4),  # K
+                        0,  # B
                     ]
             case (
                 MMAType.F32_32x32x16_F8
@@ -236,6 +241,7 @@ def mma_index_offset(self, mma_type: Optional[MMAType]):
                     ),  # M
                     lane % 32,  # N
                     8 * floor(lane / 32),  # K
+                    0,  # B
                 ]
                 if mma_type == MMAType.F32_32x32x16_K4_F8:
                     offset = [
@@ -252,6 +258,7 @@ def mma_index_offset(self, mma_type: Optional[MMAType]):
                         (8 * floor(GPR_NUM / 4))
                         + 4 * floor(lane / 32)
                         + (GPR_NUM % 4),  # K
+                        0,  # B
                     ]
             case _:
                 raise ValueError("Unsupported MMA type")
@@ -313,22 +320,26 @@ def apply_mma_mapping(
                     Piecewise((1, ~MMA_ACC), (4, MMA_ACC)),  # M
                     1,  # N
                     4,  # K
+                    1,  # B
                 ]
                 stride = [
                     Piecewise((1, ~MMA_ACC), (16, MMA_ACC)),  # M
                     1,  # N
                     1,  # K
+                    1,  # B
                 ]
             case MMAType.F32_32x32x8_F16 | MMAType.I32_32x32x8_I8:
                 size = [
                     Piecewise((1, ~MMA_ACC), (16, MMA_ACC)),  # M
                     1,  # N
                     4,  # K
+                    1,  # B
                 ]
                 stride = [
                     Piecewise((1, ~MMA_ACC), (32, MMA_ACC)),  # M
                     1,  # N
                     1,  # K
+                    1,  # B
                 ]
             case (
                 MMAType.F32_16x16x32_F8
@@ -340,11 +351,13 @@ def apply_mma_mapping(
                     Piecewise((1, ~MMA_ACC), (4, MMA_ACC)),  # M
                     1,  # N
                     8,  # K
+                    1,  # B
                 ]
                 stride = [
                     Piecewise((1, ~MMA_ACC), (16, MMA_ACC)),  # M
                     1,  # N
                     1,  # K
+                    1,  # B
                 ]
             case (
                 MMAType.F32_32x32x16_F8
@@ -356,11 +369,13 @@ def apply_mma_mapping(
                     Piecewise((1, ~MMA_ACC), (16, MMA_ACC)),  # M
                     1,  # N
                     8,  # K
+                    1,  # B
                 ]
                 stride = [
                     Piecewise((1, ~MMA_ACC), (32, MMA_ACC)),  # M
                     1,  # N
                     1,  # K
+                    1,  # B
                 ]
             case _:
                 raise ValueError("Unsupported MMA type")
diff --git a/iree/turbine/kernel/wave/expansion/expansion.py b/iree/turbine/kernel/wave/expansion/expansion.py
index 88a695be0..5b6162a89 100644
--- a/iree/turbine/kernel/wave/expansion/expansion.py
+++ b/iree/turbine/kernel/wave/expansion/expansion.py
@@ -400,7 +400,6 @@ def add_get_results(trace: CapturedTrace):
             get_result = get_custom(
                 GetResult(reduction.fx_node, 0).add_to_graph(reduction.graph)
             )
-            get_result.vector_shapes = reduction.init_args[0].vector_shapes
             reduction.replace_all_uses_with_except(get_result, [get_result])
 
 
@@ -731,8 +730,6 @@ def expand_graph(
     to the root of the graph.
     """
 
-    add_get_results(trace)
-
     leaf_ops = [get_custom(node) for node in reversed(trace.walk(is_leaf_node))]
     if not leaf_ops:
         final_op = get_custom(trace.get_root_graph()._root.prev)
diff --git a/iree/turbine/kernel/wave/expansion/expansion_utils.py b/iree/turbine/kernel/wave/expansion/expansion_utils.py
index de0aada25..21d2ec8e0 100644
--- a/iree/turbine/kernel/wave/expansion/expansion_utils.py
+++ b/iree/turbine/kernel/wave/expansion/expansion_utils.py
@@ -90,8 +90,9 @@ def get_dim_scaling(
                 or (tile_size / wave_count) % vector_size != 0
             ):
                 raise ValueError(
-                    f"Tile size must be divisible by wave count and vector size, got: "
-                    f"tile_size={tile_size}, wave_count={wave_count}, vector_size={vector_size}"
+                    f"Tile size for dim {constraint.dim} must be divisible by "
+                    f"wave count and vector size, got: tile_size={tile_size}, "
+                    f"wave_count={wave_count}, vector_size={vector_size}"
                 )
             dim_scaling[constraint.dim] = tile_size // wave_count // vector_size
 
diff --git a/iree/turbine/kernel/wave/utils/general_utils.py b/iree/turbine/kernel/wave/utils/general_utils.py
index e2558bae0..01445bec2 100644
--- a/iree/turbine/kernel/wave/utils/general_utils.py
+++ b/iree/turbine/kernel/wave/utils/general_utils.py
@@ -297,7 +297,9 @@ def get_largest_index_and_size(
 
     sorted_values = sorted(
         [
+            # Call simplify_index to avoid comparing constants with sympy values.
             (i, dim, subs_idxc(index.size))
+            # (i, dim, simplify_index(subs_idxc(index.size)))
             for i, (dim, index) in enumerate(indices.items())
         ],
         # x[0] is the index of the dimension.
diff --git a/iree/turbine/kernel/wave/utils/mma_utils.py b/iree/turbine/kernel/wave/utils/mma_utils.py
index 3b6deaed1..1ce98daf5 100644
--- a/iree/turbine/kernel/wave/utils/mma_utils.py
+++ b/iree/turbine/kernel/wave/utils/mma_utils.py
@@ -82,6 +82,15 @@ def is_mma(node):
         mapping[custom][m] = MMAOperand.M
         mapping[custom][n] = MMAOperand.N
         mapping[custom][k] = MMAOperand.K
+        for b in set(acc_shape) - set(
+            (
+                m,
+                n,
+                k,
+            )
+        ):
+            mapping[custom][b] = MMAOperand.B
+
         custom.vector_shapes = {
             m: hardware_constraint.mma_matrix_shapes(custom.mma_type)[0],
             n: hardware_constraint.mma_matrix_shapes(custom.mma_type)[1],
@@ -190,6 +199,8 @@ def simplify_index(index: IndexExpr) -> IndexExpr:
     Simplifies the index by applying the following bindings:
         - MMA acc_index bindings so the index of the MMA node is the acc_index.
     """
+    if isinstance(index, int):
+        return index
     mapping = {MMA_LHS: 0, MMA_RHS: 0, MMA_ACC: 1}
     return subs_idxc(index.subs(mapping))
 
diff --git a/iree/turbine/kernel/wave/wave.py b/iree/turbine/kernel/wave/wave.py
index ec17be283..60a011e45 100644
--- a/iree/turbine/kernel/wave/wave.py
+++ b/iree/turbine/kernel/wave/wave.py
@@ -47,7 +47,7 @@
 from .compile_options import WaveCompileOptions
 from .decompose_reduce_ops import decompose_reduce_ops
 from .decompose_vmma_ops import decompose_vmma_ops
-from .expansion.expansion import expand_graph
+from .expansion.expansion import expand_graph, add_get_results
 from .global_to_shared_gathers import global_to_shared_gathers
 from .hoisting import hoist_loop_invariant_ops
 from .minimize_global_loads import minimize_global_loads
@@ -439,8 +439,7 @@ def compile_to_mlir(
     def build_initial_pass_pipeline(
         self,
         trace: CapturedTrace,
-        print_ir_before: Sequence[str] = [],
-        print_ir_after: Sequence[str] = [],
+        options: WaveCompileOptions,
     ):
         idxc = IndexingContext.current()
 
@@ -459,20 +458,106 @@ def substitute_vector_shapes():
             partial(self.initialize_workgroup_constraints, trace),
             finalize_indices,
             substitute_vector_shapes,
+            partial(add_get_results, trace),
             partial(infer_types, trace),
             partial(promote_placeholders, trace, self.constraints),
             partial(
                 set_node_indices,
                 trace,
                 self.constraints,
-                print_ir_before,
-                print_ir_after,
+                options.print_ir_before,
+                options.print_ir_after,
             ),
             partial(expand_graph, trace, self.constraints),
             partial(set_post_expansion_indices, trace, self.constraints),
             partial(remove_chained_getresult, trace),
         ]
 
+    def build_optimization_pass_pipeline(
+        self,
+        trace: CapturedTrace,
+        options: WaveCompileOptions,
+    ):
+        return [
+            partial(decompose_vmma_ops, trace, self.constraints),
+            partial(hoist_loop_invariant_ops, trace, self.constraints),
+            partial(global_to_shared_gathers, trace, self.constraints),
+            partial(minimize_global_loads, trace, self.constraints),
+            partial(reuse_shared_allocs, trace),
+            partial(apply_shared_memory_indexing_corrections, trace, self.constraints),
+        ]
+
+    def build_partitioning_pass_pipeline(
+        self,
+        trace: CapturedTrace,
+        options: WaveCompileOptions,
+    ):
+        return [
+            partial(partition_ops_with_gpr_offsets, trace, self.constraints),
+            partial(partition_strided_operators, trace, self.constraints),
+            partial(remove_chained_extractslice, trace),
+        ]
+
+    def build_reduction_pass_pipeline(
+        self,
+        trace: CapturedTrace,
+        options: WaveCompileOptions,
+    ):
+        # Schedule the reduction ops.
+        # Scheduling should always be used with use_scheduling_barriers=True,
+        # as this is the only way we can ensure that LLVM enforces our desired schedule.
+        # However, due a bug in LLVM, you will need to patch your local LLVM repo
+        # with the following commit: https://github.com/kerbowa/llvm-project/commit/ee52732cddae42deed2e3387a83b20ec05860b4e
+        # Specifically:
+        # git fetch https://github.com/kerbowa/llvm-project.git ee52732cddae42deed2e3387a83b20ec05860b4e
+        # git cherry-pick ee52732cddae42deed2e3387a83b20ec05860b4e
+        # [Manually resolve conflicts consistent with the PR]
+        return [
+            partial(decompose_reduce_ops, trace, self.constraints),
+            partial(
+                schedule_graph,
+                trace,
+                self.constraints,
+                options.use_scheduling_barriers,
+                options.schedule,
+            ),
+        ]
+
+    def build_shared_memory_pass_pipeline(
+        self,
+        trace: CapturedTrace,
+        options: WaveCompileOptions,
+    ):
+        return [
+            # Align sizes to WG/Tile sizes
+            # This pass changes indexing keys, which can interfere with other passes,
+            # so it should be called close to the end of pipeline.
+            partial(align_index_sizes, trace, self.constraints),
+            partial(add_shared_memory_barriers, trace),
+            partial(compute_shared_memory_usage, trace, options.kernel_launch_info),
+        ]
+
+    def build_full_pass_pipeline(
+        self,
+        trace: CapturedTrace,
+        options: WaveCompileOptions,
+    ):
+        # Initial passes, pre-optimization.
+        graph_passes = self.build_initial_pass_pipeline(trace, options)
+
+        # Optimizations.
+        graph_passes += self.build_optimization_pass_pipeline(trace, options)
+
+        # Partition strided operators.
+        graph_passes += self.build_partitioning_pass_pipeline(trace, options)
+
+        # Reduction decomposition and scheduling.
+        graph_passes += self.build_reduction_pass_pipeline(trace, options)
+
+        # Shared memory passes.
+        graph_passes += self.build_shared_memory_pass_pipeline(trace, options)
+        return graph_passes
+
     def _trace_and_get_kernel_signature(
         self,
         options: WaveCompileOptions,
@@ -515,62 +600,11 @@ def _trace_and_get_kernel_signature(
             print(f"***After trace/Before first pass***\n")
             print_trace(trace)
 
-        # Initial passes, pre-optimization.
-        graph_passes = self.build_initial_pass_pipeline(
-            trace, print_ir_before, print_ir_after
-        )
-
-        # Optimizations.
-        graph_passes += [
-            partial(decompose_vmma_ops, trace, self.constraints),
-            partial(hoist_loop_invariant_ops, trace, self.constraints),
-            partial(global_to_shared_gathers, trace, self.constraints),
-            partial(minimize_global_loads, trace, self.constraints),
-            partial(reuse_shared_allocs, trace),
-            partial(apply_shared_memory_indexing_corrections, trace, self.constraints),
-        ]
-
-        # Partition strided operators.
-        graph_passes += [
-            partial(partition_ops_with_gpr_offsets, trace, self.constraints),
-            partial(partition_strided_operators, trace, self.constraints),
-            partial(remove_chained_extractslice, trace),
-        ]
-
-        graph_passes += [partial(decompose_reduce_ops, trace, self.constraints)]
-
-        # Schedule the reduction ops.
-        # Scheduling should always be used with use_scheduling_barriers=True,
-        # as this is the only way we can ensure that LLVM enforces our desired schedule.
-        # However, due a bug in LLVM, you will need to patch your local LLVM repo
-        # with the following commit: https://github.com/kerbowa/llvm-project/commit/ee52732cddae42deed2e3387a83b20ec05860b4e
-        # Specifically:
-        # git fetch https://github.com/kerbowa/llvm-project.git ee52732cddae42deed2e3387a83b20ec05860b4e
-        # git cherry-pick ee52732cddae42deed2e3387a83b20ec05860b4e
-        # [Manually resolve conflicts consistent with the PR]
-        scheduling_type = options.schedule
-        use_scheduling_barriers = options.use_scheduling_barriers
-        graph_passes.append(
-            partial(
-                schedule_graph,
-                trace,
-                self.constraints,
-                use_scheduling_barriers,
-                scheduling_type,
-            )
-        )
-
-        graph_passes += [
-            # Align sizes to WG/Tile sizes
-            # This pass changes indexing keys, which can interfere with other passes,
-            # so it should be called close to the end of pipeline.
-            partial(align_index_sizes, trace, self.constraints),
-            partial(add_shared_memory_barriers, trace),
-            partial(compute_shared_memory_usage, trace, options.kernel_launch_info),
-        ]
+        # Create the pass pipeline.
+        graph_passes = self.build_full_pass_pipeline(trace, options)
 
         for p in graph_passes:
-            try_apply_pass(p, trace, print_ir_before, print_ir_after)
+            try_apply_pass(p, trace, options.print_ir_before, options.print_ir_after)
 
         if "all" in print_ir_after or "last" in print_ir_after:
             # Take advantage of Python leaking loop variables
diff --git a/lit_tests/kernel/wave/barriers.py b/lit_tests/kernel/wave/barriers.py
index 6a729f1eb..2ae970f19 100644
--- a/lit_tests/kernel/wave/barriers.py
+++ b/lit_tests/kernel/wave/barriers.py
@@ -11,7 +11,7 @@
 from iree.turbine.kernel.wave.promotion import promote_node, promote_placeholders
 from iree.turbine.kernel.wave.barriers import add_shared_memory_barriers
 from iree.turbine.kernel.wave.hoisting import hoist_loop_invariant_ops
-from iree.turbine.kernel.wave.expansion.expansion import expand_graph
+from iree.turbine.kernel.wave.expansion.expansion import expand_graph, add_get_results
 from iree.turbine.kernel.wave.type_inference import infer_types
 from iree.turbine.kernel.lang.global_symbols import *
 from iree.turbine.kernel._support.tracing import CapturedTrace
@@ -93,6 +93,7 @@ def test_read_write_equal_sizes():
         graph: fx.Graph = trace.get_root_graph()
         read_node = get_read_nodes(graph)[0]
         IndexingContext.current().finalize()
+        add_get_results(trace)
         infer_types(trace)
         promote_node(read_node, None, SHARED_ADDRESS_SPACE, constraints)
         set_node_indices(trace, constraints)
@@ -180,6 +181,7 @@ def test_gemm():
         graph: fx.Graph = trace.get_subgraph("region_0")
         IndexingContext.current().finalize()
         initialize_iter_args(trace)
+        add_get_results(trace)
         infer_types(trace)
         read_nodes = get_read_nodes(graph)
         for read_node in read_nodes:
diff --git a/lit_tests/kernel/wave/expansion.py b/lit_tests/kernel/wave/expansion.py
index 574c65f35..663a56ad4 100644
--- a/lit_tests/kernel/wave/expansion.py
+++ b/lit_tests/kernel/wave/expansion.py
@@ -4,7 +4,7 @@
 import iree.turbine.kernel as tk
 import iree.turbine.kernel.lang as tkl
 import iree.turbine.kernel.wave as tkw
-from iree.turbine.kernel.wave.expansion.expansion import expand_graph
+from iree.turbine.kernel.wave.expansion.expansion import expand_graph, add_get_results
 from iree.turbine.kernel.wave.type_inference import infer_types
 from iree.turbine.kernel.wave.analysis.index_sequence_analysis import (
     set_node_indices,
@@ -77,6 +77,7 @@ def test_read_write_equal_sizes():
     ):
         graph = read_write_same_size()
         IndexingContext.current().finalize()
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -159,6 +160,7 @@ def test_read_write():
     ):
         graph = read_write_different_dims()
         IndexingContext.current().finalize()
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -241,6 +243,7 @@ def test_write_in_iterate():
         graph = write_in_iterate()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -323,6 +326,7 @@ def test_no_writes():
         graph = no_writes()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -363,6 +367,7 @@ def test_gemm():
         graph = gemm()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -554,6 +559,7 @@ def test_batched_gemm():
         graph = batched_gemm()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -736,6 +742,7 @@ def test_gemm_non_direct_acc():
         graph = gemm_non_direct_acc()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -804,6 +811,7 @@ def test_tiled_max():
         graph = tiled_max()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -837,6 +845,7 @@ def test_gemm_iterate_expansion_only():
         graph = gemm()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -983,6 +992,7 @@ def test_attention():
         graph = attention()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
@@ -1079,6 +1089,7 @@ def py_arithmetic_different_dims():
         graph = py_arithmetic_different_dims()
         IndexingContext.current().finalize()
         infer_types(graph)
+        add_get_results(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
         set_post_expansion_indices(graph, constraints)
@@ -1185,6 +1196,7 @@ def test_chained_gemm_32x32x8():
         graph = chained_gemm_32x32x8()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
diff --git a/lit_tests/kernel/wave/gemm.py b/lit_tests/kernel/wave/gemm.py
index c5d72fbe5..6fa2dc559 100644
--- a/lit_tests/kernel/wave/gemm.py
+++ b/lit_tests/kernel/wave/gemm.py
@@ -132,6 +132,78 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
     # CHECK-COUNT-4:    vector.store
 
 
+@run_test
+def test_gemm_bias():
+    constraints: list[tkw.Constraint] = [tkw.WorkgroupConstraint(M, BLOCK_M, 0)]
+    constraints += [tkw.WorkgroupConstraint(N, BLOCK_N, 1)]
+    constraints += [tkw.TilingConstraint(K, BLOCK_K)]
+    constraints += [tkw.WaveConstraint(M, BLOCK_M / 2)]
+    constraints += [tkw.WaveConstraint(N, BLOCK_N / 2)]
+
+    constraints += [
+        tkw.HardwareConstraint(
+            threads_per_wave=64,
+            waves_per_block=(2, 2, 1),
+            mma_type=tkw.MMAType.F32_16x16x16_F16,
+        )
+    ]
+
+    @tkw.wave(constraints)
+    def gemm_bias(
+        a: tkl.Memory[M, K, ADDRESS_SPACE, tkl.f16],
+        b: tkl.Memory[N, K, ADDRESS_SPACE, tkl.f16],
+        bias: tkl.Memory[M, N, ADDRESS_SPACE_0, tkl.f32],
+        c: tkl.Memory[M, N, ADDRESS_SPACE_0, tkl.f32],
+    ):
+        c_reg = tkl.Register[M, N, tkl.f32](0.0)
+
+        @tkw.reduction(K, init_args=[c_reg])
+        def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
+            a_reg = tkw.read(a)
+            b_reg = tkw.read(b)
+            acc = tkw.mma(a_reg, b_reg, acc)
+            return acc
+
+        bias_reg = tkw.read(bias)
+        result = repeat + bias_reg
+        tkw.write(result, c)
+
+    options = WaveCompileOptions(
+        subs={
+            M: 64,
+            N: 128,
+            K: 64,
+            BLOCK_M: 32,
+            BLOCK_N: 32,
+            BLOCK_K: 16,
+            LOAD_ELEMS_PER_THREAD: 4,
+            STORE_ELEMS_PER_THREAD: 4,
+            ADDRESS_SPACE: SHARED_ADDRESS_SPACE,
+            ADDRESS_SPACE_0: GLOBAL_ADDRESS_SPACE,
+        },
+        canonicalize=True,
+        compile_to_mlir=True,
+    )
+    gemm_bias = wave_compile(options, gemm_bias)
+    print(gemm_bias.asm)
+
+    # CHECK-LABEL:     func.func @gemm_bias
+    # CHECK:             scf.for
+    # CHECK-COUNT-1:       vector.load
+    # CHECK-COUNT-1:       vector.store
+    # CHECK-COUNT-1:       vector.load
+    # CHECK-COUNT-1:       vector.store
+    # CHECK-COUNT-2:       vector.load
+    # CHECK-COUNT-1:       amdgpu.mfma
+    # CHECK:             scf.yield
+    # Load transposed bias from global
+    # CHECK-COUNT-4:       vector.load
+    # Add result to bias
+    # CHECK-COUNT-1:       arith.addf
+    # Store result
+    # CHECK-COUNT-4:       vector.store
+
+
 @run_test
 def test_cdna2_int_gemm():
     constraints: list[tkw.Constraint] = [tkw.WorkgroupConstraint(M, BLOCK_M, 0)]
diff --git a/lit_tests/kernel/wave/index_sequence_analysis.py b/lit_tests/kernel/wave/index_sequence_analysis.py
index e68cc60f0..805f1abb8 100644
--- a/lit_tests/kernel/wave/index_sequence_analysis.py
+++ b/lit_tests/kernel/wave/index_sequence_analysis.py
@@ -6,7 +6,7 @@
 import iree.turbine.kernel.wave as tkw
 from iree.turbine.kernel.wave.promotion import promote_placeholders
 from iree.turbine.kernel.wave.hoisting import hoist_loop_invariant_ops
-from iree.turbine.kernel.wave.expansion.expansion import expand_graph
+from iree.turbine.kernel.wave.expansion.expansion import expand_graph, add_get_results
 from iree.turbine.kernel.wave.type_inference import infer_types
 from iree.turbine.kernel.lang.global_symbols import *
 from iree.turbine.kernel._support.tracing import CapturedTrace
@@ -95,6 +95,7 @@ def test_gemm():
         trace: CapturedTrace = gemm()
         IndexingContext.current().finalize()
         initialize_iter_args(trace)
+        add_get_results(trace)
         infer_types(trace)
         promote_placeholders(trace, constraints)
         set_node_indices(trace, constraints)
diff --git a/lit_tests/kernel/wave/minimize_global_loads.py b/lit_tests/kernel/wave/minimize_global_loads.py
index f6b436c30..49b63b18f 100644
--- a/lit_tests/kernel/wave/minimize_global_loads.py
+++ b/lit_tests/kernel/wave/minimize_global_loads.py
@@ -7,7 +7,7 @@
 from iree.turbine.kernel.wave.promotion import promote_placeholders
 from iree.turbine.kernel.wave.hoisting import hoist_loop_invariant_ops
 from iree.turbine.kernel.wave.barriers import add_shared_memory_barriers
-from iree.turbine.kernel.wave.expansion.expansion import expand_graph
+from iree.turbine.kernel.wave.expansion.expansion import expand_graph, add_get_results
 from iree.turbine.kernel.wave.type_inference import infer_types
 from iree.turbine.kernel.lang.global_symbols import *
 from iree.turbine.kernel._support.tracing import CapturedTrace
@@ -95,6 +95,7 @@ def test_gemm():
         visualize = False
         IndexingContext.current().finalize()
         initialize_iter_args(trace)
+        add_get_results(trace)
         infer_types(trace)
         promote_placeholders(trace, constraints)
         set_node_indices(trace, constraints)
diff --git a/lit_tests/kernel/wave/moe.py b/lit_tests/kernel/wave/moe.py
new file mode 100644
index 000000000..a9db92cb5
--- /dev/null
+++ b/lit_tests/kernel/wave/moe.py
@@ -0,0 +1,452 @@
+# RUN: python %s | FileCheck %s
+
+import copy
+import logging
+from typing import Sequence
+
+import torch
+import torch.nn.functional as F
+from torch.testing import assert_close
+
+import iree.turbine.kernel as tk
+import iree.turbine.kernel.lang as tkl
+import iree.turbine.kernel.wave as tkw
+from iree.turbine.kernel.compiler.ir import Context, Location, Module
+from iree.turbine.kernel.wave.type_inference import infer_types
+from iree.turbine.kernel.wave.wave import LaunchableWave
+from iree.turbine.kernel.lang.global_symbols import *
+from iree.turbine.kernel._support.indexing import IndexingContext
+from iree.turbine.kernel._support.tracing import CapturedTrace
+from iree.turbine.kernel.wave.compile import WaveCompileOptions, wave_compile
+from iree.turbine.kernel.wave.utils.mma_utils import (
+    get_mfma_load_elems_per_thread,
+    get_mfma_store_elems_per_thread,
+)
+from iree.turbine.kernel.wave.utils.general_utils import (
+    get_default_scheduling_params,
+    run_test,
+)
+from iree.turbine.kernel.wave.utils.print_utils import print_trace, try_apply_pass
+from iree.turbine.kernel.wave.utils.compile_utils import (
+    set_default_compile_config,
+)
+from iree.turbine.kernel.wave.utils.run_utils import (
+    set_default_run_config,
+)
+
+DOC = """
+# MoE TCTK
+This is a TCTK version of the MoE kernel. It is looking at the problem from first principles.
+
+Let's have fun with this one
+
+# Torch Naive MoE
+# https://github.com/sgl-project/sglang/blob/7e3bb5270524d38bd98b93a22441fa693c3fa64c/test/srt/test_fused_moe.py#L47
+# Triton MoE
+# https://github.com/sgl-project/sglang/blob/ffa1b3e318c9d1342a5e430eb04df609e22a3775/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py#L1159
+    def torch_naive_moe(self, a, w1, w2, score, topk):
+        # where
+        # a: Tensor[B, D1],
+        # W1: Tensor[E, 2 * N, D1],
+        # W2: Tensor[E, D2, N],
+        # score: Tensor[B, E], topk: int
+        # out: Tensor[B, TOPK, D2]
+        B, D = a.shape
+        a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+        out = torch.zeros(B * topk,
+                          w2.shape[1],
+                          dtype=a.dtype,
+                          device=a.device)
+        score = torch.softmax(score, dim=-1, dtype=torch.float32)
+        topk_weights, topk_ids = torch.topk(score, topk)
+        topk_weights = topk_weights.view(-1)
+        topk_ids = topk_ids.view(-1)
+        for i in range(w1.shape[0]):
+            mask = topk_ids == i
+            if mask.sum():
+                out[mask] = silu_and_mul(
+                    a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1)
+        return (out.view(B, -1, w2.shape[1]) *
+                topk_weights.view(B, -1, 1).to(out.dtype)).sum(dim=1)
+
+## TCTK notation
+
+# score: Tensor[B, E],
+topk_score[B, TOPK], topk_idx[B, TOPK] = \
+    torch.topk(score[B, E])
+
+# Note: softmax normalization only needed on the TOPK subset
+m_s[B] max= score[B, :]
+score[B, :] -= m_s[B] # for softmax stability
+sum_exp[B] += torch.exp(score[B, :]) # for actual softmax normalization
+topk_score[B, TOPK] = exp(topk_score[B, TOPK] - m_s[B]) / sum_exp[B]
+
+TMP[TOPK, B, 2 * N] = A[B, D1] @ W1[subset(E by topk_idx[B, TOPK]), 2 * N, D1].transpose(0, 1)
+TMP_2[TOPK, B, N] = SILU(TMP[TOPK, B, :N]) * TMP[TOPK, B, N:]
+TMP_3[TOPK, B, D2] = TMP_2[TOPK, B, N] @ W2[subset(E by topk_idx[B, TOPK]), D2, N].transpose(0, 1)
+
+Instead we'll use TOPK, B, D2 for the output tensor because TK does not like
+1-broadcast dimensions in another form.
+RESULT[B, TOPK, D2] = TMP_3[TOPK, B, D2] * topk_score[B, TOPK]
+"""
+
+# Static shapes
+vB = 16
+vN = 32
+vD1 = 64
+vD2 = 128
+vTOPK = 5
+vE = 256
+
+# Symbols
+B, N, D1, D2, TOPK = (
+    tkl.sym.B,
+    tkl.sym.N,
+    tkl.sym.D1,
+    tkl.sym.D2,
+    tkl.sym.TOPK,
+)
+
+BLOCK_B, BLOCK_N, BLOCK_D1, BLOCK_D2, BLOCK_TOPK = (
+    tkl.sym.BLOCK_B,
+    tkl.sym.BLOCK_N,
+    tkl.sym.BLOCK_D1,
+    tkl.sym.BLOCK_D2,
+    tkl.sym.BLOCK_TOPK,
+)
+
+LOAD_ELEMS_PER_THREAD, STORE_ELEMS_PER_THREAD = (
+    tkl.sym.LOAD_TOKS_PER_THREAD,
+    tkl.sym.STORE_ELEMS_PER_THREAD,
+)
+
+ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE
+
+
+def lit_harness(build_constraints_fun, get_kernel_fun, test_config):
+    with tk.gen.TestLaunchContext(
+        test_config["static_symbols"] if "static_symbols" in test_config else {}
+    ):
+        executable_kernel = get_kernel_fun(build_constraints_fun(test_config))
+        trace: CapturedTrace = executable_kernel._trace()
+        options = WaveCompileOptions(
+            subs=test_config["static_symbols"]
+            if "static_symbols" in test_config
+            else {},
+        )
+        options.print_ir_before = ["all"]
+        idxc: IndexingContext = IndexingContext.current()
+        graph_passes = executable_kernel.build_full_pass_pipeline(trace, options)
+        for p in graph_passes:
+            try_apply_pass(p, trace, print_ir_before=["all"])
+
+        executable_kernel.infer_grid_shape(idxc)
+
+        options = set_default_compile_config(options)
+        options.canonicalize = True
+        with Context() as context:
+            (
+                mb,
+                trace,
+                exe,
+                kernel_sig,
+                entrypoint_name,
+            ) = executable_kernel.compile_to_mlir(trace, context, options=options)
+            print(mb.module_op)
+
+
+def build_block_constraints(test_config) -> Sequence[tkw.Constraint]:
+    constraints: list[tkw.Constraint] = []
+    constraints += [
+        tkw.WorkgroupConstraint(B, BLOCK_B, 2),
+        tkw.WaveConstraint(B, 1),
+    ]
+    constraints += [
+        tkw.WorkgroupConstraint(TOPK, BLOCK_TOPK, 1),
+        tkw.WaveConstraint(TOPK, 1),
+    ]
+    constraints += [
+        tkw.WorkgroupConstraint(D2, BLOCK_D2, 0),
+        tkw.WaveConstraint(D2, 1),
+    ]
+    constraints += [tkw.TilingConstraint(N, BLOCK_N)]
+
+    constraints += [
+        tkw.HardwareConstraint(
+            waves_per_block=(1, 1, 1),
+            threads_per_wave=64,
+            # One must always specify mma_type or vector_shapes.
+            mma_type=tkw.MMAType.F32_16x16x16_F16,
+            vector_shapes=test_config["vector_shapes"]
+            if "vector_shapes" in test_config
+            else {},
+        )
+    ]
+    return constraints
+
+
+def create_test_config(mma_variant: tkw.MMAType = tkw.MMAType.F32_16x16x16_F16):
+    # fmt: off
+    return {
+        "static_symbols": {
+            ### Problem sizes.
+            N: vN,
+            # D1: vD1,
+            # z, y, x
+            B:       vB, # z
+            TOPK: vTOPK, # y
+            D2:     vD2, # x
+            ### Block sizes.
+            # tiling
+            BLOCK_N: 16,
+            # z, y, x
+            BLOCK_B:   16, # z
+            BLOCK_TOPK: 1, # y
+            BLOCK_D2:  16, # x
+            ### L/S sizes (ideally omitted).
+            LOAD_ELEMS_PER_THREAD: get_mfma_load_elems_per_thread(mma_variant),
+            STORE_ELEMS_PER_THREAD: get_mfma_store_elems_per_thread(mma_variant),
+        },
+        # Need to specify vector_shape explicitly because somehow this does
+        # not get propagated.
+        "vector_shapes": {
+            # N: 1,
+            # D1: 16,  # TODO: connected to MFMA op type
+            # z, y, x
+            B:    16, # z
+            # Warning: it is necessary to set TOPK == 0 for proper batch
+            # dimension processing, otherwise there is a risk of miscompile.
+            TOPK:  0, # y
+            D2:   16, # x
+        },
+        "canonicalize": {True},
+    }
+    # fmt: on
+
+
+# fmt: off
+# Note: W2 really has torch.Tensor.shape [E, D2, N] but we want to index it with
+# indices [B, TOPK, D2, N].
+# We don't want to introduce index E because we'd get the cartesian product
+# [E, B, TOPK], which is not what we want.
+# So we just use TOPK to index into W2 (alternatively we could use B).
+w2_layout = tkl.MemoryLayout(shape=(vE, D2, N,))
+
+### Current
+x, y, z, k = [tkw.IndexMapping.iterator(i) for i in range(4)]
+d0 = tkw.IndexMapping.dynamic_val(0)
+offset_mapping_w2 = tkw.IndexMapping(
+    num_iterators=4,
+    inputs= {D2: x, TOPK: d0, B: z, N: k},
+    outputs={D2: x, TOPK:  y, B: z, N: k},
+    # offset_mapping_w2 is fed by:
+    #    `expert_id = tkw.read(TOPK_IDS) # : [B, TOPK]`
+    # where expert_id is d0.
+    # We need to match d0 to its indirect read `tkw.read(TOPK_IDS) # : [B, TOPK]`
+    dynamic_val_mappings={B: z, TOPK: y},
+)
+
+# ### The following sytax would be preferred.
+# x, y, z, k = [tkw.IndexMapping.iterator(i) for i in range(4)]
+# d0 = tkw.IndexMapping.dynamic_val(0)
+# offset_mapping_w2 = tkw.IndexMapping(
+#     shape     = [D2, TOPK, B, N],
+#     iterators = [ x,    y, z, k],
+#     inputs    = [ x,   d0, z, k],
+#     output    = [ x,    y, z, k],
+#     # Match d0 to its indirect read `tkw.read(TOPK_IDS) # : [B, TOPK]` which is
+#     # indexed by [z, y]
+#     dynamic_val_mappings = {d0 : [z, y]},
+# )
+# fmt: on
+
+
+def fused_moe_kernel(
+    TMP_2: tkl.Memory[B, TOPK, N, ADDRESS_SPACE, tkl.f16],
+    W2: tkl.Memory[TOPK, D2, N, ADDRESS_SPACE, tkl.f16, w2_layout],
+    TOPK_IDS: tkl.Memory[B, TOPK, ADDRESS_SPACE, tkl.i64],
+    TOPK_WEIGHTS: tkl.Memory[B, TOPK, ADDRESS_SPACE, tkl.f32],
+    RESULT: tkl.Memory[B, TOPK, D2, ADDRESS_SPACE, tkl.f32],
+):
+    res_reg = tkl.Register[TOPK, B, D2, tkl.f32](0.0)
+
+    # fmt: off
+    @tkw.iterate(N, init_args=[res_reg])
+    def repeat(acc: tkl.Register[TOPK, B, D2, tkl.f32]) -> tkl.Register[TOPK, B, D2, tkl.f32]:
+        ###
+        # TMP_3[TOPK, B, D2] = TMP_2[TOPK, B, N:]
+        #   @ W2[subset(E by topk_idx[B, TOPK]), D2, N].transpose(0, 1)
+        ###
+        # elements_per_thread=LOAD_ELEMS_PER_THREAD, automatically derived by the
+        # system from the mma op for all reads below.
+        tmp_2_reg = tkw.read(TMP_2,)                            # : [TOPK, B, N]
+        expert_id = tkw.read(TOPK_IDS,)                         # : [B, TOPK]
+        w2_reg = tkw.read(
+            W2,
+            mapping=offset_mapping_w2,
+            mapping_dynamic_vals=(expert_id,),
+        )                                                       # : [TOPK, D2, N] but indexed as [E=(B, TOPK), D2, N] and E expert_id
+        acc = tkw.mma(tmp_2_reg, w2_reg, acc)                   # : [TOPK, B, N] * [TOPK, D2, N] -> [TOPK, B, D2]
+        return acc
+
+    res = repeat
+
+    ### res combined with read and automatic mma inference does not work.
+    ### errors out with "ValueError: index 16 is out of bounds for array with size 16"
+    ### Instead roundtrip through memory.
+    ### elements_per_thread=STORE_ELEMS_PER_THREAD, automatically derived by the
+    ### system from the mma op above.
+    res = tkw.permute(res, target_shape=[B, TOPK, D2])                      # : [B, TOPK, D2]
+    tkw.write(res, RESULT,)                                                 # : [B, TOPK, D2]
+
+    # ###
+    # # RESULT[B, TOPK, D2] = TMP_3[TOPK, B, D2] * topk_score[B, TOPK]
+    # ###
+    # ### Elementwise part, use elements_per_thread=4
+    # res = tkw.read(RESULT, elements_per_thread=4,)                          # : [B, TOPK, D2]
+    # topk_weights = tkw.read(TOPK_WEIGHTS, elements_per_thread=4,)           # : [B, TOPK]
+    # topk_weights = tkw.broadcast(topk_weights, target_shape=[B, TOPK, D2])  # : [B, TOPK, D2]
+    # res = res * tkw.cast(topk_weights, tkl.f32)                             # : [B, TOPK, D2]
+    # tkw.write(res, RESULT, elements_per_thread=4,)                          # : [B, TOPK, D2]
+    # # fmt: on
+
+
+# Note: W2 really has torch.Tensor.shape [E, D2, N] but we want to index it
+# with indices [B, TOPK, D2, N].
+# We don't want to introduce index E because we'd get the cartesian product
+# [E, B, TOPK], which is not what we want.
+# So we just use TOPK to index into W2 (alternatively we could use B).
+def get_fused_moe_kernel(constraints):
+    @tkw.wave(constraints)
+    def fused_moe_kernel_executable(
+        TMP_2: tkl.Memory[B, TOPK, N, ADDRESS_SPACE, tkl.f16],
+        W2: tkl.Memory[TOPK, D2, N, ADDRESS_SPACE, tkl.f16, w2_layout],
+        TOPK_IDS: tkl.Memory[B, TOPK, ADDRESS_SPACE, tkl.i64],
+        TOPK_WEIGHTS: tkl.Memory[B, TOPK, ADDRESS_SPACE, tkl.f32],
+        RESULT: tkl.Memory[B, TOPK, D2, ADDRESS_SPACE, tkl.f32],
+    ):
+        return fused_moe_kernel(TMP_2, W2, TOPK_IDS, TOPK_WEIGHTS, RESULT)
+
+    return fused_moe_kernel_executable
+
+
+def silu_and_mul(x: torch.Tensor):
+    d = x.shape[-1] // 2
+    return F.silu(x[..., :d]) * x[..., d:]
+
+
+def torch_naive_moe(a, w1, w2, score, topk, result_dtype):
+    # where
+    # a: Tensor[B, D1],
+    # W1: Tensor[E, 2 * N, D1],
+    # W2: Tensor[E, D2, N],
+    # score: Tensor[B, E], topk: int
+    # out: Tensor[B, TOPK, D2]
+    B, D = a.shape
+    a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+    out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score, topk)
+    topk_weights = topk_weights.view(-1)
+    topk_ids = topk_ids.view(-1)
+    for i in range(w1.shape[0]):
+        mask = topk_ids == i
+        if mask.sum():
+            out[mask] = silu_and_mul(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
+                0, 1
+            )
+    return out.view(B, -1, w2.shape[1]) * topk_weights.view(B, -1, 1).to(result_dtype)
+    # .sum(dim=1)
+
+
+def torch_naive_moe_step_1(a, w1, w2, score, topk):
+    # where
+    # a: Tensor[B, D1],
+    # W1: Tensor[E, 2 * N, D1],
+    # W2: Tensor[E, D2, N],
+    # score: Tensor[B, E], topk: int
+    # out: Tensor[B, TOPK, D2]
+    B, D = a.shape
+    a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+    out_tmp = torch.zeros(B * topk, w1.shape[1] // 2, dtype=a.dtype, device=a.device)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weights, topk_ids = torch.topk(score, topk)
+    topk_weights = topk_weights.view(-1)
+    topk_ids = topk_ids.view(-1)
+    for i in range(w1.shape[0]):
+        mask = topk_ids == i
+        if mask.sum():
+            out_tmp[mask] = silu_and_mul(a[mask] @ w1[i].transpose(0, 1))
+    return topk_weights, topk_ids, out_tmp
+
+
+def torch_naive_moe_step_2(w1, w2, topk_weights, topk_ids, out_tmp):
+    out = torch.zeros(
+        out_tmp.shape[0], w2.shape[1], dtype=out_tmp.dtype, device=out_tmp.device
+    )
+    for i in range(w1.shape[0]):
+        mask = topk_ids == i
+        if mask.sum():
+            out[mask] = out_tmp[mask] @ w2[i].transpose(0, 1)
+    return out
+
+
+def torch_naive_moe_step_3(topk_weights, out, result_dtype):
+    return out * topk_weights.view(out.shape[0], -1, 1).to(result_dtype)
+    # .sum(dim=1)
+
+
+if __name__ == "__main__":
+
+    @run_test
+    def static_correct_1():
+        test_config = copy.deepcopy(create_test_config())
+        lit_harness(build_block_constraints, get_fused_moe_kernel, test_config)
+
+        a = torch.randn(vB, vD1, dtype=torch.float16).cuda()
+        w1 = torch.randn(vE, 2 * vN, vD1, dtype=torch.float16).cuda()
+        w2 = torch.randn(vE, vD2, vN, dtype=torch.float16).cuda()
+        score = torch.randn(vB, vE).cuda()
+        topk = vTOPK
+        reference = torch_naive_moe(a, w1, w2, score, topk, result_dtype=torch.float32)
+
+        cmp_params = dict(atol=3e-3, rtol=3e-3, check_dtype=False)
+        ref_full_1 = torch_naive_moe(a, w1, w2, score, topk, result_dtype=torch.float32)
+        topk_weights, topk_ids, out_tmp = torch_naive_moe_step_1(a, w1, w2, score, topk)
+        out_tmp_2 = torch_naive_moe_step_2(w1, w2, topk_weights, topk_ids, out_tmp)
+        # Reshape out_tmp_2 from 2-D to 3-D.
+        out_tmp_2 = out_tmp_2.view(vB, -1, w2.shape[1])
+        ref_full_2 = torch_naive_moe_step_3(
+            topk_weights, out_tmp_2, result_dtype=torch.float32
+        )
+        assert_close(ref_full_1, ref_full_2, **cmp_params)
+
+        from torch.profiler import profile, ProfilerActivity
+
+        with profile(activities=[ProfilerActivity.CUDA]) as prof:
+            test_config = create_test_config()
+            executable_kernel = get_fused_moe_kernel(
+                build_block_constraints(test_config)
+            )
+
+            options = WaveCompileOptions(
+                subs=test_config["static_symbols"]
+                if "static_symbols" in test_config
+                else {},
+            )
+            options = set_default_compile_config(options)
+            options.canonicalize = True
+            options = set_default_run_config(options)
+            executable_kernel = wave_compile(options, executable_kernel)
+
+            result = torch.zeros_like(ref_full_1).cuda()
+            executable_kernel(
+                out_tmp.view(vB, vTOPK, vN),
+                w2,
+                topk_ids.view(vB, vTOPK),
+                topk_weights.view(vB, vTOPK),
+                result,
+            )
+            assert_close(result, out_tmp_2, **cmp_params)
+
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
diff --git a/lit_tests/kernel/wave/scheduling.py b/lit_tests/kernel/wave/scheduling.py
index 014b420ec..03a8fa616 100644
--- a/lit_tests/kernel/wave/scheduling.py
+++ b/lit_tests/kernel/wave/scheduling.py
@@ -6,7 +6,7 @@
 import iree.turbine.kernel.wave as tkw
 from iree.turbine.kernel.wave.promotion import promote_placeholders
 from iree.turbine.kernel.wave.hoisting import hoist_loop_invariant_ops
-from iree.turbine.kernel.wave.expansion.expansion import expand_graph
+from iree.turbine.kernel.wave.expansion.expansion import expand_graph, add_get_results
 from iree.turbine.kernel.wave.type_inference import infer_types
 from iree.turbine.kernel.wave.scheduling.schedule import SchedulingType
 from iree.turbine.kernel.lang.global_symbols import *
@@ -100,6 +100,7 @@ def test_gemm_pipelined():
         trace: CapturedTrace = gemm_pipelined()
         IndexingContext.current().finalize()
         initialize_iter_args(trace)
+        add_get_results(trace)
         infer_types(trace)
         promote_placeholders(trace, constraints)
         set_node_indices(trace, constraints)
diff --git a/tests/kernel/wave/scheduling_test.py b/tests/kernel/wave/scheduling_test.py
index d2cb6ff49..073b344c2 100644
--- a/tests/kernel/wave/scheduling_test.py
+++ b/tests/kernel/wave/scheduling_test.py
@@ -29,7 +29,7 @@
 from iree.turbine.kernel._support.indexing import IndexingContext
 from iree.turbine.kernel.wave.promotion import promote_placeholders
 from iree.turbine.kernel.wave.hoisting import hoist_loop_invariant_ops
-from iree.turbine.kernel.wave.expansion.expansion import expand_graph
+from iree.turbine.kernel.wave.expansion.expansion import expand_graph, add_get_results
 from iree.turbine.kernel.wave.type_inference import infer_types
 from iree.turbine.kernel.wave.minimize_global_loads import minimize_global_loads
 from iree.turbine.kernel.wave.scheduling.schedule import schedule_graph
@@ -285,6 +285,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
             trace: CapturedTrace = gemm()
             IndexingContext.current().finalize()
             initialize_iter_args(trace)
+            add_get_results(trace)
             infer_types(trace)
             promote_placeholders(trace, constraints)
             hoist_loop_invariant_ops(trace, constraints)
diff --git a/tests/kernel/wave/visualization_test.py b/tests/kernel/wave/visualization_test.py
index d14ab5ce5..ebff1b28b 100644
--- a/tests/kernel/wave/visualization_test.py
+++ b/tests/kernel/wave/visualization_test.py
@@ -12,7 +12,7 @@
 import iree.turbine.kernel as tk
 import iree.turbine.kernel.lang as tkl
 import iree.turbine.kernel.wave as tkw
-from iree.turbine.kernel.wave.expansion.expansion import expand_graph
+from iree.turbine.kernel.wave.expansion.expansion import expand_graph, add_get_results
 from iree.turbine.kernel.wave.type_inference import infer_types
 from iree.turbine.kernel._support.tracing import CapturedTrace
 from iree.turbine.kernel._support.indexing import IndexingContext
@@ -96,6 +96,7 @@ def test_gemm():
         graph = gemm()
         IndexingContext.current().finalize()
         initialize_iter_args(graph)
+        add_get_results(graph)
         infer_types(graph)
         set_node_indices(graph, constraints)
         expand_graph(graph, constraints)
diff --git a/tests/kernel/wave/wave_gemm_test.py b/tests/kernel/wave/wave_gemm_test.py
index efa0982b7..0c81d3fa5 100644
--- a/tests/kernel/wave/wave_gemm_test.py
+++ b/tests/kernel/wave/wave_gemm_test.py
@@ -360,6 +360,147 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
     assert_close(c, iree_ref, atol=2e-4, rtol=3e-4, check_device=False)
 
 
+@require_e2e
+@pytest.mark.parametrize("shape", get_test_shapes("test_gemm"))
+@pytest.mark.parametrize("enable_scheduling", [SchedulingType.NONE])
+@param_bool("dynamic_dims", "dyn")
+@pytest.mark.parametrize(
+    "mfma_variant",
+    [
+        MMAType.F32_16x16x16_F16,
+        MMAType.F32_32x32x8_F16,
+    ],
+)
+def testGemmBias(
+    shape: tuple[int],
+    enable_scheduling: SchedulingType,
+    dynamic_dims: bool,
+    mfma_variant: MMAType,
+    request,
+):
+    run_bench = request.config.getoption("--runperf")
+    dump_perf = request.config.getoption("--dump-perf-files-path")
+    # Input sizes
+    M = tkl.sym.M
+    N = tkl.sym.N
+    K = tkl.sym.K
+    # Workgroup tile sizes
+    BLOCK_M = tkl.sym.BLOCK_M
+    BLOCK_N = tkl.sym.BLOCK_N
+    BLOCK_K = tkl.sym.BLOCK_K
+    # Address space (for GPU, shared(1) or global(0))
+    ADDRESS_SPACE = tkl.sym.ADDRESS_SPACE
+
+    # Expose user-constraints
+    constraints: list[tkw.Constraint] = [tkw.WorkgroupConstraint(M, BLOCK_M, 0)]
+    constraints += [tkw.WorkgroupConstraint(N, BLOCK_N, 1)]
+    constraints += [tkw.TilingConstraint(K, BLOCK_K)]
+    constraints += [tkw.WaveConstraint(M, BLOCK_M / 2)]
+    constraints += [tkw.WaveConstraint(N, BLOCK_N / 2)]
+
+    constraints += [
+        tkw.HardwareConstraint(
+            threads_per_wave=64, waves_per_block=(2, 2, 1), mma_type=mfma_variant
+        )
+    ]
+
+    # With dynamic dimensions, we need to add an assumption on how big
+    # the reduction dimension is to determine whether we can schedule or not.
+    if dynamic_dims:
+        constraints += [tkw.Assumption(K > BLOCK_K * 4)]
+
+    # Wave-level micro-kernel.
+    # Since warps are not directly addressable, there is no
+    # explicit notion of a warp id (like a workgroup or thread id).
+    # This kernel uses the input sizes M, N, K throughout, as the tiling
+    # and data movement strategy is determined during the compilation process.
+    # These can be influenced by introducing constraints.
+    @tkw.wave(constraints)
+    def gemm_bias(
+        a: tkl.Memory[M, K, ADDRESS_SPACE, tkl.f16],
+        b: tkl.Memory[N, K, ADDRESS_SPACE, tkl.f16],
+        bias: tkl.Memory[M, N, GLOBAL_ADDRESS_SPACE, tkl.f32],
+        c: tkl.Memory[M, N, GLOBAL_ADDRESS_SPACE, tkl.f32],
+    ):
+        c_reg = tkl.Register[M, N, tkl.f32](0.0)
+
+        # This microkernel encodes the fact that if the reduction
+        # dimension were tiled, then we would need to materialize a loop.
+        @tkw.reduction(K, init_args=[c_reg])
+        def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
+            a_reg = tkw.read(a)
+            b_reg = tkw.read(b)
+            acc = tkw.mma(a_reg, b_reg, acc)
+            return acc
+
+        # repeat represents the results of the loop
+        bias_reg = tkw.read(bias)
+        result = repeat + bias_reg
+        tkw.write(result, c)
+
+    hyperparams = {
+        ADDRESS_SPACE: SHARED_ADDRESS_SPACE,
+        BLOCK_M: 64,
+        BLOCK_N: 64,
+        BLOCK_K: 32,
+        M: shape[0],
+        N: shape[1],
+        K: shape[2],
+    }
+    hyperparams.update(get_default_scheduling_params())
+
+    dynamic_symbols = []
+    dynamic_symbols_map = {}
+    if dynamic_dims:
+        dynamic_symbols_map[M] = hyperparams[M]
+        dynamic_symbols_map[N] = hyperparams[N]
+        dynamic_symbols_map[K] = hyperparams[K]
+        dynamic_symbols.append(M)
+        dynamic_symbols.append(N)
+        dynamic_symbols.append(K)
+        del hyperparams[M]
+        del hyperparams[N]
+        del hyperparams[K]
+
+    perf_filename = request.node.name + ".json"
+    options = WaveCompileOptions(
+        subs=hyperparams,
+        canonicalize=True,
+        run_bench=run_bench,
+        schedule=enable_scheduling,
+        use_scheduling_barriers=enable_scheduling_barriers,
+        dynamic_symbols=dynamic_symbols,
+        dynamic_symbols_map=dynamic_symbols_map,
+        benchmark_batch_size=10,
+        benchmark_repetitions=3,
+        benchmark_results_file=(
+            os.path.join(dump_perf, "tk_" + perf_filename) if dump_perf else None
+        ),
+    )
+    options = set_default_run_config(options)
+    gemm_bias = wave_compile(options, gemm_bias)
+
+    a = device_randn(shape[0], shape[2], dtype=torch.float16)
+    b = device_randn(shape[1], shape[2], dtype=torch.float16)
+    c = device_zeros(shape[0], shape[1], dtype=torch.float32)
+    bias = device_randn(shape[0], shape[1], dtype=torch.float32)
+    asm = gemm_bias(a, b, bias, c)
+
+    if dump_generated_mlir:
+        filename = f"wave_gemm_{'x'.join(map(str, shape))}.mlir"
+        with open(filename, "w") as f:
+            f.write(asm)
+
+    if run_bench:
+        if dump_perf is not None:
+            options.benchmark_results_file = os.path.join(
+                dump_perf, "iree_" + perf_filename
+            )
+    iree_ref = device_zeros(shape[0], shape[1], dtype=torch.float32)
+    generate_iree_ref("mmt", [a, b], [iree_ref], options)
+    assert_close(c, iree_ref + bias, check_device=False)
+
+
 @require_e2e
 @require_cdna2
 @pytest.mark.parametrize("shape", get_test_shapes("test_gemm"))