From e4cd6c5f9887f0118a43e3e5484130f45291fce9 Mon Sep 17 00:00:00 2001 From: Michael C Li Date: Thu, 3 Sep 2026 19:04:14 +0100 Subject: [PATCH] feat(moe): group fixed routes and reuse expert weights --- .github/workflows/ci.yml | 3 + aten/plena/isa_compiler.py | 2 + aten/plena/isa_matrix.py | 24 +- aten/plena/memory.py | 28 +- aten/plena/memory_state.py | 4 + aten/plena/program_routed_moe.py | 453 +++++++++++++++++++++++-- aten/plena/program_tensors.py | 2 + aten/tests/test_moe_expert_grouping.py | 204 +++++++++++ 8 files changed, 692 insertions(+), 28 deletions(-) create mode 100644 aten/tests/test_moe_expert_grouping.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3732abc..90ca5ee 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/aten/plena/isa_compiler.py b/aten/plena/isa_compiler.py index 7d1402d..894c1f9 100644 --- a/aten/plena/isa_compiler.py +++ b/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): diff --git a/aten/plena/isa_matrix.py b/aten/plena/isa_matrix.py index 0f26b68..84adddb 100644 --- a/aten/plena/isa_matrix.py +++ b/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,6 +1123,7 @@ 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 @@ -1122,6 +1131,8 @@ def _emit_vram_sub_projection_to( 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( diff --git a/aten/plena/memory.py b/aten/plena/memory.py index 7dff198..fee2878 100644 --- a/aten/plena/memory.py +++ b/aten/plena/memory.py @@ -154,7 +154,9 @@ 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. """ @@ -162,6 +164,7 @@ class MatrixBlockLayout: 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: diff --git a/aten/plena/memory_state.py b/aten/plena/memory_state.py index ce35cac..de36adb 100644 --- a/aten/plena/memory_state.py +++ b/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 diff --git a/aten/plena/program_routed_moe.py b/aten/plena/program_routed_moe.py index 9720c6c..89b6966 100644 --- a/aten/plena/program_routed_moe.py +++ b/aten/plena/program_routed_moe.py @@ -786,6 +786,16 @@ def _moe_dynamic_load_sub_matrix_col_v0( gp_base=gp_base, name=name, ) + # Address selection has its own stage marker. Hand attribution back + # before setup and H_PREFETCH_M so transfer cycles and bytes are not + # silently charged to ``expert_weight_address``. + asm.comment( + moe_stage_marker( + "expert_weight_prefetch", + f"issue dynamic HBM weight tile: template={weight_template.name}, " + f"pair={pair_idx}, col={col_idx}", + ) + ) self._emit_hbm_prefetch_setup(asm, layout, gp_scale, gp_stride) self._emit_hbm_subblock_sequence( asm, @@ -876,8 +886,28 @@ def moe_dynamic_linear_projection_v0( expert_base_table_int_base: int | None = None, name: str, physical_shape: tuple[int, int] | None = None, + reuse_weight_across_row_blocks: bool = False, + weight_panel_mode: str = "blocking", + panel_k_tiles: int | None = None, ) -> VRAMMatrixVar: - """Tiled linear projection with runtime expert-id weight selection.""" + """Tiled linear projection with runtime expert-id weight selection. + + ``reuse_weight_across_row_blocks`` is the fixed-route grouped-execution + contract. A weight column tile is loaded into Matrix SRAM once and is + consumed by every token-row block belonging to the same expert before + the next tile replaces it. The default remains the historical + load-per-row-block behavior so existing programs are unchanged. + + ``weight_panel_mode="pingpong"`` reserves two disjoint Matrix-SRAM + panels. It emits the next fill before consuming the current panel and + reuses a slot only after its previous panel has been consumed. + Functional execution remains dependency-safe; a dependency-aware + simulator may overlap a pending fill with work on the ready panel. + """ + if weight_panel_mode not in ("blocking", "pingpong"): + raise ValueError(f"unsupported weight_panel_mode={weight_panel_mode!r}") + if weight_panel_mode != "blocking" and not reuse_weight_across_row_blocks: + raise ValueError("buffered weight panels require grouped row-block reuse") mlen = self.mlen rows, _k_total = input_var.shape _weight_rows, out_features = weight_template.shape @@ -899,7 +929,13 @@ def moe_dynamic_linear_projection_v0( num_row_blocks = math.ceil(physical_rows / mlen) num_col_blocks = math.ceil(physical_out_features / mlen) num_k_tiles = math.ceil(physical_k / mlen) - max_k_tiles = self.mram_tile_capacity + max_k_tiles = self.mram_tile_capacity if panel_k_tiles is None else int(panel_k_tiles) + if max_k_tiles <= 0: + raise ValueError(f"panel_k_tiles must be positive, got {max_k_tiles}") + if max_k_tiles > self.mram_tile_capacity: + raise ValueError( + f"panel_k_tiles={max_k_tiles} exceeds mram_tile_capacity={self.mram_tile_capacity}" + ) output = self.alloc( name, @@ -909,28 +945,156 @@ def moe_dynamic_linear_projection_v0( physical_shape=(physical_rows, physical_out_features), ) - def emit_projection(row_idx, col_idx, target, target_row_idx, target_col_idx, **k_split) -> None: - self.moe_dynamic_vram_sub_projection_to_v0( - input_var, - row_idx, - weight_template, - col_idx, - target, - target_row_idx, - target_col_idx, - expert_indices_int_base=expert_indices_int_base, - pair_idx=pair_idx, - table_base=table_base, - per_expert_stride=per_expert_stride, - expert_base_table_int_base=expert_base_table_int_base, - name=f"{name}_pair{pair_idx}", + if weight_panel_mode != "blocking": + panel_tiles = 4 if panel_k_tiles is None else int(panel_k_tiles) + if panel_tiles <= 0: + raise ValueError(f"panel_k_tiles must be positive, got {panel_tiles}") + panel_depth = 2 + if panel_depth * panel_tiles > self.mram_tile_capacity: + raise ValueError( + "buffered panels exceed Matrix-SRAM capacity: " + f"{panel_depth}*{panel_tiles} > mram_tile_capacity={self.mram_tile_capacity}" + ) + + super().reset_mram() + panel_elems = panel_tiles * mlen * mlen + panel_bases = tuple( + self.mram_allocator.allocate(f"{name}_panel{slot}", panel_elems) + for slot in range(panel_depth) + ) + panels = [] + for k_chunk_idx, k_block_start in enumerate(range(0, num_k_tiles, panel_tiles)): + k_block_count = min(panel_tiles, num_k_tiles - k_block_start) + for col_idx in range(num_col_blocks): + panels.append((k_chunk_idx, k_block_start, k_block_count, col_idx)) + + def prefetch_panel(panel, mram_start_addr: int) -> None: + _chunk_idx, k_block_start, k_block_count, col_idx = panel + self._moe_dynamic_load_sub_matrix_col_v0( + weight_template=weight_template, + col_idx=col_idx, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=table_base, + per_expert_stride=per_expert_stride, + expert_base_table_int_base=expert_base_table_int_base, + mram_start_addr=mram_start_addr, + k_block_start=k_block_start, + k_block_count=k_block_count, + name=f"{name}_pair{pair_idx}_pingpong", + ) + + def consume_panel( + panel, + mram_start_addr: int, + target: VRAMMatrixVar, + row_idx: int, + target_row_idx: int, + target_col_idx: int, + ) -> None: + _chunk_idx, k_block_start, k_block_count, col_idx = panel + self._emit( + IsaBuilder().comment( + moe_stage_marker( + "expert_projection", + f"{name}: pingpong resident weight pair={pair_idx}, col={col_idx}, " + f"row_block={row_idx}, mram={mram_start_addr}", + ) + ) + ) + super(ProgramRoutedMoeMixin, self).vram_sub_projection_to( + vram_mat_name=input_var.name, + vram_row_idx=row_idx, + mram_mat_name=weight_template.name, + mram_col_idx=col_idx, + target_matrix=target.name, + target_row_idx=target_row_idx, + target_col_idx=target_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + resident_mram_start_addr=mram_start_addr, + ) + + temp = self.alloc(f"{name}_temp", mlen, mlen) if num_k_tiles > panel_tiles else None + initial_fill = min(panel_depth, len(panels)) + for panel_idx in range(initial_fill): + prefetch_panel(panels[panel_idx], panel_bases[panel_idx]) + for panel_idx, panel in enumerate(panels): + current_base = panel_bases[panel_idx % panel_depth] + k_chunk_idx, _k_start, _k_count, col_idx = panel + for row_idx in range(num_row_blocks): + if k_chunk_idx == 0: + consume_panel(panel, current_base, output, row_idx, row_idx, col_idx) + else: + assert temp is not None + consume_panel(panel, current_base, temp, row_idx, 0, 0) + self.vram_block_add_to(output, row_idx, col_idx, temp, 0, 0, output, row_idx, col_idx) + next_panel_idx = panel_idx + panel_depth + if next_panel_idx < len(panels): + prefetch_panel(panels[next_panel_idx], current_base) + if temp is not None: + self.free_tensor(temp) + return output + + def emit_projection( + row_idx, + col_idx, + target, + target_row_idx, + target_col_idx, + *, + load_weight: bool, + **k_split, + ) -> None: + if load_weight: + self.moe_dynamic_vram_sub_projection_to_v0( + input_var, + row_idx, + weight_template, + col_idx, + target, + target_row_idx, + target_col_idx, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=table_base, + per_expert_stride=per_expert_stride, + expert_base_table_int_base=expert_base_table_int_base, + name=f"{name}_pair{pair_idx}", + **k_split, + ) + return + + self._emit( + IsaBuilder().comment( + moe_stage_marker( + "expert_projection", + f"{name}: resident weight pair={pair_idx}, col={col_idx}, row_block={row_idx}", + ) + ) + ) + super(ProgramRoutedMoeMixin, self).vram_sub_projection_to( + vram_mat_name=input_var.name, + vram_row_idx=row_idx, + mram_mat_name=weight_template.name, + mram_col_idx=col_idx, + target_matrix=target.name, + target_row_idx=target_row_idx, + target_col_idx=target_col_idx, **k_split, ) if num_k_tiles <= max_k_tiles: for col_idx in range(num_col_blocks): for row_idx in range(num_row_blocks): - emit_projection(row_idx, col_idx, output, row_idx, col_idx) + emit_projection( + row_idx, + col_idx, + output, + row_idx, + col_idx, + load_weight=not reuse_weight_across_row_blocks or row_idx == 0, + ) return output temp = self.alloc(f"{name}_temp", mlen, mlen) @@ -940,9 +1104,25 @@ def emit_projection(row_idx, col_idx, target, target_row_idx, target_col_idx, ** for col_idx in range(num_col_blocks): for row_idx in range(num_row_blocks): if k_chunk_idx == 0: - emit_projection(row_idx, col_idx, output, row_idx, col_idx, **k_split) + emit_projection( + row_idx, + col_idx, + output, + row_idx, + col_idx, + load_weight=not reuse_weight_across_row_blocks or row_idx == 0, + **k_split, + ) else: - emit_projection(row_idx, col_idx, temp, 0, 0, **k_split) + emit_projection( + row_idx, + col_idx, + temp, + 0, + 0, + load_weight=not reuse_weight_across_row_blocks or row_idx == 0, + **k_split, + ) self.vram_block_add_to(output, row_idx, col_idx, temp, 0, 0, output, row_idx, col_idx) self.free_tensor(temp) return output @@ -1227,6 +1407,168 @@ def moe_dynamic_expert_pair_v0( self.vram_mul(out, route, num_rows=rows) return out + def moe_dynamic_expert_group_v0( + self, + x: VRAMMatrixVar, + weights: ExpertWeights, + *, + weight_table_bases: tuple[int, int, int], + weight_table_strides: tuple[int, int, int], + expert_indices_int_base: int, + weights_fp_base: int, + pair_indices: Sequence[int], + bias_tables: ExpertBiases | None, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + zero_row: FPVar | None = None, + route_fp_scratch: FPVar | None = None, + policy_name: str = "gpt_oss", + activation_policy: str = "gpt_oss_clamp_gated", + weight_panel_mode: str = "blocking", + panel_k_tiles: int | None = None, + name: str = "moe_dynamic_expert_group", + ) -> VRAMMatrixVar: + """Run one fixed-route expert for a compact group of token rows. + + The caller guarantees that every entry in ``pair_indices`` names the + same expert in integer SRAM. The first pair resolves the expert's HBM + base once; its gate/up/down tiles remain in Matrix SRAM while all + compact token rows consume them. Distinct per-pair route weights are + then materialized row by row before scatter-add. + + This helper intentionally does not sort device-selected routing. It is + legal for fixed-route/trace replay where grouping is known at compile + time; a runtime router needs a separate sort/segment mechanism before + it may call the same execution primitive. + """ + pair_list = [int(pair) for pair in pair_indices] + if rows <= 0: + raise ValueError(f"{name}: rows must be positive") + if len(pair_list) != rows: + raise ValueError(f"{name}: pair_indices={len(pair_list)} must equal rows={rows}") + if len(set(pair_list)) != len(pair_list) or min(pair_list) < 0: + raise ValueError(f"{name}: pair_indices must be distinct non-negative values") + if x.shape[0] != rows: + raise ValueError(f"{name}: compact x rows={x.shape[0]} must equal rows={rows}") + if bias_tables is not None and rows > self.blen: + raise ValueError( + f"{name}: grouped execution with expert bias supports at most " + f"BLEN={self.blen} rows, got rows={rows}" + ) + + representative_pair = pair_list[0] + w_gate, w_up, w_down = weights + gate_bias_table, up_bias_table, down_bias_table = bias_tables or (None, None, None) + gate_base, up_base, down_base = weight_table_bases + gate_stride, up_stride, down_stride = weight_table_strides + projection_rows = max(self.mlen, x.physical_shape[0], math.ceil(rows / self.blen) * self.blen) + + def project(input_var, weight, table_base, stride, suffix, output_shape): + return self.moe_dynamic_linear_projection_v0( + input_var, + weight, + expert_indices_int_base=expert_indices_int_base, + pair_idx=representative_pair, + table_base=table_base, + per_expert_stride=stride, + name=f"{name}_{suffix}", + physical_shape=output_shape, + reuse_weight_across_row_blocks=True, + weight_panel_mode=weight_panel_mode, + panel_k_tiles=panel_k_tiles, + ) + + gate = project( + x, + w_gate, + gate_base, + gate_stride, + "gate", + (projection_rows, w_gate.physical_shape[1]), + ) + up = project( + x, + w_up, + up_base, + up_stride, + "up", + (projection_rows, w_up.physical_shape[1]), + ) + if gate_bias_table is not None: + self.moe_add_dynamic_expert_bias_v0( + gate, + gate_bias_table, + expert_indices_int_base=expert_indices_int_base, + pair_idx=representative_pair, + rows=rows, + width=intermediate, + name=f"{name}_gate_bias", + ) + if up_bias_table is not None: + self.moe_add_dynamic_expert_bias_v0( + up, + up_bias_table, + expert_indices_int_base=expert_indices_int_base, + pair_idx=representative_pair, + rows=rows, + width=intermediate, + name=f"{name}_up_bias", + ) + + hidden = self.moe_expert_activation_v0( + gate, + up, + rows=rows, + intermediate=intermediate, + constants=constants, + activation_policy=activation_policy, + policy_name=policy_name, + stage="expert_activation", + name=name, + ) + out = project( + hidden, + w_down, + down_base, + down_stride, + "out", + (projection_rows, w_down.physical_shape[1]), + ) + if down_bias_table is not None: + self.moe_add_dynamic_expert_bias_v0( + out, + down_bias_table, + expert_indices_int_base=expert_indices_int_base, + pair_idx=representative_pair, + rows=rows, + width=w_down.physical_shape[1], + name=f"{name}_down_bias", + ) + + route = self.moe_materialize_route_weights_for_active_rows_v0( + weights_fp_base=weights_fp_base, + pair_indices=pair_list, + active_rows=list(range(rows)), + rows=rows, + hidden=w_down.physical_shape[1], + zero_row=zero_row, + fp_scratch=route_fp_scratch, + policy_name=policy_name, + stage="expert_route_weight", + name=f"{name}_route", + ) + self._emit( + IsaBuilder().comment( + moe_stage_marker( + "expert_route_weight", + f"[{policy_name}] apply grouped route weights: rows={rows}, representative_pair={representative_pair}", + ) + ) + ) + self.vram_mul(out, route, num_rows=rows) + return out + def moe_gather_token_rows_from_hbm_v0( self, x_input: InputVar, @@ -1403,6 +1745,77 @@ def moe_gather_token_rows_from_vram_v0( return gathered + def moe_gather_token_rows_compact_from_vram_v0( + self, + x: VRAMMatrixVar, + *, + token_indices: Sequence[int], + hidden: int, + zero_row: FPVar | None = None, + policy_name: str = "gpt_oss", + name: str = "moe_grouped_x_vram", + ) -> VRAMMatrixVar: + """Gather one expert's token rows densely for weight-reuse execution. + + Unlike :meth:`moe_gather_token_rows_from_vram_v0`, which reserves one + BLEN-row slot per routed pair, this mapping places token ``j`` at row + ``j``. Padding is restricted to the final BLEN tile, so one resident + expert weight tile serves all valid rows in the group. + """ + self._ensure_vram_sub_matrix_registered(x) + if hidden % self.mlen != 0: + raise ValueError(f"compact gather hidden={hidden} must be divisible by MLEN={self.mlen}") + if hidden > x.shape[1]: + raise ValueError(f"compact gather hidden={hidden} exceeds x width={x.shape[1]}") + + token_list = [int(token) for token in token_indices] + if not token_list: + raise ValueError("compact gather token_indices must be non-empty") + if len(set(token_list)) != len(token_list): + raise ValueError("compact gather requires unique token rows for one expert") + if min(token_list) < 0 or max(token_list) >= x.physical_shape[0]: + raise ValueError(f"compact gather token_indices {token_list} exceed x physical rows={x.physical_shape[0]}") + + rows = len(token_list) + physical_rows = max(self.blen, math.ceil(rows / self.blen) * self.blen) + gathered = self.alloc( + name, + rows=rows, + cols=hidden, + strict=False, + physical_shape=(physical_rows, hidden), + ) + self.moe_true_zero_vram_rows_v0( + gathered, + rows=list(range(physical_rows)), + hidden=hidden, + zero_row=zero_row, + policy_name=policy_name, + stage="gather", + name=f"{name}_zero", + ) + + num_col_blocks = hidden // self.mlen + gp_dst, gp_src = self._reg.allocate_gp(2) + try: + asm = IsaBuilder().comment( + moe_stage_marker( + "gather", + f"[{policy_name}] compact expert-major VRAM gather: rows={rows}, hidden={hidden}", + ) + ) + for active_row, token_idx in enumerate(token_list): + for col_block in range(num_col_blocks): + dst_addr = self._vram_matrix_row_addr(gathered, active_row, col_block) + src_addr = self._vram_matrix_row_addr(x, token_idx, col_block) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), dst_addr) + asm.instr("S_ADDI_INT", gp(gp_src), gp(0), src_addr) + asm.instr("V_ADD_VV", gp(gp_dst), gp(gp_dst), gp(gp_src), 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_dst, gp_src]) + return gathered + def moe_true_zero_vram_rows_v0( self, matrix: VRAMMatrixVar, diff --git a/aten/plena/program_tensors.py b/aten/plena/program_tensors.py index c2b6099..9ccb34d 100644 --- a/aten/plena/program_tensors.py +++ b/aten/plena/program_tensors.py @@ -19,6 +19,7 @@ def input( prestaged_vram_addr: int | None = None, physical_shape: tuple[int, int] | None = None, real_data_ratio: float | None = None, + hbm_storage_order: str = "row_major", ) -> InputVar: """ Declare an input tensor (in HBM). @@ -62,6 +63,7 @@ def input( shape=shape, physical_shape=physical_shape, real_data_ratio=real_data_ratio, + storage_order=hbm_storage_order, ) return var diff --git a/aten/tests/test_moe_expert_grouping.py b/aten/tests/test_moe_expert_grouping.py new file mode 100644 index 0000000..1a8f4a9 --- /dev/null +++ b/aten/tests/test_moe_expert_grouping.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +from compiler.aten.plena import PlenaCompiler +from compiler.aten.plena.memory import MatrixBlockLayout + + +def test_hbm_storage_order_preserves_row_major_default() -> None: + layout = MatrixBlockLayout("w", (128, 128), block_size=64) + assert layout.storage_order == "row_major" + assert layout.hbm_row_stride_elements == 128 + assert layout.get_sub_block(0, 1).hbm_offset == 64 + assert layout.get_sub_block(1, 0).hbm_offset == 8192 + + +def test_tile_major_places_each_matrix_tile_contiguously() -> None: + layout = MatrixBlockLayout( + "w", (128, 128), block_size=64, storage_order="tile_major" + ) + assert layout.hbm_row_stride_elements == 64 + assert layout.get_sub_block(0, 0).hbm_offset == 0 + assert layout.get_sub_block(0, 1).hbm_offset == 4096 + assert layout.get_sub_block(1, 0).hbm_offset == 8192 + assert layout.get_sub_block(1, 1).hbm_offset == 12288 + + +def test_tile_major_rejects_partial_physical_tiles() -> None: + try: + MatrixBlockLayout("w", (65, 64), block_size=64, storage_order="tile_major") + except ValueError as exc: + assert "multiples of block_size=64" in str(exc) + else: + raise AssertionError("tile-major layout must reject partial physical tiles") + + +def _program() -> tuple[ + PlenaCompiler, object, tuple[object, object, object], tuple[object, ...] +]: + prog = PlenaCompiler(mlen=64, blen=4, real_data_ratio=1.125) + x = prog.alloc("x", rows=4, cols=64, strict=False, physical_shape=(4, 64)) + weights = ( + prog.input("w_gate", shape=(64, 64), hbm_storage_order="tile_major"), + prog.input("w_up", shape=(64, 64), hbm_storage_order="tile_major"), + prog.input("w_down", shape=(64, 64), hbm_storage_order="tile_major"), + ) + zero = prog.fp_var("zero", size=1) + limit_pos = prog.fp_var("limit_pos", size=4) + limit_neg = prog.fp_var("limit_neg", size=4) + one = prog.fp_var("one", size=4) + neg_one = prog.fp_var("neg_one", size=4) + return prog, x, weights, (zero, limit_pos, limit_neg, one, neg_one) + + +def _emit_group(pair_indices: list[int]) -> str: + prog, x, weights, constants = _program() + gathered = prog.moe_gather_token_rows_compact_from_vram_v0( + x, + token_indices=list(range(len(pair_indices))), + hidden=64, + zero_row=constants[0], + name="grouped_gather", + ) + route_scratch = prog.fp_var("route_scratch", size=64) + prog.moe_dynamic_expert_group_v0( + gathered, + weights, + weight_table_bases=(0, 4096, 8192), + weight_table_strides=(4096, 4096, 4096), + expert_indices_int_base=0, + weights_fp_base=128, + pair_indices=pair_indices, + bias_tables=None, + rows=len(pair_indices), + intermediate=64, + constants=constants, + zero_row=constants[0], + route_fp_scratch=route_scratch, + activation_policy="standard_swiglu", + name="grouped_expert", + ) + return prog.compile() + + +def test_compact_gather_packs_rows_without_blen_holes() -> None: + asm = _emit_group([1, 4, 7]) + assert "compact expert-major VRAM gather: rows=3" in asm + assert "apply grouped route weights: rows=3" in asm + assert "active_row=0" in asm + assert "active_row=1" in asm + assert "active_row=2" in asm + + +def test_group_fetches_one_weight_triple_for_multiple_pairs() -> None: + asm = _emit_group([1, 4, 7]) + prefetch_markers = [ + line for line in asm.splitlines() if "dynamic HBM weight prefetch" in line + ] + assert len(prefetch_markers) == 3 + assert all("pair=1" in line for line in prefetch_markers) + + +def test_hbm_prefetch_is_attributed_after_dynamic_address_calculation() -> None: + lines = _emit_group([1, 4, 7]).splitlines() + for line_index, line in enumerate(lines): + if not line.startswith("H_PREFETCH_M"): + continue + latest_stage = next( + previous + for previous in reversed(lines[:line_index]) + if previous.startswith("; @stage=") + ) + assert latest_stage.startswith("; @stage=expert_weight_prefetch") + + +def test_group_rejects_duplicate_pair_indices() -> None: + try: + _emit_group([1, 1]) + except ValueError as exc: + assert "distinct non-negative" in str(exc) + else: + raise AssertionError("duplicate pair indices must be rejected") + + +def _emit_buffered_projection( + *, + mode: str = "pingpong", + mram_tile_capacity: int = 8, +) -> str: + prog = PlenaCompiler( + mlen=64, + blen=4, + real_data_ratio=1.125, + mram_tile_capacity=mram_tile_capacity, + ) + # Six four-tile panels force both the depth-two and depth-four paths to + # consume a live slot before wrapping around and refilling it. + k = 768 + x = prog.alloc("x", rows=4, cols=k, strict=False, physical_shape=(64, k)) + weight = prog.input("w", shape=(k, 128), hbm_storage_order="tile_major") + prog.moe_dynamic_linear_projection_v0( + x, + weight, + expert_indices_int_base=0, + pair_idx=0, + table_base=0, + per_expert_stride=k * 128, + name="buffered_projection", + reuse_weight_across_row_blocks=True, + weight_panel_mode=mode, + panel_k_tiles=4, + ) + return prog.compile() + + +def test_pingpong_prefetches_next_panel_before_consuming_current_panel() -> None: + lines = _emit_buffered_projection().splitlines() + prefetch_markers = [ + index + for index, line in enumerate(lines) + if "dynamic HBM weight prefetch" in line + ] + projection_markers = [ + index for index, line in enumerate(lines) if "pingpong resident weight" in line + ] + assert len(prefetch_markers) == 6 + assert len(projection_markers) == 6 + assert prefetch_markers[0] < prefetch_markers[1] < projection_markers[0] + assert projection_markers[0] < prefetch_markers[2] < projection_markers[1] + assert "mram=0" in lines[projection_markers[0]] + assert "mram=16384" in lines[projection_markers[1]] + assert "mram=0" in lines[projection_markers[2]] + + +def test_pingpong_requires_two_four_tile_panels() -> None: + try: + _emit_buffered_projection(mram_tile_capacity=4) + except ValueError as exc: + assert "exceed Matrix-SRAM capacity" in str(exc) + else: + raise AssertionError("pingpong must reject insufficient Matrix-SRAM capacity") + + +def test_blocking_panel_size_is_independent_of_total_mram_capacity() -> None: + prog = PlenaCompiler( + mlen=64, + blen=4, + real_data_ratio=1.125, + mram_tile_capacity=16, + ) + x = prog.alloc("x", rows=4, cols=512, strict=False, physical_shape=(64, 512)) + weight = prog.input("w", shape=(512, 64), hbm_storage_order="tile_major") + prog.moe_dynamic_linear_projection_v0( + x, + weight, + expert_indices_int_base=0, + pair_idx=0, + table_base=0, + per_expert_stride=512 * 64, + name="fixed_blocking_panel", + reuse_weight_across_row_blocks=True, + weight_panel_mode="blocking", + panel_k_tiles=4, + ) + asm = prog.compile() + assert asm.count("dynamic HBM weight prefetch") == 2