feature: add routed MoE compiler substrate

README.md

@@ -1 +1,10 @@
-# PLENA Compiler  
+# PLENA Compiler
+
+## MoE code organization
+
+- `aten/plena/program_routed_moe.py` contains reusable routed-MoE lowering
+  helpers: router logits, V_TOPK selection, dynamic expert-weight addressing,
+  routed gather/scatter, expert activation, and combine.
+- `aten/models/gpt_oss/` contains GPT-OSS-specific reference semantics and
+  real-checkpoint loading utilities used to validate that substrate.
+- ISA, assembler, and hardware documentation remain in `assembler/` and `doc/`.

asm_templates/im2col_asm.py

@@ -116,8 +116,8 @@ def im2col_asm(
     )
     if num_tiles > 1:
         assert vlen % K == 0, (
-            f"multi-tile im2col with V_SHIFT_V requires K ({K}) | VLEN ({vlen}); "
-            "K-groups would cross tile boundaries and V_SHIFT_V cannot right-shift across tile"
+            f"multi-tile im2col with V_SHFT_V requires K ({K}) | VLEN ({vlen}); "
+            "K-groups would cross tile boundaries and V_SHFT_V cannot right-shift across tile"
         )
 
     # W_padded: each input row is stored with W_padded elements in HBM so
