Group fixed-route MoE tokens and reuse expert weight tiles

.github/workflows/ci.yml

@@ -58,6 +58,9 @@ jobs:
       - name: Run the stage-terminator guard
         run: PYTHONPATH=. python3 -m pytest aten/tests/test_moe_stage_terminator.py -v
 
+      - name: Run fixed-route expert-grouping compiler guard
+        run: PYTHONPATH=. python3 -m pytest aten/tests/test_moe_expert_grouping.py -v
+
   unit-tests:
     name: Generator unit tests
     runs-on: ubuntu-latest

aten/plena/isa_compiler.py

@@ -386,6 +386,7 @@ def add_hbm_object(
         shape: tuple[int, int],
         physical_shape: tuple[int, int] | None = None,
         real_data_ratio: float = 1.125,
+        storage_order: str = "row_major",
     ):
         """Register an HBM object and build its HBM layout.
 
@@ -400,6 +401,7 @@ def add_hbm_object(
             physical_shape=physical_shape,
             hbm_addr=hbm_addr,
             real_data_ratio=real_data_ratio,
+            storage_order=storage_order,
         )
 
     def free_hbm_object(self, name: str, strict: bool = False):

aten/plena/isa_matrix.py

@@ -55,7 +55,12 @@ def _emit_hbm_prefetch_setup(
         if set_scale:
             asm.instr("S_ADDI_INT", gp(gp_scale), gp(0), rows * cols)
             asm.instr("C_SET_SCALE_REG", gp(gp_scale))
-        asm.instr("S_ADDI_INT", gp(gp_stride), gp(0), cols * hbm_element_bytes)
+        asm.instr(
+            "S_ADDI_INT",
+            gp(gp_stride),
+            gp(0),
+            layout.hbm_row_stride_elements * hbm_element_bytes,
+        )
         asm.instr("C_SET_STRIDE_REG", gp(gp_stride))
 
     def _emit_hbm_subblock_prefetch(
@@ -308,6 +313,7 @@ def vram_sub_projection_asm(
         k_block_start: int = 0,
         k_block_count: int | None = None,
         unroll: bool | None = None,
+        resident_mram_start_addr: int | None = None,
     ) -> str:
         """Emit VRAM[row][:] @ MRAM[:][col] projection."""
         gp_regs = self._default_projection_gp_regs(gp_regs)
@@ -327,10 +333,12 @@ def vram_sub_projection_asm(
         valid_rows = vram_row_blocks[k_block_start].valid_shape[0] if vram_row_blocks[k_block_start].valid_shape else self.mlen
         row_loop_count = max(1, math.ceil(valid_rows / self.blen))
         vram_row_start_addr = vram_row_blocks[k_block_start].vram_addr
-        mram_col_start_addr = self._loaded_mram_start(
-            mram_col_blocks,
-            lambda block: f"{mram_mat_name}[{block.row_idx}][{mram_col_idx}]",
-        )
+        mram_col_start_addr = resident_mram_start_addr
+        if mram_col_start_addr is None:
+            mram_col_start_addr = self._loaded_mram_start(
+                mram_col_blocks,
+                lambda block: f"{mram_mat_name}[{block.row_idx}][{mram_col_idx}]",
+            )
 
         header_lines = [
             f"; VRAM Sub Projection: {vram_mat_name}[{vram_row_idx}][:] @ {mram_mat_name}[:][{mram_col_idx}]",
@@ -1115,13 +1123,16 @@ def _emit_vram_sub_projection_to(
         target_col_idx: int,
         k_block_start: int = 0,
         k_block_count: int | None = None,
+        resident_mram_start_addr: int | None = None,
     ) -> str:
         result_vram_addr, target_base_addr, target_rows = self._target_tile_addr(
             target_matrix, target_row_idx, target_col_idx
         )
         gp_regs = self.register_allocator.allocate_gp(9)
 
         if transposed:
+            if resident_mram_start_addr is not None:
+                raise ValueError("resident MRAM override is only supported for non-transposed projections")
             isa_code = f"; VRAM Sub Projection T To: {vram_mat_name}[{vram_row_idx}][:] @ {mram_mat_name}[{mram_idx}][:]^T -> {target_matrix}[{target_row_idx}][{target_col_idx}]\n"
             asm = self.vram_sub_projection_T_asm(
                 vram_mat_name=vram_mat_name,
@@ -1142,6 +1153,7 @@ def _emit_vram_sub_projection_to(
                 gp_regs=gp_regs,
                 k_block_start=k_block_start,
                 k_block_count=k_block_count,
+                resident_mram_start_addr=resident_mram_start_addr,
             )
         isa_code += f"; Target VRAM addr: {result_vram_addr} (base={target_base_addr}, offset=col*{target_rows}*{self.mlen} + row*{self.mlen}*{self.mlen})\n"
         isa_code += asm
@@ -1160,6 +1172,7 @@ def vram_sub_projection_to(
         target_col_idx: int,
         k_block_start: int = 0,
         k_block_count: int | None = None,
+        resident_mram_start_addr: int | None = None,
     ) -> str:
         """
         Sub-block multiplication:
