From b71035b3d1d8bcbea008ee7718dc5880e74cb7e9 Mon Sep 17 00:00:00 2001 From: Michael Date: Thu, 25 Jun 2026 23:54:38 +0100 Subject: [PATCH 1/3] feat: add routed MoE compiler substrate --- asm_templates/store_act_asm.py | 18 +- assembler/assembly_to_binary.py | 15 +- assembler/parser.py | 4 +- assembler/tests/test_vector_rmask_handling.py | 20 + aten/gpt_oss_moe.py | 470 +++++ aten/gpt_oss_real_layer0_utils.py | 157 ++ aten/plena/compiler.py | 2 + aten/plena/isa_attention.py | 52 +- aten/plena/isa_compiler.py | 7 +- aten/plena/isa_matrix.py | 362 +++- aten/plena/isa_tile_rows.py | 24 + aten/plena/program_attention.py | 135 +- aten/plena/program_fp_tile_ops.py | 45 +- aten/plena/program_gpt_oss_moe.py | 1812 +++++++++++++++++ aten/plena/program_matrix_ops.py | 374 +++- aten/plena/program_tensors.py | 25 +- aten/reference.py | 2 +- aten/tests/test_gpt_oss_moe_assertions.py | 82 + aten/tests/test_gpt_oss_moe_reference.py | 321 +++ aten/tests/test_plena_compiler.py | 119 +- doc/operation.svh | 7 +- doc/plena_isa_spec.md | 42 + 22 files changed, 4053 insertions(+), 42 deletions(-) create mode 100644 aten/gpt_oss_moe.py create mode 100644 aten/gpt_oss_real_layer0_utils.py create mode 100644 aten/plena/program_gpt_oss_moe.py create mode 100644 aten/tests/test_gpt_oss_moe_assertions.py create mode 100644 aten/tests/test_gpt_oss_moe_reference.py diff --git a/asm_templates/store_act_asm.py b/asm_templates/store_act_asm.py index 114d274..21a1a0a 100644 --- a/asm_templates/store_act_asm.py +++ b/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 diff --git a/assembler/assembly_to_binary.py b/assembler/assembly_to_binary.py index 4ea6746..17a3c43 100644 --- a/assembler/assembly_to_binary.py +++ b/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( { diff --git a/assembler/parser.py b/assembler/parser.py index 409bb07..486c1ba 100644 --- a/assembler/parser.py +++ b/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): diff --git a/assembler/tests/test_vector_rmask_handling.py b/assembler/tests/test_vector_rmask_handling.py index b479884..7d7bade 100644 --- a/assembler/tests/test_vector_rmask_handling.py +++ b/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() diff --git a/aten/gpt_oss_moe.py b/aten/gpt_oss_moe.py new file mode 100644 index 0000000..22a585e --- /dev/null +++ b/aten/gpt_oss_moe.py @@ -0,0 +1,470 @@ +"""GPT-OSS MoE reference helpers for PLENA bring-up. + +These helpers intentionally separate the GPT-OSS mathematical contract from +PLENA's current HBM precision. Golden A uses high-precision weights and fixed +GPT-OSS semantics. Golden B reuses Golden A's routing decision and applies the +current PLENA MXFP8 weight model to expert weights only. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class GptOssMoeResult: + output: torch.Tensor + router_logits: torch.Tensor + topk_indices: torch.Tensor + topk_weights: torch.Tensor + gate_up_preact: torch.Tensor + + +@dataclass(frozen=True) +class ClampStats: + gate_min: float + gate_max: float + up_min: float + up_max: float + inactive: bool + + +@dataclass(frozen=True) +class SplitGateUp: + gate_weight: torch.Tensor + up_weight: torch.Tensor + gate_bias: torch.Tensor | None + up_bias: torch.Tensor | None + + +@dataclass(frozen=True) +class CompareStats: + rel_rms: float + atol: float + rtol: float + allclose: bool + pass_rate: float + max_abs_error: float + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def _maybe_bf16(x: torch.Tensor, enabled: bool) -> torch.Tensor: + return _bf16(x) if enabled else x.float() + + +def quantize_to_plena_mxfp8(tensor: torch.Tensor) -> torch.Tensor: + """Quantize/dequantize through PLENA's current HBM MXFP8 format.""" + from plena_quant.quantizer.hardware_quantizer.mxfp import _mx_fp_quantize_hardware + + orig_shape = tensor.shape + tensor_2d = tensor.float().reshape(-1, tensor.shape[-1]) + bm_x, _, _, _ = _mx_fp_quantize_hardware( + tensor_2d, + width=8, + exponent_width=4, + exponent_bias_width=8, + block_size=[1, 8], + ) + return bm_x.reshape(orig_shape) + + +def split_packed_gate_up( + gate_up_weight: torch.Tensor, + gate_up_bias: torch.Tensor | None = None, +) -> SplitGateUp: + """Split HF/OpenAI GPT-OSS packed gate/up tensors into PLENA tensors. + + HF stores both expert projections in one interleaved tensor: + + gate_up_proj[..., 0::2] -> gate + gate_up_proj[..., 1::2] -> up + + The packed bias uses the same convention. Do not split this tensor into + contiguous halves; doing so is shape-compatible but numerically wrong. + """ + if gate_up_weight.shape[-1] % 2 != 0: + raise ValueError(f"gate_up_weight last dim must be even, got {gate_up_weight.shape[-1]}") + if gate_up_bias is not None and gate_up_bias.shape[-1] != gate_up_weight.shape[-1]: + raise ValueError( + "gate_up_bias last dim must match packed gate_up_weight last dim " + f"({gate_up_bias.shape[-1]} != {gate_up_weight.shape[-1]})" + ) + + return SplitGateUp( + gate_weight=gate_up_weight[..., 0::2].contiguous(), + up_weight=gate_up_weight[..., 1::2].contiguous(), + gate_bias=None if gate_up_bias is None else gate_up_bias[..., 0::2].contiguous(), + up_bias=None if gate_up_bias is None else gate_up_bias[..., 1::2].contiguous(), + ) + + +def compare_stats( + actual: torch.Tensor, + reference: torch.Tensor, + *, + rtol: float, + atol_scale: float = 0.01, +) -> CompareStats: + """Return GPT-OSS real-layer comparison metrics. + + ``rel_rms`` is ||actual-reference||_2 / ||reference||_2. The elementwise + check uses ``atol = atol_scale * std(reference)`` plus the caller-provided + relative tolerance. + """ + actual_f = actual.float() + reference_f = reference.float() + diff = actual_f - reference_f + denom = torch.linalg.vector_norm(reference_f).clamp_min(1e-12) + rel_rms = float((torch.linalg.vector_norm(diff) / denom).item()) + atol = float((reference_f.std(unbiased=False) * atol_scale).item()) + allowed = atol + rtol * reference_f.abs() + passed = diff.abs() <= allowed + return CompareStats( + rel_rms=rel_rms, + atol=atol, + rtol=rtol, + allclose=bool(passed.all().item()), + pass_rate=float(passed.float().mean().item()), + max_abs_error=float(diff.abs().max().item()), + ) + + +def assert_compare_within( + actual: torch.Tensor, + reference: torch.Tensor, + *, + name: str, + max_rel_rms: float, + rtol: float, + atol_scale: float = 0.01, +) -> CompareStats: + stats = compare_stats(actual, reference, rtol=rtol, atol_scale=atol_scale) + if stats.rel_rms > max_rel_rms or not stats.allclose: + raise AssertionError( + f"{name} failed: rel_rms={stats.rel_rms:.6g} (limit {max_rel_rms:.6g}), " + f"allclose={stats.allclose}, pass_rate={stats.pass_rate:.2%}, " + f"atol={stats.atol:.6g}, rtol={stats.rtol:.6g}, max_abs={stats.max_abs_error:.6g}" + ) + return stats + + +def assert_gap_in_band( + actual: torch.Tensor, + reference: torch.Tensor, + *, + name: str, + min_rel_rms: float, + max_rel_rms: float, +) -> CompareStats: + stats = compare_stats(actual, reference, rtol=max_rel_rms) + if not (min_rel_rms <= stats.rel_rms <= max_rel_rms): + raise AssertionError( + f"{name} gap outside expected band: rel_rms={stats.rel_rms:.6g}, " + f"expected [{min_rel_rms:.6g}, {max_rel_rms:.6g}]" + ) + return stats + + +def gpt_oss_swiglu( + gate_up: torch.Tensor, + *, + limit: float = 7.0, + alpha: float = 1.702, + bf16_intermediates: bool = False, + apply_clamp: bool = True, +) -> torch.Tensor: + """GPT-OSS expert activation. + + The gate/up projection is interleaved in the final dimension: + even lanes are gate, odd lanes are up. GPT-OSS clamps gate only on the + upper side and clamps up on both sides before applying the gated activation: + + gate * sigmoid(alpha * gate) * (up + 1) + """ + gate = gate_up[..., ::2].float() + up = gate_up[..., 1::2].float() + + if bf16_intermediates: + gate = _bf16(gate) + up = _bf16(up) + + if apply_clamp: + gate = torch.clamp(gate, max=limit) + up = torch.clamp(up, min=-limit, max=limit) + if bf16_intermediates: + gate = _bf16(gate) + up = _bf16(up) + + scaled_gate = _maybe_bf16(gate * alpha, bf16_intermediates) + sigmoid = _maybe_bf16(torch.sigmoid(scaled_gate), bf16_intermediates) + glu = _maybe_bf16(gate * sigmoid, bf16_intermediates) + up_plus_one = _maybe_bf16(up + 1.0, bf16_intermediates) + return _maybe_bf16(up_plus_one * glu, bf16_intermediates) + + +def clamp_stats(gate_up_preact: torch.Tensor, *, limit: float = 7.0) -> ClampStats: + """Return whether GPT-OSS clamp would be a no-op for these preactivations.""" + gate = gate_up_preact[..., ::2].float() + up = gate_up_preact[..., 1::2].float() + gate_min = float(gate.min().item()) + gate_max = float(gate.max().item()) + up_min = float(up.min().item()) + up_max = float(up.max().item()) + return ClampStats( + gate_min=gate_min, + gate_max=gate_max, + up_min=up_min, + up_max=up_max, + inactive=gate_max <= limit and up_min >= -limit and up_max <= limit, + ) + + +def assert_clamp_inactive(gate_up_preact: torch.Tensor, *, limit: float = 7.0) -> ClampStats: + """Validate the route-A smoke-test precondition. + + Gate only has an upper clamp in GPT-OSS; up has a two-sided clamp. A smoke + test may skip clamp only when those exact clamp sites would be no-ops. + """ + stats = clamp_stats(gate_up_preact, limit=limit) + if not stats.inactive: + raise AssertionError( + "GPT-OSS clamp would be active: " + f"gate range [{stats.gate_min:.6g}, {stats.gate_max:.6g}], " + f"up range [{stats.up_min:.6g}, {stats.up_max:.6g}], limit={limit}" + ) + return stats + + +def gpt_oss_moe_golden_a( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_bias: torch.Tensor, + down_weight: torch.Tensor, + down_bias: torch.Tensor | None = None, + *, + experts_per_token: int = 4, + swiglu_limit: float = 7.0, + swiglu_alpha: float = 1.702, + bf16_intermediates: bool = True, + apply_clamp: bool = True, +) -> GptOssMoeResult: + """Pure GPT-OSS MoE math with high-precision weights. + + Shapes use PLENA-friendly row-major matrices: + x: [tokens, hidden] + router_weight: [hidden, experts] + router_bias: [experts] + gate_up_weight: [experts, hidden, 2 * intermediate] + gate_up_bias: [experts, 2 * intermediate] + down_weight: [experts, intermediate, hidden] + down_bias: [experts, hidden] or None + + This function does not use HF's default MXFP4 execution path. If weights + come from a quantized checkpoint, callers must pass already-dequantized + tensors when they want a high-precision Golden A. + """ + x_ref = _maybe_bf16(x, bf16_intermediates) + router_logits = x_ref.float() @ router_weight.float() + router_bias.float() + topk_values, topk_indices = torch.topk(router_logits, k=experts_per_token, dim=-1) + topk_weights = torch.softmax(topk_values.float(), dim=-1) + if bf16_intermediates: + topk_weights = _bf16(topk_weights) + + output, gate_up_preact = _run_selected_experts( + x_ref, + topk_indices, + topk_weights, + gate_up_weight, + gate_up_bias, + down_weight, + down_bias, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + bf16_intermediates=bf16_intermediates, + apply_clamp=apply_clamp, + ) + return GptOssMoeResult(output, router_logits, topk_indices, topk_weights, gate_up_preact) + + +def gpt_oss_moe_golden_b_plena_mxfp8( + x: torch.Tensor, + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_bias: torch.Tensor, + down_weight: torch.Tensor, + down_bias: torch.Tensor | None = None, + *, + swiglu_limit: float = 7.0, + swiglu_alpha: float = 1.702, + bf16_intermediates: bool = True, + apply_clamp: bool = True, +) -> GptOssMoeResult: + """PLENA-aware Golden B with fixed Golden-A routing. + + Router logits and top-k are deliberately not recomputed here. This isolates + expert compute plus weighted combine from the fragile discrete routing + decision. Expert matrix weights are quantized through PLENA's current HBM + MXFP8 model; router precision is outside this function by construction. + """ + x_ref = _maybe_bf16(x, bf16_intermediates) + gate_up_q = quantize_to_plena_mxfp8(gate_up_weight) + down_q = quantize_to_plena_mxfp8(down_weight) + weights_ref = _maybe_bf16(topk_weights, bf16_intermediates) + output, gate_up_preact = _run_selected_experts( + x_ref, + topk_indices, + weights_ref, + gate_up_q, + gate_up_bias, + down_q, + down_bias, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + bf16_intermediates=bf16_intermediates, + apply_clamp=apply_clamp, + ) + router_logits = torch.empty(x.shape[0], 0, dtype=torch.float32, device=x.device) + return GptOssMoeResult(output, router_logits, topk_indices, weights_ref, gate_up_preact) + + +def gpt_oss_moe_fixed_routing_host_smoke( + x: torch.Tensor, + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_bias: torch.Tensor, + down_weight: torch.Tensor, + down_bias: torch.Tensor | None = None, + *, + quantize_expert_weights_to_plena_mxfp8: bool = False, + swiglu_limit: float = 7.0, + swiglu_alpha: float = 1.702, + bf16_intermediates: bool = True, + apply_clamp: bool = True, +) -> GptOssMoeResult: + """Host-side fixed-routing MoE smoke model. + + This mirrors the v0 dataflow deliberately: routing decisions are fixed by + the caller, tokens are grouped by expert on the host, each selected expert + is evaluated, and weighted outputs are accumulated back to token order. It + exists to test wiring separately from the fragile top-k decision. + """ + x_ref = _maybe_bf16(x, bf16_intermediates) + weights_ref = _maybe_bf16(topk_weights, bf16_intermediates) + if quantize_expert_weights_to_plena_mxfp8: + gate_up_weight = quantize_to_plena_mxfp8(gate_up_weight) + down_weight = quantize_to_plena_mxfp8(down_weight) + + tokens, hidden = x.shape + experts_per_token = topk_indices.shape[1] + gate_up_width = gate_up_weight.shape[-1] + # Match HF's standalone GptOssExperts CPU fallback for unquantized Golden A: + # expert-id sorted loop, BF16 matmul outputs, and BF16 index_add_ + # accumulation. The quantized Golden-B smoke intentionally keeps the older + # FP32 matmul + explicit BF16-round model so it stays bit-equivalent to + # gpt_oss_moe_golden_b_plena_mxfp8. + hf_cpu_equiv = bf16_intermediates and not quantize_expert_weights_to_plena_mxfp8 + accum_dtype = torch.bfloat16 if hf_cpu_equiv else torch.float32 + next_states = torch.zeros(tokens, hidden, dtype=accum_dtype, device=x.device) + gate_up_preact = torch.zeros( + tokens, + experts_per_token, + gate_up_width, + dtype=torch.float32, + device=x.device, + ) + + for expert_idx_tensor in torch.unique(topk_indices): + expert_idx = int(expert_idx_tensor.item()) + token_idx, top_k_pos = torch.where(topk_indices == expert_idx) + if token_idx.numel() == 0: + continue + + if hf_cpu_equiv: + current_state = x_ref[token_idx].to(torch.bfloat16) + gate_up = ( + current_state @ gate_up_weight[expert_idx].to(torch.bfloat16) + + gate_up_bias[expert_idx].to(torch.bfloat16) + ) + else: + current_state = x_ref[token_idx] + gate_up = current_state.float() @ gate_up_weight[expert_idx].float() + gate_up_bias[expert_idx].float() + gate_up = _maybe_bf16(gate_up, bf16_intermediates) + gate_up_preact[token_idx, top_k_pos] = gate_up.float() + + activated = gpt_oss_swiglu( + gate_up, + limit=swiglu_limit, + alpha=swiglu_alpha, + bf16_intermediates=bf16_intermediates, + apply_clamp=apply_clamp, + ) + if hf_cpu_equiv: + expert_out = activated.to(torch.bfloat16) @ down_weight[expert_idx].to(torch.bfloat16) + if down_bias is not None: + expert_out = expert_out + down_bias[expert_idx].to(torch.bfloat16) + else: + expert_out = activated.float() @ down_weight[expert_idx].float() + if down_bias is not None: + expert_out = expert_out + down_bias[expert_idx].float() + expert_out = _maybe_bf16(expert_out, bf16_intermediates) + + weighted = expert_out.to(accum_dtype) * weights_ref[token_idx, top_k_pos].to(accum_dtype).unsqueeze(-1) + next_states.index_add_(0, token_idx, weighted.to(accum_dtype)) + + next_states = _maybe_bf16(next_states, bf16_intermediates) + router_logits = torch.empty(tokens, 0, dtype=torch.float32, device=x.device) + return GptOssMoeResult(next_states, router_logits, topk_indices, weights_ref, gate_up_preact) + + +def _run_selected_experts( + x: torch.Tensor, + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + gate_up_weight: torch.Tensor, + gate_up_bias: torch.Tensor, + down_weight: torch.Tensor, + down_bias: torch.Tensor | None, + *, + swiglu_limit: float, + swiglu_alpha: float, + bf16_intermediates: bool, + apply_clamp: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + tokens, hidden = x.shape + experts_per_token = topk_indices.shape[1] + intermediate = down_weight.shape[1] + + selected_gate_up_w = gate_up_weight[topk_indices].float() + selected_gate_up_b = gate_up_bias[topk_indices].float() + gate_up = torch.einsum("th,tkeh->tke", x.float(), selected_gate_up_w.transpose(-1, -2)) + gate_up = gate_up + selected_gate_up_b + gate_up = _maybe_bf16(gate_up, bf16_intermediates) + + activated = gpt_oss_swiglu( + gate_up, + limit=swiglu_limit, + alpha=swiglu_alpha, + bf16_intermediates=bf16_intermediates, + apply_clamp=apply_clamp, + ) + assert activated.shape == (tokens, experts_per_token, intermediate) + + selected_down_w = down_weight[topk_indices].float() + expert_out = torch.einsum("tki,tkih->tkh", activated.float(), selected_down_w) + if down_bias is not None: + expert_out = expert_out + down_bias[topk_indices].float() + expert_out = _maybe_bf16(expert_out, bf16_intermediates) + + combined = (expert_out.float() * topk_weights.float().unsqueeze(-1)).sum(dim=1) + combined = _maybe_bf16(combined, bf16_intermediates) + assert combined.shape == (tokens, hidden) + return combined, gate_up diff --git a/aten/gpt_oss_real_layer0_utils.py b/aten/gpt_oss_real_layer0_utils.py new file mode 100644 index 0000000..bad9f39 --- /dev/null +++ b/aten/gpt_oss_real_layer0_utils.py @@ -0,0 +1,157 @@ +"""Reusable GPT-OSS real-layer loading helpers. + +These helpers are shared by the correctness tests that need a real GPT-OSS +layer without making those tests depend on the standalone checkpoint probe +script. They intentionally use cached HuggingFace files only; callers decide +which tensors to load and how to compare them. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +from functools import lru_cache +from pathlib import Path + +import torch +from huggingface_hub import hf_hub_download +from safetensors import safe_open + + +REPO_ID = "openai/gpt-oss-20b" +SHARD0 = "model-00000-of-00002.safetensors" +SHARD2 = "model-00002-of-00002.safetensors" +SHARDS = (SHARD0, SHARD2) +_TORCHVISION_NMS_LIB = None + + +def ensure_torchvision_nms_schema() -> None: + """Define a minimal torchvision::nms schema when torchvision C ops are absent.""" + global _TORCHVISION_NMS_LIB + if _TORCHVISION_NMS_LIB is not None: + return + try: + if importlib.util.find_spec("torchvision._C") is not None: + return + except Exception: + pass + try: + _TORCHVISION_NMS_LIB = torch.library.Library("torchvision", "DEF") + _TORCHVISION_NMS_LIB.define("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor") + except Exception: + pass + + +def cached_file(filename: str) -> Path: + return Path(hf_hub_download(REPO_ID, filename, local_files_only=True)) + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def snapshot_revision(path: Path) -> str | None: + parts = path.parts + try: + idx = parts.index("snapshots") + except ValueError: + return None + if idx + 1 >= len(parts): + return None + return parts[idx + 1] + + +def load_json(filename: str) -> dict: + return json.loads(cached_file(filename).read_text()) + + +@lru_cache(maxsize=1) +def weight_map() -> dict[str, str]: + index_path = cached_file("model.safetensors.index.json") + return json.loads(index_path.read_text())["weight_map"] + + +def load_sharded_tensor(tensor_name: str) -> torch.Tensor: + filename = weight_map().get(tensor_name) + candidates = (filename,) if filename is not None else SHARDS + for candidate in candidates: + if candidate is None: + continue + shard_path = cached_file(candidate) + with safe_open(shard_path, framework="pt", device="cpu") as f: + if tensor_name in f.keys(): + return f.get_tensor(tensor_name) + raise KeyError(f"tensor {tensor_name!r} not found in cached GPT-OSS shards") + + +def load_token_ids(prompt: str, max_tokens: int) -> torch.Tensor: + ensure_torchvision_nms_schema() + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(REPO_ID, local_files_only=True) + token_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids[0] + if token_ids.numel() == 0: + raise ValueError("Prompt tokenized to zero tokens") + return token_ids[:max_tokens].to(torch.long) + + +def slice_rows(shard_path: Path, tensor_name: str, rows: torch.Tensor) -> torch.Tensor: + pieces = [] + with safe_open(shard_path, framework="pt", device="cpu") as f: + tensor_slice = f.get_slice(tensor_name) + for row in rows.tolist(): + pieces.append(tensor_slice[int(row) : int(row) + 1]) + return torch.cat(pieces, dim=0) + + +def rms_norm(hidden_states: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + input_dtype = hidden_states.dtype + x = hidden_states.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + return (weight.float() * x).to(input_dtype) + + +def load_layer_tensors(layer_index: int) -> dict[str, torch.Tensor]: + from transformers.integrations.mxfp4 import convert_moe_packed_tensors + + prefix = f"model.layers.{layer_index}" + gate_blocks = load_sharded_tensor(f"{prefix}.mlp.experts.gate_up_proj_blocks") + gate_scales = load_sharded_tensor(f"{prefix}.mlp.experts.gate_up_proj_scales") + down_blocks = load_sharded_tensor(f"{prefix}.mlp.experts.down_proj_blocks") + down_scales = load_sharded_tensor(f"{prefix}.mlp.experts.down_proj_scales") + tensors = { + "router_weight": load_sharded_tensor(f"{prefix}.mlp.router.weight").to(torch.bfloat16), + "router_bias": load_sharded_tensor(f"{prefix}.mlp.router.bias").to(torch.bfloat16), + "gate_up_bias": load_sharded_tensor(f"{prefix}.mlp.experts.gate_up_proj_bias").to(torch.bfloat16), + "down_bias": load_sharded_tensor(f"{prefix}.mlp.experts.down_proj_bias").to(torch.bfloat16), + "input_layernorm_weight": load_sharded_tensor(f"{prefix}.input_layernorm.weight").to(torch.bfloat16), + } + + tensors["gate_up_weight"] = convert_moe_packed_tensors(gate_blocks, gate_scales).to(torch.bfloat16) + tensors["down_weight"] = convert_moe_packed_tensors(down_blocks, down_scales).to(torch.bfloat16) + return tensors + + +def load_layer0_tensors() -> dict[str, torch.Tensor]: + return load_layer_tensors(0) + + +def build_hf_mlp(config, tensors: dict[str, torch.Tensor]): + ensure_torchvision_nms_schema() + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssMLP + + mlp = GptOssMLP(config).eval().to(torch.bfloat16) + with torch.no_grad(): + mlp.router.weight.copy_(tensors["router_weight"]) + mlp.router.bias.copy_(tensors["router_bias"]) + mlp.experts.gate_up_proj.copy_(tensors["gate_up_weight"]) + mlp.experts.gate_up_proj_bias.copy_(tensors["gate_up_bias"]) + mlp.experts.down_proj.copy_(tensors["down_weight"]) + mlp.experts.down_proj_bias.copy_(tensors["down_bias"]) + return mlp diff --git a/aten/plena/compiler.py b/aten/plena/compiler.py index e0ab4cf..6940246 100644 --- a/aten/plena/compiler.py +++ b/aten/plena/compiler.py @@ -8,6 +8,7 @@ from compiler.aten.plena.isa_compiler import IsaCompiler from compiler.aten.plena.program_attention import ProgramAttentionMixin from compiler.aten.plena.program_fp_tile_ops import ProgramFPTileOpsMixin +from compiler.aten.plena.program_gpt_oss_moe import ProgramGptOssMoeMixin from compiler.aten.plena.program_matrix_ops import ProgramMatrixOpsMixin from compiler.aten.plena.program_tensors import ProgramTensorMixin from compiler.aten.plena.vars import FPVar, InputVar, TensorVar @@ -56,6 +57,7 @@ class PlenaCompiler( ProgramTensorMixin, ProgramFPTileOpsMixin, ProgramMatrixOpsMixin, + ProgramGptOssMoeMixin, ProgramAttentionMixin, IsaCompiler, ): diff --git a/aten/plena/isa_attention.py b/aten/plena/isa_attention.py index f6f1340..5ad4019 100644 --- a/aten/plena/isa_attention.py +++ b/aten/plena/isa_attention.py @@ -20,6 +20,7 @@ def _online_softmax_asm( scale: float = 1.0, rows: int | None = None, valid_cols: int | None = None, + sink_address: int | None = None, ) -> str: """ Online Softmax Computation. @@ -43,6 +44,7 @@ def _online_softmax_asm( m_start_address=m_start_address, scale=scale, valid_cols=valid_cols, + sink_address=sink_address, ) gp_regs = self.register_allocator.allocate_gp(5) @@ -62,6 +64,8 @@ def _online_softmax_asm( fp_sum_p = 4 # f4: sum(P) fp_scale = 5 # f5: scale factor fp_row_max = 6 # f6: current row max (temporary) + fp_sink = 7 # f7: GPT-OSS attention sink logit + fp_sink_exp = fp_row_max # reuse f6 after row max is folded into m_curr lines = [] lines.append("; === Online Softmax ===") @@ -84,6 +88,8 @@ def _online_softmax_asm( # scale factor is pre-loaded at FP SRAM addr 1 by the flash-attention driver. if scale != 1.0: lines.append(f"S_LD_FP f{fp_scale}, gp0, 1") + if sink_address is not None: + lines.append(f"S_LD_FP f{fp_sink}, gp0, {sink_address}") loop_rows = mlen if rows is None else rows lines.append(f"C_LOOP_START gp{gp_loop}, {loop_rows}") @@ -100,9 +106,14 @@ def _online_softmax_asm( # m_curr = max(row_max, m_old) — online softmax must retain the running max. lines.append(f"S_MAX_FP f{fp_m_old}, f{fp_row_max}, f{fp_m_old}") + if sink_address is not None: + lines.append(f"S_MAX_FP f{fp_m_old}, f{fp_sink}, f{fp_m_old}") lines.append(f"S_SUB_FP f{fp_m_res}, f{fp_m_res}, f{fp_m_old}") lines.append(f"S_EXP_FP f{fp_m_res}, f{fp_m_res}, 0") + if sink_address is not None: + lines.append(f"S_SUB_FP f{fp_sink_exp}, f{fp_sink}, f{fp_m_old}") + lines.append(f"S_EXP_FP f{fp_sink_exp}, f{fp_sink_exp}, 0") lines.append(f"S_ST_FP f{fp_m_res}, gp{gp_m_res_addr}, 0") lines.append(f"S_ST_FP f{fp_m_old}, gp{gp_m_addr}, 0") @@ -117,6 +128,8 @@ def _online_softmax_asm( lines.append(f"S_MUL_FP f{fp_l_old}, f{fp_l_old}, f{fp_m_res}") lines.append(f"S_ADD_FP f{fp_l_old}, f{fp_l_old}, f{fp_sum_p}") + if sink_address is not None: + lines.append(f"S_ADD_FP f{fp_l_old}, f{fp_l_old}, f{fp_sink_exp}") lines.append(f"S_ST_FP f{fp_l_old}, gp{gp_l_addr}, 0") @@ -136,6 +149,7 @@ def _online_softmax_asm_unrolled( m_start_address: int, scale: float = 1.0, valid_cols: int | None = None, + sink_address: int | None = None, ) -> str: """Legacy Python-unrolled online softmax emission, kept for A/B comparisons.""" gp_regs = self.register_allocator.allocate_gp(5) @@ -151,6 +165,8 @@ def _online_softmax_asm_unrolled( fp_sum_p = 4 fp_scale = 5 fp_row_max = 6 + fp_sink = 7 + fp_sink_exp = fp_row_max lines = [] lines.append("; === Online Softmax ===") @@ -161,6 +177,8 @@ def _online_softmax_asm_unrolled( if scale != 1.0: lines.append(f"S_LD_FP f{fp_scale}, gp0, 1") + if sink_address is not None: + lines.append(f"S_LD_FP f{fp_sink}, gp0, {sink_address}") mask_en = 0 if valid_cols is not None and valid_cols < mlen: @@ -184,9 +202,14 @@ def _online_softmax_asm_unrolled( # m_curr = max(row_max, m_old) — online softmax must retain the running max. lines.append(f"S_MAX_FP f{fp_m_old}, f{fp_row_max}, f{fp_m_old}") + if sink_address is not None: + lines.append(f"S_MAX_FP f{fp_m_old}, f{fp_sink}, f{fp_m_old}") lines.append(f"S_SUB_FP f{fp_m_res}, f{fp_m_res}, f{fp_m_old}") lines.append(f"S_EXP_FP f{fp_m_res}, f{fp_m_res}, 0") + if sink_address is not None: + lines.append(f"S_SUB_FP f{fp_sink_exp}, f{fp_sink}, f{fp_m_old}") + lines.append(f"S_EXP_FP f{fp_sink_exp}, f{fp_sink_exp}, 0") lines.append(f"S_ST_FP f{fp_m_res}, gp{gp_m_res_addr}, {row}") lines.append(f"S_ST_FP f{fp_m_old}, gp{gp_m_addr}, {row}") @@ -201,6 +224,8 @@ def _online_softmax_asm_unrolled( lines.append(f"S_MUL_FP f{fp_l_old}, f{fp_l_old}, f{fp_m_res}") lines.append(f"S_ADD_FP f{fp_l_old}, f{fp_l_old}, f{fp_sum_p}") + if sink_address is not None: + lines.append(f"S_ADD_FP f{fp_l_old}, f{fp_l_old}, f{fp_sink_exp}") lines.append(f"S_ST_FP f{fp_l_old}, gp{gp_l_addr}, {row}") @@ -220,6 +245,8 @@ def _pv_multiply_asm( pv_address: int, rows: int | None = None, pv_physical_rows: int | None = None, + v_hbm_row_stride: int | None = None, + v_hbm_element_bytes: int = 1, ) -> str: """ Compute PV = P @ V via M_MM. @@ -251,6 +278,8 @@ def _pv_multiply_asm( v_hbm_offset=v_hbm_offset, pv_address=pv_address, pv_physical_rows=pv_physical_rows, + v_hbm_row_stride=v_hbm_row_stride, + v_hbm_element_bytes=v_hbm_element_bytes, ) gp_regs = self.register_allocator.allocate_gp(8) @@ -274,7 +303,13 @@ def _pv_multiply_asm( lines.append(f"; V split into {num_v_col_blocks} column blocks of width {mlen}") lines.append("; Storage layout: (batch, mlen, hidden/mlen), column-block major") - # STRIDE was set to mlen by the flash-attention driver — do not overwrite it here. + # H_PREFETCH_M stride is byte-based. MXFP8 is 1 byte/element; Plain + # BF16 KeyValue loads need 2 bytes/element. + if v_hbm_row_stride is None: + v_hbm_row_stride = mlen + lines.append(f"S_ADDI_INT gp{gp_stride}, gp0, {v_hbm_row_stride * v_hbm_element_bytes}") + lines.append(f"C_SET_STRIDE_REG gp{gp_stride}") + # M_MM_WO requires a nonzero stride reg (gp0=0 would be interpreted as stride=1). # With column-block-major storage, consecutive rows within a column block are # adjacent, so the writeback stride = 1. @@ -288,7 +323,7 @@ def _pv_multiply_asm( # Prefetch V[:, v_col_block*mlen:(v_col_block+1)*mlen] (mlen × mlen) to MSRAM. # V is row-major in HBM: V[row, col] at offset row*head_dim + col, so the # column-block base offset = v_hbm_offset + v_col_block * mlen (elements). - v_block_hbm_offset = v_hbm_offset + v_col_block * mlen + v_block_hbm_offset = (v_hbm_offset + v_col_block * mlen) * v_hbm_element_bytes lines.append(f"S_ADDI_INT gp{gp_v}, gp0, 0") lines.extend(load_large_int(gp_hbm, v_block_hbm_offset)) lines.append(f"H_PREFETCH_M gp{gp_v}, gp{gp_hbm}, a{v_hbm_offset_reg}, 1, 1") @@ -321,6 +356,8 @@ def _pv_multiply_asm_unrolled( v_hbm_offset: int, pv_address: int, pv_physical_rows: int | None = None, + v_hbm_row_stride: int | None = None, + v_hbm_element_bytes: int = 1, ) -> str: """Legacy Python-unrolled P @ V emission, kept for A/B comparisons.""" if pv_physical_rows is None: @@ -340,13 +377,17 @@ def _pv_multiply_asm_unrolled( lines.append("; M_MM: (blen, mlen) @ (mlen, blen) -> (blen, blen), K=mlen in one shot") lines.append(f"; V split into {num_v_col_blocks} column blocks of width {mlen}") lines.append("; Storage layout: (batch, mlen, hidden/mlen), column-block major") + if v_hbm_row_stride is None: + v_hbm_row_stride = mlen + lines.append(f"S_ADDI_INT gp{gp_stride}, gp0, {v_hbm_row_stride * v_hbm_element_bytes}") + lines.append(f"C_SET_STRIDE_REG gp{gp_stride}") lines.append(f"S_ADDI_INT gp{gp_stride}, gp0, 1") for v_col_block in range(num_v_col_blocks): lines.append( f"; --- V column block {v_col_block} (columns {v_col_block * mlen} to {(v_col_block + 1) * mlen - 1}) ---" ) - v_block_hbm_offset = v_hbm_offset + v_col_block * mlen + v_block_hbm_offset = (v_hbm_offset + v_col_block * mlen) * v_hbm_element_bytes lines.append(f"S_ADDI_INT gp{gp_v}, gp0, 0") lines.extend(load_large_int(gp_hbm, v_block_hbm_offset)) lines.append(f"H_PREFETCH_M gp{gp_v}, gp{gp_hbm}, a{v_hbm_offset_reg}, 1, 1") @@ -738,6 +779,7 @@ def online_softmax_block( scale: float, rows: int | None = None, valid_cols: int | None = None, + sink_address: int | None = None, ) -> str: """ Run Online Softmax on one S block. @@ -760,6 +802,7 @@ def online_softmax_block( scale=scale, rows=rows, valid_cols=valid_cols, + sink_address=sink_address, ) return self._emit(isa_code) @@ -772,6 +815,7 @@ def compute_pv( pv_matrix: str, head_dim: int, rows: int | None = None, + v_hbm_element_bytes: int = 1, ) -> str: """ Compute PV = P @ V[k_idx]. @@ -811,6 +855,8 @@ def compute_pv( pv_address=pv_address, rows=rows, pv_physical_rows=pv_info.physical_shape[0], + v_hbm_row_stride=physical_head_dim, + v_hbm_element_bytes=v_hbm_element_bytes, ) self.register_allocator.free_gp(gp_regs) diff --git a/aten/plena/isa_compiler.py b/aten/plena/isa_compiler.py index 3a6c8d8..b77aae7 100644 --- a/aten/plena/isa_compiler.py +++ b/aten/plena/isa_compiler.py @@ -132,6 +132,8 @@ def store_to_hbm( vlen: int = 64, precision: int = 0, # 0 = Activation, 1 = KeyValue store_amount: int | None = None, # HBM_V_Writeback_Amount + hbm_element_bytes: int = 1, + hbm_real_data_ratio: float | None = None, ) -> str: """ Write tensor from VRAM back to HBM. @@ -195,13 +197,15 @@ def store_to_hbm( hbm_addr_reg=hbm_addr_reg, stride_size=hidden_size, store_amount=store_amount, + precision=precision, + hbm_element_bytes=hbm_element_bytes, ) if tensor_info.hbm_addr < 0 or tensor_info.hbm_addr != hbm_addr: tensor_info.hbm_addr = hbm_addr # HBM stores the MXFP-expanded size (logical size × real_data_ratio). size = batch_size * hidden_size - tensor_info.hbm_size = int(size * self.real_data_ratio) + tensor_info.hbm_size = int(size * (hbm_real_data_ratio or self.real_data_ratio)) finally: self.register_allocator.free_gp(gp_regs) if need_free_addr: @@ -213,6 +217,7 @@ def store_to_hbm( hbm_addr=hbm_addr, shape=tensor_info.shape, physical_shape=(batch_size, hidden_size), + real_data_ratio=hbm_real_data_ratio or self.real_data_ratio, ) return self._emit(isa_code) diff --git a/aten/plena/isa_matrix.py b/aten/plena/isa_matrix.py index a65e90a..9117e8f 100644 --- a/aten/plena/isa_matrix.py +++ b/aten/plena/isa_matrix.py @@ -41,11 +41,21 @@ def reset_mram(self) -> str: def _default_hbm_gp_regs(self, gp_regs: list[int] | None) -> list[int]: return [1, 2, 3] if gp_regs is None else gp_regs - def _emit_hbm_prefetch_setup(self, asm: IsaBuilder, layout, gp_scale: int, gp_stride: int) -> None: + def _emit_hbm_prefetch_setup( + self, + asm: IsaBuilder, + layout, + gp_scale: int, + gp_stride: int, + *, + set_scale: bool = True, + hbm_element_bytes: int = 1, + ) -> None: rows, cols = layout.physical_shape or layout.full_shape - 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) + 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("C_SET_STRIDE_REG", gp(gp_stride)) def _emit_hbm_subblock_prefetch( @@ -59,15 +69,17 @@ def _emit_hbm_subblock_prefetch( gp_scale: int, gp_mram: int, comment: str | None = None, + precision: int = 0, + hbm_element_bytes: int = 1, ) -> None: sub_block = layout.get_sub_block(row_idx, col_idx) - hbm_offset = sub_block.hbm_offset + hbm_offset = sub_block.hbm_offset * hbm_element_bytes sub_block.mram_addr = mram_addr asm.comment(comment if comment is not None else f"SubBlock [{row_idx}][{col_idx}]: HBM offset = {hbm_offset}") asm.instr("S_ADDI_INT", gp(gp_mram), gp(0), mram_addr) asm.instr("S_ADDI_INT", gp(gp_scale), gp(0), hbm_offset) - asm.instr("H_PREFETCH_M", gp(gp_mram), gp(gp_scale), areg(hbm_addr_reg), 1, 0) + asm.instr("H_PREFETCH_M", gp(gp_mram), gp(gp_scale), areg(hbm_addr_reg), 1, precision) def _emit_hbm_subblock_sequence( self, @@ -78,6 +90,8 @@ def _emit_hbm_subblock_sequence( hbm_addr_reg: int, gp_scale: int, gp_mram: int, + precision: int = 0, + hbm_element_bytes: int = 1, ) -> None: mram_addr = mram_start_addr block_size = self.mlen * self.mlen @@ -91,6 +105,8 @@ def _emit_hbm_subblock_sequence( hbm_addr_reg, gp_scale, gp_mram, + precision=precision, + hbm_element_bytes=hbm_element_bytes, ) mram_addr += block_size @@ -102,6 +118,9 @@ def load_sub_matrix_asm( mram_dest_addr: int, hbm_addr_reg: int = 1, gp_regs: list[int] | None = None, + precision: int = 0, + set_scale: bool = True, + hbm_element_bytes: int = 1, ) -> str: """Emit HBM->MRAM prefetch for one mlen x mlen sub-block.""" gp_regs = self._default_hbm_gp_regs(gp_regs) @@ -112,7 +131,14 @@ def load_sub_matrix_asm( gp_scale = gp_regs[0] gp_stride = gp_regs[1] gp_mram = gp_regs[2] - self._emit_hbm_prefetch_setup(asm, layout, gp_scale, gp_stride) + self._emit_hbm_prefetch_setup( + asm, + layout, + gp_scale, + gp_stride, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) hbm_offset = layout.get_sub_block(row_idx, col_idx).hbm_offset self._emit_hbm_subblock_prefetch( asm, @@ -124,6 +150,8 @@ def load_sub_matrix_asm( gp_scale, gp_mram, comment=f"HBM offset: {hbm_offset} (precomputed)", + precision=precision, + hbm_element_bytes=hbm_element_bytes, ) return asm.render() @@ -135,6 +163,9 @@ def load_row_sub_matrices_asm( mram_start_addr: int, hbm_addr_reg: int = 1, gp_regs: list[int] | None = None, + precision: int = 0, + set_scale: bool = True, + hbm_element_bytes: int = 1, ) -> str: """Emit HBM->MRAM prefetches for one block row.""" gp_regs = self._default_hbm_gp_regs(gp_regs) @@ -146,7 +177,14 @@ def load_row_sub_matrices_asm( gp_scale = gp_regs[0] gp_stride = gp_regs[1] gp_mram = gp_regs[2] - self._emit_hbm_prefetch_setup(asm, layout, gp_scale, gp_stride) + self._emit_hbm_prefetch_setup( + asm, + layout, + gp_scale, + gp_stride, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) self._emit_hbm_subblock_sequence( asm, @@ -156,6 +194,8 @@ def load_row_sub_matrices_asm( hbm_addr_reg, gp_scale, gp_mram, + precision=precision, + hbm_element_bytes=hbm_element_bytes, ) return asm.render() @@ -169,6 +209,9 @@ def load_col_sub_matrices_asm( gp_regs: list[int] | None = None, k_block_start: int = 0, k_block_count: int | None = None, + precision: int = 0, + set_scale: bool = True, + hbm_element_bytes: int = 1, ) -> str: """Emit HBM->MRAM prefetches for one block column or K-split slice.""" gp_regs = self._default_hbm_gp_regs(gp_regs) @@ -181,7 +224,14 @@ def load_col_sub_matrices_asm( gp_scale = gp_regs[0] gp_stride = gp_regs[1] gp_mram = gp_regs[2] - self._emit_hbm_prefetch_setup(asm, layout, gp_scale, gp_stride) + self._emit_hbm_prefetch_setup( + asm, + layout, + gp_scale, + gp_stride, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) effective_count = k_block_count if k_block_count is not None else num_row_blocks self._emit_hbm_subblock_sequence( @@ -192,6 +242,8 @@ def load_col_sub_matrices_asm( hbm_addr_reg, gp_scale, gp_mram, + precision=precision, + hbm_element_bytes=hbm_element_bytes, ) return asm.render() @@ -302,6 +354,253 @@ def vram_sub_projection_asm( row_loop_count=row_loop_count, ) + def vram_sub_projection_microtile_accumulate_asm( + self, + vram_mat_name: str, + vram_row_idx: int, + mram_mat_name: str, + mram_col_idx: int, + result_vram_addr: int, + micro_row_idx: int, + micro_col_idx: int, + *, + k_block_start: int, + k_block_count: int, + write_out: bool, + gp_regs: list[int] | None = None, + ) -> str: + """Emit M_MM for one 4x4 microtile, optionally flushing with M_MM_WO. + + This is used by the Qwen router precision path. Wide routers cannot + fit the full K dimension in MRAM, and the old path flushed each K chunk + to BF16 VRAM before adding the next chunk. Here the compiler keeps the + matrix-machine accumulator live for one output microtile across + multiple MRAM reloads, then writes it once after the final chunk. + """ + gp_regs = [1, 2, 3] if gp_regs is None else gp_regs + if len(gp_regs) < 3: + raise ValueError(f"vram_sub_projection_microtile_accumulate_asm requires 3 gp regs, got {len(gp_regs)}") + gp_act, gp_mat, gp_result = gp_regs[:3] + + vram_layout, mram_layout, vram_row_blocks = self._projection_context( + vram_mat_name, vram_row_idx, mram_mat_name + ) + mram_col_blocks = mram_layout.get_col_blocks(mram_col_idx)[k_block_start : k_block_start + k_block_count] + if len(mram_col_blocks) != k_block_count: + raise ValueError( + f"Dimension mismatch: expected {k_block_count} MRAM blocks, got {len(mram_col_blocks)}" + ) + + tiles_per_mlen = self.mlen // self.blen + if micro_col_idx < 0 or micro_col_idx >= tiles_per_mlen: + raise ValueError(f"micro_col_idx={micro_col_idx} outside 0..{tiles_per_mlen - 1}") + + full_batch = (vram_layout.physical_shape or vram_layout.full_shape)[0] + 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 = min(tiles_per_mlen, math.ceil(valid_rows / self.blen)) + if micro_row_idx < 0 or micro_row_idx >= row_loop_count: + raise ValueError(f"micro_row_idx={micro_row_idx} outside 0..{row_loop_count - 1}") + + vram_hidden_block_stride = full_batch * self.mlen + mram_hidden_block_stride = self.mlen * self.mlen + output_row_stride = self.blen * self.mlen + mat_col_stride = self.blen + + vram_row_start_addr = vram_row_blocks[k_block_start].vram_addr + micro_row_idx * output_row_stride + mram_col_start_addr = self._loaded_mram_start( + mram_col_blocks, + lambda block: f"{mram_mat_name}[{block.row_idx}][{mram_col_idx}]", + ) + micro_col_idx * mat_col_stride + result_addr = result_vram_addr + micro_col_idx * self.blen + micro_row_idx * output_row_stride + + lines = [ + f"; VRAM Sub Projection microtile accumulate: {vram_mat_name}[{vram_row_idx}], " + f"{mram_mat_name}[:][{mram_col_idx}], row4={micro_row_idx}, col4={micro_col_idx}, " + f"k=[{k_block_start}, {k_block_start + k_block_count}), write_out={write_out}" + ] + for ih in range(k_block_count): + act_addr = vram_row_start_addr + ih * vram_hidden_block_stride + mat_addr = mram_col_start_addr + ih * mram_hidden_block_stride + lines.extend(load_large_int(gp_act, act_addr)) + lines.extend(load_large_int(gp_mat, mat_addr)) + lines.append(f"M_MM 0, gp{gp_mat}, gp{gp_act}") + + if write_out: + lines.extend(load_large_int(gp_result, result_addr)) + lines.append(f"M_MM_WO gp{gp_result}, gp0, 0") + + return "\n".join(lines) + "\n" + + def vram_sub_projection_microtile_accumulate_to( + self, + vram_mat_name: str, + vram_row_idx: int, + mram_mat_name: str, + mram_col_idx: int, + target_matrix: str, + target_row_idx: int, + target_col_idx: int, + micro_row_idx: int, + micro_col_idx: int, + *, + k_block_start: int, + k_block_count: int, + write_out: bool, + ) -> 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(3) + try: + asm = self.vram_sub_projection_microtile_accumulate_asm( + vram_mat_name=vram_mat_name, + vram_row_idx=vram_row_idx, + mram_mat_name=mram_mat_name, + mram_col_idx=mram_col_idx, + result_vram_addr=result_vram_addr, + micro_row_idx=micro_row_idx, + micro_col_idx=micro_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + write_out=write_out, + gp_regs=gp_regs, + ) + finally: + self.register_allocator.free_gp(gp_regs) + return self._emit(asm) + + def vram_sub_projection_packed_skinny_microtile_accumulate_asm( + self, + vram_mat_name: str, + vram_row_idx: int, + packed_mram_mat_name: str, + packed_group_idx: int, + packed_col_idx: int, + result_vram_addr: int, + micro_row_idx: int, + micro_col_idx: int, + *, + k_block_start: int, + k_block_count: int, + write_out: bool, + gp_regs: list[int] | None = None, + ) -> str: + """Emit one 4x4 projection microtile from a packed-skinny MRAM tile. + + The packed-skinny convention is intentionally narrow and opt-in: one + full ``mlen x mlen`` MRAM tile stores several skinny weight slices for a + single output micro-column. Slice ``i`` lives in columns + ``i*blen:(i+1)*blen`` of that tile. This lets a router probe keep up to + ``mlen / blen`` K-blocks in one MRAM cell without changing M_MM's + existing full-tile read contract. + """ + gp_regs = [1, 2, 3] if gp_regs is None else gp_regs + if len(gp_regs) < 3: + raise ValueError( + "vram_sub_projection_packed_skinny_microtile_accumulate_asm " + f"requires 3 gp regs, got {len(gp_regs)}" + ) + gp_act, gp_mat, gp_result = gp_regs[:3] + if k_block_count <= 0: + raise ValueError(f"k_block_count must be > 0, got {k_block_count}") + + tiles_per_mlen = self.mlen // self.blen + if k_block_count > tiles_per_mlen: + raise ValueError( + f"packed skinny tile can hold at most {tiles_per_mlen} slices, got {k_block_count}" + ) + if micro_col_idx < 0 or micro_col_idx >= tiles_per_mlen: + raise ValueError(f"micro_col_idx={micro_col_idx} outside 0..{tiles_per_mlen - 1}") + + vram_layout, packed_layout, vram_row_blocks = self._projection_context( + vram_mat_name, + vram_row_idx, + packed_mram_mat_name, + ) + if k_block_start < 0 or k_block_start + k_block_count > len(vram_row_blocks): + raise ValueError( + f"k range [{k_block_start}, {k_block_start + k_block_count}) outside " + f"VRAM hidden blocks 0..{len(vram_row_blocks)}" + ) + + packed_block = packed_layout.get_sub_block(packed_group_idx, packed_col_idx) + if packed_block.mram_addr is None: + raise RuntimeError( + f"Packed skinny block {packed_mram_mat_name}[{packed_group_idx}][{packed_col_idx}] " + "not loaded to MRAM" + ) + + full_batch = (vram_layout.physical_shape or vram_layout.full_shape)[0] + 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 = min(tiles_per_mlen, math.ceil(valid_rows / self.blen)) + if micro_row_idx < 0 or micro_row_idx >= row_loop_count: + raise ValueError(f"micro_row_idx={micro_row_idx} outside 0..{row_loop_count - 1}") + + vram_hidden_block_stride = full_batch * self.mlen + output_row_stride = self.blen * self.mlen + vram_row_start_addr = vram_row_blocks[k_block_start].vram_addr + micro_row_idx * output_row_stride + result_addr = result_vram_addr + micro_col_idx * self.blen + micro_row_idx * output_row_stride + + lines = [ + f"; VRAM Sub Projection packed skinny microtile: {vram_mat_name}[{vram_row_idx}], " + f"{packed_mram_mat_name}[{packed_group_idx}][{packed_col_idx}], " + f"row4={micro_row_idx}, col4={micro_col_idx}, " + f"k=[{k_block_start}, {k_block_start + k_block_count}), write_out={write_out}", + f"; Packed skinny convention: slice i is at MRAM tile columns i*{self.blen}..(i+1)*{self.blen}", + ] + for ih in range(k_block_count): + act_addr = vram_row_start_addr + ih * vram_hidden_block_stride + mat_addr = packed_block.mram_addr + ih * self.blen + lines.extend(load_large_int(gp_act, act_addr)) + lines.extend(load_large_int(gp_mat, mat_addr)) + lines.append(f"M_MM 0, gp{gp_mat}, gp{gp_act}") + + if write_out: + lines.extend(load_large_int(gp_result, result_addr)) + lines.append(f"M_MM_WO gp{gp_result}, gp0, 0") + + return "\n".join(lines) + "\n" + + def vram_sub_projection_packed_skinny_microtile_accumulate_to( + self, + vram_mat_name: str, + vram_row_idx: int, + packed_mram_mat_name: str, + packed_group_idx: int, + packed_col_idx: int, + target_matrix: str, + target_row_idx: int, + target_col_idx: int, + micro_row_idx: int, + micro_col_idx: int, + *, + k_block_start: int, + k_block_count: int, + write_out: bool, + ) -> 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(3) + try: + asm = self.vram_sub_projection_packed_skinny_microtile_accumulate_asm( + vram_mat_name=vram_mat_name, + vram_row_idx=vram_row_idx, + packed_mram_mat_name=packed_mram_mat_name, + packed_group_idx=packed_group_idx, + packed_col_idx=packed_col_idx, + result_vram_addr=result_vram_addr, + micro_row_idx=micro_row_idx, + micro_col_idx=micro_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + write_out=write_out, + gp_regs=gp_regs, + ) + finally: + self.register_allocator.free_gp(gp_regs) + return self._emit(asm) + def vram_sub_projection_T_asm( self, vram_mat_name: str, @@ -418,6 +717,9 @@ def load_sub_matrix_row( name: str, row_idx: int, mram_start_addr: int | None = None, + precision: int = 0, + set_scale: bool = True, + hbm_element_bytes: int = 1, ) -> str: """Load entire row sub-blocks from HBM to MRAM: matrix[row_idx][:].""" layout = self.get_hbm_layout(name) @@ -437,6 +739,42 @@ def load_sub_matrix_row( mram_start_addr=mram_start_addr, hbm_addr_reg=addr_reg, gp_regs=gp_regs, + precision=precision, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ), + ) + + def load_sub_matrix( + self, + name: str, + row_idx: int, + col_idx: int, + mram_dest_addr: int | None = None, + precision: int = 0, + set_scale: bool = True, + hbm_element_bytes: int = 1, + ) -> str: + """Load one HBM sub-block into one MRAM tile.""" + layout = self.get_hbm_layout(name) + block_size = self.mlen * self.mlen + + if mram_dest_addr is None: + mram_dest_addr = self.mram_allocator.allocate(f"{name}[{row_idx}][{col_idx}]", block_size) + + return self._emit_hbm_matrix_load( + layout, + 3, + lambda addr_reg, gp_regs: self.load_sub_matrix_asm( + name=name, + row_idx=row_idx, + col_idx=col_idx, + mram_dest_addr=mram_dest_addr, + hbm_addr_reg=addr_reg, + gp_regs=gp_regs, + precision=precision, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, ), ) @@ -447,6 +785,9 @@ def load_sub_matrix_col( mram_start_addr: int | None = None, k_block_start: int = 0, k_block_count: int | None = None, + precision: int = 0, + set_scale: bool = True, + hbm_element_bytes: int = 1, ) -> str: """ Load entire column sub-blocks from HBM to MRAM: matrix[:][col_idx]. @@ -472,6 +813,9 @@ def load_sub_matrix_col( gp_regs=gp_regs, k_block_start=k_block_start, k_block_count=k_block_count, + precision=precision, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, ), ) diff --git a/aten/plena/isa_tile_rows.py b/aten/plena/isa_tile_rows.py index a225209..8d56f53 100644 --- a/aten/plena/isa_tile_rows.py +++ b/aten/plena/isa_tile_rows.py @@ -94,6 +94,24 @@ def tile_row_mul_fp( ) -> str: return self._tile_row_single_matrix_op("tile_row_mul_fp_asm", matrix_name, row_map, tile_row_idx, tile_col_idx) + def tile_row_max_fp( + self, + matrix_name: str, + row_map: list[tuple[int, int]], + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self._tile_row_single_matrix_op("tile_row_max_fp_asm", matrix_name, row_map, tile_row_idx, tile_col_idx) + + def tile_row_min_fp( + self, + matrix_name: str, + row_map: list[tuple[int, int]], + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self._tile_row_single_matrix_op("tile_row_min_fp_asm", matrix_name, row_map, tile_row_idx, tile_col_idx) + def tile_row_add_fp( self, matrix_name: str, @@ -380,6 +398,12 @@ def tile_row_sub_fp_asm(self, vram_addr: int, row_map: list[tuple[int, int]]) -> def tile_row_mul_fp_asm(self, vram_addr: int, row_map: list[tuple[int, int]]) -> str: return self._emit_tile_row_fp_scalar("Mul", "V_MUL_VF", vram_addr, row_map, opcode_extra_args=(0,)) + def tile_row_max_fp_asm(self, vram_addr: int, row_map: list[tuple[int, int]]) -> str: + return self._emit_tile_row_fp_scalar("Max", "V_MAX_VF", vram_addr, row_map, opcode_extra_args=(0,)) + + def tile_row_min_fp_asm(self, vram_addr: int, row_map: list[tuple[int, int]]) -> str: + return self._emit_tile_row_fp_scalar("Min", "V_MIN_VF", vram_addr, row_map, opcode_extra_args=(0,)) + def tile_row_add_fp_asm(self, vram_addr: int, row_map: list[tuple[int, int]]) -> str: return self._emit_tile_row_fp_scalar("Add", "V_ADD_VF", vram_addr, row_map, opcode_extra_args=(0,)) diff --git a/aten/plena/program_attention.py b/aten/plena/program_attention.py index d1df9fa..6e86084 100644 --- a/aten/plena/program_attention.py +++ b/aten/plena/program_attention.py @@ -58,6 +58,78 @@ def _build_valid_col_mask(self, name: str, valid_cols: int) -> VRAMMatrixVar: self.emit("\n".join(lines) + "\n") return mask + def _build_causal_score_mask(self, name: str) -> VRAMMatrixVar: + """Materialize an MLEN x MLEN causal score mask. + + The mask is 0 on and below the diagonal and -inf above it. This is the + same single-tile triangular mask that the MHA path uses when callers pass + an explicit VRAM mask. Multi-sequence-tile causal geometry remains owned + by the existing MHA tile-skipping logic; packed GQA still rejects + seq_len > MLEN. + """ + mask = self.alloc(name, self.mlen, self.mlen) + mask_addr = self.get_vram_addr(mask.name) + fp_scratch_base = self._ONLINE_SOFTMAX_FPSRAM_BASE + gp_mask, gp_fp = self.register_allocator.allocate_gp(2) + + lines = [ + f"; === Build causal score mask: MLEN={self.mlen} ===", + f"S_ADDI_INT gp{gp_fp}, gp0, {fp_scratch_base}", + "S_LD_FP f7, gp0, 2", + ] + lines.extend(load_large_int(gp_mask, mask_addr)) + for row in range(self.mlen): + for col in range(self.mlen): + value_reg = "f0" if col <= row else "f7" + lines.append(f"S_ST_FP {value_reg}, gp{gp_fp}, {col}") + lines.extend( + [ + f"S_MAP_V_FP gp{gp_mask}, gp{gp_fp}, 0", + f"S_ADDI_INT gp{gp_mask}, gp{gp_mask}, {self.mlen}", + ] + ) + + self.register_allocator.free_gp([gp_mask, gp_fp]) + self.emit("\n".join(lines) + "\n") + return mask + + def _build_sliding_causal_score_mask(self, name: str, sliding_window: int) -> VRAMMatrixVar: + """Materialize a single-tile GPT-OSS sliding causal score mask. + + Visible columns satisfy ``col <= row`` and ``col > row - sliding_window``. + This is only the short-sequence/single-tile mask. Multi-tile sliding + remains part of the existing sequence-tile debt. + """ + if sliding_window <= 0: + raise ValueError(f"sliding_window must be positive, got {sliding_window}") + mask = self.alloc(name, self.mlen, self.mlen) + mask_addr = self.get_vram_addr(mask.name) + fp_scratch_base = self._ONLINE_SOFTMAX_FPSRAM_BASE + gp_mask, gp_fp = self.register_allocator.allocate_gp(2) + + lines = [ + f"; === Build sliding causal score mask: MLEN={self.mlen}, window={sliding_window} ===", + f"S_ADDI_INT gp{gp_fp}, gp0, {fp_scratch_base}", + "S_LD_FP f7, gp0, 2", + ] + lines.extend(load_large_int(gp_mask, mask_addr)) + for row in range(self.mlen): + min_visible_col = row - sliding_window + 1 + for col in range(self.mlen): + visible = col <= row and col >= min_visible_col + value_reg = "f0" if visible else "f7" + lines.append(f"S_ST_FP {value_reg}, gp{gp_fp}, {col}") + lines.extend( + [ + f"S_MAP_V_FP gp{gp_mask}, gp{gp_fp}, 0", + f"S_ADDI_INT gp{gp_mask}, gp{gp_mask}, {self.mlen}", + ] + ) + + self.register_allocator.free_gp([gp_mask, gp_fp]) + self.emit("\n".join(lines) + "\n") + return mask + def flash_attention( self, Q, @@ -87,8 +159,6 @@ def flash_attention( if h_qkv is None: raise ValueError("GQA mode requires h_qkv to be specified") - if causal_mask is not None: - raise NotImplementedError("causal_mask is not yet supported for GQA flash attention") return self._flash_attention_gqa_fused( Q, K, @@ -97,6 +167,7 @@ def flash_attention( hq, hkv, h_qkv, + causal_mask=causal_mask, batch_size=batch_size, seq_len=seq_len, kv_seq_len=kv_seq_len, @@ -162,6 +233,8 @@ def _flash_attention_mha( if scale is None: scale = 1.0 / math.sqrt(head_dim) + if causal_mask is True: + causal_mask = self._build_causal_score_mask("_mha_causal_mask") num_q_blocks = math.ceil(seq_len / mlen) num_k_blocks = math.ceil(kv_seq_len / mlen) @@ -301,6 +374,7 @@ def _flash_attention_gqa_fused( hkv, h_qkv, *, + causal_mask=None, batch_size: int = 1, seq_len: int | None = None, kv_seq_len: int | None = None, @@ -358,6 +432,8 @@ def _flash_attention_gqa_fused( if scale is None: scale = 1.0 / math.sqrt(h_qkv) + if causal_mask is True: + causal_mask = self._build_causal_score_mask("_gqa_causal_mask") if s_q > mlen or s_kv > mlen: raise NotImplementedError("Packed GQA lowering currently supports one sequence tile.") @@ -435,7 +511,7 @@ def _flash_attention_gqa_fused( scratch_base_address=scratch_addr, broadcast_amount=broadcast_amount, scale=scale, - causal_mask=None, + causal_mask=causal_mask, output_head_base=0, k_idx=batch_idx * k_row_blocks_per_batch, valid_cols=s_kv, @@ -461,6 +537,9 @@ def _emit_packed_qkt_to_s( k_idx: int, s_base_address: int, k_head_offset: int = 0, + k_matrix_precision: str | int = "weights", + k_set_scale: bool = True, + k_hbm_element_bytes: int = 1, ) -> None: """Emit packed QK^T with M_BTMM/M_BMM_WO into head-major S tiles.""" self._ensure_hbm_sub_matrix_registered(K) @@ -475,6 +554,9 @@ def _emit_packed_qkt_to_s( mram_dest_addr=0, hbm_addr_reg=addr_reg, gp_regs=gp_regs, + precision=0 if k_matrix_precision in ("weights", "weight", 0) else 1, + set_scale=k_set_scale, + hbm_element_bytes=k_hbm_element_bytes, ), ) @@ -616,6 +698,11 @@ def _emit_packed_attention_group_internal( output_head_base: int = 0, k_idx: int = 0, valid_cols: int | None = None, + sink_base_address: int | None = None, + k_matrix_precision: str | int = "weights", + k_set_scale: bool = True, + k_hbm_element_bytes: int = 1, + v_hbm_element_bytes: int = 1, ) -> None: """Compiler-owned packed-head flash attention for one KV group.""" seq_len, q_width = Q_group.shape @@ -676,6 +763,9 @@ def _emit_packed_attention_group_internal( q_idx=0, k_idx=k_idx, s_base_address=scratch_base_address, + k_matrix_precision=k_matrix_precision, + k_set_scale=k_set_scale, + k_hbm_element_bytes=k_hbm_element_bytes, ) active_cols = valid_cols or seq_len @@ -704,8 +794,23 @@ def _emit_packed_attention_group_internal( if valid_col_mask is not None or isinstance(causal_mask, VRAMMatrixVar) else active_cols ) - self.online_softmax_block(s_head, softmax_scale, rows=rows, valid_cols=softmax_valid_cols) - self.compute_pv(s_head, V, k_idx, pv, head_slot_dim, rows=rows) + sink_address = None if sink_base_address is None else sink_base_address + output_head_base + head + self.online_softmax_block( + s_head, + softmax_scale, + rows=rows, + valid_cols=softmax_valid_cols, + sink_address=sink_address, + ) + self.compute_pv( + s_head, + V, + k_idx, + pv, + head_slot_dim, + rows=rows, + v_hbm_element_bytes=v_hbm_element_bytes, + ) self.scale_o_row(o_head, 0, rows=rows) self.vram_add(o_head, pv, num_rows=rows) self.final_scale_o(0, o_head, rows=rows) @@ -772,6 +877,8 @@ def flash_attention_packed_groups_looped( raise ValueError(f"Packed attention requires HLEN={head_slot_dim}, got {self.hlen}") if scale is None: scale = 1.0 / math.sqrt(head_slot_dim) + if causal_mask is True: + causal_mask = self._build_causal_score_mask("_packed_loop_causal_mask") for K, V in kv_pairs: self._ensure_hbm_sub_matrix_registered(K) @@ -917,6 +1024,12 @@ def flash_attention_packed_group( causal_mask: bool | VRAMMatrixVar | None = True, k_idx: int = 0, valid_cols: int | None = None, + sink_base_address: int | None = None, + output_head_base: int = 0, + k_matrix_precision: str | int = "weights", + k_set_scale: bool = True, + k_hbm_element_bytes: int = 1, + v_hbm_element_bytes: int = 1, ) -> None: """Emit one KV group's packed-head flash-attention body. @@ -938,6 +1051,8 @@ def flash_attention_packed_group( ) if scale is None: scale = 1.0 / math.sqrt(head_slot_dim) + if causal_mask is True: + causal_mask = self._build_causal_score_mask("_packed_group_causal_mask") self._emit_packed_attention_group_internal( Q_group=Q_group, @@ -951,8 +1066,14 @@ def flash_attention_packed_group( broadcast_amount=broadcast_amount, scale=scale, causal_mask=causal_mask, + output_head_base=output_head_base, k_idx=k_idx, valid_cols=valid_cols, + sink_base_address=sink_base_address, + k_matrix_precision=k_matrix_precision, + k_set_scale=k_set_scale, + k_hbm_element_bytes=k_hbm_element_bytes, + v_hbm_element_bytes=v_hbm_element_bytes, ) def init_online_softmax(self, q_idx: int, o_matrix: VRAMMatrixVar, rows: int | None = None): @@ -974,6 +1095,7 @@ def online_softmax_block( scale: float, rows: int | None = None, valid_cols: int | None = None, + sink_address: int | None = None, ): """Perform Online Softmax on S block""" super().online_softmax_block( @@ -981,6 +1103,7 @@ def online_softmax_block( scale=scale, rows=rows, valid_cols=valid_cols, + sink_address=sink_address, ) def compute_pv( @@ -991,6 +1114,7 @@ def compute_pv( pv_matrix: VRAMMatrixVar, head_dim: int, rows: int | None = None, + v_hbm_element_bytes: int = 1, ): """Compute PV = P @ V[k_idx] where P is stored in s_block.""" if not isinstance(s_block, VRAMMatrixVar): @@ -1008,6 +1132,7 @@ def compute_pv( pv_matrix=pv_matrix.name, head_dim=head_dim, rows=rows, + v_hbm_element_bytes=v_hbm_element_bytes, ) def scale_o_row(self, o_matrix: VRAMMatrixVar, q_idx: int, rows: int | None = None): diff --git a/aten/plena/program_fp_tile_ops.py b/aten/plena/program_fp_tile_ops.py index ba4e064..61b04aa 100644 --- a/aten/plena/program_fp_tile_ops.py +++ b/aten/plena/program_fp_tile_ops.py @@ -162,15 +162,17 @@ def tile_row_exp( source: VRAMMatrixVar, row_idx: int | None = None, rows: Iterable[int] | None = None, + tile_col_idx: int = 0, ): - super().tile_row_exp(source.name, self._resolve_rows(row_idx=row_idx, rows=rows)) + super().tile_row_exp(source.name, self._resolve_rows(row_idx=row_idx, rows=rows), tile_col_idx=tile_col_idx) def tile_row_reci( self, source: VRAMMatrixVar, rows: Iterable[int] | None = None, + tile_col_idx: int = 0, ): - super().tile_row_reci(source.name, self._default_rows(rows)) + super().tile_row_reci(source.name, self._default_rows(rows), tile_col_idx=tile_col_idx) def tile_row_sub_fp( self, @@ -180,9 +182,10 @@ def tile_row_sub_fp( rows: Iterable[int] | None = None, fpram_offset: int = 0, fpram_base_offset: int = 0, + tile_col_idx: int = 0, ): return self._tile_row_fp_scalar( - "tile_row_sub_fp", source, fpram_addr, row_idx, rows, fpram_offset, fpram_base_offset + "tile_row_sub_fp", source, fpram_addr, row_idx, rows, fpram_offset, fpram_base_offset, tile_col_idx ) def tile_row_mul_fp( @@ -193,9 +196,38 @@ def tile_row_mul_fp( rows: Iterable[int] | None = None, fpram_offset: int = 0, fpram_base_offset: int = 0, + tile_col_idx: int = 0, ): return self._tile_row_fp_scalar( - "tile_row_mul_fp", source, fpram_addr, row_idx, rows, fpram_offset, fpram_base_offset + "tile_row_mul_fp", source, fpram_addr, row_idx, rows, fpram_offset, fpram_base_offset, tile_col_idx + ) + + def tile_row_max_fp( + self, + source: VRAMMatrixVar, + fpram_addr: int | FPVar, + row_idx: int | None = None, + rows: Iterable[int] | None = None, + fpram_offset: int = 0, + fpram_base_offset: int = 0, + tile_col_idx: int = 0, + ): + return self._tile_row_fp_scalar( + "tile_row_max_fp", source, fpram_addr, row_idx, rows, fpram_offset, fpram_base_offset, tile_col_idx + ) + + def tile_row_min_fp( + self, + source: VRAMMatrixVar, + fpram_addr: int | FPVar, + row_idx: int | None = None, + rows: Iterable[int] | None = None, + fpram_offset: int = 0, + fpram_base_offset: int = 0, + tile_col_idx: int = 0, + ): + return self._tile_row_fp_scalar( + "tile_row_min_fp", source, fpram_addr, row_idx, rows, fpram_offset, fpram_base_offset, tile_col_idx ) def tile_row_add_fp( @@ -203,9 +235,10 @@ def tile_row_add_fp( source: VRAMMatrixVar, fp_var: FPVar, rows: Iterable[int] | None = None, + tile_col_idx: int = 0, ): resolved_rows = self._default_rows(rows) - super().tile_row_add_fp(source.name, [(row, fp_var[row]) for row in resolved_rows]) + super().tile_row_add_fp(source.name, [(row, fp_var[row]) for row in resolved_rows], tile_col_idx=tile_col_idx) def _tile_row_binary(self, isa_method: str, dst: VRAMMatrixVar, src: VRAMMatrixVar, rows: Iterable[int] | None): return getattr(super(), isa_method)(dst.name, src.name, self._default_rows(rows)) @@ -219,6 +252,7 @@ def _tile_row_fp_scalar( rows: Iterable[int] | None, fpram_offset: int, fpram_base_offset: int, + tile_col_idx: int = 0, ): return getattr(super(), isa_method)( source.name, @@ -229,6 +263,7 @@ def _tile_row_fp_scalar( single_offset=fpram_offset, base_offset=fpram_base_offset, ), + tile_col_idx=tile_col_idx, ) def tile_row_add( diff --git a/aten/plena/program_gpt_oss_moe.py b/aten/plena/program_gpt_oss_moe.py new file mode 100644 index 0000000..5d7f5c0 --- /dev/null +++ b/aten/plena/program_gpt_oss_moe.py @@ -0,0 +1,1812 @@ +"""Routed-MoE v0 program-builder helpers.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence + +from compiler.aten.isa_builder import IsaBuilder, addr as areg, fp, gp +from compiler.aten.plena.vars import FPVar, InputVar, VRAMMatrixVar + +GptOssFPConstants = tuple[FPVar, FPVar, FPVar, FPVar, FPVar] +ExpertWeights = tuple[InputVar, InputVar, InputVar] +ExpertBiases = tuple[VRAMMatrixVar | None, VRAMMatrixVar | None, VRAMMatrixVar | None] + + +class ProgramGptOssMoeMixin: + """Routed-MoE v0 emit helpers used by GPT-OSS and Qwen bring-up. + + The helpers intentionally keep routing policy explicit: GPT-OSS uses + high-precision router logits plus softmax-after-topk; Qwen adapters reuse + the same substrate with 128-way/top-8 selection. Expert gate/up projections + are split instead of using packed even/odd gate_up because current + vector-scalar min/max ops apply to whole rows, not alternating lanes. + """ + + def _validate_gpt_oss_constants(self, constants: GptOssFPConstants, rows: int) -> None: + zero, limit_pos, limit_neg, one, neg_alpha = constants + if zero.address != 0: + raise ValueError("vram_fill_zero assumes FPRAM f0 is preloaded with zero") + for var in (limit_pos, limit_neg, one, neg_alpha): + if var.size < rows: + raise ValueError(f"FPVar {var.name} size={var.size} is smaller than rows={rows}") + + def _validate_standard_swiglu_constants(self, constants: GptOssFPConstants, rows: int) -> None: + zero, _unused_pos, _unused_neg, one, neg_one = constants + if zero.address != 0: + raise ValueError("vram_fill_zero assumes FPRAM f0 is preloaded with zero") + for var in (one, neg_one): + if var.size < rows: + raise ValueError(f"FPVar {var.name} size={var.size} is smaller than rows={rows}") + + def _vram_matrix_row_addr(self, matrix: VRAMMatrixVar, row_idx: int, tile_col_idx: int = 0) -> int: + row_block = row_idx // self.mlen + row_in_block = row_idx % self.mlen + return self.get_vram_tile_addr(matrix.name, row_block, tile_col_idx) + row_in_block * self.mlen + + def gpt_oss_router_logits_bf16_v0( + self, + x: VRAMMatrixVar, + router_weight_rows: VRAMMatrixVar, + *, + rows: int, + hidden: int, + num_experts: int, + name: str = "gpt_oss_router_logits", + ) -> VRAMMatrixVar: + """Emit high-precision GPT-OSS router logits using BF16 vector dot products. + + This path intentionally avoids ``linear_projection`` and all HBM/MX + prefetch machinery. ``x`` and ``router_weight_rows`` must already be + resident in BF16 VRAM. The router is the only GPT-OSS MoE v0 path that + uses this non-MX lowering; expert projections continue to use the MXFP8 + matrix path. + + ``router_weight_rows`` is laid out as ``[num_experts, hidden]`` so every + expert's hidden vector can be multiplied row-wise against a token row. + """ + if hidden % self.mlen != 0: + raise ValueError(f"router hidden={hidden} must be divisible by MLEN={self.mlen}") + if rows > x.shape[0]: + raise ValueError(f"router rows={rows} exceeds x rows={x.shape[0]}") + if hidden > x.shape[1]: + raise ValueError(f"router hidden={hidden} exceeds x width={x.shape[1]}") + if num_experts > router_weight_rows.shape[0] or hidden > router_weight_rows.shape[1]: + raise ValueError( + "router_weight_rows must have shape at least " + f"({num_experts}, {hidden}), got {router_weight_rows.shape}" + ) + + expert_blocks = math.ceil(num_experts / self.mlen) + logical_logit_rows = rows if expert_blocks == 1 else rows * expert_blocks + logical_logit_cols = num_experts if expert_blocks == 1 else self.mlen + physical_rows = max(self.blen, math.ceil(logical_logit_rows / self.blen) * self.blen) + logits = self.alloc( + name, + rows=logical_logit_rows, + cols=logical_logit_cols, + strict=False, + physical_shape=(physical_rows, self.mlen), + ) + scratch = self.alloc( + f"{name}_dot_scratch", + rows=1, + cols=self.mlen, + strict=False, + physical_shape=(1, self.mlen), + ) + fp_scratch = self.fp_var(f"{name}_fp_scratch", size=expert_blocks * self.mlen) + + scratch_addr = self.get_vram_addr(scratch.name) + k_blocks = hidden // self.mlen + gp_x, gp_w, gp_scratch, gp_fp, gp_out, gp_loop = self._reg.allocate_gp(6) + fp_acc = self.allocate_fp_reg(1)[0] + try: + asm = IsaBuilder().comment( + f"GPT-OSS router BF16 vector-dot logits: rows={rows}, hidden={hidden}, experts={num_experts}" + ) + asm.instr("S_ADDI_INT", gp(gp_scratch), gp(0), scratch_addr) + for token_idx in range(rows): + fp_base = fp_scratch.address + + asm.comment(f"Router token {token_idx}: clear FPRAM logits scratch") + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_base) + asm.instr("C_LOOP_START", gp(gp_loop), expert_blocks * self.mlen) + asm.instr("S_ST_FP", fp(0), gp(gp_fp), 0) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(gp_fp), 1) + asm.instr("C_LOOP_END", gp(gp_loop)) + + for expert_idx in range(num_experts): + x_addr = self._vram_matrix_row_addr(x, token_idx, 0) + w_addr = self._vram_matrix_row_addr(router_weight_rows, expert_idx, 0) + x_step = x.physical_shape[0] * self.mlen + w_step = router_weight_rows.physical_shape[0] * self.mlen + + asm.comment(f"Router dot token {token_idx}, expert {expert_idx}") + asm.instr("S_ADD_FP", fp(fp_acc), fp(0), fp(0)) + asm.instr("S_ADDI_INT", gp(gp_x), gp(0), x_addr) + asm.instr("S_ADDI_INT", gp(gp_w), gp(0), w_addr) + asm.instr("C_LOOP_START", gp(gp_loop), k_blocks) + asm.instr("V_MUL_VV", gp(gp_scratch), gp(gp_x), gp(gp_w), 0) + asm.instr("V_RED_SUM", fp(fp_acc), gp(gp_scratch), 0, 0) + asm.instr("S_ADDI_INT", gp(gp_x), gp(gp_x), x_step) + asm.instr("S_ADDI_INT", gp(gp_w), gp(gp_w), w_step) + asm.instr("C_LOOP_END", gp(gp_loop)) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_base + expert_idx) + asm.instr("S_ST_FP", fp(fp_acc), gp(gp_fp), 0) + + asm.comment(f"Router token {token_idx}: map FPRAM logits scratch to contiguous VRAM rows") + for expert_block in range(expert_blocks): + out_row = token_idx if expert_blocks == 1 else token_idx * expert_blocks + expert_block + out_addr = self._vram_matrix_row_addr(logits, out_row, 0) + asm.instr("S_ADDI_INT", gp(gp_out), gp(0), out_addr) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_base + expert_block * self.mlen) + asm.instr("S_MAP_V_FP", gp(gp_out), gp(gp_fp), 0) + self._emit(asm) + finally: + self.free_fp_reg([fp_acc]) + self._reg.free_gp([gp_x, gp_w, gp_scratch, gp_fp, gp_out, gp_loop]) + + return logits + + def qwen3_router_logits_matrix_bf16_rowpacked_v0( + self, + x: VRAMMatrixVar, + router_weight_matrix: InputVar, + *, + rows: int, + hidden: int, + num_experts: int, + mram_tile_capacity: int = 4, + stream_k_accum: bool = True, + name: str = "qwen3_router_logits_matrix", + ) -> VRAMMatrixVar: + """Emit Qwen router logits through the BF16 matrix path in V_TOPK layout. + + The matrix machine naturally produces a ``[rows, experts]`` tensor. The + existing V_TOPK ABI for 128 experts expects each token to occupy + contiguous MLEN-wide rows: ``token0/block0, token0/block1, token1/...``. + This helper keeps the numerically better matrix accumulation while + packing the result into that token-major row layout. + """ + if hidden % self.mlen != 0: + raise ValueError(f"router hidden={hidden} must be divisible by MLEN={self.mlen}") + if rows > x.shape[0]: + raise ValueError(f"router rows={rows} exceeds x rows={x.shape[0]}") + if hidden > x.shape[1]: + raise ValueError(f"router hidden={hidden} exceeds x width={x.shape[1]}") + if router_weight_matrix.shape[0] < hidden or router_weight_matrix.shape[1] < num_experts: + raise ValueError( + "router_weight_matrix must have shape at least " + f"({hidden}, {num_experts}), got {router_weight_matrix.shape}" + ) + if mram_tile_capacity <= 0: + raise ValueError(f"mram_tile_capacity must be > 0, got {mram_tile_capacity}") + + expert_blocks = math.ceil(num_experts / self.mlen) + physical_rows = max(self.blen, math.ceil(rows / self.blen) * self.blen) + physical_experts = expert_blocks * self.mlen + logical_logit_rows = rows if expert_blocks == 1 else rows * expert_blocks + physical_logit_rows = max(self.blen, math.ceil(logical_logit_rows / self.blen) * self.blen) + + old_capacity = self.mram_tile_capacity + self.mram_tile_capacity = mram_tile_capacity + try: + if stream_k_accum: + matrix_logits = self.linear_projection_bf16_stream_k_accum( + x, + router_weight_matrix, + name=f"{name}_matrix", + physical_shape=(physical_rows, physical_experts), + max_k_tiles=mram_tile_capacity, + ) + else: + matrix_logits = self.linear_projection_bf16( + x, + router_weight_matrix, + name=f"{name}_matrix", + physical_shape=(physical_rows, physical_experts), + ) + finally: + self.mram_tile_capacity = old_capacity + + if expert_blocks == 1: + return matrix_logits + + packed_logits = self.alloc( + name, + rows=logical_logit_rows, + cols=self.mlen, + strict=False, + physical_shape=(physical_logit_rows, self.mlen), + ) + + gp_dst, gp_src = self._reg.allocate_gp(2) + try: + asm = IsaBuilder().comment( + f"Qwen router matrix logits pack: rows={rows}, experts={num_experts}, blocks={expert_blocks}" + ) + for token_idx in range(rows): + for expert_block in range(expert_blocks): + src_addr = self._vram_matrix_row_addr(matrix_logits, token_idx, expert_block) + dst_row = token_idx * expert_blocks + expert_block + dst_addr = self._vram_matrix_row_addr(packed_logits, dst_row, 0) + 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_VF", gp(gp_dst), gp(gp_src), fp(0), 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_dst, gp_src]) + + self.free_tensor(matrix_logits) + return packed_logits + + def qwen3_router_logits_packed_skinny_bf16_rowpacked_v0( + self, + x: VRAMMatrixVar, + router_weight_packed_skinny: InputVar, + *, + rows: int, + hidden: int, + num_experts: int, + k_tiles_per_packed_tile: int = 8, + name: str = "qwen3_router_logits_packed_skinny", + ) -> VRAMMatrixVar: + """Emit Qwen router logits from a packed-skinny BF16 HBM table. + + This is the integration form of the packed-skinny router probe: the + weight table packs several skinny K slices into one full MRAM tile, so + cap8-equivalent accumulation can be expressed under the existing cap4 + MRAM contract. The result is returned in the existing V_TOPK token- + major ABI when ``num_experts`` spans multiple MLEN rows. + """ + if hidden % self.mlen != 0: + raise ValueError(f"router hidden={hidden} must be divisible by MLEN={self.mlen}") + if rows > x.shape[0]: + raise ValueError(f"router rows={rows} exceeds x rows={x.shape[0]}") + if rows > self.mlen: + raise NotImplementedError( + "packed-skinny Qwen router currently supports one sequence row-block; " + f"got rows={rows}, MLEN={self.mlen}" + ) + if hidden > x.shape[1]: + raise ValueError(f"router hidden={hidden} exceeds x width={x.shape[1]}") + if k_tiles_per_packed_tile <= 0: + raise ValueError(f"k_tiles_per_packed_tile must be > 0, got {k_tiles_per_packed_tile}") + + expert_blocks = math.ceil(num_experts / self.mlen) + physical_rows = max(self.blen, math.ceil(rows / self.blen) * self.blen) + physical_experts = expert_blocks * self.mlen + logical_logit_rows = rows if expert_blocks == 1 else rows * expert_blocks + physical_logit_rows = max(self.blen, math.ceil(logical_logit_rows / self.blen) * self.blen) + tiles_per_mlen = self.mlen // self.blen + + required_col_blocks = expert_blocks * math.ceil(self.mlen / self.blen) + if router_weight_packed_skinny.physical_shape[1] < required_col_blocks * self.mlen: + raise ValueError( + "router_weight_packed_skinny physical width is too small for " + f"{expert_blocks} output blocks: got {router_weight_packed_skinny.physical_shape}" + ) + + matrix_logits = self.alloc( + f"{name}_matrix", + rows=rows, + cols=num_experts, + strict=False, + physical_shape=(physical_rows, physical_experts), + ) + for expert_block in range(expert_blocks): + self.vram_sub_projection_packed_skinny_stream_k_accum_to( + x, + 0, + router_weight_packed_skinny, + expert_block * tiles_per_mlen, + matrix_logits, + 0, + expert_block, + max_k_tiles_per_packed_tile=k_tiles_per_packed_tile, + matrix_precision="keyvalue", + set_scale=False, + hbm_element_bytes=2, + ) + + if expert_blocks == 1: + return matrix_logits + + packed_logits = self.alloc( + name, + rows=logical_logit_rows, + cols=self.mlen, + strict=False, + physical_shape=(physical_logit_rows, self.mlen), + ) + + gp_dst, gp_src = self._reg.allocate_gp(2) + try: + asm = IsaBuilder().comment( + f"Qwen router packed-skinny logits pack: rows={rows}, experts={num_experts}, blocks={expert_blocks}" + ) + for token_idx in range(rows): + for expert_block in range(expert_blocks): + src_addr = self._vram_matrix_row_addr(matrix_logits, token_idx, expert_block) + dst_row = token_idx * expert_blocks + expert_block + dst_addr = self._vram_matrix_row_addr(packed_logits, dst_row, 0) + 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_VF", gp(gp_dst), gp(gp_src), fp(0), 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_dst, gp_src]) + + self.free_tensor(matrix_logits) + return packed_logits + + def gpt_oss_router_topk_softmax_v0( + self, + logits: VRAMMatrixVar, + *, + token_idx: int, + weights_fp_base: int, + indices_int_base: int, + num_experts: int = 32, + top_k: int = 4, + name: str = "gpt_oss_router_topk", + ) -> None: + """Emit V_TOPK for one router-logit row. + + V_TOPK v0 reads one BF16 router-logit row from VRAM, performs a + linear-scan top-k with low-index tie break, stores the selected expert + ids to INT SRAM, and stores the softmax-over-selected weights to FP + SRAM. The instruction intentionally keeps router/top-k on the BF16 + path and does not touch MX scale state. + """ + if token_idx < 0 or token_idx >= logits.shape[0]: + raise ValueError(f"token_idx={token_idx} outside logits rows={logits.shape[0]}") + expert_blocks = math.ceil(num_experts / self.mlen) + if expert_blocks == 1 and logits.shape[1] < num_experts: + raise ValueError(f"V_TOPK expects at least {num_experts} logits, got {logits.shape[1]}") + if expert_blocks > 1 and logits.shape[1] < self.mlen: + raise ValueError(f"V_TOPK expects MLEN-wide logit rows, got width={logits.shape[1]}") + required_rows = token_idx + 1 if expert_blocks == 1 else (token_idx + 1) * expert_blocks + if logits.shape[0] < required_rows: + raise ValueError( + f"V_TOPK expects token {token_idx} to occupy {expert_blocks} contiguous logit rows, " + f"got logits shape={logits.shape}" + ) + policy_rmask = {(32, 4): 0, (128, 8): 1}.get((num_experts, top_k)) + if policy_rmask is None: + raise NotImplementedError(f"V_TOPK policy unsupported for num_experts={num_experts}, top_k={top_k}") + + gp_weights, gp_logits, gp_indices = self._reg.allocate_gp(3) + try: + asm = IsaBuilder().comment( + f"Routed-MoE V_TOPK {name}: token={token_idx}, experts={num_experts}, top_k={top_k}, " + f"weights_fp={weights_fp_base}, indices_int={indices_int_base}" + ) + asm.instr("S_ADDI_INT", gp(gp_weights), gp(0), weights_fp_base) + asm.instr( + "S_ADDI_INT", + gp(gp_logits), + gp(0), + self._vram_matrix_row_addr(logits, token_idx if expert_blocks == 1 else token_idx * expert_blocks, 0), + ) + asm.instr("S_ADDI_INT", gp(gp_indices), gp(0), indices_int_base) + asm.instr("V_TOPK", gp(gp_weights), gp(gp_logits), gp(gp_indices), policy_rmask) + self._emit(asm) + finally: + self._reg.free_gp([gp_weights, gp_logits, gp_indices]) + + def _emit_expert_id_to_weight_base_v0( + self, + asm: IsaBuilder, + *, + expert_indices_int_base: int, + pair_idx: int, + table_base: int, + per_expert_stride: int, + addr_reg: int, + gp_table: int, + gp_expert: int, + gp_stride: int, + gp_offset: int, + gp_base: int, + name: str, + ) -> None: + """Emit the shared true-expert-id -> HBM-base address calculation.""" + if per_expert_stride <= 0: + raise ValueError(f"{name}: per_expert_stride must be positive, got {per_expert_stride}") + asm.comment( + f"{name}: expert_id_to_weight_base pair={pair_idx}, " + f"table_base={table_base}, stride={per_expert_stride}" + ) + asm.instr("S_ADDI_INT", gp(gp_table), gp(0), expert_indices_int_base) + asm.instr("S_LD_INT", gp(gp_expert), gp(gp_table), pair_idx) + asm.instr("S_ADDI_INT", gp(gp_stride), gp(0), per_expert_stride) + asm.instr("S_MUL_INT", gp(gp_offset), gp(gp_expert), gp(gp_stride)) + asm.instr("S_ADDI_INT", gp(gp_base), gp(0), table_base) + asm.instr("S_ADD_INT", gp(gp_base), gp(gp_base), gp(gp_offset)) + asm.instr("C_SET_ADDR_REG", areg(addr_reg), gp(0), gp(gp_base)) + + def _emit_expert_id_to_weight_base_table_v0( + self, + asm: IsaBuilder, + *, + expert_indices_int_base: int, + expert_base_table_int_base: int, + pair_idx: int, + addr_reg: int, + gp_table: int, + gp_expert: int, + gp_base: int, + name: str, + ) -> None: + """Emit expert-id -> HBM-base lookup through an IntSRAM base table.""" + asm.comment( + f"{name}: expert_id_to_weight_base_table pair={pair_idx}, " + f"base_table_int={expert_base_table_int_base}" + ) + asm.instr("S_ADDI_INT", gp(gp_table), gp(0), expert_indices_int_base) + asm.instr("S_LD_INT", gp(gp_expert), gp(gp_table), pair_idx) + asm.instr("S_LD_INT", gp(gp_base), gp(gp_expert), expert_base_table_int_base) + asm.instr("C_SET_ADDR_REG", areg(addr_reg), gp(0), gp(gp_base)) + + def gpt_oss_expert_id_to_weight_base_v0( + self, + *, + expert_indices_int_base: int, + pair_idx: int, + table_base: int, + per_expert_stride: int, + addr_reg: int, + name: str = "gpt_oss_expert_id_to_weight_base", + ) -> None: + """Public helper for Step6: set ``addr_reg`` to true expert HBM base. + + ``expert_indices_int_base[pair_idx]`` must contain a true GPT-OSS expert + id in ``[0, 31]``. This helper is the only supported address contract for + dynamic expert weights; callers must not remap to host-compressed expert + indices. + """ + gp_table, gp_expert, gp_stride, gp_offset, gp_base = self._reg.allocate_gp(5) + try: + asm = IsaBuilder() + self._emit_expert_id_to_weight_base_v0( + asm, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=table_base, + per_expert_stride=per_expert_stride, + addr_reg=addr_reg, + gp_table=gp_table, + gp_expert=gp_expert, + gp_stride=gp_stride, + gp_offset=gp_offset, + gp_base=gp_base, + name=name, + ) + self._emit(asm) + finally: + self._reg.free_gp([gp_table, gp_expert, gp_stride, gp_offset, gp_base]) + + def _gpt_oss_dynamic_load_sub_matrix_col_v0( + self, + *, + weight_template: InputVar, + col_idx: int, + expert_indices_int_base: int, + pair_idx: int, + table_base: int, + per_expert_stride: int, + expert_base_table_int_base: int | None = None, + mram_start_addr: int | None = None, + k_block_start: int = 0, + k_block_count: int | None = None, + name: str = "gpt_oss_dynamic_weight_load", + ) -> None: + """Load one weight column tile using runtime true expert id addressing.""" + self._ensure_hbm_sub_matrix_registered(weight_template) + layout = self.get_hbm_layout(weight_template.name) + num_row_blocks = layout.num_row_blocks + block_size = self.mlen * self.mlen + effective_count = k_block_count if k_block_count is not None else num_row_blocks + if mram_start_addr is None: + mram_start_addr = self.mram_allocator.allocate( + f"{name}_{weight_template.name}_pair{pair_idx}_col{col_idx}", + effective_count * block_size, + ) + + gp_table, gp_expert, gp_expert_stride, gp_expert_offset, gp_base, gp_scale, gp_stride, gp_mram = ( + self._reg.allocate_gp(8) + ) + addr_reg = self._reg.allocate_addr(1)[0] + try: + asm = IsaBuilder().comment( + f"GPT-OSS dynamic HBM weight prefetch: template={weight_template.name}, " + f"pair={pair_idx}, col={col_idx}" + ) + if expert_base_table_int_base is None: + self._emit_expert_id_to_weight_base_v0( + asm, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=table_base, + per_expert_stride=per_expert_stride, + addr_reg=addr_reg, + gp_table=gp_table, + gp_expert=gp_expert, + gp_stride=gp_expert_stride, + gp_offset=gp_expert_offset, + gp_base=gp_base, + name=name, + ) + else: + self._emit_expert_id_to_weight_base_table_v0( + asm, + expert_indices_int_base=expert_indices_int_base, + expert_base_table_int_base=expert_base_table_int_base, + pair_idx=pair_idx, + addr_reg=addr_reg, + gp_table=gp_table, + gp_expert=gp_expert, + gp_base=gp_base, + name=name, + ) + self._emit_hbm_prefetch_setup(asm, layout, gp_scale, gp_stride) + self._emit_hbm_subblock_sequence( + asm, + layout, + ((row_idx, col_idx) for row_idx in range(k_block_start, k_block_start + effective_count)), + mram_start_addr, + addr_reg, + gp_scale, + gp_mram, + ) + self._emit(asm) + finally: + self._reg.free_gp( + [gp_table, gp_expert, gp_expert_stride, gp_expert_offset, gp_base, gp_scale, gp_stride, gp_mram] + ) + self._reg.free_addr([addr_reg]) + + def gpt_oss_dynamic_vram_sub_projection_to_v0( + self, + vram_matrix: VRAMMatrixVar, + vram_row_idx: int, + weight_template: InputVar, + weight_col_idx: int, + target: VRAMMatrixVar, + target_row_idx: int, + target_col_idx: int, + *, + expert_indices_int_base: int, + pair_idx: int, + table_base: int, + per_expert_stride: int, + expert_base_table_int_base: int | None = None, + auto_reset_mram: bool = True, + k_block_start: int = 0, + k_block_count: int | None = None, + name: str = "gpt_oss_dynamic_projection", + ) -> None: + """Projection tile where the HBM weight base comes from V_TOPK expert id.""" + vram_matrix = self._require_var(vram_matrix, VRAMMatrixVar, "vram_matrix") + weight_template = self._require_var(weight_template, InputVar, "weight_template") + target = self._require_var(target, VRAMMatrixVar, "target") + self._ensure_vram_sub_matrix_registered(vram_matrix) + self._ensure_hbm_sub_matrix_registered(weight_template) + if auto_reset_mram: + super().reset_mram() + self._gpt_oss_dynamic_load_sub_matrix_col_v0( + weight_template=weight_template, + col_idx=weight_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, + k_block_start=k_block_start, + k_block_count=k_block_count, + name=name, + ) + super().vram_sub_projection_to( + vram_mat_name=vram_matrix.name, + vram_row_idx=vram_row_idx, + mram_mat_name=weight_template.name, + mram_col_idx=weight_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, + ) + + def gpt_oss_dynamic_linear_projection_v0( + self, + input_var: VRAMMatrixVar, + weight_template: InputVar, + *, + expert_indices_int_base: int, + pair_idx: int, + table_base: int, + per_expert_stride: int, + expert_base_table_int_base: int | None = None, + name: str, + physical_shape: tuple[int, int] | None = None, + ) -> VRAMMatrixVar: + """Tiled linear projection with runtime expert-id weight selection.""" + mlen = self.mlen + rows, _k_total = input_var.shape + _weight_rows, out_features = weight_template.shape + if physical_shape is None: + # K-split accumulation uses 64x64 block adds. Routed dynamic + # projections often have only 4/8 logical rows, so keep outputs + # tile-backed to prevent block-add accumulation from walking into + # the next column block. + physical_rows = max(self.mlen, input_var.physical_shape[0], math.ceil(rows / self.blen) * self.blen) + physical_out_features = weight_template.physical_shape[1] + else: + physical_rows, physical_out_features = physical_shape + if physical_rows < rows or physical_out_features < out_features: + raise ValueError( + f"physical_shape {physical_shape} cannot be smaller than logical output {(rows, out_features)}" + ) + + physical_k = max(input_var.physical_shape[1], weight_template.physical_shape[0]) + 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 + + output = self.alloc( + name, + rows, + out_features, + strict=False, + 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.gpt_oss_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, + ) + + 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) + return output + + temp = self.alloc(f"{name}_temp", mlen, mlen) + for k_chunk_idx, k_block_start in enumerate(range(0, num_k_tiles, max_k_tiles)): + k_block_count = min(max_k_tiles, num_k_tiles - k_block_start) + k_split = {"k_block_start": k_block_start, "k_block_count": k_block_count} + 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) + else: + emit_projection(row_idx, col_idx, temp, 0, 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 + + def gpt_oss_add_dynamic_expert_bias_v0( + self, + dst: VRAMMatrixVar, + bias_table: VRAMMatrixVar, + *, + expert_indices_int_base: int, + pair_idx: int, + rows: int, + width: int, + name: str = "gpt_oss_dynamic_bias", + ) -> None: + """Add BF16 bias selected by true expert id from a VRAM bias table.""" + if width % self.mlen != 0: + raise ValueError(f"{name}: width={width} must be divisible by MLEN={self.mlen}") + if rows > self.blen: + raise ValueError(f"{name}: v0 expects one routed pair slot (rows<=BLEN), got rows={rows}") + self._ensure_vram_sub_matrix_registered(dst) + self._ensure_vram_sub_matrix_registered(bias_table) + num_col_blocks = width // self.mlen + bias_rows = bias_table.physical_shape[0] + expert_row_stride = self.blen * self.mlen + + gp_table, gp_expert, gp_stride, gp_expert_offset, gp_src_base, gp_src, gp_dst = self._reg.allocate_gp(7) + try: + asm = IsaBuilder().comment(f"GPT-OSS dynamic expert bias add {name}: pair={pair_idx}, rows={rows}") + asm.instr("S_ADDI_INT", gp(gp_table), gp(0), expert_indices_int_base) + asm.instr("S_LD_INT", gp(gp_expert), gp(gp_table), pair_idx) + asm.instr("S_ADDI_INT", gp(gp_stride), gp(0), expert_row_stride) + asm.instr("S_MUL_INT", gp(gp_expert_offset), gp(gp_expert), gp(gp_stride)) + for col_block in range(num_col_blocks): + src_col_base = self._vram_matrix_row_addr(bias_table, 0, col_block) + for row_idx in range(rows): + dst_addr = self._vram_matrix_row_addr(dst, row_idx, col_block) + asm.instr("S_ADDI_INT", gp(gp_src_base), gp(0), src_col_base + row_idx * self.mlen) + asm.instr("S_ADD_INT", gp(gp_src), gp(gp_src_base), gp(gp_expert_offset)) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), dst_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_table, gp_expert, gp_stride, gp_expert_offset, gp_src_base, gp_src, gp_dst]) + + def gpt_oss_materialize_topk_route_weight_v0( + self, + *, + weights_fp_base: int, + pair_idx: int, + rows: int, + hidden: int, + zero_row: FPVar | None = None, + fp_scratch: FPVar | None = None, + name: str = "gpt_oss_device_route_weight", + ) -> VRAMMatrixVar: + """Expand device V_TOPK scalar weight into a VRAM route matrix.""" + if rows > self.blen: + raise ValueError(f"{name}: v0 expects one routed pair slot (rows<=BLEN), got rows={rows}") + if hidden % self.mlen != 0: + raise ValueError(f"{name}: hidden={hidden} must be divisible by MLEN={self.mlen}") + route = self.alloc(name, rows=rows, cols=hidden, strict=False, physical_shape=(self.blen, hidden)) + self.gpt_oss_true_zero_vram_rows_v0( + route, + rows=list(range(self.blen)), + hidden=hidden, + zero_row=zero_row, + name=f"{name}_zero", + ) + fp_scratch = fp_scratch or self.fp_var(f"{name}_fp_row", size=self.mlen) + gp_dst, gp_fp = self._reg.allocate_gp(2) + try: + for col_block in range(hidden // self.mlen): + self.fpvar_fill_from_fpram_asm(fp_scratch.address, weights_fp_base + pair_idx, self.mlen) + asm = IsaBuilder().comment( + f"GPT-OSS materialize route weight pair={pair_idx}, col_block={col_block}" + ) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), self._vram_matrix_row_addr(route, 0, col_block)) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_scratch.address) + asm.instr("S_MAP_V_FP", gp(gp_dst), gp(gp_fp), 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_dst, gp_fp]) + return route + + def gpt_oss_materialize_route_weights_for_active_rows_v0( + self, + *, + weights_fp_base: int, + pair_indices: Sequence[int], + active_rows: Sequence[int], + rows: int, + hidden: int, + zero_row: FPVar | None = None, + fp_scratch: FPVar | None = None, + name: str = "gpt_oss_device_route_weights_grouped", + ) -> VRAMMatrixVar: + """Expand selected scalar route weights into specific active VRAM rows.""" + if len(pair_indices) != len(active_rows): + raise ValueError( + f"{name}: pair_indices={len(pair_indices)} active_rows={len(active_rows)} length mismatch" + ) + if rows <= 0: + raise ValueError(f"{name}: rows must be positive") + if hidden % self.mlen != 0: + raise ValueError(f"{name}: hidden={hidden} must be divisible by MLEN={self.mlen}") + active_list = [int(row) for row in active_rows] + pair_list = [int(pair) for pair in pair_indices] + physical_rows = max(self.blen, math.ceil(rows / self.blen) * self.blen) + if active_list and (min(active_list) < 0 or max(active_list) >= physical_rows): + raise ValueError(f"{name}: active rows {active_list} exceed physical rows={physical_rows}") + + route = self.alloc(name, rows=rows, cols=hidden, strict=False, physical_shape=(physical_rows, hidden)) + self.gpt_oss_true_zero_vram_rows_v0( + route, + rows=list(range(physical_rows)), + hidden=hidden, + zero_row=zero_row, + name=f"{name}_zero", + ) + fp_scratch = fp_scratch or self.fp_var(f"{name}_fp_row", size=self.mlen) + gp_dst, gp_fp = self._reg.allocate_gp(2) + try: + for pair_idx, active_row in zip(pair_list, active_list, strict=True): + for col_block in range(hidden // self.mlen): + self.fpvar_fill_from_fpram_asm(fp_scratch.address, weights_fp_base + pair_idx, self.mlen) + asm = IsaBuilder().comment( + f"GPT-OSS materialize route weight pair={pair_idx}, active_row={active_row}, col_block={col_block}" + ) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), self._vram_matrix_row_addr(route, active_row, col_block)) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_scratch.address) + asm.instr("S_MAP_V_FP", gp(gp_dst), gp(gp_fp), 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_dst, gp_fp]) + return route + + def gpt_oss_dynamic_expert_pair_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_idx: int, + bias_tables: ExpertBiases | None, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + zero_row: FPVar | None = None, + route_fp_scratch: FPVar | None = None, + activation_policy: str = "gpt_oss_clamp_gated", + name: str = "gpt_oss_dynamic_expert_pair", + ) -> VRAMMatrixVar: + """Run one routed pair using true expert id from device V_TOPK output.""" + 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) + + gate = self.gpt_oss_dynamic_linear_projection_v0( + x, + w_gate, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=gate_base, + per_expert_stride=gate_stride, + name=f"{name}_gate", + physical_shape=(projection_rows, w_gate.physical_shape[1]), + ) + up = self.gpt_oss_dynamic_linear_projection_v0( + x, + w_up, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=up_base, + per_expert_stride=up_stride, + name=f"{name}_up", + physical_shape=(projection_rows, w_up.physical_shape[1]), + ) + if gate_bias_table is not None: + self.gpt_oss_add_dynamic_expert_bias_v0( + gate, + gate_bias_table, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + rows=rows, + width=intermediate, + name=f"{name}_gate_bias", + ) + if up_bias_table is not None: + self.gpt_oss_add_dynamic_expert_bias_v0( + up, + up_bias_table, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + 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, + name=name, + ) + out = self.gpt_oss_dynamic_linear_projection_v0( + hidden, + w_down, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=down_base, + per_expert_stride=down_stride, + name=f"{name}_out", + physical_shape=(projection_rows, w_down.physical_shape[1]), + ) + if down_bias_table is not None: + self.gpt_oss_add_dynamic_expert_bias_v0( + out, + down_bias_table, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + rows=rows, + width=w_down.physical_shape[1], + name=f"{name}_down_bias", + ) + route = self.gpt_oss_materialize_topk_route_weight_v0( + weights_fp_base=weights_fp_base, + pair_idx=pair_idx, + rows=rows, + hidden=w_down.physical_shape[1], + zero_row=zero_row, + fp_scratch=route_fp_scratch, + name=f"{name}_route", + ) + self.vram_mul(out, route, num_rows=rows) + return out + + def gpt_oss_gather_token_rows_from_hbm_v0( + self, + x_input: InputVar, + *, + token_offsets_int_base: int, + pair_count: int, + hidden: int, + zero_row: FPVar | None = None, + name: str = "gpt_oss_gathered_x", + ) -> VRAMMatrixVar: + """Gather routed token rows from HBM into compact BF16 VRAM rows. + + ``token_offsets_int_base`` points into int SRAM. Entry ``i`` contains + the element offset of the source token row inside ``x_input``'s HBM + element stream, i.e. ``token_index * hidden``. The loop count remains + compile-time fixed; only the HBM row offset is loaded at runtime. + + H_PREFETCH_V transfers four VLEN rows per call. Each routed pair + therefore owns a four-row slot. The active row is the first row of the + slot, while the remaining rows are cleared after prefetch. This is + intentionally wasteful but keeps the first L2 correctness path exact + under the current ISA and avoids a copy-through-vector-ALU rounding + step. + """ + if pair_count <= 0: + raise ValueError("pair_count must be positive") + if hidden % self.mlen != 0: + raise ValueError(f"gather hidden={hidden} must be divisible by MLEN={self.mlen}") + if hidden > x_input.shape[1]: + raise ValueError(f"gather hidden={hidden} exceeds x_input width={x_input.shape[1]}") + + logical_rows = pair_count * self.blen + physical_rows = max(self.blen, math.ceil(logical_rows / self.blen) * self.blen) + gathered = self.alloc( + name, + rows=logical_rows, + cols=hidden, + strict=False, + physical_shape=(physical_rows, hidden), + ) + + x_rows, x_cols = x_input.physical_shape + if x_cols != hidden: + raise ValueError(f"gather currently expects x_input physical width {hidden}, got {x_cols}") + num_col_blocks = hidden // self.mlen + + gp_table, gp_token_offset, gp_col, gp_offset, gp_dst, gp_scale, gp_stride = self._reg.allocate_gp(7) + addr_reg = self._reg.allocate_addr(1)[0] + try: + asm = IsaBuilder().comment( + f"GPT-OSS gather token rows from HBM: pairs={pair_count}, hidden={hidden}, slot_rows={self.blen}" + ) + asm.instr("S_ADDI_INT", gp(gp_table), gp(0), token_offsets_int_base) + asm.instr("S_ADDI_INT", gp(gp_scale), gp(0), x_rows * x_cols) + asm.instr("C_SET_SCALE_REG", gp(gp_scale)) + asm.instr("S_ADDI_INT", gp(gp_stride), gp(0), x_cols) + asm.instr("C_SET_STRIDE_REG", gp(gp_stride)) + asm.instr("S_ADDI_INT", gp(gp_offset), gp(0), x_input.hbm_addr) + asm.instr("C_SET_ADDR_REG", areg(addr_reg), gp(0), gp(gp_offset)) + + for pair_idx in range(pair_count): + active_row = pair_idx * self.blen + asm.comment(f"Gather pair slot {pair_idx}: dynamic token row offset from int SRAM") + asm.instr("S_LD_INT", gp(gp_token_offset), gp(gp_table), pair_idx) + for col_block in range(num_col_blocks): + col_offset = col_block * self.mlen + dst_addr = self._vram_matrix_row_addr(gathered, active_row, col_block) + asm.instr("S_ADDI_INT", gp(gp_col), gp(0), col_offset) + asm.instr("S_ADD_INT", gp(gp_offset), gp(gp_token_offset), gp(gp_col)) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), dst_addr) + asm.instr("H_PREFETCH_V", gp(gp_dst), gp(gp_offset), areg(addr_reg), 1, 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_table, gp_token_offset, gp_col, gp_offset, gp_dst, gp_scale, gp_stride]) + self._reg.free_addr([addr_reg]) + + padding_rows = [ + pair_idx * self.blen + pad_idx + for pair_idx in range(pair_count) + for pad_idx in range(1, self.blen) + ] + if padding_rows: + fp_zero_row = zero_row or self.fp_var(f"{name}_zero_row", size=self.mlen) + gp_fp, gp_dst, gp_loop = self._reg.allocate_gp(3) + try: + asm = IsaBuilder().comment(f"Clear gather padding rows with true FP zero rows: {padding_rows}") + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_zero_row.address) + asm.instr("C_LOOP_START", gp(gp_loop), self.mlen) + asm.instr("S_ST_FP", fp(0), gp(gp_fp), 0) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(gp_fp), 1) + asm.instr("C_LOOP_END", gp(gp_loop)) + for row_idx in padding_rows: + for col_block in range(num_col_blocks): + dst_addr = self._vram_matrix_row_addr(gathered, row_idx, col_block) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), dst_addr) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_zero_row.address) + asm.instr("S_MAP_V_FP", gp(gp_dst), gp(gp_fp), 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_fp, gp_dst, gp_loop]) + + return gathered + + def gpt_oss_gather_token_rows_from_vram_v0( + self, + x: VRAMMatrixVar, + *, + token_indices: Sequence[int], + hidden: int, + zero_row: FPVar | None = None, + name: str = "gpt_oss_gathered_x_vram", + ) -> VRAMMatrixVar: + """Copy routed token rows from BF16 VRAM into BLEN-row pair slots. + + This is the decoder-block counterpart to + :meth:`gpt_oss_gather_token_rows_from_hbm_v0`. A real block feeds MoE + from the VRAM-resident post-attention RMSNorm output, so this helper + must not emit HBM prefetches, ``C_SET_SCALE_REG``, or activation + quantization. Each routed pair still owns one BLEN-row slot to match + the existing dynamic expert-pair path; row 0 of each slot is active and + padding rows are written with true zeros. + """ + self._ensure_vram_sub_matrix_registered(x) + if hidden % self.mlen != 0: + raise ValueError(f"VRAM gather hidden={hidden} must be divisible by MLEN={self.mlen}") + if hidden > x.shape[1]: + raise ValueError(f"VRAM gather hidden={hidden} exceeds x width={x.shape[1]}") + + token_list = [int(token) for token in token_indices] + if not token_list: + raise ValueError("VRAM gather token_indices must be non-empty") + if min(token_list) < 0 or max(token_list) >= x.physical_shape[0]: + raise ValueError(f"VRAM gather token_indices {token_list} exceed x physical rows={x.physical_shape[0]}") + + pair_count = len(token_list) + logical_rows = pair_count * self.blen + physical_rows = max(self.blen, math.ceil(logical_rows / self.blen) * self.blen) + gathered = self.alloc( + name, + rows=logical_rows, + cols=hidden, + strict=False, + physical_shape=(physical_rows, hidden), + ) + + self.gpt_oss_true_zero_vram_rows_v0( + gathered, + rows=list(range(physical_rows)), + hidden=hidden, + zero_row=zero_row, + name=f"{name}_zero", + ) + + num_col_blocks = hidden // self.mlen + gp_dst, gp_src = self._reg.allocate_gp(2) + try: + asm = IsaBuilder().comment( + f"GPT-OSS gather token rows from VRAM: pairs={pair_count}, hidden={hidden}, slot_rows={self.blen}" + ) + for pair_idx, token_idx in enumerate(token_list): + active_row = pair_idx * self.blen + asm.comment(f"VRAM gather pair slot {pair_idx}: token row {token_idx}") + 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 gpt_oss_true_zero_vram_rows_v0( + self, + matrix: VRAMMatrixVar, + *, + rows: Sequence[int], + hidden: int, + zero_row: FPVar | None = None, + name: str = "gpt_oss_zero_rows", + ) -> None: + """Clear selected VRAM rows by mapping a true FP zero row. + + Do not use vector multiply-by-zero here: padding/gather slots can + contain NaNs, and ``NaN * 0`` remains NaN. This helper writes real + zeros through ``S_MAP_V_FP`` and is therefore safe for gather padding + and scatter accumulators. + """ + if hidden % self.mlen != 0: + raise ValueError(f"zero hidden={hidden} must be divisible by MLEN={self.mlen}") + row_list = [int(row) for row in rows] + if not row_list: + return + if min(row_list) < 0 or max(row_list) >= matrix.physical_shape[0]: + raise ValueError(f"zero rows {row_list} exceed physical rows={matrix.physical_shape[0]}") + + num_col_blocks = hidden // self.mlen + fp_zero_row = zero_row or self.fp_var(f"{name}_zero_row", size=self.mlen) + gp_fp, gp_dst, gp_loop = self._reg.allocate_gp(3) + try: + asm = IsaBuilder().comment(f"GPT-OSS true-zero VRAM rows {row_list} in {matrix.name}") + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_zero_row.address) + asm.instr("C_LOOP_START", gp(gp_loop), self.mlen) + asm.instr("S_ST_FP", fp(0), gp(gp_fp), 0) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(gp_fp), 1) + asm.instr("C_LOOP_END", gp(gp_loop)) + for row_idx in row_list: + for col_block in range(num_col_blocks): + dst_addr = self._vram_matrix_row_addr(matrix, row_idx, col_block) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), dst_addr) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_zero_row.address) + asm.instr("S_MAP_V_FP", gp(gp_dst), gp(gp_fp), 0) + self._emit(asm) + finally: + self._reg.free_gp([gp_fp, gp_dst, gp_loop]) + + def gpt_oss_scatter_add_active_rows_v0( + self, + dst: VRAMMatrixVar, + src: VRAMMatrixVar, + *, + token_indices: Sequence[int], + active_rows: Sequence[int], + hidden: int, + name: str = "gpt_oss_scatter_add", + ) -> None: + """Add routed active slot rows into final token rows in VRAM. + + ``src`` is a 4-row-slot tensor produced by gather/expert execution; + only one active row per slot is added into ``dst[token]``. This is the + VRAM-only half of L2 scatter-combine and intentionally does not emit + HBM stores. + """ + if hidden % self.mlen != 0: + raise ValueError(f"scatter hidden={hidden} must be divisible by MLEN={self.mlen}") + if len(token_indices) != len(active_rows): + raise ValueError(f"token_indices={len(token_indices)} active_rows={len(active_rows)} length mismatch") + + token_list = [int(token) for token in token_indices] + active_list = [int(row) for row in active_rows] + if token_list and (min(token_list) < 0 or max(token_list) >= dst.physical_shape[0]): + raise ValueError(f"scatter tokens {token_list} exceed dst physical rows={dst.physical_shape[0]}") + if active_list and (min(active_list) < 0 or max(active_list) >= src.physical_shape[0]): + raise ValueError(f"scatter active rows {active_list} exceed src physical rows={src.physical_shape[0]}") + + num_col_blocks = hidden // self.mlen + gp_dst, gp_src = self._reg.allocate_gp(2) + try: + asm = IsaBuilder().comment(f"GPT-OSS VRAM scatter-add {name}: {len(token_list)} active rows") + for token_idx, active_row in zip(token_list, active_list, strict=True): + for col_block in range(num_col_blocks): + dst_addr = self._vram_matrix_row_addr(dst, token_idx, col_block) + src_addr = self._vram_matrix_row_addr(src, active_row, 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]) + + def gpt_oss_clamp_gated_activation_v0( + self, + gate: VRAMMatrixVar, + up: VRAMMatrixVar, + *, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + name: str, + ) -> VRAMMatrixVar: + """Emit GPT-OSS clamp-gated activation and return hidden in ``up``. + + Computes ``(clamp(up)+1) * clamp(gate) * sigmoid(1.702 * clamp(gate))``. + The implementation uses ``exp(-1.702 * gate)`` and reciprocal to form + sigmoid. Inputs and outputs are BF16 VRAM tensors. + """ + self._validate_gpt_oss_constants(constants, rows) + _, limit_pos, limit_neg, one, neg_alpha = constants + active_rows = list(range(rows)) + physical_rows = max(self.mlen, math.ceil(rows / self.mlen) * self.mlen) + num_col_blocks = math.ceil(intermediate / self.mlen) + sigmoid = self.alloc( + f"{name}_sigmoid", + rows=rows, + cols=intermediate, + physical_shape=(physical_rows, intermediate), + strict=False, + ) + + for col_block in range(num_col_blocks): + self.tile_row_min_fp(gate, limit_pos, rows=active_rows, tile_col_idx=col_block) + self.tile_row_min_fp(up, limit_pos, rows=active_rows, tile_col_idx=col_block) + self.tile_row_max_fp(up, limit_neg, rows=active_rows, tile_col_idx=col_block) + + self.vram_fill_zero(sigmoid, rows=active_rows) + self.vram_add(sigmoid, gate, num_rows=rows) + + for col_block in range(num_col_blocks): + self.tile_row_mul_fp(sigmoid, neg_alpha, rows=active_rows, tile_col_idx=col_block) + self.tile_row_exp(sigmoid, rows=active_rows, tile_col_idx=col_block) + self.tile_row_add_fp(sigmoid, one, rows=active_rows, tile_col_idx=col_block) + self.tile_row_reci(sigmoid, rows=active_rows, tile_col_idx=col_block) + self.vram_mul(gate, sigmoid, num_rows=rows) + + for col_block in range(num_col_blocks): + self.tile_row_add_fp(up, one, rows=active_rows, tile_col_idx=col_block) + self.vram_mul(up, gate, num_rows=rows) + return up + + def standard_swiglu_activation_v0( + self, + gate: VRAMMatrixVar, + up: VRAMMatrixVar, + *, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + name: str, + ) -> VRAMMatrixVar: + """Emit standard SwiGLU activation and return hidden in ``up``. + + Computes ``silu(gate) * up = gate * sigmoid(gate) * up``. Qwen-style + experts use this non-clamped path, unlike GPT-OSS' clamp-gated variant. + Inputs and outputs are BF16 VRAM tensors. + """ + self._validate_standard_swiglu_constants(constants, rows) + _, _unused_pos, _unused_neg, one, neg_one = constants + active_rows = list(range(rows)) + physical_rows = max(self.mlen, math.ceil(rows / self.mlen) * self.mlen) + num_col_blocks = math.ceil(intermediate / self.mlen) + sigmoid = self.alloc( + f"{name}_sigmoid", + rows=rows, + cols=intermediate, + physical_shape=(physical_rows, intermediate), + strict=False, + ) + + self.vram_fill_zero(sigmoid, rows=active_rows) + self.vram_add(sigmoid, gate, num_rows=rows) + + for col_block in range(num_col_blocks): + self.tile_row_mul_fp(sigmoid, neg_one, rows=active_rows, tile_col_idx=col_block) + self.tile_row_exp(sigmoid, rows=active_rows, tile_col_idx=col_block) + self.tile_row_add_fp(sigmoid, one, rows=active_rows, tile_col_idx=col_block) + self.tile_row_reci(sigmoid, rows=active_rows, tile_col_idx=col_block) + self.vram_mul(gate, sigmoid, num_rows=rows) + self.vram_mul(up, gate, num_rows=rows) + return up + + def gpt_oss_expert_v0( + self, + x: VRAMMatrixVar, + weights: ExpertWeights, + *, + biases: ExpertBiases | None = None, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + name: str, + ) -> VRAMMatrixVar: + """Emit one GPT-OSS expert and return its output.""" + w_gate, w_up, w_down = weights + gate_bias, up_bias, down_bias = biases or (None, None, None) + # The K-split projection path accumulates partial sums with a 64x64 + # block add. Routed slots are often only 4/8/12 physical rows, so keep + # expert projection outputs tile-backed to prevent the block add from + # walking into the next column block. + projection_rows = max(self.mlen, x.physical_shape[0], math.ceil(rows / self.blen) * self.blen) + gate = self.linear_projection( + x, + w_gate, + name=f"{name}_gate", + physical_shape=(projection_rows, w_gate.physical_shape[1]), + ) + up = self.linear_projection( + x, + w_up, + name=f"{name}_up", + physical_shape=(projection_rows, w_up.physical_shape[1]), + ) + if gate_bias is not None: + self.vram_add(gate, gate_bias, num_rows=rows) + if up_bias is not None: + self.vram_add(up, up_bias, num_rows=rows) + hidden = self.gpt_oss_clamp_gated_activation_v0( + gate, + up, + rows=rows, + intermediate=intermediate, + constants=constants, + name=name, + ) + out = self.linear_projection( + hidden, + w_down, + name=f"{name}_out", + physical_shape=(projection_rows, w_down.physical_shape[1]), + ) + if down_bias is not None: + self.vram_add(out, down_bias, num_rows=rows) + return out + + def gpt_oss_moe_fixed_routing_v0( + self, + x: VRAMMatrixVar, + experts: Sequence[ExpertWeights], + route_weights: Sequence[VRAMMatrixVar], + *, + expert_biases: Sequence[ExpertBiases | None] | None = None, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + name: str = "gpt_oss_moe", + ) -> VRAMMatrixVar: + """Emit fixed-routing MoE v0 and return the combined output. + + ``route_weights`` must already be loaded in VRAM and expanded to the + expert output shape. The first expert output is used as the accumulator. + """ + if not experts: + raise ValueError("At least one expert is required") + if len(experts) != len(route_weights): + raise ValueError(f"experts={len(experts)} does not match route_weights={len(route_weights)}") + if expert_biases is not None and len(expert_biases) != len(experts): + raise ValueError(f"expert_biases={len(expert_biases)} does not match experts={len(experts)}") + + acc: VRAMMatrixVar | None = None + for idx, (weights, route) in enumerate(zip(experts, route_weights, strict=True)): + biases = None if expert_biases is None else expert_biases[idx] + expert_out = self.gpt_oss_expert_v0( + x, + weights, + biases=biases, + rows=rows, + intermediate=intermediate, + constants=constants, + name=f"{name}_expert{idx}", + ) + self.vram_mul(expert_out, route, num_rows=rows) + if acc is None: + acc = expert_out + else: + self.vram_add(acc, expert_out, num_rows=rows) + + assert acc is not None + return acc + + def _require_moe_policy_v0(self, policy_name: str, *, helper: str) -> None: + """Guard generic MoE wrappers with implemented router/expert semantics.""" + supported = {"gpt_oss", "qwen3_moe"} + if policy_name not in supported: + raise NotImplementedError( + f"{helper} currently supports policy_name in {sorted(supported)}, got {policy_name!r}" + ) + + def _require_routed_moe_substrate_policy_v0(self, policy_name: str, *, helper: str) -> None: + """Guard wrappers that are substrate movement primitives, not router semantics.""" + supported = {"gpt_oss", "qwen3_moe", "deepseek_v3"} + if policy_name not in supported: + raise NotImplementedError( + f"{helper} supports only routed MoE substrate policies {sorted(supported)}, got {policy_name!r}" + ) + + def moe_router_logits_bf16_v0( + self, + x: VRAMMatrixVar, + router_weight_rows: VRAMMatrixVar, + *, + rows: int, + hidden: int, + num_experts: int, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_router_logits", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for high-precision router logits. + + GPT-OSS and Qwen3-MoE both use a high-precision BF16 router matmul in + the current PLENA substrate. DeepSeek is intentionally not enabled here + because its shared/routed expert split needs a separate adapter. + """ + self._require_moe_policy_v0(policy_name, helper="moe_router_logits_bf16_v0") + return self.gpt_oss_router_logits_bf16_v0( + x, + router_weight_rows, + rows=rows, + hidden=hidden, + num_experts=num_experts, + name=name, + ) + + def moe_router_select_v0( + self, + logits: VRAMMatrixVar, + *, + token_idx: int, + weights_fp_base: int, + indices_int_base: int, + policy_name: str = "gpt_oss", + num_experts: int = 32, + top_k: int = 4, + name: str = "gpt_oss_router_topk", + ) -> None: + """Generic substrate wrapper for router selection and selected weights.""" + self._require_moe_policy_v0(policy_name, helper="moe_router_select_v0") + self.gpt_oss_router_topk_softmax_v0( + logits, + token_idx=token_idx, + weights_fp_base=weights_fp_base, + indices_int_base=indices_int_base, + num_experts=num_experts, + top_k=top_k, + name=name, + ) + + def moe_expert_id_to_weight_base_v0( + self, + *, + expert_indices_int_base: int, + pair_idx: int, + table_base: int, + per_expert_stride: int, + addr_reg: int, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_expert_id_to_weight_base", + ) -> None: + """Generic substrate wrapper for true expert-id to HBM base address.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_expert_id_to_weight_base_v0") + self.gpt_oss_expert_id_to_weight_base_v0( + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + table_base=table_base, + per_expert_stride=per_expert_stride, + addr_reg=addr_reg, + name=name, + ) + + def moe_dynamic_linear_projection_v0( + self, + input_var: VRAMMatrixVar, + weight_template: InputVar, + *, + expert_indices_int_base: int, + pair_idx: int, + table_base: int, + per_expert_stride: int, + expert_base_table_int_base: int | None = None, + name: str, + physical_shape: tuple[int, int] | None = None, + policy_name: str = "gpt_oss", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for runtime expert-selected projection.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_dynamic_linear_projection_v0") + return self.gpt_oss_dynamic_linear_projection_v0( + input_var, + weight_template, + 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=name, + physical_shape=physical_shape, + ) + + def moe_add_dynamic_expert_bias_v0( + self, + dst: VRAMMatrixVar, + bias_table: VRAMMatrixVar, + *, + expert_indices_int_base: int, + pair_idx: int, + rows: int, + width: int, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_dynamic_bias", + ) -> None: + """Generic substrate wrapper for runtime expert-selected bias add.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_add_dynamic_expert_bias_v0") + self.gpt_oss_add_dynamic_expert_bias_v0( + dst, + bias_table, + expert_indices_int_base=expert_indices_int_base, + pair_idx=pair_idx, + rows=rows, + width=width, + name=name, + ) + + def moe_materialize_route_weight_v0( + self, + *, + weights_fp_base: int, + pair_idx: int, + rows: int, + hidden: int, + zero_row: FPVar | None = None, + fp_scratch: FPVar | None = None, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_device_route_weight", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for selected route-weight materialization.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_materialize_route_weight_v0") + return self.gpt_oss_materialize_topk_route_weight_v0( + weights_fp_base=weights_fp_base, + pair_idx=pair_idx, + rows=rows, + hidden=hidden, + zero_row=zero_row, + fp_scratch=fp_scratch, + name=name, + ) + + def moe_materialize_route_weights_for_active_rows_v0( + self, + *, + weights_fp_base: int, + pair_indices: Sequence[int], + active_rows: Sequence[int], + rows: int, + hidden: int, + zero_row: FPVar | None = None, + fp_scratch: FPVar | None = None, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_device_route_weights_grouped", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for selected route weights in grouped active rows.""" + self._require_routed_moe_substrate_policy_v0( + policy_name, helper="moe_materialize_route_weights_for_active_rows_v0" + ) + return self.gpt_oss_materialize_route_weights_for_active_rows_v0( + weights_fp_base=weights_fp_base, + pair_indices=pair_indices, + active_rows=active_rows, + rows=rows, + hidden=hidden, + zero_row=zero_row, + fp_scratch=fp_scratch, + name=name, + ) + + def moe_dynamic_expert_pair_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_idx: 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", + name: str = "gpt_oss_dynamic_expert_pair", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for one routed expert pair.""" + self._require_moe_policy_v0(policy_name, helper="moe_dynamic_expert_pair_v0") + if activation_policy not in {"gpt_oss_clamp_gated", "standard_swiglu"}: + raise NotImplementedError( + "moe_dynamic_expert_pair_v0 supports activation_policy in " + "{'gpt_oss_clamp_gated', 'standard_swiglu'}, got " + f"{activation_policy!r}" + ) + return self.gpt_oss_dynamic_expert_pair_v0( + x, + weights, + weight_table_bases=weight_table_bases, + weight_table_strides=weight_table_strides, + expert_indices_int_base=expert_indices_int_base, + weights_fp_base=weights_fp_base, + pair_idx=pair_idx, + bias_tables=bias_tables, + rows=rows, + intermediate=intermediate, + constants=constants, + zero_row=zero_row, + route_fp_scratch=route_fp_scratch, + activation_policy=activation_policy, + name=name, + ) + + def moe_gather_token_rows_from_hbm_v0( + self, + x_input: InputVar, + *, + token_offsets_int_base: int, + pair_count: int, + hidden: int, + zero_row: FPVar | None = None, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_gathered_x", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for HBM token gather into routed slots.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_gather_token_rows_from_hbm_v0") + return self.gpt_oss_gather_token_rows_from_hbm_v0( + x_input, + token_offsets_int_base=token_offsets_int_base, + pair_count=pair_count, + hidden=hidden, + zero_row=zero_row, + name=name, + ) + + def moe_gather_token_rows_from_vram_v0( + self, + x: VRAMMatrixVar, + *, + token_indices: Sequence[int], + hidden: int, + zero_row: FPVar | None = None, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_gathered_x_vram", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for VRAM-resident token rows into routed slots.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_gather_token_rows_from_vram_v0") + return self.gpt_oss_gather_token_rows_from_vram_v0( + x, + token_indices=token_indices, + hidden=hidden, + zero_row=zero_row, + name=name, + ) + + def moe_true_zero_vram_rows_v0( + self, + matrix: VRAMMatrixVar, + *, + rows: Sequence[int], + hidden: int, + zero_row: FPVar | None = None, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_zero_rows", + ) -> None: + """Generic substrate wrapper for true-zeroing routed slot rows.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_true_zero_vram_rows_v0") + self.gpt_oss_true_zero_vram_rows_v0(matrix, rows=rows, hidden=hidden, zero_row=zero_row, name=name) + + def moe_scatter_add_active_rows_v0( + self, + dst: VRAMMatrixVar, + src: VRAMMatrixVar, + *, + token_indices: Sequence[int], + active_rows: Sequence[int], + hidden: int, + policy_name: str = "gpt_oss", + name: str = "gpt_oss_scatter_add", + ) -> None: + """Generic substrate wrapper for routed-slot scatter-add.""" + self._require_routed_moe_substrate_policy_v0(policy_name, helper="moe_scatter_add_active_rows_v0") + self.gpt_oss_scatter_add_active_rows_v0( + dst, + src, + token_indices=token_indices, + active_rows=active_rows, + hidden=hidden, + name=name, + ) + + def moe_expert_activation_v0( + self, + gate: VRAMMatrixVar, + up: VRAMMatrixVar, + *, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + activation_policy: str = "gpt_oss_clamp_gated", + name: str, + ) -> VRAMMatrixVar: + """Generic substrate wrapper for expert activation backends.""" + if activation_policy == "gpt_oss_clamp_gated": + return self.gpt_oss_clamp_gated_activation_v0( + gate, + up, + rows=rows, + intermediate=intermediate, + constants=constants, + name=name, + ) + if activation_policy == "standard_swiglu": + return self.standard_swiglu_activation_v0( + gate, + up, + rows=rows, + intermediate=intermediate, + constants=constants, + name=name, + ) + raise NotImplementedError( + "moe_expert_activation_v0 supports activation_policy in " + "{'gpt_oss_clamp_gated', 'standard_swiglu'}, got " + f"{activation_policy!r}" + ) + + def moe_expert_v0( + self, + x: VRAMMatrixVar, + weights: ExpertWeights, + *, + biases: ExpertBiases | None = None, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + policy_name: str = "gpt_oss", + activation_policy: str = "gpt_oss_clamp_gated", + name: str, + ) -> VRAMMatrixVar: + """Generic substrate wrapper for one statically selected expert.""" + self._require_moe_policy_v0(policy_name, helper="moe_expert_v0") + if activation_policy != "gpt_oss_clamp_gated": + raise NotImplementedError( + "moe_expert_v0 currently supports only activation_policy='gpt_oss_clamp_gated', " + f"got {activation_policy!r}" + ) + return self.gpt_oss_expert_v0( + x, + weights, + biases=biases, + rows=rows, + intermediate=intermediate, + constants=constants, + name=name, + ) + + def moe_fixed_routing_v0( + self, + x: VRAMMatrixVar, + experts: Sequence[ExpertWeights], + route_weights: Sequence[VRAMMatrixVar], + *, + expert_biases: Sequence[ExpertBiases | None] | None = None, + rows: int, + intermediate: int, + constants: GptOssFPConstants, + policy_name: str = "gpt_oss", + activation_policy: str = "gpt_oss_clamp_gated", + name: str = "gpt_oss_moe", + ) -> VRAMMatrixVar: + """Generic substrate wrapper for fixed-routing MoE combine.""" + self._require_moe_policy_v0(policy_name, helper="moe_fixed_routing_v0") + if activation_policy != "gpt_oss_clamp_gated": + raise NotImplementedError( + "moe_fixed_routing_v0 currently supports only " + f"activation_policy='gpt_oss_clamp_gated', got {activation_policy!r}" + ) + return self.gpt_oss_moe_fixed_routing_v0( + x, + experts, + route_weights, + expert_biases=expert_biases, + rows=rows, + intermediate=intermediate, + constants=constants, + name=name, + ) + + +__all__ = ["ProgramGptOssMoeMixin"] diff --git a/aten/plena/program_matrix_ops.py b/aten/plena/program_matrix_ops.py index 9862f46..5a5f3bc 100644 --- a/aten/plena/program_matrix_ops.py +++ b/aten/plena/program_matrix_ops.py @@ -17,6 +17,19 @@ def _iter_k_chunks(num_k_tiles: int, max_k_tiles: int): k_start = k_end +def _matrix_precision_code(matrix_precision: str | int) -> int: + if isinstance(matrix_precision, int): + if matrix_precision not in (0, 1): + raise ValueError(f"matrix_precision int must be 0 or 1, got {matrix_precision}") + return matrix_precision + normalized = matrix_precision.lower().replace("-", "_") + if normalized in {"weight", "weights", "hbm_m_weight", "hbm_m_weight_type"}: + return 0 + if normalized in {"kv", "keyvalue", "key_value", "hbm_m_kv", "hbm_m_kv_type"}: + return 1 + raise ValueError(f"unknown matrix_precision={matrix_precision!r}") + + class ProgramMatrixOpsMixin: # ======================================================================== # Matrix Projection and VRAM Operations @@ -74,6 +87,9 @@ def vram_sub_projection_to( auto_reset_mram: bool = True, k_block_start: int = 0, k_block_count: int | None = None, + matrix_precision: str | int = "weights", + set_scale: bool = True, + hbm_element_bytes: int = 1, ): """ target[target_row_idx][target_col_idx] = vram_matrix[vram_row_idx][:] @ mram_input[:][mram_col_idx] @@ -85,6 +101,9 @@ def vram_sub_projection_to( col_idx=mram_col_idx, k_block_start=k_block_start, k_block_count=k_block_count, + precision=_matrix_precision_code(matrix_precision), + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, ) super().vram_sub_projection_to( vram_mat_name=vram_matrix.name, @@ -108,12 +127,21 @@ def vram_sub_projection_T_to( target_row_idx: int, target_col_idx: int, auto_reset_mram: bool = True, + matrix_precision: str | int = "weights", + set_scale: bool = True, + hbm_element_bytes: int = 1, ): """ target[target_row_idx][target_col_idx] = vram_matrix[vram_row_idx][:] @ mram_input[mram_row_idx][:]^T """ vram_matrix, mram_input, target = self._prepare_projection(vram_matrix, mram_input, target, auto_reset_mram) - super().load_sub_matrix_row(name=mram_input.name, row_idx=mram_row_idx) + super().load_sub_matrix_row( + name=mram_input.name, + row_idx=mram_row_idx, + precision=_matrix_precision_code(matrix_precision), + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) super().vram_sub_projection_T_to( vram_mat_name=vram_matrix.name, vram_row_idx=vram_row_idx, @@ -124,12 +152,172 @@ def vram_sub_projection_T_to( target_col_idx=target_col_idx, ) + def vram_sub_projection_stream_k_accum_to( + self, + vram_matrix: VRAMMatrixVar, + vram_row_idx: int, + mram_input: InputVar, + mram_col_idx: int, + target: VRAMMatrixVar, + target_row_idx: int, + target_col_idx: int, + *, + max_k_tiles: int, + matrix_precision: str | int = "keyvalue", + set_scale: bool = False, + hbm_element_bytes: int = 2, + ): + """Project one output tile while keeping K chunks in the FP32 accumulator. + + The ordinary wide-K path materializes each chunk to BF16 VRAM and then + adds chunks with vector ops. That is fine for most projections but can + flip Qwen router top-k near rank boundaries. This helper is deliberately + narrow: it reloads MRAM per 4x4 output microtile and only writes once, + preserving the matrix-machine accumulator across K chunks. + """ + vram_matrix, mram_input, target = self._prepare_projection( + vram_matrix, mram_input, target, auto_reset_mram=True + ) + if max_k_tiles <= 0: + raise ValueError(f"max_k_tiles must be > 0, got {max_k_tiles}") + + vram_layout = self.vram_matrices[vram_matrix.name] + vram_row_blocks = vram_layout.get_row_blocks(vram_row_idx) + physical_k = max(vram_matrix.physical_shape[1], mram_input.physical_shape[0]) + num_k_tiles = math.ceil(physical_k / self.mlen) + tiles_per_mlen = self.mlen // self.blen + valid_rows = vram_row_blocks[0].valid_shape[0] if vram_row_blocks[0].valid_shape else self.mlen + row_loop_count = min(tiles_per_mlen, max(1, math.ceil(valid_rows / self.blen))) + chunks = list(_iter_k_chunks(num_k_tiles, max_k_tiles)) + + for micro_col_idx in range(tiles_per_mlen): + for micro_row_idx in range(row_loop_count): + for chunk_idx, (k_block_start, k_block_count) in enumerate(chunks): + super().reset_mram() + super().load_sub_matrix_col( + name=mram_input.name, + col_idx=mram_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + precision=_matrix_precision_code(matrix_precision), + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) + super().vram_sub_projection_microtile_accumulate_to( + vram_mat_name=vram_matrix.name, + vram_row_idx=vram_row_idx, + mram_mat_name=mram_input.name, + mram_col_idx=mram_col_idx, + target_matrix=target.name, + target_row_idx=target_row_idx, + target_col_idx=target_col_idx, + micro_row_idx=micro_row_idx, + micro_col_idx=micro_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + write_out=(chunk_idx == len(chunks) - 1), + ) + + def vram_sub_projection_packed_skinny_stream_k_accum_to( + self, + vram_matrix: VRAMMatrixVar, + vram_row_idx: int, + packed_mram_input: InputVar, + packed_col_base_idx: int, + target: VRAMMatrixVar, + target_row_idx: int, + target_col_idx: int, + *, + max_k_tiles_per_packed_tile: int, + matrix_precision: str | int = "keyvalue", + set_scale: bool = False, + hbm_element_bytes: int = 2, + ): + """Compile-only packed-skinny projection probe. + + ``packed_mram_input`` is not a normal weight matrix. It is expected to + contain one full HBM tile per output micro-column and K group. Within a + tile, consecutive skinny K slices occupy columns + ``0:blen, blen:2*blen, ...``. The helper proves cap8-equivalent router + scheduling can fit in one MRAM tile per K group while preserving M_MM's + existing full-tile contract. + """ + vram_matrix, packed_mram_input, target = self._prepare_projection( + vram_matrix, + packed_mram_input, + target, + auto_reset_mram=True, + ) + if max_k_tiles_per_packed_tile <= 0: + raise ValueError( + f"max_k_tiles_per_packed_tile must be > 0, got {max_k_tiles_per_packed_tile}" + ) + + vram_layout = self.vram_matrices[vram_matrix.name] + vram_row_blocks = vram_layout.get_row_blocks(vram_row_idx) + num_k_tiles = len(vram_row_blocks) + tiles_per_mlen = self.mlen // self.blen + if max_k_tiles_per_packed_tile > tiles_per_mlen: + raise ValueError( + f"packed tile can hold at most {tiles_per_mlen} skinny slices, " + f"got {max_k_tiles_per_packed_tile}" + ) + + packed_layout = self.hbm_matrices[packed_mram_input.name] + chunks = list(_iter_k_chunks(num_k_tiles, max_k_tiles_per_packed_tile)) + if packed_layout.num_row_blocks < len(chunks): + raise ValueError( + f"packed_mram_input has {packed_layout.num_row_blocks} row groups, " + f"but {len(chunks)} are needed" + ) + if packed_layout.num_col_blocks < packed_col_base_idx + tiles_per_mlen: + raise ValueError( + f"packed_mram_input has {packed_layout.num_col_blocks} col blocks, " + f"but base {packed_col_base_idx} plus {tiles_per_mlen} micro-columns are needed" + ) + + valid_rows = vram_row_blocks[0].valid_shape[0] if vram_row_blocks[0].valid_shape else self.mlen + row_loop_count = min(tiles_per_mlen, max(1, math.ceil(valid_rows / self.blen))) + + for micro_col_idx in range(tiles_per_mlen): + packed_col_idx = packed_col_base_idx + micro_col_idx + for micro_row_idx in range(row_loop_count): + for group_idx, (k_block_start, k_block_count) in enumerate(chunks): + super().reset_mram() + super().load_sub_matrix( + name=packed_mram_input.name, + row_idx=group_idx, + col_idx=packed_col_idx, + mram_dest_addr=0, + precision=_matrix_precision_code(matrix_precision), + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, + ) + super().vram_sub_projection_packed_skinny_microtile_accumulate_to( + vram_mat_name=vram_matrix.name, + vram_row_idx=vram_row_idx, + packed_mram_mat_name=packed_mram_input.name, + packed_group_idx=group_idx, + packed_col_idx=packed_col_idx, + target_matrix=target.name, + target_row_idx=target_row_idx, + target_col_idx=target_col_idx, + micro_row_idx=micro_row_idx, + micro_col_idx=micro_col_idx, + k_block_start=k_block_start, + k_block_count=k_block_count, + write_out=(group_idx == len(chunks) - 1), + ) + def linear_projection( self, input_var: VRAMMatrixVar, weight_var: InputVar, name: str = "linear_out", physical_shape: tuple[int, int] | None = None, + matrix_precision: str | int = "weights", + set_scale: bool = True, + hbm_element_bytes: int = 1, ): """Emit tiled PLENA linear projection, including K-split accumulation.""" mlen = self.mlen @@ -171,6 +359,9 @@ def emit_projection(row_idx, col_idx, target, target_row_idx, target_col_idx, ** target, target_row_idx, target_col_idx, + matrix_precision=matrix_precision, + set_scale=set_scale, + hbm_element_bytes=hbm_element_bytes, **k_split, ) @@ -208,6 +399,187 @@ def emit_projection(row_idx, col_idx, target, target_row_idx, target_col_idx, ** self.free_tensor(temp) return output + def linear_projection_bf16_stream_k_accum( + self, + input_var: VRAMMatrixVar, + weight_var: InputVar, + name: str = "linear_out_bf16_stream_k_accum", + physical_shape: tuple[int, int] | None = None, + max_k_tiles: int | None = None, + ): + """BF16 projection with cross-K-chunk matrix accumulator retention.""" + mlen = self.mlen + rows, _k_total = input_var.shape + _weight_rows, out_features = weight_var.shape + if physical_shape is None: + physical_rows = max(input_var.physical_shape[0], math.ceil(rows / self.blen) * self.blen) + physical_out_features = weight_var.physical_shape[1] + else: + physical_rows, physical_out_features = physical_shape + if physical_rows < rows or physical_out_features < out_features: + raise ValueError( + f"physical_shape {physical_shape} cannot be smaller than " + f"logical output {(rows, out_features)}" + ) + + physical_k = max(input_var.physical_shape[1], weight_var.physical_shape[0]) + 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_tiles = self.mram_tile_capacity if max_k_tiles is None else max_k_tiles + + output = self.alloc( + name, + rows, + out_features, + strict=False, + physical_shape=(physical_rows, physical_out_features), + ) + + if num_k_tiles <= max_tiles: + for col_idx in range(num_col_blocks): + for row_idx in range(num_row_blocks): + self.vram_sub_projection_to( + input_var, + row_idx, + weight_var, + col_idx, + output, + row_idx, + col_idx, + matrix_precision="keyvalue", + set_scale=False, + hbm_element_bytes=2, + ) + return output + + for col_idx in range(num_col_blocks): + for row_idx in range(num_row_blocks): + self.vram_sub_projection_stream_k_accum_to( + input_var, + row_idx, + weight_var, + col_idx, + output, + row_idx, + col_idx, + max_k_tiles=max_tiles, + matrix_precision="keyvalue", + set_scale=False, + hbm_element_bytes=2, + ) + return output + + def linear_projection_bf16( + self, + input_var: VRAMMatrixVar, + weight_var: InputVar, + name: str = "linear_out_bf16", + physical_shape: tuple[int, int] | None = None, + ): + """Emit a high-precision BF16 matrix projection through HBM_M_KV_TYPE. + + The build must configure HBM_M_KV_TYPE as a Plain BF16 type for tensors + that use this path. No C_SET_SCALE_REG is emitted because plain BF16 + has no MX scale stream. + """ + return self.linear_projection( + input_var, + weight_var, + name=name, + physical_shape=physical_shape, + matrix_precision="keyvalue", + set_scale=False, + hbm_element_bytes=2, + ) + + def linear_projection_bias_bf16( + self, + input_var: VRAMMatrixVar, + weight_var: InputVar, + bias_var: VRAMMatrixVar | None = None, + name: str = "linear_out_bf16", + physical_shape: tuple[int, int] | None = None, + bias_rows: int | None = None, + ): + """Emit a BF16 projection and optional BF16 VRAM bias add. + + This is the shared high-precision projection substrate for router and + attention projections. It deliberately keeps projection/bias separate + from RoPE so callers can do one wide Q/K projection and then apply + model-specific RoPE on per-head VRAM views without duplicating the + projection path. + """ + out = self.linear_projection_bf16( + input_var, + weight_var, + name=name, + physical_shape=physical_shape, + ) + if bias_var is not None: + self.vram_add(out, bias_var, num_rows=bias_rows) + return out + + def runtime_rope_projection_bf16( + self, + x_var: VRAMMatrixVar, + rotate_weight_var: InputVar, + cos_var: VRAMMatrixVar, + sin_var: VRAMMatrixVar, + name: str = "runtime_rope_rot", + ) -> VRAMMatrixVar: + """Apply RoPE to a runtime projection output. + + Existing RoPE expects ``rotate_half(x)`` to already exist in VRAM. This + helper computes that rotate-half tensor through the BF16 HBM_M_KV + projection path, applies RoPE in-place to ``x_var``, then releases the + temporary. The rotate matrix can be model-specific, which keeps the + substrate generic for GPT-OSS/Qwen/DeepSeek adapters. + """ + x_rot = self.linear_projection_bf16( + x_var, + rotate_weight_var, + name=name, + physical_shape=x_var.physical_shape, + ) + self.rope(x_var, x_rot, cos_var, sin_var) + self.free_tensor(x_rot) + return x_var + + def head_runtime_rope_bf16( + self, + head_var: VRAMMatrixVar, + rotate_weight_var: InputVar, + cos_var: VRAMMatrixVar, + sin_var: VRAMMatrixVar, + *, + norm_weight_var: VRAMMatrixVar | None = None, + eps_offset: int | None = None, + reci_hid_offset: int | None = None, + num_rows: int | None = None, + name: str = "head_runtime_rope", + ) -> VRAMMatrixVar: + """Apply optional per-head RMSNorm and runtime RoPE to a BF16 head. + + GPT-OSS uses projection+bias+RoPE, while Qwen-style adapters also need + per-head Q/K RMSNorm before RoPE. Keeping this sequence in one helper + prevents each model harness from hand-rolling a slightly different + projection-to-attention path. The projection itself remains BF16 and + emits no MX scale setup. + """ + if norm_weight_var is not None: + if eps_offset is None or reci_hid_offset is None: + raise ValueError("head_runtime_rope_bf16 requires eps/reci offsets when norm_weight_var is set") + self.rms_norm(head_var, eps_offset=eps_offset, reci_hid_offset=reci_hid_offset) + self.vram_mul(head_var, norm_weight_var, num_rows=num_rows) + return self.runtime_rope_projection_bf16( + head_var, + rotate_weight_var, + cos_var, + sin_var, + name=name, + ) + def linear( self, input_var: VRAMMatrixVar, diff --git a/aten/plena/program_tensors.py b/aten/plena/program_tensors.py index 52a5c2e..6df43e4 100644 --- a/aten/plena/program_tensors.py +++ b/aten/plena/program_tensors.py @@ -18,6 +18,7 @@ def input( hbm_addr: int | None = None, prestaged_vram_addr: int | None = None, physical_shape: tuple[int, int] | None = None, + real_data_ratio: float | None = None, ) -> InputVar: """ Declare an input tensor (in HBM). @@ -35,9 +36,12 @@ def input( Returns: InputVar proxy object """ + if real_data_ratio is None: + real_data_ratio = self.real_data_ratio + h, w = physical_shape or shape size = h * w - hbm_size = int(size * self.real_data_ratio) + hbm_size = int(size * real_data_ratio) if hbm_addr is None: hbm_addr = self._allocate_hbm(hbm_size) @@ -57,7 +61,7 @@ def input( hbm_addr=hbm_addr, shape=shape, physical_shape=physical_shape, - real_data_ratio=self.real_data_ratio, + real_data_ratio=real_data_ratio, ) return var @@ -131,7 +135,15 @@ def load_batch( # Store Operations # ======================================================================== - def store(self, tensor_var, name: str | None = None, hbm_addr: int | None = None) -> InputVar: + def store( + self, + tensor_var, + name: str | None = None, + hbm_addr: int | None = None, + precision: int = 0, + hbm_element_bytes: int = 1, + real_data_ratio: float | None = None, + ) -> InputVar: """ Write tensor from VRAM back to HBM. @@ -147,11 +159,11 @@ def store(self, tensor_var, name: str | None = None, hbm_addr: int | None = None if hbm_addr is None: h, w = tensor_var.physical_shape size = h * w - hbm_size = int(size * self.real_data_ratio) + hbm_size = int(size * (real_data_ratio or self.real_data_ratio)) hbm_addr = self._allocate_hbm(hbm_size) else: h, w = tensor_var.physical_shape - hbm_size = int(h * w * self.real_data_ratio) + hbm_size = int(h * w * (real_data_ratio or self.real_data_ratio)) super().store_to_hbm( tensor_name=tensor_var.name, # internal name for symbol table lookup @@ -159,6 +171,9 @@ def store(self, tensor_var, name: str | None = None, hbm_addr: int | None = None hbm_object_name=internal_name, vlen=self.mlen, store_amount=self.hbm_v_writeback_amount, + precision=precision, + hbm_element_bytes=hbm_element_bytes, + hbm_real_data_ratio=real_data_ratio, ) var = InputVar( diff --git a/aten/reference.py b/aten/reference.py index 78128cd..283934e 100644 --- a/aten/reference.py +++ b/aten/reference.py @@ -10,7 +10,7 @@ import torch.nn.functional as F from compiler.aten.model_extract import LayerWeights, ModelConfig -from quant.quantizer.hardware_quantizer.mxfp import _mx_fp_quantize_hardware +from plena_quant.quantizer.hardware_quantizer.mxfp import _mx_fp_quantize_hardware _HW_MAX_K_TILES = 4 diff --git a/aten/tests/test_gpt_oss_moe_assertions.py b/aten/tests/test_gpt_oss_moe_assertions.py new file mode 100644 index 0000000..669792f --- /dev/null +++ b/aten/tests/test_gpt_oss_moe_assertions.py @@ -0,0 +1,82 @@ +"""Tests for GPT-OSS real-layer comparison thresholds.""" + +from __future__ import annotations + +import pytest +import torch + +from aten.gpt_oss_moe import assert_compare_within, assert_gap_in_band, compare_stats + + +def _reference_tensor() -> torch.Tensor: + return torch.linspace(-3.0, 3.0, 37, dtype=torch.float32).reshape(1, 37) + + +def test_hf_golden_a_threshold_accepts_tiny_math_noise(): + ref = _reference_tensor() + actual = ref + ref.std(unbiased=False) * 1e-4 + + stats = assert_compare_within( + actual, + ref, + name="HF_vs_GoldenA", + max_rel_rms=1e-2, + rtol=1e-2, + ) + + assert stats.rel_rms < 1e-3 + assert stats.allclose + + +def test_emulator_golden_b_threshold_accepts_two_percent_bound(): + ref = _reference_tensor() + actual = ref * 1.004 + + stats = assert_compare_within( + actual, + ref, + name="emu_vs_GoldenB", + max_rel_rms=0.02, + rtol=0.02, + ) + + assert stats.rel_rms < 0.005 + assert stats.allclose + + +def test_compare_threshold_rejects_large_lowering_error(): + ref = _reference_tensor() + actual = ref * 1.10 + + with pytest.raises(AssertionError, match="emu_vs_GoldenB failed"): + assert_compare_within( + actual, + ref, + name="emu_vs_GoldenB", + max_rel_rms=0.02, + rtol=0.02, + ) + + +def test_a_b_gap_band_is_recordable_and_checked(): + golden_a = _reference_tensor() + golden_b = golden_a * 0.98 + + stats = assert_gap_in_band( + golden_b, + golden_a, + name="GoldenA_vs_GoldenB", + min_rel_rms=0.005, + max_rel_rms=0.15, + ) + + assert 0.005 <= stats.rel_rms <= 0.15 + + +def test_compare_stats_uses_one_percent_reference_std_as_atol(): + ref = _reference_tensor() + stats = compare_stats(ref, ref, rtol=0.02) + + assert stats.atol == pytest.approx(float(ref.std(unbiased=False).item()) * 0.01) + assert stats.rel_rms == 0.0 + assert stats.pass_rate == 1.0 diff --git a/aten/tests/test_gpt_oss_moe_reference.py b/aten/tests/test_gpt_oss_moe_reference.py new file mode 100644 index 0000000..66d66c6 --- /dev/null +++ b/aten/tests/test_gpt_oss_moe_reference.py @@ -0,0 +1,321 @@ +"""Unit tests for GPT-OSS MoE Golden A/B helpers. + +Run: + PYTHONPATH=. python3 aten/tests/test_gpt_oss_moe_reference.py +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +_THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(_THIS_DIR))) +for _path in [ + os.path.join(_REPO_ROOT, "PLENA_Tools"), + os.path.join(_REPO_ROOT, "PLENA_Simulator", "PLENA_Tools"), +]: + if os.path.isdir(_path) and _path not in sys.path: + sys.path.insert(0, _path) + +from aten.gpt_oss_moe import ( + assert_clamp_inactive, + clamp_stats, + gpt_oss_moe_fixed_routing_host_smoke, + gpt_oss_moe_golden_a, + gpt_oss_moe_golden_b_plena_mxfp8, + gpt_oss_swiglu, + split_packed_gate_up, +) + + +def _tiny_inputs(scale: float = 0.05): + torch.manual_seed(7) + tokens, hidden, experts, inter = 3, 8, 5, 6 + x = torch.randn(tokens, hidden) * scale + router_w = torch.randn(hidden, experts) * scale + router_b = torch.randn(experts) * scale + gate_up_w = torch.randn(experts, hidden, 2 * inter) * scale + gate_up_b = torch.randn(experts, 2 * inter) * scale + down_w = torch.randn(experts, inter, hidden) * scale + down_b = torch.randn(experts, hidden) * scale + return x, router_w, router_b, gate_up_w, gate_up_b, down_w, down_b + + +def _load_hf_gpt_oss_mlp(): + # Some local torch/torchvision builds fail while importing transformers + # because torchvision registers a fake nms kernel before defining its op. + # Defining only the schema is enough for this CPU-only GPT-OSS MLP test. + try: + lib = torch.library.Library("torchvision", "DEF") + lib.define("nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor") + except Exception: + pass + + try: + from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssMLP + except Exception as exc: + pytest.skip(f"HF GPT-OSS MLP import unavailable in this environment: {exc}") + return GptOssConfig, GptOssMLP + + +def test_bf16_represents_gpt_oss_clamp_boundary_exactly(): + seven = torch.tensor([7.0], dtype=torch.bfloat16).float() + minus_seven = torch.tensor([-7.0], dtype=torch.bfloat16).float() + assert seven.item() == 7.0 + assert minus_seven.item() == -7.0 + + +def test_gpt_oss_swiglu_matches_manual_formula(): + gate_up = torch.tensor([[[8.0, -8.0, -2.0, 3.0]]]) + out = gpt_oss_swiglu(gate_up, bf16_intermediates=False) + + gate = torch.tensor([[[7.0, -2.0]]]) + up = torch.tensor([[[-7.0, 3.0]]]) + expected = (up + 1.0) * gate * torch.sigmoid(1.702 * gate) + assert torch.allclose(out, expected, atol=1e-6) + + +def test_split_packed_gate_up_uses_even_odd_weight_and_bias(): + weight = torch.arange(2 * 3 * 8, dtype=torch.float32).reshape(2, 3, 8) + bias = torch.arange(2 * 8, dtype=torch.float32).reshape(2, 8) + + split = split_packed_gate_up(weight, bias) + + assert torch.equal(split.gate_weight, weight[..., 0::2]) + assert torch.equal(split.up_weight, weight[..., 1::2]) + assert torch.equal(split.gate_bias, bias[..., 0::2]) + assert torch.equal(split.up_bias, bias[..., 1::2]) + assert not torch.equal(split.gate_weight, weight[..., :4]) + assert not torch.equal(split.up_weight, weight[..., 4:]) + + +def test_split_packed_gate_up_reconstructs_hf_projection_lanes(): + torch.manual_seed(11) + tokens, hidden, inter = 4, 5, 7 + # One-hot rows make this an exact lane-selection check rather than a + # floating-point matmul associativity test. + x = torch.eye(hidden, dtype=torch.float32)[:tokens] + packed_weight = torch.randn(hidden, 2 * inter) + packed_bias = torch.randn(2 * inter) + + packed = x @ packed_weight + packed_bias + split = split_packed_gate_up(packed_weight, packed_bias) + gate = x @ split.gate_weight + split.gate_bias + up = x @ split.up_weight + split.up_bias + + assert torch.allclose(gate, packed[..., 0::2], atol=0.0, rtol=0.0) + assert torch.allclose(up, packed[..., 1::2], atol=0.0, rtol=0.0) + + +def test_golden_a_matches_hf_gpt_oss_mlp_tiny(): + GptOssConfig, GptOssMLP = _load_hf_gpt_oss_mlp() + args = _tiny_inputs(scale=0.04) + x, router_w, router_b, gate_up_w, gate_up_b, down_w, down_b = args + + config = GptOssConfig( + num_hidden_layers=1, + num_local_experts=gate_up_w.shape[0], + hidden_size=x.shape[1], + intermediate_size=down_w.shape[1], + num_attention_heads=1, + num_key_value_heads=1, + head_dim=x.shape[1], + num_experts_per_tok=2, + vocab_size=16, + ) + hf_mlp = GptOssMLP(config).eval() + with torch.no_grad(): + hf_mlp.router.weight.copy_(router_w.T.contiguous()) + hf_mlp.router.bias.copy_(router_b) + hf_mlp.experts.gate_up_proj.copy_(gate_up_w) + hf_mlp.experts.gate_up_proj_bias.copy_(gate_up_b) + hf_mlp.experts.down_proj.copy_(down_w) + hf_mlp.experts.down_proj_bias.copy_(down_b) + + hf_router_logits = torch.nn.functional.linear(x, hf_mlp.router.weight, hf_mlp.router.bias) + hf_scores, hf_indices = hf_mlp.router(x) + hf_selected_scores = hf_scores.gather(1, hf_indices) + hf_expert_out = hf_mlp.experts(x.unsqueeze(0), hf_indices, hf_scores).squeeze(0) + hf_mlp_out, hf_mlp_scores = hf_mlp(x.unsqueeze(0)) + hf_mlp_selected_scores = hf_mlp_scores.gather(1, hf_indices) + + golden = gpt_oss_moe_golden_a( + *args, + experts_per_token=2, + bf16_intermediates=False, + ) + + assert torch.equal(golden.topk_indices, hf_indices) + assert torch.allclose(golden.router_logits, hf_router_logits, atol=1e-6, rtol=1e-6) + assert torch.allclose(golden.topk_weights, hf_selected_scores, atol=1e-6, rtol=1e-6) + assert torch.allclose(golden.topk_weights, hf_mlp_selected_scores, atol=1e-6, rtol=1e-6) + assert torch.allclose(golden.output, hf_expert_out, atol=1e-6, rtol=1e-6) + assert torch.allclose(golden.output, hf_mlp_out.squeeze(0), atol=1e-6, rtol=1e-6) + + +def test_clamp_inactive_smoke_precondition_and_unclamped_match(): + args = _tiny_inputs(scale=0.02) + result = gpt_oss_moe_golden_a(*args, experts_per_token=2, bf16_intermediates=True) + stats = assert_clamp_inactive(result.gate_up_preact) + + unclamped = gpt_oss_moe_golden_a( + *args, + experts_per_token=2, + bf16_intermediates=True, + apply_clamp=False, + ) + assert stats.inactive + assert torch.equal(result.topk_indices, unclamped.topk_indices) + assert torch.allclose(result.output.float(), unclamped.output.float(), atol=0.0, rtol=0.0) + + +def test_fixed_routing_wiring_smoke_matches_golden_a_without_clamp_or_quantization(): + args = _tiny_inputs(scale=0.02) + golden = gpt_oss_moe_golden_a( + *args, + experts_per_token=2, + bf16_intermediates=False, + apply_clamp=True, + ) + assert_clamp_inactive(golden.gate_up_preact) + x, _, _, gate_up_w, gate_up_b, down_w, down_b = args + + smoke = gpt_oss_moe_fixed_routing_host_smoke( + x, + golden.topk_indices, + golden.topk_weights, + gate_up_w, + gate_up_b, + down_w, + down_b, + bf16_intermediates=False, + apply_clamp=False, + ) + + assert torch.allclose(smoke.output, golden.output, atol=1e-6, rtol=1e-6) + + +def test_fixed_routing_clamp_active_smoke_matches_golden_a_before_quantization(): + args = list(_tiny_inputs(scale=0.04)) + # Force an active clamp site without changing routing. Golden A and the + # fixed-routing smoke should still agree when both apply exact clamp. + args[1] = torch.zeros_like(args[1]) + args[2] = torch.tensor([10.0, 9.0, 0.0, 0.0, 0.0]) + args[3] = args[3].clone() + args[4] = args[4].clone() + args[3][0, :, 0] = 0.0 + args[4][0, 0] = 8.0 + args[3][0, :, 1] = 0.0 + args[4][0, 1] = -8.0 + args = tuple(args) + + golden = gpt_oss_moe_golden_a( + *args, + experts_per_token=2, + bf16_intermediates=False, + apply_clamp=True, + ) + stats = clamp_stats(golden.gate_up_preact) + assert not stats.inactive + x, _, _, gate_up_w, gate_up_b, down_w, down_b = args + + smoke = gpt_oss_moe_fixed_routing_host_smoke( + x, + golden.topk_indices, + golden.topk_weights, + gate_up_w, + gate_up_b, + down_w, + down_b, + bf16_intermediates=False, + apply_clamp=True, + ) + + assert torch.allclose(smoke.output, golden.output, atol=1e-6, rtol=1e-6) + + +def test_fixed_routing_quantized_smoke_matches_golden_b(): + args = _tiny_inputs(scale=0.08) + golden_a = gpt_oss_moe_golden_a( + *args, + experts_per_token=2, + bf16_intermediates=True, + apply_clamp=True, + ) + x, _, _, gate_up_w, gate_up_b, down_w, down_b = args + golden_b = gpt_oss_moe_golden_b_plena_mxfp8( + x, + golden_a.topk_indices, + golden_a.topk_weights, + gate_up_w, + gate_up_b, + down_w, + down_b, + bf16_intermediates=True, + apply_clamp=True, + ) + + smoke = gpt_oss_moe_fixed_routing_host_smoke( + x, + golden_a.topk_indices, + golden_a.topk_weights, + gate_up_w, + gate_up_b, + down_w, + down_b, + quantize_expert_weights_to_plena_mxfp8=True, + bf16_intermediates=True, + apply_clamp=True, + ) + + assert torch.allclose(smoke.output.float(), golden_b.output.float(), atol=0.0, rtol=0.0) + + +def test_clamp_inactive_check_fails_when_gate_or_up_hits_limit(): + preact = torch.tensor([[[8.0, 0.0, 0.0, -8.0]]]) + try: + assert_clamp_inactive(preact) + except AssertionError as exc: + assert "clamp would be active" in str(exc) + else: + raise AssertionError("assert_clamp_inactive should reject active clamp inputs") + + +def test_golden_b_uses_fixed_routing_and_plena_mxfp8_expert_weights(): + args = _tiny_inputs(scale=0.08) + a = gpt_oss_moe_golden_a(*args, experts_per_token=2, bf16_intermediates=True) + x, _, _, gate_up_w, gate_up_b, down_w, down_b = args + + b = gpt_oss_moe_golden_b_plena_mxfp8( + x, + a.topk_indices, + a.topk_weights, + gate_up_w, + gate_up_b, + down_w, + down_b, + bf16_intermediates=True, + ) + + assert torch.equal(a.topk_indices, b.topk_indices) + assert b.output.shape == a.output.shape + assert b.router_logits.numel() == 0 + assert not torch.equal(a.output, b.output), "MXFP8 expert quantization should affect this seeded case" + + +if __name__ == "__main__": + test_bf16_represents_gpt_oss_clamp_boundary_exactly() + test_gpt_oss_swiglu_matches_manual_formula() + test_golden_a_matches_hf_gpt_oss_mlp_tiny() + test_clamp_inactive_smoke_precondition_and_unclamped_match() + test_fixed_routing_wiring_smoke_matches_golden_a_without_clamp_or_quantization() + test_fixed_routing_clamp_active_smoke_matches_golden_a_before_quantization() + test_fixed_routing_quantized_smoke_matches_golden_b() + test_clamp_inactive_check_fails_when_gate_or_up_hits_limit() + test_golden_b_uses_fixed_routing_and_plena_mxfp8_expert_weights() + print("PASS test_gpt_oss_moe_reference") diff --git a/aten/tests/test_plena_compiler.py b/aten/tests/test_plena_compiler.py index 22bf3ab..608f557 100644 --- a/aten/tests/test_plena_compiler.py +++ b/aten/tests/test_plena_compiler.py @@ -74,6 +74,24 @@ def test_fpvar_helper_uses_canonical_emit_path(): print(" PASS test_fpvar_helper_uses_canonical_emit_path") +def test_tile_row_minmax_fp_helpers_emit_vector_scalar_clamp_ops(): + """GPT-OSS clamp helpers should lower to V_MIN_VF/V_MAX_VF.""" + from compiler.aten.plena import PlenaCompiler + + prog = PlenaCompiler(mlen=8, blen=2) + x = prog.alloc("X", 8, 8) + limit = prog.fp_var("limit", size=1) + + prog.tile_row_min_fp(x, limit.address, rows=[0, 1]) + prog.tile_row_max_fp(x, limit.address, rows=[0, 1]) + code = prog.get_code() + + assert "V_MIN_VF" in code + assert "V_MAX_VF" in code + assert "S_LD_FP" in code + print(" PASS test_tile_row_minmax_fp_helpers_emit_vector_scalar_clamp_ops") + + def test_hbm_load_helper_uses_typed_legalization(): """Converted HBM load helpers should legalize typed large immediates.""" from compiler.aten.plena import IsaCompiler @@ -208,7 +226,7 @@ def test_compiler_threads_runtime_memory_geometry(): assert prog.mram_capacity_elems == 3 * tile_elems assert prog.mram_allocator.alignment == tile_elems assert prog.mram_allocator.total_size == 3 * tile_elems - assert prog.vram_allocator.alignment == 256 + assert prog.vram_allocator.alignment == tile_elems print(" PASS test_compiler_threads_runtime_memory_geometry") @@ -235,6 +253,103 @@ def test_linear_projection_uses_runtime_mram_tile_capacity(): print(" PASS test_linear_projection_uses_runtime_mram_tile_capacity") +def test_packed_skinny_stream_k_probe_compiles_cap8_under_cap4_mram(): + """Packed-skinny router probe keeps eight K slices in one MRAM tile.""" + from compiler.aten.plena import PlenaCompiler + + mlen = 128 + blen = 4 + tiles_per_mlen = mlen // blen + hidden = 2048 + num_k_tiles = hidden // mlen + max_k_tiles_per_packed_tile = 8 + num_groups = num_k_tiles // max_k_tiles_per_packed_tile + + prog = PlenaCompiler(mlen=mlen, blen=blen, mram_tile_capacity=4) + x_input = prog.input( + "X", + shape=(4, hidden), + physical_shape=(4, hidden), + real_data_ratio=1.0, + ) + x = prog.load_batch(x_input, name="X") + packed_weight = prog.input( + "W_router_packed_skinny_probe", + shape=(num_groups * mlen, tiles_per_mlen * mlen), + physical_shape=(num_groups * mlen, tiles_per_mlen * mlen), + real_data_ratio=1.0, + ) + logits = prog.alloc("router_logits", 4, 128, strict=False, physical_shape=(4, 128)) + + prog.vram_sub_projection_packed_skinny_stream_k_accum_to( + x, + 0, + packed_weight, + 0, + logits, + 0, + 0, + max_k_tiles_per_packed_tile=max_k_tiles_per_packed_tile, + ) + code = prog.get_code() + + # One HBM->MRAM tile per (micro-column, K group). A non-packed cap8 + # implementation would need eight full tiles live and overflow cap4 MRAM. + assert code.count("H_PREFETCH_M") == tiles_per_mlen * num_groups + assert code.count("M_MM 0,") == tiles_per_mlen * num_k_tiles + assert code.count("M_MM_WO") == tiles_per_mlen + assert "VRAM Sub Projection packed skinny microtile" in code + for offset in range(0, max_k_tiles_per_packed_tile * blen, blen): + assert f"S_ADDI_INT gp2, gp0, {offset}" in code + + print(" PASS test_packed_skinny_stream_k_probe_compiles_cap8_under_cap4_mram") + + +def test_qwen_packed_skinny_router_rowpacked_compiles_for_128_experts(): + """Integrated Qwen packed-skinny router emits V_TOPK rowpacked logits.""" + from compiler.aten.plena import PlenaCompiler + + mlen = 64 + blen = 4 + rows = 4 + hidden = 128 + num_experts = 128 + k_tiles_per_packed_tile = 8 + num_k_tiles = hidden // mlen + num_groups = 1 + output_blocks = num_experts // mlen + microcols = num_experts // blen + + prog = PlenaCompiler(mlen=mlen, blen=blen, mram_tile_capacity=4) + x_input = prog.input("X", shape=(rows, hidden), physical_shape=(rows, hidden), real_data_ratio=1.0) + x = prog.load_batch(x_input, name="X") + packed_weight = prog.input( + "W_router_packed_skinny_qwen", + shape=(num_groups * mlen, microcols * mlen), + physical_shape=(num_groups * mlen, microcols * mlen), + real_data_ratio=1.0, + ) + + logits = prog.qwen3_router_logits_packed_skinny_bf16_rowpacked_v0( + x, + packed_weight, + rows=rows, + hidden=hidden, + num_experts=num_experts, + k_tiles_per_packed_tile=k_tiles_per_packed_tile, + ) + code = prog.compile() + + assert logits.shape == (rows * output_blocks, mlen) + assert code.count("H_PREFETCH_M") == microcols * num_groups + assert code.count("M_MM 0,") == microcols * num_k_tiles + assert code.count("M_MM_WO") == microcols + assert code.count("V_ADD_VF") >= rows * output_blocks + assert "Qwen router packed-skinny logits pack" in code + + print(" PASS test_qwen_packed_skinny_router_rowpacked_compiles_for_128_experts") + + def test_vram_layout_tracks_logical_and_physical_shape(): """Layouts should keep native logical rows while allocating physical BLEN row storage.""" from compiler.aten.plena import PlenaCompiler @@ -812,6 +927,7 @@ def test_native_compile_assembles(): test_isa_builder_legalizes_large_absolute_immediates, test_isa_builder_legalizes_relative_large_immediates, test_fpvar_helper_uses_canonical_emit_path, + test_tile_row_minmax_fp_helpers_emit_vector_scalar_clamp_ops, test_hbm_load_helper_uses_typed_legalization, test_vram_fill_zero_all_column_blocks, test_vram_add_all_column_blocks, @@ -820,6 +936,7 @@ def test_native_compile_assembles(): test_mram_allocator_scales_with_runtime_mlen, test_compiler_threads_runtime_memory_geometry, test_linear_projection_uses_runtime_mram_tile_capacity, + test_packed_skinny_stream_k_probe_compiles_cap8_under_cap4_mram, test_vram_layout_tracks_logical_and_physical_shape, test_partial_row_linear_uses_one_blen_row_group, test_ffn_workspace_uses_allocator_and_avoids_rope_tables, diff --git a/doc/operation.svh b/doc/operation.svh index aed6e7d..694d2da 100644 --- a/doc/operation.svh +++ b/doc/operation.svh @@ -176,7 +176,10 @@ typedef enum logic [instruction_pkg::OPCODE_WIDTH - 1:0] { V_PS_V = 6'h31, V_SHFT_V = 6'h32, C_HADAMARD_TRANSFORM = 6'h33, - C_BREAK = 6'h34 + C_BREAK = 6'h34, + V_MAX_VF = 6'h35, + V_MIN_VF = 6'h36, + V_TOPK = 6'h37 } CUSTOM_ISA_OPCODE; typedef enum logic [2:0] { @@ -222,4 +225,4 @@ typedef struct { logic update_v_waddr; } OP_BUNDLE; -`endif \ No newline at end of file +`endif diff --git a/doc/plena_isa_spec.md b/doc/plena_isa_spec.md index 9d16666..787f122 100644 --- a/doc/plena_isa_spec.md +++ b/doc/plena_isa_spec.md @@ -287,6 +287,48 @@ Similar to `V_ADD_VV`, but performs element-wise multiplication. Similar to `V_ADD_VF`, but performs element-wise multiplication. +### V_MAX_VF + +**Format:** `V_MAX_VF rd, rs1, fp2, rmask` + +**Operation:** `Vector[gp_reg] = max(Vector[gp_reg], Broadcast(fp_reg))` + +**Description:** + +Element-wise maximum between a vector and a scalar. This is the lower-bound clamp primitive used by GPT-OSS expert activation, e.g. `max(up, -7.0)`. + +### V_MIN_VF + +**Format:** `V_MIN_VF rd, rs1, fp2, rmask` + +**Operation:** `Vector[gp_reg] = min(Vector[gp_reg], Broadcast(fp_reg))` + +**Description:** + +Element-wise minimum between a vector and a scalar. This is the upper-bound clamp primitive used by GPT-OSS expert activation, e.g. `min(gate, 7.0)` and `min(up, 7.0)`. + +### V_TOPK + +**Format:** `V_TOPK rd, rs1, rs2, rmask` + +**Operation:** Routed-MoE router top-k helper over one BF16 logits row. + +**Description:** + +`rs1` contains the Vector SRAM address of a router-logit row. The instruction +scans logits in descending logit order and breaks exact ties by smaller expert +index. `rmask` selects the routed-MoE policy: + +- `rmask=0`: scan 32 logits, select top-4, and write weights/indices to + `FP_MEM[gp_reg + 0..3]` / `INT_MEM[gp_reg + 0..3]`. +- `rmask=1`: scan 128 logits, select top-8, and write weights/indices to + `FP_MEM[gp_reg + 0..7]` / `INT_MEM[gp_reg + 0..7]`. + +Weights are softmax-over-selected logits. + +This is a correctness-first v0 linear scan. Bitonic/performance top-k remains a +future optimization. + ### V_EXP_V **Format:** `V_EXP_V rd, rs1, rmask` From e252dd985e6eb914644c4ab6c5a27480f6ae0afb Mon Sep 17 00:00:00 2001 From: Michael Date: Mon, 29 Jun 2026 15:59:37 +0100 Subject: [PATCH 2/3] chore: organize routed MoE compiler modules --- README.md | 11 ++++++++++- aten/models/__init__.py | 1 + aten/models/gpt_oss/__init__.py | 1 + .../gpt_oss/moe_reference.py} | 0 .../gpt_oss/real_layer_utils.py} | 0 aten/plena/compiler.py | 4 ++-- .../{program_gpt_oss_moe.py => program_routed_moe.py} | 4 ++-- aten/tests/test_gpt_oss_moe_assertions.py | 2 +- aten/tests/test_gpt_oss_moe_reference.py | 2 +- 9 files changed, 18 insertions(+), 7 deletions(-) create mode 100644 aten/models/__init__.py create mode 100644 aten/models/gpt_oss/__init__.py rename aten/{gpt_oss_moe.py => models/gpt_oss/moe_reference.py} (100%) rename aten/{gpt_oss_real_layer0_utils.py => models/gpt_oss/real_layer_utils.py} (100%) rename aten/plena/{program_gpt_oss_moe.py => program_routed_moe.py} (99%) diff --git a/README.md b/README.md index 4d0348e..201fc26 100644 --- a/README.md +++ b/README.md @@ -1 +1,10 @@ -# PLENA Compiler \ No newline at end of file +# 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/`. diff --git a/aten/models/__init__.py b/aten/models/__init__.py new file mode 100644 index 0000000..9267049 --- /dev/null +++ b/aten/models/__init__.py @@ -0,0 +1 @@ +"""Model-specific reference helpers for ATen/PLENA bring-up.""" diff --git a/aten/models/gpt_oss/__init__.py b/aten/models/gpt_oss/__init__.py new file mode 100644 index 0000000..cc275e0 --- /dev/null +++ b/aten/models/gpt_oss/__init__.py @@ -0,0 +1 @@ +"""GPT-OSS model semantics and reference utilities.""" diff --git a/aten/gpt_oss_moe.py b/aten/models/gpt_oss/moe_reference.py similarity index 100% rename from aten/gpt_oss_moe.py rename to aten/models/gpt_oss/moe_reference.py diff --git a/aten/gpt_oss_real_layer0_utils.py b/aten/models/gpt_oss/real_layer_utils.py similarity index 100% rename from aten/gpt_oss_real_layer0_utils.py rename to aten/models/gpt_oss/real_layer_utils.py diff --git a/aten/plena/compiler.py b/aten/plena/compiler.py index 6940246..08f0295 100644 --- a/aten/plena/compiler.py +++ b/aten/plena/compiler.py @@ -8,7 +8,7 @@ from compiler.aten.plena.isa_compiler import IsaCompiler from compiler.aten.plena.program_attention import ProgramAttentionMixin from compiler.aten.plena.program_fp_tile_ops import ProgramFPTileOpsMixin -from compiler.aten.plena.program_gpt_oss_moe import ProgramGptOssMoeMixin +from compiler.aten.plena.program_routed_moe import ProgramRoutedMoeMixin from compiler.aten.plena.program_matrix_ops import ProgramMatrixOpsMixin from compiler.aten.plena.program_tensors import ProgramTensorMixin from compiler.aten.plena.vars import FPVar, InputVar, TensorVar @@ -57,7 +57,7 @@ class PlenaCompiler( ProgramTensorMixin, ProgramFPTileOpsMixin, ProgramMatrixOpsMixin, - ProgramGptOssMoeMixin, + ProgramRoutedMoeMixin, ProgramAttentionMixin, IsaCompiler, ): diff --git a/aten/plena/program_gpt_oss_moe.py b/aten/plena/program_routed_moe.py similarity index 99% rename from aten/plena/program_gpt_oss_moe.py rename to aten/plena/program_routed_moe.py index 5d7f5c0..bcf5e53 100644 --- a/aten/plena/program_gpt_oss_moe.py +++ b/aten/plena/program_routed_moe.py @@ -13,7 +13,7 @@ ExpertBiases = tuple[VRAMMatrixVar | None, VRAMMatrixVar | None, VRAMMatrixVar | None] -class ProgramGptOssMoeMixin: +class ProgramRoutedMoeMixin: """Routed-MoE v0 emit helpers used by GPT-OSS and Qwen bring-up. The helpers intentionally keep routing policy explicit: GPT-OSS uses @@ -1809,4 +1809,4 @@ def moe_fixed_routing_v0( ) -__all__ = ["ProgramGptOssMoeMixin"] +__all__ = ["ProgramRoutedMoeMixin"] diff --git a/aten/tests/test_gpt_oss_moe_assertions.py b/aten/tests/test_gpt_oss_moe_assertions.py index 669792f..8f3af08 100644 --- a/aten/tests/test_gpt_oss_moe_assertions.py +++ b/aten/tests/test_gpt_oss_moe_assertions.py @@ -5,7 +5,7 @@ import pytest import torch -from aten.gpt_oss_moe import assert_compare_within, assert_gap_in_band, compare_stats +from aten.models.gpt_oss.moe_reference import assert_compare_within, assert_gap_in_band, compare_stats def _reference_tensor() -> torch.Tensor: diff --git a/aten/tests/test_gpt_oss_moe_reference.py b/aten/tests/test_gpt_oss_moe_reference.py index 66d66c6..0f2e9f1 100644 --- a/aten/tests/test_gpt_oss_moe_reference.py +++ b/aten/tests/test_gpt_oss_moe_reference.py @@ -21,7 +21,7 @@ if os.path.isdir(_path) and _path not in sys.path: sys.path.insert(0, _path) -from aten.gpt_oss_moe import ( +from aten.models.gpt_oss.moe_reference import ( assert_clamp_inactive, clamp_stats, gpt_oss_moe_fixed_routing_host_smoke, From 54a68e6969651eb5c1c3397ac59bed6585c08806 Mon Sep 17 00:00:00 2001 From: Michael C Li Date: Wed, 1 Jul 2026 03:16:42 +0100 Subject: [PATCH 3/3] Align V_SHFT_V opcode with RTL ISA --- asm_templates/im2col_asm.py | 8 ++++---- aten/ops/plena/conv_ops.py | 6 +++--- aten/plena/program_attention.py | 4 ++-- doc/plena_isa_spec.md | 4 ++-- generator/passes/code_gen.py | 2 +- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/asm_templates/im2col_asm.py b/asm_templates/im2col_asm.py index 6edd6a5..db014b1 100644 --- a/asm_templates/im2col_asm.py +++ b/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") diff --git a/aten/ops/plena/conv_ops.py b/aten/ops/plena/conv_ops.py index e4ec74c..1b0ae50 100644 --- a/aten/ops/plena/conv_ops.py +++ b/aten/ops/plena/conv_ops.py @@ -5,8 +5,8 @@ (H_PREFETCH_V + V_MUL_VV + V_RED_SUM + S_ST_FP + S_MAP_V_FP). Slower because each patch element is extracted scalar-by-scalar through FP_SRAM. - Opt-in (use_shift=True, or CONV_USE_SHIFT=1 env var): uses the - V_SHIFT_V opcode for vector-level patch placement. Faster for long - patches. Requires the emulator to support V_SHIFT_V (opcode 0x31, + V_SHFT_V opcode for vector-level patch placement. Faster for long + patches. Requires the emulator to support V_SHFT_V (opcode 0x32, LSB-first right-shift per main.rs fix by George Wu, commit 24eb011). After im2col the systolic matmul uses the compiler's standard linear path. @@ -42,7 +42,7 @@ def conv2d_plena( ): """ PLENA backend: hardware im2col + systolic matmul. - Default path avoids V_SHIFT_V; use_shift=True opts into the vector-shift path. + Default path avoids V_SHFT_V; use_shift=True opts into the vector-shift path. Args: prog: PlenaCompiler instance. diff --git a/aten/plena/program_attention.py b/aten/plena/program_attention.py index 6e86084..e748571 100644 --- a/aten/plena/program_attention.py +++ b/aten/plena/program_attention.py @@ -642,7 +642,7 @@ def _pack_o_head_to_output( *load_large_int(gp_scratch, scratch_address), *load_large_int(gp_shift, shift), f"C_LOOP_START gp{gp_loop}, {rows}", - f"V_SHIFT_V gp{gp_scratch}, gp{gp_src}, gp{gp_shift}", + f"V_SHFT_V gp{gp_scratch}, gp{gp_src}, gp{gp_shift}", f"V_ADD_VV gp{gp_dst}, gp{gp_dst}, gp{gp_scratch}, 0", f"S_ADDI_INT gp{gp_src}, gp{gp_src}, {self.mlen}", f"S_ADDI_INT gp{gp_dst}, gp{gp_dst}, {self.mlen}", @@ -672,7 +672,7 @@ def _pack_o_head_to_output_dynamic( *load_large_int(gp_scratch, scratch_address), *load_large_int(gp_shift, shift), f"C_LOOP_START gp{gp_loop}, {rows}", - f"V_SHIFT_V gp{gp_scratch}, gp{gp_src}, gp{gp_shift}", + f"V_SHFT_V gp{gp_scratch}, gp{gp_src}, gp{gp_shift}", f"V_ADD_VV gp{gp_dst}, gp{gp_dst}, gp{gp_scratch}, 0", f"S_ADDI_INT gp{gp_src}, gp{gp_src}, {self.mlen}", f"S_ADDI_INT gp{gp_dst}, gp{gp_dst}, {self.mlen}", diff --git a/doc/plena_isa_spec.md b/doc/plena_isa_spec.md index 787f122..055f5a0 100644 --- a/doc/plena_isa_spec.md +++ b/doc/plena_isa_spec.md @@ -752,9 +752,9 @@ End of a hardware loop. If the loop counter (in register `rd`) is greater than 0 ## Extensions -### V_SHIFT_V +### V_SHFT_V -**Format:** `V_SHIFT_V rd, rs1, rs2` +**Format:** `V_SHFT_V rd, rs1, rs2` **Operation:** `Vector[gp_reg] = ElementShift(Vector[gp_reg], gp_reg)` diff --git a/generator/passes/code_gen.py b/generator/passes/code_gen.py index 81b4227..47ee236 100644 --- a/generator/passes/code_gen.py +++ b/generator/passes/code_gen.py @@ -484,7 +484,7 @@ def _generate_conv2d_code( conv_stride = dims.get("stride", patch_size) # Check whether every output column produces a 64-aligned HBM pixel - # offset. When it does, the fast V_SHIFT_V template can be used; + # offset. When it does, the fast V_SHFT_V template can be used; # otherwise fall back to the no-shift (basis-vector) template which # handles arbitrary alignment via per-element extraction. _all_cols_aligned = all((ow * conv_stride) % 64 == 0 for ow in range(OW))