@@ -147,7 +147,7 @@ def im2col_asm(
 
     lines: list[str] = []
     lines.append("; ============================================================")
-    lines.append("; im2col (with V_SHIFT_V): NCHW input in HBM -> im2col matrix in VRAM")
+    lines.append("; im2col (with V_SHFT_V): NCHW input in HBM -> im2col matrix in VRAM")
     lines.append(f";   input shape : (1, {C_in}, {H}, {W})  in HBM (W_padded={W_padded})")
     lines.append(f";   kernel      : {K}x{K},  OH={OH}, OW={OW}, stride={stride}")
     lines.append(
@@ -266,7 +266,7 @@ def im2col_asm(
                     # Shift right to local position within tile.
                     if local_shift > 0:
                         lines.append(f"S_ADDI_INT gp{shift_reg}, gp0, {local_shift}")
-                        lines.append(f"V_SHIFT_V gp{scratch_reg}, gp{scratch_reg}, gp{shift_reg}")
+                        lines.append(f"V_SHFT_V gp{scratch_reg}, gp{scratch_reg}, gp{shift_reg}")
 
                     # Accumulate: accum += shifted scratch row.
                     lines.append(f"V_ADD_VV gp{acc_reg}, gp{acc_reg}, gp{scratch_reg}, 0")

asm_templates/store_act_asm.py

@@ -14,6 +14,8 @@ def store_act_asm(
     hbm_addr_reg: int,
     stride_size: int | None = None,
     store_amount: int = 4,
+    precision: int = 0,
+    hbm_element_bytes: int = 1,
 ) -> str:
     """Store activation from VRAM back to HBM (reverse of preload_act_asm).
 
@@ -29,6 +31,8 @@ def store_act_asm(
     inner_loop_register = alive_registers[4]
 
     stride_len = hidden_size if stride_size is None else stride_size
+    if hbm_element_bytes <= 0:
+        raise ValueError(f"hbm_element_bytes must be > 0, got {hbm_element_bytes}")
     store_amount_per_hidden = math.ceil(hidden_size / vlen)
 
     # Initialize VRAM source address
@@ -38,19 +42,19 @@ def store_act_asm(
     # HBM MX formats store scale bytes after the element payload.  H_STORE_V
     # uses C_SET_SCALE_REG as the scale-section base offset; do not inherit a
     # stale value from a previous HBM load/store.
-    generated_code += _load_large_int(set_stride_register, batch * hidden_size)
+    generated_code += _load_large_int(set_stride_register, batch * hidden_size * hbm_element_bytes)
     generated_code += f"C_SET_SCALE_REG gp{set_stride_register}\n"
 
     if batch == 1:
         # Simple case: no stride needed, store sequentially
         elements_per_store = vlen * store_amount
         for i in range(math.ceil(hidden_size / elements_per_store)):
-            generated_code += f"H_STORE_V gp{vram_reg}, gp{hbm_offset_reg}, a{hbm_addr_reg}, 0, 0\n"
+            generated_code += f"H_STORE_V gp{vram_reg}, gp{hbm_offset_reg}, a{hbm_addr_reg}, 0, {precision}\n"
             generated_code += f"S_ADDI_INT gp{vram_reg}, gp{vram_reg}, {elements_per_store}\n"
-            generated_code += f"S_ADDI_INT gp{hbm_offset_reg}, gp{hbm_offset_reg}, {elements_per_store}\n"
+            generated_code += f"S_ADDI_INT gp{hbm_offset_reg}, gp{hbm_offset_reg}, {elements_per_store * hbm_element_bytes}\n"
     else:
         # Set stride register (HBM row stride = hidden_size)
-        generated_code += _load_large_int(set_stride_register, stride_len)
+        generated_code += _load_large_int(set_stride_register, stride_len * hbm_element_bytes)
         generated_code += f"C_SET_STRIDE_REG gp{set_stride_register}\n"
         hbm_base_reg = set_stride_register  # reuse after stride is set
 
@@ -62,15 +66,15 @@ def store_act_asm(
             # Inner loop: iterate over batch blocks
             generated_code += f"C_LOOP_START gp{inner_loop_register}, {math.ceil(batch / store_amount)}\n"
 
-        generated_code += f"H_STORE_V gp{vram_reg}, gp{hbm_base_reg}, a{hbm_addr_reg}, 1, 0\n"
+        generated_code += f"H_STORE_V gp{vram_reg}, gp{hbm_base_reg}, a{hbm_addr_reg}, 1, {precision}\n"
         generated_code += f"S_ADDI_INT gp{vram_reg}, gp{vram_reg}, {vlen * store_amount}\n"
 
         if batch > store_amount:
-            generated_code += f"S_ADDI_INT gp{hbm_base_reg}, gp{hbm_base_reg}, {hidden_size * store_amount}\n"
+            generated_code += f"S_ADDI_INT gp{hbm_base_reg}, gp{hbm_base_reg}, {hidden_size * store_amount * hbm_element_bytes}\n"
             generated_code += f"C_LOOP_END gp{inner_loop_register}\n"
 
         # Move to next column block in HBM
-        generated_code += f"S_ADDI_INT gp{hbm_offset_reg}, gp{hbm_offset_reg}, {vlen}\n"
+        generated_code += f"S_ADDI_INT gp{hbm_offset_reg}, gp{hbm_offset_reg}, {vlen * hbm_element_bytes}\n"
         generated_code += f"C_LOOP_END gp{outer_loop_register}\n"
 
     return generated_code

assembler/assembly_to_binary.py

@@ -8,7 +8,20 @@
 # instruction (millions of times for large programs). Membership is O(1) and identical
 # to the previous `opcode in [ ... ]` list literals.
 _RMASK_VECTOR_OPS = frozenset(
-    {"V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF", "V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"}
+    {
+        "V_ADD_VV",
+        "V_ADD_VF",
+        "V_MUL_VV",
+        "V_SUB_VV",
+        "V_MUL_VF",
+        "V_EXP_V",
+        "V_RECI_V",
+        "V_RED_SUM",
+        "V_RED_MAX",
+        "V_MAX_VF",
+        "V_MIN_VF",
+        "V_TOPK",
+    }
 )
 _IMM_RS1_RD_OPS = frozenset(
     {

assembler/parser.py

@@ -96,7 +96,9 @@ def __repr__(self):
 # Hoisted to module scope: these were previously re-created for every line of the
 # .asm (millions of times for large programs), which dominated sim_env re-parse time.
 vector_masked_unary_or_reduction_ops = frozenset({"V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"})
-vector_masked_binary_ops = frozenset({"V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF"})
+vector_masked_binary_ops = frozenset(
+    {"V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF", "V_MAX_VF", "V_MIN_VF", "V_TOPK"}
+)
 
 
 def _parse_operand(operand):

assembler/tests/test_vector_rmask_handling.py

@@ -23,6 +23,26 @@ def test_encoder_defaults_missing_rmask_to_zero(self):
 
         self.assertEqual(self.asm._convert_to_binary(missing_mask), self.asm._convert_to_binary(explicit_mask))
 
+    def test_vector_scalar_minmax_encode_like_masked_vector_ops(self):
+        max_instr = Instruction("V_MAX_VF", 1, 2, 3, 0, None, None, None)
+        min_instr = Instruction("V_MIN_VF", 1, 2, 3, 0, None, None, None)
+
+        max_binary = self.asm._convert_to_binary(max_instr)
+        min_binary = self.asm._convert_to_binary(min_instr)
+
+        self.assertEqual(max_binary & 0x3F, self.asm.isa_definitions["V_MAX_VF"])
+        self.assertEqual(min_binary & 0x3F, self.asm.isa_definitions["V_MIN_VF"])
+
+    def test_v_topk_encodes_like_masked_vector_op(self):
+        instr = Instruction("V_TOPK", 1, 2, 3, 0, None, None, None)
+        binary = self.asm._convert_to_binary(instr)
+
+        self.assertEqual(binary & 0x3F, self.asm.isa_definitions["V_TOPK"])
+        self.assertEqual((binary >> 6) & 0xF, 1)
+        self.assertEqual((binary >> 10) & 0xF, 2)
+        self.assertEqual((binary >> 14) & 0xF, 3)
+        self.assertEqual((binary >> 18) & 0xF, 0)
+
 
 if __name__ == "__main__":
     unittest.main()

aten/models/__init__.py

@@ -0,0 +1 @@
+"""Model-specific reference helpers for ATen/PLENA bring-up."""

aten/models/gpt_oss/__init__.py

@@ -0,0 +1 @@
+"""GPT-OSS model semantics and reference utilities."""