@@ -1177,6 +1190,7 @@ def vram_sub_projection_to(
             target_col_idx=target_col_idx,
             k_block_start=k_block_start,
             k_block_count=k_block_count,
+            resident_mram_start_addr=resident_mram_start_addr,
         )
 
     def vram_sub_projection_T_to(

aten/plena/memory.py

@@ -154,14 +154,17 @@ class MatrixBlockLayout:
     """
     Block layout information for large matrices.
 
-    HBM storage: [rows, cols] row-major contiguous, stride=cols per row.
+    HBM storage is either row-major or tile-major. In tile-major mode each
+    ``block_size x block_size`` tile is contiguous and tiles are ordered by
+    ``(row_block, col_block)``.
     MRAM storage: [batch, mlen, hidden/mlen] column-block major.
     """
 
     name: str
     full_shape: tuple[int, int]  # Logical matrix shape (rows, cols)
     block_size: int = MLEN  # Sub-block size (default 64)
     physical_shape: tuple[int, int] | None = None  # Backing storage shape, if padded
+    storage_order: str = "row_major"
 
     num_row_blocks: int = 0
     num_col_blocks: int = 0
@@ -184,12 +187,24 @@ def __post_init__(self):
         self.physical_shape = (physical_rows, physical_cols)
         self.num_row_blocks = math.ceil(physical_rows / self.block_size)
         self.num_col_blocks = math.ceil(physical_cols / self.block_size)
+        if self.storage_order not in ("row_major", "tile_major"):
+            raise ValueError(f"Unsupported HBM matrix storage order: {self.storage_order!r}")
+        if self.storage_order == "tile_major" and (
+            physical_rows % self.block_size != 0 or physical_cols % self.block_size != 0
+        ):
+            raise ValueError(
+                "tile_major HBM storage requires physical rows and columns to be "
+                f"multiples of block_size={self.block_size}, got {self.physical_shape}"
+            )
 
         # Create information for all sub-blocks (pre-calculate addresses)
         for r in range(self.num_row_blocks):
             for c in range(self.num_col_blocks):
-                # HBM offset (row-major): sub-block (r,c) starts at r*block_size*cols + c*block_size
-                hbm_offset = r * self.block_size * physical_cols + c * self.block_size
+                if self.storage_order == "tile_major":
+                    tile_index = r * self.num_col_blocks + c
+                    hbm_offset = tile_index * self.block_size * self.block_size
+                else:
+                    hbm_offset = r * self.block_size * physical_cols + c * self.block_size
                 valid_rows = max(0, min(self.block_size, rows - r * self.block_size))
                 valid_cols = max(0, min(self.block_size, cols - c * self.block_size))
 
@@ -204,6 +219,13 @@ def __post_init__(self):
                 )
                 self.sub_blocks[(r, c)] = sub_info
 
+    @property
+    def hbm_row_stride_elements(self) -> int:
+        """Element stride between rows while loading one matrix tile."""
+        if self.storage_order == "tile_major":
+            return self.block_size
+        return self.physical_shape[1]
+
     def get_sub_block(self, row_idx: int, col_idx: int) -> SubMatrixInfo:
         """Get specified sub-block"""
         if (row_idx, col_idx) not in self.sub_blocks:

aten/plena/memory_state.py

@@ -129,6 +129,7 @@ def add_hbm_object(
         kind: str = "HBMObject",
         real_data_ratio: float = 1.125,
         strict: bool = False,
+        storage_order: str = "row_major",
     ) -> MemoryObjectInfo:
         del dtype, kind
         self.register_matrix(
@@ -138,6 +139,7 @@ def add_hbm_object(
             hbm_base_addr=hbm_addr,
             real_data_ratio=real_data_ratio,
             strict=strict,
+            storage_order=storage_order,
         )
         return self[name]
 
@@ -239,6 +241,7 @@ def register_matrix(
         physical_shape: tuple[int, int] | None = None,
         real_data_ratio: float = 1.125,
         strict: bool = True,
+        storage_order: str = "row_major",
     ) -> MatrixBlockLayout:
         """Register an HBM matrix and derive its mlen block layout."""
         del real_data_ratio  # HBM size is now row-aligned (see hbm_tensor_size)
@@ -261,6 +264,7 @@ def register_matrix(
             block_size=self.mlen,
             hbm_base_addr=hbm_base_addr,
             hbm_size=hbm_size,
+            storage_order=storage_order,
         )
 
         self.hbm_matrices[name] = layout