diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3732abc..0c5fb62 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -58,6 +58,29 @@ jobs: - name: Run the stage-terminator guard run: PYTHONPATH=. python3 -m pytest aten/tests/test_moe_stage_terminator.py -v + moe-export-tests: + name: MoE numerical exporter + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/checkout@v4 + with: + repository: AICrossSim/PLENA_Tools + ref: 6b31fe00d99564909c5aa68d68ec9a5d41981806 + path: PLENA_Tools + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + - name: Install numerical test dependencies + run: | + pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cpu + pip install numpy bitstring pytest pyyaml + - name: Validate encoded bytes and independent numerical reference + run: | + PYTHONPATH=.:PLENA_Tools python -m pytest \ + aten/tests/test_moe_normal_export.py \ + aten/tests/test_moe_full_shape_export.py -q + unit-tests: name: Generator unit tests runs-on: ubuntu-latest diff --git a/aten/plena/moe_full_shape_export.py b/aten/plena/moe_full_shape_export.py new file mode 100644 index 0000000..eed278b --- /dev/null +++ b/aten/plena/moe_full_shape_export.py @@ -0,0 +1,179 @@ +"""Stream full-dimension, synthetic-value MoE fixtures through the real codec. + +The NumPy oracle vectorizes independent outputs only: every dot product still +uses ascending K, separate FP32 multiply/add, and the V0 BF16 boundaries. +No BLAS summation order or dense model inference is implied. +""" +from __future__ import annotations +import hashlib +import inspect +import math +from pathlib import Path +import numpy as np +import torch +import moe_normal_export as base + + +def round_bf16(values): + values = np.asarray(values, dtype=np.float32) + if not np.isfinite(values).all(): + raise ValueError("nonfinite FP32 intermediate") + bits = values.view(np.uint32) + rounded = ((bits + np.uint32(0x7fff) + ((bits >> 16) & 1)) >> 16).astype(np.uint16) + result = (rounded.astype(np.uint32) << 16).view(np.float32) + if not np.isfinite(result).all(): + raise ValueError("BF16 overflow") + return result + + +def generated_matrix(rows, cols, seed, activation=False): + r = np.arange(rows, dtype=np.uint32)[:, None] + c = np.arange(cols, dtype=np.uint32)[None, :] + mixed = ((r + 1) * np.uint32(2654435761)) ^ ((c + seed + 1) * np.uint32(2246822519)) + mixed ^= mixed >> 13 + values = (mixed % 31).astype(np.float32) - 15 + scale = 1 / 64 if activation else 2.0 ** (-math.ceil(math.log2(math.sqrt(cols))) - 5) + return values * np.float32(scale) + + +def pack_e4m3_vectorized(signed_exponent, signed_mantissa): + """Same arithmetic as plena_quant.pack_fp_to_bin, with a batched check. + + The reference packer's per-element Python assertions dominate full weight + export. This replaces only that loop, retaining its FP32 packing arithmetic. + Compatibility is tested over every finite E4M3 encoding and quantized tails. + """ + exponent = signed_exponent + 7 + if not torch.isfinite(exponent).all() or not torch.isfinite(signed_mantissa).all(): + raise ValueError('nonfinite codec fields') + if torch.any((exponent < 0) | (exponent > 15)): + raise ValueError('exponent outside E4M3') + sign = (signed_mantissa < 0).to(torch.int64) + mantissa = torch.abs(signed_mantissa) + fraction = torch.where(exponent == 0, mantissa, mantissa - 1) * 8 + return (sign * 128 + exponent * 8 + fraction).int() + + +def encode_array(values): + from plena_quant.mxfp.quantizer import _mx_fp_quantize_hardware + values = np.asarray(values, dtype=np.float32) + if values.ndim != 2 or min(values.shape) == 0 or not np.isfinite(values).all(): + raise ValueError("weights require a finite nonempty matrix") + n, k = values.shape + stride = ((k + 7) // 8) * 8 + padded = np.pad(values, ((0, 0), (0, stride-k))) + with torch.no_grad(): + _, exp, man, scales = _mx_fp_quantize_hardware( + torch.from_numpy(padded), width=8, exponent_width=4, + exponent_bias_width=8, block_size=[8]) + elements = pack_e4m3_vectorized(exp, man) + elements = elements.cpu().numpy().astype(np.uint8).reshape(n, stride) + raw_scales = scales.cpu().numpy().reshape(n, stride // 8) + if np.any((raw_scales < 0) | (raw_scales > 254)) or np.any((elements & 0x78) == 0x78): + raise ValueError("nonfinite encoded weight") + scales = raw_scales.astype(np.uint8) + # Independently decode using the scalar reference's format definitions. + lookup = np.array([base.decode_element(i, 127) if (i & 0x78) != 0x78 else 0 + for i in range(256)], dtype=np.float32) + scale_values = np.ldexp(np.ones(scales.shape, dtype=np.float32), scales.astype(np.int32)-127) + scale_values[scales == 0] = 0 + decoded = round_bf16(lookup[elements] * np.repeat(scale_values, 8, axis=1))[:, :k].copy() + return elements.tobytes(), scales.tobytes(), decoded, dict( + rows=n, cols=k, element_row_stride=stride, scale_row_stride=stride // 8) + + +def gemm_reference(x, weights): + x, weights = np.asarray(x, dtype=np.float32), np.asarray(weights, dtype=np.float32) + if x.shape[1] != weights.shape[1]: + raise ValueError("GEMM reduction dimensions differ") + accum = np.zeros((x.shape[0], weights.shape[0]), dtype=np.float32) + product = np.empty_like(accum) + for k in range(x.shape[1]): + np.multiply(x[:, k, None], weights[None, :, k], out=product) + np.add(accum, product, out=accum) + return round_bf16(accum) + + +def swiglu_reference(gate, up): + # Match the scalar Python oracle's double exp -> FP32 boundary. Vector + # library approximations to exp need not have identical rounding. + exponential = np.fromiter((base.f32(math.exp(-float(g))) for g in gate.flat), + dtype=np.float32, count=gate.size).reshape(gate.shape) + return round_bf16(np.multiply(np.divide(gate, np.add(np.float32(1), exponential)), up)) + + +def export_full_shape(output_dir, *, input_dim, expert_hidden_dim, tokens, + routes, name, provenance=None): + """Generate synthetic values at actual dimensions; routes are caller data. + + The regenerated FP32 arrays are bound by a streaming SHA-256 and the exact + generator source. Only active routed experts are exported. No shared model + gate/expert, router, residual or full model execution is represented. + """ + if min(input_dim, expert_hidden_dim, tokens) <= 0: + raise ValueError("positive dimensions and tokens are required") + output = Path(output_dir); output.mkdir(parents=True, exist_ok=True) + ids = {r['expert'] for r in routes} + if not ids: + raise ValueError("full-shape campaign requires active experts") + groups = base.group_routes(routes, tokens, ids) + normalized = sorted([r for g in groups for r in g['routes']], key=lambda r:(r['token'],r['slot'])) + x_raw = generated_matrix(tokens, input_dim, 91, activation=True) + x = round_bf16(x_raw) + source_hash = hashlib.sha256(x_raw.tobytes()) + route_index = {(r['token'],r['slot']):i for i,r in enumerate(normalized)} + contributions = np.zeros((len(normalized), input_dim), dtype=np.float32) + experts = [] + image_path = output / 'weights.bin' + image_hash = hashlib.sha256() + with image_path.open('wb') as image: + def append(payload): + padding = bytes((-image.tell()) % 64) + image.write(padding); image_hash.update(padding) + address = image.tell(); image.write(payload); image_hash.update(payload) + return address + for group in groups: + expert = group['expert']; views = {'id':expert}; decoded = {} + for stage, n, k, offset in [('gate',expert_hidden_dim,input_dim,1), + ('up',expert_hidden_dim,input_dim,2), + ('down',input_dim,expert_hidden_dim,3)]: + values = generated_matrix(n, k, expert*3+offset) + source_hash.update(values.tobytes()) + elements, scales, decoded[stage], shape = encode_array(values) + views[stage] = dict(shape, element_base=append(elements), scale_base=append(scales)) + inputs = x[[r['token'] for r in group['routes']]] + gate = gemm_reference(inputs, decoded['gate']) + up = gemm_reference(inputs, decoded['up']) + y = gemm_reference(swiglu_reference(gate, up), decoded['down']) + for r, values in zip(group['routes'], y): + contributions[route_index[(r['token'],r['slot'])]] = values + experts.append(views) + append(b'') + sums = np.zeros_like(x) + for r, values in zip(normalized, contributions): + sums[r['token']] += np.float32(r['weight']) * values + final = round_bf16(sums) + import plena_quant.mxfp.quantizer as quantizer + import plena_quant.mxfp.utils as packer + import plena_quant.common.minifloat as minifloat + import plena_quant.common.hardware_utils as hardware_utils + import plena_quant.common.utils as common_utils + sources = [Path(__file__), Path(base.__file__), *(Path(inspect.getfile(m)) for m in + [quantizer, packer, minifloat, hardware_utils, common_utils])] + metadata = dict(weight_format=base.FORMAT, weight_layout='output_major_N_K', + block_size=8, scale_axis='per_row_reduction_K', hbm_sha256=image_hash.hexdigest(), + synthetic_source_arrays_sha256=source_hash.hexdigest(), quantizer_torch_version=torch.__version__, + sources=[dict(path=str(p.resolve()),sha256=hashlib.sha256(p.read_bytes()).hexdigest()) for p in sources], + evidence_scope='full matrix dimensions; synthetic inputs/weights; caller fixed routed experts; not full-model inference', + provenance=provenance or {}) + workload = dict(schema_version=1, name=name, input_dim=input_dim, expert_hidden_dim=expert_hidden_dim, + hbm_file='weights.bin', inputs_bf16=(x.view(np.uint32)>>16).tolist(), routes=normalized, + experts=experts, shared_expert=None, grouped_routes=groups, metadata=metadata) + workload_bytes = base._json_bytes(workload) + golden = dict(schema_version=1, name=name, semantics='ascending K separate FP32 multiply/add; BF16 stages; scalar exp; stable token/slot combine', + output_bf16=(final.view(np.uint32)>>16).tolist(), output_f32=final.tolist(), + pre_round_output_f32=sums.tolist(), workload_sha256=hashlib.sha256(workload_bytes).hexdigest(), + hbm_sha256=image_hash.hexdigest()) + (output/'workload.json').write_bytes(workload_bytes) + (output/'golden.json').write_bytes(base._json_bytes(golden)) + return workload, golden diff --git a/aten/plena/moe_normal_export.py b/aten/plena/moe_normal_export.py new file mode 100644 index 0000000..80461bf --- /dev/null +++ b/aten/plena/moe_normal_export.py @@ -0,0 +1,347 @@ +"""Export ready-input MoE fixtures for the dual normal-buffer execution harness. + +This is a manifest exporter, not a router or an ISA lowering. Weights use the +repository's PLENA E4M3/E8M0 block-8 codec (including its non-IEEE subnormals), +with output-major [N,K] rows. The independent scalar reference decodes the +exported bytes, accumulates GEMMs in ascending K with FP32 multiply/add, and +rounds each GEMM and SwiGLU output to BF16. Combine follows token/slot order. + +Run this file with PLENA_Tools on PYTHONPATH and an interpreter with torch. +""" +from __future__ import annotations + +import argparse +import hashlib +import inspect +import json +import math +from pathlib import Path +import struct +from typing import Any, Mapping, Sequence + +BLOCK_SIZE = 8 +FORMAT = "plena_e4m3_e8m0_block8" + + +def f32(value: float) -> float: + """One explicit FP32 rounding; no fused multiply/add in the reference.""" + try: + return struct.unpack(" int: + """Round FP32 to BF16, round-to-nearest-even (finite fixture values only).""" + value = f32(value) + if not math.isfinite(value): + raise ValueError("non-finite BF16 value or arithmetic overflow") + bits = struct.unpack("> 16) & 1)) >> 16 + if (rounded & 0x7F80) == 0x7F80: + raise ValueError("BF16 rounding overflow") + return rounded + + +def bf16_value(bits: int) -> float: + if not isinstance(bits, int) or not 0 <= bits <= 0xFFFF: + raise ValueError("BF16 bits must be uint16") + return struct.unpack(" float: + return bf16_value(bf16_bits(value)) + + +def _matrix(values: Any, label: str) -> list[list[float]]: + if hasattr(values, "detach"): + values = values.detach().cpu().tolist() + elif hasattr(values, "tolist"): + values = values.tolist() + rows = [[f32(x) for x in row] for row in values] + if not rows or not rows[0] or any(len(r) != len(rows[0]) for r in rows): + raise ValueError(f"{label} must be a nonempty rectangular matrix") + if not all(math.isfinite(x) for r in rows for x in r): + raise ValueError(f"{label} contains non-finite or FP32-overflow values") + return rows + + +def _sha256(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _json_bytes(value: Any) -> bytes: + return (json.dumps(value, sort_keys=True, indent=2, allow_nan=False) + "\n").encode() + + +def encode_matrix(values: Any) -> tuple[bytes, bytes, dict[str, int]]: + """Use the actual PLENA_Tools quantizer; pad each reduction row to block 8. + + Element and scale payloads are separate. Padding belongs to its own row and + never shares a block with the next output channel. No transpose/requantize + is performed by the runtime. + """ + import torch + try: + from plena_quant.mxfp.quantizer import _mx_fp_quantize_hardware + from plena_quant.mxfp.utils import pack_fp_to_bin + except ImportError as exc: + raise ImportError("Place PLENA_Tools on PYTHONPATH for the actual PLENA codec") from exc + + rows = _matrix(values, "weight") + n, k = len(rows), len(rows[0]) + stride = ((k + BLOCK_SIZE - 1) // BLOCK_SIZE) * BLOCK_SIZE + padded = torch.tensor([r + [0.0] * (stride - k) for r in rows], dtype=torch.float32) + with torch.no_grad(): + _, exponents, mantissas, scales = _mx_fp_quantize_hardware( + padded, width=8, exponent_width=4, exponent_bias_width=8, block_size=[BLOCK_SIZE] + ) + elements = pack_fp_to_bin(exponents, mantissas, exp_width=4, man_width=3) + element_list = [int(x) for x in elements.flatten().tolist()] + scale_list = [int(x) for x in scales.flatten().tolist()] + if any(not 0 <= x <= 254 for x in scale_list): + raise ValueError("weight scale is outside finite PLENA E8M0 range") + if any((x & 0x78) == 0x78 for x in element_list): + raise ValueError("weight contains a non-finite encoded E4M3 element") + return bytes(element_list), bytes(scale_list), { + "rows": n, "cols": k, "element_row_stride": stride, + "scale_row_stride": stride // BLOCK_SIZE, + } + + +def decode_element(element: int, scale: int) -> float: + """Independent scalar interpretation of quantize::FpType's local format.""" + exponent, mantissa = (element >> 3) & 15, element & 7 + if exponent == 15 or scale == 255: + raise ValueError("non-finite PLENA weight encoding") + if scale == 0: + return 0.0 # existing Rust E8M0 cast treats exponent-zero/mantissa-zero as zero + magnitude = (mantissa / 8.0 if exponent == 0 else 1.0 + mantissa / 8.0) + value = math.ldexp(magnitude, exponent - 7 + scale - 127) + return _bf16(-value if element & 128 else value) + + +def decode_matrix(hbm: bytes, region: Mapping[str, int]) -> list[list[float]]: + rows, cols = region["rows"], region["cols"] + return [[decode_element( + hbm[region["element_base"] + n * region["element_row_stride"] + k], + hbm[region["scale_base"] + n * region["scale_row_stride"] + k // BLOCK_SIZE], + ) for k in range(cols)] for n in range(rows)] + + +def group_routes(routes: Sequence[Mapping[str, Any]], num_tokens: int, + expert_ids: set[int]) -> list[dict[str, Any]]: + """Canonical grouping preserves every (token, slot, expert, weight) tuple.""" + seen = set() + by_expert: dict[int, list[dict[str, Any]]] = {} + for route in routes: + for key in ("token", "slot", "expert"): + if isinstance(route[key], bool) or not isinstance(route[key], int): + raise ValueError(f"route {key} must be an integer") + token, slot, expert = route["token"], route["slot"], route["expert"] + weight = f32(route["weight"]) + if not 0 <= token < num_tokens or slot < 0 or expert not in expert_ids: + raise ValueError("route has out-of-range token/slot or unknown expert") + if not math.isfinite(weight): + raise ValueError("route weight must be finite FP32") + if (token, slot) in seen: + raise ValueError("duplicate token/slot route") + seen.add((token, slot)) + by_expert.setdefault(expert, []).append({ + "token": token, "slot": slot, "expert": expert, "weight": weight, + }) + return [{"expert": expert, "routes": sorted(group, key=lambda r: (r["token"], r["slot"]))} + for expert, group in sorted(by_expert.items())] + + +def _gemm_row(x: Sequence[float], weights: Sequence[Sequence[float]]) -> list[float]: + out = [] + for row in weights: + acc = 0.0 + for a, b in zip(x, row, strict=True): + acc = f32(acc + f32(a * b)) + out.append(_bf16(acc)) + return out + + +def _swiglu(gate: float, up: float) -> float: + try: + exponential = f32(math.exp(-gate)) + except OverflowError: + exponential = math.inf + return _bf16(f32(f32(gate / f32(1.0 + exponential)) * up)) + + +def numerical_reference(workload: Mapping[str, Any], hbm: bytes) -> dict[str, Any]: + """Execute exported data in a scalar oracle independent of Rust scheduling.""" + inputs = [[bf16_value(x) for x in row] for row in workload["inputs_bf16"]] + weights = {e["id"]: {stage: decode_matrix(hbm, e[stage]) for stage in ("gate", "up", "down")} + for e in workload["experts"]} + group_routes(workload["routes"], len(inputs), set(weights)) + cache: dict[tuple[int, int], list[float]] = {} + + def expert_output(token: int, expert: int) -> list[float]: + key = token, expert + if key not in cache: + matrices = weights[expert] + gate = _gemm_row(inputs[token], matrices["gate"]) + up = _gemm_row(inputs[token], matrices["up"]) + z = [_swiglu(g, u) for g, u in zip(gate, up, strict=True)] + cache[key] = _gemm_row(z, matrices["down"]) + return cache[key] + + accum = [[0.0] * workload["input_dim"] for _ in inputs] + contributions = [] + for route in sorted(workload["routes"], key=lambda r: (r["token"], r["slot"])): + token = route["token"] + values = expert_output(token, route["expert"]) + for d, value in enumerate(values): + accum[token][d] = f32(accum[token][d] + f32(f32(route["weight"]) * value)) + contributions.append({**route, "output_bf16": [bf16_bits(x) for x in values]}) + shared = workload.get("shared_expert") + if shared is not None: + for token in range(len(inputs)): + for d, value in enumerate(expert_output(token, shared["expert"])): + accum[token][d] = f32(accum[token][d] + f32(f32(shared["weight"]) * value)) + output_bits = [[bf16_bits(x) for x in row] for row in accum] + return { + "schema_version": 1, "name": workload["name"], + "semantics": "ascending-K FP32 mul/add; BF16 RNE GEMM and SwiGLU; stable token/slot FP32 combine; shared last; BF16 output", + "output_bf16": output_bits, + "output_f32": [[bf16_value(x) for x in row] for row in output_bits], + "pre_round_output_f32": accum, "route_contributions": contributions, + } + + +def export_workload(output_dir: str | Path, *, inputs: Any, + experts: Sequence[Mapping[str, Any]], routes: Sequence[Mapping[str, Any]], + name: str = "moe_normal", shared_expert: Mapping[str, Any] | None = None, + provenance: Mapping[str, Any] | None = None) -> dict[str, Any]: + """Export caller-supplied arrays and fixed route decisions, without executing a router. + + ``experts`` contains ``id, gate[E,D], up[E,D], down[D,E]``. An optional + ``shared_expert={expert: id, weight: float}`` references an entry in experts. + No biases, shared gating network, residual, attention or model loading is + implied. Runtime grouping and the numerical reference retain duplicate + token/expert assignments when their route slots differ. + """ + import torch + import plena_quant.mxfp.quantizer as quantizer + import plena_quant.mxfp.utils as packer + import plena_quant.common.minifloat as minifloat + import plena_quant.common.hardware_utils as hardware_utils + import plena_quant.common.utils as common_utils + + x = _matrix(inputs, "inputs") + d = len(x[0]) + normalized = [] + ids: set[int] = set() + hidden = None + for expert in experts: + expert_id = expert["id"] + if isinstance(expert_id, bool) or not isinstance(expert_id, int) or expert_id < 0 or expert_id in ids: + raise ValueError("expert IDs must be distinct nonnegative integers") + ids.add(expert_id) + matrices = {stage: _matrix(expert[stage], f"expert {expert_id} {stage}") + for stage in ("gate", "up", "down")} + e = len(matrices["gate"]) + if hidden is None: + hidden = e + if (e != hidden or len(matrices["gate"][0]) != d or len(matrices["up"]) != e + or len(matrices["up"][0]) != d or len(matrices["down"]) != d + or len(matrices["down"][0]) != e): + raise ValueError("all experts must have gate/up[E,D] and down[D,E] with a common E") + normalized.append({"id": expert_id, **matrices}) + if hidden is None: + raise ValueError("at least one expert is required") + grouped = group_routes(routes, len(x), ids) + normalized_routes = sorted([r for group in grouped for r in group["routes"]], + key=lambda r: (r["token"], r["slot"])) + shared = None + if shared_expert is not None: + shared = {"expert": shared_expert["expert"], "weight": f32(shared_expert["weight"])} + if (isinstance(shared["expert"], bool) or not isinstance(shared["expert"], int) + or shared["expert"] not in ids or not math.isfinite(shared["weight"])): + raise ValueError("invalid shared expert or weight") + hbm = bytearray() + + def append_aligned(payload: bytes) -> int: + hbm.extend(bytes((-len(hbm)) % 64)) + base = len(hbm) + hbm.extend(payload) + return base + + regions = [] + for expert in sorted(normalized, key=lambda e: e["id"]): + views: dict[str, Any] = {"id": expert["id"]} + for stage in ("gate", "up", "down"): + elements, scales, shape = encode_matrix(expert[stage]) + views[stage] = {**shape, "element_base": append_aligned(elements), + "scale_base": append_aligned(scales)} + regions.append(views) + # Pad the final request granule so a real 64B load never leaves the image. + hbm.extend(bytes((-len(hbm)) % 64)) + source_arrays = {"inputs_fp32": x, "experts_fp32": normalized, + "routes": normalized_routes, "shared_expert": shared} + sources = [Path(__file__), *(Path(inspect.getfile(module)) for module in + (quantizer, packer, minifloat, hardware_utils, common_utils))] + metadata = { + "evidence_scope": "caller_supplied_ready_inputs_and_routes; no runtime_router_or_full_model_claim", + "weight_format": FORMAT, "weight_layout": "output_major_N_K", + "quantizer_torch_version": torch.__version__, + "scale_axis": "per_row_reduction_K", "block_size": BLOCK_SIZE, + "hbm_sha256": _sha256(hbm), "source_arrays_sha256": _sha256(_json_bytes(source_arrays)), + "sources": [{"path": str(p.resolve()), "sha256": _sha256(p.read_bytes())} for p in sources], + "provenance": dict(provenance or {"scope": "caller_supplied_arrays"}), + } + workload = { + "schema_version": 1, "name": name, "input_dim": d, "expert_hidden_dim": hidden, + "hbm_file": "weights.bin", + "inputs_bf16": [[bf16_bits(v) for v in row] for row in x], + "routes": normalized_routes, "experts": regions, "shared_expert": shared, + "grouped_routes": grouped, "metadata": metadata, + } + golden = numerical_reference(workload, bytes(hbm)) + golden["workload_sha256"] = _sha256(_json_bytes(workload)) + golden["hbm_sha256"] = metadata["hbm_sha256"] + output = Path(output_dir) + output.mkdir(parents=True, exist_ok=True) + paths = {"workload": output / "workload.json", "hbm": output / "weights.bin", + "golden": output / "golden.json", "source": output / "source_arrays.json"} + paths["workload"].write_bytes(_json_bytes(workload)) + paths["hbm"].write_bytes(hbm) + paths["golden"].write_bytes(_json_bytes(golden)) + paths["source"].write_bytes(_json_bytes(source_arrays)) + return {"workload": workload, "golden": golden, "paths": {k: str(v.resolve()) for k, v in paths.items()}} + + +def demo_arrays() -> dict[str, Any]: + """Deterministic nonzero fixture: M=6,2,1,0; D=11/E=13 K and N tails.""" + d, e, t = 11, 13, 9 + inputs = [[((token * 7 + k * 3) % 19 - 9) / 16 for k in range(d)] for token in range(t)] + experts = [] + for expert in range(4): + matrices = {} + for stage, rows, cols, offset in (("gate", e, d, 1), ("up", e, d, 3), ("down", d, e, 5)): + matrices[stage] = [[((n * 7 + k * 11 + expert * 3 + offset) % 23 - 11) + * (2.0 ** (((n + k // 8 + expert) % 3) - 6)) + for k in range(cols)] for n in range(rows)] + experts.append({"id": expert, **matrices}) + routes = [{"token": token, "slot": 0, "expert": 0 if token < 6 else 1 if token < 8 else 2, + "weight": 0.625 + (token % 3) * 0.125} for token in range(t)] + return {"inputs": inputs, "experts": experts, "routes": routes} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--shared", action="store_true", help="also execute expert 3 for every token, weight 0.25") + args = parser.parse_args() + result = export_workload(args.output_dir, **demo_arrays(), name="synthetic_moe_normal_tails", + shared_expert={"expert": 3, "weight": 0.25} if args.shared else None, + provenance={"scope": "synthetic_shapes_inputs_routes_weights", "generator": "demo_arrays"}) + print(json.dumps(result["paths"], indent=2)) + + +if __name__ == "__main__": + main() diff --git a/aten/tests/test_moe_full_shape_export.py b/aten/tests/test_moe_full_shape_export.py new file mode 100644 index 0000000..433061e --- /dev/null +++ b/aten/tests/test_moe_full_shape_export.py @@ -0,0 +1,55 @@ +import importlib.util +from pathlib import Path +import sys +import numpy as np +import pytest +import torch + +SOURCE = Path(__file__).resolve().parents[1] / 'plena' +sys.path.insert(0, str(SOURCE)) +import moe_normal_export as scalar +import moe_full_shape_export as full + + +def test_packing_matches_reference_for_every_finite_e4m3_encoding(): + from plena_quant.mxfp.utils import pack_fp_to_bin + # Signed zero follows the existing packer's canonical positive-zero rule. + fields = [(e-7, (-1 if s else 1) * ((m/8) if e == 0 else 1+m/8)) + for s in range(2) for e in range(15) for m in range(8)] + exponent = torch.tensor([e for e,m in fields], dtype=torch.float32) + mantissa = torch.tensor([m for e,m in fields], dtype=torch.float32) + assert torch.equal(full.pack_e4m3_vectorized(exponent,mantissa), + pack_fp_to_bin(exponent,mantissa,exp_width=4,man_width=3)) + with pytest.raises(ValueError): + full.pack_e4m3_vectorized(torch.tensor([9.]), torch.tensor([1.])) + + +def test_vectorized_codec_matches_existing_exporter_and_scalar_decode(): + values = full.generated_matrix(13, 19, 27) + element, scale, decoded, shape = full.encode_array(values) + old_element, old_scale, old_shape = scalar.encode_matrix(values) + assert (element, scale, shape) == (old_element, old_scale, old_shape) + for row in range(13): + for k in range(19): + assert decoded[row,k] == scalar.decode_element(element[row*24+k],scale[row*3+k//8]) + + +def test_streamed_full_shape_oracle_matches_independent_scalar(tmp_path): + routes = [dict(token=t, slot=s, expert=(t+s)%3, weight=0.25*(s+1)) for t in range(5) for s in range(2)] + workload, golden = full.export_full_shape(tmp_path,input_dim=17,expert_hidden_dim=19, + tokens=5,routes=routes,name='oracle_check') + reference = scalar.numerical_reference(workload,(tmp_path/'weights.bin').read_bytes()) + assert golden['output_bf16'] == reference['output_bf16'] + assert golden['pre_round_output_f32'] == reference['pre_round_output_f32'] + assert any(value != 0 for row in golden['output_f32'] for value in row) + + +def test_gemm_preserves_ascending_k_cancellation_order(): + x = np.array([[1,1,1]],dtype=np.float32) + w = np.array([[2**24,1,-2**24]],dtype=np.float32) + assert full.gemm_reference(x,w).item() == 0 + + +def test_nonfinite_values_are_rejected(): + with pytest.raises(ValueError): full.round_bf16(np.array([np.inf],dtype=np.float32)) + with pytest.raises(ValueError): full.encode_array(np.array([[np.nan]],dtype=np.float32)) diff --git a/aten/tests/test_moe_normal_export.py b/aten/tests/test_moe_normal_export.py new file mode 100644 index 0000000..f4cb761 --- /dev/null +++ b/aten/tests/test_moe_normal_export.py @@ -0,0 +1,165 @@ +"""Byte-contract, numerical and routing checks for the normal-buffer exporter.""" +import hashlib +import json +import math + +import pytest + +from compiler.aten.plena.moe_normal_export import ( + bf16_bits, bf16_value, decode_element, decode_matrix, demo_arrays, + encode_matrix, export_workload, group_routes, numerical_reference, +) + + +def _unit_arrays(): + return {"inputs": [[1.0]], + "experts": [{"id": 3, "gate": [[1.0]], "up": [[1.0]], "down": [[1.0]]}], + "routes": [{"token": 0, "slot": 0, "expert": 3, "weight": 1.0}]} + + +def test_actual_plena_codec_known_bytes_rounding_subnormals_and_tail(): + elements, scales, shape = encode_matrix([ + [1, 1.064, 1.09375, -1, 0, 2**-10, 2**-7, 2**-6, 2], + ]) + # 1.064 rounds down with PLENA's quarter-truncated mantissa rounding. + # 2^-7 saturates the local subnormal mantissa, unlike standard E4M3. + assert elements == bytes([0x38, 0x38, 0x39, 0xB8, 0, 1, 7, 8, 0x38] + [0] * 7) + assert scales == bytes([127, 128]) + assert shape == {"rows": 1, "cols": 9, "element_row_stride": 16, "scale_row_stride": 2} + assert decode_element(0x07, 127) == 7 / 1024 + assert decode_element(0x87, 127) == -7 / 1024 + assert decode_element(0x38, 128) == 2 + assert decode_element(0x38, 0) == 0 + with pytest.raises(ValueError, match="non-finite"): + decode_element(0x78, 127) + + +def test_row_boundaries_never_share_scale_blocks(): + elements, scales, shape = encode_matrix([[1] * 9, [8] * 9]) + assert list(scales) == [127, 127, 130, 130] + assert elements[9:16] == bytes(7) + assert elements[25:32] == bytes(7) + region = {**shape, "element_base": 0, "scale_base": len(elements)} + assert decode_matrix(elements + scales, region) == [[1] * 9, [8] * 9] + + +def test_bf16_round_nearest_even_is_explicit(): + assert bf16_bits(1 + 1 / 256) == 0x3F80 + assert bf16_bits(1 + 3 / 256) == 0x3F82 + assert bf16_value(0xBF80) == -1.0 + for value in (math.nan, math.inf, 1e100): + with pytest.raises(ValueError): + bf16_bits(value) + + +def test_hand_computed_nonzero_swiglu_fixture(tmp_path): + result = export_workload(tmp_path, **_unit_arrays()) + # SiLU(1) is 0.731058..., rounded BF16 is 187/256; unit Up and Down. + assert result["golden"]["output_bf16"] == [[0x3F3B]] + assert result["golden"]["output_f32"] == [[187 / 256]] + assert result["golden"]["pre_round_output_f32"] == [[187 / 256]] + + +def test_grouping_retains_weights_slots_duplicates_and_zero_weights(): + routes = [ + {"token": 2, "slot": 7, "expert": 5, "weight": 0.0}, + {"token": 0, "slot": 1, "expert": 3, "weight": 0.75}, + {"token": 0, "slot": 0, "expert": 3, "weight": 0.25}, + ] + grouped = group_routes(routes, 3, {3, 5, 8}) + assert [g["expert"] for g in grouped] == [3, 5] + assert grouped[0]["routes"] == [routes[2], routes[1]] + assert grouped[1]["routes"] == [routes[0]] + assert sum(len(g["routes"]) for g in grouped) == len(routes) + + +@pytest.mark.parametrize("routes, message", [ + ([{"token": 0, "slot": 0, "expert": 3, "weight": 1}] * 2, "duplicate"), + ([{"token": 1, "slot": 0, "expert": 3, "weight": 1}], "out-of-range"), + ([{"token": 0, "slot": -1, "expert": 3, "weight": 1}], "out-of-range"), + ([{"token": 0, "slot": 0, "expert": 4, "weight": 1}], "unknown"), + ([{"token": 0, "slot": 0, "expert": 3, "weight": math.nan}], "finite"), + ([{"token": 0.0, "slot": 0, "expert": 3, "weight": 1}], "integer"), +]) +def test_invalid_routes_rejected(routes, message): + with pytest.raises(ValueError, match=message): + group_routes(routes, 1, {3}) + + +def test_stable_slot_combine_and_same_expert_multiple_routes(tmp_path): + arrays = _unit_arrays() + arrays["routes"] = [{"token": 0, "slot": slot, "expert": 3, "weight": weight} + for slot, weight in enumerate((0.25, 0.75))] + first = export_workload(tmp_path / "a", **arrays) + arrays["routes"].reverse() + second = export_workload(tmp_path / "b", **arrays) + assert first["golden"]["output_bf16"] == [[0x3F3B]] + assert first["golden"] == second["golden"] + assert len(first["golden"]["route_contributions"]) == 2 + + +def test_shared_only_and_unrouted_tokens(tmp_path): + arrays = _unit_arrays() + arrays["routes"] = [] + arrays["inputs"] = [[1.0], [0.0]] + no_routes = export_workload(tmp_path / "empty", **arrays) + assert no_routes["golden"]["output_bf16"] == [[0], [0]] + shared = export_workload(tmp_path / "shared", **arrays, + shared_expert={"expert": 3, "weight": 1.0}) + assert shared["golden"]["output_bf16"] == [[0x3F3B], [0]] + + +def test_demo_manifest_roundtrip_alignment_tails_scales_and_provenance(tmp_path): + arrays = demo_arrays() + result = export_workload(tmp_path, **arrays, + provenance={"scope": "synthetic_test", "source_sha256": "test-source"}) + manifest = json.loads((tmp_path / "workload.json").read_text()) + golden = json.loads((tmp_path / "golden.json").read_text()) + hbm = (tmp_path / "weights.bin").read_bytes() + assert manifest == result["workload"] + assert [len(g["routes"]) for g in manifest["grouped_routes"]] == [6, 2, 1] + assert len(hbm) % 64 == 0 + assert len(manifest["experts"]) == 4 # one deliberately unrouted expert + scale_values = set() + for expert in manifest["experts"]: + for stage in ("gate", "up", "down"): + region = expert[stage] + assert region["element_base"] % 64 == region["scale_base"] % 64 == 0 + assert region["cols"] % 8 != 0 + assert region["rows"] % 4 != 0 + start = region["scale_base"] + scale_values.update(hbm[start:start + region["rows"] * region["scale_row_stride"]]) + assert len(scale_values) > 1 + assert any(x != 0 for row in golden["output_bf16"] for x in row) + assert manifest["metadata"]["hbm_sha256"] == hashlib.sha256(hbm).hexdigest() + assert manifest["metadata"]["source_arrays_sha256"] == hashlib.sha256((tmp_path / "source_arrays.json").read_bytes()).hexdigest() + assert golden["workload_sha256"] == hashlib.sha256((tmp_path / "workload.json").read_bytes()).hexdigest() + assert manifest["metadata"]["provenance"]["scope"] == "synthetic_test" + assert numerical_reference(manifest, hbm)["output_bf16"] == golden["output_bf16"] + + +def test_shape_and_weight_validation_before_export(tmp_path): + arrays = _unit_arrays() + arrays["experts"][0]["up"] = [[1.0, 2.0]] + with pytest.raises(ValueError, match="gate/up"): + export_workload(tmp_path, **arrays) + assert not (tmp_path / "workload.json").exists() + with pytest.raises(ValueError, match="non-finite"): + encode_matrix([[math.inf]]) + + +def test_combine_uses_fp32_and_canonical_slot_order(tmp_path): + arrays = _unit_arrays() + # The middle contribution is lost at this FP32 magnitude. Slot order is + # observable; a double-precision sum or completion-order sum is incorrect. + arrays["routes"] = [{"token": 0, "slot": slot, "expert": 3, "weight": weight} + for slot, weight in enumerate((2**28, 1, -(2**28)))] + result = export_workload(tmp_path / "canonical", **arrays) + assert result["golden"]["output_bf16"] == [[0]] + arrays["routes"].reverse() + reversed_input = export_workload(tmp_path / "shuffled", **arrays) + assert reversed_input["golden"]["output_bf16"] == [[0]] + # Actually change the slot order: large terms cancel before the small one. + arrays["routes"][0]["slot"], arrays["routes"][1]["slot"] = 1, 2 + changed_slots = export_workload(tmp_path / "changed_slots", **arrays) + assert changed_slots["golden"]["output_bf16"] == [[0x3F3B]] diff --git a/doc/moe_normal_review.md b/doc/moe_normal_review.md new file mode 100644 index 0000000..b930afd --- /dev/null +++ b/doc/moe_normal_review.md @@ -0,0 +1,60 @@ +# MoE normal-buffer exporter: review guide + +This change supplies the input and independent numerical reference for the paired +Simulator MoE experiment. It exports fixed routes; it does not add a runtime +router or lower dual-core execution into the existing instruction stream. + +## What to review + +1. `aten/plena/moe_normal_export.py`: validate inputs/routes, group routes by + expert while retaining token/slot/weight, encode weights, and emit a versioned + workload plus an independently computed golden output. +2. `aten/plena/moe_full_shape_export.py`: generate deterministic nonzero arrays + at specified model dimensions and vectorize independent oracle outputs. Each + dot product still accumulates in ascending K with separate FP32 multiply/add. +3. The two corresponding tests in `aten/tests`: codec bytes, row padding, + malformed routes, duplicate expert selections, shared experts and numerical + agreement between the scalar and vectorized paths. + +## Data contract + +`X[T,D]` enters an expert with gate/up weights `[F,D]` and down weights `[D,F]`. +Those are physical output-major weight rows; logical computation is +`gate=X×Wgateᵀ`, `up=X×Wupᵀ`, `Z=SiLU(gate)×up`, `Y=Z×Wdownᵀ`. +The transposed mathematical notation does not require a runtime transpose buffer. +Each projection and SwiGLU output rounds to BF16; route-weighted combine uses +stable token/slot order with FP32 accumulation and final BF16 rounding. + +Weights use the actual local PLENA E4M3/E8M0 codec, with 8 elements per scale. +Elements and scales occupy separate address ranges, with explicit row strides +and tail padding. The oracle decodes the exported bytes before computing. +This local codec is not a claim of interchangeability with every MX format. + +Files: `workload.json` (shapes/routes/addresses), `weights.bin` (encoded bytes), +`golden.json` (expected output), and `source_arrays.json` for the small exporter. +Hashes bind the workload to the weight image and record its provenance. + +## Run + +Use Python 3.10+ with torch, NumPy, bitstring, PyYAML and pytest installed. Set `PLENA_TOOLS` to a +checkout of the Simulator's pinned PLENA_Tools submodule. From this repository: + +```sh +PYTHONPATH="$PWD:$PLENA_TOOLS" python -m pytest \ + aten/tests/test_moe_normal_export.py aten/tests/test_moe_full_shape_export.py -q +PYTHONPATH="$PLENA_TOOLS" python aten/plena/moe_normal_export.py \ + --output-dir /tmp/plena-moe-fixture --shared +``` + +The paired Simulator PR contains the single/dual configurations and a smoke +command that invokes this exporter, runs Rust twice per configuration and checks +exact BF16 output, resource limits, HBM counters and repeat determinism. + +## Scope and result + +The review checkout passes **21 exporter tests**. The previous performance +campaign used D=2048 with F=512 (Qwen) or F=1408 (DeepSeek), token windows 8/32, +archived routes and synthetic nonzero weights/activations. The best tested dual +configuration took 8.88% more simulated time than the optimized single baseline; +this exporter makes the comparison reproducible, not faster by itself. +See the paired Simulator review for core/SRAM dimensions and timing boundaries.