diff --git a/.github/workflows/transactional_emulator.yml b/.github/workflows/transactional_emulator.yml index e63788de2..ed63acefd 100644 --- a/.github/workflows/transactional_emulator.yml +++ b/.github/workflows/transactional_emulator.yml @@ -33,7 +33,13 @@ jobs: rust-lint: runs-on: ubuntu-latest steps: + # Submodules: clippy compiles `--all-targets`, and one test guard reads the + # compiler's `doc/operation.svh` through `include_str!` to check that every + # opcode the compiler emits decodes here. Without the submodule that is a + # compile error, not a skipped test. - uses: actions/checkout@v4 + with: + submodules: recursive # Pinned: with `-D warnings`, a rolling channel reds unrelated PRs. - name: Install Rust toolchain @@ -147,3 +153,24 @@ jobs: - name: Run emulator timing gates run: | nix develop --command bash -c "just test-timing-gates" + + # Mamba-2 / selective SSM. Synthetic like the tests above -- no checkpoint, + # no HF libs. Inputs are drawn exactly representable in MX-FP8 so the + # reported error is the kernel's rather than the input quantisation's. + # + # All four cases are bit-exact (max absolute error 0.000000), so any + # regression shows up as a hard failure rather than as tolerance drift. + - name: Run Mamba-2 stage tests + run: | + nix develop --command bash -c "just test-mamba2-all" + + # The chunked-SSD reference vs the plain recurrence, plus the cross-repo + # stage-marker contract. Pure torch. + - name: Run Mamba-2 reference cross-check + run: | + nix develop --command bash -c "just test-mamba2-reference" + + # Bandwidth term + the two KV-cost bugs it exposed in the analytic model. + - name: Run analytic performance-model regression tests + run: | + nix develop --command bash -c "just test-perf-model" diff --git a/PLENA_Compiler b/PLENA_Compiler index d89ad594c..27119a91a 160000 --- a/PLENA_Compiler +++ b/PLENA_Compiler @@ -1 +1 @@ -Subproject commit d89ad594c798daa54f63f914aebad5317653489a +Subproject commit 27119a91ad5b9ae77c9a1e3eb082ece434c93ba2 diff --git a/analytic_models/__init__.py b/analytic_models/__init__.py index e17a5ab2e..8676ffee0 100644 --- a/analytic_models/__init__.py +++ b/analytic_models/__init__.py @@ -1,3 +1,10 @@ -from . import memory, performance, utilisation +"""PLENA analytic models. -__all__ = ["memory", "performance", "utilisation"] +Only the `performance` package exists in this checkout; the previous +`from . import memory, performance, utilisation` raised ImportError on any +`import analytic_models` because `memory` and `utilisation` were never added. +""" + +from . import performance + +__all__ = ["performance"] diff --git a/analytic_models/performance/customISA_lib.json b/analytic_models/performance/customISA_lib.json index eb3f7800c..fb900b650 100644 --- a/analytic_models/performance/customISA_lib.json +++ b/analytic_models/performance/customISA_lib.json @@ -93,6 +93,10 @@ "alone" : "6 + 3 + VECTOR_MUL_CYCLES + 1", "pipelined" : "VECTOR_MUL_CYCLES" }, + "V_FMA_VF" : { + "alone" : "6 + 3 + VECTOR_MUL_CYCLES + 1", + "pipelined" : "VECTOR_MUL_CYCLES" + }, "V_EXP_V" : { "alone" : "6 + 3 + VECTOR_EXP_CYCLES + 1", @@ -119,6 +123,31 @@ "pipelined" : " 1 + VECTOR_MAX_CYCLES" }, + "V_MAX_VF" : { + "alone" : "6 + 3 + VECTOR_MAX_CYCLES + 1", + "pipelined" : "1 + VECTOR_MAX_CYCLES" + }, + + "V_MIN_VF" : { + "alone" : "6 + 3 + VECTOR_MAX_CYCLES + 1", + "pipelined" : "1 + VECTOR_MAX_CYCLES" + }, + + "V_PS_V" : { + "alone" : "6 + 3 + VECTOR_PREFIX_SCAN_CYCLES + 1", + "pipelined" : "VECTOR_PREFIX_SCAN_CYCLES" + }, + + "V_SHFT_V" : { + "alone" : "6 + 3 + VECTOR_SHIFT_CYCLES + 1", + "pipelined" : "VECTOR_SHIFT_CYCLES" + }, + + "V_SOFTPLUS_V" : { + "alone" : "6 + 3 + VECTOR_SOFTPLUS_CYCLES + 1", + "pipelined" : "1 + VECTOR_SOFTPLUS_CYCLES" + }, + "S_BASIC" : { "alone" : " SCALAR_FP_BASIC_CYCLES + 1", "pipelined" : "SCALAR_FP_BASIC_CYCLES" @@ -174,6 +203,11 @@ "pipelined" : "1" }, + "S_MAP_FP_V" : { + "alone" : "6 + 1", + "pipelined" : "1" + }, + "S_ADD_INT" : { "alone" : "2 + SCALAR_INT_BASIC_CYCLES + 1", "pipelined" : "SCALAR_INT_BASIC_CYCLES" diff --git a/analytic_models/performance/hybrid_model.py b/analytic_models/performance/hybrid_model.py new file mode 100644 index 000000000..82f8be844 --- /dev/null +++ b/analytic_models/performance/hybrid_model.py @@ -0,0 +1,753 @@ +"""Hybrid-model performance driver: one cost per layer, by that layer's type. + +`llama_model.py`, `gpt_oss_model.py` and `mamba2_model.py` all assume one block +shape repeated `num_hidden_layers` times -- the assumption is literally the line +`overall_exe_cycle += block_cycles * self.num_hidden_layers`. A hybrid model has +no such block. Nemotron-3 interleaves Mamba-2, attention and MLP/MoE blocks; +Kimi Linear interleaves KDA and full attention three-to-one inside every layer. + +`PerfModel` was already operator-level and model-agnostic, so what was missing +was only this: a driver that reads a per-layer type list and dispatches. + +Why it matters beyond running two more models +--------------------------------------------- +The claim a hybrid accelerator argument rests on is that some operator +contributes a small share of a model's FLOPs and a large share of its decode +latency. Neither `llama_model.py` nor `mamba2_model.py` can express that, +because in a homogeneous stack both shares are 100% by construction. This driver +reports **FLOPs, cycles and HBM bytes per operator**, which is the form that +claim has to be read off. + +FLOPs are counted here rather than taken from `PerfModel`, deliberately. They are +a property of the model, the cost model is a property of the machine, and the +whole point is to compare the two -- deriving both from the same code would make +the comparison circular. + +Layer-type encodings, both read from the config and both verified against the +model's own source +------------------------------------------------------------------------------ +**NemotronH** (`nemotron_h`) carries `hybrid_override_pattern`, one character per +layer, each layer being a *single* block -- a mixer or a feed-forward, not both. +From `configuration_nemotron_h.py`:: + + pattern_mapping = {"M": "mamba", "E": "moe", "*": "attention"} # and "-": "mlp" + +**Kimi Linear** (`kimi_linear`) carries `linear_attn_config` with `kda_layers` +and `full_attn_layers`. Those indices are **1-based against a 0-based +`layer_idx`** -- `is_kda_layer` is `(layer_idx + 1) in kda_layers` -- and every +layer additionally has a feed-forward, dense for the first +`first_k_dense_replace` layers and MoE after:: + + num_experts is not None and layer_idx >= first_k_dense_replace + and layer_idx % moe_layer_freq == 0 + +Both are checked for completeness when parsed: an index that appears twice, or a +layer that appears in neither list, raises rather than being silently dropped. + +Two things this driver does not fix +----------------------------------- +**The `MLEN` disagreement.** `plena_settings.toml` declares +`[ANALYTIC.CONFIG.MLEN] = 2048` and `[TRANSACTIONAL.CONFIG.MLEN] = 64`, with +`BLEN` 128 against 4. Those describe different machines. `--sweep` walks between +them rather than picking one, but the absolute cycle counts still belong to +whichever point is being read. + +**And width is free in this model.** `MLEN`/`VLEN` enter cycles only through +instruction counts -- `ceil(work / MLEN)` -- while the latencies are constants: +`V_FMA_VF` bills `VECTOR_MUL_CYCLES` whatever `VLEN` is. So a wider machine here +covers 32x the data per instruction at the same cost, with no area term anywhere. +That makes absolute speedups against width near-tautological and they should not +be quoted as findings. What survives is the *comparison between operators*, which +is priced by the same idealisation on both sides: when one saturates on width and +another does not, that difference is a property of the operators. Every number +this driver is used for should be of that shape -- a ratio, or a turning point, +never an absolute speedup. + +**The clock.** Cycles are cycles. `runtime_config.rs` hard-codes `PERIOD` at 1 ns +with no stated basis, so this driver reports no times at all rather than +multiplying by a number it cannot defend. + +Usage:: + + python hybrid_model.py --model kimi-linear-48b-a3b --model-lib ./doc/Model_Lib \\ + --config ./plena_settings.toml --isa-lib ./customISA_lib.json + python hybrid_model.py --list-models --model-lib ./doc/Model_Lib +""" + +from __future__ import annotations + +import argparse +import json +from itertools import pairwise +from dataclasses import dataclass, field +from pathlib import Path + +try: # script-style invocation, matching the sibling drivers + from perf_model import PerfModel, load_hardware_config_from_toml +except ImportError: # package-style import + from analytic_models.performance.perf_model import PerfModel, load_hardware_config_from_toml + +#: `hybrid_override_pattern` characters, from NemotronH's own `pattern_mapping` +#: plus the `-` the docstring names ("M: Mamba2, *: Attention, -: MLP"). +_NEMOTRON_PATTERN = {"M": "mamba", "*": "attention", "-": "mlp", "E": "moe"} + +#: Operators a layer's mixer slot can hold. `None` means the layer has no mixer, +#: which is the normal case for NemotronH's feed-forward-only layers. +MIXERS = ("mamba", "attention", "kda") +FFNS = ("mlp", "moe") + + +@dataclass(frozen=True) +class LayerSpec: + """One layer: at most one mixer and at most one feed-forward. + + Both slots are optional because the two families disagree about what a layer + is. A NemotronH layer is exactly one of the two; a Kimi Linear layer is one + of each. Modelling that with two optional slots rather than a single "type" + string is what keeps the per-operator totals honest -- a Kimi layer really + does run a mixer *and* an FFN, and folding it into one label would lose one + of them. + """ + + index: int + mixer: str | None + ffn: str | None + + def __post_init__(self) -> None: + if self.mixer is not None and self.mixer not in MIXERS: + raise ValueError(f"layer {self.index}: unknown mixer {self.mixer!r}") + if self.ffn is not None and self.ffn not in FFNS: + raise ValueError(f"layer {self.index}: unknown ffn {self.ffn!r}") + if self.mixer is None and self.ffn is None: + raise ValueError(f"layer {self.index} has neither a mixer nor a feed-forward") + + +def layer_plan(config: dict) -> list[LayerSpec]: + """One :class:`LayerSpec` per layer, from whichever encoding the config uses. + + Raises on any config it cannot read rather than falling back to a + homogeneous stack: a hybrid model silently modelled as uniform is exactly + the failure this driver exists to remove. + """ + model_type = config.get("model_type") + num_layers = config["num_hidden_layers"] + + if model_type == "nemotron_h": + pattern = config["hybrid_override_pattern"] + if len(pattern) != num_layers: + raise ValueError( + f"hybrid_override_pattern is {len(pattern)} characters but " + f"num_hidden_layers is {num_layers}; one of them is stale" + ) + unknown = sorted(set(pattern) - set(_NEMOTRON_PATTERN)) + if unknown: + raise ValueError(f"unknown hybrid_override_pattern characters: {unknown}") + plan = [] + for i, ch in enumerate(pattern): + kind = _NEMOTRON_PATTERN[ch] + if kind in FFNS: + plan.append(LayerSpec(i, None, kind)) + else: + plan.append(LayerSpec(i, kind, None)) + return plan + + if model_type == "kimi_linear": + linear = config.get("linear_attn_config") + if linear is None: + raise ValueError("kimi_linear config without linear_attn_config") + kda = set(linear["kda_layers"]) + full = set(linear["full_attn_layers"]) + both = sorted(kda & full) + if both: + raise ValueError(f"layers listed as both KDA and full attention: {both}") + # 1-based, per `is_kda_layer`: `(layer_idx + 1) in kda_layers`. + missing = sorted(set(range(1, num_layers + 1)) - kda - full) + if missing: + raise ValueError( + f"layers in neither kda_layers nor full_attn_layers: {missing}. " + f"Defaulting them to either would silently mis-attribute the mixer" + ) + first_dense = config.get("first_k_dense_replace", 0) + freq = config.get("moe_layer_freq", 1) or 1 + has_experts = config.get("num_experts") is not None + plan = [] + for i in range(num_layers): + mixer = "kda" if (i + 1) in kda else "attention" + is_moe = has_experts and i >= first_dense and i % freq == 0 + plan.append(LayerSpec(i, mixer, "moe" if is_moe else "mlp")) + return plan + + raise ValueError( + f"model_type {model_type!r} carries no layer-type encoding this driver " + f"knows. Add it here rather than assuming a uniform stack" + ) + + +@dataclass +class OperatorCost: + """Accumulated cost of one operator across every layer that runs it.""" + + cycles: int = 0 + bytes_: float = 0.0 + flops: float = 0.0 + layers: int = 0 + + +@dataclass +class HybridModel: + """Per-operator cost of one hybrid model on the analytic PLENA model.""" + + config: dict + perf: PerfModel + seq_len: int = 4096 + batch_size: int = 1 + kda_chunk: int = 16 + #: Price a prefetch that takes a row count instead of a whole-block count. + #: No such instruction exists; this exists to put a number on not having it. + row_granular_prefetch: bool = False + plan: list[LayerSpec] = field(init=False) + + def __post_init__(self) -> None: + self.plan = layer_plan(self.config) + c = self.config + self.hidden_size = c["hidden_size"] + self.num_layers = c["num_hidden_layers"] + self.vocab_size = c["vocab_size"] + + # -- shapes, per family -------------------------------------------------- + + @property + def _attn(self) -> dict: + c = self.config + heads = c["num_attention_heads"] + kv_heads = c.get("num_key_value_heads", heads) + head_dim = c.get("head_dim") or c.get("attention_head_dim") or self.hidden_size // heads + return {"heads": heads, "kv_heads": kv_heads, "head_dim": head_dim} + + @property + def _kda(self) -> dict: + linear = self.config["linear_attn_config"] + return { + "heads": linear["num_heads"], + "head_dim": linear["head_dim"], + "conv_kernel": linear["short_conv_kernel_size"], + } + + @property + def _mamba(self) -> dict: + c = self.config + heads = c["mamba_num_heads"] + head_dim = c.get("mamba_head_dim") or (c["hidden_size"] * c.get("expand", 2)) // heads + return { + "heads": heads, + "head_dim": head_dim, + "state_size": c.get("ssm_state_size") or c["state_size"], + "n_groups": c.get("n_groups", 1), + "conv_kernel": c.get("conv_kernel", 4), + "chunk_size": c.get("chunk_size", 256), + } + + @property + def _ffn(self) -> dict: + c = self.config + # NemotronH's MLP is `down(act(up(x)))` -- two matmuls, no gate. Kimi's is + # `down(act(gate(x)) * up(x))` -- three. Read from the source of each, + # not assumed: a gated count applied to NemotronH overstates every + # feed-forward layer in the model by 50%. + gated = c.get("model_type") != "nemotron_h" + return { + "intermediate": c["intermediate_size"], + "moe_intermediate": c.get("moe_intermediate_size", c["intermediate_size"]), + "num_experts": c.get("num_experts") or c.get("n_routed_experts") or 0, + "top_k": c.get("num_experts_per_token") or c.get("num_experts_per_tok") or 0, + "shared": c.get("num_shared_experts", 0) or 0, + "matmuls": 3 if gated else 2, + } + + # -- FLOPs, counted independently of the cost model ---------------------- + + @staticmethod + def _gemm_flops(m: int, n: int, k: int) -> float: + """Multiply-accumulate counted as two flops, the usual convention.""" + return 2.0 * m * n * k + + def _mixer_flops(self, kind: str, tokens: int, kv_len: int) -> float: + h = self.hidden_size + if kind == "attention": + a = self._attn + q_out = a["heads"] * a["head_dim"] + kv_out = a["kv_heads"] * a["head_dim"] + flops = self._gemm_flops(tokens, q_out + 2 * kv_out, h) # qkv + flops += self._gemm_flops(tokens, h, q_out) # o proj + # scores and the value read, both over the whole visible context + flops += 2 * self._gemm_flops(tokens, kv_len, q_out) + return flops + if kind == "kda": + k = self._kda + width = k["heads"] * k["head_dim"] + flops = self._gemm_flops(tokens, 3 * width, h) # q, k, v + flops += self._gemm_flops(tokens, h, width) # out proj + flops += 2.0 * tokens * width * k["conv_kernel"] # causal conv + # The recurrence: predict, error and update each touch the whole + # [key, value] state once per token. + flops += 3 * 2.0 * tokens * k["heads"] * k["head_dim"] * k["head_dim"] + return flops + if kind == "mamba": + m = self._mamba + inner = m["heads"] * m["head_dim"] + proj_out = 2 * inner + 2 * m["n_groups"] * m["state_size"] + m["heads"] + flops = self._gemm_flops(tokens, proj_out, h) # in_proj + flops += self._gemm_flops(tokens, h, inner) # out_proj + conv_dim = inner + 2 * m["n_groups"] * m["state_size"] + flops += 2.0 * tokens * conv_dim * m["conv_kernel"] + flops += 3 * 2.0 * tokens * m["heads"] * m["head_dim"] * m["state_size"] + return flops + raise ValueError(f"no FLOP model for mixer {kind!r}") + + def _ffn_flops(self, kind: str, tokens: int) -> float: + f = self._ffn + h = self.hidden_size + if kind == "mlp": + return f["matmuls"] * self._gemm_flops(tokens, f["intermediate"], h) + active = f["top_k"] + f["shared"] + flops = active * f["matmuls"] * self._gemm_flops(tokens, f["moe_intermediate"], h) + flops += self._gemm_flops(tokens, f["num_experts"], h) # router + return flops + + # -- dispatch ------------------------------------------------------------ + + def _stage(self, fn, *args, **kwargs) -> tuple[int, float]: + """Run one PerfModel stage and return its (cycles, HBM bytes). + + `PerfModel` accumulates traffic globally, so the counters are reset + around each call. That is the only way to get a per-operator byte count + out of it without changing its interface. + """ + before = self.perf.traffic_bytes + cycles = fn(*args, **kwargs) + return cycles, self.perf.traffic_bytes - before + + def _mixer_cost(self, kind: str, mode: str, tokens: int, kv_len: int) -> tuple[int, float]: + p, b = self.perf, self.batch_size + if kind == "attention": + a = self._attn + c1, b1 = self._stage( + p.projection, self.hidden_size, a["heads"], a["kv_heads"], a["head_dim"], tokens, b, mode + ) + c2, b2 = self._stage(p.flash_attention, a["heads"], a["kv_heads"], a["head_dim"], tokens, kv_len, b, mode) + return c1 + c2, b1 + b2 + if kind == "kda": + k = self._kda + width = k["heads"] * k["head_dim"] + cycles, byts = self._stage( + p.projection, self.hidden_size, 3 * k["heads"], k["heads"], k["head_dim"], tokens, b, mode + ) + c, bb = self._stage(p.causal_conv1d, width, k["conv_kernel"], tokens, b, mode) + cycles, byts = cycles + c, byts + bb + if mode == "prefill": + c, bb = self._stage( + p.kda_chunk_prefill, + k["heads"], + k["head_dim"], + k["head_dim"], + self.kda_chunk, + tokens, + b, + row_granular_prefetch=self.row_granular_prefetch, + ) + else: + c, bb = self._stage(p.kda_recurrence_decode, k["heads"], k["head_dim"], k["head_dim"], b) + cycles, byts = cycles + c, byts + bb + c, bb = self._stage(p.gated_rms_norm, width, tokens, b, mode) + return cycles + c, byts + bb + if kind == "mamba": + m = self._mamba + inner = m["heads"] * m["head_dim"] + conv_dim = inner + 2 * m["n_groups"] * m["state_size"] + cycles, byts = self._stage( + p.linear, + self.hidden_size, + 2 * inner + 2 * m["n_groups"] * m["state_size"] + m["heads"], + tokens, + b, + mode, + ) + c, bb = self._stage(p.causal_conv1d, conv_dim, m["conv_kernel"], tokens, b, mode) + cycles, byts = cycles + c, byts + bb + c, bb = self._stage(p.dt_activation, m["heads"], tokens, b, mode) + cycles, byts = cycles + c, byts + bb + if mode == "prefill": + c, bb = self._stage( + p.ssd_chunk_scan, + m["heads"], + m["head_dim"], + m["state_size"], + m["n_groups"], + m["chunk_size"], + tokens, + b, + ) + else: + c, bb = self._stage( + p.ssd_recurrence_decode, m["heads"], m["head_dim"], m["state_size"], m["n_groups"], b + ) + cycles, byts = cycles + c, byts + bb + c, bb = self._stage(p.gated_rms_norm, inner, tokens, b, mode) + cycles, byts = cycles + c, byts + bb + c, bb = self._stage(p.linear, inner, self.hidden_size, tokens, b, mode) + return cycles + c, byts + bb + raise ValueError(f"no cost model for mixer {kind!r}") + + def _ffn_cost(self, kind: str, mode: str, tokens: int) -> tuple[int, float]: + f, p, b = self._ffn, self.perf, self.batch_size + if kind == "mlp": + return self._stage(p.feed_forward, self.hidden_size, f["intermediate"], tokens, b, mode) + return self._stage( + p.mlp_moe, + self.hidden_size, + tokens, + b, + f["num_experts"], + f["top_k"] + f["shared"], + f["moe_intermediate"], + mode, + ) + + def per_operator(self, mode: str = "decode") -> dict[str, OperatorCost]: + """FLOPs, cycles and HBM bytes for every operator, summed over layers.""" + if mode not in ("prefill", "decode"): + raise ValueError(f"mode must be prefill or decode, got {mode!r}") + tokens = self.seq_len if mode == "prefill" else 1 + # In decode the attention read spans the whole context; the recurrent + # mixers do not, which is the asymmetry the whole comparison is about. + kv_len = self.seq_len + + out: dict[str, OperatorCost] = {} + + def add(name: str, cycles: int, byts: float, flops: float) -> None: + e = out.setdefault(name, OperatorCost()) + e.cycles += cycles + e.bytes_ += byts + e.flops += flops + e.layers += 1 + + for layer in self.plan: + if layer.mixer is not None: + c, bb = self._mixer_cost(layer.mixer, mode, tokens, kv_len) + add(layer.mixer, c, bb, self._mixer_flops(layer.mixer, tokens * self.batch_size, kv_len)) + if layer.ffn is not None: + c, bb = self._ffn_cost(layer.ffn, mode, tokens) + add(layer.ffn, c, bb, self._ffn_flops(layer.ffn, tokens * self.batch_size)) + # Two norms and two residuals per layer whichever slots it fills. + c, bb = self._stage(self.perf.rms_layer, self.hidden_size, tokens, self.batch_size, mode) + c2, bb2 = self._stage(self.perf.residual, self.hidden_size, tokens, self.batch_size, mode) + add("norm+residual", c + c2, bb + bb2, 0.0) + + c, bb = self._stage(self.perf.lm_head, self.hidden_size, self.vocab_size, self.batch_size) + add("lm_head", c, bb, self._gemm_flops(self.batch_size, self.vocab_size, self.hidden_size)) + return out + + # -- reporting ----------------------------------------------------------- + + def layer_census(self) -> dict[str, int]: + census: dict[str, int] = {} + for layer in self.plan: + for slot in (layer.mixer, layer.ffn): + if slot is not None: + census[slot] = census.get(slot, 0) + 1 + return census + + def report(self, mode: str = "decode") -> str: + costs = self.per_operator(mode) + total_c = sum(e.cycles for e in costs.values()) or 1 + total_f = sum(e.flops for e in costs.values()) or 1.0 + total_b = sum(e.bytes_ for e in costs.values()) or 1.0 + + lines = [ + f"{'operator':<16}{'layers':>7}{'GFLOP':>12}{'%FLOP':>8}" + f"{'Mcycles':>12}{'%cycles':>9}{'MB':>12}{'%bytes':>8}" + ] + lines.append("-" * 85) + for name, e in sorted(costs.items(), key=lambda kv: -kv[1].cycles): + lines.append( + f"{name:<16}{e.layers:>7}{e.flops / 1e9:>12.3f}{100 * e.flops / total_f:>7.1f}%" + f"{e.cycles / 1e6:>12.3f}{100 * e.cycles / total_c:>8.1f}%" + f"{e.bytes_ / 1e6:>12.3f}{100 * e.bytes_ / total_b:>7.1f}%" + ) + lines.append("-" * 85) + lines.append( + f"{'total':<16}{'':>7}{total_f / 1e9:>12.3f}{100.0:>7.1f}%" + f"{total_c / 1e6:>12.3f}{100.0:>8.1f}%{total_b / 1e6:>12.3f}{100.0:>7.1f}%" + ) + return "\n".join(lines) + + +#: Coherent machine points between the two shapes `plena_settings.toml` declares. +#: +#: `MLEN`/`BLEN`/`HLEN`/`VLEN` are not independent -- `MLEN % BLEN == 0`, +#: `VLEN >= BLEN`, and `flash_attention` divides by `MLEN // HLEN`, which is a +#: crash rather than a wrong answer when it rounds to zero. Overriding one field +#: and leaving the rest is what makes that happen, so the sweep moves them +#: together and the endpoints are the two declared configurations: +#: `TRANSACTIONAL` at 64/4/16/64 and `ANALYTIC` at 2048/128/128/2048. +MACHINE_POINTS: tuple[tuple[int, int, int, int], ...] = ( + (64, 4, 16, 64), + (256, 16, 32, 256), + (512, 32, 64, 512), + (1024, 64, 128, 1024), + (2048, 128, 128, 2048), +) + +#: HBM bandwidth in bytes per cycle. The shipped `HBM_WIDTH` of 512 is one point +#: on this axis, not a measurement -- `perf_model.py` derives bandwidth from it +#: as "one HBM row per cycle", which is an assumption about the interface rather +#: than a number anyone has read off silicon. +BANDWIDTH_POINTS: tuple[int, ...] = (64, 128, 256, 512, 1024, 2048, 4096) + + +def sweep( + config: dict, + isa_lib: Path, + settings: Path, + *, + seq_len: int, + batch_size: int, + kda_chunk: int, + mode: str, + operator: str, + row_granular_prefetch: bool = False, +) -> list[dict]: + """`operator`'s FLOP share against its cycle share, over machine and bandwidth. + + The ratio of the two is the quantity a hybrid accelerator argument turns on: + a value near 1 says the operator costs what its arithmetic says it should, + and a value far from 1 says the machine is spending time somewhere the FLOP + count cannot see. + + Both axes are swept because neither is known. `MLEN` is declared twice in + `plena_settings.toml` with a 32x disagreement, and bandwidth is derived from + a row width rather than measured. Reporting a single ratio at one arbitrary + point on that plane would be reporting the plane's shape as if it were a + property of the design. + """ + rows = [] + for mlen, blen, hlen, vlen in MACHINE_POINTS: + for bw in BANDWIDTH_POINTS: + hw = load_hardware_config_from_toml(str(settings)) + hw.MLEN, hw.BLEN, hw.HLEN, hw.VLEN = mlen, blen, hlen, vlen + hw.HBM_BANDWIDTH_BYTES_PER_CYCLE = float(bw) + model = HybridModel( + config=config, + perf=PerfModel(hw, str(isa_lib)), + seq_len=seq_len, + batch_size=batch_size, + kda_chunk=kda_chunk, + row_granular_prefetch=row_granular_prefetch, + ) + costs = model.per_operator(mode) + if operator not in costs: + raise SystemExit(f"{operator!r} is not an operator of this model; it runs {sorted(costs)}") + total_f = sum(e.flops for e in costs.values()) or 1.0 + total_c = sum(e.cycles for e in costs.values()) or 1 + e = costs[operator] + traffic = model.perf.traffic_summary() + mem_bound = traffic["memory_only_cycles"] / max(1, traffic["compute_only_cycles"]) + rows.append( + { + "mlen": mlen, + "bandwidth": bw, + "flop_share": e.flops / total_f, + "cycle_share": e.cycles / total_c, + "ratio": (e.cycles / total_c) / max(1e-12, e.flops / total_f), + "memory_over_compute": mem_bound, + "total_cycles": total_c, + "operator_cycles": e.cycles, + } + ) + return rows + + +def format_sweep(rows: list[dict], operator: str) -> str: + """The ratio as a grid, plus how memory-bound the whole model is at each point. + + `mem/compute` is the ratio of the two sides of the roofline summed over every + stage. Above 1 the model is memory-bound overall, and on-chip width stops + changing anything -- which is the reading that decides whether a wider + machine is worth anything at all at that bandwidth. + """ + bws = sorted({r["bandwidth"] for r in rows}) + mlens = sorted({r["mlen"] for r in rows}) + grid = {(r["mlen"], r["bandwidth"]): r for r in rows} + + out = [f"cycle share / FLOP share for `{operator}` (1.00 = costs what its arithmetic says)", ""] + out.append(" MLEN " + "".join(f"{b:>9}" for b in bws) + " B/cycle") + out.append(" " + "-" * (5 + 9 * len(bws))) + for m in mlens: + out.append(f"{m:>6} " + "".join(f"{grid[(m, b)]['ratio']:>9.2f}" for b in bws)) + out.append("") + out.append("whole-model memory cycles / compute cycles (>1 = memory-bound overall)") + out.append("") + out.append(" MLEN " + "".join(f"{b:>9}" for b in bws) + " B/cycle") + out.append(" " + "-" * (5 + 9 * len(bws))) + for m in mlens: + out.append(f"{m:>6} " + "".join(f"{grid[(m, b)]['memory_over_compute']:>9.2f}" for b in bws)) + return "\n".join(out) + + +def format_width_return(rows: list[dict], operator: str) -> str: + """What another doubling of machine width is worth, and where it stops being worth anything. + + The ratio grid says an operator's *share* of the time gets worse on a wider + machine. It does not say whether the machine is faster, which is the question + a width decision actually turns on -- a wider machine helps the dense + operators even while it wastes lanes on the recurrent one, and the two have to + be weighed against each other rather than read separately. + + So this reports total decode cycles against width, and the marginal return of + each doubling. There is no area model here, so "wider is better" is trivially + true wherever it still helps at all; what is not trivial, and what this + reports, is the width past which it stops helping. Beyond that point the extra + lanes are spent on an operand that cannot fill them. + """ + bws = sorted({r["bandwidth"] for r in rows}) + mlens = sorted({r["mlen"] for r in rows}) + grid = {(r["mlen"], r["bandwidth"]): r for r in rows} + + out = ["total decode cycles, in millions", ""] + out.append(" MLEN " + "".join(f"{b:>9}" for b in bws) + " B/cycle") + out.append(" " + "-" * (5 + 9 * len(bws))) + for m in mlens: + out.append(f"{m:>6} " + "".join(f"{grid[(m, b)]['total_cycles'] / 1e6:>9.2f}" for b in bws)) + + out += ["", "marginal return of doubling MLEN: % of total cycles removed", ""] + out.append(" MLEN " + "".join(f"{b:>9}" for b in bws) + " B/cycle") + out.append(" " + "-" * (5 + 9 * len(bws))) + for prev, m in pairwise(mlens): + cells = [] + for b in bws: + before = grid[(prev, b)]["total_cycles"] + after = grid[(m, b)]["total_cycles"] + cells.append(f"{100 * (before - after) / max(1, before):>8.1f}%") + out.append(f"{prev:>4}->{m:<1}" + "".join(cells)) + + out += ["", f"share of that width spent on `{operator}`, which cannot use it", ""] + out.append(" MLEN " + "".join(f"{b:>9}" for b in bws) + " B/cycle") + out.append(" " + "-" * (5 + 9 * len(bws))) + for m in mlens: + out.append(f"{m:>6} " + "".join(f"{100 * grid[(m, b)]['cycle_share']:>8.1f}%" for b in bws)) + return "\n".join(out) + + +def load_model_config(name: str, model_lib: Path) -> dict: + path = model_lib / f"{name}.json" + if not path.exists(): + raise SystemExit(f"no config {path}; --list-models shows what is available") + return json.loads(path.read_text()) + + +def list_hybrid_models(model_lib: Path) -> list[tuple[str, str, dict[str, int]]]: + """Every config in the library this driver can read, with its layer census.""" + found = [] + for path in sorted(model_lib.glob("*.json")): + try: + config = json.loads(path.read_text()) + plan = layer_plan(config) + except Exception: + continue + census: dict[str, int] = {} + for layer in plan: + for slot in (layer.mixer, layer.ffn): + if slot is not None: + census[slot] = census.get(slot, 0) + 1 + found.append((path.stem, config.get("model_type", "?"), census)) + return found + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model", help="config stem in --model-lib, e.g. kimi-linear-48b-a3b") + ap.add_argument("--model-lib", type=Path, default=Path("doc/Model_Lib")) + ap.add_argument("--config", type=Path, default=Path("plena_settings.toml")) + ap.add_argument("--isa-lib", type=Path, default=Path("customISA_lib.json")) + ap.add_argument("--seq-len", type=int, default=4096) + ap.add_argument("--batch-size", type=int, default=1) + ap.add_argument("--kda-chunk", type=int, default=16, help="bf16 range bounds this at 17") + ap.add_argument("--mode", choices=("prefill", "decode", "both"), default="both") + ap.add_argument("--list-models", action="store_true") + ap.add_argument( + "--sweep", + metavar="OPERATOR", + help="sweep MLEN and HBM bandwidth, reporting OPERATOR's cycle share " + "against its FLOP share. Neither axis is known: MLEN is declared twice " + "with a 32x disagreement and bandwidth is derived from a row width, so " + "a single point would report the plane rather than the design.", + ) + ap.add_argument( + "--row-granular-prefetch", + action="store_true", + help="price an instruction set whose prefetch takes a row count rather " + "than a whole-block count. No such instruction exists; this exists to " + "put a number on not having it.", + ) + args = ap.parse_args() + + if args.list_models: + rows = list_hybrid_models(args.model_lib) + if not rows: + raise SystemExit(f"no hybrid configs found under {args.model_lib}") + for stem, model_type, census in rows: + census_str = ", ".join(f"{k}x{v}" for k, v in sorted(census.items())) + print(f" {stem:<34}{model_type:<16}{census_str}") + return + + if not args.model: + raise SystemExit("--model is required (or --list-models)") + + config = load_model_config(args.model, args.model_lib) + hw = load_hardware_config_from_toml(str(args.config)) + perf = PerfModel(hw, str(args.isa_lib)) + model = HybridModel( + config=config, + perf=perf, + seq_len=args.seq_len, + batch_size=args.batch_size, + kda_chunk=args.kda_chunk, + row_granular_prefetch=args.row_granular_prefetch, + ) + + census = ", ".join(f"{k} x{v}" for k, v in sorted(model.layer_census().items())) + print(f"\n{args.model} ({config.get('model_type')})") + print(f" hidden {model.hidden_size}, {model.num_layers} layers: {census}") + print(f" seq_len {args.seq_len}, batch {args.batch_size}, MLEN {perf.mlen}, VLEN {perf.vlen}") + print( + " NOTE cycles are uncalibrated and come from the ANALYTIC config, whose " + f"MLEN ({perf.mlen}) disagrees with the TRANSACTIONAL one. Read the ratios." + ) + if args.sweep: + mode = "decode" if args.mode == "both" else args.mode + print(f"\n--- sweep, {mode} ---\n") + rows = sweep( + config, + args.isa_lib, + args.config, + seq_len=args.seq_len, + batch_size=args.batch_size, + kda_chunk=args.kda_chunk, + mode=mode, + operator=args.sweep, + row_granular_prefetch=args.row_granular_prefetch, + ) + print(format_sweep(rows, args.sweep)) + print() + print(format_width_return(rows, args.sweep)) + print() + return + + for mode in ("prefill", "decode") if args.mode == "both" else (args.mode,): + print(f"\n--- {mode} ---") + print(model.report(mode)) + print() + + +if __name__ == "__main__": + main() diff --git a/analytic_models/performance/mamba2_model.py b/analytic_models/performance/mamba2_model.py new file mode 100644 index 000000000..581c1292e --- /dev/null +++ b/analytic_models/performance/mamba2_model.py @@ -0,0 +1,413 @@ +""" +Mamba-2 Performance Model for PLENA Simulator. + +Uses PerfModel from perf_model.py for per-layer cycle counting and computes +overall inference performance (prefill time, decode time, TPS, TTFT). + +Structure of one Mamba-2 mixer block (following HF `Mamba2Mixer`): + + x_norm = RMSNorm(x) + z, xBC, dt = in_proj(x_norm) hidden -> 2*d_inner + 2*n_groups*N + H + xBC = SiLU(causal_conv1d(xBC)) over conv_dim = d_inner + 2*n_groups*N + dt = clamp(softplus(dt + dt_bias), dt_min, dt_max) + y = SSD(x, dt, A, B, C) chunked scan in prefill, + single-step recurrence in decode + y = RMSNorm(y * SiLU(z)) gated norm + out = out_proj(y) d_inner -> hidden + x = x + out residual + +There is no attention and no KV cache: the entire history is compressed into a +fixed `num_heads * head_dim * state_size` recurrent state (plus a +`conv_dim * (conv_kernel-1)` conv window), so decode cost is O(1) in context +length. Whether that translates into a measured win depends on the memory +bandwidth term in perf_model.py — with no bandwidth term the comparison is +meaningless, because the growing KV cache costs nothing to read. + +Usage: + python mamba2_model.py --model mamba2-2.7b --model-lib ./Model_Lib \ + --config ./plena_settings.toml --isa-lib ./customISA_lib.json + python mamba2_model.py --list-models --model-lib ./Model_Lib +""" + +import argparse +import json +from pathlib import Path + +try: # script-style invocation (matches llama_model.py / gpt_oss_model.py) + from perf_model import PerfModel, load_hardware_config_from_toml +except ImportError: # package-style import + from analytic_models.performance.perf_model import PerfModel, load_hardware_config_from_toml + + +class Mamba2Model: + """ + Mamba-2 architecture performance model. + + Uses PerfModel for per-layer cycle counting and computes overall inference + performance (prefill time, decode time, TPS, TTFT). + """ + + def __init__( + self, + model_config_path: str, + hardware_config, + custom_isa_path: str, + batch_size: int = 1, + input_seq_len: int = 2048, + output_seq_len: int = 128, + device_num: int = 1, + frequency_hz: float = 1e9, + ): + with open(model_config_path) as f: + model_param = json.load(f) + + self.hidden_size = model_param["hidden_size"] + self.num_hidden_layers = model_param["num_hidden_layers"] + self.vocab_size = model_param["vocab_size"] + + # SSM geometry. HF's Mamba2Config names these exactly this way. + self.state_size = model_param.get("state_size", 128) + self.expand = model_param.get("expand", 2) + self.d_inner = model_param.get("intermediate_size", self.expand * self.hidden_size) + self.head_dim = model_param.get("head_dim", 64) + self.num_heads = model_param.get("num_heads", self.d_inner // self.head_dim) + self.n_groups = model_param.get("n_groups", 1) + self.conv_kernel = model_param.get("conv_kernel", 4) + self.chunk_size = model_param.get("chunk_size", 256) + + if self.num_heads * self.head_dim != self.d_inner: + raise ValueError( + f"num_heads*head_dim ({self.num_heads}*{self.head_dim}) != d_inner ({self.d_inner}); " + "check head_dim / expand / intermediate_size in the model config" + ) + + # in_proj emits [z | x | B | C | dt] + self.in_proj_out = 2 * self.d_inner + 2 * self.n_groups * self.state_size + self.num_heads + # conv1d runs over [x | B | C] only + self.conv_dim = self.d_inner + 2 * self.n_groups * self.state_size + + self.input_seq_len = input_seq_len + self.output_seq_len = output_seq_len + + self.frequency = frequency_hz + self.hardware_config = hardware_config + self.batch_size = batch_size + self.device_batch_size = max(1, batch_size // device_num) + self.device_num = device_num + + self.perf = PerfModel(hardware_config, custom_isa_path) + + # ------------------------------------------------------------------------- + # Reporting helpers + # ------------------------------------------------------------------------- + + @property + def state_elems_per_layer(self) -> int: + """Recurrent SSM state size per sequence, per layer (elements).""" + return self.num_heads * self.head_dim * self.state_size + + @property + def conv_state_elems_per_layer(self) -> int: + return self.conv_dim * (self.conv_kernel - 1) + + def print_config(self): + """Print model and hardware configuration.""" + state_bytes = (self.state_elems_per_layer + self.conv_state_elems_per_layer) * self.perf.state_bytes + print("=" * 60) + print("Mamba-2 Model Configuration") + print("=" * 60) + print(f"Hidden size: {self.hidden_size}") + print(f"Num hidden layers: {self.num_hidden_layers}") + print(f"d_inner (expand={self.expand}): {self.d_inner}") + print(f"Num heads: {self.num_heads}") + print(f"Head dim: {self.head_dim}") + print(f"State size (N): {self.state_size}") + print(f"Num groups: {self.n_groups}") + print(f"Conv kernel: {self.conv_kernel}") + print(f"Conv dim: {self.conv_dim}") + print(f"Chunk size: {self.chunk_size}") + print(f"in_proj out features: {self.in_proj_out}") + print(f"Vocab size: {self.vocab_size}") + print("-" * 60) + print("Recurrent State (context-independent)") + print("-" * 60) + print(f"SSM state / layer: {self.state_elems_per_layer:,} elems") + print(f"Conv state / layer: {self.conv_state_elems_per_layer:,} elems") + print(f"Total state (all L): {state_bytes * self.num_hidden_layers / 1e6:.2f} MB per sequence") + print("-" * 60) + print("Inference Settings") + print("-" * 60) + print(f"Batch size: {self.batch_size}") + print(f"Input seq len: {self.input_seq_len}") + print(f"Output seq len: {self.output_seq_len}") + print(f"Device num: {self.device_num}") + print("-" * 60) + print("Hardware Config") + print("-" * 60) + print(f"MLEN: {self.perf.mlen}, BLEN: {self.perf.blen}, VLEN: {self.perf.vlen}") + print(f"HBM bandwidth: {self.perf.hbm_bytes_per_cycle:.1f} B/cycle") + print(f"bytes/elem weight={self.perf.weight_bytes} act={self.perf.act_bytes} state={self.perf.state_bytes}") + print("=" * 60) + + # ------------------------------------------------------------------------- + # One mixer block + # ------------------------------------------------------------------------- + + def _block_cycles(self, seq_len: int, mode: str) -> dict: + """Cycle counts for one Mamba-2 block, broken down by stage.""" + b = self.device_batch_size + parts = {} + + # pre-mixer RMSNorm + parts["rms"] = self.perf.rms_layer(self.hidden_size, seq_len, b, mode) + + # in_proj: hidden -> 2*d_inner + 2*n_groups*N + num_heads + parts["in_proj"] = self.perf.linear(self.hidden_size, self.in_proj_out, seq_len, b, mode) + + # causal depthwise conv1d + SiLU over [x | B | C] + parts["conv1d"] = self.perf.causal_conv1d(self.conv_dim, self.conv_kernel, seq_len, b, mode) + + # dt = clamp(softplus(dt + dt_bias)) + parts["dt"] = self.perf.dt_activation(self.num_heads, seq_len, b, mode) + + # the selective scan itself + if mode == "prefill": + parts["ssd"] = self.perf.ssd_chunk_scan( + num_heads=self.num_heads, + head_dim=self.head_dim, + state_size=self.state_size, + n_groups=self.n_groups, + chunk_size=self.chunk_size, + seq_len=seq_len, + batch_size=b, + ) + else: + parts["ssd"] = self.perf.ssd_recurrence_decode( + num_heads=self.num_heads, + head_dim=self.head_dim, + state_size=self.state_size, + n_groups=self.n_groups, + batch_size=b, + ) + + # gated RMSNorm on the d_inner-wide output + parts["gated_norm"] = self.perf.gated_rms_norm(self.d_inner, seq_len, b, mode) + + # out_proj: d_inner -> hidden + parts["out_proj"] = self.perf.linear(self.d_inner, self.hidden_size, seq_len, b, mode) + + # residual + parts["residual"] = self.perf.residual(self.hidden_size, seq_len, b, mode) + + return parts + + # ------------------------------------------------------------------------- + # Phase-level timing + # ------------------------------------------------------------------------- + + def compute_prefill_time(self, verbose: bool = True) -> float: + """Compute prefill phase execution time in seconds.""" + mode = "prefill" + overall_exe_cycle = self.perf.embeddings(self.hidden_size, self.input_seq_len, self.device_batch_size, mode) + + parts = self._block_cycles(self.input_seq_len, mode) + block_cycles = sum(parts.values()) + overall_exe_cycle += block_cycles * self.num_hidden_layers + + # final norm + LM head + overall_exe_cycle += self.perf.rms_layer(self.hidden_size, self.input_seq_len, self.device_batch_size, mode) + overall_exe_cycle += self.perf.lm_head(self.hidden_size, self.vocab_size, self.device_batch_size) + + execution_time = overall_exe_cycle / self.frequency + + if verbose: + print("\nPrefill Execution Distribution:") + for name, cyc in parts.items(): + print(f" {name:<12} {cyc / block_cycles * 100:6.2f}%") + print(f"\n Total cycles: {overall_exe_cycle:,}") + + return execution_time + + def compute_decode_time(self, output_token_size: int, verbose: bool = True) -> float: + """Compute decode phase execution time in seconds. + + Unlike an attention model there is no kv_size to advance: every decode + step costs exactly the same regardless of how much context precedes it. + We therefore evaluate one step and multiply. + """ + mode = "decode" + # Every layer is evaluated explicitly (rather than one block x num_layers) + # so that PerfModel's HBM traffic counters hold one full token's traffic + # after this call — that is what makes the state-vs-KV comparison auditable. + parts: dict[str, int] = {} + for _ in range(self.num_hidden_layers): + for name, cyc in self._block_cycles(1, mode).items(): + parts[name] = parts.get(name, 0) + cyc + per_step = sum(parts.values()) + per_step += self.perf.rms_layer(self.hidden_size, 1, self.device_batch_size, mode) + per_step += self.perf.lm_head(self.hidden_size, self.vocab_size, self.device_batch_size) + + overall_inst_num = per_step * output_token_size + # Factor of 2 matches llama_model.py / gpt_oss_model.py: instruction issue + + # memory access pipeline stages in decode mode. Kept identical so the two + # architectures are compared on the same footing. + overall_exe_cycle = overall_inst_num * 2 + execution_time = overall_exe_cycle / self.frequency + + if verbose: + block_total = sum(parts.values()) + print("\nDecode Execution Distribution (mixer blocks):") + for name, cyc in parts.items(): + print(f" {name:<12} {cyc / block_total * 100:6.2f}%") + + return execution_time + + def compute_performance(self, verbose: bool = True) -> tuple: + """ + Compute overall inference performance. + + Returns: + tuple: (TTFT in seconds, TPS) + """ + prefill_time = self.compute_prefill_time(verbose) + first_token_decode = self.compute_decode_time(1, verbose=False) + ttft = (prefill_time + first_token_decode) / self.device_num + + decode_time = self.compute_decode_time(self.output_seq_len, verbose) + tps = (self.batch_size * self.output_seq_len) / decode_time + + return ttft, tps + + +# ============================================================================= +# Model Library Utilities +# ============================================================================= + + +def list_available_models(model_lib_path: Path) -> list: + """List all available model configs in Model_Lib.""" + if not model_lib_path.exists(): + return [] + return sorted([f.stem for f in model_lib_path.glob("*.json")]) + + +def resolve_model_path(model_name: str, model_lib_path: Path) -> Path: + """Resolve model name to full path.""" + model_path = model_lib_path / f"{model_name}.json" + if not model_path.exists(): + available = list_available_models(model_lib_path) + raise FileNotFoundError(f"Model '{model_name}' not found.\nAvailable models: {', '.join(available)}") + return model_path + + +# ============================================================================= +# CLI Entry Point +# ============================================================================= + + +def main(): + parser = argparse.ArgumentParser( + description="Mamba-2 Performance Model - Compute TPS and TTFT for SSM inference on PLENA", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python mamba2_model.py --model mamba2-2.7b --model-lib ./Model_Lib --config ./plena_settings.toml --isa-lib ./customISA_lib.json + python mamba2_model.py --model-path ./mamba2-2.7b.json --config ./plena_settings.toml --isa-lib ./customISA_lib.json + python mamba2_model.py --list-models --model-lib ./Model_Lib +""", + ) + + model_group = parser.add_mutually_exclusive_group(required=True) + model_group.add_argument("--model", "-m", help="Model name from Model_Lib") + model_group.add_argument("--model-path", help="Full path to model config JSON") + model_group.add_argument("--list-models", "-l", action="store_true", help="List available models") + model_group.add_argument( + "--task-file", "-t", help="Path to task JSON file specifying model, batch, input_seq, output_seq" + ) + + parser.add_argument( + "--model-lib", required=False, help="Path to Model_Lib directory (required for --model and --list-models)" + ) + parser.add_argument("--config", "-c", required=False, help="Path to hardware config TOML (required for inference)") + parser.add_argument("--isa-lib", required=False, help="Path to customISA_lib.json (required for inference)") + parser.add_argument("--batch-size", "-b", type=int, default=4, help="Batch size (default: 4)") + parser.add_argument("--input-seq", "-i", type=int, default=2048, help="Input sequence length (default: 2048)") + parser.add_argument("--output-seq", "-o", type=int, default=1024, help="Output sequence length (default: 1024)") + parser.add_argument("--device-num", "-d", type=int, default=1, help="Number of devices (default: 1)") + parser.add_argument("--json", action="store_true", help="Output results as JSON") + parser.add_argument("--quiet", "-q", action="store_true", help="Suppress detailed output") + + args = parser.parse_args() + + if args.list_models: + if not args.model_lib: + parser.error("--model-lib is required for --list-models") + print("Available models:") + for model in list_available_models(Path(args.model_lib)): + print(f" {model}") + return + + if args.task_file: + with open(args.task_file) as f: + task = json.load(f) + args.model = task.get("model") + args.batch_size = task.get("batch_size", args.batch_size) + args.input_seq = task.get("input_seq", args.input_seq) + args.output_seq = task.get("output_seq", args.output_seq) + args.device_num = task.get("device_num", args.device_num) + if not args.model: + parser.error("Task file must specify 'model'") + + if not args.config: + parser.error("--config is required for inference") + if not args.isa_lib: + parser.error("--isa-lib is required for inference") + + if args.model or args.task_file: + if not args.model_lib: + parser.error("--model-lib is required when using --model") + model_path = str(resolve_model_path(args.model, Path(args.model_lib))) + else: + model_path = args.model_path + + hardware_config = load_hardware_config_from_toml(args.config) + + model = Mamba2Model( + model_config_path=model_path, + hardware_config=hardware_config, + custom_isa_path=args.isa_lib, + batch_size=args.batch_size, + input_seq_len=args.input_seq, + output_seq_len=args.output_seq, + device_num=args.device_num, + ) + + if not args.quiet: + model.print_config() + + ttft, tps = model.compute_performance(verbose=not args.quiet) + + if args.json: + result = { + "model": args.model or args.model_path, + "architecture": "mamba2", + "batch_size": args.batch_size, + "input_seq_len": args.input_seq, + "output_seq_len": args.output_seq, + "device_num": args.device_num, + "ttft_seconds": ttft, + "ttft_ms": ttft * 1000, + "tps": tps, + } + print(json.dumps(result, indent=2)) + else: + print("\n" + "=" * 60) + print("Performance Results") + print("=" * 60) + print(f"TTFT (Time to First Token): {ttft:.6f} s ({ttft * 1000:.3f} ms)") + print(f"TPS (Tokens Per Second): {tps:.2f}") + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/analytic_models/performance/perf_model.py b/analytic_models/performance/perf_model.py index 70a7b26a2..dd6d79d6a 100644 --- a/analytic_models/performance/perf_model.py +++ b/analytic_models/performance/perf_model.py @@ -2,7 +2,26 @@ PLENA Hardware Performance Model. Provides per-layer hardware latency modeling using instruction latencies. -This module is used by llama_model.py for LLM-level performance estimation. +This module is used by llama_model.py / gpt_oss_model.py / mamba2_model.py for +LLM-level performance estimation. + +Two cost axes are modelled per layer stage: + +1. *Compute / instruction-issue cycles* — tile counts multiplied by the + ``pipelined`` cycle cost of each opcode in ``customISA_lib.json``. +2. *Memory cycles* — the bytes each stage has to move across the HBM interface + divided by ``HBM_BANDWIDTH_BYTES_PER_CYCLE`` (see ``PerfModel`` docstring). + +A stage costs ``max(compute_cycles, mem_cycles)``, not the sum. Rationale: the +PLENA datapath double-buffers HBM prefetch against systolic/vector execution +(``H_PREFETCH_*`` issues asynchronously and the consumer stalls only if the data +has not landed), so a perfectly pipelined stage runs at whichever of the two +limits is worse. ``max`` is the standard roofline bound; summing would model a +machine that never overlaps DMA with compute and would roughly double every +memory-bound stage. The cost of ``max`` is that it ignores prologue/epilogue +ramp (the first tile's fetch is genuinely exposed), so the model is optimistic +by roughly one tile-fetch per stage. Set ``enable_bandwidth=False`` to recover +the pre-bandwidth (compute-only) behaviour exactly. """ import json @@ -11,6 +30,59 @@ import toml from pydantic import BaseModel, Field, model_validator +# ============================================================================= +# Numeric format parsing ([ANALYTIC.PRECISION.*]) +# ============================================================================= + + +def _scalar_type_bits(type_spec: dict) -> int: + """Bit width of a single scalar element described by a PRECISION type table. + + Handles the two shapes used in plena_settings.toml: + {type="Fp", sign=, exponent=, mantissa=} + {type="Int", width=} + """ + kind = str(type_spec.get("type", "Fp")).lower() + if kind == "int": + return int(type_spec["width"]) + if kind == "fp": + return int(bool(type_spec.get("sign", False))) + int(type_spec["exponent"]) + int(type_spec["mantissa"]) + raise ValueError(f"Unknown PRECISION element type: {type_spec.get('type')!r}") + + +def _entry_bytes_per_elem(entry: dict) -> float: + """Bytes per element for one [ANALYTIC.PRECISION.] entry. + + ``format = "Plain"`` -> DATA_TYPE bits / 8. + ``format = "Mx"`` -> a block of ``block`` elements shares ONE scale, so the + amortised width is (elem_bits*block + scale_bits) / block. + MXFP8 at block=8 with an 8-bit scale is therefore + (8*8 + 8) / (8*8) = 1.125 bytes/elem, not 1.0. + bare type table -> (e.g. SCALAR_FP) treated as Plain. + """ + fmt = entry.get("format") + if fmt == "Mx": + block = int(entry.get("block", 1)) + elem_bits = _scalar_type_bits(entry["ELEM"]) + scale_bits = _scalar_type_bits(entry["SCALE"]) + return (elem_bits * block + scale_bits) / (8.0 * block) + if fmt == "Plain": + return _scalar_type_bits(entry["DATA_TYPE"]) / 8.0 + if "type" in entry: # bare type table, e.g. [ANALYTIC.PRECISION.SCALAR_FP] + return _scalar_type_bits(entry) / 8.0 + raise ValueError(f"Unrecognised PRECISION entry: {entry!r}") + + +def parse_precision_bytes(precision_section: dict) -> dict[str, float]: + """Parse a whole [.PRECISION] section into {name: bytes_per_element}.""" + out: dict[str, float] = {} + for name, entry in precision_section.items(): + if not isinstance(entry, dict): + continue + out[name] = _entry_bytes_per_elem(entry) + return out + + # ============================================================================= # Hardware Configuration Schema # ============================================================================= @@ -28,6 +100,21 @@ class HardwareConfig(BaseModel): # Memory configuration VECTOR_SRAM_SIZE: int = Field(gt=0, description="Vector SRAM size in elements") HBM_V_Prefetch_Amount: int = Field(gt=0, description="HBM vector prefetch amount") + HBM_V_Writeback_Amount: int = Field(default=16, gt=0, description="HBM vector writeback amount") + + # HBM interface. HBM_WIDTH is a row/burst width in BYTES — this matches every + # other consumer in the tree (address_alloc.py's `hbm_row_width: int = 512 + # # bytes`, PLENA_Tools/memory_mapping/rand_gen.py's `HBM_WIDTH // 4` floats + # per line). One row per cycle is the assumed peak. + HBM_WIDTH: int = Field(default=512, gt=0, description="HBM row width in bytes") + HBM_BANDWIDTH_BYTES_PER_CYCLE: float = Field( + default=0.0, + ge=0.0, + description="Explicit HBM bandwidth override in bytes/cycle. 0 = derive from HBM_WIDTH.", + ) + + # bytes-per-element per memory class, parsed from [.PRECISION.*] + PRECISION_BYTES: dict[str, float] = Field(default_factory=dict, description="Bytes per element by memory class") # Allow extra fields for latency parameters (dynamically loaded) model_config = {"extra": "allow"} @@ -109,6 +196,11 @@ def load_hardware_config_from_toml(toml_path: str) -> HardwareConfig: elif "value" in val: config_dict[param_name] = val["value"] + # Extract PRECISION section -> bytes per element per memory class. + # Previously ignored entirely, which left the model with no notion of how + # many bytes a tensor actually occupies in HBM. + config_dict["PRECISION_BYTES"] = parse_precision_bytes(analytic_data.get("PRECISION", {})) + return HardwareConfig(**config_dict) @@ -158,22 +250,50 @@ class PerfModel: On init, builds pipelined instruction latency from customISA_lib.json using hardware config values (MLEN, BLEN, VLEN, HLEN, latency params). + Memory-bandwidth model + ---------------------- + Peak HBM bandwidth is ``HBM_BANDWIDTH_BYTES_PER_CYCLE`` if set in the config, + otherwise ``HBM_WIDTH`` bytes/cycle (one HBM row per cycle). With the shipped + analytic config that is 512 B/cycle, i.e. 512 GB/s at the 1 GHz clock + ``llama_model.py`` assumes. + + Bytes per element come from ``[ANALYTIC.PRECISION.*]``; MX formats amortise + their shared block scale, so MXFP8 at block=8 is 1.125 B/elem. + + Each layer stage returns ``max(compute_cycles, mem_cycles)`` — see the module + docstring for why ``max`` and not a sum, and what that approximation costs. + + Assumptions worth knowing before trusting a number: + * Weights are streamed from HBM on *every* invocation of a stage. At batch 1 + decode this is the dominant term and is realistic; at large batch, a real + scheduler would amortise one weight fetch over the whole batch, and this + model does too (weight bytes do not scale with batch). + * No cache/SRAM residency model for weights across layers: a model small + enough to stay resident in on-chip SRAM would be over-charged here. + * ``max`` per stage, summed across stages: it cannot express overlap + *between* stages (prefetching layer N+1's weights during layer N). + Attributes: config: Validated HardwareConfig instr: Validated InstructionLatency (access via instr["M_MM"]) mlen, blen, vlen, hlen: Hardware dimensions + hbm_bytes_per_cycle: Peak HBM bandwidth used by the roofline + weight_bytes / kv_bytes / act_bytes / state_bytes: bytes per element """ config: HardwareConfig instr: InstructionLatency - def __init__(self, hardware_config: HardwareConfig, custom_isa_path: str): + def __init__(self, hardware_config: HardwareConfig, custom_isa_path: str, enable_bandwidth: bool = True): """ Initialize PerfModel. Args: hardware_config: Validated hardware configuration custom_isa_path: Path to customISA_lib.json + enable_bandwidth: If False, every stage returns its compute-only cycle + count (the pre-bandwidth behaviour). Used by the regression test + that pins backward compatibility. """ self.config = hardware_config self.mlen = hardware_config.MLEN @@ -182,14 +302,90 @@ def __init__(self, hardware_config: HardwareConfig, custom_isa_path: str): self.hlen = hardware_config.HLEN self.vector_sram_size = hardware_config.VECTOR_SRAM_SIZE * self.vlen self.prefetch_v_amount = hardware_config.HBM_V_Prefetch_Amount + # BUGFIX: stores are writebacks and must use the writeback amount. The + # previous code used HBM_V_Prefetch_Amount for H_STORE_V while + # HBM_V_Writeback_Amount sat unused in plena_settings.toml. + self.writeback_v_amount = hardware_config.HBM_V_Writeback_Amount + + # ---- memory bandwidth / precision ----------------------------------- + self.enable_bandwidth = enable_bandwidth + self.precision_bytes: dict[str, float] = dict(hardware_config.PRECISION_BYTES) + # bf16 (2 B) fallbacks if a config predates the PRECISION section + self.weight_bytes = self.precision_bytes.get("HBM_M_WEIGHT_TYPE", 2.0) + self.kv_bytes = self.precision_bytes.get("HBM_M_KV_TYPE", 2.0) + self.act_bytes = self.precision_bytes.get("HBM_V_ACT_TYPE", 2.0) + # Recurrent SSM state lives on the vector side; HBM_V_KV_TYPE is the + # closest declared class (same MX format as the attention KV cache). + self.state_bytes = self.precision_bytes.get("HBM_V_KV_TYPE", self.kv_bytes) + + bw_override = float(getattr(hardware_config, "HBM_BANDWIDTH_BYTES_PER_CYCLE", 0.0) or 0.0) + self.hbm_bytes_per_cycle = bw_override if bw_override > 0 else float(hardware_config.HBM_WIDTH) + + self.reset_traffic() # Build validated instruction latencies self.instr = build_pipelined_latency(hardware_config, custom_isa_path) + # ------------------------------------------------------------------------- + # Bandwidth / roofline helpers + # ------------------------------------------------------------------------- + + def reset_traffic(self) -> None: + """Zero the traffic/roofline accounting counters.""" + self.traffic_bytes = 0.0 + self.compute_only_cycles = 0 + self.memory_only_cycles = 0 + self.roofline_cycles = 0 + self.memory_bound_stages = 0 + self.total_stages = 0 + + def mem_cycles(self, num_bytes: float) -> int: + """Cycles to move ``num_bytes`` across the HBM interface at peak bandwidth.""" + if num_bytes <= 0: + return 0 + return math.ceil(num_bytes / self.hbm_bytes_per_cycle) + + def _roofline(self, compute_cycles: int, mem_bytes: float = 0.0) -> int: + """Combine a stage's compute cost and its HBM traffic into one cycle count. + + Returns ``max(compute, memory)`` (see module docstring). Also accumulates + traffic statistics so callers can report how memory-bound a phase is. + """ + compute_cycles = int(compute_cycles) + mc = self.mem_cycles(mem_bytes) + + self.traffic_bytes += mem_bytes + self.compute_only_cycles += compute_cycles + self.memory_only_cycles += mc + self.total_stages += 1 + if mc > compute_cycles: + self.memory_bound_stages += 1 + + result = max(compute_cycles, mc) if self.enable_bandwidth else compute_cycles + self.roofline_cycles += result + return result + + def traffic_summary(self) -> dict: + """Snapshot of the traffic counters accumulated since ``reset_traffic()``.""" + return { + "hbm_bytes_per_cycle": self.hbm_bytes_per_cycle, + "traffic_bytes": self.traffic_bytes, + "compute_only_cycles": self.compute_only_cycles, + "memory_only_cycles": self.memory_only_cycles, + "roofline_cycles": self.roofline_cycles, + "memory_bound_stages": self.memory_bound_stages, + "total_stages": self.total_stages, + } + # ------------------------------------------------------------------------- # Helpers # ------------------------------------------------------------------------- + def _spill_elems(self, hidden_size: int, seq_len: int, batch_size: int) -> int: + """Activation elements that do NOT fit in vector SRAM and spill to HBM.""" + total = hidden_size * seq_len * batch_size + return max(0, total - self.vector_sram_size) + def _effective_btmm(self, eff_rows: int, eff_cols: int) -> int: """Effective BTMM cost for a tile of actual size eff_rows × eff_cols. @@ -210,30 +406,29 @@ def rms_layer(self, hidden_size: int, seq_len: int, batch_size: int, mode: str = loop_inst_num = 8 loop_num = hidden_size // self.vlen overall_cycles = 0 + mem_bytes = 0.0 if mode == "prefill": compute_cycle = ( setting_inst_num * self.instr["S_BASIC"] + loop_num * loop_inst_num * seq_len * self.instr["V_BASIC"] * batch_size ) - if hidden_size * seq_len * batch_size > self.vector_sram_size: + spill = self._spill_elems(hidden_size, seq_len, batch_size) + if spill: overall_cycles += ( - compute_cycle - + ( - (hidden_size * seq_len * batch_size - self.vector_sram_size) - // (self.vlen * self.prefetch_v_amount) - ) - * self.instr["H_PREFETCH_V"] - * 2 + compute_cycle + (spill // (self.vlen * self.prefetch_v_amount)) * self.instr["H_PREFETCH_V"] * 2 ) + # spilled activations are read back and written out again + mem_bytes = 2 * spill * self.act_bytes else: overall_cycles += compute_cycle else: # decode + # A single token's activations always fit in vector SRAM -> no HBM traffic. overall_cycles = ( setting_inst_num * self.instr["S_BASIC"] + loop_num * loop_inst_num * self.instr["V_BASIC"] * batch_size ) - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def projection( self, @@ -250,6 +445,13 @@ def projection( overall_cycles = 0 bs_dim = seq_len * batch_size + # HBM traffic shared by both modes: the Q/K/V weight matrices are streamed + # once per invocation, and the newly produced K,V are appended to the cache. + # (o_proj is not modelled here — the callers never charge for it.) + weight_elems = hidden_size * hidden_size + 2 * hidden_size * (num_kv_heads * head_dim) + kv_written_elems = 2 * batch_size * (seq_len if mode == "prefill" else 1) * num_kv_heads * head_dim + mem_bytes = weight_elems * self.weight_bytes + kv_written_elems * self.kv_bytes + if mode == "prefill": # Projection of Q compute_cycle += ( @@ -280,22 +482,22 @@ def projection( * self.instr["M_MM"] ) # Activation (Q) Manipulation - if hidden_size * seq_len * batch_size > self.vector_sram_size: + spill = self._spill_elems(hidden_size, seq_len, batch_size) + if spill: overall_cycles += ( - compute_cycle - + ( - (hidden_size * seq_len * batch_size - self.vector_sram_size) - // (self.vlen * self.prefetch_v_amount) - ) - * self.instr["H_PREFETCH_V"] - * 2 + compute_cycle + (spill // (self.vlen * self.prefetch_v_amount)) * self.instr["H_PREFETCH_V"] * 2 ) + mem_bytes += 2 * spill * self.act_bytes else: overall_cycles += compute_cycle # Store K,V Cache + # BUGFIX (a): math.ceil, not floor. Floor division silently charged ZERO + # cycles whenever the cache slice was smaller than one writeback burst. + # BUGFIX (b): a store is metered by HBM_V_Writeback_Amount, not the + # prefetch amount. overall_cycles += ( 2 - * ((batch_size * seq_len * num_kv_heads * head_dim) // (self.vlen * self.prefetch_v_amount)) + * math.ceil((batch_size * seq_len * num_kv_heads * head_dim) / (self.vlen * self.writeback_v_amount)) * self.instr["H_STORE_V"] ) else: # decode @@ -329,13 +531,16 @@ def projection( * self.instr["M_MM"] ) overall_cycles += compute_cycle - # Store K,V Cache + # Store K,V Cache — same two bugfixes as the prefill path. At the + # default config this term used to evaluate to + # (4 * 8 * 128) // (2048 * 16) == 4096 // 32768 == 0 + # i.e. the decode cache write was free. overall_cycles += ( 2 - * ((batch_size * num_kv_heads * head_dim) // (self.vlen * self.prefetch_v_amount)) + * math.ceil((batch_size * num_kv_heads * head_dim) / (self.vlen * self.writeback_v_amount)) * self.instr["H_STORE_V"] ) - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def flash_attention( self, @@ -355,6 +560,16 @@ def flash_attention( inner_q_head_loop = num_attention_heads // num_kv_heads tr = math.ceil(seq_len / self.mlen) tc = math.ceil(kv_size / self.mlen) + + # HBM traffic: the KV cache is streamed once per Q block (FlashAttention's + # outer loop is over Q, inner over KV), so tr passes over the whole cache. + # In decode tr == 1, giving exactly one pass — the term that makes decode + # attention cost grow linearly with context. + # Approximation: causal masking would remove ~half the tiles in prefill; + # neither the compute nor the memory term models that (both charge tr*tc). + kv_elems = 2 * batch_size * kv_size * num_kv_heads * head_dim + mem_bytes = tr * kv_elems * self.kv_bytes + if mode == "prefill": # Effective BTMM cost — scales with actual tile fill, not full MLEN×MLEN eff_btmm = self._effective_btmm(min(seq_len, self.mlen), min(kv_size, self.mlen)) @@ -396,7 +611,7 @@ def flash_attention( inner_compute_cycles += (2 * self.instr["V_BASIC"] + 1) * inner_q_head_loop overall_cycles = inner_compute_cycles * tr * tc * kv_head_loop * batch_size - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def self_attention( self, @@ -415,6 +630,10 @@ def self_attention( kv_head_loop = num_kv_heads inner_q_head_loop = num_attention_heads // num_kv_heads + # Unfused attention materialises the full S/P matrix, but this model does + # not charge for spilling it; only the KV cache read is counted (one pass). + mem_bytes = 2 * batch_size * kv_size * num_kv_heads * head_dim * self.kv_bytes + if mode == "prefill": # S = Q (seq_len, num_attention_heads, head_dim) @ K^T (seq_len, num_kv_heads, head_dim) = (num_attention_heads, seq_len, seq_len) eff_btmm = self._effective_btmm(min(seq_len, self.mlen), min(kv_size, self.mlen)) @@ -477,7 +696,7 @@ def self_attention( * num_attention_heads ) overall_cycles = single_batch_compute_cycles * batch_size - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def mlp_moe( self, @@ -498,6 +717,13 @@ def mlp_moe( """ overall_cycles = 0 + # HBM traffic: router weights plus the weights of every expert actually + # touched. Prefill touches all of them; decode touches at most + # batch*top_k distinct experts (capped by num_experts). + expert_weight_elems = 3 * hidden_size * intermediate_size + touched_experts = num_experts if mode == "prefill" else min(num_experts, max(1, batch_size * expert_per_token)) + mem_bytes = (hidden_size * num_experts + touched_experts * expert_weight_elems) * self.weight_bytes + if mode == "prefill": # Total tokens being processed total_tokens = batch_size * seq_len @@ -616,7 +842,7 @@ def mlp_moe( * (self.instr["V_MUL_VV"] + self.instr["V_ADD_VV"]) ) - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def sliding_window_attention( self, @@ -647,6 +873,12 @@ def sliding_window_attention( kv_head_loop = num_kv_heads inner_q_head_loop = num_attention_heads // num_kv_heads + # HBM traffic: only the in-window slice of the KV cache is ever read, which + # is what bounds sliding-window decode to O(1) memory per token. + window_kv = min(seq_len if mode == "prefill" else kv_size, sliding_window_size) + tr = math.ceil(seq_len / self.mlen) if mode == "prefill" else 1 + mem_bytes = tr * 2 * batch_size * window_kv * num_kv_heads * head_dim * self.kv_bytes + if mode == "prefill": # Effective attention window per query position # Each query at position i attends to keys in range [max(0, i - sliding_window_size + 1), i] @@ -758,31 +990,28 @@ def sliding_window_attention( ) overall_cycles = single_batch_compute_cycles * batch_size - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def residual(self, hidden_size: int, seq_len: int, batch_size: int, mode: str = "prefill") -> int: """Residual connection cycle count.""" iteration = hidden_size // self.vlen overall_cycles = 0 + mem_bytes = 0.0 if mode == "prefill": compute_cycle = (self.instr["V_ADD_VV"] + 3) * seq_len * iteration * batch_size - if hidden_size * seq_len * batch_size > self.vector_sram_size: + spill = self._spill_elems(hidden_size, seq_len, batch_size) + if spill: overall_cycles += ( - compute_cycle - + ( - (hidden_size * seq_len * batch_size - self.vector_sram_size) - // (self.vlen * self.prefetch_v_amount) - ) - * self.instr["H_PREFETCH_V"] - * 2 + compute_cycle + (spill // (self.vlen * self.prefetch_v_amount)) * self.instr["H_PREFETCH_V"] * 2 ) + mem_bytes = 2 * spill * self.act_bytes else: overall_cycles += compute_cycle else: compute_cycle = (self.instr["V_ADD_VV"] + 3) * iteration * batch_size overall_cycles += compute_cycle - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def feed_forward( self, hidden_size: int, intermediate_size: int, seq_len: int, batch_size: int, mode: str = "prefill" @@ -790,6 +1019,9 @@ def feed_forward( """Feed-forward (MLP) layer cycle count.""" overall_cycles = 0 + # HBM traffic: gate + up + down weight matrices, streamed once per call. + mem_bytes = 3 * hidden_size * intermediate_size * self.weight_bytes + if mode == "prefill": # Upsize Linear and Gate overall_cycles += ( @@ -829,19 +1061,20 @@ def feed_forward( * self.instr["M_MM"] ) - return overall_cycles + return self._roofline(overall_cycles, mem_bytes) def embeddings(self, hidden_size: int, seq_len: int, batch_size: int, mode: str = "prefill") -> int: """Embedding layer cycle count.""" setting_inst_num = 3 overall_cycles = setting_inst_num * self.instr["S_BASIC"] + tokens = seq_len * batch_size if mode == "prefill" else batch_size - if mode == "prefill": - overall_cycles += seq_len * batch_size * math.ceil(hidden_size / self.vlen) * self.instr["H_PREFETCH_V"] - else: # decode - overall_cycles += batch_size * math.ceil(hidden_size / self.vlen) * self.instr["H_PREFETCH_V"] + # HBM traffic: one gathered embedding row per token. + mem_bytes = tokens * hidden_size * self.act_bytes - return overall_cycles + overall_cycles += tokens * math.ceil(hidden_size / self.vlen) * self.instr["H_PREFETCH_V"] + + return self._roofline(overall_cycles, mem_bytes) def lm_head(self, hidden_size: int, vocab_size: int, batch_size: int) -> int: """LM head cycle count (linear projection to vocab).""" @@ -856,7 +1089,10 @@ def lm_head(self, hidden_size: int, vocab_size: int, batch_size: int) -> int: * self.instr["M_MM"] ) - return overall_cycles + # HBM traffic: the (hidden x vocab) unembedding matrix. + mem_bytes = hidden_size * vocab_size * self.weight_bytes + + return self._roofline(overall_cycles, mem_bytes) def lm_head_full_seq(self, hidden_size: int, vocab_size: int, seq_len: int, batch_size: int) -> int: """LM head cycle count over full sequence (used by LLaDA: all positions need logits).""" @@ -871,7 +1107,9 @@ def lm_head_full_seq(self, hidden_size: int, vocab_size: int, seq_len: int, batc * self.instr["M_MM"] ) - return overall_cycles + mem_bytes = hidden_size * vocab_size * self.weight_bytes + + return self._roofline(overall_cycles, mem_bytes) def softmax_full_seq(self, vocab_size: int, seq_len: int, batch_size: int) -> int: """Softmax over vocab for all sequence positions (used by LLaDA for confidence scoring).""" @@ -879,4 +1117,458 @@ def softmax_full_seq(self, vocab_size: int, seq_len: int, batch_size: int) -> in loop_num = math.ceil(vocab_size / self.vlen) # softmax: max-reduce + exp + sum + divide = ~6 V_ ops per chunk overall_cycles = batch_size * seq_len * loop_num * 6 * self.instr["V_BASIC"] - return overall_cycles + # The full logit tensor is far larger than vector SRAM: read once, write once. + mem_bytes = 2 * batch_size * seq_len * vocab_size * self.act_bytes + return self._roofline(overall_cycles, mem_bytes) + + # ------------------------------------------------------------------------- + # Mamba-2 / selective state-space primitives + # ------------------------------------------------------------------------- + + def linear(self, in_features: int, out_features: int, seq_len: int, batch_size: int, mode: str = "prefill") -> int: + """Generic dense projection [tokens, in_features] @ [in_features, out_features]. + + Used for Mamba-2's in_proj / out_proj, which are plain GEMMs with no RoPE + and no KV cache, so `projection()` (attention-shaped) does not apply. + """ + tokens = seq_len * batch_size if mode == "prefill" else batch_size + overall_cycles = ( + math.ceil(tokens / self.blen) + * math.ceil(in_features / self.mlen) + * math.ceil(out_features / self.blen) + * self.instr["M_MM"] + ) + mem_bytes = in_features * out_features * self.weight_bytes + return self._roofline(overall_cycles, mem_bytes) + + def causal_conv1d( + self, + conv_dim: int, + conv_kernel: int, + seq_len: int, + batch_size: int, + mode: str = "prefill", + activation: bool = True, + ) -> int: + """Causal depthwise conv1d over ``conv_dim`` channels with kernel ``conv_kernel``. + + Depthwise means there is no channel mixing, so this is NOT a GEMM: it is + ``conv_kernel`` fused multiply-adds per (channel, token), vectorised across + channels on the vector unit. In decode the (conv_kernel-1)-deep rolling + window of past inputs is part of the recurrent cache: it is read, shifted + and written back every token. + """ + tokens = seq_len * batch_size if mode == "prefill" else batch_size + chunks = math.ceil(conv_dim / self.vlen) + + # kernel taps: multiply-accumulate per tap, then bias. The depthwise weight + # for tap k is a per-channel *vector*, hence V_MUL_VV not V_MUL_VF. + overall_cycles = tokens * chunks * conv_kernel * (self.instr["V_MUL_VV"] + self.instr["V_ADD_VV"]) + overall_cycles += tokens * chunks * self.instr["V_ADD_VF"] + if activation: # SiLU, same 6-op budget the FFN uses + overall_cycles += tokens * chunks * 6 * self.instr["V_BASIC"] + + conv_state_elems = batch_size * conv_dim * (conv_kernel - 1) + if mode == "prefill": + # window slides inside the vector unit; only the final state is written out + mem_bytes = conv_dim * conv_kernel * self.weight_bytes + conv_state_elems * self.state_bytes + else: + # shift the rolling window by one position, then read + write it + overall_cycles += tokens * chunks * self.instr["V_SHFT_V"] + mem_bytes = conv_dim * conv_kernel * self.weight_bytes + 2 * conv_state_elems * self.state_bytes + + return self._roofline(overall_cycles, mem_bytes) + + def dt_activation(self, num_heads: int, seq_len: int, batch_size: int, mode: str = "prefill") -> int: + """dt = clamp(softplus(dt_raw + dt_bias), dt_min, dt_max), one value per head per token. + + softplus is a single V_SOFTPLUS_V (ISA 0x39); the clamp is V_MAX_VF + V_MIN_VF. + Purely on-chip: dt is produced by in_proj and consumed by the scan. + + dt is [tokens, num_heads] and contiguous, so it is charged as + ceil(tokens*num_heads / VLEN) full-width vector ops rather than one + (mostly empty) op per token — num_heads is typically far below VLEN. + """ + tokens = seq_len * batch_size if mode == "prefill" else batch_size + vecs = math.ceil(tokens * num_heads / self.vlen) + overall_cycles = vecs * ( + self.instr["V_ADD_VF"] + self.instr["V_SOFTPLUS_V"] + self.instr["V_MAX_VF"] + self.instr["V_MIN_VF"] + ) + return self._roofline(overall_cycles, 0.0) + + def ssd_chunk_scan( + self, + num_heads: int, + head_dim: int, + state_size: int, + n_groups: int, + chunk_size: int, + seq_len: int, + batch_size: int, + ) -> int: + """Mamba-2 SSD chunked scan — the prefill form of the selective recurrence. + + Per chunk of ``chunk_size`` tokens the algorithm is (Dao & Gu, 2024 §6): + + cs = cumsum(dt * A) within the chunk [chunk] per head + G = C @ B^T [chunk, chunk] per group + L = exp(segsum(cs)) with a causal mask [chunk, chunk] per head + Y_intra = (L o G) @ X [chunk, head_dim] per head + h = decay*h_prev + (B*decay)^T @ X_scaled [state_size, head_dim] per head + Y_inter = C @ h_prev [chunk, head_dim] per head + Y = Y_intra + Y_inter + D*X + + G is computed once per *group* (B and C are shared across the heads of a + group); everything else is per head. + + Vector-lane convention: elementwise work over a dense tile is charged as + ceil(elems/VLEN) full-width ops. Work that needs a *different broadcast + scalar per row* (the decay subtraction, the X rescale) stays row-serial, + because each row needs its own S_MAP_FP_V extraction first — which is why + S_MAP_FP_V, not the matmuls, dominates this kernel at chunk_size=256. + """ + num_chunks = math.ceil(seq_len / chunk_size) + chunk = min(chunk_size, seq_len) + + c_blocks = math.ceil(chunk / self.blen) + p_blocks = math.ceil(head_dim / self.blen) + n_blocks_m = math.ceil(state_size / self.mlen) + c_blocks_m = math.ceil(chunk / self.mlen) + p_vec = math.ceil(head_dim / self.vlen) + c_vec = math.ceil(chunk / self.vlen) + + # --- per group --------------------------------------------------------- + # G = C @ B^T : [chunk, state_size] @ [state_size, chunk] + per_group = c_blocks * n_blocks_m * c_blocks * self.instr["M_MM"] + + # --- per head ---------------------------------------------------------- + per_head = 0 + # cumulative sum of dt*A inside the chunk (prefix scan), then one broadcast + # scalar per row of the decay matrix (S_MAP_FP_V, ISA 0x3A) + per_head += c_vec * (self.instr["V_MUL_VV"] + self.instr["V_PS_V"]) + per_head += chunk * self.instr["S_MAP_FP_V"] + # decay mask L = exp(cs_i - cs_j), causal: [chunk, chunk]. + # subtraction is row-serial (per-row broadcast cs_i), exp is packed. + per_head += chunk * c_vec * self.instr["V_SUB_VF"] + per_head += math.ceil(chunk * chunk / self.vlen) * self.instr["V_EXP_V"] + # L o G : dense [chunk, chunk] elementwise + per_head += math.ceil(chunk * chunk / self.vlen) * self.instr["V_MUL_VV"] + # Y_intra = (L o G) @ X : [chunk, chunk] @ [chunk, head_dim] + per_head += c_blocks * c_blocks_m * p_blocks * self.instr["M_MM"] + # X_scaled = X * exp(cs_chunk - cs_t): one broadcast scalar per row + per_head += chunk * self.instr["S_MAP_FP_V"] + chunk * p_vec * self.instr["V_MUL_VF"] + # state update h_new = B_scaled^T @ X_scaled : [state_size, chunk] @ [chunk, head_dim] + per_head += math.ceil(state_size / self.blen) * c_blocks_m * p_blocks * self.instr["M_TMM"] + # h = h_prev * decay + h_new over [state_size, head_dim] + per_head += math.ceil(state_size * head_dim / self.vlen) * (self.instr["V_MUL_VF"] + self.instr["V_ADD_VV"]) + # Y_inter = C @ h_prev : [chunk, state_size] @ [state_size, head_dim] + per_head += c_blocks * n_blocks_m * p_blocks * self.instr["M_MM"] + # Y = Y_intra + Y_inter + D*X : dense [chunk, head_dim] elementwise + per_head += math.ceil(chunk * head_dim / self.vlen) * (2 * self.instr["V_ADD_VV"] + self.instr["V_MUL_VF"]) + + overall_cycles = num_chunks * batch_size * (n_groups * per_group + num_heads * per_head) + + # --- HBM traffic ------------------------------------------------------- + # X, B, C, dt for the whole sequence, plus the final recurrent state. + # Only the part that does not fit in vector SRAM actually goes to HBM. + d_inner = num_heads * head_dim + act_elems = batch_size * seq_len * (d_inner + 2 * n_groups * state_size + num_heads) + spill = max(0, act_elems - self.vector_sram_size) + state_elems = batch_size * num_heads * head_dim * state_size + mem_bytes = 2 * spill * self.act_bytes + state_elems * self.state_bytes + + return self._roofline(overall_cycles, mem_bytes) + + def ssd_recurrence_decode( + self, + num_heads: int, + head_dim: int, + state_size: int, + n_groups: int, + batch_size: int, + ) -> int: + """Mamba-2 single-token recurrence: h = dA*h + dB*x ; y = C.h + D*x. + + The state is ``num_heads * head_dim * state_size`` elements and is read AND + written for every token — the O(1)-in-context term that replaces the + linearly growing KV cache. This is where the bandwidth model matters: with + no memory term the state update looks free. + + Approximation: the C.h contraction is charged head-serially as one M_MV per + head over a [head_dim, state_size] tile. On a machine with MLEN=2048 and + state_size=128 that leaves the matrix unit mostly idle; a real kernel would + batch heads into the array, so this term is pessimistic by up to MLEN/state_size. + """ + state_vec = math.ceil(head_dim * state_size / self.vlen) + + # dA = exp(dt * A), one scalar per head + overall_cycles = math.ceil(num_heads / self.vlen) * (self.instr["V_MUL_VV"] + self.instr["V_EXP_V"]) + # dB*x outer product, dA*h decay, and the add — over the [head_dim, state_size] + # state block of each head + overall_cycles += num_heads * state_vec * (2 * self.instr["V_MUL_VF"] + self.instr["V_ADD_VV"]) + # y = C.h : [head_dim, state_size] @ [state_size] + overall_cycles += ( + num_heads * math.ceil(head_dim / self.blen) * math.ceil(state_size / self.mlen) * self.instr["M_MV"] + ) + # y += D*x + overall_cycles += ( + num_heads * math.ceil(head_dim / self.vlen) * (self.instr["V_MUL_VF"] + self.instr["V_ADD_VV"]) + ) + overall_cycles *= batch_size + + # state read + state write, every token + state_elems = batch_size * num_heads * head_dim * state_size + mem_bytes = 2 * state_elems * self.state_bytes + + return self._roofline(overall_cycles, mem_bytes) + + # ------------------------------------------------------------------------- + # KDA (Kimi Delta Attention) — the gated delta rule + # ------------------------------------------------------------------------- + # + # Both methods below are structural derivations in the same style as the SSD + # pair above, but unlike those they have a **measured** counterpart: the KDA + # lowering exists, compiles, and runs as machine code on the transactional + # emulator, so `test_kda_stage_instruction_mix.py` checks these formulas + # against the compiler's real instruction counts at MLEN=64 rather than + # leaving them unfalsifiable. Where the two disagree the formula is wrong. + # + # KDA's decay is **channel-wise on the key axis** -- one scalar per (token, + # key), not one per token. That is what makes the state update a per-key + # sweep rather than a single scaled add, and it is the single biggest + # difference from Mamba-2's cost shape. + + def kda_recurrence_decode( + self, + num_heads: int, + key_dim: int, + value_dim: int, + batch_size: int, + ) -> int: + """KDA single-token gated delta rule, per token. + + With state ``S[key, value]`` held transposed (the layout the lowering + uses so every sweep is a row progression and therefore a hardware loop):: + + a = exp(gate) per key channel + k_hat = k * a ; q_hat = q * a + o = S^T q_hat contraction over key + e = beta * (v - S^T k_hat) + S = S * a + k_hat (x) e + + Each of the three sweeps walks ``key_dim`` rows of ``value_dim`` lanes. + The predict and update sweeps are `V_FMA_VF` — one instruction per row, + which is what makes the cost constant in ``key_dim`` per row rather than + the copy/multiply/add triple it replaced. + + **The scalar overhead is not a rounding term and is billed here.** Every + sweep iteration advances its pointers with an explicit `S_ADDI_INT` per + operand and reloads its FPRAM scalar with `S_LD_FP`. On the compiled + kernel that is 50% and 25% of the dynamic instruction stream against + 18% for the arithmetic itself, so a model counting only the FMAs would + be optimistic by roughly 4x in instruction count. Post-increment + addressing would remove most of it; until the ISA has it, this is the + machine that exists. + + Gates, L2 normalisation and the output norm are billed separately by the + caller, the same split `ssd_recurrence_decode` uses for `dt`. + """ + value_vec = math.ceil(value_dim / self.vlen) + + # One sweep iteration: the vector op, its FPRAM scalar, and the pointer + # arithmetic for both operands. + sweep_row = self.instr["S_LD_FP"] + 2 * self.instr["S_ADDI_INT"] + fma_row = self.instr["V_FMA_VF"] * value_vec + sweep_row + mul_row = self.instr["V_MUL_VF"] * value_vec + sweep_row + + # a = exp(gate), one value per (head, key) + overall_cycles = math.ceil(num_heads * key_dim / self.vlen) * (self.instr["V_MUL_VV"] + self.instr["V_EXP_V"]) + # k_hat and q_hat: two elementwise products over the key-width vectors + overall_cycles += 2 * math.ceil(num_heads * key_dim / self.vlen) * self.instr["V_MUL_VV"] + # predict o = S^T q_hat, and the error term e's S^T k_hat: two sweeps + overall_cycles += 2 * num_heads * key_dim * fma_row + # state decay S *= a, then the rank-1 accumulate S += k_hat (x) e + overall_cycles += num_heads * key_dim * (mul_row + fma_row) + overall_cycles *= batch_size + + # The state is read and written every token. This is the O(1)-in-context + # term that replaces a growing KV cache -- and the reason the bandwidth + # model matters: with no memory term the update looks free. + state_elems = batch_size * num_heads * key_dim * value_dim + mem_bytes = 2 * state_elems * self.state_bytes + + return self._roofline(overall_cycles, mem_bytes) + + def kda_chunk_prefill( + self, + num_heads: int, + key_dim: int, + value_dim: int, + chunk_size: int, + seq_len: int, + batch_size: int, + row_granular_prefetch: bool = False, + ) -> int: + """KDA chunked prefill: seven matrix products per chunk, per head. + + The chunk form collapses ``C`` sequential rank-1 updates into:: + + M = tril(k_hat @ k_tilde^T, -1) [C, C] + N = tril(q_hat @ k_tilde^T) [C, C], diagonal included + T = (I + tril(diag(beta) M, -1))^-1 diag(beta) + W = V - k_hat @ S_0^T [C, value] + E = T @ W + out = scale * (q_hat @ S_0^T + N @ E) + S_C = A_C * S_0 + E^T @ (k * A_C / A) + + **Every one of the seven has a dynamic second operand** -- there is no + weight here, only activations and state -- so each is spilled to HBM and + re-prefetched into MRAM, because MRAM is writable only by + `H_PREFETCH_M`. That spill traffic is the dominant memory term and is + billed below; a model that assumed a weight-resident operand would + under-count it by the number of products. + + The UT transform is a forward substitution, one `V_FMA_VF` sweep per row + of the chunk, so it costs ``C`` sweeps and not ``C^2`` instructions. + + ``chunk_size`` is bounded by bf16 range, not by tiling: ``1/A`` reaches + ``exp(chunk * |gate_lower_bound|)`` and overflows past 17 at Kimi's + gate floor of -5. 16 is what the lowering uses. + """ + chunk = min(chunk_size, seq_len) + num_chunks = math.ceil(seq_len / chunk_size) + + c_blocks = math.ceil(chunk / self.blen) + c_blocks_m = math.ceil(chunk / self.mlen) + k_blocks_m = math.ceil(key_dim / self.mlen) + v_blocks_m = math.ceil(value_dim / self.mlen) + k_vec = math.ceil(key_dim / self.vlen) + v_vec = math.ceil(value_dim / self.vlen) + c_vec = math.ceil(chunk / self.vlen) + + # --- per chunk, per head ---------------------------------------------- + # The cumulative decay is a running product, not exp of a running sum: + # the cumulative log-decay reaches chunk * gate_lower_bound = -80, where + # bf16's ulp is 0.31, so exponentiating a stored sum costs 17% relative + # error on A. That makes it a sequential scan of `chunk` steps, not one + # matmul against a triangular ones matrix the way Mamba-2's cumsum is. + # The scan stages the previous timestep at the current row before the + # multiply, because the row ops take one row list for both operands. + # That staging copy is a zero-fill plus an add, per step, per key block. + per_chunk = ( + chunk * k_blocks_m * (self.instr["V_MUL_VV"] + self.instr["V_ADD_VV"] + 2 * self.instr["S_ADDI_INT"]) + ) + # k_tilde = k / A, then k_hat / q_hat + per_chunk += chunk * k_vec * (self.instr["V_RECI_V"] + 3 * self.instr["V_MUL_VV"]) + + # The two grams, contracting over key: [C, key] @ [C, key]^T -> [C, C] + per_chunk += 2 * c_blocks * k_blocks_m * c_blocks_m * self.instr["M_TMM"] + # The causal mask on N + per_chunk += chunk * c_vec * self.instr["V_MUL_VV"] + + # UT transform: one FMA sweep per row, row i walking j < i + per_chunk += sum( + i * (self.instr["V_FMA_VF"] * c_vec + self.instr["S_LD_FP"] + 2 * self.instr["S_ADDI_INT"]) + for i in range(chunk) + ) + per_chunk += chunk * (self.instr["S_MAP_FP_V"] + c_vec * self.instr["V_MUL_VF"]) + + # W^T = v^T - S_0 @ k_hat^T ; E^T = W^T @ T^T + per_chunk += 2 * math.ceil(value_dim / self.blen) * c_blocks_m * k_blocks_m * self.instr["M_TMM"] + per_chunk += value_dim * c_vec * self.instr["V_SUB_VV"] + # out = scale * (q_hat @ S_0^T + N @ E) + per_chunk += 2 * c_blocks * v_blocks_m * k_blocks_m * self.instr["M_TMM"] + per_chunk += chunk * v_vec * (self.instr["V_ADD_VV"] + self.instr["V_MUL_VF"]) + # S_C = A_C * S_0 + E^T @ k_end + per_chunk += math.ceil(value_dim / self.blen) * c_blocks_m * k_blocks_m * self.instr["M_MM"] + per_chunk += value_dim * k_vec * (self.instr["V_MUL_VV"] + self.instr["V_ADD_VV"]) + + # The six spills are instructions too, not only bytes. Each one zero-fills + # the rows past its live data (the prefetch pulls a whole mlen x mlen + # block regardless), stores the tile a writeback-amount of rows at a + # time, and is prefetched back into MRAM. Leaving this out is what made a + # first version of this method 30% under the compiler's measured count + # while the decode form was within 2%. + rows_per_store = max(1, self.writeback_v_amount) + + def spill(col_blocks: int, live_rows: int) -> int: + """One spill of a tile `col_blocks` column blocks wide. + + The width matters and is not a detail. A VRAM matrix wider than + `mlen` is column-block-major, and both the zero-fill and the store + walk every block -- so a `[chunk, key_dim]` tile at `key_dim` 128 + against `mlen` 64 costs twice what the same tile one block wide + does. Billing every spill at one block under-counted the key axis by + half, which showed up as the model scaling with `key_dim` at half + the rate the compiler does. + + **The height is where the machine's width leaks in.** Today the + prefetch pulls a whole `mlen x mlen` MRAM block -- `k_block_count` + selects whole blocks and cannot trim a partial one -- so the tile has + to be `mlen` rows tall and every row past `live_rows` must be zeroed + before the store, or the matmul contracts against whatever was there. + The fill therefore scales with `mlen` while the data scales with + `chunk`, and KDA's chunk cannot grow to meet it: `1/A` overflows bf16 + past 17. + + `row_granular_prefetch` models the instruction set that does not have + that constraint -- a prefetch taking a row count. Then only the live + rows are stored and none need zeroing, and the term stops tracking + the machine's width. It is a hypothetical: no such instruction + exists, and this switch is here to price it rather than to assume it. + """ + rows = live_rows if row_granular_prefetch else self.mlen + fill = 0 if row_granular_prefetch else self.mlen * self.instr["V_MUL_VF"] + return col_blocks * ( + fill + math.ceil(rows / rows_per_store) * self.instr["H_STORE_V"] + self.instr["H_PREFETCH_M"] + ) + + # k_tilde, k_hat and k_end are [chunk, key]; t_mat is [chunk, chunk]; + # the state is [value, key] and the error [value, chunk]. + per_chunk += ( + 3 * spill(k_blocks_m, chunk) + + spill(c_blocks_m, chunk) + + spill(k_blocks_m, value_dim) + + spill(c_blocks_m, value_dim) + ) + + # `A_C` is broadcast down a tile by doubling the filled span, which is + # ceil(log2(n)) + 1 copies rather than n. Twice per chunk. + per_chunk += ( + 2 * (math.ceil(math.log2(max(chunk, value_dim))) + 1) * (self.instr["V_MUL_VF"] + self.instr["V_ADD_VV"]) + ) + + overall_cycles = per_chunk * num_chunks * num_heads * batch_size + + # --- memory ------------------------------------------------------------ + # Six spills per chunk per head, each written once and prefetched once. + # Sized from the tiles the lowering actually spills: k_tilde, k_hat and + # k_end are [C, key]; t_mat is [C, C]; state and err are [value, *]. + spill_elems = 3 * chunk * key_dim + chunk * chunk + value_dim * key_dim + value_dim * chunk + mem_bytes = 2 * spill_elems * self.act_bytes * num_chunks * num_heads * batch_size + # The carried state, read at the start of each chunk and written at its end. + mem_bytes += 2 * num_chunks * num_heads * key_dim * value_dim * self.state_bytes * batch_size + + return self._roofline(overall_cycles, mem_bytes) + + def gated_rms_norm(self, dim: int, seq_len: int, batch_size: int, mode: str = "prefill") -> int: + """RMSNorm(x * silu(z)) — Mamba-2's output norm, gated by the z branch.""" + tokens = seq_len * batch_size if mode == "prefill" else batch_size + chunks = math.ceil(dim / self.vlen) + setting_inst_num = 10 + loop_inst_num = 8 + + overall_cycles = setting_inst_num * self.instr["S_BASIC"] + # SiLU on the gate branch + the elementwise product + overall_cycles += tokens * chunks * (6 * self.instr["V_BASIC"] + self.instr["V_MUL_VV"]) + # the RMSNorm itself, same 8-op-per-chunk budget as rms_layer() + overall_cycles += tokens * chunks * loop_inst_num * self.instr["V_BASIC"] + + mem_bytes = 0.0 + if mode == "prefill": + spill = self._spill_elems(dim, seq_len, batch_size) + if spill: + overall_cycles += (spill // (self.vlen * self.prefetch_v_amount)) * self.instr["H_PREFETCH_V"] * 2 + mem_bytes = 2 * spill * self.act_bytes + + return self._roofline(overall_cycles, mem_bytes) diff --git a/analytic_models/performance/vector_sram_banking.py b/analytic_models/performance/vector_sram_banking.py new file mode 100644 index 000000000..bfdb45aad --- /dev/null +++ b/analytic_models/performance/vector_sram_banking.py @@ -0,0 +1,239 @@ +"""Would a banked Vector SRAM stall on the static Mamba/KDA sweeps? + +The emulator does not model banking: `VectorSram` is a flat `Vec>`, +one entry per row, so every access costs the same whatever its address. The +question is therefore analytic, and it is worth answering before anyone builds +banking into the RTL: **do the sweeps this lowering emits have strides that would +conflict?** + +The model +--------- +A vector row is `mlen` words. Under row-interleaved banking with `B` banks, row +`r` lives in bank `r % B`, and a sweep that advances by `stride` rows visits +banks `r, r+stride, r+2*stride, ...` (mod B). It touches `B / gcd(stride, B)` +distinct banks before repeating, so the fraction of the array it can keep busy +is `gcd(stride, B) / B` -- and the worst case, `stride % B == 0`, pins every +access to one bank. + +Word-interleaved banking is the other plausible mapping, and under it a whole +`mlen`-word row spans `min(mlen, B)` banks, so a row-sequential sweep never +conflicts regardless of stride. That mapping makes this whole question moot; +the study is about the row-interleaved one. + +Where the strides come from +--------------------------- +Not from a model of the algorithm -- from the emitted assembly. Every hardware +loop this lowering produces advances its pointers with +`S_ADDI_INT gpN, gpN, step`, and `step` is in words. :func:`sweep_strides` reads +them out of a compiled program, so the study tracks the code rather than a +description of it. +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass + +#: `S_ADDI_INT gpN, gpN, step` -- a pointer advancing by itself is a sweep step. +#: The initialising form is `S_ADDI_INT gpN, gp0, addr`, which this does not match. +_STEP = re.compile(r"S_ADDI_INT\s+gp(\d+)\s*,\s*gp(\d+)\s*,\s*(-?\d+)") + +#: `C_LOOP_START gpN, count` +_LOOP = re.compile(r"C_LOOP_START\s+gp(\d+)\s*,\s*(\d+)") + + +@dataclass(frozen=True) +class Sweep: + """One hardware loop's pointer advance, in rows.""" + + stride_rows: int + trips: int + + def distinct_banks(self, banks: int) -> int: + """Banks this sweep visits before it repeats.""" + if self.stride_rows == 0: + return 1 # a pinned pointer sits on one bank + return banks // math.gcd(abs(self.stride_rows), banks) + + def conflict_fraction(self, banks: int) -> float: + """Share of accesses that land on an already-busy bank. + + ``0.0`` means the sweep walks every bank in turn; ``1 - 1/banks`` is the + worst case, every access on one bank. + """ + distinct = self.distinct_banks(banks) + return 1.0 - distinct / banks + + +def sweep_strides(asm: str, mlen: int) -> list[Sweep]: + """Extract each hardware loop's pointer stride, in rows, from ``asm``. + + A loop body may advance several pointers; the one that matters for banking + is the largest stride, because that is what determines how far apart + consecutive accesses land. A pinned pointer (no advance) is reported as + stride 0, which is the worst case under row interleaving. + """ + sweeps: list[Sweep] = [] + lines = [ln.strip() for ln in asm.splitlines() if ln.strip()] + i = 0 + while i < len(lines): + loop = _LOOP.search(lines[i]) + if not loop: + i += 1 + continue + loop_reg, trips = int(loop.group(1)), int(loop.group(2)) + steps: list[int] = [] + j = i + 1 + depth = 1 + while j < len(lines) and depth: + if _LOOP.search(lines[j]): + depth += 1 + elif f"C_LOOP_END gp{loop_reg}" in lines[j]: + depth -= 1 + if not depth: + break + m = _STEP.search(lines[j]) + if m and m.group(1) == m.group(2): + steps.append(int(m.group(3))) + j += 1 + # Pointers that advance by whole rows are the vector-SRAM walkers; + # anything else is an FPRAM pointer and does not touch this array. + row_steps = [s // mlen for s in steps if s and s % mlen == 0] + sweeps.append(Sweep(max(row_steps, default=0), trips)) + i = j + 1 + return sweeps + + +def report(named: dict[str, str], mlen: int, bank_counts=(2, 4, 8, 16, 32)) -> str: + """A table of stalled-access fractions per kernel per bank count.""" + out = [ + f"Vector SRAM banking, row-interleaved, mlen={mlen}", + "", + f"{'kernel':<28} {'loops':>6} {'worst stride':>13} " + " ".join(f"{'B=' + str(b):>7}" for b in bank_counts), + "-" * (28 + 7 + 14 + 8 * len(bank_counts)), + ] + for name, asm in named.items(): + sweeps = sweep_strides(asm, mlen) + if not sweeps: + out.append(f"{name:<28} {'0':>6} (no hardware loops)") + continue + # Weight by trip count: a 128-trip conflicted loop matters more than a + # 2-trip one. + total = sum(s.trips for s in sweeps) or 1 + cells = [] + for b in bank_counts: + stalled = sum(s.trips * s.conflict_fraction(b) for s in sweeps) + cells.append(f"{stalled / total:>7.1%}") + worst = max(sweeps, key=lambda s: s.conflict_fraction(max(bank_counts))) + out.append(f"{name:<28} {len(sweeps):>6} {worst.stride_rows:>13} " + " ".join(cells)) + return "\n".join(out) + + +def _lowered_kernels(mlen: int = 64) -> dict[str, str]: + """Compile the kernels this study is about.""" + import sys + from pathlib import Path + + root = Path(__file__).resolve().parents[2] + for candidate in (root / "PLENA_Compiler", root.parent / "PLENA_Compiler"): + if (candidate / "aten" / "plena" / "compiler.py").exists(): + sys.path.insert(0, str(candidate)) + break + + from compiler.aten.models.kda.shape import KdaShape + from compiler.aten.plena import PlenaCompiler + from compiler.aten.plena.program_kda_common import kda_state_rows, kda_vector_rows + from compiler.aten.plena.program_kda_gates import kda_head_blocks, kda_key_blocks + from compiler.aten.plena.program_kda_mixer import KdaMixerBuffers + from compiler.aten.plena.program_mamba_common import Mamba2Shape + + out: dict[str, str] = {} + up = lambda n: ((n + mlen - 1) // mlen) * mlen # noqa: E731 + + # KDA: the recurrence and the whole mixer, at Kimi K3's shape. + shape = KdaShape.kimi_k3() + kb = kda_key_blocks(shape, mlen) + p = PlenaCompiler(mlen=mlen, blen=4) + a = lambda n, r: p.alloc(n, up(r), mlen, strict=False) # noqa: E731 + decay = p.fp_var("decay", size=shape.key_dim) + buffers = KdaMixerBuffers( + q=a("q", shape.num_heads * kb), + k=a("k", shape.num_heads * kb), + v=a("v", kda_vector_rows(shape, mlen)), + gate=a("g", shape.num_heads * kb), + dt_bias=a("dtb", shape.num_heads * kb), + beta_logit=a("bl", kda_head_blocks(shape, mlen)), + state=a("st", kda_state_rows(shape, mlen)), + out=a("o", kda_vector_rows(shape, mlen)), + pred=a("pr", kda_vector_rows(shape, mlen)), + err=a("er", kda_vector_rows(shape, mlen)), + sq_scratch=a("sq", shape.num_heads * kb), + decay_fp=decay, + q_hat_fp=decay, + k_hat_fp=p.fp_var("kh", size=shape.key_dim), + beta_fp=p.fp_var("b", size=kda_head_blocks(shape, mlen) * mlen), + part_fp=p.fp_var("pt", size=kb), + acc_fp=p.fp_var("ac", size=1), + output_scale_fp=p.fp_var("os", size=1), + rate_fp=p.fp_var("rt", size=shape.num_heads), + lower_bound_fp=p.fp_var("lb", size=1), + consts=p.kda_fp_constants(), + ) + mark = len(p.get_code()) + p.kda_decode_step_v0( + state=buffers.state, + q_fp=buffers.q_hat_fp, + k_fp=buffers.k_hat_fp, + decay_fp=buffers.decay_fp, + beta_fp=buffers.beta_fp, + v=buffers.v, + o=buffers.out, + pred=buffers.pred, + err=buffers.err, + shape=shape, + output_scale_fp=buffers.output_scale_fp, + head_rows=[0], + fp_head_stride=0, + ) + out["KDA recurrence (Kimi K3)"] = p.get_code()[mark:] + + mark = len(p.get_code()) + p.kda_mixer_step_v0(buffers=buffers, shape=shape, head_rows=[0]) + out["KDA mixer, one head"] = p.get_code()[mark:] + + # Mamba: the decode step. + m_shape = Mamba2Shape( + hidden_size=mlen, + num_heads=1, + head_dim=mlen, + state_size=mlen, + n_groups=1, + conv_kernel=4, + chunk_size=16, + seq_len=1, + ) + q = PlenaCompiler(mlen=mlen, blen=4) + b = lambda n, r: q.alloc(n, up(r), mlen, strict=False) # noqa: E731 + consts = q.mamba_fp_constants() + mark = len(q.get_code()) + q.ssm_decode_step_v0( + state=b("st", m_shape.state_size), + x=b("x", 1), + y=b("y", 1), + b_fp=q.fp_var("b", size=m_shape.state_size), + c_fp=q.fp_var("c", size=m_shape.state_size), + da_fp=q.fp_var("da", size=1), + dt_fp=q.fp_var("dt", size=1), + d_fp=q.fp_var("d", size=1), + scratch=b("sc", 1), + shape=m_shape, + consts=consts, + ) + out["Mamba-2 decode step"] = q.get_code()[mark:] + return out + + +if __name__ == "__main__": # pragma: no cover - a report, not a test + MLEN = 64 + print(report(_lowered_kernels(MLEN), MLEN)) diff --git a/analytic_models/test_kda_stage_calibration.py b/analytic_models/test_kda_stage_calibration.py new file mode 100644 index 000000000..192dc7085 --- /dev/null +++ b/analytic_models/test_kda_stage_calibration.py @@ -0,0 +1,246 @@ +"""The KDA cost model, checked against the compiler's real instruction counts. + +Every other stage in `perf_model.py` is a hand-derived formula with nothing to +check it against: the lowering it describes does not exist, so the formula can be +wrong in any direction and no test will say so. The KDA stages are the exception. +`PLENA_Compiler` emits this kernel, `test_kda_prefill_structure.py` counts the +instructions it emits, and those counts are reproduced below as the oracle. + +What is compared, and why it is instructions and not cycles +----------------------------------------------------------- +`PerfModel` returns cycles, which are latencies times counts, and the latencies +are themselves unvalidated. Setting every opcode latency to 1 turns the same +formula into a count of instructions issued, which the compiler can be asked for +directly. So this pins the part that can be checked -- how many of each thing the +kernel does, and how that scales -- and leaves the part that cannot. + +The hardware shape is forced to `MLEN 64 / BLEN 4 / VLEN 64`, the TRANSACTIONAL +config the measurements were taken at, not the ANALYTIC `MLEN 2048` the driver +runs with. Those two disagree (see `hybrid_model.py`'s module docstring); this +file deliberately uses the one the oracle came from. + +Where it lands +-------------- +Prefill ratios run **0.87 to 1.14** across the grid, centred on 1.0; decode is +0.98. The residual is bounded rather than tuned to zero -- the terms that would +close it are bookkeeping the formula does not model in detail (register spills, +immediate legalisation, loop prologues), and a fudge factor would make the +number agree without making the model more correct. + +One structural miss was found this way and fixed rather than absorbed. The model +scaled with `key_dim` at half the compiler's rate, worst at `key_dim 128, +value_dim 64` (0.75). The cause was real: a VRAM matrix wider than `mlen` is +column-block-major, so a spill's zero-fill and store walk every block, and +`per_spill` was billing one. Fixing that took the range from 0.75-0.96 to +0.87-1.14. **That is the whole value of having an oracle** -- the formula was +wrong in a specific, findable way, and nothing else here could have said so. + +What this covers, and what it does not +-------------------------------------- +Zeroing a term and re-running gives what each is worth. Against the compiler's +2,426 (prefill, chunk 16, 128 x 128) and 2,616 (decode): + + prefill, whole spill term removed 0.70 caught + prefill, UT forward substitution removed 0.91 NOT caught + prefill, all seven matmuls removed 0.91 NOT caught + decode, scalar sweep overhead removed 0.39 caught + decode, the FMA itself removed 0.69 caught + +**The instruction-count oracle is blind to the matmuls and the UT transform.** +Each matmul is a single `M_TMM` issue and the substitution is ~120 sweeps out of +2,450, so removing either moves the count by 9% and stays inside any bound loose +enough for the real residual -- while both are large in *cycles*, which is what +the driver actually reports. Those terms have the same unvalidated status as +every other stage in `perf_model.py`. What is pinned here is the overall +magnitude, the scaling in all three axes, and the two terms that dominate the +count: the spills and the scalar sweep overhead. + +The decode form is the tighter of the two, and only because that overhead is +billed. Half the compiled kernel's dynamic stream is `S_ADDI_INT` pointer +arithmetic and a quarter is `S_LD_FP`, against 18% for the arithmetic -- a model +counting only the vector ops lands at 0.39. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT / "analytic_models" / "performance")) + +from perf_model import PerfModel, load_hardware_config_from_toml # noqa: E402 + +#: Static instruction counts emitted by `PLENA_Compiler` for one head, measured +#: 2026-08-26 and pinned by `aten/tests/test_kda_prefill_structure.py`'s +#: `PREFILL_STATIC_MAX` at the same shapes. +MEASURED_PREFILL_STATIC = { + (4, 64, 64): 818, + (4, 128, 64): 1237, + (4, 128, 128): 1470, + (8, 64, 64): 1039, + (8, 128, 64): 1547, + (8, 128, 128): 1792, + (16, 64, 64): 1476, + (16, 128, 64): 2157, + (16, 128, 128): 2426, +} + +#: Dynamic instruction count for one KDA decode step at Kimi's per-head shape, +#: from `doc/static_path_measurements.md`. +MEASURED_DECODE_DYNAMIC = 2616 + +#: The model must stay inside these of the compiler. Measured range is 0.87-1.14, +#: so this leaves under 10% of margin on each side -- tight enough that removing +#: the spill term (0.70) fails, loose enough not to be brittle against the known +#: residual. See the module docstring for what these bounds cannot see. +RATIO_MIN, RATIO_MAX = 0.80, 1.20 + + +class _UnitLatency: + """Every opcode costs 1, so `cycles` reads as `instructions issued`.""" + + def __getitem__(self, key: str) -> int: + return 1 + + def __contains__(self, key: str) -> bool: + return True + + +def _compiler_shaped_model() -> PerfModel: + hw = load_hardware_config_from_toml(str(ROOT / "plena_settings.toml")) + # The TRANSACTIONAL shape, which is what the oracle was measured at. + hw.MLEN, hw.BLEN, hw.VLEN = 64, 4, 64 + perf = PerfModel(hw, str(ROOT / "analytic_models" / "performance" / "customISA_lib.json"), enable_bandwidth=False) + perf.instr = _UnitLatency() + return perf + + +@pytest.mark.parametrize("shape", sorted(MEASURED_PREFILL_STATIC)) +def test_prefill_instruction_count_tracks_the_compiler(shape) -> None: + chunk, key_dim, value_dim = shape + perf = _compiler_shaped_model() + modelled = perf.kda_chunk_prefill( + num_heads=1, + key_dim=key_dim, + value_dim=value_dim, + chunk_size=chunk, + seq_len=chunk, + batch_size=1, + ) + measured = MEASURED_PREFILL_STATIC[shape] + ratio = modelled / measured + assert RATIO_MIN <= ratio <= RATIO_MAX, ( + f"chunk {chunk}, key {key_dim}, value {value_dim}: modelled {modelled} " + f"against the compiler's {measured} is a ratio of {ratio:.2f}, outside " + f"[{RATIO_MIN}, {RATIO_MAX}]" + ) + + +def test_decode_instruction_count_tracks_the_compiler() -> None: + """The decode form, which lands at 0.98. + + It does so only because the scalar sweep overhead is billed: half the + compiled kernel's dynamic stream is `S_ADDI_INT` pointer arithmetic and a + quarter is `S_LD_FP`, against 18% for the arithmetic. Counting only the + vector ops -- the natural thing to write -- lands at 0.39. + """ + perf = _compiler_shaped_model() + modelled = perf.kda_recurrence_decode(num_heads=1, key_dim=128, value_dim=128, batch_size=1) + ratio = modelled / MEASURED_DECODE_DYNAMIC + assert 0.90 <= ratio <= 1.10, ( + f"modelled {modelled} against the compiler's {MEASURED_DECODE_DYNAMIC} is a ratio of {ratio:.2f}" + ) + + +def test_prefill_scales_with_both_axes_and_the_chunk() -> None: + """The magnitude could be right by luck; the scaling could not. + + Doubling the key axis, the value axis or the chunk must each raise the count, + and none of them may leave it unchanged -- an axis a formula ignores is the + failure mode a single-shape check cannot see. + """ + perf = _compiler_shaped_model() + + def cost(chunk: int, key_dim: int, value_dim: int) -> int: + return perf.kda_chunk_prefill( + num_heads=1, + key_dim=key_dim, + value_dim=value_dim, + chunk_size=chunk, + seq_len=chunk, + batch_size=1, + ) + + base = cost(8, 64, 64) + assert cost(8, 128, 64) > base, "the key axis must cost something" + assert cost(8, 64, 128) > base, "the value axis must cost something" + assert cost(16, 64, 64) > base, "the chunk must cost something" + + +def test_decode_state_traffic_is_the_whole_state_read_and_written() -> None: + """The O(1)-in-context term, which is the point of a recurrent mixer. + + With no memory term the state update looks free, and the comparison against + a growing KV cache becomes meaningless. Checked as bytes rather than cycles + because `_roofline` takes the max of the two and the compute side would hide + a wrong byte count at this shape. + """ + hw = load_hardware_config_from_toml(str(ROOT / "plena_settings.toml")) + perf = PerfModel(hw, str(ROOT / "analytic_models" / "performance" / "customISA_lib.json")) + heads, key_dim, value_dim = 32, 128, 128 + + perf.reset_traffic() + perf.kda_recurrence_decode(num_heads=heads, key_dim=key_dim, value_dim=value_dim, batch_size=1) + expected = 2 * heads * key_dim * value_dim * perf.state_bytes + assert perf.traffic_bytes == pytest.approx(expected), ( + "decode must read and write the whole recurrent state exactly once" + ) + + +def test_row_granular_prefetch_is_off_by_default_and_changes_nothing() -> None: + """The hypothetical must not move the calibrated number. + + `row_granular_prefetch` prices an instruction that does not exist. If it + leaked into the default it would move this file's oracle comparison against + a compiler that emits the instruction set that does, and the calibration + would be measuring the wrong machine. + """ + perf = _compiler_shaped_model() + shape = dict(num_heads=1, key_dim=128, value_dim=128, chunk_size=16, seq_len=16, batch_size=1) + assert perf.kda_chunk_prefill(**shape) == perf.kda_chunk_prefill(**shape, row_granular_prefetch=False) + + +def test_row_granular_prefetch_pays_in_proportion_to_machine_width() -> None: + """What the instruction is worth, and the reason it is worth anything. + + The fill it removes scales with `mlen` while the data it protects scales + with `chunk`, so the saving has to grow with width -- a factor that did not + would mean the term was being modelled as something other than a per-row + fill over an `mlen`-tall tile. + + The bounds are loose because the point is the shape, not the digits: nearly + nothing at `mlen` 64, an order of magnitude at 2048. Anyone quoting the wide + figure as "the value of this instruction" should be made to say which + machine they mean. + """ + shape = dict(num_heads=1, key_dim=128, value_dim=128, chunk_size=16, seq_len=16, batch_size=1) + factors = {} + for mlen, blen, hlen, vlen in ((64, 4, 16, 64), (512, 32, 64, 512), (2048, 128, 128, 2048)): + hw = load_hardware_config_from_toml(str(ROOT / "plena_settings.toml")) + hw.MLEN, hw.BLEN, hw.HLEN, hw.VLEN = mlen, blen, hlen, vlen + perf = PerfModel( + hw, + str(ROOT / "analytic_models" / "performance" / "customISA_lib.json"), + enable_bandwidth=False, + ) + perf.instr = _UnitLatency() + whole = perf.kda_chunk_prefill(**shape) + rows = perf.kda_chunk_prefill(**shape, row_granular_prefetch=True) + factors[mlen] = whole / rows + + assert factors[64] < 2.0, f"at mlen 64 it should barely pay; got {factors[64]:.1f}x" + assert factors[2048] > 8.0, f"at mlen 2048 it should pay heavily; got {factors[2048]:.1f}x" + assert factors[64] < factors[512] < factors[2048], f"the saving must grow with machine width, got {factors}" diff --git a/analytic_models/test_perf_model_bandwidth.py b/analytic_models/test_perf_model_bandwidth.py new file mode 100644 index 000000000..c333150b0 --- /dev/null +++ b/analytic_models/test_perf_model_bandwidth.py @@ -0,0 +1,218 @@ +"""Regression tests for the perf_model bandwidth term, precision parsing and +the two KV-store bugs it uncovered. + +Run: python3 analytic_models/test_perf_model_bandwidth.py +(or with pytest; every check is a plain assert in a test_* function) +""" + +import math +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +PERF_DIR = REPO_ROOT / "analytic_models" / "performance" +sys.path.insert(0, str(PERF_DIR)) + +from perf_model import ( # noqa: E402 + PerfModel, + _entry_bytes_per_elem, + load_hardware_config_from_toml, + parse_precision_bytes, +) + +CONFIG = str(REPO_ROOT / "plena_settings.toml") +ISA = str(PERF_DIR / "customISA_lib.json") + + +def _model(enable_bandwidth=True): + return PerfModel(load_hardware_config_from_toml(CONFIG), ISA, enable_bandwidth=enable_bandwidth) + + +# ----------------------------------------------------------------------------- +# Precision parsing +# ----------------------------------------------------------------------------- + + +def test_mx_bytes_per_elem_includes_shared_scale(): + """MXFP8 block=8 with an 8-bit shared scale is 1.125 B/elem, not 1.0.""" + entry = { + "format": "Mx", + "block": 8, + "ELEM": {"type": "Fp", "sign": True, "exponent": 4, "mantissa": 3}, + "SCALE": {"type": "Fp", "sign": False, "exponent": 8, "mantissa": 0}, + } + assert _entry_bytes_per_elem(entry) == (8 * 8 + 8) / (8 * 8) + assert _entry_bytes_per_elem(entry) == 1.125 + + +def test_plain_fp_and_int_bytes_per_elem(): + fp = {"format": "Plain", "DATA_TYPE": {"type": "Fp", "sign": True, "exponent": 8, "mantissa": 7}} + i32 = {"format": "Plain", "DATA_TYPE": {"type": "Int", "width": 32}} + assert _entry_bytes_per_elem(fp) == 2.0 # bf16 + assert _entry_bytes_per_elem(i32) == 4.0 + + +def test_precision_section_parsed_from_toml(): + """[ANALYTIC.PRECISION.*] used to be ignored entirely.""" + cfg = load_hardware_config_from_toml(CONFIG) + pb = cfg.PRECISION_BYTES + assert pb, "PRECISION section was not parsed" + for name in ("HBM_M_WEIGHT_TYPE", "HBM_M_KV_TYPE", "HBM_V_ACT_TYPE", "HBM_V_KV_TYPE"): + assert pb[name] == 1.125, f"{name} = {pb[name]}" + assert pb["HBM_V_INT_TYPE"] == 4.0 + assert pb["MATRIX_SRAM_TYPE"] == 2.0 + assert parse_precision_bytes({}) == {} + + +# ----------------------------------------------------------------------------- +# Bandwidth term +# ----------------------------------------------------------------------------- + + +def test_bandwidth_derived_from_hbm_width(): + m = _model() + assert m.hbm_bytes_per_cycle == 512.0 # HBM_WIDTH bytes/cycle + assert m.mem_cycles(512) == 1 + assert m.mem_cycles(513) == 2 + assert m.mem_cycles(0) == 0 + + +def test_bandwidth_override_wins(): + cfg = load_hardware_config_from_toml(CONFIG) + cfg.HBM_BANDWIDTH_BYTES_PER_CYCLE = 64.0 + m = PerfModel(cfg, ISA) + assert m.hbm_bytes_per_cycle == 64.0 + + +def test_stage_takes_max_not_sum(): + """A stage costs max(compute, memory), never the sum.""" + m = _model() + m.reset_traffic() + got = m._roofline(100, 512 * 1000) # 1000 memory cycles vs 100 compute + assert got == 1000 + got = m._roofline(5000, 512) # 1 memory cycle vs 5000 compute + assert got == 5000 + assert m.memory_bound_stages == 1 + assert m.total_stages == 2 + + +def test_enable_bandwidth_false_is_backward_compatible(): + """The compute-only path must be untouched by the bandwidth work.""" + compute = _model(enable_bandwidth=False) + full = _model(enable_bandwidth=True) + args = dict(hidden_size=4096, intermediate_size=14336, seq_len=1, batch_size=1, mode="decode") + c = compute.feed_forward(**args) + f = full.feed_forward(**args) + assert f >= c, "bandwidth model must never make a stage cheaper" + assert f > c, "llama decode FFN is expected to be memory-bound at batch 1" + + +def test_ffn_decode_is_memory_bound_and_matches_hand_calc(): + m = _model() + m.reset_traffic() + m.feed_forward(4096, 14336, 1, 1, "decode") + expected_bytes = 3 * 4096 * 14336 * 1.125 + assert m.traffic_bytes == expected_bytes + assert m.memory_only_cycles == math.ceil(expected_bytes / 512) + + +# ----------------------------------------------------------------------------- +# Bug 1: floor division charged ZERO cycles for the decode KV cache write +# Bug 2: a store used HBM_V_Prefetch_Amount instead of HBM_V_Writeback_Amount +# ----------------------------------------------------------------------------- + + +def test_decode_kv_store_is_no_longer_free(): + """Reproduces the old formula and shows it evaluated to 0 at the default config.""" + m = _model() + batch_size, num_kv_heads, head_dim = 4, 8, 128 + + old_bursts = (batch_size * num_kv_heads * head_dim) // (m.vlen * m.prefetch_v_amount) + assert (batch_size * num_kv_heads * head_dim) == 4096 + assert (m.vlen * m.prefetch_v_amount) == 32768 + assert old_bursts == 0, "precondition: the old floor division really did give 0" + + new_bursts = math.ceil((batch_size * num_kv_heads * head_dim) / (m.vlen * m.writeback_v_amount)) + assert new_bursts == 1, "a decode cache write must cost at least one burst" + + +def test_store_uses_writeback_amount(): + cfg = load_hardware_config_from_toml(CONFIG) + assert cfg.HBM_V_Writeback_Amount == 16, "config value must be read, not ignored" + m = PerfModel(cfg, ISA) + assert m.writeback_v_amount == cfg.HBM_V_Writeback_Amount + # And it must be a distinct knob: changing it changes the store cost. + cfg2 = load_hardware_config_from_toml(CONFIG) + cfg2.HBM_V_Writeback_Amount = 1 + m2 = PerfModel(cfg2, ISA, enable_bandwidth=False) + m1 = PerfModel(cfg, ISA, enable_bandwidth=False) + a = m1.projection(4096, 32, 8, 128, 2048, 4, "prefill") + b = m2.projection(4096, 32, 8, 128, 2048, 4, "prefill") + assert b > a, "shrinking the writeback burst must raise the store cost" + + +# ----------------------------------------------------------------------------- +# New ISA entries +# ----------------------------------------------------------------------------- + + +def test_mamba_isa_opcodes_present(): + m = _model() + for op in ("V_SOFTPLUS_V", "S_MAP_FP_V", "V_MAX_VF", "V_MIN_VF", "V_PS_V", "V_SHFT_V"): + assert op in m.instr, f"{op} missing from customISA_lib.json" + assert m.instr[op] > 0 + # V_SOFTPLUS_V is budgeted above a bare exp: it also evaluates a logarithm. + assert m.instr["V_SOFTPLUS_V"] > m.instr["V_EXP_V"] + + +# ----------------------------------------------------------------------------- +# Mamba-2 state behaviour +# ----------------------------------------------------------------------------- + + +def test_ssm_decode_state_traffic_is_context_independent(): + """The whole point: decode traffic does not grow with how much came before.""" + m = _model() + m.reset_traffic() + m.ssd_recurrence_decode(num_heads=80, head_dim=64, state_size=128, n_groups=1, batch_size=1) + first = m.traffic_bytes + expected = 2 * 80 * 64 * 128 * 1.125 # read + write + assert first == expected + + # a thousand tokens later, the per-token cost is identical + m.reset_traffic() + m.ssd_recurrence_decode(num_heads=80, head_dim=64, state_size=128, n_groups=1, batch_size=1) + assert m.traffic_bytes == first + + +def test_attention_decode_traffic_grows_with_context(): + """Contrast case: the KV cache read scales linearly with kv_size.""" + m = _model() + m.reset_traffic() + m.flash_attention(32, 8, 128, 1, 512, 1, "decode") + short = m.traffic_bytes + m.reset_traffic() + m.flash_attention(32, 8, 128, 1, 8192, 1, "decode") + long = m.traffic_bytes + assert long == short * 16, f"{long} vs {short}" + + +def _run_all(): + failures = [] + for name, fn in sorted(globals().items()): + if name.startswith("test_") and callable(fn): + try: + fn() + print(f"PASS {name}") + except AssertionError as exc: + failures.append((name, exc)) + print(f"FAIL {name}: {exc}") + if failures: + print(f"\n{len(failures)} failure(s)") + return 1 + print("\nall tests passed") + return 0 + + +if __name__ == "__main__": + sys.exit(_run_all()) diff --git a/doc/static_path_measurements.md b/doc/static_path_measurements.md new file mode 100644 index 000000000..f58127ccb --- /dev/null +++ b/doc/static_path_measurements.md @@ -0,0 +1,1003 @@ +# Static Mamba/KDA path — measurements + +Everything here was measured on `feat/static-kda`. The two kinds of number are +kept apart on purpose, because only one of them decides anything. + +* **Hard facts** — static instruction counts, dynamic instruction counts, HBM + bytes, memory footprints. Properties of the compiled artifact. Reproducible, + and admissible as CI gates. +* **Uncalibrated model output** — cycles, µs, TPOT. Produced by a timing model + that has never been calibrated against silicon. Labelled as such on every + line, and **not compared against any other uncalibrated model output**: two + numbers from two uncalibrated models decide nothing between them. + +There is a second split inside the hard facts, and this document got it wrong +four times before 2026-08-27. A **static** count is the size of the program +image. A **dynamic** count is the number of instructions issued, with every +hardware loop expanded by its trip count. Both are exact and reproducible; +only the second says anything about cost. A sweep the `V_FMA_VF` conversion +collapsed into a hardware loop is a handful of words in the image and runs its +body up to 192 times. + +Every claim here about what something *costs*, about which kernel dominates, +or about what an alternative lowering would save is now made against the +dynamic count. The four that were not: + +| claim | from the image | from the issue stream | +|---|---|---| +| projection gather, share of a KDA layer | 0.07% | **0.94%** | +| convolutions vs mixer, share of a layer | 58% / 42% | **12% / 87%** | +| what `V_FMA_VF` bought | 97× | **3.6×** | +| Kimi K3 head vs `64 × 64`, for 4× state | 1.6× | **2.3×** | + +The image counts are not retracted — they are the right instrument for the CI +gates, which exist to catch a sweep falling off the hardware-loop path, and +that is an image failure no numeric test can see. What is retracted is reading +any of them as a cost. + +--- + +## 1. Hard facts + +### The instructions that were added, and what they bought + +Three, all ordinary fixed-function ops with their operands named in the +instruction word: + +| | | | +|---|---|---| +| `V_SOFTPLUS_V` | `0x39` | `log(1 + exp(x))` over a vector row, for Mamba's `dt` | +| `S_MAP_FP_V` | `0x3A` | one VRAM row into `VLEN` consecutive FPRAM slots | +| `V_FMA_VF` | `0x3B` | `Vector[rd] += Vector[rs1] * fp_reg` | + +`V_FMA_VF` is the one the numbers below are about. It replaced the +`copy + multiply + add` triple that both recurrences used. + +Instructions for one head, image and issue stream side by side. The image +columns are the ones this document used to carry alone: + +| | image before | image after | issued after | +|---|---|---|---| +| KDA decode, `key_dim` 8 | 505 | **76** | 174 | +| KDA decode, `key_dim` 16 | 961 | **76** | 302 | +| KDA decode, `key_dim` 64 | 3,697 | **76** | 1,070 | +| KDA decode, `key_dim` 128 | 7,345 | **76** | 2,094 | +| Mamba decode, `state_size` 4 | 208 | **54** | 77 | +| Mamba decode, `state_size` 8 | 376 | **54** | 121 | +| Mamba decode, `state_size` 16 | 712 | **54** | 209 | +| Mamba decode, `state_size` 32 | 1,384 | **54** | 385 | + +The reason the image collapsed is not the saved multiply. The triple staged +every row through a scratch row, and because it was the *same* row each +iteration the destination never formed an arithmetic progression — so the sweep +could not become a hardware loop and was unrolled in full. + +**The cost did not go from linear to constant.** This document said it did, and +put a factor of 97 on it; both statements are about the image. Expanded, the two +kernels are still exactly linear in the contracted dimension, which they have to +be — a recurrence touches every state row and an encoding cannot change how many +rows there are: + +``` +KDA: issued = 16 * key_dim + 46 +Mamba: issued = 11 * state_size + 33 +``` + +Both fit their four measured points exactly. The pre-conversion slopes were +about 57 and 42, so **what `V_FMA_VF` bought is a 3.6× smaller coefficient on a +line that stays a line** — together with an image that stops growing, which is +what makes a whole-model program image tractable and is worth having on its own +terms. + +Both facts are real and they are not the same fact. 97× is the image; 3.6× is +the time. `test_the_conversion_reduced_the_slope_not_the_linearity` in +`test_instruction_budget.py` pins the fits and asserts the two never converge +again. + +### One KDA layer, Kimi K3 shape, `mlen` 64 + +| | static (image) | dynamic (issued) | dynamic share | +|---|---|---|---| +| three convolutions | 53,757 | 59,409 | 12.1% | +| mixer (gates + normalisation + recurrence) | 39,526 | 428,622 | 87.0% | +| projection gather | 70 | 4,650 | 0.94% | +| **layer** | **93,353** | **492,681** | | +| **× 93 layers** | **8,681,829** | **45,819,333** | | + +Two columns because they answer different questions and were conflated here +until 2026-08-27. **Static** is the program image, and it is what the +`V_FMA_VF` conversion bought — an image that does not grow with `key_dim`. It +is the right number for a budget gate against accidental unrolling, and the +wrong number for pricing a kernel, because a hardware loop whose body runs 192 +times costs 192 times what the image shows. + +The larger correction is the ordering. The convolutions are 58% of the image and +**12% of the work**; the mixer is 42% of the image and **87% of the work**. This +document and the compiler's budget file both called the convolutions "the larger +half of the layer" on the strength of the image, which points optimisation effort +at the smaller kernel by a factor of seven. The cause is structural: the +convolutions are a four-tap FIR whose taps are separate weights, so the image +carries an instruction per tap per block and almost nothing loops, while the +mixer is the recurrence the `V_FMA_VF` conversion collapsed into hardware loops +running once per key block per head — a small image over a large trip count. + +The gather is the case where the difference bites: 70 static is 0.07% of the +layer, and the earlier version of this table reported that as the gather's +cost. **The gather issues 4,650 instructions, 0.94%** — thirteen times more. +The conclusion the 0.07% was reaching for still holds at 0.94% (no mechanism +that only makes the split faster can move the layer by more than a percent), +but 0.94% is a cost, not a rounding error, and it is the number any alternative +must be compared against. + +### The gather is not forced by packing, and is not forced at all + +The 0.94% invites the question of whether a different projection lowering +avoids it: five separate projections, each writing straight to where its +consumer wants it, instead of one packed tile plus a split. Both were built and +counted, at `mlen` 64, `blen` 4, through `linear_projection`: + +| | matmul static | matmul issued | `M_MM` | gather issued | +|---|---:|---:|---:|---:| +| packed | 1,090,417 | 13,876,267 | 21,560 | 4,650 | +| separate | 1,107,395 | 13,893,245 | 21,560 | 4,650 | + +Identical matmul work, 0.12% worse for the separate form from the extra +per-projection setup, and the same gather. **There is no crossover in `blen` to +find** — the premise was wrong. + +The reason is the writeback granularity. `M_MM_WO` writes a `blen × blen` +sub-tile and the loops around it cover `mlen / blen` column groups, so the +smallest thing a projection can lay down is `blen` token-rows by `mlen` lanes: +column block `c` lands at row `c * blen` however the weight matrices are +arranged. The consumers want one token's blocks as consecutive rows. **The +mismatch is exactly `blen`, and it is a property of the matrix writeback rather +than of the packing**, so no rearrangement of weights touches it. + +What removes it is not moving the data. Consumers take explicit row indices, +any arithmetic progression collapses into one hardware loop, and the step is an +`S_ADDI_INT` immediate — so a step of `blen` costs what a step of 1 costs. +`kda_conv_step_v0` now takes `x_new_row_base` / `x_new_row_stride`, and feeding +it the tall projection view removes **4,041 of the 4,650** issued instructions, +the q/k/v share, for no instructions and no hardware. Gate and beta reach a +different consumer and still gather. + +What stays is gate and beta: **1,176 issued, 0.24% of the state-engine path.** +They reach `kda_decay_scalars_v0` and `kda_beta_scalars_v0` rather than the +convolution, and the two are not symmetric. beta is single-tile work throughout +— sigmoid in place, then `tile_row_to_fpram`, which addresses FPRAM by position +in the row list rather than by row index — so a stride drops straight in, for 18 +instructions. gate opens with `tile_row_add(gate, dt_bias)`, and underneath +`_emit_tile_row_vector_op` derives one address progression and steps both +operands by it; a strided gate against a dense `dt_bias` needs that emitter to +take `(dst, src)` row pairs the way `_emit_tile_row_fma` already takes triples. +That is a change to a row-op emitter every kernel on the branch shares, for +0.24%. Laying `dt_bias` out at the same stride instead needs no ISA change but +pads it from 192 rows to 768 — 74 KiB per layer, 6.8 MiB more HBM traffic per +token across 93 layers, to save 1,176 instructions. Both trades are bad, so +they stay and the test records the price. + +### The input projection, which no budget covered + +Building it turned up the largest kernel in the layer and the fact that nothing +gated it — `test_kda_layer.py` loads the projection's result instead of +computing it, so it was never in the lowering. At Kimi K3, `mlen` 64, and the +capacity the shipped transactional config implies: + +| | static | issued | +|---|---:|---:| +| input projection | 389,717 | 4,967,367 | +| state-engine path (conv + mixer + gather) | 93,353 | 492,681 | + +**Those two rows must not be added, ranked, or turned into a percentage.** An +`M_MM` is a `64 × 64 × 64` matmul and a `V_FMA_VF` is a 64-lane vector +operation; counting them in one unit says nothing about time, and doing so +would be the same error as §1's first four. The row is here because a kernel +that size should not sit outside every budget, and it now has one. Comparing +the gather against the state-engine path stays legitimate — both are vector +work — and that comparison is the 0.94% above. + +A second thing fell out. **The compiler's `mram_tile_capacity` defaults to 4 +and the configured machine has 64.** Nothing derives one from the other: +`MatrixSram::new(tile_size = MLEN, depth = MATRIX_SRAM_SIZE)` keeps +`depth / tile_size` tiles, so MLEN 64 over `MATRIX_SRAM_SIZE` 4096 is 64 tiles, +while `PlenaCompiler` takes `mram_tile_capacity: int = 4` from its own +signature. Kimi K3's projection contracts over 112 k-tiles and +`linear_projection` re-streams the weights once per capacity-sized chunk: + +| capacity | | static | issued | +|---|---|---:|---:| +| 4 | 32 KiB — compiler default | 1,090,417 | 13,876,267 | +| 16 | 128 KiB | 524,467 | 6,680,617 | +| 64 | 512 KiB — configured machine | 389,717 | 4,967,367 | +| 112 | 896 KiB — whole contraction | 362,767 | 4,624,717 | + +**A caller who did not name a capacity got 2.79× the projection the configured +machine would run**, and nothing said so. Fixed: the default is now +`MATRIX_SRAM_SIZE // mlen`, the same arithmetic on the same two numbers as the +emulator's `MatrixSram::new`. + +Chasing that turned up why the configuration had never reached the compiler at +all. `PlenaCompiler`'s docstring said HLEN, BROADCAST_AMOUNT and the prefetch +amounts default to `plena_settings.toml`; none of them ever did, for two +independent reasons at once: + +1. The reader asked for `[BEHAVIOR].CONFIG`, and the file has no such table — + it ships `[MODE]`, `[ANALYTIC.*]` and `[TRANSACTIONAL.*]`. +2. `load_toml_config` required the third-party `toml` package and raised + `ImportError` without it. `toml` is not a declared dependency and CI + installs `pytest pyyaml`, so it failed there too — and the `except + Exception` around the call turned "cannot read the config" into "the default + is correct". + +Either alone would have been enough, and together they made a silent failure +that looked exactly like a working default. The reader now goes through +`tomllib`, and searches `BEHAVIOR` first — so a file that does define one still +wins — then the section whose `MLEN` is the one being compiled for. + +**Matching on `MLEN` rather than on `[MODE].active` is deliberate.** `active` +selects the machine the *simulator* models, not the one the compiler emits for, +and the two differ today: the shipped file is `analytic` (MLEN 2048) while every +program on this branch is compiled at 64 and executed by the transactional +emulator, whose section declares 64. Reading `active` would hand a 64-wide +compilation the 2048-wide machine's numbers. + +At `mlen` 64 the fix also moves HLEN from 64 to 16 and BROADCAST_AMOUNT from 1 +to 4. Neither reaches an emitted program — `hlen` is read only by a +packed-attention precondition whose one test sets both attributes by hand at +`mlen` 256, and `broadcast_amount` is read nowhere outside the constructor. The +prefetch and writeback amounts do reach the emitters, and are 4 in that section, +which is what they already were. Every count in this document is unchanged: the +state-engine path does not go through MRAM matmuls. + +At an `mlen` no section declares — the 8 the budget tests use — the capacity +stays 4 and nothing changes. At `mlen` 2048 the `ANALYTIC` config gives +`256 / 2048 = 0` tiles and the constructor now raises naming both figures, +rather than substituting a number nobody chose: a contraction split into +zero-tile chunks does not terminate. Nothing reads that mode today, but the +number in the file is still not a capacity anyone picked. + +This is the **state-engine path alone** — the projections, MoE blocks, norms and +embeddings are not in it. Any figure that includes those is measuring a +different thing and the two must not be put side by side. + +### Where the issue slots go, and what an addressing mode would buy + +The layer table above says a Kimi K3 layer issues 492,681 instructions. Per +opcode, **124,428 of them compute and 368,253 do not** — 25.3%. + +The shape of the recurrent sweep is what does it. Every `V_FMA_VF` sits in a +body of five, of which one computes: + +``` +S_LD_FP f1, gp3, 0 <- the scalar for this row +V_FMA_VF gp1, gp2, f1, 0 <- the only instruction that computes +S_ADDI_INT gp1, gp1, mlen <- advance the destination +S_ADDI_INT gp2, gp2, mlen <- advance the source +C_LOOP_END gp4 <- not counted; zero-overhead loop boundary +``` + +| | conv ×3 | mixer | gather | **layer** | +|---|---:|---:|---:|---:| +| issued | 59,409 | 428,622 | 4,650 | **492,681** | +| arithmetic | 18,432 | 104,456 | 1,540 | **124,428** | +| share | 31.0% | 24.4% | 33.1% | **25.3%** | + +That 25% is the figure a dedicated recurrence engine beats by construction — +one instruction covering a whole head has no instruction overhead at all. So it +is worth knowing how much of the gap is *encoding* rather than *architecture*. + +A **post-increment operand** — the pointer carrying its own stride, the way +`_emit_tile_row_fma` already carries independent row progressions — folds +`S_ADDI_INT gpN, gpN, step` into the instruction that consumes the pointer. It +only works where the register has a consumer in the same loop body, so +`self_advance_counts` splits them and prices the saving against the ones that +can really disappear. **Every one of them qualifies: 215,634 foldable, 0 not.** +86% fold into a single consumer; the rest are read more than once in the body +and would need the increment on the last use. + +| | issued | arithmetic share | vs today | +|---|---:|---:|---:| +| today | 492,681 | 25.3% | — | +| **P1** post-inc on the vector operands | 277,047 | 44.9% | **1.78×** | +| **P2** P1 + the FPRAM scalar auto-advancing | 174,129 | 71.5% | **2.83×** | + +Over 93 layers: 45.8M issued instructions falling to 16.2M. + +P2 is an addressing field on instructions that already exist. What it competes +with is a dedicated recurrence datapath with its own banked state SRAM plus a +layout engine to feed it — the two are not close in cost, though neither has +been synthesised so neither has an area number. + +**What neither proposal touches**, and this is the limit that matters: KDA's +decay is channel-wise on the key axis, so the operand is `key_dim` lanes wide — +6.2% of `VLEN` 2048 (§5). That is a **width** mismatch, not an encoding one, +and no addressing mode reaches it. The two inefficiencies are independent and +they multiply: at `VLEN` 2048, 71.5% of instructions computing on 6.2% of the +lanes is 4.4% of the machine doing work; at a `VLEN` that matches the operand +it is 71.5%. **That is the one place a purpose-built engine wins on +architecture rather than on encoding**, and an addressing mode does not answer +it — choosing the width, or changing the state layout, would. + +### Memory + +| | | +|---|---| +| FPRAM, one KDA layer, Kimi K3 | **492 slots** of 512 (`FP_SRAM_DEPTH`) | +| FPRAM, 93 layers | **492** — it is scratch, reused per layer | +| HBM, per layer | **2,336,832 bytes (2,282 KiB)** | +| HBM, 93 layers | **217,325,376 bytes (0.202 GiB)** | +| emulator `HBM_SIZE` | 16 GiB | + +The footprint fits the emulator's flat allocation with three orders of magnitude +to spare, which is why the program executes rather than being described. That is +a property worth stating on its own: the emulator preloads HBM from a flat file +starting at offset 0, so a layout whose regions are spread across a wide address +space needs an allocation and a file as large as the span, whatever the live +data comes to. + +The HBM figure counts everything a layer stages — conv history and weights, the +recurrent state, `dt_bias` and the projection tile — and excludes only the +projection *weights*, which any implementation needs and which are +model-inherent. A first version of this number was 331,776 bytes, seven times +too low, because the measuring helper put the state and the projection in VRAM +while the emulator test staged them from HBM. The two now agree. + +The FPRAM number is tight and was nearly wrong: the first accounting came to 620 +slots because `beta` was counted as one slot when `S_MAP_FP_V` writes a padded +row (128 at Kimi), and `rate` was omitted entirely. `decay` and `q_hat` are never +live at the same time, and sharing one window is what brings it to 492. + +### Numerical agreement, on the transactional emulator + +Against the CPU reference, compiled and assembled and executed as machine code: + +| | max abs error | +|---|---| +| cumulative decay | 1.95e-03 (one bf16 ulp) | +| UT transform | 1.95e-03 (one bf16 ulp) | +| prefill, one chunk of 16 | 1.13e-03 out / 1.32e-02 state | +| prefill, three chunks chained | 6.71e-04 out / 8.30e-03 state | +| state layout transpose | **0.00** | +| **assembled layer** | **2.14e-04 abs / 4.3% mean relative** | +| **four layers back to back** | **2.75e-04 abs / 4.4% mean relative** | + +The layer figures are quoted both ways on purpose. Its outputs are O(5e-3), so +an absolute error of 2e-4 is **4% relative** — and an earlier version of this +table quoted only the absolute number against a test whose `atol` was 2e-2, +which every golden value already satisfied. That test passed on an all-zero +output tile. It is relative now, and 4.3% is what bf16 delivers on the +read-out's 128-term contraction (`sqrt(128) * 2^-9 = 2.2%` before the rest of +the chain). The lowering itself is right to **5e-4 in float64**, which the ISA +interpreter checks separately — the 4% is precision, not a defect. + +--- + +## 2. The banking question + +**Answer: stalls ≈ 0. The vector path is limited by instruction issue, not by +SRAM.** + +`analytic_models/performance/vector_sram_banking.py` extracts every hardware +loop's pointer stride from the *emitted assembly* — not from a description of +the algorithm — and models row-interleaved banking, where row `r` lives in bank +`r % B` and a stride-`s` sweep visits `B / gcd(s, B)` distinct banks. + +| kernel | loops | B=2 | B=4 | B=8 | B=16 | B=32 | +|---|---|---|---|---|---|---| +| KDA recurrence (Kimi K3) | 20 | 0.6% | 0.9% | 1.0% | 1.1% | 1.1% | +| KDA mixer, one head | 48 | 1.2% | 1.8% | 2.1% | 2.3% | 2.3% | +| Mamba-2 decode step | 7 | 1.0% | 1.5% | 1.8% | 1.9% | 2.0% | + +The dominant loops are **stride 1 over 128 trips** (KDA) and stride 1 over 64 +(Mamba). A stride-1 walk visits every bank in turn, so it never conflicts. The +1–2% residue is entirely single-trip loops with a pinned pointer. + +So there is no evidence for a Vector SRAM bank-mapping mode. Note also that the +emulator does not model banking at all — `VectorSram` is a flat row vector — so +this is analytic, and a word-interleaved mapping would make the question moot +regardless. + +### What the follow-up actually is + +Half the dynamic instruction stream is pointer arithmetic: + +| kernel | dynamic | `S_ADDI_INT` | `S_LD_FP` | `V_FMA_VF` | +|---|---|---|---|---| +| KDA recurrence (Kimi K3) | 4,179 | **50.0%** | 24.6% | 18.4% | +| KDA mixer, one head | 4,398 | **50.2%** | 23.9% | 17.5% | +| Mamba-2 decode step | 730 | **46.4%** | 26.6% | 17.7% | + +Under a fifth of the stream is arithmetic. Every hardware loop advances its +pointers with an explicit `S_ADDI_INT` per pointer per iteration, and reloads its +FPRAM scalar with `S_LD_FP`. **Post-increment addressing on the vector operand +pointers, and an auto-advancing FP pointer, would together remove up to +three-quarters of the dynamic instructions on the recurrent path.** That is the +next instruction-set question, and it is a much larger effect than banking. + +--- + +## 3. Uncalibrated model output + +**Every number in this section comes from a timing model that has never been +calibrated against silicon. They are not performance claims.** + +| | UNCALIBRATED | +|---|---| +| KDA chunked prefill, one 16-token chunk, one head | 173,939 cycles | + +No `µs` or TPOT figure is given, because converting cycles to time requires a +clock the model does not have, and quoting one would give the number a +credibility it has not earned. + +--- + +## 3b. Prefill past one block: `key_dim` and `value_dim` above `mlen` + +Chunked prefill used to refuse `key_dim > mlen` outright, which is why Kimi K3's +128-wide head could not prefill at `mlen` 64. It now runs, and so does +`value_dim > mlen`, which nothing had ever exercised. + +### What changed + +The key axis moved from **rows** onto **lanes**. A `[chunk, key_dim]` tile is now +`chunk` rows spread across `key_dim / mlen` column blocks, instead of +`chunk * key_blocks` rows one block wide. + +That is what the projections want. Five of the seven products contract over key, +and `vram_sub_projection_asm_impl` already walks a VRAM operand's column blocks +at stride `physical_rows * mlen`, accumulating inside the systolic array -- so a +key axis on lanes contracts across as many blocks as it needs with **no new +instruction and no explicit sum**. No projection call changed. The row ops pay +instead, one sweep per column block rather than one over the folded rows, and +rows within a block are still consecutive so each sweep is still a hardware loop. + +### Measured on the transactional emulator, as machine code + +Every case at 100% of values inside tolerance -- see the note on the 90% rule +below for why that phrasing matters. + +| case | `key_dim` × `value_dim` | max abs error | +|---|---|---| +| cumulative decay | 64 × 64 | 1.95e-03 | +| cumulative decay | 128 × 64 | 1.95e-03 | +| state layout transpose | 64 × 64 | **0.00** | +| state layout transpose | 128 × 128 | **0.00** | +| prefill, out | 64 × 64 | 1.13e-03 | +| prefill, out | 128 × 64 | 6.83e-04 | +| prefill, out | 128 × 128 | 6.71e-04 | +| prefill, state | 64 × 64 | 1.32e-02 | +| prefill, state | 128 × 64 | 5.13e-03 | +| prefill, state | 128 × 128 | 9.28e-03 | +| three chunks chained, out | 128 × 128 | 6.10e-04 | +| three chunks chained, state | 128 × 128 | 7.32e-03 | + +`128 × 128` is Kimi K3's head at `mlen` 64: two key blocks and two value blocks. + +### Instructions, one head, `mlen` 64 + +Image, then issue stream: + +| chunk | 64 × 64 | 128 × 64 | 128 × 128 | +|---|---|---|---| +| 4 | 818 / 18,904 | 1,237 / 25,898 | 1,470 / 43,851 | +| 8 | 1,039 / 19,246 | 1,547 / 26,366 | 1,792 / 44,343 | +| 16 | **1,476** / **20,120** | 2,157 / 27,490 | **2,426** / **45,515** | + +The single-block column is **unchanged** by the layout move (1,476 against a +1,500 measurement before it), so lifting the limit cost nothing at the shape that +already worked. + +Kimi K3's head against a `64 × 64` one, for 4× the state: **1.6× by image, +2.3× by issue stream**. An earlier version of this note gave only the first and +read it as the cost of quadrupling the state. The kernel is sublinear in the +state either way, which was the point being made, but by 2.3× rather than 1.6×. + +The `chunk` column carries a second divergence worth having. Going from chunk 4 +to chunk 16 costs **1.65× of the image and 1.04× of the issue stream** — four +times the tokens per chunk for 4% more work. Prefill is spill-dominated, so the +per-chunk cost is the fixed traffic of filling and storing `mlen`-sized tiles +and barely moves with how many tokens ride along. That makes the `chunk` cap of +17 — set by bf16 range on `1/A`, not by the lowering, see §5 — more expensive +than the image suggests: every token that cannot join a chunk pays the fixed +cost again. `test_kda_prefill_structure.py` pins both ratios. + +### One real defect, and the tolerance that nearly hid it + +`out = scale * (...)` ended with a single `tile_row_mul_fp_broadcast`, and a row +op reaches **one column block per call**. At `value_dim` 128 that left the upper +half of every token unscaled -- `sqrt(key_dim)` = 11.3x too large. + +It reported **PASSED**. `check_mem`'s `allclose_pass` is `match_rate >= 90.0`, +and 94.68% of the values sat inside the case's `atol` of 5e-2 because the data is +of order 1e-3: a tolerance fifty times the signal. Splitting the comparison by +value block is what showed it -- block 0 at 6.4e-04, block 1 at 1.4e-01, a mean +relative error of 11.28 against a predicted 11.31. + +The stage test now restates the bar at **100%** on the same numbers +(`_assert_every_value_within_tolerance`). `check_mem` is shared by many callers +and its 90% rule is left alone. Under the new bar, four mutations of the +column-block loops die: + +| mutation | match rate | would the 90% rule have caught it? | +|---|---|---| +| output scale on block 0 only *(the real defect)* | 94.68% | **no** | +| `k_hat`/`q_hat` on block 0 only | 0.00% | yes | +| `A_C * state` on block 0 only | 55.79% | yes | +| `k / A` on block 0 only | 85.80% | yes | + +Only one of the four needed the stricter bar, and it is the one that actually +happened. + +The layer cases moved to `atol` 2.5e-04 with the same `rtol` 0.12 for the same +reason from the other side: three of their sixty-four outputs are near 1e-03, +bf16 delivers about 1.8e-04 absolute whatever the value, and a **pure relative** +bound cannot be met there. They had been passing at 95.31% on the 90% rule rather +than on the tolerance. The floor keeps the check real: an all-zero output scores +**1.6%** against it. + +### Still refused + +`key_dim` and `value_dim` must each be a whole multiple of `mlen` +(`kda_key_blocks` and `kda_blocks` enforce it), and `chunk` must still fit one +block. The state transpose additionally requires its output to be exactly +`key_dim` rows when there is more than one value block, because that is what +makes its column-block layout coincide with the row indexing decode uses. + +--- + +## 4. What is not measured + +* **No real checkpoint.** Every number above uses synthetic weights. Binding to + published Kimi K3 or Nemotron-3 weights needs a model config that this + repository does not have. +* **No whole model.** The projections, MoE blocks, norms and embeddings are + shared code, unchanged by this work and not instantiated here. "One layer × + 93" is a per-layer figure multiplied out, not a compiled 93-layer program. +* **No RTL.** `PLENA_RTL`'s `operation.svh` stops at `0x34`. `V_FMA_VF` joins six + opcodes (`0x35`–`0x3A`) that the compiler and emulator implement and the + hardware does not. +* **The exp model.** The emulator uses libtorch's exact `exp`; + `PLENA_Tools/plena_quant` implements a range reduction plus a 3-term Taylor + with systematic 0.5–1.5% relative error. KDA applies `exp` on the critical path + of a multiplicative recurrence, where that compounds. Passing here bounds the + lowering, not the silicon. + +--- + +## 5. Hybrid models: cost per operator + +`analytic_models/performance/hybrid_model.py` reads a per-layer type list out of +a model config and dispatches each layer to the operator it actually runs, then +reports **FLOPs, cycles and HBM bytes per operator**. The three sibling drivers +cannot do this: each assumes one block shape repeated `num_hidden_layers` times, +which is the line `overall_exe_cycle += block_cycles * self.num_hidden_layers`. + +### The configs are real, and were fetched rather than written + +| `doc/Model_Lib` entry | source | layers | +|---|---|---| +| `kimi-linear-48b-a3b` | `moonshotai/Kimi-Linear-48B-A3B-Instruct` | 20 KDA + 7 attention, 1 dense MLP + 26 MoE | +| `nemotron-3-nano-4b` | `nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16` | 21 Mamba + 4 attention + 17 MLP | +| `nemotron-3-super-120b-a12b` | `nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4` | 40 Mamba + 8 attention + 40 MoE | + +Verbatim, except that the Super's 5.7 MB `quantization_config` (per-tensor NVFP4 +scales) is stripped. Both layer-type encodings were read out of each model's own +source rather than inferred: NemotronH's `pattern_mapping = {"M": "mamba", +"E": "moe", "*": "attention"}` with `-` for MLP, and Kimi's `is_kda_layer` being +`(layer_idx + 1) in kda_layers`, one-based against a zero-based index. The +parser refuses a config it cannot read rather than falling back to a uniform +stack. + +**`KdaShape.kimi_k3()` does not match any of these.** It is `hidden 7168, 96 +heads, 93 layers`; the public Kimi Linear is `hidden 2304, 32 heads, 27 layers`. +The *per-head* shape does match — `head_dim` 128 and `short_conv_kernel_size` 4 +are exactly what the real config carries — so the kernel-level measurements in +sections 1 and 3b stand. Everything in this repository labelled "Kimi K3" at +model scale is a synthetic shape. + +### The KDA cost model is calibrated against the compiler + +Every other stage in `perf_model.py` is a formula with no oracle. `kda_chunk_prefill` +and `kda_recurrence_decode` have one: set every opcode latency to 1 and the +formula counts instructions issued, which the compiler can be asked for directly. + +| | modelled / measured | +|---|---| +| prefill, nine shapes (chunk 4-16, key and value 64-128) | **0.87 - 1.14** | +| decode, one step at 128 x 128 | **0.98** | + +Calibrating found a real error rather than confirming the formula. The model +scaled with `key_dim` at half the compiler's rate, worst at `key_dim 128, +value_dim 64` (0.75): a spill's zero-fill and store walk every **column block** +of a tile wider than `mlen`, and the term billed one. Fixing that moved the range +from 0.75-0.96 to 0.87-1.14. + +`test_kda_stage_calibration.py` pins it, and states what it cannot see: at +instruction-count granularity, removing the seven matmuls or the UT forward +substitution moves the count by only 9% — each matmul is one `M_TMM` issue — so +those terms stay as unvalidated as every other stage, while being large in +*cycles*, which is what the driver reports. + +### The roofline mismatch is real, and it has a mechanism + +**An earlier version of this section said the opposite.** It reported that the +FLOP share and the cycle share track each other and concluded that this model +shows no mismatch. That was measured at one point -- the shipped `ANALYTIC` +config, `MLEN` 2048 with bandwidth derived from `HBM_WIDTH` 512 -- on a plane +neither of whose axes is known. Sweeping both changes the answer. + +`hybrid_model.py --sweep ` reports an operator's cycle share divided by +its FLOP share. 1.00 means it costs what its arithmetic says; above 1 means the +machine spends time the FLOP count cannot see. Kimi Linear, decode, 4k context: + +| operator | `MLEN` 64, any bandwidth | `MLEN` 2048, 2048 B/cycle | +|---|---|---| +| **KDA** | 0.96 | **1.80** | +| MoE | 1.11 | 0.76 | +| attention | 0.64 | 0.63 | + +On a narrow machine the three sit within 0.6-1.1 of each other and nothing stands +out. On a wide one KDA costs **2.9x** what attention does per unit of arithmetic. +The mismatch is not a property of the workload; it is a property of the workload +**on a wide machine**. + +### The mechanism is KDA's channel-wise decay, and Mamba-2 does not share it + +The cause is vector-lane utilisation, and it is arithmetic rather than +speculation. Each vector op covers `ceil(width / VLEN) * VLEN` lanes and uses +`width` of them: + +| | operand width per vector op | at `VLEN` 64 | at `VLEN` 2048 | +|---|---|---|---| +| KDA, one key row | `value_dim` = 128 | 100% | **6.2%** | +| Mamba-2, one state block | `head_dim * state_size` = 8,192 | 100% | **100%** | +| Kimi hidden activation | 2,304 | 100% | 56.2% | +| MoE intermediate | 1,024 | 100% | 50.0% | + +**KDA's decay is per key channel**, so every key row carries its own scalar and +has to be its own vector operation -- 128 lanes wide, whatever the machine is. +Mamba-2's decay is one scalar per head, so its whole `[head_dim, state_size]` +state scales in a single op and fills any width. + +That is the same fact this project has turned on from the beginning: the +channel-wise decay is what makes KDA more expressive than a scalar-gated SSM, and +it is *also* what makes it waste a wide vector unit. The two are the same +property seen from two sides. + +Nemotron-3 Super confirms the negative half. Its recurrent operator is Mamba-2, +and its ratio moves only 0.60 to 0.84 across the whole plane while its MoE sits +flat at 1.09. **"Recurrent operators are disproportionately expensive" is not +the finding**; "an operator whose vector operand is narrower than the machine is" +is. + +### Which configuration is real still decides the number + +| context | KDA ratio, `ANALYTIC` MLEN 2048 | KDA ratio, `TRANSACTIONAL` MLEN 64 | +|---|---|---| +| 4,096 | 1.27 | **0.96** | +| 1,048,576 | 1.37 | 1.46 | + +The sweep above is the answer to this, not a workaround for it: the design +question is not "which of the two declared configurations is correct" but "how +wide should the machine be, given that width past the recurrent state's own +dimension buys the recurrent operator nothing". That is a number the sweep +produces and a single configuration cannot. + +Mixing the two configurations is still not survivable: overriding `MLEN` to 64 +while leaving `HLEN` at 128 makes `flash_attention` divide by `mlen // hlen` = 0, +which is why `MACHINE_POINTS` moves all four fields together. + +### What the earlier reading got wrong + +### The single-point reading, and why it was wrong + +Kimi Linear, decode, at the shipped `ANALYTIC` configuration only: + +| context | KDA %FLOP | %cycles | ratio | attention %FLOP | %cycles | ratio | +|---|---|---|---|---|---|---| +| 4,096 | 24.7% | 31.5% | 1.27 | 8.8% | 7.1% | 0.80 | +| 65,536 | 15.2% | 20.0% | 1.31 | 43.8% | 41.1% | 0.94 | +| 1,048,576 | 2.1% | 2.9% | 1.37 | 92.1% | 91.4% | 0.99 | + +Read alone, this says the two shares track each other and no operator is +disproportionately expensive. **Sweeping context was the wrong axis.** It changes +which operator dominates without changing what any of them costs per unit of +arithmetic, so the ratios barely move and the plane the answer actually lives on +is never visited. + +Even here the signal was present — KDA at 1.27–1.37 against attention at +0.80–0.99, a 1.6x spread — and it was read as noise around 1.0 rather than as the +narrow end of a trend that reaches 2.9x. Reporting the number that came out is +necessary but not sufficient; a number from one point on an unexplored plane is +not yet a result. + +### The two declared configurations, as the endpoints of that plane + +The same context sweep on the two configurations `plena_settings.toml` declares: + +| context | KDA ratio, `ANALYTIC` MLEN 2048 | KDA ratio, `TRANSACTIONAL` MLEN 64 | +|---|---|---| +| 4,096 | 1.27 | **0.96** | +| 65,536 | 1.31 | 1.12 | +| 1,048,576 | 1.37 | 1.46 | + +At short context the wide machine makes KDA look disproportionately expensive and +the narrow one makes it exactly proportionate. That is a **qualitative** +disagreement produced by a config field, not a modelling subtlety, and it is why +`MLEN` 2048 against 64 (and `BLEN` 128 against 4, `HLEN` 128 against 16) has to +be settled before any of these numbers mean anything. Mixing the two is not +survivable either: overriding `MLEN` to 64 while leaving `HLEN` at 128 makes +`flash_attention` divide by `mlen // hlen` = 0. + +No times are reported anywhere in this section. `runtime_config.rs` hard-codes +`PERIOD` at 1 ns with no stated basis, so cycles stay cycles. + +### What "wide" and "narrow" mean here, and the assumption inside them + +`MLEN`, `BLEN`, `HLEN` and `VLEN` are the tile sizes `doc/plena_isa_spec.md` +defines: the matrix machine's tile, the systolic array's, the partitioned +systolic array's, and the vector machine's. `M_MM` fetches a `(BLEN, MLEN)` tile +from Vector SRAM against an `(MLEN, BLEN)` tile from Matrix SRAM. So `MLEN` is +the edge one matrix instruction addresses in and `VLEN` is how many elements one +vector instruction covers. + +"Narrow" is `TRANSACTIONAL`: 64 / 4 / 16 / 64. "Wide" is `ANALYTIC`: +2048 / 128 / 128 / 2048. + +**The assumption this bakes in has to be stated, because it decides how the +tables below may be read.** In `perf_model.py` those fields enter cycles *only* +through instruction counts: + + cycles = ceil(work / MLEN) x instruction latency + +and the latency is a constant. `V_MUL_VF` and `V_FMA_VF` bill +`VECTOR_MUL_CYCLES` from the TOML whatever `VLEN` is; `M_MM` bills `BLEN`. So a +wider machine here means **one instruction covering 32x the data at the same +cost** -- 32x the throughput, free, and with no area term anywhere in the model. + +That makes the *absolute* speedups below near-tautological. "Widening from 64 to +2048 makes decode 15x faster" is mostly a restatement of the assumption, not a +finding, and it should not be quoted as one. + +**What survives the assumption is the comparison between operators.** KDA and +Mamba-2 are priced by the same idealisation, so when one saturates on width at +23.9% marginal return and the other is still at 44.0%, the difference is a +property of the two operators and not of the free throughput both were granted. +Every result in this section is of that shape: a ratio between operators, or a +turning point, never an absolute speedup. + +**A second inconsistency, found while writing this.** The transactional +emulator's matrix core is a hard-coded `big_4x1024` -- 4,096 PEs, in +`matrix_core.rs`, independent of `MLEN` entirely. So the two halves model width +by different mechanisms: the analytic side by instruction count, the emulator +side by a fixed PE array. That is separate from the `MLEN` disagreement in +`plena_settings.toml` and equally unresolved. + +### How wide should the machine be + +The ratio grid says an operator's *share* of the time gets worse on a wider +machine. It does not say whether the machine is faster, which is what a width +decision turns on. `--sweep` reports that too. + +**Width and bandwidth have to scale together, or width buys nothing.** Marginal +return of doubling `MLEN`, as a percentage of total decode cycles removed, Kimi +Linear: + +| doubling | 64 B/cyc | 256 | 1024 | 4096 | +|---|---|---|---|---| +| 64 → 256 | 14.4% | 74.6% | 75.5% | 75.5% | +| 256 → 512 | 0.0% | 6.2% | 43.0% | 43.0% | +| 512 → 1024 | 0.0% | 0.0% | 36.8% | 36.8% | +| 1024 → 2048 | 0.0% | 0.0% | 13.4% | 23.9% | + +The zeros are the finding. Below the diagonal a wider machine removes **no** +cycles at all: the memory server is already the limit and the extra lanes idle. +On Nemotron-3 Super the same cells go slightly **negative** (−0.1% to −0.5%) -- +a wider machine is fractionally slower when bandwidth does not follow it. + +**Where width stops paying depends on which recurrent operator the model uses.** +At matched bandwidth, the last doubling (1024 → 2048) returns: + +| | marginal return | recurrent operator's share of cycles, 64 → 2048 | +|---|---|---| +| Kimi Linear (KDA) | **23.9%** | 23.9% → **44.5%** | +| Nemotron-3 Super (Mamba-2) | **44.0%** | 10.7% → 15.1% | + +The Mamba model is still returning nearly half per doubling at the widest point +tested. The KDA model has decayed to a quarter, and the reason is visible in the +second column: as the machine widens, KDA goes from a quarter of decode time to +**nearly half**, while contributing 24.7% of the FLOPs and running at **6.2% +vector-lane utilisation**. Mamba-2's share barely moves, because its state block +fills any width. + +So at `MLEN` 2048, about **45% of Kimi Linear's decode time is spent by an +operator using 6.2% of the lanes it is charged for**. That is the concrete +inefficiency a heterogeneous-operator substrate would target, and it is the +number this driver exists to produce. + +**What this does not say.** There is no area model here, so "wider is better" +wherever it still helps is trivially true and not a recommendation. What is not +trivial is the shape: the returns decay at a rate set by the narrowest operand in +the model, and two hybrid models with the same layer census and different +recurrent operators sit at different points on that curve. + +--- + +### Prefill and decode want different machines, and by a lot + +The width question has a different answer in each mode. Total cycles at matched +bandwidth, Kimi Linear: + +| `MLEN` | decode (Mcycles) | prefill (Mcycles) | +|---|---|---| +| 256 | 14.20 | 3,831 | +| 512 | 8.10 | 1,856 | +| 1024 | 5.12 | **1,742** | +| 2048 | **3.89** | 2,650 | + +**Decode's optimum is the widest point tested; prefill's is 1024, and 2048 is +52% worse than that.** A single fixed width cannot be right for both, which is a +direct measurement of the thing a mode-switching fabric would be for -- and, read +the other way, a direct cost of not having one. + +**Nemotron-3 Super does not have this problem.** Same sweep, same axes: + +| `MLEN` | decode, marginal return | prefill, total (Mcycles) | +|---|---|---| +| 1024 | | 2,808 | +| 2048 | **+44.0%** | **1,145** (+59.2%) | + +Both modes still want the widest machine tested. So the divergence is not a +property of hybrid models, or of recurrent mixers, or of prefill: it is a +property of **KDA**, and the next section says why. + +### The mechanism is `chunk` against `mlen`, and it is a numerical constraint + +Decomposing KDA prefill by machine width, one chunk and one head, instructions +issued: + +| `MLEN` | total | without the spill term | spill share | +|---|---|---|---| +| 64 | 2,450 | 1,696 | 30.8% | +| 512 | 4,390 | 1,072 | 75.6% | +| 2048 | 14,173 | **1,063** | **92.5%** | + +**The real work gets cheaper as the machine widens** -- 1,696 down to 1,063, +which is what wider tiles should do. The spill term grows from 754 instructions +to 13,110, and swamps it. + +Every spill zero-fills the rows past its live data, because `load_sub_matrix_*` +prefetches a whole `mlen x mlen` MRAM block unconditionally: `k_block_count` +selects whole blocks and cannot trim a partial one. So the fill scales with +`mlen` while the live data scales with `chunk` -- and `chunk` **cannot grow to +meet it**. `kda_chunk_check_range` caps it at 17, because `1/A` reaches +`exp(chunk * |gate_lower_bound|)` and overflows bf16 past that. + +At `mlen` 2048 that is 2,048 rows zeroed to protect 16 rows of data, six times +per chunk. The bf16 range of the reciprocal decay is what sets the numerator and +the machine's tile size sets the denominator, and nothing in the lowering can +reconcile them. + +Mamba-2's chunk is **256** -- `nemotron-3-super-120b-a12b.json` says so, and +nothing bounds it numerically, because the SSD scan has no reciprocal decay to +overflow. Sixteen times more of every tile is live, which is why its prefill +keeps improving where KDA's turns around. + +So both of KDA's costs on a wide machine come from the same design choice this +project has circled from the start. The decay is **channel-wise on the key +axis**, which makes each decode sweep 128 lanes wide (6.2% utilisation at `VLEN` +2048) and forces the chunk form to divide by a per-channel cumulative decay, +whose reciprocal caps the chunk at 17 and leaves prefill's tiles 99% padding. The +same property that makes KDA more expressive than a scalar-gated SSM is what +costs it in both modes, by two different routes. + +### The ISA consequence, priced rather than asserted + +A prefetch taking a row count instead of a whole-block count would remove the +fill and store only the live rows. `kda_chunk_prefill(row_granular_prefetch=True)` +models exactly that -- no such instruction exists, and the switch is there to +price it, not to assume it. `--row-granular-prefetch` carries it through the +driver. + +Instructions issued, one chunk and one head: + +| `mlen` | whole-block prefetch | row-granular | factor | +|---|---|---|---| +| 64 | 2,450 | 1,801 | 1.4x | +| 512 | 4,390 | 1,146 | 3.8x | +| 2048 | 14,173 | 1,137 | **12.5x** | + +**The proposal is worth almost nothing on a narrow machine and a great deal on a +wide one**, which is the shape you would expect from a term that scales with +`mlen`. Anyone reading 12.5x as the value of this instruction on the machine they +have should check which column they are in first. + +*(An earlier draft of this section said 13x, from 14,173 down to 1,063. That was +the count with the whole spill term zeroed. Row-granular prefetch still stores +the live rows, so the floor is 1,137 and the factor is 12.5x. Predicting a number +and then measuring it is how the difference showed up.)* + +### And it makes the two modes agree on width + +Total prefill cycles, Kimi Linear, at matched bandwidth: + +| `MLEN` | whole-block prefetch | row-granular | +|---|---|---| +| 512 | 1,856 | 1,325 | +| 1024 | **1,742** (optimum) | 676 | +| 2048 | 2,650 (52% worse) | **515** (optimum) | + +The optimum **moves from 1024 to 2048**, and the last doubling returns **+23.9%** +instead of −52.1%. Decode's return at that same doubling is also +23.9% -- with +the fill removed, both modes are limited by the same dense work, and the width +disagreement disappears. + +So the mode divergence measured above is not a fact about KDA's algorithm. It is +a fact about **KDA's algorithm on this instruction set**: the chunk is capped at +17 by bf16 range, the prefetch granularity is a whole block, and the product of +those two is the whole effect. One of them is a numerical constraint that cannot +be moved. The other is an instruction-set choice that can. + +This is the opposite shape from a descriptor-driven state engine: it asks for +*less* to be implied by one instruction, not more. Nothing here implements it. + +--- + +--- + +## 6. The clock, and what it is not + +`runtime_config.rs` carried `pub(crate) const PERIOD: Duration = +Duration::from_nanos(1)` — a bare constant, in a source file, with no stated +basis. Every second, microsecond, TPOT and token-per-joule figure the emulator +can produce is that number's consequence. + +It is now `[.CONFIG.CLOCK_PERIOD_PS]`, defaulted to the same 1000 ps so +nothing moves, with the assumption written where it can be read: + +> THIS IS AN ASSUMPTION, NOT A MEASUREMENT. No RTL has been synthesised, so no +> critical path has set a frequency. + +Three things changed beyond moving the value. + +**It is printed at startup, with its provenance.** Every emulator run now begins: + +``` +Clock: 1000 ps (1.000 GHz) from CLOCK_PERIOD_PS -- an assumption, not a +synthesised frequency. DRAM model: 1000 ps (1x). +``` + +A number that appears in a report and a number that appears in a log line have +different half-lives. This one is now impossible to use without seeing where it +came from. + +**The relationship to the DRAM clock is asserted, not assumed.** +`stage_profile.rs` had already recorded the hazard in a comment: a cycle-domain +comparison there *"only held because the DRAM tCK happened to equal PERIOD; any +preset or frequency change made it fail"*. It does happen to be equal — the HBM2 +preset at 2000 MBPS is a 1 ns command clock — and that is a property of the +preset, not of the design. `runner` now requires one period to be a whole +multiple of the other and refuses to start otherwise: + +``` +accelerator clock is 700 ps (1.429 GHz, from CLOCK_PERIOD_PS) and the DRAM +model's is 1000 ps; neither divides the other, so ticks never line up and any +cycle count derived from both is off by a fraction nothing reports +``` + +Verified by setting it to 700 and watching it refuse. + +**No times are reported anywhere in this document.** Making the clock +configurable does not make it known. Until a synthesis run sets a frequency, a +microsecond figure here would be a cycle count multiplied by a number someone +chose, and section 5's sweep is the shape of that: cycles are the quantity the +model produces, and cycles are what it reports. + +### What is still open on this + +`analytic_models/performance/` has no clock at all — it reports cycles and stops, +which is the right behaviour for now but means there is no single place where a +frequency would be applied when one exists. The two halves also still disagree +about bandwidth: the analytic side derives it from `HBM_WIDTH` 512 as one row per +cycle, and the transactional side runs an HBM2 preset. Section 5 sweeps the +former; nothing yet reconciles the two. diff --git a/docs/superpowers/plans/2026-08-25-static-mamba-kda.md b/docs/superpowers/plans/2026-08-25-static-mamba-kda.md new file mode 100644 index 000000000..8da926747 --- /dev/null +++ b/docs/superpowers/plans/2026-08-25-static-mamba-kda.md @@ -0,0 +1,2181 @@ +# Static Mamba-2 + KDA on PLENA Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Run Nemotron-3 Mamba-2 and Kimi-K3 KDA on PLENA — **prefill and decode, whole model, on real checkpoint weights** — using only statically-scheduled instructions: no memory-resident descriptors, no command queues, no residency cache. + +**Architecture:** Recurrent state is an ordinary tensor: it lives in Vector SRAM at compiler-chosen addresses and moves with `H_PREFETCH_V` / `H_STORE_V`, exactly like a KV cache. The recurrence is a loop of row-wise vector ops. Both Mamba-2 and KDA store per-head state with the **contracted axis as the row axis** (Mamba: `state_size`; KDA: `key_dim`), which turns every step of both recurrences — decay, rank-1 update, and state contraction — into one broadcast-scalar operation per row. + +**Phasing.** Functionality first, optimisation second, each measured before the next is designed: + +| Phase | What lands | New opcodes | Why here | +|---|---|---:|---| +| **0** (Tasks 1–4) | KDA **decode** runs end to end, validated against the golden | **0** | Proving KDA needs no new instruction is itself the result worth being able to state. Uses the same `copy + multiply + add` triple the Mamba kernel uses today. | +| **1** (Tasks 5–8) | `V_FMA_VF` (`rd += rs1 * fp2`), both decode kernels converted onto it | 1 (`0x3B`) | Removing the scratch row makes each sweep a single arithmetic-progression row walk, which the existing emitter turns into a **hardware loop** instead of an unrolled block. Phase 0 supplies the measured baseline. | +| **2** (Tasks 9–10) | KDA **chunked prefill** on the Matrix Engine | 0 | Decode alone is half the function. Prefill is the half that had never been mapped onto the Matrix Engine. | +| **3** (Tasks 11–13) | Whole model lowered, **executed end to end in Rust**, bound to **real checkpoint weights** | 0 | The largest credibility gap in the project: no whole model has ever been run through the emulator here, and none has ever used real weights. | +| **4** (Tasks 14–15) | Banking study, uncalibrated latency report, no-descriptor invariant confirmed | 0 | Measurement and confirmation, last. | + +KDA needs **no new instruction** to work: `V_SOFTPLUS_V` is Mamba's (`dt`), and `S_MAP_FP_V` — the one primitive the recurrence genuinely requires, to get per-key scalars from a VRAM row into FP registers — already ships on `feat/mamba2-support`. `V_FMA_VF` is therefore a measured optimisation, never a prerequisite. + +**Tech Stack:** Python 3 (compiler / `aten` lowering, pytest), Rust (transactional emulator, `cargo test`), SystemVerilog headers (`doc/operation.svh`, shared ISA definition), `just` recipes for the test matrix. + +## Global Constraints + +- **Base branches.** Compiler work branches from `feat/mamba2-support` (Compiler `f5eb36a`). Emulator work branches from `feat/mamba2-support` (Simulator `65cf6f09`). Do **not** branch from any descriptor-driven line of work. +- **New branch names.** Compiler: `feat/static-kda`. Simulator: `feat/static-kda`. +- **Opcode allocation.** `V_FMA_VF = 6'h3B`. `0x39` (`V_SOFTPLUS_V`) and `0x3A` (`S_MAP_FP_V`) are already taken by `feat/mamba2-support`. Do not use `0x3D` or `0x3F` — leaving them unallocated keeps the two histories distinguishable. +- **No new memory-resident control structures.** No instruction may read a descriptor, command block, or any control word from HBM. Every instruction is fully defined by its 32-bit word plus sticky control registers. +- **No new runtime traps.** Shape, precision, and address violations are compile-time errors in Python or assembler errors — never a status code returned by hardware. +- **State precision is BF16 for Mamba-2, FP32 for KDA**, matching the profiled runtime configurations. Never store recurrent state as MX8: the docstring at `aten/plena/program_ssm_recurrent.py:104-115` gives the error-amplification argument; it applies unchanged to KDA and must be preserved in any new docstring. +- **Every state load and its matching store must agree on `storage_precision` and `hbm_element_bytes`.** Mismatch produces a wrong answer rather than an error — see `aten/plena/program_ssm_recurrent.py:120-126`. +- **Every new row-walking emitter must go through `self._row_progression(rows)` and emit a hardware loop when it returns non-`None`.** Falling back to the unrolled path for a long walk is the single biggest static-footprint regression available; it is the single biggest contributor to a bloated whole-model static footprint. +- **`MAX_LOOP_DEPTH = 4`** (`PLENA_RTL/src/definitions/configuration.svh:35`). No emitted loop nest may exceed 4 levels. +- **All emitters allocate GP registers via `self._reg.allocate_gp(n)` inside a `try` / `finally: self._reg.free_gp(gp_regs)`.** Never hardcode a GP index — `mamba_rsqrt_fpram` (`program_mamba_common.py:568-596`) documents why: the allocator hands registers out descending, so a hardcoded index silently clobbers a live value. +- **Test commands.** Compiler: `pytest` from the Compiler root. Emulator: `cargo test` from `transactional_emulator/`, plus `cargo clippy --all-targets -- -D warnings`. +- **Commit style.** Conventional commits (`feat:`, `fix:`, `test:`, `refactor:`, `docs:`), matching existing history. + +## Out of scope + +Named here so a reviewer does not read them as gaps. Each is a separate subsystem needing its own plan. + +- **Post-increment vector addressing.** See the Reference section: emitted instruction count is dominated by address arithmetic, not by arithmetic. A sticky "advance by STRIDE after each V-type op" mode would cut the inner loop from 6 instructions to 2. That is the highest-value follow-up, and it is a general ISA improvement rather than a Mamba/KDA one. +- **Vector SRAM bank mapping.** Task 14 measures whether row-walking a head tile conflicts. Only if it does should a bank-mapping mode be designed, in its own plan. + +--- + +## Reference: the two recurrences + +Stated in the exact form the lowering implements. Task authors should not re-derive them. + +**Mamba-2 decode step**, per head `h`, group `g = h // heads_per_group`, state `S[n, :]` for `n` in `0..state_size`, each row `head_dim` wide: + +``` +xs[:] = dt[h] * x[h, :] +S[n, :] = dA[h] * S[n, :] + B[g, n] * xs[:] for each n +y[h, :] = sum_n C[g, n] * S[n, :] +y[h, :] += D[h] * x[h, :] +``` + +**KDA decode step**, per head `h`, state stored **transposed** as `T[k, :]` for `k` in `0..key_dim`, each row `value_dim` wide (`T[k, v] == S[v, k]` in the row-major reference): + +``` +q̂ = normalize(q[h]); k̂ = normalize(k[h]) # L2 over key_dim +beta = sigmoid(beta_proj[h]) +decay[k] = exp(gate_lower_bound * sigmoid(exp(a_log[h]) * (gate[h, k] + dt_bias[h, k]))) + +# sweep 1: decay, then predict +T[k, :] = decay[k] * T[k, :] for each k +pred[:] = sum_k k̂[k] * T[k, :] + +err[:] = beta * (v[h, :] - pred[:]) + +# sweep 2: update, then read out (state must be updated before it is read) +T[k, :] += err[:] * k̂[k] for each k +out[:] = sum_k q̂[k] * T[k, :] +o[h, :] = output_scale * out[:] +``` + +The transposition is the load-bearing design decision. In the reference orientation (`state[v][k]`, with `key` contiguous) the per-key decay is strided and both contractions run within a row, needing `V_MUL_VV` + `V_RED_SUM` per row. Transposed, `key` becomes the row axis and all four steps reduce to one broadcast-scalar op per row. + +### Instruction accounting — read this before setting any budget + +The emitters do **not** emit one instruction per vector operation. `_emit_tile_row_fp_scalar` (`isa_tile_rows.py:371-409`) shows the real hardware-loop body: + +```asm +C_LOOP_START gp_loop, row_count +S_LD_FP f1, gp_fp, 0 + gp_src, gp_src, f1, 0 +S_ADDI_INT gp_src, gp_src, row_step * mlen +S_ADDI_INT gp_fp, gp_fp, fp_step +C_LOOP_END gp_loop +``` + +So a per-row-scalar sweep costs **5 dynamic instructions per row**, of which 1 is arithmetic. A `V_FMA_VF` sweep needs a third pointer and costs **6 per row**. Two consequences: + +**1. The win from `V_FMA_VF` is mostly static, not dynamic.** The `copy + multiply + add` per state row — what Mamba does today and what Phase 0's KDA will do — is three separate emitter calls with `rows=[1]`, so each falls to the *unrolled* path: roughly 12 fully-expanded instructions per state row, every one a distinct word in the program image. With FMA each sweep becomes one arithmetic progression and therefore one hardware loop: + +| | static instr / layer | dynamic instr / state row | +|---|---:|---:| +| Mamba-2 today (unrolled triples) | ~98,000 | ~12 | +| Mamba-2 with `V_FMA_VF` (3 hardware loops per head) | **~1,900** | 17 | + +A ~50× static reduction. That is what attacks the 23.663 MiB / 44.490 MiB program images directly, and it is the reason `V_FMA_VF` earns an opcode. + +**2. Dynamic cost is dominated by address arithmetic, and no latency number in this project is calibrated.** Each row costs 5–6 emitted instructions of which 1 is arithmetic. The lever is post-increment addressing (out of scope above). A state coprocessor would *not* fix this — it has the same 64 lanes (`2 head_lanes × 32 fma_lanes` = 64, identical to `VLEN=64`); it only hides the addressing inside a private sequencer. + +**Do not convert instruction counts into microseconds and then into a decision.** Every PLENA latency available anywhere in this project is an uncalibrated analytic model output. A µs figure of the form `bytes ÷ bytes-per-cycle` is a roofline the model produces by construction, not a measurement, and two such figures decide nothing between them. + +**Two different kinds of number, and only one is admissible as a gate:** + +| Quantity | Status | Use | +|---|---|---| +| **Static instruction count** | Hard fact — a property of the compiled artifact | **Gate on it.** Tasks 8 and 11. | +| Dynamic instruction count | Hard fact, given a fixed program and token count | Report it. | +| Cycles / µs / TPOT | Uncalibrated model output on both branches | **Report only, labelled uncalibrated. Never a gate, never a comparison against another uncalibrated model.** | + +Task 3 records the Phase 0 static count, Task 8 Step 9 records the Phase 1 count, and **their ratio is the justification for spending opcode `0x3B`** — not the estimates in this section. + +--- + +## File Structure + +**Simulator (`PLENA_Simulator`, branch `feat/static-kda`)** + +| Path | Responsibility | +|---|---| +| `transactional_emulator/src/op.rs` | Modify: decode `0x3B` into `Opcode::V_FMA_VF` | +| `transactional_emulator/src/vector_machine.rs` | Modify: add `fma_scalar` next to `mul_scalar` | +| `transactional_emulator/src/accelerator/dispatch.rs` | Modify: execute `V_FMA_VF`; add to the VRAM-read operand list | +| `analytic_models/reference/kimi_k3_kda.py` | Create: KDA golden, single oracle for both repos | +| `analytic_models/reference/test_kimi_k3_kda.py` | Create: reference self-test | +| `analytic_models/performance/vector_sram_banking.py` | Create (Task 14): bank model retargeted at Vector SRAM | +| `transactional_emulator/testbench/models/kda_static_test.py` | Create: end-to-end KDA layer, emulator vs golden | +| `transactional_emulator/testbench/models/whole_model_test.py` | Create (Task 12): first layer to last, both models | +| `transactional_emulator/testbench/models/checkpoint_test.py` | Create (Task 13): real published weights vs HuggingFace | +| `doc/checkpoint_validation.md` | Create (Task 13): exactly what was validated, and what was not | +| `doc/static_path_measurements.md` | Create (Task 14): hard facts and uncalibrated model output, kept apart | + +**Compiler (`PLENA_Compiler`, branch `feat/static-kda`)** + +| Path | Responsibility | +|---|---| +| `doc/operation.svh` | Modify: `V_FMA_VF = 6'h3B` | +| `doc/plena_isa_spec.md` | Modify: document `V_FMA_VF` | +| `assembler/assembly_to_binary.py` | Modify: encode `V_FMA_VF` | +| `aten/plena/isa_tile_rows.py` | Modify: `_emit_tile_row_fma` + `tile_row_fma_fp_asm` + `tile_row_fma_fp_broadcast` | +| `aten/plena/program_ssm_recurrent.py` | Modify: rewrite `ssm_decode_step_v0` onto FMA sweeps | +| `aten/plena/program_kda_common.py` | Create: `ProgramKdaCommonMixin` — state load/store, conv roll, L2 normalize | +| `aten/plena/program_kda_recurrent.py` | Create: the two-sweep KDA decode step | +| `aten/plena/program_kda_layer.py` | Create: full KDA decode layer | +| `aten/plena/program_kda_chunk.py` | Create (Task 9): chunk decay cumsum and the UT transform | +| `aten/plena/program_kda_prefill.py` | Create (Task 10): chunked prefill on the Matrix Engine | +| `aten/models/kda/reference.py` | Create: re-export of the Simulator's golden, incl. the single `KdaShape` | +| `aten/native_ops.yaml` | Modify: register `kda_decode` for CPU and PLENA backends | +| `aten/ops/cpu/kda_ops.py`, `aten/ops/plena/kda_ops.py` | Create: backend entry points | +| `aten/tests/test_v_fma_vf.py` | Create: encoding + emitter tests | +| `aten/tests/test_kda_recurrent.py` | Create: per-head and per-layer numeric tests | +| `aten/tests/test_instruction_budget.py` | Create: measured **static** instruction-count gates | +| `aten/tests/test_kda_prefill.py` | Create (Tasks 9–10): UT transform, prefill/decode equivalence | +| `aten/tests/test_whole_model_lowering.py` | Create (Task 11): 52- and 93-layer lowering | + +**`KdaShape` is defined exactly once**, in `analytic_models/reference/kimi_k3_kda.py`, and re-exported by `aten/models/kda/reference.py`. `program_kda_common.py` imports it; it must not define its own. + +--- + +## Phase 0 — KDA decode runs at all + +**Zero new instructions.** Everything here is lowering work on top of what `feat/mamba2-support` already ships. Nothing in this phase touches `doc/operation.svh`, the decoder, or the assembler, so it carries no RTL implication and no ISA risk. At the end of Phase 0, KDA decode works; prefill is Phase 2. + +--- + +### Task 1: KDA golden reference + +Establishes the reference maths so the static lowering has an oracle independent of any descriptor machinery. + +**Files:** +- Create: `analytic_models/reference/kimi_k3_kda.py`, `analytic_models/reference/test_kimi_k3_kda.py` (Simulator) +- Create: `aten/models/kda/__init__.py`, `aten/models/kda/reference.py` (Compiler) + +**Interfaces:** +- Consumes: nothing. +- Produces: + - `KdaShape(num_heads: int, key_dim: int, value_dim: int, conv_kernel: int, heads_per_group: int = 1)` with `state_rows -> num_heads * key_dim` and `head_base(head) -> head * key_dim`. **This is the only definition of `KdaShape` in either repo.** + - `kda_decode_step(state, q, k, v, gate, beta_logit, a_log, dt_bias, gate_lower_bound, output_scale) -> tuple[Tensor, Tensor]` returning `(new_state, output)`, with `state` in the **reference** orientation `[heads, value_dim, key_dim]`. + - `kda_layer_reference(weights, tokens, shape) -> Tensor`. + +- [ ] **Step 1: Establish the reference** + +```bash +# Write analytic_models/reference/kimi_k3_kda.py and its self-test: the KDA +# forward in torch, transposed-state formulation, no PLENA concepts. +``` + +- [ ] **Step 2: Run the reference self-test** + +Run: `pytest analytic_models/reference/test_kimi_k3_kda.py -v` +Expected: PASS. If an import fails, bring that module across too — but nothing that imports from a runtime state engine. + +- [ ] **Step 3: Add `KdaShape` if the reference module lacks it** + +```python +from dataclasses import dataclass + + +@dataclass(frozen=True) +class KdaShape: + """KDA head geometry. + + `state_rows` is expressed in the TRANSPOSED layout the PLENA lowering + uses -- one row per (head, key) -- because that is the orientation in + which every step of the recurrence is a single broadcast-scalar op. The + reference functions in this module keep the [heads, value_dim, key_dim] + orientation; the lowering transposes on load. + """ + + num_heads: int + key_dim: int + value_dim: int + conv_kernel: int + heads_per_group: int = 1 + + @property + def state_rows(self) -> int: + return self.num_heads * self.key_dim + + def head_base(self, head: int) -> int: + return head * self.key_dim +``` + +- [ ] **Step 4: Write the failing transposition-contract test** + +Append to `analytic_models/reference/test_kimi_k3_kda.py`: + +```python +def test_decode_step_matches_the_transposed_formulation(): + """T[k, v] == S[v, k]. This is the contract between the reference and + aten/plena/program_kda_recurrent.py.""" + import torch + + torch.manual_seed(0) + heads, key_dim, value_dim = 2, 8, 8 + state = torch.randn(heads, value_dim, key_dim, dtype=torch.float32) + q = torch.randn(heads, key_dim); k = torch.randn(heads, key_dim) + v = torch.randn(heads, value_dim) + gate = torch.randn(heads, key_dim) + beta_logit = torch.randn(heads) + a_log = torch.randn(heads); dt_bias = torch.randn(heads, key_dim) + + new_state, out = kda_decode_step( + state, q, k, v, gate, beta_logit, a_log, dt_bias, + gate_lower_bound=-5.0, output_scale=1.0, + ) + + t = state.transpose(1, 2).clone() # [heads, key, value] + qn = torch.nn.functional.normalize(q, dim=-1) + kn = torch.nn.functional.normalize(k, dim=-1) + beta = torch.sigmoid(beta_logit) + decay = torch.exp(-5.0 * torch.sigmoid(a_log[:, None].exp() * (gate + dt_bias))) + t = t * decay[:, :, None] + pred = (kn[:, :, None] * t).sum(dim=1) + err = beta[:, None] * (v - pred) + t = t + err[:, None, :] * kn[:, :, None] + out_t = (qn[:, :, None] * t).sum(dim=1) + + torch.testing.assert_close(t.transpose(1, 2), new_state, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(out_t, out, rtol=1e-5, atol=1e-6) +``` + +- [ ] **Step 5: Run it** + +Run: `pytest analytic_models/reference/test_kimi_k3_kda.py::test_decode_step_matches_the_transposed_formulation -v` +Expected: PASS. + +If it fails, the reference orders decay, prediction, or normalisation differently than the replay above. **Fix the replay to match the reference**, then update this plan's Reference section to the corrected ordering — the lowering must follow the reference, not the other way round. + +- [ ] **Step 6: Expose it to the Compiler** + +Create `aten/models/kda/__init__.py` (empty) and `aten/models/kda/reference.py`: + +```python +"""KDA golden reference. + +The maths lives in the Simulator repo (analytic_models/reference/kimi_k3_kda.py) +so the compiler tests and the emulator tests share one oracle instead of two +drifting copies. This module only re-exports it. +""" + +from analytic_models.reference.kimi_k3_kda import ( # noqa: F401 + KdaShape, + kda_decode_step, + kda_layer_reference, +) + +__all__ = ["KdaShape", "kda_decode_step", "kda_layer_reference"] +``` + +If `kda_layer_reference` does not exist in the reference module, write it there now — a whole-layer forward taking the weight dict from Task 4 — and add a self-test for it before moving on. + +- [ ] **Step 7: Commit (both repos)** + +```bash +# Simulator +git add analytic_models/reference/ +git commit -m "test(kda): add the KDA golden reference and pin the transposed formulation" +# Compiler +git add aten/models/kda/ +git commit -m "test(kda): re-export the KDA golden reference for compiler tests" +``` + +--- + +--- + +### Task 2: KDA state layout and data movement + +**Files:** +- Create: `aten/plena/program_kda_common.py` +- Create: `aten/tests/test_kda_recurrent.py` + +**Interfaces:** +- Consumes: `KdaShape` from `aten.models.kda.reference` (Task 1); `pin_hbm_region`, `pinned_hbm_region`, `load_batch`, `store` from `ProgramSSMRecurrentMixin` (`program_ssm_recurrent.py:63-136`); `mamba_row_copy`, `mamba_block_copy` (`program_mamba_common.py:325-342`); `tile_row_mul`, `tile_row_sum` (`isa_tile_rows.py:188, 52`); `mamba_rsqrt_fpram` (`program_mamba_common.py:568`). +- Produces `ProgramKdaCommonMixin` with: + - `kda_load_state_v0(name: str, shape: KdaShape, hbm_addr: int) -> VRAMMatrixVar` — `[state_rows, value_dim]`, FP32 + - `kda_store_state_v0(state, name: str, hbm_addr: int) -> InputVar` + - `kda_conv_state_roll_v0(conv_state, new_row_src, new_row_idx, shape)` + - `kda_l2_normalize_v0(vec, row: int, sq_scratch, sq_row: int, acc_fp, consts)` + +- [ ] **Step 1: Write the failing state tests** + +Create `aten/tests/test_kda_recurrent.py`: + +```python +"""KDA static lowering: state layout, decode step, and full layer.""" + +import torch + +from aten.models.kda.reference import KdaShape + + +def test_state_is_loaded_in_the_transposed_layout(kda_program): + shape = KdaShape(num_heads=2, key_dim=4, value_dim=kda_program.mlen, conv_kernel=4) + addr = kda_program.pin_hbm_region( + "kda_state", shape.state_rows * shape.value_dim, hbm_element_bytes=4 + ) + state = kda_program.kda_load_state_v0("kda_state", shape, addr) + assert state.shape == (shape.num_heads * shape.key_dim, shape.value_dim) + + +def test_state_store_mirrors_the_load_precision(kda_program): + """A load/store precision mismatch is a wrong answer, not an error.""" + shape = KdaShape(num_heads=1, key_dim=4, value_dim=kda_program.mlen, conv_kernel=4) + addr = kda_program.pin_hbm_region( + "s", shape.state_rows * shape.value_dim, hbm_element_bytes=4 + ) + state = kda_program.kda_load_state_v0("s", shape, addr) + stored = kda_program.kda_store_state_v0(state, "s", addr) + assert stored.hbm_element_bytes == 4 + assert stored.precision == 1 + + +def test_conv_roll_shifts_history_and_appends(kda_program): + shape = KdaShape(num_heads=1, key_dim=4, value_dim=kda_program.mlen, conv_kernel=4) + conv = kda_program.alloc_vram("conv", (shape.conv_kernel - 1, shape.key_dim)) + new = kda_program.alloc_vram("new", (1, shape.key_dim)) + asm = kda_program.capture_asm( + lambda: kda_program.kda_conv_state_roll_v0(conv, new, 0, shape) + ) + body = [l for l in asm.splitlines() if l.strip() and not l.strip().startswith(";")] + assert body, "roll emitted nothing" +``` + +Build the `kda_program` fixture, and `alloc_vram` / `capture_asm`, from whatever `test_mamba2_reference.py` already uses — reuse its fixture wholesale and mix in `ProgramKdaCommonMixin`. Do not invent a new program builder. + +- [ ] **Step 2: Run and confirm it fails** + +Run: `pytest aten/tests/test_kda_recurrent.py -v` +Expected: FAIL — no module `aten.plena.program_kda_common`. + +- [ ] **Step 3: Implement the module** + +Create `aten/plena/program_kda_common.py`: + +```python +"""KDA (Kimi Delta Attention) state layout and data movement. + +State layout +------------ +Per head the recurrent state is stored ``[key_dim, value_dim]`` -- the +*transpose* of the reference's ``[value_dim, key_dim]``. This is deliberate. +In the reference orientation the per-key decay is strided and both +contractions run within a row, so each needs ``V_MUL_VV`` + ``V_RED_SUM``. +Transposed, ``key`` becomes the row axis and all four steps of the recurrence +reduce to one broadcast-scalar operation per row: + + decay T[k, :] *= decay[k] + predict pred[:] += k_hat[k] * T[k, :] + update T[k, :] += err[:] * k_hat[k] + read out out[:] += q_hat[k] * T[k, :] + +Each is then a single arithmetic row progression over the state tile, which is +what lets Phase 1 turn every one of them into a hardware loop once V_FMA_VF +removes the scratch row. Phase 0 still pays copy + multiply + add. + +Precision +--------- +FP32, not MX8. The argument in ``ssm_load_state_v0``'s docstring applies +unchanged: the state is a multiplicative accumulator carried across the whole +sequence, so quantisation error is amplified by ``1 / sqrt(1 - lambda^2)``. +KDA's decay is per-key and driven toward 1 by ``gate_lower_bound``, which is +exactly the long-memory regime where the amplification is worst. The profiled +Kimi runtime also uses FP32 recurrent state with BF16 convolution state, and +this lowering matches that split. +""" + +from __future__ import annotations + +from aten.models.kda.reference import KdaShape + +__all__ = ["KdaShape", "ProgramKdaCommonMixin"] + + +class ProgramKdaCommonMixin: + def kda_load_state_v0(self, name: str, shape: KdaShape, hbm_addr: int): + """Prefetch the pinned FP32 state tile into VRAM, transposed.""" + self.emit_comment( + f"stage=kda_state_load {name} [{shape.state_rows},{shape.value_dim}]" + ) + var = self.input( + name, + (shape.state_rows, shape.value_dim), + hbm_addr=hbm_addr, + real_data_ratio=1.0, + ) + return self.load_batch( + var, name=f"{name}_vram", storage_precision=3, precision=1 + ) + + def kda_store_state_v0(self, state, name: str, hbm_addr: int): + """Write the updated state back over the same pinned HBM range. + + Must mirror kda_load_state_v0 exactly -- same precision class, same + bytes per element -- or the state read back is not the state written, + and the failure mode is a wrong answer rather than an error. + """ + self.emit_comment(f"stage=kda_state_store {name}") + return self.store( + state, + name=name, + hbm_addr=hbm_addr, + precision=1, + hbm_element_bytes=4, + real_data_ratio=1.0, + ) + + def kda_conv_state_roll_v0(self, conv_state, new_row_src, new_row_idx, shape: KdaShape): + """Shift the q/k/v conv1d history by one timestep and append the new one. + + Structurally identical to ssm_conv_state_roll_v0: a physical copy, + because address immediates are baked at ASM-gen time and a runtime + ring pointer is not expressible. + """ + history = shape.conv_kernel - 1 + if history <= 0: + return conv_state + if conv_state.shape[0] < history: + raise ValueError(f"conv_state needs {history} rows, has {conv_state.shape[0]}") + self.emit_comment(f"stage=kda_conv1d roll history={history}") + for i in range(history - 1): + self.mamba_row_copy(conv_state, i, conv_state, i + 1) + self.mamba_row_copy(conv_state, history - 1, new_row_src, new_row_idx) + return conv_state + + def kda_l2_normalize_v0(self, vec, row: int, sq_scratch, sq_row: int, acc_fp, consts): + """In-place L2 normalisation of one VRAM row. + + Same sequence mamba_gated_rmsnorm_v0 uses for RMSNorm, with the mean + scale set to 1: square, V_RED_SUM into FPRAM, then mamba_rsqrt_fpram, + which computes 1/sqrt(acc * reci_group + eps). `consts.reci_group` + must be 1.0 here -- L2 is sqrt(sum), not sqrt(mean). + """ + self.emit_comment(f"stage=kda_normalize row={row}") + self.mamba_row_copy(sq_scratch, sq_row, vec, row) + self.tile_row_mul(sq_scratch, vec, rows=[sq_row]) + self.tile_row_sum(acc_fp, sq_scratch, rows=[sq_row], target_base_offset=0) + self.mamba_rsqrt_fpram(acc_fp, consts, count=1) + self.tile_row_mul_fp_broadcast(vec, acc_fp, rows=[row], fpram_offset=0) + return vec +``` + +Check `tile_row_mul`'s and `tile_row_sum`'s real signatures at `isa_tile_rows.py:188` and `:52` and match them; the calls above follow `mamba_gated_rmsnorm_v0` (`program_mamba_common.py:543-556`). + +- [ ] **Step 4: Run the tests and confirm they pass** + +Run: `pytest aten/tests/test_kda_recurrent.py -v` +Expected: PASS, three tests. + +- [ ] **Step 5: Commit** + +```bash +git add aten/plena/program_kda_common.py aten/tests/test_kda_recurrent.py +git commit -m "feat(kda): transposed FP32 state layout, conv roll, and L2 normalize" +``` + +--- + +--- + +### Task 3: The KDA decode step (existing instructions only) + +**Files:** +- Create: `aten/plena/program_kda_recurrent.py` +- Modify: `aten/tests/test_kda_recurrent.py` + +**Interfaces:** +- Consumes: `KdaShape`, `ProgramKdaCommonMixin` (Task 2); `tile_row_mul_fp_broadcast` (`isa_tile_rows.py:258`); `tile_row_sub` (`isa_tile_rows.py:167`); `vram_fill_zero` (`isa_tile_rows.py:272`); `mamba_row_copy`, `mamba_row_add` (`program_mamba_common.py:325-337`). +- Produces: `ProgramKdaRecurrentMixin.kda_decode_step_v0(*, state, scratch, q_fp, k_fp, decay_fp, beta_fp, v, o, pred, err, shape, output_scale_fp, head_rows=None) -> VRAMMatrixVar` returning `o`. + +**Existing instructions only.** This task uses the same `copy + multiply + add` triple that `ssm_decode_step_v0` uses today. It is deliberately the un-optimised form: Phase 0 exists to prove that the static ISA expresses KDA at all, and to produce the baseline that Phase 1's `V_FMA_VF` is measured against. Task 8 converts both kernels together. + + `q_fp` / `k_fp` / `decay_fp` are FPRAM arrays holding `q̂[h, k]`, `k̂[h, k]`, `decay[h, k]`, `k`-major within each head. `beta_fp` holds `beta[h]`. `v`, `o`, `pred`, `err` are `[num_heads, value_dim]` VRAM tiles. + +**Column blocking.** `value_dim` may exceed `VLEN`. One emitter row is `self.mlen` wide, so a state row spans `blocks = value_dim // mlen` emitter rows. Sweeping all of them in one walk would need the FPRAM cursor to advance every `blocks` rows, which is not an arithmetic progression. Instead emit **one hardware loop per column block**, with the row walk stepping by `blocks` and the FPRAM cursor stepping by 1. `value_dim % mlen == 0` is required and asserted. + +- [ ] **Step 1: Write the failing tests** + +Append to `aten/tests/test_kda_recurrent.py`: + +```python +def test_single_head_decode_matches_the_golden(kda_program, run_on_emulator): + from aten.models.kda.reference import kda_decode_step + + torch.manual_seed(7) + # value_dim MUST be a multiple of the vector width -- the lowering asserts + # it, because a partial row would leave the lanes past value_dim polluting + # the accumulation. Derive it rather than hardcoding. + shape = KdaShape( + num_heads=1, key_dim=8, value_dim=kda_program.mlen, conv_kernel=4 + ) + + state = torch.randn(1, shape.value_dim, shape.key_dim, dtype=torch.float32) + q = torch.randn(1, shape.key_dim); k = torch.randn(1, shape.key_dim) + v = torch.randn(1, shape.value_dim) + gate = torch.randn(1, shape.key_dim); beta_logit = torch.randn(1) + a_log = torch.randn(1); dt_bias = torch.randn(1, shape.key_dim) + + expected_state, expected_out = kda_decode_step( + state, q, k, v, gate, beta_logit, a_log, dt_bias, + gate_lower_bound=-5.0, output_scale=1.0, + ) + + actual_state, actual_out = run_on_emulator( + kda_program, shape, + state=state.transpose(1, 2).reshape(shape.state_rows, shape.value_dim), + q=q, k=k, v=v, gate=gate, beta_logit=beta_logit, + a_log=a_log, dt_bias=dt_bias, + ) + + torch.testing.assert_close(actual_out, expected_out, rtol=2e-3, atol=2e-3) + torch.testing.assert_close( + actual_state.reshape(1, shape.key_dim, shape.value_dim).transpose(1, 2), + expected_state, rtol=2e-3, atol=2e-3, + ) + + +def test_decode_uses_only_instructions_that_already_exist(kda_program): + """Phase 0 introduces no new opcode. This test is the gate on that.""" + shape = KdaShape( + num_heads=1, key_dim=8, value_dim=2 * kda_program.mlen, conv_kernel=4 + ) + asm = kda_program.capture_asm(lambda: _emit_decode(kda_program, shape)) + body = [l for l in asm.splitlines() if l.strip() and not l.strip().startswith(";")] + assert body, "decode emitted nothing" + assert not any("V_FMA_VF" in l for l in body), ( + "Phase 0 must not depend on V_FMA_VF; that arrives in Task 7" + ) + mnemonics = {l.split()[0] for l in body} + allowed = { + "S_ADDI_INT", "S_LD_FP", "V_MUL_VF", "V_MUL_VV", "V_ADD_VV", "V_SUB_VV", + "C_LOOP_START", "C_LOOP_END", "S_MAP_FP_V", "S_MAP_V_FP", + } + assert mnemonics <= allowed, f"unexpected mnemonics: {mnemonics - allowed}" + + +def test_decode_records_its_phase0_instruction_count(kda_program): + """Not a budget -- a baseline. Task 8 must beat this number, and the + amount by which it does is the measured case for spending an opcode.""" + shape = KdaShape( + num_heads=1, key_dim=8, value_dim=kda_program.mlen, conv_kernel=4 + ) + asm = kda_program.capture_asm(lambda: _emit_decode(kda_program, shape)) + body = [l for l in asm.splitlines() if l.strip() and not l.strip().startswith(";")] + print(f"PHASE0_KDA_DECODE_STATIC_INSTR_PER_HEAD={len(body)}") + assert len(body) > 0 + + +def test_multi_head_decode_matches_the_golden(kda_program, run_on_emulator): + """Four heads: catches a per-head FPRAM offset bug that one head hides.""" + from aten.models.kda.reference import kda_decode_step + + torch.manual_seed(9) + shape = KdaShape( + num_heads=4, key_dim=8, value_dim=kda_program.mlen, conv_kernel=4 + ) + state = torch.randn(4, shape.value_dim, shape.key_dim, dtype=torch.float32) + q = torch.randn(4, shape.key_dim); k = torch.randn(4, shape.key_dim) + v = torch.randn(4, shape.value_dim); gate = torch.randn(4, shape.key_dim) + beta_logit = torch.randn(4); a_log = torch.randn(4) + dt_bias = torch.randn(4, shape.key_dim) + + _, expected_out = kda_decode_step( + state, q, k, v, gate, beta_logit, a_log, dt_bias, + gate_lower_bound=-5.0, output_scale=1.0, + ) + _, actual_out = run_on_emulator( + kda_program, shape, + state=state.transpose(1, 2).reshape(shape.state_rows, shape.value_dim), + q=q, k=k, v=v, gate=gate, beta_logit=beta_logit, + a_log=a_log, dt_bias=dt_bias, + ) + torch.testing.assert_close(actual_out, expected_out, rtol=2e-3, atol=2e-3) +``` + +Add `_emit_decode(program, shape)` as a small local helper in the test file that allocates the VRAM and FPRAM operands and calls `kda_decode_step_v0`. Build `run_on_emulator` from the harness `test_mamba2_reference.py` already uses; do not write a new emulator driver. + +- [ ] **Step 2: Run and confirm it fails** + +Run: `pytest aten/tests/test_kda_recurrent.py -k decode -v` +Expected: FAIL — no module `aten.plena.program_kda_recurrent`. + +- [ ] **Step 3: Implement the decode step** + +Create `aten/plena/program_kda_recurrent.py`: + +```python +"""The KDA recurrent decode step: two sweeps over the transposed state tile. + +Sweep 1 applies the per-key decay and accumulates the prediction. The error is +then a value-length vector. Sweep 2 applies the rank-1 update and reads the +output out of the *updated* state -- the reference reads the updated value +inside the same loop (``reduced += updated * q``), so update and read-out +cannot be reordered into separate passes over the state. + +State is read twice and written once. That is inherent to the delta rule, not +an artefact of this lowering. + +Phase 0 form: this module uses only instructions that already exist, with the +same copy + multiply + add triple ssm_decode_step_v0 uses today. Task 8 +converts it -- and the Mamba kernel -- onto V_FMA_VF together, so the win is a +measured before/after on two real programs rather than an estimate. +""" + +from __future__ import annotations + +from typing import Sequence + +from aten.models.kda.reference import KdaShape + +__all__ = ["ProgramKdaRecurrentMixin"] + + +class ProgramKdaRecurrentMixin: + def kda_decode_step_v0( + self, + *, + state, + scratch, + q_fp, + k_fp, + decay_fp, + beta_fp, + v, + o, + pred, + err, + shape: KdaShape, + output_scale_fp, + head_rows: Sequence[int] | None = None, + ): + if shape.value_dim % self.mlen: + raise ValueError( + f"value_dim {shape.value_dim} must be a multiple of " + f"the vector width {self.mlen}" + ) + blocks = shape.value_dim // self.mlen + heads = list(range(shape.num_heads)) if head_rows is None else list(head_rows) + k_dim = shape.key_dim + self.emit_comment( + f"stage=kda_state_update decode heads={len(heads)} key_dim={k_dim}" + ) + + for h in heads: + base = shape.head_base(h) + fp_base = h * k_dim + # One emitter row per (key, column block). Column block `c` of key + # `j` lives at row (base + j) * blocks + c, so a walk over keys at + # fixed `c` steps by `blocks` -- an arithmetic progression, which + # is what keeps each sweep a single hardware loop. + def key_rows(c: int) -> list[int]: + return [(base + j) * blocks + c for j in range(k_dim)] + + for c in range(blocks): + rows = key_rows(c) + acc = h * blocks + c + + # --- sweep 1: decay, then predict ------------------------ + # decay is already a single broadcast over an arithmetic row + # progression, so this one is a hardware loop even in Phase 0. + self.tile_row_mul_fp_broadcast( + state, decay_fp, rows=rows, fpram_offset=fp_base + ) + # predict: copy + multiply + add per key, via a scratch row. + # The scratch breaks the row progression, so this unrolls -- + # that is exactly the cost Task 7 removes. + self.vram_fill_zero(pred, rows=[acc]) + for i, r in enumerate(rows): + self.mamba_row_copy(scratch, 0, state, r) + self.tile_row_mul_fp_broadcast( + scratch, k_fp, rows=[0], fpram_offset=fp_base + i + ) + self.mamba_row_add(pred, acc, scratch, 0) + + # --- error: err[h, :] = beta_h * (v[h, :] - pred[h, :]) ------- + self.emit_comment(f"stage=kda_error head={h}") + err_rows = [h * blocks + c for c in range(blocks)] + for r in err_rows: + self.mamba_row_copy(err, r, v, r) + self.tile_row_sub(err, pred, rows=err_rows) + self.tile_row_mul_fp_broadcast(err, beta_fp, rows=err_rows, fpram_offset=h) + + for c in range(blocks): + rows = key_rows(c) + acc = h * blocks + c + + # --- sweep 2: update, then read out ----------------------- + # Every key's update touches only that key's row, so updating + # all rows and then reducing over all rows is identical to the + # reference's interleaved `updated = ...; reduced += updated*q`. + for i, r in enumerate(rows): + self.mamba_row_copy(scratch, 0, err, acc) + self.tile_row_mul_fp_broadcast( + scratch, k_fp, rows=[0], fpram_offset=fp_base + i + ) + self.mamba_row_add(state, r, scratch, 0) + + self.vram_fill_zero(o, rows=[acc]) + for i, r in enumerate(rows): + self.mamba_row_copy(scratch, 0, state, r) + self.tile_row_mul_fp_broadcast( + scratch, q_fp, rows=[0], fpram_offset=fp_base + i + ) + self.mamba_row_add(o, acc, scratch, 0) + + self.tile_row_mul_fp_broadcast( + o, output_scale_fp, rows=[acc], fpram_offset=0 + ) + + return o +``` + +Confirm `tile_row_sub`'s signature at `isa_tile_rows.py:167` — if it takes `(dst, src, rows=...)` the call above is correct; if it takes row offsets like `vram_add`, adapt. + +- [ ] **Step 4: Run the single-head numeric test** + +Run: `pytest aten/tests/test_kda_recurrent.py -k single_head -v` +Expected: PASS. + +Diagnosis if it fails: output right but state wrong → the update sweep ran after read-out. Both off by a constant → `output_scale` applied twice, or normalisation skipped. Output right for key 0 only → the FPRAM cursor is not advancing (check `fpram_offset` and Task 3 Step 4). + +- [ ] **Step 5: Run the loop-structure and multi-head tests** + +Run: `pytest aten/tests/test_kda_recurrent.py -k "hardware_loops or multi_head" -v` +Expected: PASS, both. + +- [ ] **Step 6: Commit** + +```bash +git add aten/plena/program_kda_recurrent.py aten/tests/test_kda_recurrent.py +git commit -m "feat(kda): two-sweep recurrent decode step on FMA hardware loops" +``` + +--- + +--- + +### Task 4: The full KDA layer + +**Files:** +- Create: `aten/plena/program_kda_layer.py` +- Create: `aten/ops/cpu/kda_ops.py`, `aten/ops/plena/kda_ops.py` +- Modify: `aten/native_ops.yaml` +- Modify: `aten/tests/test_kda_recurrent.py` + +**Interfaces:** +- Consumes: everything from Tasks 2 and 3; the projection and RMSNorm helpers in `program_mamba_common.py` and `program_tensors.py`; `kda_layer_reference` (Task 1). +- Produces: `ProgramKdaLayerMixin.kda_layer_v0(*, hidden, weights, shape, hbm) -> VRAMMatrixVar`, emitting the stage sequence `kda_qkv_projection` → `kda_decay_beta_projection` → `kda_conv1d` → `kda_normalize` → `kda_state_update` → `kda_output_gate` → `kda_output_gate_rmsnorm` → `kda_out_projection`. + +- [ ] **Step 1: Write the failing full-layer tests** + +Append to `aten/tests/test_kda_recurrent.py`: + +```python +def _random_kda_weights(shape): + """Non-degenerate weights. + + Every tensor is randn-derived on purpose. A connected + tests set W_kda_decay_a, W_kda_decay_b and W_kda_beta to zeros, which made + decay and the delta-rule subtraction constants -- so the two paths that + distinguish KDA from a plain linear attention were never exercised. + """ + d = shape.num_heads * shape.value_dim + return { + "W_kda_q": torch.randn(d, shape.num_heads * shape.key_dim) * 0.05, + "W_kda_k": torch.randn(d, shape.num_heads * shape.key_dim) * 0.05, + "W_kda_v": torch.randn(d, d) * 0.05, + "W_kda_gate": torch.randn(d, d) * 0.05, + "W_kda_out": torch.randn(d, d) * 0.05, + "W_kda_decay_a": torch.randn(d, shape.num_heads * shape.key_dim) * 0.05, + "W_kda_decay_b": torch.randn( + shape.num_heads * shape.key_dim, shape.num_heads * shape.key_dim + ) * 0.05, + "W_kda_beta": torch.randn(d, shape.num_heads) * 0.05, + "W_kda_norm": torch.ones(1, shape.value_dim), + "Q_CONV_WEIGHT": torch.randn(shape.num_heads * shape.key_dim, shape.conv_kernel) * 0.3, + "K_CONV_WEIGHT": torch.randn(shape.num_heads * shape.key_dim, shape.conv_kernel) * 0.3, + "V_CONV_WEIGHT": torch.randn(d, shape.conv_kernel) * 0.3, + "A_LOG": torch.randn(shape.num_heads) * 0.5, + "DT_BIAS": torch.randn(shape.num_heads, shape.key_dim) * 0.5, + } + + +def test_weights_used_by_the_layer_test_are_not_degenerate(): + """Guard the failure mode a descriptor-driven state engine invites.""" + w = _random_kda_weights(KdaShape(4, 16, 64, 4)) + for key in ("W_kda_decay_a", "W_kda_decay_b", "W_kda_beta", "A_LOG", "DT_BIAS"): + assert w[key].abs().max() > 1e-3, f"{key} is degenerate" + + +def test_full_layer_matches_the_golden_over_four_tokens(kda_program, run_layer_on_emulator): + """Four sequential tokens: catches state-carry bugs one token hides.""" + from aten.models.kda.reference import kda_layer_reference + + torch.manual_seed(11) + shape = KdaShape(num_heads=4, key_dim=16, value_dim=kda_program.mlen, conv_kernel=4) + weights = _random_kda_weights(shape) + tokens = torch.randn(4, shape.num_heads * shape.value_dim) + + expected = kda_layer_reference(weights, tokens, shape) + actual = run_layer_on_emulator(kda_program, weights, tokens, shape) + + torch.testing.assert_close(actual, expected, rtol=5e-3, atol=5e-3) +``` + +- [ ] **Step 2: Run and confirm it fails** + +Run: `pytest aten/tests/test_kda_recurrent.py -k full_layer -v` +Expected: FAIL — no module `aten.plena.program_kda_layer`. + +- [ ] **Step 3: Implement the layer** + +Create `aten/plena/program_kda_layer.py`, following the Mamba layer in `program_mamba_common.py` step for step: project, scatter into per-head rows, convolve with `kda_conv_state_roll_v0`, L2-normalise q and k with `kda_l2_normalize_v0`, move the per-head scalars into FPRAM with `ssm_decode_scalars_to_fpram_v0`, run `kda_decode_step_v0`, apply the sigmoid output gate, RMSNorm with `W_kda_norm`, project out. + +Emit a `stage=` marker comment at each boundary using the same convention as `mamba_stage_marker`, so the emulator's stage profiler attributes the work. + +- [ ] **Step 4: Register the op** + +Add a `kda_decode` entry to `aten/native_ops.yaml` following the `mamba2` entries `feat/mamba2-support` added, and create `aten/ops/cpu/kda_ops.py` and `aten/ops/plena/kda_ops.py` mirroring `aten/ops/cpu/mamba_ops.py` and `aten/ops/plena/mamba_ops.py`. + +- [ ] **Step 5: Run the full-layer test and confirm it passes** + +Run: `pytest aten/tests/test_kda_recurrent.py -k "full_layer or degenerate" -v` +Expected: PASS, both. + +- [ ] **Step 6: Run the whole compiler suite** + +Run: `pytest` +Expected: PASS, no regressions in the Mamba tests. + +- [ ] **Step 7: Commit** + +```bash +git add aten/plena/program_kda_layer.py aten/ops/cpu/kda_ops.py aten/ops/plena/kda_ops.py aten/native_ops.yaml aten/tests/test_kda_recurrent.py +git commit -m "feat(kda): full KDA layer lowering, validated against the golden over four tokens" +``` + +--- + +--- + +## Phase 1 — One instruction, measured against a real baseline + +Phase 0 gave us two working programs. Now add `V_FMA_VF` and convert both kernels onto it, so the instruction's value is a measured before/after rather than an estimate. + +--- + +### Task 5: `V_FMA_VF` in the emulator + +Adds the instruction to the reference implementation first, so the compiler has something to test against. + +**Files:** +- Modify: `transactional_emulator/src/op.rs` (enum near line 105, decode near line 448) +- Modify: `transactional_emulator/src/vector_machine.rs` (after `mul_scalar`, which ends at line 137) +- Modify: `transactional_emulator/src/accelerator/dispatch.rs` (dispatch near line 263, operand list near line 870) +- Test: inline `#[cfg(test)]` in `vector_machine.rs` and `op.rs` + +**Interfaces:** +- Consumes: nothing. +- Produces: `Opcode::V_FMA_VF { rd: u8, rs1: u8, rs2: u8, rmask: u8 }` decoded from opcode `0x3B`; `VectorMachine::fma_scalar(&self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32)` with semantics `vram[vd] = vram[vd] + vram[vs1] * f`. + +- [ ] **Step 1: Write the failing decode test** + +Add to the `#[cfg(test)] mod tests` block in `transactional_emulator/src/op.rs`: + +```rust +#[test] +fn v_fma_vf_decodes_like_v_mul_vf() { + // R-type layout: opcode[5:0], rd[9:6], rs1[13:10], rs2[17:14], rs3[21:18] + let word: u32 = 0x3B | (3 << 6) | (4 << 10) | (2 << 14); + match Opcode::decode(word) { + Opcode::V_FMA_VF { rd, rs1, rs2, rmask } => { + assert_eq!((rd, rs1, rs2, rmask), (3, 4, 2, 0)); + } + other => panic!("expected V_FMA_VF, got {other:?}"), + } +} +``` + +If the surrounding tests build words with a helper rather than by hand, use that helper instead and keep the same field values. + +- [ ] **Step 2: Run it and confirm it fails** + +Run: `cargo test --manifest-path transactional_emulator/Cargo.toml v_fma_vf_decodes_like_v_mul_vf` +Expected: FAIL — `V_FMA_VF` is not a variant of `Opcode`. + +- [ ] **Step 3: Add the enum variant and the decode arm** + +In `op.rs`, immediately after the `V_MUL_VF` variant: + +```rust + /// `Vector[rd] += Vector[rs1] * fp_reg` — fused multiply-add, one + /// rounding. Unlike every other V-type op, `rd` is read as well as written. + V_FMA_VF { + rd: u8, + rs1: u8, + rs2: u8, + rmask: u8, + }, +``` + +In the decode `match`, after the `0x3A` arm: + +```rust + 0x3B => Self::V_FMA_VF { + rd, + rs1, + rs2, + rmask: rs3, + }, +``` + +- [ ] **Step 4: Run the decode test and confirm it passes** + +Run: `cargo test --manifest-path transactional_emulator/Cargo.toml v_fma_vf_decodes_like_v_mul_vf` +Expected: PASS. + +- [ ] **Step 5: Write the failing execution test** + +Add to the `#[cfg(test)] mod tests` block in `vector_machine.rs`. Construct the machine and read/write rows using whatever the nearest existing `#[tokio::test]` in that file already does — do not invent helpers: + +```rust +#[tokio::test] +async fn fma_scalar_accumulates_into_the_destination() { + let vm = /* same construction as the neighbouring vector_machine test */; + /* write [1.0, 2.0, 3.0, 4.0] at row 0 (destination) */ + /* write [10.0, 20.0, 30.0, 40.0] at row 1 (source) */ + vm.fma_scalar(0, 1, 0.5, 0, u32::MAX).await; + /* assert row 0 == [6.0, 12.0, 18.0, 24.0] */ + /* assert row 1 is unchanged == [10.0, 20.0, 30.0, 40.0] */ +} +``` + +- [ ] **Step 6: Run it and confirm it fails** + +Run: `cargo test --manifest-path transactional_emulator/Cargo.toml fma_scalar_accumulates` +Expected: FAIL — no method `fma_scalar`. + +- [ ] **Step 7: Implement `fma_scalar`** + +In `vector_machine.rs`, directly after `mul_scalar`. It mirrors `mul_scalar` exactly, except that it also reads `vd`: + +```rust + /// `Vector[vd] += Vector[vs1] * f`. + /// + /// Mirrors `mul_scalar` but reads the destination too. The accumulate is + /// what lets one instruction carry a rank-1 state update or a state + /// contraction that otherwise costs copy + multiply + add, and — because + /// it removes the scratch row — lets a whole sweep become one arithmetic + /// row progression, which the compiler turns into a hardware loop. + pub(crate) async fn fma_scalar(&self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32) { + let a = self.vram.read(vs1).await; + let d = self.vram.read(vd).await; + if rmask == 0 { + let c = QuantTensor::quantize( + d.as_tensor() + a.as_tensor() * (f as f64), + d.data_type(), + ); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } else { + let result = d.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let addend = a.as_tensor().narrow(0, start, end - start) * (f as f64); + let updated = &sliced + &addend; + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, d.data_type()); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } + } +``` + +- [ ] **Step 8: Run the execution test and confirm it passes** + +Run: `cargo test --manifest-path transactional_emulator/Cargo.toml fma_scalar_accumulates` +Expected: PASS. + +- [ ] **Step 9: Wire it into dispatch** + +In `dispatch.rs`, after the `V_MUL_VF` arm: + +```rust + op::Opcode::V_FMA_VF { + rd, + rs1, + rs2, + rmask, + } => { + let mask = self.resolve_v_mask(*rmask); + self.v_machine + .fma_scalar( + self.reg_file.read_gp(*rd), + self.reg_file.read_gp(*rs1), + self.reg_file.read_fp(*rs2).into(), + *rmask, + mask, + ) + .await; + } +``` + +In the VRAM-read operand list near line 870, add: + +```rust + | op::Opcode::V_FMA_VF { rs1, .. } +``` + +- [ ] **Step 10: Confirm the whole suite and lints are clean** + +Run: `cargo test --manifest-path transactional_emulator/Cargo.toml` +Expected: PASS, no new failures. + +Run: `cargo clippy --manifest-path transactional_emulator/Cargo.toml --all-targets -- -D warnings` +Expected: no warnings. + +- [ ] **Step 11: Commit** + +```bash +git add transactional_emulator/src/op.rs transactional_emulator/src/vector_machine.rs transactional_emulator/src/accelerator/dispatch.rs +git commit -m "feat(isa): V_FMA_VF, a fused broadcast multiply-add on the vector machine" +``` + +--- + +--- + +### Task 6: `V_FMA_VF` in the ISA definition and assembler + +**Files:** +- Modify: `doc/operation.svh` (`CUSTOM_ISA_OPCODE`, after `S_MAP_FP_V = 6'h3A`) +- Modify: `doc/plena_isa_spec.md` (after the `### V_MUL_VF` section, near line 290) +- Modify: `assembler/assembly_to_binary.py` +- Test: `aten/tests/test_v_fma_vf.py` + +**Interfaces:** +- Consumes: `Opcode::V_FMA_VF` decoding `0x3B` (Task 5). +- Produces: assembly mnemonic `V_FMA_VF rd, rs1, fp2, rmask` encoding to a word whose low 6 bits are `0x3B`. + +- [ ] **Step 1: Write the failing encoding test** + +Create `aten/tests/test_v_fma_vf.py`: + +```python +"""V_FMA_VF encodes into the same R-type slots as V_MUL_VF.""" + +import tempfile +from pathlib import Path + +from assembler.assembly_to_binary import AssemblyToBinary +from assembler.parser import parse_asm_file + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _assemble(text: str) -> list[int]: + with tempfile.NamedTemporaryFile("w", suffix=".asm", delete=False) as handle: + handle.write(text) + path = Path(handle.name) + try: + converter = AssemblyToBinary( + str(REPO_ROOT / "doc" / "operation.svh"), + str(REPO_ROOT / "doc" / "configuration.svh"), + ) + return [converter._convert_to_binary(i) for i in parse_asm_file(str(path))] + finally: + path.unlink(missing_ok=True) + + +def test_v_fma_vf_opcode_and_operand_slots(): + (word,) = _assemble("V_FMA_VF gp3, gp4, f2, 0\n") + assert word & 0x3F == 0x3B + assert (word >> 6) & 0xF == 3 # rd + assert (word >> 10) & 0xF == 4 # rs1 + assert (word >> 14) & 0xF == 2 # fp2 + assert (word >> 18) & 0xF == 0 # rmask + + +def test_v_fma_vf_matches_v_mul_vf_layout(): + (fma,) = _assemble("V_FMA_VF gp5, gp6, f1, 0\n") + (mul,) = _assemble("V_MUL_VF gp5, gp6, f1, 0\n") + assert fma >> 6 == mul >> 6, "only the opcode field may differ" +``` + +- [ ] **Step 2: Run it and confirm it fails** + +Run: `pytest aten/tests/test_v_fma_vf.py -v` +Expected: FAIL — the assembler does not know `V_FMA_VF`. + +- [ ] **Step 3: Add the opcode to `doc/operation.svh`** + +In `CUSTOM_ISA_OPCODE`, put a comma after `S_MAP_FP_V` and append: + +```systemverilog + V_SOFTPLUS_V = 6'h39, + S_MAP_FP_V = 6'h3A, + V_FMA_VF = 6'h3B // Vector[rd] += Vector[rs1] * fp_reg +``` + +- [ ] **Step 4: Teach the assembler the mnemonic** + +In `assembler/assembly_to_binary.py`, find where `V_MUL_VF` is handled (the 4-operand `rd, rs1, fp2, rmask` form) and add `V_FMA_VF` to exactly that handling. Do not add a new code path — if `V_MUL_VF` appears in a name list, append `"V_FMA_VF"` to that list. + +- [ ] **Step 5: Run the encoding tests and confirm they pass** + +Run: `pytest aten/tests/test_v_fma_vf.py -v` +Expected: PASS, both tests. + +- [ ] **Step 6: Document it in the ISA spec** + +In `doc/plena_isa_spec.md`, immediately after `### V_MUL_VF`: + +````markdown +### V_FMA_VF + +**Format:** `V_FMA_VF rd, rs1, fp2, rmask` + +**Operation:** `Vector[gp_reg] += Vector[gp_reg] * fp_reg` + +**Description:** + +Fused broadcast multiply-add. Multiplies a (VLEN, 1) vector by a scalar from +`fp_reg` and accumulates into the destination vector, with a single +rounding. + +Unlike every other V-type instruction, `rd` is a **source as well as a +destination**. The recurrent-state kernels are the motivating case: a rank-1 +state update (`S[n,:] += B[n] * xs[:]`) and a state contraction +(`y[:] += C[n] * S[n,:]`) are otherwise three instructions each — copy, +multiply, add — plus a scratch row. Removing the scratch is what lets a whole +sweep be expressed as one arithmetic row progression, and therefore as one +hardware loop rather than an unrolled block. The same shape appears in +attention and FFN accumulation. + +**Operands:** +- `rd`: Register containing the destination/accumulator Vector SRAM address +- `rs1`: Register containing the source Vector SRAM address +- `fp2`: FP register index holding the broadcast scalar +- `rmask`: `0` = all lanes; otherwise use `V_MASK_REG` + +**Example:** +```asm +S_ADDI_INT gp3, gp0, 512 ; accumulator row +S_ADDI_INT gp4, gp0, 1024 ; source row +V_FMA_VF gp3, gp4, f2, 0 ; Vector[512..] += Vector[1024..] * f2 +``` +```` + +- [ ] **Step 7: Confirm the cross-repo opcode tables agree** + +Run `grep -n "6'h3B" doc/operation.svh` in the Compiler and `grep -n "0x3B" transactional_emulator/src/op.rs` in the Simulator. +Expected: both name `V_FMA_VF`. + +- [ ] **Step 8: Commit** + +```bash +git add doc/operation.svh doc/plena_isa_spec.md assembler/assembly_to_binary.py aten/tests/test_v_fma_vf.py +git commit -m "feat(isa): declare and assemble V_FMA_VF at 0x3B" +``` + +--- + +--- + +### Task 7: The FMA row-sweep emitter + +The whole plan's static-footprint claim rests on this task emitting a hardware loop. + +**Files:** +- Modify: `aten/plena/isa_tile_rows.py` — add `_emit_tile_row_fma` next to `_emit_tile_row_fp_scalar` (line 371), and the two public wrappers next to `tile_row_mul_fp_asm` (line 513) and `tile_row_mul_fp_broadcast_asm` (line 534) +- Test: `aten/tests/test_v_fma_vf.py` (extend) + +**Interfaces:** +- Consumes: the `V_FMA_VF` mnemonic (Task 6); the existing `self._reg`, `self._row_progression`, `self._emit`, `self.mlen`, `IsaBuilder`, `gp`, `fp`. +- Produces: + - `tile_row_fma_fp_asm(self, dst_addr: int, src_addr: int, row_map: list[tuple[int, int, int]]) -> str` where each entry is `(dst_row, src_row, fpram_addr)`. + - `tile_row_fma_fp_broadcast_asm(self, dst_addr: int, src_addr: int, fpram_base: int, dst_rows: list[int], src_rows: list[int]) -> str` — walks `fpram_base + i` for entry `i`. + + Semantics: `dst[dst_row] += src[src_row] * FPRAM[fpram_addr]`, entries applied in order. + +- [ ] **Step 1: Write the failing emitter tests** + +Append to `aten/tests/test_v_fma_vf.py`. `tile_program` should be built the same way the nearest existing test that calls `tile_row_mul_fp_broadcast_asm` builds its program object: + +```python +def _body(asm: str) -> list[str]: + return [l.strip() for l in asm.splitlines() + if l.strip() and not l.strip().startswith(";")] + + +def test_fma_sweep_over_a_progression_emits_a_hardware_loop(tile_program): + """A 128-row sweep must be a loop, not 128 unrolled blocks. + + This is the whole static-footprint argument: the copy/multiply/add + predecessor needed a scratch row per step, which broke the progression and + forced the unrolled path. + """ + asm = tile_program.tile_row_fma_fp_broadcast_asm( + dst_addr=0, + src_addr=8192, + fpram_base=16, + dst_rows=list(range(128)), + src_rows=[0] * 128, + ) + body = _body(asm) + assert sum("C_LOOP_START" in l for l in body) == 1 + assert sum("C_LOOP_END" in l for l in body) == 1 + assert sum("V_FMA_VF" in l for l in body) == 1, "loop body holds one FMA" + assert len(body) < 12, f"expected a compact loop, got {len(body)} instructions" + + +def test_fma_sweep_holds_a_constant_destination(tile_program): + """The contraction walks the source and pins the destination.""" + asm = tile_program.tile_row_fma_fp_broadcast_asm( + dst_addr=0, + src_addr=8192, + fpram_base=16, + dst_rows=[3] * 64, + src_rows=list(range(64)), + ) + body = _body(asm) + assert sum("C_LOOP_START" in l for l in body) == 1 + assert sum("V_FMA_VF" in l for l in body) == 1 + + +def test_fma_uses_no_scratch_and_no_separate_multiply(tile_program): + asm = tile_program.tile_row_fma_fp_broadcast_asm( + dst_addr=0, src_addr=8192, fpram_base=16, + dst_rows=[0, 1, 2, 3], src_rows=[0, 1, 2, 3], + ) + body = _body(asm) + assert not any("V_ADD_VV" in l for l in body) + assert not any("V_MUL_VF" in l for l in body) + + +def test_fma_rejects_mismatched_row_counts(tile_program): + import pytest + with pytest.raises(ValueError, match="row counts"): + tile_program.tile_row_fma_fp_broadcast_asm( + dst_addr=0, src_addr=8192, fpram_base=16, + dst_rows=[0, 1], src_rows=[0], + ) +``` + +- [ ] **Step 2: Run and confirm they fail** + +Run: `pytest aten/tests/test_v_fma_vf.py -k fma -v` +Expected: FAIL — no attribute `tile_row_fma_fp_broadcast_asm`. + +- [ ] **Step 3: Implement the emitter** + +In `aten/plena/isa_tile_rows.py`, immediately after `_emit_tile_row_fp_scalar`. This mirrors that method exactly, with a third pointer and two independent row progressions: + +```python + def _emit_tile_row_fma( + self, + dst_addr: int, + src_addr: int, + row_map: list[tuple[int, int, int]], + ) -> str: + """`dst[d] += src[s] * FPRAM[f]` for each `(d, s, f)` in `row_map`. + + Four pointers, so four GP registers. When both row walks and the FPRAM + walk are arithmetic progressions this collapses to one hardware loop -- + the case that matters, since every recurrent sweep either steps a state + tile with a pinned accumulator or the reverse. A step of 0 is a valid + progression, which is how the pinned side is expressed. + """ + gp_regs = self._reg.allocate_gp(4) + gp_dst, gp_src, gp_fp, gp_loop = gp_regs + try: + asm = IsaBuilder().comment( + f"Tile Row FMA: VRAM[{dst_addr}] += VRAM[{src_addr}] * FPRAM" + ) + dst_rows = [d for d, _, _ in row_map] + src_rows = [s for _, s, _ in row_map] + fp_addrs = [f for _, _, f in row_map] + + dst_prog = self._row_progression(dst_rows) + src_prog = self._row_progression(src_rows) + fp_prog = self._row_progression(fp_addrs) + + if dst_prog is not None and src_prog is not None and fp_prog is not None: + dst_start, count, dst_step = dst_prog + src_start, _, src_step = src_prog + fp_start, _, fp_step = fp_prog + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), dst_addr + dst_start * self.mlen) + asm.instr("S_ADDI_INT", gp(gp_src), gp(0), src_addr + src_start * self.mlen) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fp_start) + asm.instr("C_LOOP_START", gp(gp_loop), count) + asm.instr("S_LD_FP", fp(1), gp(gp_fp), 0) + asm.instr("V_FMA_VF", gp(gp_dst), gp(gp_src), fp(1), 0) + if dst_step: + asm.instr("S_ADDI_INT", gp(gp_dst), gp(gp_dst), dst_step * self.mlen) + if src_step: + asm.instr("S_ADDI_INT", gp(gp_src), gp(gp_src), src_step * self.mlen) + if fp_step: + asm.instr("S_ADDI_INT", gp(gp_fp), gp(gp_fp), fp_step) + asm.instr("C_LOOP_END", gp(gp_loop)) + else: + for dst_row, src_row, fpram_addr in row_map: + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), dst_addr + dst_row * self.mlen) + asm.instr("S_ADDI_INT", gp(gp_src), gp(0), src_addr + src_row * self.mlen) + asm.instr("S_ADDI_INT", gp(gp_fp), gp(0), fpram_addr) + asm.instr("S_LD_FP", fp(1), gp(gp_fp), 0) + asm.instr("V_FMA_VF", gp(gp_dst), gp(gp_src), fp(1), 0) + + return self._emit(asm) + finally: + self._reg.free_gp(gp_regs) + + def tile_row_fma_fp_asm( + self, + dst_addr: int, + src_addr: int, + row_map: list[tuple[int, int, int]], + ) -> str: + return self._emit_tile_row_fma(dst_addr, src_addr, row_map) + + def tile_row_fma_fp_broadcast_asm( + self, + dst_addr: int, + src_addr: int, + fpram_base: int, + dst_rows: list[int], + src_rows: list[int], + ) -> str: + """Walk one FPRAM slot per row pair, starting at `fpram_base`.""" + if len(dst_rows) != len(src_rows): + raise ValueError( + f"tile_row_fma row counts differ: " + f"{len(dst_rows)} destinations, {len(src_rows)} sources" + ) + row_map = [ + (d, s, fpram_base + i) + for i, (d, s) in enumerate(zip(dst_rows, src_rows)) + ] + return self._emit_tile_row_fma(dst_addr, src_addr, row_map) +``` + +Confirm `fp(1)` is the right scratch FP register by checking that `_emit_tile_row_fp_scalar` uses it; if it uses a different one, match that. + +- [ ] **Step 4: Confirm `_row_progression` accepts a constant list** + +Run: + +```bash +python -c " +from aten.plena.isa_tile_rows import IsaTileRowsMixin as M +print(M._row_progression(None, [3]*8)) +" +``` +Expected: a `(start, count, step)` tuple with `step == 0`. If it returns `None` for a constant list, fix `_row_progression` to treat a constant run as a step-0 progression, and add a test for that — the pinned-accumulator sweeps depend on it. Adjust the import path / call convention to whatever the class actually is. + +- [ ] **Step 5: Add the `VRAMMatrixVar`-level wrapper** + +Next to `tile_row_mul_fp_broadcast` (line 258), following that method's signature style: + +```python + def tile_row_fma_fp_broadcast( + self, + dst, + src, + fp, + *, + dst_rows: list[int], + src_rows: list[int], + fpram_offset: int, + ): + """`dst[dst_rows[i]] += src[src_rows[i]] * FPRAM[fp + fpram_offset + i]`.""" + return self.tile_row_fma_fp_broadcast_asm( + dst_addr=dst.address, + src_addr=src.address, + fpram_base=fp.address + fpram_offset, + dst_rows=dst_rows, + src_rows=src_rows, + ) +``` + +Match how `tile_row_mul_fp_broadcast` resolves `.address` on its arguments and follow it exactly. + +- [ ] **Step 6: Run and confirm all four tests pass** + +Run: `pytest aten/tests/test_v_fma_vf.py -v` +Expected: PASS, six tests total. + +- [ ] **Step 7: Commit** + +```bash +git add aten/plena/isa_tile_rows.py aten/tests/test_v_fma_vf.py +git commit -m "feat(lowering): FMA row sweeps that collapse to a hardware loop" +``` + +--- + +--- + +### Task 8: Collapse both decode steps onto FMA sweeps + +Proves the emitter on a kernel that already has a passing numeric test. + +**Files:** +- Modify: `aten/plena/program_ssm_recurrent.py:143-211` (`ssm_decode_step_v0`) +- Modify: `aten/plena/program_kda_recurrent.py` (`kda_decode_step_v0`, from Task 3) +- Create: `aten/tests/test_instruction_budget.py` +- Test: `aten/tests/test_mamba2_reference.py` and `aten/tests/test_kda_recurrent.py` (both existing, both must stay green) + +**Interfaces:** +- Consumes: `tile_row_fma_fp_broadcast` (Task 7). +- Produces: `ssm_decode_step_v0` and `kda_decode_step_v0`, both with unchanged numerical output. `kda_decode_step_v0` keeps its signature; its `scratch` parameter becomes unused and is dropped in Step 8. Both emit hardware loops instead of unrolled triples. + +Doing both kernels in one task is deliberate: they share the emitter, so a bug in it shows up twice, and the measured before/after covers both programs at once. + +- [ ] **Step 1: Record the current static count as a baseline** + +Create `aten/tests/test_instruction_budget.py`: + +```python +"""Static instruction-count gates. + +The only guard against a lowering change that is numerically correct but +quietly unrolls a sweep that used to be a hardware loop. Unrolling is what +makes a whole-model program image tens of megabytes. + +Budgets are MEASURED, not derived: run the test, read the reported count, +set the constant ~10% above it, and note the date. Raising a budget requires +a line in docs/superpowers/plans/2026-08-25-static-mamba-kda.md saying why. +""" + +# Measured 2026-08-25 after Task 8. Three hardware loops per head plus +# per-head scalar setup. +MAMBA_DECODE_STATIC_INSTR_PER_HEAD_MAX = 40 + + +def test_mamba_decode_step_static_instruction_count(mamba_decode_program): + per_head = ( + mamba_decode_program.instruction_count // mamba_decode_program.num_heads + ) + assert per_head <= MAMBA_DECODE_STATIC_INSTR_PER_HEAD_MAX, ( + f"{per_head} static instructions per head exceeds " + f"{MAMBA_DECODE_STATIC_INSTR_PER_HEAD_MAX}; a sweep probably fell off " + f"the hardware-loop path" + ) +``` + +Build the `mamba_decode_program` fixture from whatever `aten/tests/test_mamba2_reference.py` already uses to construct a decode program. + +- [ ] **Step 2: Run it and confirm it fails, and record the number** + +Run: `pytest aten/tests/test_instruction_budget.py -v` +Expected: FAIL, reporting roughly 1,100–1,200 static instructions per head (the unrolled `9 * state_size` pattern). Write the observed number into a comment in the test file. + +- [ ] **Step 3: Rewrite the update and contraction loops** + +In `ssm_decode_step_v0`, replace the whole per-head body: + +```python + for h in heads: + group = h // shape.heads_per_group + base = h * n_state + state_rows = list(range(base, base + n_state)) + + # xs = dt_h * x[h] -- fold dt into x once per head so the per-row + # scalars below are only B[g, n]. + self.mamba_row_copy(scratch, 0, x, h) + self.tile_row_mul_fp_broadcast(scratch, dt_fp, rows=[0], fpram_offset=h) + + # decay: h[h, n, :] *= dA_h (one hardware loop, fp step 0) + self.tile_row_mul_fp_broadcast(state, da_fp, rows=state_rows, fpram_offset=h) + + # rank-1 update: h[h, n, :] += B[g, n] * xs + # dst walks the state tile, src is pinned to scratch row 0. + self.tile_row_fma_fp_broadcast( + state, scratch, b_fp, + dst_rows=state_rows, + src_rows=[0] * n_state, + fpram_offset=group * n_state, + ) + + # contraction: y[h] = sum_n C[g, n] * h[h, n, :] + # dst is pinned to y row h, src walks the state tile. + self.vram_fill_zero(y, rows=[h]) + self.tile_row_fma_fp_broadcast( + y, state, c_fp, + dst_rows=[h] * n_state, + src_rows=state_rows, + fpram_offset=group * n_state, + ) + + # skip: y[h] += D_h * x[h] + self.emit_comment(mamba_stage_marker("mamba_skip", f"head={h}")) + self.tile_row_fma_fp_broadcast( + y, x, d_fp, dst_rows=[h], src_rows=[h], fpram_offset=h + ) + + return y +``` + +- [ ] **Step 4: Run the numeric test and confirm the answer is unchanged** + +Run: `pytest aten/tests/test_mamba2_reference.py -v` +Expected: PASS, unchanged. + +If it fails: the maths is identical, so the bug is ordering. Check that `b_fp` / `c_fp` are laid out `n`-major within each group (the docstring at `program_ssm_recurrent.py:157-165` states this), and that `vram_fill_zero(y, rows=[h])` still runs *before* the contraction. + +- [ ] **Step 5: Run the budget test and confirm it passes** + +Run: `pytest aten/tests/test_instruction_budget.py -v` +Expected: PASS. If the measured count is above 40 but the emitted asm does contain 3 `C_LOOP_START`s per head, raise the constant to measured + 10% and note it. + +- [ ] **Step 6: Confirm the sweeps really are loops** + +Run: `pytest aten/tests/test_mamba2_reference.py -v -k asm` or dump one head's assembly and check by eye: + +```bash +grep -c "C_LOOP_START" +``` +Expected: 3 per head. If it is 0, `_row_progression` rejected one of the row lists — most likely the pinned `[0] * n_state` (see Task 7 Step 4). + +- [ ] **Step 7: Convert the KDA decode step** + +In `aten/plena/program_kda_recurrent.py`, replace each of the three unrolled +`for i, r in enumerate(rows)` blocks with one FMA sweep. `dst` below is the +pinned accumulator row list, `rows` the state-tile walk: + +```python + # sweep 1: predict + self.vram_fill_zero(pred, rows=[acc]) + self.tile_row_fma_fp_broadcast( + pred, state, k_fp, + dst_rows=[acc] * k_dim, src_rows=rows, fpram_offset=fp_base, + ) + + # sweep 2: update, then read out + self.tile_row_fma_fp_broadcast( + state, err, k_fp, + dst_rows=rows, src_rows=[acc] * k_dim, fpram_offset=fp_base, + ) + self.vram_fill_zero(o, rows=[acc]) + self.tile_row_fma_fp_broadcast( + o, state, q_fp, + dst_rows=[acc] * k_dim, src_rows=rows, fpram_offset=fp_base, + ) +``` + +The decay sweep is already a broadcast over a progression and does not change. + +- [ ] **Step 8: Run the KDA tests and confirm the answer is unchanged** + +Run: `pytest aten/tests/test_kda_recurrent.py -v` +Expected: PASS, unchanged numerically. + +`test_decode_uses_only_instructions_that_already_exist` from Task 3 will now +fail, correctly — Phase 1 is where `V_FMA_VF` enters. Delete that test and add +`V_FMA_VF` to the allowed set in nothing else; the count baseline test from +Task 3 stays and should now report a much smaller number. + +Also drop the now-unused `scratch` parameter from `kda_decode_step_v0` and from +its call sites in `program_kda_layer.py`. + +- [ ] **Step 9: Record the measured Phase 0 → Phase 1 improvement** + +Run: `pytest aten/tests/test_kda_recurrent.py -k records_its_phase0 -s` +Expected: prints the new per-head static count. + +Write both numbers — the Task 3 baseline and this one — into a comment in +`aten/tests/test_instruction_budget.py`. **This ratio is the measured +justification for spending opcode `0x3B`.** If it is below about 3×, say so +plainly rather than proceeding on the assumption that the instruction earned +its place. + +- [ ] **Step 10: Commit** + +```bash +git add aten/plena/program_ssm_recurrent.py aten/plena/program_kda_recurrent.py aten/plena/program_kda_layer.py aten/tests/test_instruction_budget.py aten/tests/test_kda_recurrent.py +git commit -m "perf: fold both recurrent decode steps into FMA hardware loops" +``` + +--- + +--- + +## Phase 2 — KDA chunked prefill + +Decode alone is half the function. Prefill for KDA is the half that had never been mapped onto the PLENA Matrix Engine at all. + +This phase mirrors what `feat/mamba2-support` already did for Mamba-2 in `aten/plena/program_ssd.py`, whose module docstring (lines 1-45) is the model to follow: identify which chunk-level products drop onto the existing flash-attention Matrix templates unchanged, and isolate the one that does not. + +**The structural difference from Mamba.** Mamba-2's chunked form is four matmuls, three of which map directly onto `M_MM` / `M_TMM`. The gated delta rule does not chunk that simply: within a chunk the update is a product of rank-1 projectors `∏(I − β_t k_t k_tᵀ)`, and collapsing it requires the WY / UT transform, whose core is the inverse of a `[chunk, chunk]` lower-triangular matrix. That is why `chunk = 16` matters — a 16×16 triangular solve by forward substitution is 16 rank-1 updates, which the vector unit does directly. Do not raise the chunk size without redoing this analysis. + +--- + +### Task 9: KDA chunk primitives — decay cumsum and the UT transform + +**Files:** +- Create: `aten/plena/program_kda_chunk.py` +- Test: `aten/tests/test_kda_prefill.py` + +**Interfaces:** +- Consumes: `KdaShape` (Task 1); `ssd_lower_triangular_ones` and `ssd_chunk_cumsum_v0` patterns from `aten/plena/program_ssd.py:177-269`; `tile_row_fma_fp_broadcast` (Task 7). +- Produces `ProgramKdaChunkMixin` with: + - `kda_chunk_decay_cumsum_v0(decay_log, cs_fp, *, chunk, shape) -> FPVar` — per-key cumulative log-decay within the chunk, into FPRAM. + - `kda_ut_transform_v0(k, beta_fp, t_out, scratch, *, chunk, shape) -> VRAMMatrixVar` — builds `T = (I + tril(diag(β) K Kᵀ, −1))⁻¹ diag(β)`, `[chunk, chunk]`, by forward substitution. + +- [ ] **Step 1: Write the failing UT-transform test against a dense reference** + +Create `aten/tests/test_kda_prefill.py`: + +```python +"""KDA chunked prefill: chunk primitives, then the full prefill layer.""" + +import torch + +from aten.models.kda.reference import KdaShape + + +def _dense_ut(k: torch.Tensor, beta: torch.Tensor) -> torch.Tensor: + """T = (I + tril(diag(beta) K K^T, -1))^-1 diag(beta), computed densely.""" + chunk = k.shape[0] + a = torch.tril(torch.diag(beta) @ (k @ k.T), diagonal=-1) + return torch.linalg.inv(torch.eye(chunk) + a) @ torch.diag(beta) + + +def test_ut_transform_matches_a_dense_inverse(kda_program, run_on_emulator): + torch.manual_seed(3) + chunk = 16 + shape = KdaShape(num_heads=1, key_dim=8, value_dim=kda_program.mlen, conv_kernel=4) + k = torch.nn.functional.normalize(torch.randn(chunk, shape.key_dim), dim=-1) + beta = torch.sigmoid(torch.randn(chunk)) + + expected = _dense_ut(k, beta) + actual = run_on_emulator( + kda_program, "kda_ut_transform_v0", chunk=chunk, shape=shape, k=k, beta=beta + ) + torch.testing.assert_close(actual, expected, rtol=5e-3, atol=5e-3) + + +def test_ut_transform_is_lower_triangular(kda_program, run_on_emulator): + """Forward substitution must never write above the diagonal.""" + torch.manual_seed(4) + chunk = 16 + shape = KdaShape(num_heads=1, key_dim=8, value_dim=kda_program.mlen, conv_kernel=4) + k = torch.nn.functional.normalize(torch.randn(chunk, shape.key_dim), dim=-1) + beta = torch.sigmoid(torch.randn(chunk)) + actual = run_on_emulator( + kda_program, "kda_ut_transform_v0", chunk=chunk, shape=shape, k=k, beta=beta + ) + assert torch.triu(actual, diagonal=1).abs().max() < 1e-6 +``` + +- [ ] **Step 2: Run and confirm it fails** + +Run: `pytest aten/tests/test_kda_prefill.py -v` +Expected: FAIL — no module `aten.plena.program_kda_chunk`. + +- [ ] **Step 3: Implement the UT transform by forward substitution** + +Create `aten/plena/program_kda_chunk.py`. The solve is row-by-row: row `i` of `T` is +`β_i·e_i − Σ_{j i` — the inner loop bound is `range(i)`, not `range(chunk)`. + +- [ ] **Step 5: Commit** + +```bash +git add aten/plena/program_kda_chunk.py aten/tests/test_kda_prefill.py +git commit -m "feat(kda): chunk decay cumsum and the UT transform by forward substitution" +``` + +--- + +### Task 10: The KDA chunked prefill layer + +**Files:** +- Create: `aten/plena/program_kda_prefill.py` +- Modify: `aten/tests/test_kda_prefill.py`, `analytic_models/reference/kimi_k3_kda.py` (add `kda_prefill_reference`) + +**Interfaces:** +- Consumes: `ProgramKdaChunkMixin` (Task 9); `kda_decode_step_v0` (Tasks 3, 8) for the equivalence test; the Matrix helpers `ssd_transposed_projection_v0` and `ssd_chunk_head_v0` patterns (`program_ssd.py:321-436`). +- Produces: `ProgramKdaPrefillMixin.kda_prefill_v0(*, hidden, weights, shape, chunk, hbm) -> tuple[VRAMMatrixVar, VRAMMatrixVar]` returning `(output, final_state)`. + +**Chunk size is an ABI constant.** `chunk = 16` for prefill; decode's physical row capacity is `BLEN = 4`. The two must not be conflated: giving decode prefill's row pitch creates 4–16× zero-padded compute and address holes. + +- [ ] **Step 1: Write the failing equivalence test** + +The strongest available correctness statement, and the one `program_ssd.py` uses for Mamba: **a chunked prefill over S tokens must produce the same final state and the same outputs as S sequential decode steps.** + +```python +def test_prefill_equals_sequential_decode(kda_program, run_layer_on_emulator): + """The chunked path and the recurrent path must agree exactly. + + This is the only test that can catch a UT-transform sign error, a decay + cumsum off-by-one, or an intra/inter-chunk mixup -- each of which produces + plausible-looking output on its own. + """ + from aten.models.kda.reference import kda_layer_reference + + torch.manual_seed(13) + shape = KdaShape(num_heads=2, key_dim=16, value_dim=kda_program.mlen, conv_kernel=4) + weights = _random_kda_weights(shape) + tokens = torch.randn(32, shape.num_heads * shape.value_dim) # 2 chunks of 16 + + seq_out, seq_state = kda_layer_reference( + weights, tokens, shape, return_state=True + ) + pre_out, pre_state = run_layer_on_emulator( + kda_program, weights, tokens, shape, mode="prefill", chunk=16 + ) + + torch.testing.assert_close(pre_out, seq_out, rtol=5e-3, atol=5e-3) + torch.testing.assert_close(pre_state, seq_state, rtol=5e-3, atol=5e-3) + + +def test_prefill_then_decode_carries_state(kda_program, run_layer_on_emulator): + """Prefill 32 tokens, then decode 4 more, against 36 sequential steps.""" + from aten.models.kda.reference import kda_layer_reference + + torch.manual_seed(17) + shape = KdaShape(num_heads=2, key_dim=16, value_dim=kda_program.mlen, conv_kernel=4) + weights = _random_kda_weights(shape) + tokens = torch.randn(36, shape.num_heads * shape.value_dim) + + expected, _ = kda_layer_reference(weights, tokens, shape, return_state=True) + actual = run_layer_on_emulator( + kda_program, weights, tokens, shape, + mode="prefill_then_decode", chunk=16, prefill_len=32, + ) + torch.testing.assert_close(actual, expected, rtol=5e-3, atol=5e-3) +``` + +Import `_random_kda_weights` from `aten/tests/test_kda_recurrent.py` rather than +duplicating it; extend `kda_layer_reference` with `return_state` and a +`kda_prefill_reference` in the Simulator's reference module first, with its own +self-test, exactly as Task 1 Step 6 required. + +- [ ] **Step 2: Run and confirm it fails** + +Run: `pytest aten/tests/test_kda_prefill.py -k prefill_equals -v` +Expected: FAIL — no module `aten.plena.program_kda_prefill`. + +- [ ] **Step 3: Implement the prefill layer** + +Create `aten/plena/program_kda_prefill.py`. Per chunk, per head: + +1. `A = tril(diag(β) K Kᵀ, −1)` — `M_TMM`, the `qkt.py` template with `key_dim` substituted for `head_dim`. +2. `T` — `kda_ut_transform_v0` (Task 9). +3. `W = T K`, `U = T V` — two `M_MM`. +4. Intra-chunk output: `Y_intra = tril(Q Kᵀ ⊙ decay_mask) U` — `M_TMM` then `M_MM`, mirroring `ssd_chunk_head_v0`. +5. Inter-chunk output: `Y_inter = Q S_prev`, rows scaled by the decay cumsum — mirrors `ssd_inter_chunk_output_v0`. +6. State update: `S_new = decay_chunk ⊙ S_prev + Wᵀ (U − W S_prev)`. + +Step 6 contracts over the **row** axis of both operands, which neither `M_MM` nor +`M_TMM` expresses — the identical obstacle `program_ssd.py:38-45` documents for +Mamba's `h[n,p] = Σ_t B[t,n] X[t,p]`. Use the same fix: stage the transposed +operand at its source via the `ssd_transposed_projection_v0` pattern +(`program_ssd.py:321-366`). Do not invent a transpose instruction. + +- [ ] **Step 4: Run the equivalence test** + +Run: `pytest aten/tests/test_kda_prefill.py -k prefill_equals -v` +Expected: PASS. + +Diagnosis: outputs match but final state does not → step 6's transposed staging. First chunk matches, second does not → the inter-chunk decay scaling in step 5. Everything off by a per-row factor → the decay cumsum is exclusive where it should be inclusive. + +- [ ] **Step 5: Run the prefill-then-decode test** + +Run: `pytest aten/tests/test_kda_prefill.py -k then_decode -v` +Expected: PASS. This is the gate on prefill and decode agreeing on the state layout. + +- [ ] **Step 6: Confirm the full suite** + +Run: `pytest` +Expected: PASS. + +- [ ] **Step 7: Commit** + +```bash +git add aten/plena/program_kda_prefill.py aten/tests/test_kda_prefill.py +git commit -m "feat(kda): chunked prefill on the Matrix Engine, equivalent to sequential decode" +``` + +--- + +## Phase 3 — Whole model, executed, on real weights + +This is the phase that closes the project's two largest credibility gaps: whole-model artifacts that are never bound to a real checkpoint, and never executed first layer to last. + +Generating machine code is not the same as running it. Until Task 12 passes, "PLENA runs Mamba-2 and KDA" is not a statement anyone should make. + +--- + +### Task 11: Whole-model lowering and static instruction gates + +**Files:** +- Create: `aten/tests/test_instruction_budget.py` additions +- Create: `aten/tests/test_whole_model_lowering.py` +- Modify: `.github/workflows/` (both repos) + +**Interfaces:** +- Consumes: `kda_layer_v0` (Task 4), `kda_prefill_v0` (Task 10), `ssm_decode_step_v0` (Task 8). +- Produces: full Nemotron-3 (52 layers) and Kimi-K3 (93 layers) lowered programs, plus CI gates on their static instruction counts. + +- [ ] **Step 1: Write the failing whole-model lowering test** + +```python +"""Whole-model lowering. Generating the program is the precondition for +Task 12 executing it; passing this test alone proves nothing about numerics.""" + +def test_nemotron3_52_layers_lower_and_assemble(nemotron_model_program): + words = nemotron_model_program.assemble() + assert len(words) > 0 + assert all(0 <= w <= 0xFFFF_FFFF for w in words) + print(f"NEMOTRON_STATIC_INSTR={len(words)} BYTES={len(words) * 4}") + + +def test_kimi_k3_93_layers_lower_and_assemble(kimi_model_program): + words = kimi_model_program.assemble() + assert len(words) > 0 + assert all(0 <= w <= 0xFFFF_FFFF for w in words) + print(f"KIMI_STATIC_INSTR={len(words)} BYTES={len(words) * 4}") +``` + +- [ ] **Step 2: Run, and record both counts** + +Run: `pytest aten/tests/test_whole_model_lowering.py -v -s` +Expected: PASS, printing two counts. + +The recurrent layers should be hardware loops rather than unrolled triples. **If the count is larger than the loop structure predicts, a sweep is unrolled** — find it with `grep -c C_LOOP_START` on the dumped layer before proceeding. + +- [ ] **Step 3: Add the static gates** + +Append to `aten/tests/test_instruction_budget.py`, with the ceilings set to the Step 2 measurements plus 10% and the date recorded: + +```python +# MEASURED 2026-__-__. Static instruction count is a property of the compiled +# artifact, not a model output -- it is the only quantity in this project +# admissible as a gate. See the plan's "Instruction accounting" section. +NEMOTRON_STATIC_INSTR_MAX = None # fill in from Task 11 Step 2 +KIMI_STATIC_INSTR_MAX = None # fill in from Task 11 Step 2 +``` + +- [ ] **Step 4: Re-run and confirm the gates pass** + +Run: `pytest aten/tests/test_instruction_budget.py -v` +Expected: PASS. + +- [ ] **Step 5: Wire into CI and prove it runs** + +Run: `grep -rn "test_whole_model_lowering\|test_instruction_budget" .github/workflows/` +Expected: at least one hit each. A guard that is defined but never invoked is the failure mode this step exists to prevent — a guard that is defined but never invoked stays green forever without executing once. + +- [ ] **Step 6: Commit** + +```bash +git add aten/tests/test_whole_model_lowering.py aten/tests/test_instruction_budget.py .github/workflows/ +git commit -m "test: lower both whole models and gate their static instruction counts" +``` + +--- + +### Task 12: Execute a whole model end to end in the emulator + +This has never been done here. It is the difference between "we generated machine code" and "PLENA runs the model". + +**Files:** +- Create: `transactional_emulator/testbench/models/whole_model_test.py` +- Modify: `justfile`, `.github/workflows/` + +**Interfaces:** +- Consumes: the lowered programs from Task 11. +- Produces: a first-layer-to-last-layer emulator run of both models with synthetic weights, numerically compared against a full PyTorch reference forward. + +- [ ] **Step 1: Write the failing end-to-end test** + +```python +"""First layer to last, in the emulator, against a full reference forward. + +Start with a reduced layer count so the test is runnable in CI, then run the +full depth as a nightly. A shape-reduced model exercises every code path a +full one does; what it cannot catch is HBM address-space exhaustion, which +Step 4 checks separately. +""" + +def test_nemotron3_reduced_depth_matches_reference(whole_model_harness): + out = whole_model_harness.run("nemotron3", layers=4, tokens=4) + torch.testing.assert_close(out.actual, out.expected, rtol=1e-2, atol=1e-2) + + +def test_kimi_k3_reduced_depth_matches_reference(whole_model_harness): + out = whole_model_harness.run("kimi_k3", layers=4, tokens=4) + torch.testing.assert_close(out.actual, out.expected, rtol=1e-2, atol=1e-2) +``` + +- [ ] **Step 2: Run and confirm it fails** + +Run: `python transactional_emulator/testbench/models/whole_model_test.py` +Expected: FAIL — the harness does not exist. + +- [ ] **Step 3: Build the harness and get 4 layers passing** + +Reuse the connected-test harness on `feat/mamba2-support` for program loading, HBM image construction, and result extraction. The new part is chaining layer outputs into layer inputs rather than testing one layer in isolation. + +- [ ] **Step 4: Check the HBM address space** + +Run the harness with the full layer count and assert the required HBM allocation is within the emulator's configured `HBM_SIZE` (`plena_settings.toml`, currently 16 GiB). + +This is a real wall: the emulator preloads HBM from a flat file starting at offset 0, so the allocation and the file both have to span every address the program touches. A layout whose regions sit far apart needs a span-sized allocation whatever the live data comes to, and stops being executable long before it stops being describable. If ours does that, pack the regions into one bounded arena. + +- [ ] **Step 5: Run full depth as a nightly** + +Add a `just test-whole-model-full` recipe running 52 and 93 layers, wired to a scheduled workflow rather than to every push. + +- [ ] **Step 6: Commit** + +```bash +git add transactional_emulator/testbench/models/whole_model_test.py justfile .github/workflows/ +git commit -m "test: execute both whole models first layer to last in the emulator" +``` + +--- + +### Task 13: Bind to real checkpoint weights + +**Files:** +- Create: `transactional_emulator/testbench/models/checkpoint_test.py` +- Modify: `justfile`, `doc/` (record the result) + +**Interfaces:** +- Consumes: the harness from Task 12. +- Produces: both models running on real published weights, compared against the HuggingFace reference forward. + +- [ ] **Step 1: Write the failing checkpoint test** + +```python +"""Real weights, not randn. + +Synthetic weights validate the datapath; they cannot validate the layout +contract -- a transposed projection, a wrong head-group mapping, or a swapped +q/k ordering all produce plausible numbers on symmetric random inputs and +wrong numbers on trained ones. +""" + +import pytest + + +@pytest.mark.checkpoint +def test_nemotron3_layer0_matches_huggingface(checkpoint_harness): + out = checkpoint_harness.run_layer("nemotron3", layer=0, tokens=8) + torch.testing.assert_close(out.actual, out.expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.checkpoint +def test_kimi_k3_layer0_matches_huggingface(checkpoint_harness): + out = checkpoint_harness.run_layer("kimi_k3", layer=0, tokens=8) + torch.testing.assert_close(out.actual, out.expected, rtol=2e-2, atol=2e-2) +``` + +Mark them `checkpoint` and deselect by default, so contributors without the +weights are not blocked; run them in a scheduled workflow that has them. + +- [ ] **Step 2: Run and confirm it fails** + +Run: `pytest -m checkpoint -v` +Expected: FAIL — no harness. + +- [ ] **Step 3: Load one real layer and get it passing** + +Load the layer's weights from the published checkpoint, run the HF reference +forward for that layer, lower and run the same layer on the emulator, compare. +Start with layer 0 of each model. + +Expect this step to surface layout bugs rather than arithmetic bugs. When one +appears, fix the **lowering**, and add a synthetic test that would have caught +it — a synthetic case with a deliberately asymmetric weight, so the regression +is guarded without needing the checkpoint. + +- [ ] **Step 4: Extend to the whole model at reduced depth** + +Run layers 0–3 chained, on real weights, against the HF forward. + +- [ ] **Step 5: Record the result** + +Write a short `doc/checkpoint_validation.md` stating exactly what was validated: +which checkpoint revision, which layers, which tolerance, prefill or decode. +State plainly what was **not** validated. This document is what licenses the +claim "PLENA runs Nemotron-3 and Kimi-K3"; keep it narrower than the claim. + +- [ ] **Step 6: Commit** + +```bash +git add transactional_emulator/testbench/models/checkpoint_test.py justfile doc/checkpoint_validation.md +git commit -m "test: validate both models against real checkpoint weights" +``` + +--- + +## Phase 4 — Measure, then retire + +Measurement and cleanup, last. No new mechanism is designed here. + +--- + +### Task 14: Banking study and the uncalibrated latency report + +**Files:** +- Create: `analytic_models/performance/vector_sram_banking.py` +- Create: `doc/static_path_measurements.md` + +**Interfaces:** +- Consumes: green Tasks 11–13. +- Produces: a recorded answer to the Vector SRAM banking question, and a latency report that is honest about its own status. + +- [ ] **Step 1: Retarget the bank-conflict model** + +A bank model that consumes only `(rows, columns)` is algorithm-independent. Port it into `analytic_models/performance/vector_sram_banking.py`, driven by the Vector SRAM row/column geometry of the lowered Mamba and KDA tiles. + +- [ ] **Step 2: Run it and record the answer** + +Run: `python -m analytic_models.performance.vector_sram_banking` +Expected: stall-cycle counts for the Mamba and KDA state sweeps. + +Record in `doc/static_path_measurements.md`: +- Stalls ≈ 0 → the vector path is limited by instruction issue, and the follow-up is post-increment addressing, not SRAM work. +- Stalls material → that is the evidence for a Vector SRAM bank-mapping mode. Separate plan. + +- [ ] **Step 3: Write the measurement report** + +`doc/static_path_measurements.md` must separate the two kinds of number explicitly, with the table from this plan's "Instruction accounting" section: + +- **Hard facts:** static instruction counts (Task 11), dynamic instruction counts, HBM bytes moved. +- **Uncalibrated model output:** cycles, µs, TPOT — labelled as such on every line, and not compared against any other uncalibrated model output; two such numbers decide nothing between them. + +- [ ] **Step 4: Commit** + +```bash +git add analytic_models/performance/vector_sram_banking.py doc/static_path_measurements.md +git commit -m "docs: banking study and an explicitly uncalibrated latency report" +``` + +--- + +### Task 15: Confirm the tree carries no runtime state machinery + +The static path is defined by what it does *not* have: no descriptors read at +run time, no command queue, no residency cache. This task is the check that it +stayed that way, not a removal -- this line of work was cut from `main`, which +has none of it, so there has never been anything to delete. + +**Interfaces:** +- Consumes: green Tasks 11-14. +- Produces: a tree with no descriptor fetch, no queue, no residency cache. + +- [ ] **Step 1: Confirm nothing introduces runtime state machinery** + +```bash +grep -rniE "descriptor|residency|state_queue|StateDescriptor" \ + --include="*.rs" --include="*.py" --include="*.svh" --include="*.json" . \ + | grep -v "^./docs/superpowers/" +``` +Expected: no hit that describes a run-time fetch. `aten/tests/test_no_state_engine.py` +is the CI form of this and must be green. + +- [ ] **Step 2: Confirm the emulator builds and passes** + +Run: `cargo test --manifest-path transactional_emulator/Cargo.toml` +Expected: PASS. + +Run: `cargo clippy --manifest-path transactional_emulator/Cargo.toml --all-targets -- -D warnings` +Expected: no warnings. +## What is deliberately not carried over + +Recorded so a reviewer does not read these as oversights. + +| Dropped | Why | +|---|---| +| 256-byte HBM descriptor | Every field was a compile-time constant; the indirection carried no runtime information and cost a serially-dependent HBM read per command. | +| 16 async queues, events, `FENCE`, completion records | Never exercised by any emitted program; the event map had no eviction path, so it was not implementable as an RTL spec. | +| `StateCache` residency table | Residency is a compile-time allocation. The compiler already emitted explicit spill code; the runtime table only re-checked what the compiler guaranteed. | +| `cache_hit_rate` metric | It was `!descriptor.is_streaming()` replayed — a pure function of the program image, not a measured effect. | +| 7 runtime status codes | All collapsed to a `panic!` at the dispatch site. Shape and precision violations are now Python or assembler errors. | +| Named layout modes carried in a descriptor | The projection co-layout finding is real and worth publishing, but it is a Matrix-writeback question, not a state-engine one. It needs its own plan, gated on Task 14's banking result. | + +## What is carried over + +| Kept | Where it lands | +|---|---| +| KDA golden reference and its self-test | Task 1 | +| Mixed-precision error study (Mamba MX8 2.07% vs KDA 0.35% at 32K) | Justifies the FP32/BF16 split in Task 2's docstring | +| Bank-conflict model | Task 14, retargeted at Vector SRAM | +| Full-system HBM-bandwidth DSE | Untouched; out of scope for this plan | +| Multi-token state-carry testing with a real convolution kernel and non-zero decay | Task 4's four-token test, plus the explicit degeneracy guard | diff --git a/docs/superpowers/progress/2026-08-25-static-mamba-kda.md b/docs/superpowers/progress/2026-08-25-static-mamba-kda.md new file mode 100644 index 000000000..838585631 --- /dev/null +++ b/docs/superpowers/progress/2026-08-25-static-mamba-kda.md @@ -0,0 +1,1985 @@ +# Static Mamba-2 + KDA — execution log + +Plan: [`../plans/2026-08-25-static-mamba-kda.md`](../plans/2026-08-25-static-mamba-kda.md) +Branches: `feat/static-kda` in `PLENA_Simulator` (from `65cf6f09`) and in the +`PLENA_Compiler` submodule (from `f5eb36a`), both branched off `feat/mamba2-support`. + +## Environment + +Not provisioned on this machine at the start. Recorded because every number below +depends on it. + +| | | +|---|---| +| Python | `uv` 0.12.5, `.venv` at `PLENA_Simulator/.venv`, Python 3.12 | +| torch | 2.7.1+cu126 (via `uv sync`; CPU execution) | +| Test command | `uv run pytest`, with `PYTHONPATH=$SIM:$SIM/PLENA_Compiler:$SIM/PLENA_Tools:$SIM/transactional_emulator/testbench` | +| Rust / emulator | **`nix develop`** — the flake supplies the toolchain, Ramulator2 and libtorch | + +**Correction.** I first ran `cargo` outside the Nix dev shell, hit +`unable to find library -lramulator`, and recorded it as "a C++ dependency that +is not built" blocking whole-model execution. That was wrong. `flake.nix:78` +pulls `customPkgs.ramulator2` into `buildInputs` and pre-fetches libtorch +(`:39`), `transactional_emulator/pkgs/` carries both derivations, and Nix was +already installed. Inside the shell: + +``` +Ramulator2: library at /nix/store/...-ramulator-2.1-unstable-2026-07-30/lib +Libtorch: /nix/store/...-libtorch-2.7.0/lib +cargo test -> 106 passed; 0 failed +``` + +The rustup toolchain I installed is redundant; the flake provides its own. +**Emulator work is not blocked.** + +## Baseline — `feat/mamba2-support`, before any change + +`pytest PLENA_Compiler/aten/tests` at `f5eb36a`, verified in a clean worktree: + +**3 failed, pre-existing, not caused by this work:** + +- `test_moe_stage_attribution.py::test_stage_parameters_have_no_default` — three + existing emitters give `stage` a default (`program_mamba_common.py:362`, + `program_ssd.py:321`, `program_ssm_recurrent.py:213`) +- `test_plena_compiler.py::test_qwen_packed_skinny_router_rowpacked_compiles_for_128_experts` +- `test_quantization_ablation.py::test_mxfp8_is_sole_gap_source` + +Any run below reporting exactly these three is at baseline. + +--- + +## Task 1 — KDA golden reference ✅ + +**Result: 3 failed (all baseline), 112 passed.** New: `test_kda_reference.py`, 9 passed. + +### What landed + +| File | | +|---|---| +| `PLENA_Compiler/aten/models/kda/reference.py` | 355 lines, the torch KDA forward | +| `PLENA_Compiler/aten/models/kda/state_precision.py` | 95 lines, a dependency of the above | +| `PLENA_Compiler/aten/models/kda/__init__.py` | new | +| `PLENA_Compiler/aten/tests/test_kda_reference.py` | 253 lines + 1 contract test | +| `PLENA_Compiler/.github/workflows/ci.yml` | wired the new test in | + +### Two corrections to the plan, made while executing + +**1. The reference belongs in the Compiler, not the Simulator.** The plan put it at +`PLENA_Simulator/analytic_models/reference/` and had the Compiler import it. That +direction does not exist: `transactional_emulator/testbench/` imports +`aten.models.gpt_oss.*` from the Compiler, and the Compiler imports nothing from +the Simulator. `aten/models/mamba2/reference.py` is the precedent. Moved to +`aten/models/kda/`; the Simulator can reach it because `PLENA_Compiler` is on its +`PYTHONPATH` (`justfile:135`). + +**2. The reference API is not what the plan assumed.** Real signature: + +```python +kda_step(q, k, v, gate, beta_logit, state: KdaState, a_log, dt_bias, shape: KdaShape, + *, scale=None, state_storage=StateStorage.FP32) -> tuple[Tensor, KdaState] +``` + +Batched `[batch, heads, dim]`; `gate_lower_bound` lives on `KdaShape` (default +`-5.0`); `scale` defaults to `1/sqrt(key_dim)`. + +Most important: normalisation is **`x * rsqrt(sum(x²) + 1e-6)`**, matching FlashKDA's +kernel, *not* `torch.nn.functional.normalize`, which clamps instead. The plan's draft +test used `F.normalize` and would have disagreed. Downstream lowering must use the +rsqrt form. + +### The contract test + +`test_decode_step_matches_the_transposed_formulation` pins `T[k, v] == S[v, k]` — +the load-bearing design decision of the whole plan. It replays the reference in the +transposed orientation the lowering will use and asserts both the output and the +final state agree to `rtol=1e-5, atol=1e-6`. Without it a lowering bug could hide +behind a plausible reimplementation of the maths. + +### Deviation from the plan worth noting + +The reference package's `__init__.py` also pulled in +`nemotron3_mamba.py`. Not carried across: a second Mamba reference alongside the +Compiler's existing `aten/models/mamba2/reference.py` is exactly the two-drifting- +copies problem the plan warns about. The KDA-only `__init__` records why. + +### Guard that did its job + +`test_every_test_file_here_is_wired_into_ci` failed the moment the new test file +appeared and was not in `ci.yml` — caught before commit. This is the fail-closed +guard style the review recommended, working as intended. + +### Review round — `a5753d2` + +A subagent ran mutation testing against the commit. **Six functional defects, all +false-greens from degenerate inputs.** The transposition contract itself was sound +(nine mutations to `reference.py` correctly turned it red), but the inputs it ran +on made four classes of bug invisible. + +| # | Defect | Mutation that stayed green | Fix | +|---|---|---|---| +| D1 | `batch=1`; no test in the suite compared numerics at batch>1 | `beta`, `k_normalized`, `log_decay` each recomputed from batch 0 | `batch=2` | +| D2 | `key_dim == value_dim == 8` | output scale `1/sqrt(key_dim)` → `sqrt(value_dim)` | `key_dim=8, value_dim=5` | +| D3 | `q`/`k` at unit variance: `‖x‖²≈O(1)`, so the `1e-6` inside `rsqrt` perturbs by ~2e-7, **below the test's own `atol=1e-6`** | dropping the epsilon; swapping in `F.normalize` | scale `q`/`k` by `0.05` | +| D4 | `activate_log_decay` tested with `a_log=zeros`, `dt_bias=zeros`, and used as its own oracle everywhere else | dropping `dt_bias`; dropping `exp(a_log)` | new `test_log_decay_matches_its_closed_form` | +| D5 | `test_sequence_matches_repeated_steps` only asserted shape + `isfinite` | dropping the carried state; reversing token order | compare against repeated `kda_step` | +| D6 | `state_precision.py` present, its tests not — MX8_B128 path uncovered | — | ported both tests | + +Every fix was verified by re-running the mutation it targets. All six mutants now die. + +**D3 is the one to remember.** The comment in the test asserted that the FlashKDA +`rsqrt(sum + eps)` convention differs from `F.normalize` — and it does, but not at +an amplitude the test could see. A correct claim, an input that could not test it. + +Also switched the package from bare `aten.` to `compiler.aten.`, matching +`aten/models/mamba2/__init__.py` and 37 other files. Both resolve to the same +directory via `compiler/__init__.py`'s `__path__` hack but produce *distinct module +objects*, so `KdaShape` imported both ways would be two classes — breaking +`isinstance` and dataclass equality as soon as the lowering imports it. + +**After fixes: 3 failed (baseline), 115 passed.** `test_kda_reference.py`: 12 passed, +also verified under the CI invocation (`cd PLENA_Compiler && PYTHONPATH=. pytest`). + +--- + +## Task 2 — KDA state layout and data movement ✅ + +**Result: 3 failed (baseline), 135 passed.** `test_kda_lowering.py` 14, `test_kda_stage_contract.py` 6. + +### What landed + +`aten/plena/program_kda_common.py` — `ProgramKdaCommonMixin`, wired into `PlenaCompiler`'s +MRO after `ProgramSSMRecurrentMixin`. State residency (`kda_pin_state_v0`, +`kda_load_state_v0`, `kda_store_state_v0`), conv history roll, L2 normalisation, +FPRAM constants. Plus `aten/models/kda/shape.py` and two new test files. + +State is `[num_heads * key_dim, value_dim]` — head `h` owns rows +`[h*key_dim, (h+1)*key_dim)`. Data movement is the existing H-type path; no +descriptor, no queue, no residency table. + +### Review round — `055223c` / `e4110d47` + +A subagent ran mutation testing. **13 mutations survived the first version**, plus +three defects no test could have caught. All fixed; all 13 now fail the suite. + +**1. I broke a CI job.** `compiler.py` imports the KDA lowering at module scope → +`reference.py` → `torch`. So every `compiler.aten.plena.*` import began requiring +torch, breaking `moe-stage-guard`, which installs only pytest and pyyaml by design. +Verified with a `meta_path` hook blocking torch: clean at parent `a5753d2`, broken at +`34692ce`, clean now. Fix: `KdaShape` moved to a torch-free `shape.py`, and +`aten/models/kda/__init__.py` re-exports lazily via PEP 562. + +*(My first reproduction wrongly showed "OK" — I used the pre-3.12 `find_module` +protocol, which Python 3.12 ignores. The reviewer's finding was right and mine was +the broken experiment.)* + +**2. The FP32 precision claim was false, and `4` was wrong under every config.** +`precision=1` selects the keyvalue *class*; the width comes from the active +PRECISION table, and the shipped `HBM_V_KV_TYPE` is `format = "Mx"` with e4m3 +elements — 1 byte plus a scale stream. `storage_precision` only feeds compiler-side +address arithmetic. There is no FP32 path through that class at all. + +Corrected to **BF16, 2 bytes** on load, store and the pinned reservation, and the +module now *requires* Plain BF16 via `kda_require_state_precision_v0` rather than +asserting it — reusing the guard `f5eb36a` added for this exact bug on the Mamba +path. This also fixes the under-reserved pinned region: under Plain BF16 no scale +stream is written, so the reservation is correct *because* of the guard. + +**3. The stage markers were not markers.** Four emitters wrote `stage=kda_...` +with no `@`. `extract_stage_tag` matches `"@stage="`, so those comments were +invisible and every KDA instruction billed to the **preceding** stage — silently, +producing a profile that looks correct. Added `KDA_STAGES` + `kda_stage_marker`, +ten matching `StageKind` variants in the emulator (`cargo check` clean), and a +cross-repo contract test that also greps for hand-formatted markers. + +**4. Four tests asserted nothing.** + +| Test | Why it could not fail | +|---|---| +| `test_state_store_mirrors_the_load_precision` | `stored.hbm_addr == addr` is a tautology — `store()` builds the InputVar from the address it was given | +| `test_pinned_state_region_does_not_collide...` | compared an **element count** against a **byte address** | +| `test_conv_roll_shifts_history_and_appends` | counted `V_ADD_VV`; a reversed shift and a wrong append source both emit 3 | +| `test_l2_normalize_uses_the_scalar_sqrt_path` | `any(opcode in ...)` — `V_MUL_VF` comes from the zeroing idiom, `V_RED_SUM` from any surviving reduction | + +Rewritten around a small `_trace()` that resolves `S_ADDI_INT gpN, gp0, imm` so +assertions name **resolved operand addresses**. One mutation (block-copy src/dst +swap, which zeroes the vector) survived even the repair — every operand I checked +was unchanged — and needed an explicit assertion on the copy's direction. + +**5. Smaller functional fixes.** `mamba_fp_constants` now takes no shape (it never +read one), removing a `Mamba2Shape.__new__` uninitialised-frozen-dataclass hack. +`kda_l2_normalize_v0` rejects a scratch aliased to its input (would silently return +zeros), an undersized accumulator, and a vector wider than one column tile (the +`tile_col_idx=0` default would normalise by a partial norm — and `KdaShape.kimi_k3()` +has `key_dim=128` against a default `mlen=64`). + +### Lesson carried forward + +Both review rounds found the same shape of bug: **the assertion was right, the +setup could not exercise it**. Task 1 was degenerate inputs; Task 2 was +opcode-presence instead of operand identity. Tests for later tasks assert resolved +addresses and counts from the start. + +--- + +## Task 3 — The KDA decode step (existing instructions only) ✅ + +**Result: 3 failed (baseline), 160 passed.** `test_kda_decode_step.py` 12, +`test_isa_interpreter.py` 12. + +### The claim Phase 0 exists to make, now demonstrated + +`kda_decode_step_v0` reproduces `kda_step`'s output **and** updated state across +four geometries, **using no instruction that did not already exist**. Seven +opcodes: `S_ADDI_INT`, `S_LD_FP`, `V_MUL_VF`, `V_ADD_VV`, `V_SUB_VV`, +`C_LOOP_START`, `C_LOOP_END`. + +**Phase 0 baseline for Task 8: `static=505, dynamic=540`** (1 head, key_dim=8). + +### How it is validated without the emulator + +`aten/tests/isa_interpreter.py` — a ~200-line interpreter for exactly those seven +opcodes — executes the emitted assembly, so the tests compare numbers against the +CPU reference rather than asserting which opcodes appear. It raises on anything +outside its subset, so a lowering that starts emitting something new fails loudly. + +This is a second implementation of ISA semantics and could drift; the whole-model +emulator run (Task 12) is what finally cross-checks it. It is scoped to opcodes +whose meaning is one line each. + +### Review round — `8c23ffc` + +18 of 20 lowering mutations were caught. The two that were not, plus three oracle +defects: + +| Defect | Why it mattered | +|---|---| +| **Both accumulator zeroings deletable** | Every test ran *one* step on a freshly-zeroed `Machine`. A real decode loop reuses `pred`/`o` every token. Reviewer measured token-1 error at 1.9e-2 / 7.6e-2 | +| **`doc/plena_isa_spec.md` had `V_SUB_VV` backwards** | Spec said `rd = rs2 - rs1`; hardware does `rs1 - rs2`. Subtraction does not commute — **"fixing" the interpreter to match the spec would have flipped KDA's error sign** | +| **`gp0`/`f0` are not hardwired zero** | They are ordinary writable registers; `RegisterAllocator` reserving them is the entire guarantee. The oracle asserted hardware behaviour that does not exist | +| **No VLEN-alignment check in the oracle** | The emulator's vector SRAM asserts it; unaligned addresses panic there, passed silently here | +| **The oracle had no tests of its own** | Mutating `V_MUL_VF` to read its destination left every decode test green — the lowering only emits the in-place form | + +Fixes: a two-token test that carries state on chip and rewrites only FPRAM between +tokens; the spec corrected with a note; the register model made faithful with the +invariant guarded where it lives (a test that the lowering writes neither `gp0` +nor `f0`); an alignment guard; and `test_isa_interpreter.py`, which exercises every +opcode in a form the lowering does not produce. All four oracle mutations now fail. + +### Known restriction, carried forward + +**`value_dim` must not exceed `mlen`.** The `tile_row_*` family defaults to +`tile_col_idx=0` while `vram_fill_zero` walks every column block and +`vram_matrix_add`'s wide path needs mlen-aligned row offsets, which cross-row +single-row copies do not have. Asserted, not implied. + +**Kimi K3 is `value_dim=128` against a default `mlen=64`, so this must be lifted +before any whole-model run.** It is a change to the shared `vram_*` helpers, not +to KDA. + +### Recurring lesson, third round running + +Every review has found the same shape: **the assertion was right, the setup could +not exercise it.** Degenerate inputs (Task 1), opcode presence instead of operand +identity (Task 2), a single step instead of two (Task 3). The generalisation is +that a test must be run against a deliberately broken implementation before it is +trusted — which is now the working habit, not an afterthought. + +--- + +## Interlude — unblocking the emulator found a bug in Task 2 ✅ + +Running `cargo test` for the first time (see the environment correction above) +immediately failed one test: + +``` +stage_marker_names_match_the_compiler_vocabulary + StageKind has stages no compiler marker produces: + ["kda_qkv_proj", "kda_conv1d", ... "kda_state_store"] +``` + +`e4110d47` added ten KDA variants to `StageKind` but not to the guard that reads +the vocabulary *back* out of the compiler's Python. The guard is bidirectional and +unions several compiler modules — MoE from `program_routed_moe.py`, Mamba from +`program_mamba_common.py` — because one `StageKind` is fed by all of them. KDA was +the third and was missing. + +`cargo check` could not catch this: it is a test, not a compile error. Fixed in +`6271e982`; verified bidirectional by dropping `kda_readout` from the Python +vocabulary and confirming the Rust test names exactly that stage. + +**Emulator suite: 106 passed, 0 failed.** + +The lesson is narrower than it looks: I recorded an environment failure as a +project blocker without checking how the project builds. The `flake.nix` and +`docker/Dockerfile` both said so in the first ten lines. + +--- + +## Interlude 2 — column-block folding ✅ `993d863` / `595e812` + +Lifts the `value_dim <= mlen` restriction that Task 3 asserted against, **without +touching any shared helper**. + +### The choice + +Three options were on the table: + +| | verdict | +|---|---| +| Add `tile_col_idx` to the six helpers that lack it | Rejected. Touches emitters Mamba and attention are using, for a bug **no current caller triggers** — `Mamba2Shape.validate` already constrains Mamba to `== mlen`. And its error text names this exact fix as *"deliberately not done yet"*, so it would reverse a decision the codebase made on purpose. | +| Loop `tile_col_idx` inside KDA only | Rejected. Measured: only `tile_row_mul_fp` accepts it. `tile_row_sum`, `tile_row_mul`, `tile_row_sub`, `vram_fill_zero`, `mamba_row_copy`, `mamba_row_add` do not — so the copy/multiply/add triple could not be told which block to work on. | +| **Fold the block into the row** | **Chosen.** `row = (head * blocks + block) * key_dim + key`. Every tile is exactly one block wide, no helper needs a column argument, and fixing `(head, block)` leaves keys at unit stride so each sweep is still one hardware loop. | + +At `value_dim == mlen` the folded layout is byte-identical to the previous one — +verified by diffing the emitted assembly against `8c23ffc`: **776 lines each, every +instruction line identical**, differing only in stage-marker comments that gained +` block=0`. + +### Why it is sound + +Nothing in the KDA recurrence contracts over `value`. Decay, rank-1 update and +read-out are elementwise along it; prediction and read-out contract over `key` at a +fixed value lane. So block `c` of `pred`/`out` reads only block `c` of the state, +and blocks are independent. The reviewer confirmed this against the emitted code by +resolving every VRAM operand back to `(tile, row)` and grouping by stage marker: +**zero cross-block or cross-head accesses**, and the state rows touched are exactly +the full tile. + +### Performance data + +MLEN = 8, one decode token: + +| heads | key_dim | blocks | static | dynamic | static / (head·block) | +|---:|---:|---:|---:|---:|---:| +| 1 | 8 | 1 | 505 | 540 | 505 | +| 1 | 8 | 2 | 1010 | 1080 | 505 | +| 1 | 8 | 4 | 2020 | 2160 | 505 | +| 2 | 4 | 1 | 554 | 584 | 277 | +| 2 | 4 | 2 | 1108 | 1168 | 277 | +| 3 | 5 | 3 | 3006 | 3186 | 334 | + +**Exactly linear in blocks, zero per-block overhead.** Per-`(head, block)` cost +depends only on `key_dim`. Dynamic exceeds static by ~7%, from the `C_LOOP` bodies +in the decay and fill-zero sweeps. + +### Review round — `595e812` + +No correctness bug; 17 of 20 mutations killed. The three survivors shared one blind +spot: **nothing pinned row *ordering* or residency *sizing* at `blocks > 1`.** + +| Survivor | Why it mattered | +|---|---| +| `kda_state_row` head/block swapped | Still a bijection, still unit-stride, so every layout test accepted it. But it *is* the HBM byte layout of the pinned state — a cross-component contract as soon as prefill or a checkpoint loader writes an initial state | +| `kda_pin_state_v0` dropping the blocks factor | Under-reserves, letting a later `input()`/`store()` land on the carried state. All four residency tests ran at `value_dim == mlen`, where old and new formulas coincide | +| `kda_head_base` ignoring its `block` argument | Same root cause | + +Parametrising the residency tests over 1/2/3 blocks was **not sufficient on its +own**: the collision test compares against the size `pin` recorded, so an +under-reservation is self-consistent. What kills it is the round-trip test, which +computes the required size independently. + +**Seeding/readback symmetry held up.** Five harness mutations; only a *globally +consistent* relabelling of the block axis survived, and that is an equivalent +convention, not a bug — the gather side is anchored to the reference in true +`[heads, key, value]` order. + +### The trap this uncovered for Task 4 + +`kda_l2_normalize_v0`'s guard told callers to "block it with `kda_vector_row` +first". **That is actively wrong.** KDA normalises `q` and `k`, which are +*key*-width, and Kimi K3 is `key_dim = 128` against `mlen = 64` — but the L2 norm +**contracts over the key axis**, so the two halves of one `q` share one norm. +Following that hint normalises each half by its own partial norm: finite, +plausible, wrong. + +Row folding solves `value_dim` precisely *because* nothing contracts over value. It +cannot solve `key_dim`. + +And the obvious workaround does not work either: `tile_row_sum` emits +`S_ADD_FP f1, f0, f0` before each `V_RED_SUM`, so a second call **overwrites** the +FPRAM slot rather than adding to it — even though the `V_RED_SUM` instruction +itself accumulates into its FP register. Verified by reading the emitted assembly. + +**Task 4 must therefore reduce each block into its own FPRAM slot and combine them +with scalar `S_ADD_FP` before the rsqrt** — `blocks - 1` extra scalar adds per +vector. + +--- + +## Task 4a — The causal conv step ✅ `de3eaae` / `08a7a2c` + +**Result: 3 failed (baseline), 194 passed.** `test_kda_conv.py` 9, +`test_isa_interpreter.py` 19. + +KDA runs three short depthwise convolutions — q, k, v — before the recurrence. +Validated numerically against `_causal_conv_step` across four geometries including +multi-block, plus direct assertions that the history *shifts* rather than being +overwritten and that perturbing one channel block moves only that block. + +### Reuse, and where it stopped + +`mamba_conv1d_v0` is the same kernel and makes the same argument about `V_MUL_VF` +being wrong for a per-lane tap weight. Two differences made it its own emitter: +Mamba convolves a sequence held in `x` (so the shift is a free row offset), while +KDA decode has a **carried** state (a real row copy every step); and KDA's channel +counts are ~12,288, so the block index has to reach the inner loop and +`mamba_conv1d_v0` has nowhere to put it. + +Adding a base-row argument there would have avoided the duplication — **rejected +for the same reason as the `tile_row_*` question**: that emitter is on Mamba's +decode and prefill paths. + +### The oracle refused, as designed + +`V_MUL_VV`, `V_EXP_V`, `V_RECI_V`, `V_ADD_VF` were outside its modelled subset, so +it raised instead of guessing. Extended deliberately with those plus `V_SUB_VF`, +**semantics taken from the emulator, not the spec** — `V_EXP_V` clamps to +`[-88, 88]` *before* exp, so an unclamped model could not return `inf` and have a +comparison pass for the wrong reason. + +### Review round — `08a7a2c` + +| Defect | Why it mattered | +|---|---| +| **`out` zero-fill unanchored** | Same class as `8c23ffc`. The harness built a fresh zero-initialised `Machine` and never wrote `out`, so "the program clears it" and "it happened to start at zero" were indistinguishable. `out` is persistent and reused by q/k/v | +| **SiLU billed to a Mamba stage** | `mamba_silu_v0` emitted `@stage=mamba_gated_norm` *inside* the `kda_conv1d` region — and markers are sticky, so everything after it billed there too. ~8% of conv instructions, ×3 convs ×61 layers | +| **Stage guard scanned one file** | `test_kda_stage_contract.py` hardcoded `program_kda_common.py`, so the grep guards passed vacuously on the new module | +| **`V_RECI_V` diverged on zero** | Oracle raised where `tensor.reciprocal()` returns `inf` | +| **Alias guard omitted `weight`/`bias`** | `mamba_row_copy` is zero-then-add; it would wipe them mid-loop | +| **`kda_conv_state_row` formula unanchored** | Emitter and both harness sides funnel through it, so any bijection cancels | + +`mamba_silu_v0` now takes an optional `marker` overriding the emitted string +verbatim — additive, default-preserving, so Mamba is untouched. + +### Performance + +mlen=8, `apply_silu=True`: **`static = 39 + 104·blocks`** at kernel 4, and `+22` +per unit of kernel per block. Exactly linear; SiLU amortises to `~31 + 8·blocks`. + +**The conv body is fully unrolled.** Every history shift and tap multiply is a +separate row op with `num_rows=1`, which `_row_progression` cannot fold — only the +SiLU tail becomes a hardware loop. At Kimi K3 scale (12,288 channels ÷ 64 = 192 +blocks, kernel 4) that is **~20,000 static instructions per conv, ×3 per layer**. + +The layout was chosen so the history shift is a run of adjacent row copies, but +nothing folds them yet — **the adjacency buys nothing today.** And `V_FMA_VF` will +not help: the tap multiply is by a per-channel *vector* (`V_MUL_VV`), not a +broadcast scalar. Folding the conv would need either a block-copy-based shift or a +`V_FMA_VV`. Recorded for the Phase 1 measurement rather than acted on. + +--- + +## Task 4b(i) — Cross-block L2 normalisation ✅ `f57a45f` / `9671ed1` + +**Result: 3 failed (baseline), 211 passed.** + +The trap Task 4a's review identified, now solved. Everywhere else in KDA a wide +vector splits across rows because nothing contracts over `value`. **`q` and `k` +are key-width and the norm contracts over `key`**, so the halves of one `q` share +one norm — the column-block folding does not apply. + +The obvious workaround does not work either: `tile_row_sum` emits +`S_ADD_FP f1, f0, f0` before each `V_RED_SUM`, so a second call into the same slot +overwrites it. So each block reduces into its own slot and the slots fold with +scalar `S_ADD_FP`. Partial sums are laid out **block-major**, `part[c*vectors + v]`, +so the fold is `blocks-1` contiguous FPRAM adds of length `vectors` rather than +`vectors` separate reductions. + +### The `_fpram_row_map` trap, second occurrence + +With >1 row it walks `base_offset + i`; with exactly one it **ignores +`base_offset`** and uses `single_offset`. At `vectors == 1` every block's sum +landed in slot 0 and only the last survived — invisible at `vectors > 1`, which is +why the parametrisation runs 1, 2, 3 and 4 vectors. Both call sites now go through +`_kda_row_sum` / `_kda_scale_rows`, which own the branch. `_kda_scale_rows` moved +into `program_kda_common` so the common module no longer depends on a method +defined in a specialisation. + +### Review round — `9671ed1` + +**The lowering was clean — all 13 mutations caught.** The failures were all in the +**oracle**, and they were exactly the semantics this task had gone to the emulator +to establish: + +| Survivor | Why the lowering cannot reach it | +|---|---| +| `V_RED_SUM` overwrite vs accumulate | Every emitted `V_RED_SUM` is preceded by `S_ADD_FP f1, f0, f0`, so f1 always enters at zero | +| `f0` write guard (×3) | `RegisterAllocator` hands out f1..f7, so f0 is never a destination | +| `S_SUB_FP` direction | Emitted **zero** times by any KDA path | +| `S_ST_FP` immediate | Non-zero only under `ATEN_OPS_UNROLL=1` | +| `S_RECI_FP(0)`, `S_SQRT_FP(neg)` | Inputs are always positive here | +| `V_RED_SUM` alignment check | Addresses are always aligned | + +All eight now pinned by direct unit tests. Two genuine oracle corrections fell out: +`S_RECI_FP(-0.0)` returned `+inf` where bf16 division gives `-inf` (Python's +`a == 0.0` is true for `-0.0`), and `V_RED_SUM` accepted a mask it does not model. + +### Performance + +One-line fix worth recording: `mamba_block_copy` was copying the **padded** tile +height, not the live rows. VRAM row counts must be a multiple of `mlen`, so at +`mlen=64` with two live rows it copied 64. + +| Kimi shape (mlen=64, 1 vector, 2 blocks) | dynamic | +|---|---:| +| before | 598 | +| after | **94** | + +**6.4×**, and the padding never contributed to the result. `kda_l2_normalize_v0` +has the same pattern — an existing convention worth revisiting. + +Scaling: static `+29` per block, `+11–13` per vector (`mamba_rsqrt_fpram` is +Python-unrolled at 9 instructions per vector). + +### Flagged, not fixed — shared code + +`tile_row_max_asm` calls `_emit_tile_row_reduce` **without** +`clear_accumulator=True`, while `tile_row_sum_asm` passes it. `V_RED_MAX` +accumulates the same way `V_RED_SUM` does, so a multi-row `tile_row_max` returns a +**running maximum** seeded with whatever the previous user left in `f1`. That is +Mamba and attention code — spawned as its own task rather than fixed inline. + +--- + +## Task 4b(ii) — Decay and beta scalars ✅ `6329a62` / `ff7fa6b` + +**Result: 3 failed (baseline), 225 passed.** + +`decay[h,k] = exp(lower_bound · sigmoid(rate[h] · (gate + dt_bias)))` and +`beta[h] = sigmoid(beta_logit)`, both landing in FPRAM via `S_MAP_FP_V`. +`rate = exp(a_log)` depends only on a weight, so the host precomputes it rather +than paying a vector `exp` per token. + +Sigmoid is four vector ops in place with no scratch. `mamba_silu_v0` runs the same +sequence but needs scratch because it keeps `x` alive for a final multiply. + +**The FPRAM addresses fall out of the row order rather than being computed.** +`tile_row_to_fpram` puts row `i` at `base + i·mlen`, so iterating `(head, key +block)` lands head `h`'s block `kb` at `h·key_dim + kb·mlen` — exactly what the +recurrence indexes. A test asserts that slot by slot rather than trusting it. + +### Review round — `ff7fa6b` + +**No surviving lowering mutation** (17 of 18 red; the one green was a verified +semantic no-op). But one real out-of-bounds write: + +`kda_beta_scalars_v0` checked `beta_fp.size >= num_heads`. **`S_MAP_FP_V` moves a +whole `mlen` row**, so the write is `rows · mlen`. At Kimi's 96 heads against +`mlen` 64 that is **32 slots written outside the allocation**. Nothing downstream +catches it: `_resolve_fpram_addr` bounds-checks only the base offset, and the +emulator's scalar SRAM knows the file but not allocation boundaries. + +My own commit message contained the tell: *"harmless because they are never +indexed"* — that argument is about **reads**. The lanes are never read. They are +written. And the test masked it by sizing `beta_fp` with `_rows_up(num_heads)` — +the padded size the emitter should have required but didn't. + +The oracle would not have caught it either: **Python slice assignment past the end +extends the list** rather than failing, so an OOB `S_MAP_FP_V` silently landed at +the wrong address and stayed green. Now refuses, matching the emulator's assert. + +--- + +## ⚠ Design constraint found here — this decides the mixer's shape + +**FPRAM is 512 slots in RTL** (`FP_SRAM_DEPTH`, `configuration.svh:30`); 1024 in +the compiler's allocator. + +The current scalar layout wants `decay_fp`, `q_hat` and `k_hat` at +`num_heads · key_dim` each: + +| | Kimi K3 | +|---|---:| +| `decay_fp` | 12,288 | +| `q_hat` | 12,288 | +| `k_hat` | 12,288 | +| `beta_fp` | 96 | +| **total** | **36,960** | +| **available** | **512** | + +**72× over.** The reviewer reproduced the overflow at (96 heads, 2 key blocks) +even at `mlen=8`. + +So the mixer **cannot** compute every head's scalars and then run the recurrence. +It has to **stream per head**: produce one head's decay / q̂ / k̂, run that head's +recurrence, reuse the window. That needs `3·key_dim + 1 = 385` slots, which fits. + +`kda_decode_step_v0` already takes `head_rows` for exactly this — the per-head +path was built in Task 3 and is about to earn its keep. + +### Performance + +Static counts scale **linearly in `heads × key_blocks`, ~12 per row** — not +constant. `_arith_progression` returns `None` for a constant sequence (step 0 with +count > 1 would spin a hardware loop forever), and the three broadcast-constant +steps in the sigmoid have identical FPRAM addresses per row, so they unroll. +Kimi's gate path is ~2.4k static instructions per token, mostly that unrolling. +**Preloading the constants into registers once is a clear optimisation**, recorded +for the Phase 1 measurement. + +--- + +## Task 4c — The assembled per-head mixer step — next + +Committed as `85be73f`. 246 → passed, 3 pre-existing failures. Reviewed below. + +### Review — four surviving mutations, and the FPRAM budget was wrong + +The review found the number this whole design rests on was miscounted. + +**1. The mixer did not fit FP_SRAM_DEPTH.** `kda_mixer_fpram_slots` returned +`3·key_dim + 1` and the test asserted `385 ≤ 512`. Two omissions: + +* `beta_fp` is not one slot. `S_MAP_FP_V` moves a whole `mlen` row, so it costs + `ceil(num_heads/mlen)·mlen` — **128** at Kimi K3, not 1. (Task 4b(ii) fixed the + out-of-bounds *write* this causes but did not carry the size into the budget.) +* `rate_fp` is indexed by head number, so it costs `num_heads` — **96**. + +Real high-water was **620 against 512**. And `3·key_dim + beta = 512` on its own, +so no rearrangement of the rest could have saved it. + +The fix is a liveness observation: **`decay` and `q_hat` are never live at the +same time.** Decay is dead once the prediction has run; `q_hat` is not read until +the read-out. So `kda_decode_step_v0` split into `kda_decode_predict_v0` and +`kda_decode_update_v0`, the mixer calls the two halves, and normalises `q` +*between* them — into the window decay just vacated. The wrapper still exists and +`test_the_two_halves_emit_the_same_program_as_the_whole_step` pins it. + +| | slots | +|---|---| +| `k_hat` | 128 | +| `decay_or_q_hat` (shared) | 128 | +| `beta` | 128 | +| `rate` | 96 | +| `part`/`acc`/`scale`/`lower` | 5 | +| `consts` | 7 | +| **total** | **492** | +| **FP_SRAM_DEPTH** | **512** | + +20 slots of headroom. `kda_mixer_fpram_slots` now returns the itemised dict, and +`test_fpram_accounting_matches_what_is_actually_allocated` checks it against the +allocator's real high-water mark — because `FPRAMAllocator` defaults to **1024** +and will not catch an overflow itself. That 1024-vs-512 disagreement between +`memory.py` and `configuration.svh` is still open; the SystemVerilog is +authoritative and both copies of it (compiler `doc/`, `PLENA_RTL/src/`) say 512. + +**2. `head_rows` gave wrong answers, and nothing caught it.** Mutating the mixer +to ignore `head_rows` entirely left the suite green: the test named +`test_heads_can_be_lowered_individually` never passed the argument. Actually using +it was worse than untested — lowering three heads in three calls gave head 0 the +right answer and heads 1–2 an error of 1.2e-2, because `kda_beta_scalars_v0` ran +at the top of *every* call and consumes `beta_logit` in place. Sigmoid is not +idempotent. + +`kda_beta_scalars_v0` moved out of the mixer; the caller runs it once. The test +now lowers per head and asserts bit-identical output against one call for all +heads, and a second test checks that lowering head 0 leaves the other heads' +state untouched. + +**3. The idempotency docstring was backwards.** It warned about the L2 +normalisation, which is the one in-place step that *is* idempotent (a unit vector +normalises to itself, 2.4e-7), and omitted the two that are not: `beta_logit` +(0.43) and `gate` (0.90). Rewritten with the measured numbers. + +**4. A head subset bounded the tile by row *count*, not row *number*.** With +`heads=[12]` the row list has `key_blocks` entries but indexes row 12, so a +12-row check passed a tile of 8 rows and the emitter then read and wrote four +rows past the allocation — silently, into the next VRAM object. Bound is now +`max(rows) + 1`. + +**5. The squares copy was quadratic in head count.** `kda_l2_normalize_blocked_v0` +copied rows `0 .. first_row + n` instead of the `n` live ones, and the mixer walks +`first_row` across every head. At Kimi K3 that is **18,624 rows moved for 384 rows +of live data, 48×**, at ~80 instructions per row. `mamba_block_copy` already took +`src_row_offset`/`dst_row_offset`; using them makes the cost flat in `first_row` +(2,374 → 2,401 instructions across a 256-row offset — the residue is the wider +address immediate). Pinned by `test_the_squares_copy_is_only_as_long_as_the_slice`. + +Also fixed: both `fp_head_stride` guards were untested (no test passed the +parameter at all), and the `sq_scratch` `first_row` bound was untested because +`vec`'s bound always fired first. + +### Mutation results after the fixes + +All ten go red, including the four that previously survived: + +| | mutation | | +|---|---|---| +| M1 | mixer ignores `head_rows` | 2 failed | +| M2 | drop the `stride==0` multi-head guard | 1 failed | +| M3 | drop the negative-stride guard | 1 failed | +| M4 | `sq_scratch` bound drops `first_row` | 1 failed | +| M5 | gates bound back to `len(rows)` | 1 failed | +| M6 | block copy back to row 0 | 1 failed | +| M7 | beta counted as 1 slot | 2 failed | +| M8 | rate left out of the budget | 2 failed | +| M9 | `q` normalised before the predict half | 9 failed | +| M10 | update half reads decay as `q` | 8 failed | + +### Performance + +The 723 recorded above for Task 4c is right. I briefly "corrected" it to 27,243 +and that correction was wrong: `get_code()` returns assembly **text**, so +`len(get_code())` counts characters. The codebase counts instructions as +non-blank non-comment lines -- see `test_records_the_phase0_instruction_count`. + +One head, `key = value = mlen = 8`: + +| stage | static | +|---|---| +| beta (once, all heads) | 33 | +| normalize x2 | 110 | +| `tile_row_to_fpram` x2 | 14 | +| decay | 61 | +| **predict half** | **188** | +| **update half** | **317** | +| **mixer total** | **723** | + +| shape | static | per head | +|---|---|---| +| 1 head, key=value=8 | 723 | 723 | +| 3 heads, key=value=8 | 2,103 | 701 | +| 2 heads, key=16 value=8 | 2,469 | 1,234 | +| 3 heads, key=value=16 | 6,570 | 2,190 | +| 6 heads, key=value=16 | 13,107 | 2,184 | + +**The recurrence is 70% of the mixer** (505 of 723), and it is all the +`copy + multiply + add` triple per key row -- exactly what `V_FMA_VF` collapses. +That is the number Task 8 is measured against. + +Per-head cost is flat in head count (701 at 3 heads against 723 at 1), which is +what the copy fix bought; before it, it grew. + +--- + +## Phase 1, Task 5 — `V_FMA_VF` in the emulator — next + +## Task 5 — `V_FMA_VF` in the emulator — done (`df2ae9fa`) + +`Vector[rd] += Vector[rs1] * fp_reg` at 0x3B. Two places the plan +under-specified, both found by writing the test first: + +* **The masked path starts from the destination, not the source.** `mul_scalar` + starts from the source because there the source *is* what the result is built + from. For an accumulate, a masked-off head must keep what the destination had. + Pinned by `fma_scalar_leaves_masked_off_heads_holding_the_destination`. +* **Timing gets its own arm.** FMA touches two vector rows, like the VV ops, not + the one row the rest of the VF family does. The plan put it in the single-read + group, which would have under-reported exactly the traffic Task 8 is judged on. + +`test_mamba_opcodes_do_not_collide_with_existing_encodings` failed, correctly -- +it pinned 0x3B as free. Boundary moved to 0x3C. 109 passed. + +`cargo clippy` and `cargo fmt` are **not installed** in the Nix dev shell, so the +plan's lint steps could not run. + +## Task 6 — the ISA definition, assembler and spec — done (`93bc504`, `84ef4848`) + +Mnemonic `V_FMA_VF rd, rs1, fp2, rmask`, encoding identical to `V_MUL_VF` except +the opcode field. + +**PLENA_RTL's `operation.svh` stops at 0x34.** 0x35 through 0x3A were already +ahead of the RTL before this work, so `V_FMA_VF` joins six existing opcodes the +compiler and emulator implement and the hardware does not. Not a regression, but +it is the honest status: this instruction needs RTL work that does not exist. + +### A guard, and two pre-existing gaps it found + +`decode` carried three comments saying encodings "must stay in sync with +PLENA_Compiler's `doc/operation.svh`" -- enforced by author diligence alone, +which is the arrangement that let FPRAM 1024-vs-512 drift unnoticed. Now a test +parses the header (the submodule is checked out recursively in every CI job) and +asserts nothing in it decodes to `Invalid`. + +It found `V_PS_V` (0x31) and `C_HADAMARD_TRANSFORM` (0x33): declared by PLENA, +implemented nowhere. Both known -- `program_ssd.py` already documents that +`V_PS_V` "assembles silently and then decodes to Invalid, so emitting it is worse +than not emitting it" -- but that lived in one docstring. They are now an +explicit exemption list carrying the reason. Verified by deleting the 0x3B arm: +the guard names the opcode and fails. + +## Task 7 — the FMA row-sweep emitter — done + +`_emit_tile_row_fma` with four GP pointers, plus `tile_row_fma_fp_asm`, +`tile_row_fma_fp_sweep_asm` and `tile_row_fma_fp_broadcast_asm`. + +**Named `_sweep`, not `_broadcast` as the plan said.** In `isa_tile_rows.py` +`_broadcast` means one FPRAM slot applied to *every* row; the plan's function +walks `fpram_base + i`, which is the opposite. Using the plan's name would have +made every call site read backwards. Both now exist, with the file's meanings. + +### `_arith_progression` refused a step of 0, and its stated reason was wrong + +The comment read "Constant sequence (step=0, count>1) would cause infinite HW +loop". That is not what happens: `C_LOOP_START` takes its trip count as an +immediate and `C_LOOP_END` decrements a dedicated loop register +(`loop_state.rs`), neither of which the address step touches. The single-element +case has always returned step 0 and looped fine. + +It matters because **a pinned side is exactly a step of 0**, and every recurrent +contraction has one -- prediction and read-out walk the state rows against one +accumulator row. Removing the restriction breaks nothing (265 passed, same 3 +pre-existing failures) and changes no existing output. + +**It also buys nothing yet, and it is worth being clear about that.** The mixer +is 723 instructions before and after. The current sweeps loop in *Python* -- +`mamba_row_copy` is called once per row with a single row argument, so there is +no row list for `_row_progression` to see. Step-0 support is a **precondition** +for Task 8, which passes whole row lists to the FMA emitter; it is not a win on +its own. + +### The oracle had to learn the opcode + +`isa_interpreter.py` correctly refused `V_FMA_VF`, so it is now modelled -- as +one expression, `dst + src * f`, matching `fma_scalar`'s single quantisation +rather than a multiply followed by an add. `test_unmodelled_opcode_raises` used +`V_FMA_VF` as its example of something unmodelled and now uses `V_MAX_VF`. Two +new oracle tests pin the accumulate and the `dst == src` case. + +265 passed, 3 pre-existing failures. + +--- + +### Review of Task 5 — no bug in the instruction, three gaps around it + +The arithmetic, the masked-path base, the timing classification and the +cross-repo encoding all checked out. What did not: + +* **The profiler charged `V_FMA_VF` one cycle instead of five.** + `isa_analysis.py`'s `instruction_cycles` is a prefix cascade ending in a + catch-all `return 1`, and `V_FMA_VF` does not start with `"V_MUL"`. The + emulator charges `VECTOR_MUL_CYCLES`. A **5x under-count per FMA, in the one + tool Phase 4 measures with**, on the one opcode being measured. It was also + missing from `selected_opcodes`, so the recurrence would not have appeared in + the per-opcode breakdown at all. Both fixed. +* **`analytic_models/performance/customISA_lib.json` had no entry**, and + `perf_model.py:147` raises `KeyError` on a miss rather than defaulting. Added, + matching `V_MUL_VF`'s cycle expressions. +* **The dispatch arm and the timing arm were pinned by nothing.** Swapping + `rd`/`rs1` in dispatch -- which makes the instruction compute + `V[rs1] += V[rd] * f` -- left 110/110 green, because *nothing in the emulator + executes a decoded instruction*; every test calls the machine methods + directly. Same for the timing arm, the commit's own headline justification. + +The timing arm now has a real test. The dispatch arm has a source-parsing guard, +the technique `stage_profile.rs` already uses, because closing it properly needs +an `Accelerator` built in a test with Ramulator and a `MatrixMachine`. **Every +opcode has this gap**, not just the new one -- spawned as separate work. + +Two numbers worth keeping from the review: + +* **FMA is 1.20x more accurate than the triple in RMS** (bf16, VLEN 64, 64 + trials): max_abs 3.89e-3 against 5.81e-3. Direct difference 7.81e-3, about 2 + bf16 ulp, growing to 6.25e-2 over a 64-step sweep. +* **No test observes that difference.** `isa_interpreter.py` is float64, where + the FMA and the triple are algebraically identical, and nothing runs KDA + numerics through the transactional emulator. The difference is real, it + favours the FMA, and it is currently unmeasured. + +Also noted: `QuantTensor::quantize` is a **no-op** (`tensor.rs:128`, "TODO: add +actual quantization") -- it only attaches a type tag, and all rounding happens at +`VectorSram::write`. So "one rounding versus two" is true, but it is the SRAM +store that rounds. + +--- + +## Task 8 — both recurrences on `V_FMA_VF` — done + +### `ssm_decode_step_v0` had no test at all + +The conversion landed and the whole suite stayed green. `grep -rl +ssm_decode_step_v0` returns the emitter, its caller in `aten/ops/plena/mamba_ops.py`, +and a docstring mention. **No test.** A rewrite of Mamba's state update and +output contraction was therefore unverified. + +`aten/tests/test_ssm_decode_step.py` now runs it against +`mamba2_recurrent_reference` through the ISA interpreter -- five geometries +including multi-group, plus a seeded-garbage check that the output accumulator is +actually cleared. Eight tests. + +### The measurement + +Static instructions, one head: + +| | before | after | +|---|---|---| +| KDA decode, key_dim 8 | 505 | **76** | +| KDA decode, key_dim 16 | 961 | **76** | +| KDA decode, key_dim 64 | 3,697 | **76** | +| KDA decode, key_dim 128 | 7,345 | **76** | +| Mamba decode, state 4 | 208 | **54** | +| Mamba decode, state 8 | 376 | **54** | +| Mamba decode, state 16 | 712 | **54** | +| Mamba decode, state 32 | 1,384 | **54** | + +**The Mamba "before" column was wrong when first written** (200 / 352 / 656). I +measured it by reverting only the conversion, at a HEAD that already carried the +step-0 relaxation -- which had itself shrunk `ssm_decode_step_v0` from 712 to +656. The true pre-Phase-1 numbers, at `ede7454`, are above. The KDA figures were +measured the same way and are unaffected, because the KDA recurrence had no +step-0 progression to gain from. + +**The cost went from linear in the contracted dimension to constant.** At Kimi +K3's key_dim of 128 that is 7,345 -> 76, a factor of **97**. That, and not the +saved multiply, is the case for spending opcode 0x3B: the triple needed a scratch +row per step, and because it was the *same* row every iteration the destination +never formed a progression, so the sweep could not become a hardware loop. + +Whole mixer, per token: + +`kda_mixer_step_v0` alone, **not** counting the one-off `kda_beta_scalars_v0` +(33 instructions) that the caller now runs separately: + +| shape | after | +|---|---| +| 1 head, key=8 value=8 | 241 | +| 1 head, key=16 value=8 | 306 | +| 1 head, key=16 value=16 | 378 | +| 2 heads, key=8 value=8 | 482 | +| 3 heads, key=16 value=16 | 1,134 | +| 6 heads, key=16 value=16 | 2,268 | + +Exactly linear in head count at a fixed shape, which is what the per-head +streaming is for. The earlier version of this table mixed `kda_beta_scalars_v0` +into some rows and compared across differing `value_dim`, so its "before" column +was not comparable row to row; the decode-step figures above are the ones to +quote. + +`test_instruction_budget.py` gates both kernels, and +`test_the_cost_is_flat_in_the_contracted_dimension` asserts the flatness +directly. A budget alone would still pass if the cost grew and someone raised +the number. + +### The scratch tile is gone + +`kda_decode_step_v0` no longer takes one: the FMA accumulates in place, which is +most of what it bought. `KdaMixerBuffers` loses the field too. The aliasing check +that still matters is `pred` and `err` against `state`. + +Seven mutations of the conversion, all red: dst/src swapped in both kernels' +rank-1 updates, the contraction reading the wrong FPRAM array, the D skip +losing its accumulate, the read-out losing `fp_base`, and the emitter emitting +`V_MUL_VF` instead (32 failures). + +284 passed, 3 pre-existing failures. Emulator 112 passed. + +--- + +## Task 9 — KDA chunk primitives — done + +### The plan's formula was wrong, and the reference has no chunked path to check against + +`kda_state_engine_prefill` is a loop of `kda_state_engine_step`. So the only +oracle for a chunked form is "it must equal the sequential one", and the algebra +had to be derived and checked before any of it was worth lowering. That was done +in plain torch first. + +**KDA's decay is channel-wise on the key axis** -- `reference.py` applies +`exp(log_decay)[:, :, None, :]`, one factor per key channel, not one scalar per +timestep. The textbook chunking assumes a scalar decay, and the plan's +`_dense_ut` follows it: `tril(diag(beta) @ (K @ K.T), -1)`, with no decay +weighting. Against the sequential reference that form is **wrong by 1.8e-1**. +The correct gram is `M[t,s] = sum_key k_t k_s A_t/A_s`, which reproduces the +sequential output to **4.5e-8** and the state to 1.6e-7. + +### Three numerical findings, all measured + +**1. `chunk * |gate_lower_bound|` must stay under 88.7.** `M` and `N` are formed +as a matmul of `k*A` against `k/A`, and `1/A` reaches `exp(chunk*|lower_bound|)`. +bf16 tops out at 3.39e38, `ln` 88.72. At Kimi K3's -5 the **last chunk that works +is 17**; chunk 18 is `inf`, and then `nan` through the whole solve. The plan said +"do not raise the chunk size without redoing this analysis" -- this is the +number, and it is now an assertion rather than a comment. + +**2. The cumulative decay must be a running product, not `exp` of a running +sum.** Following `program_ssd.py`'s cumsum was the obvious move, and it is wrong +here: `c` reaches -80, where bf16's ulp is 0.31, and `exp(0.31)` is a 36% +relative error on `A`. Multiplying per-step decays keeps every intermediate in +`[e^-5, 1]`, where bf16's error is relative rather than magnified. + +| route | output error | state error | +|---|---|---| +| cumsum in f32 (the algebra's own floor) | 8.1e-05 | 9.5e-04 | +| cumsum in bf16 | 3.7e-04 | **1.1e-02** | +| **cumprod in bf16** | **5.9e-05** | **2.4e-03** | + +**3. The matmul form's cancellation does not reach the output.** Forming `M` from +bf16 `k*A` and `k/A` gives up to **315% relative error on M's smallest entries** +-- which looked disqualifying. It is not: those entries contribute nothing, and +end-to-end the matmul form and an exactly-computed one agree to 6e-5. So the +matrix engine can do this contraction after all, and the vector-unit fallback +(C^2 contractions, no cheaper than just running decode 16 times) is not needed. + +### What was built + +`aten/plena/program_kda_chunk.py`: + +* `kda_chunk_decay_cumprod_v0` -- running product per key channel, in place. + `mamba_block_copy`'s independent row offsets stage one whole timestep at a + time, so the serial scan needs no new shared emitter. +* `kda_ut_transform_v0` -- `T = (I + tril(diag(beta) M, -1))^-1 diag(beta)` by + forward substitution. `T[i] = beta_i (e_i - sum_{j mlen` is refused, +`key_blocks != 1` is refused, and the test pinned `value_dim == mlen` — so +`t_blocks`, `key_blocks` and `val_blocks` were all 1 and *any* block-index error +was undetectable. + +`test_kda_prefill_structure.py` parses the emitted assembly, where each +projection announces its operands and indices, and pins the seven products in +order. Seven mutations, all red — including dropping the spill tail zero-fill, +which is **invisible on the emulator**: every spilled tile is freshly allocated +so its tail is already zero, and the state projection's other operand happens to +be zero there too. One accident masking a whole class of over-read, now pinned as +an invariant. + +### 6. Prefill and decode hold the state transposed relative to each other + +decode is `[key, value]` (row progressions become hardware loops; shared with +Mamba). prefill is `[value, key]` (all seven products land on the projection +primitives without a transpose). Both deliberate — and at Kimi K3 +`key_dim == value_dim == 128`, so the shapes match and passing one to the other +is a finite plausible wrong answer. + +`kda_prefill_state_to_decode_layout_v0` converts at the boundary, one projection +against a staged identity. Verified on the emulator: **exact, zero error**. + +### Performance + +Broadcasting one row across a tile was a copy per row — 64 at `mlen` 64, 128 at +Kimi's `value_dim`. Doubling the filled span makes it `ceil(log2(n)) + 1`: +prefill at chunk 16 goes **1,803 → 1,476** static instructions. + +Against decode, at `mlen` 64, one head, `key_dim = value_dim = 64`: + +| | static | per token | +|---|---|---| +| prefill, chunk 16 | 1,476 | 92.2 | +| prefill, chunk 8 | 1,039 | 129.9 | +| decode × 16 | 1,152 | 72 | + +**Prefill is still 1.28× more static instructions than running decode 16 times.** +That comparison is not the case for prefill, though — instruction count is not +cycles. Prefill's seven `M_MM_WO` each drive 64³ MACs through the systolic array +while decode's are cheap `V_FMA_VF`. One chunk of 16 tokens measures +`sim_latency_cycles = 173,939`; the decode comparison at the same `mlen` is +Task 14's job and is the number that decides it. + +328 passed, 3 pre-existing failures. Emulator: 112 passed, and the KDA stage +tests `cumprod`, `ut`, `prefill_out`, `prefill_state`, `prefill_chain_out`, +`prefill_chain_state`, `state_transpose` all pass. + +--- + +## Task 11 — the whole-model KDA lowering — next + +## Task 11 — whole-model lowering — done + +### I claimed a blocker. It was not one. + +My first pass said the layer could not be assembled without giving +`kda_conv_step_v0` and `kda_l2_normalize_blocked_v0` a column-block index, and +recommended that as the fix. An adversarial pass took the claim apart and built +a working layer instead. Both halves of my premise were wrong: + +* **`kda_conv_step_v0` does not walk feature blocks as rows.** Of its six tile + operands only `x_new` comes from the projection, and it is read at exactly one + place — a single row index. Everything else is a dense tile the compiler + allocates. +* **`kda_l2_normalize_blocked_v0` is not on the projection path at all.** The + mixer's `q` and `k` are the *convolutions'* output, which is dense — my own + `test_conv_channel_blocks_and_mixer_key_blocks_are_the_same_rows` pins that + seam. The normalisation never sees a projection tile. + +And "`tile_col_idx` already threads through the `tile_row_*` family" was true of +the ISA layer and false of the Var layer the KDA emitters actually call, where +every wrapper drops it. So my recommended option was not "add a parameter" but +"build a cross-column-block copy primitive". Measured, it cost **+25%** on the +normalisation and lost a cross-block hardware loop — to solve a problem that is +not on the path. + +My cost estimate for the alternative was wrong by two orders of magnitude too: I +said a scatter would be ~700 block operations per layer; done naively it is +7,674 static instructions (6.7% of the layer), and done properly it is **70**. + +### What the layer actually needed + +Column block `c` sits at `base + c * physical_rows * mlen` and rows within a +block are linear, so a projection tile's bytes already **are** a dense +`[blocks * physical_rows, mlen]` tile — feature block `c`'s token `t` at row +`c * stride + t`. `vram_tall_view` names it; `kda_gather_projection_v0` moves a +section into place with **one `V_FMA_VF` sweep**, because both the source rows +(step `stride`) and the destination rows (step 1) are progressions. + +**14 static instructions per section, independent of block count.** A 96-head +gather costs exactly what a 2-head one costs — pinned by a test. All five +sections of a Kimi K3 layer: **70**. + +`test_the_assembled_layer_matches_kda_state_engine_step` now runs gather → three +convolutions → gates → recurrence against the reference's own boundary. **Zero +emitter changes were needed.** + +### The address bug both reviewers found + +`vram_column_block_view` used `base + c * mlen * mlen`. The real stride is +`c * physical_rows * mlen` (`memory.py`'s `col_block_base`); they coincide only +when a tile is exactly `mlen` rows tall, which was the one shape I tested. A +projection allocated `strict=False` is padded to `blen` rows, so a real decode +projection was off by `mlen / blen` and the view pointed into a different tensor. + +**Both the wrong formula and the right one passed the whole suite.** The test +seeded and read at the view's *own* address, so it was self-consistent with any +address the view invented. It now anchors against `_tile_addr` on the parent, at +four shapes including two row blocks and a `strict=False` padded one, and checks +a neighbouring block is untouched. The old formula now fails four tests. + +Also fixed: the duplicate-name hazard (`register_vram_matrix` overwrites silently +and views resolve by name at emit time, so with a default prefix the first two +splits in a program collide — and the layer runs 93 times); reading the var's +`physical_shape` where the layout is authoritative; and validating `key_width` +where the invariant is per *head* — 4 heads of `key_dim` 6 gives a `key_width` +divisible by `mlen` with three heads starting mid-block. `kda_key_blocks` already +raised exactly that error and was imported but never called. + +`alloc_at` already existed and did what I had reimplemented; the view now goes +through it. + +### Measured, Kimi K3 at mlen 64 + +| | static | share | +|---|---|---| +| three convolutions | **53,757** | 58% | +| mixer | **39,526** | 42% | +| gather | **70** | 0.07% | +| **layer** | **93,353** | | +| **× 93 layers** | **8,681,829** | | + +**The convolutions are the larger half and were ungated** — the budget covered +42% of what it was meant to protect. Now gated at 59,000, with a whole-layer gate +at 103,000 and a test that the gather stays a rounding error. + +A whole-model figure that includes projections, MoE and embeddings is not comparable; +8.68M here is the state-engine path only, without projections, MoE or +embeddings, so the two are still not comparable. + +### Prefill is closer than I said + +`kda_l2_normalize_v0`, `mamba_conv1d_v0` and `kda_chunk_decay_cumprod_v0` all run +directly on a column-block view today. What prefill still lacks is a `v_t` +transpose, a beta stride, a per-token decay activation, and an assembler — and +above all **`key_dim <= mlen`**, which Kimi K3 (128 against 64) fails. That is +the real prefill gap, and it is larger than the seam I had been calling one. + +*(Task 16 lifts the `key_dim <= mlen` limit, and `value_dim <= mlen` with it.)* + +351 passed, 3 pre-existing failures. Emulator 112 passed. + + +## Task 12 — the layer executed in the emulator, first layer to last + +This path had never been run as machine code. It now is. + +### On the transactional emulator, against `kda_state_engine_step` + +| case | result | max abs error | +|---|---|---| +| `layer` — one assembled layer | **PASS** | 2.14e-04 | +| `layer_chain` — four, back to back | **PASS** | 2.75e-04 | + +Gather, three convolutions, the gates and the recurrence, compiled, assembled, +and executed by the Rust emulator — not the float64 ISA interpreter. Three heads, +`key_dim` 128, so `q` and `k` each span two blocks. + +Four layers give 2.75e-04 against one layer's 2.14e-04: composing layers costs +essentially nothing, because each carries its own state and the error does not +chain. + +### Two things only depth could find + +**1. The second layer collided on view names.** `kda_layer_from_projected_v0` +named its tall view with a fixed prefix, so layer 1 and layer 2 asked for the +same VRAM object. The duplicate-name guard added in the Task 11 review turned +that into an error at compile time; without it, +`register_vram_matrix` overwrites silently and views resolve by name at emit +time, so **layer 1's views would have been repointed at layer 2's projection**. +Fixed by naming views after the projection tile, which is unique per layer — so +stacking needs nothing from the caller. + +**2. FPRAM is scratch, not storage.** Allocating the mixer's window per layer +overflowed the compiler's optimistic 1024 slots **at the fourth layer** — and the +hardware has 512. One window serves the whole model, the same reuse the mixer +already does per head. Pinned: FPRAM use must not grow with depth. + +### The HBM address space — the constraint that decides executability + +The emulator preloads HBM from a flat file starting at offset 0, so the +allocation and the file both have to span every address the program touches. A +layout whose regions sit far apart needs a span-sized allocation whatever the +live data comes to, and stops being executable long before it stops being +describable. Addressing for the state-engine path, excluding the projection +weights any implementation needs: + +| | | +|---|---| +| per layer | **2,336,832 bytes** (2,282 KiB) | +| 93 layers | **217,325,376 bytes** (0.202 GiB) | +| emulator `HBM_SIZE` | 16 GiB | + +**Two orders of magnitude below the wall.** (The first version of this said 0.33 +MB per layer — seven times too low, because the measuring helper put the state +and the projection in VRAM while the emulator test staged them from HBM. See the +review below.) + +### What Task 12 does not claim + +This is the KDA state-engine path first layer to last, not a whole model: the +input and output projections, the MoE blocks, the norms and the embeddings are +shared code with no Nemotron-3 or Kimi-K3 config in this repository to +instantiate them from. Task 13's real-checkpoint binding needs that config +first, and building it is a larger piece of work than anything in Phases 0–2. + +354 passed, 3 pre-existing failures. + +## Task 14 — the banking study and the measurement report — done + +### The banking question has a decisive answer: stalls ≈ 0 + +`analytic_models/performance/vector_sram_banking.py` extracts every hardware +loop's pointer stride **from the emitted assembly** — not from a description of +the algorithm — and models row-interleaved banking. + +| kernel | loops | B=2 | B=4 | B=8 | B=16 | B=32 | +|---|---|---|---|---|---|---| +| KDA recurrence (Kimi K3) | 20 | 0.6% | 0.9% | 1.0% | 1.1% | 1.1% | +| KDA mixer, one head | 48 | 1.2% | 1.8% | 2.1% | 2.3% | 2.3% | +| Mamba-2 decode step | 7 | 1.0% | 1.5% | 1.8% | 1.9% | 2.0% | + +The dominant loops are **stride 1 over 128 trips**, which visits every bank in +turn and never conflicts. The 1–2% residue is entirely single-trip loops with a +pinned pointer. **No evidence for a Vector SRAM bank-mapping mode.** + +Worth noting the emulator does not model banking at all — `VectorSram` is a flat +row vector — so this is analytic, and a word-interleaved mapping would moot the +question regardless. + +### The follow-up the plan predicted, now quantified + +| kernel | dynamic | `S_ADDI_INT` | `S_LD_FP` | `V_FMA_VF` | +|---|---|---|---|---| +| KDA recurrence (Kimi K3) | 4,179 | **50.0%** | 24.6% | 18.4% | +| KDA mixer, one head | 4,398 | **50.2%** | 23.9% | 17.5% | +| Mamba-2 decode step | 730 | **46.4%** | 26.6% | 17.7% | + +**Half the dynamic stream is pointer arithmetic and under a fifth is +arithmetic.** Post-increment addressing on the vector pointers, plus an +auto-advancing FP pointer, would remove up to three-quarters of the dynamic +instructions on the recurrent path — a far larger effect than banking, and the +next instruction-set question worth asking. + +`doc/static_path_measurements.md` collects everything, with hard facts and +uncalibrated model output kept apart and **no comparison drawn against any other +uncalibrated latency figure**, because two of those decide nothing between them. + +## Task 15 — nothing to retire, and a guard so it stays that way + +There is no descriptor machinery, no command queue, no residency cache, +no `spec/`, no descriptor, no queue, no residency cache. This branch was cut from +`a4b3e7de` on main, which never had any of them — the plan's deletion steps have +nothing to delete. + +That is the answer to the question this whole project started from. The static +path needs **three** new instructions -- `V_SOFTPLUS_V` (`0x39`), `S_MAP_FP_V` +(`0x3A`) and `V_FMA_VF` (`0x3B`) -- and every one is an ordinary fixed-function +op with its operands named in the instruction word. Everything else — the transposed state layout, the per-head +streaming, the shared FPRAM window, the chunked prefill, the projection gather — +is compiler work. + +`aten/tests/test_no_state_engine.py` makes that a property rather than an +observation: no run-time state-engine vocabulary in any tracked source file, +those directories absent, `0x3D` and `0x3F` still free, and the three opcodes +this work adds each pinned to its slot. A grep is a check that ran once; a test +is one that keeps running. + +Two docstrings in the CPU reference still named the boundary after a +descriptor-driven instruction. Reworded — the boundary is the right place to +compare against, and the name was for something this design does not have. The +state dataclass moved with them, from `KdaXState` to `KdaRecurrentState`. + +358 passed, 3 pre-existing failures. + +## Task 13 — real checkpoint weights — done for Mamba-2, blocked for KDA + +Whole-model +artifacts were *"never bound to a real checkpoint"*. One is bound now. + +### `AntonV/mamba2-130m-hf`, layer 0 + +A published Mamba-2: 24 layers, `hidden_size` 768, `state_size` 128, 24 heads of +`head_dim` 64. `ssm_decode_step_v0` runs on its real `A_log`, `D` and `dt_bias` +and matches `mamba2_recurrent_reference` to **rtol 1e-4 / atol 1e-5**, on both +the output and the carried state. + +What real weights test that `randn` cannot: + +* **Structure.** Layer 0's `A` is not a spread — 23 of 24 heads sit between + −0.73 and −0.27, and one sits at **−5.06**, about ten times the rest. A tight + cluster plus a single outlier is what training produces and a symmetric draw + does not, and that outlier head is the one most likely to expose a decay that + underflows. Pinned as a test. +* **Layout.** `in_proj` packs `z | x | B | C | dt` into one 3352-wide tensor, and + `conv1d.weight` is `[channels, 1, kernel]` rather than `[channels, kernel]` — + the reference accepts both, so a lowering that assumed the wrong one would + silently transpose a 4-tap filter. Pinned. + +I first asserted `A.min() < -10` on the strength of nothing; the real value is +−5.06. Corrected to assert the structure that is actually there. + +The test **skips when the checkpoint is not cached**, so CI never needs network. + +### KDA's half is blocked, and on what + +No published KDA checkpoint was found whose parameter semantics match +`aten/models/kda/reference.py` closely enough to bind without guessing. The +gated-delta-net models on the Hub (`m-a-p/*-GatedDeltaNet-*`, +`linear-moe-hub/Gated-Deltanet-*`, `Idiap/gated-deltanet-swa-1.4B-30B`) are the +right family, but matching a checkpoint's exact gate parameterisation to this +reference is a verification task of its own — binding the wrong semantics would +produce a green test that proves nothing, which is worse than no test. + +That is the remaining gap, and it is bounded: one checkpoint, read carefully +against `activate_log_decay` and `kda_step`. + +361 passed, 3 pre-existing failures. + +--- + +## Where the plan stands + +| Phase | | | +|---|---|---| +| 0 — KDA decode, no new opcode | Tasks 1–4 | **done** | +| 1 — `V_FMA_VF`, measured | Tasks 5–8 | **done** | +| 2 — chunked prefill | Tasks 9–10 | **done** | +| 3 — whole model, executed | Tasks 11–12 | **done** | +| 3 — real checkpoint | Task 13 | **Mamba done, KDA blocked on a checkpoint** | +| 4 — measure, retire | Tasks 14–15 | **done** | + +The two things the project set out to establish both hold, and both are now +demonstrated rather than argued: + +**KDA and Mamba-2 need no descriptor machinery.** Three new instructions -- +`V_SOFTPLUS_V`, `S_MAP_FP_V` and `V_FMA_VF`, at `0x39`-`0x3B` -- each an +ordinary fixed-function op. No descriptors read at run time, no queue, no +residency cache — and a test that keeps it that way. + +*(An earlier version of this record said "one new instruction". That counted +only `V_FMA_VF`, the one the performance argument turns on, and omitted the two +the Mamba path needs. Relative to `main` this work adds three.)* + +**The static form is smaller, not merely equivalent.** The recurrences went from +linear in the contracted dimension to constant — 7,345 static instructions to 76 +at Kimi K3's `key_dim` — because removing the scratch row let the sweeps become +hardware loops. The whole path's HBM footprint is 0.029 GiB, three orders of +magnitude inside the emulator's flat allocation. + +What is still open, in order of size: a Nemotron-3 / Kimi-K3 model config (which +Tasks 11–13 all touch), prefill for `key_dim > mlen`, and post-increment +addressing — which the measurements say is worth more than anything else on the +list, since half the dynamic instruction stream is pointer arithmetic. + +*(`key_dim > mlen` is closed by Task 16, below. The model config and +post-increment addressing are still open.)* + + +## Review of Task 12 — the emulator case could not fail + +The most serious finding in this project so far, and it was mine. + +### `case_layer`'s tolerance was larger than its signal + +`atol=2e-2` against a golden whose largest value is **0.01215**. Every one of +the 192 compared values sat inside the absolute bound, so **writing zeros into +the output tile scored 100%**. Verified, along with three more that all passed: + +| mutation | before | after | +|---|---|---| +| zero the output tile | PASS | **FAIL** | +| skip the `v` convolution | PASS | **FAIL** | +| swap the `q` and `k` gather sections | PASS | **FAIL** (6.25% match) | +| zero every layer but the last (`layer_chain`) | PASS | **FAIL** | + +It is `atol=0` with a relative bound now. The honest numbers, which the +absolute figure was hiding: **mean relative error 4.3%**, and 96.35% of lanes +inside 12% (the harness passes at ≥90% of lanes, not all). + +4.3% is what bf16 gives — the read-out contracts 128 terms, so error grows like +`sqrt(128) * 2^-9 = 2.2%` before the rest of the chain. **The lowering is right +to 5e-4 in float64**, which `test_the_assembled_layer_matches_kda_state_engine_step` +checks on the ISA interpreter. The emulator case adds bf16 realism, not +verification power — and quoting only its absolute error implied otherwise. + +### `layer_chain` tested one layer + +`target = outs[-1]` with `golden = golden_rows[-1]`, so layers 0–2 were never +compared. Zeroing all their staged inputs left the verdict **bit-identical**. + +So **"the error does not chain" was not a measurement** — nothing was chained +(the four layers are independent by construction), and the earlier layers were +outside the comparison anyway. Worse, the code to pack every layer existed but +was gated on `if heads == 1 else None`, and the only invocation the commit +recorded was `--num-heads 3`. Every layer is compared now, at every head count. + +Still not compared: the updated conv history and recurrent state. For a decode +token the state is the layer's only persistent effect, so a mixer that produces +the right output and corrupts the state would pass. + +### The FPRAM test proved something else + +No KDA emitter allocates FPRAM at all — every slot comes from the caller, and +the test's own helper shares the window. So it could not catch the failure its +commit described, which is a *caller* allocating per layer. + +It now asserts the number that actually binds: **492 slots at 93 layers against +`FP_SRAM_DEPTH` 512**. The compiler's allocator defaults to 1024 and would not +catch an overflow, which is why the hardware figure is asserted directly. + +### The HBM figure counted one seventh of the traffic + +The helper put the recurrent state, `dt_bias` and the projection tile in **VRAM** +while `case_layer` staged them from **HBM** — two tests disagreeing about where a +layer's state lives. Counted consistently, a layer is **2,336,832 bytes**, not +331,776, and **4.5× over the gate I had set**. 93 layers is 0.202 GiB, which +still sits three orders of magnitude inside the emulator's flat allocation. + +### Cycles, newly measured + +| | cycles | +|---|---| +| `layer` (1 token, 3 heads, `key_dim` 128) | **23,715** | +| `layer_chain` (4 layers) | **96,114** (24,029/layer, linear) | + +By stage marker, inside the real program: `kda_state_update` + `kda_readout` = +**24.9%**, convolutions 13.8%, normalisation 9.2%, gather **0.9%**. The 38% +attributed to "other" is HBM prefetch of this harness's per-layer staged inputs, +which is a property of the harness rather than of a decode loop. + +361 passed, 3 pre-existing failures. + +## Task 16 — prefill past one block: `key_dim` and `value_dim` above `mlen` + +The open item Task 12's summary put second is closed. Kimi K3's 128-wide head +prefills at `mlen` 64, and so does its 128-wide value axis, which nothing had +ever run. + +### The fix is a layout choice, not new machinery + +The key axis moved from **rows** onto **lanes**: a `[chunk, key_dim]` tile is now +`chunk` rows across `key_dim / mlen` column blocks rather than +`chunk * key_blocks` rows one block wide. + +That is what the systolic array already wanted. Five of the seven products +contract over key, and `vram_sub_projection_asm_impl` walks a VRAM operand's +column blocks at stride `physical_rows * mlen`, accumulating in the array — so a +key axis on lanes contracts across as many blocks as it needs with no extra +instruction and no explicit sum. **Not one projection call changed.** What +changed is the tile widths and the row ops, which now run once per column block. + +The alternative — keeping the folded rows and computing the gram of the folded +tile — computes the `kb != kb'` cross terms and throws half of them away, then +needs a strided gather to recover the diagonal. It was never worth writing down. + +### Verified as machine code, every case at 100% of values in tolerance + +| case | `key_dim` × `value_dim` | max abs error | +|---|---|---| +| state layout transpose | 128 × 128 | **0.00** | +| prefill, out | 128 × 128 | 6.71e-04 | +| prefill, state | 128 × 128 | 9.28e-03 | +| three chunks chained, out | 128 × 128 | 6.10e-04 | +| three chunks chained, state | 128 × 128 | 7.32e-03 | + +Static instructions at chunk 16, one head: 1,476 at 64 × 64 — **unchanged** by +the layout move, so the lift cost nothing at the shape that already worked — and +2,426 at Kimi K3's 128 × 128, which is 1.6x the instructions for 4x the state. + +### The defect this turned up, and the tolerance that reported it as a pass + +`out = scale * (...)` ended in a single `tile_row_mul_fp_broadcast`, and a row op +reaches one column block per call. At `value_dim` 128 the upper half of every +token went unscaled — `sqrt(key_dim)` = 11.3x too large. + +The case printed **PASSED**. `check_mem`'s `allclose_pass` is +`match_rate >= 90.0`, and 94.68% of values sat inside an `atol` of 5e-2 on data +of order 1e-3 — a tolerance fifty times the signal. Comparing the value blocks +separately is what showed it: block 0 at 6.4e-04, block 1 at 1.4e-01, mean +relative error 11.28 against a predicted 11.31. + +So the stage test now restates the bar at 100% on the same numbers. +`check_mem` is shared and its rule is left alone. Under the new bar four +mutations of the column-block loops die, and **only one of them needed it**: + +| mutation | match rate | caught by the 90% rule? | +|---|---|---| +| output scale on block 0 only *(the real defect)* | 94.68% | **no** | +| `k_hat`/`q_hat` on block 0 only | 0.00% | yes | +| `A_C * state` on block 0 only | 55.79% | yes | +| `k / A` on block 0 only | 85.80% | yes | + +This is the second time in this project a tolerance larger than the signal has +hidden a real error, after `case_layer`'s. The layer cases moved to `atol` +2.5e-04 alongside their `rtol` 0.12 for the mirror-image reason: three of their +sixty-four outputs are near 1e-03 and bf16 delivers ~1.8e-04 absolute whatever +the value, so a *pure relative* bound is unmeetable there. They had been passing +at 95.31% on the 90% rule, not on the tolerance. The floor keeps the check real — +an all-zero output scores 1.6% against it. + +### Two smaller things fixed on the way + +* `kda_prefill_spill_v0` zero-filled `range(live_rows, mlen)`. For a tile taller + than one block — the state at `value_dim` 128 — that range is empty, so the + dead rows past `mlen` were never zeroed. It is `range(live_rows, tile.shape[0])` + now. +* `contrib` was one scratch tile serving three products of three different + shapes. They are `contrib`, `readout_contrib` and `state_contrib` now. Sizing + one tile to the maximum would have passed `vram_add`'s width check by + coincidence whenever `key_dim >= chunk`, and added the wrong lanes when not. + +`kda_prefill_tile_shapes` is now the single source of the tile shapes, consulted +by the emitter's validation *and* by every caller. They used to be written +separately — the emitter kept a hand-written list of row counts and the callers +allocated `mlen x mlen` for everything, which agreed only at the one shape anyone +ran. That divergence is how the spill over-read got in. + +377 passed, 2 pre-existing failures. All 17 emulator stage cases at 100%. + +### Still refused + +`key_dim` and `value_dim` must each be a whole multiple of `mlen`, and `chunk` +must fit one block. The state transpose additionally requires its output to be +exactly `key_dim` rows when there is more than one value block, since that is +what makes its column-block layout coincide with decode's row indexing. + + +## Correction — one of the "3 pre-existing failures" was mine + +Every task record above closes with "N passed, **3** pre-existing failures". Two +of those three are genuinely pre-existing. The third was introduced by this +branch, and calling it pre-existing was wrong. CI is what caught it: the +`moe-stage-guard` job failed on the PR, and it passes on `origin/main`. + +`test_moe_stage_attribution.py::test_stage_parameters_have_no_default` forbids an +emitter from giving its `stage` parameter a default, because a caller that +forgets the argument is then billed to the wrong stage silently rather than +failing. Three emitters **this branch adds** gave it one: + + program_mamba_common.py mamba_silu_v0(stage='mamba_gated_norm') + program_ssd.py ssd_transposed_projection_v0(stage='mamba_in_proj') + program_ssm_recurrent.py ssm_decode_scalars_to_fpram_v0(stage='mamba_dt') + +Checked against `origin/main` in a separate worktree: the guard passes there (9 +passed) and all three files are additions on this branch. The other two failures +(`test_qwen_packed_skinny_router_rowpacked_compiles_for_128_experts`, +`test_mxfp8_is_sole_gap_source`) do fail on `origin/main`, so those labels stand. + +The fix. `mamba_silu_v0` takes a **required `marker`** instead of a defaulted +`stage`: it serves both the Mamba and the KDA vocabulary, so no default is right +for both, and its own KDA caller was already working around the default by +passing `marker=`. Markers are sticky, so accepting `mamba_gated_norm` there +billed KDA's SiLU *and everything emitted after it* to a Mamba stage — which is +precisely the failure the guard exists to prevent, sitting in the tree +unreported. The other two emitters have no callers at all; `stage` is now a +required keyword there, so whoever wires them has to name one. + +**Why the local runs never showed this.** The local suite reported the same three +failures every time and I read the count rather than the names, then carried +"3 pre-existing" forward from record to record without re-deriving it. The check +that settles it — run the failing test against `origin/main` — takes one command +and I did not run it until CI forced the question. A failure inherited from the +base branch and a failure you introduced look identical in a summary line. + +## Correction — the second CI failure was also mine + +`unit-tests` ("Generator unit tests") installs `torch` and `transformers`. Its +original steps run the generator tests as plain scripts, so it never needed +pytest. This branch adds sixteen steps to that job that invoke +`python3 -m pytest`, and every one of them failed with `No module named pytest`. + +The job now installs pytest. Nothing was wrong with the tests themselves — they +pass locally and in the `moe-stage-guard` job, which installs pytest explicitly. +Adding a step to a job without checking what that job has installed is the whole +of the mistake. + +## Correction — the parent repo's CI never ran at all, for two more reasons + +The compiler PR's checks were the visible failure. The parent's were worse: all +three jobs failed, twice at `actions/checkout` and once at formatting, so the +emulator was never built or tested in CI on this branch even once. + +### `actions/checkout` died on the `PLENA_Tools` submodule + +`08db1bb9` is on no branch of `AICrossSim/PLENA_Tools`, so +`git submodule update --init --recursive` cannot fetch it and the job ends +before any step runs. This branch has been carrying that gitlink since +`6ec1e259`, and it has been noted in the PR description as "the submodule will +show as broken here" — which undersold it. It does not show as broken. **It +makes the branch impossible to check out**, for CI and for any reviewer. + +The bump is now dropped: the gitlink is back to `6b31fe00`, which is what +`origin/main` points at. That is safe because nothing here needs the newer +commit, and it was checked rather than assumed — all six KDA stage cases +(`cumprod`, `ut`, `state_transpose`, `prefill_out`, `prefill_state`, `layer`) +pass against `6b31fe00`. The only `check_mem.py` change between the two is a +NaN/Inf hard-failure rule, and this branch's stage tests carry a stricter bar of +their own. The `physical_rows` / `col_block_stride` support the multi-block +comparison depends on is already present at `6b31fe00`. + +If that Tools commit is wanted, the fix is to push it to +`AICrossSim/PLENA_Tools` and bump again — not to carry an unreachable pin. + +### `cargo fmt --all -- --check` failed on three hunks of this branch's Rust + +Two in `op.rs`, one in `vector_machine.rs`'s `fma_scalar`. The Nix dev shell +ships `rustfmt` but not the `cargo fmt` subcommand, so `cargo fmt` reports +"not installed for the toolchain" and I took that to mean formatting could not +be checked locally. It could: `rustfmt --edition 2024 --check ` runs fine +in the same shell and reproduces exactly what CI reports. The files are +formatted now. + +`cargo clippy` is genuinely unavailable in that shell, and CI's clippy step never +ran because the formatting step failed first. So clippy is still unverified +here. + +## Correction — the simulator CI was not "stuck", the PR had conflicts + +I reported that the parent repo's Actions was backed up, on the evidence that no +workflow runs were being created for the branch while two `main` runs sat queued +for hours. That was the wrong conclusion from a true observation. + +`gh api .../pulls/115 --jq .mergeable_state` says **`dirty`**. A `pull_request` +workflow runs against `refs/pull//merge`, and GitHub cannot compute that ref +when the merge conflicts — so no run is created at all, which looks exactly like +a stalled queue from the outside. One field settles it and I did not check it. + +`main` had moved one commit ahead: #114, "pipelined `--timing-model` scoreboard +with real DMA overlap". It conflicts with this branch in `dispatch.rs` and +`stage_profile.rs`. + +### The merge was not mechanical + +#114 **deleted** `classify_timing_access` from `dispatch.rs` and replaced it with +an `OpAccess` model in a new `accelerator/access.rs`. Of this branch's seven +hunks in `dispatch.rs` — 139 lines, all insertions — four still applied (two +execution arms, two `resource_kind_for_opcode` arms) and three had to be +re-derived against the new model. + +The three opcodes are classified in `access.rs` now: + +* `V_FMA_VF` gets its own arm rather than joining the `V_*_VF` family, because it + reads its **destination** as well as its source. It lists + `vector(gp(rd), vector_tile)` in `reads`; grouping it with the others + under-reports its vector-SRAM traffic by a row, which is the number the FMA + conversion is judged on. +* `V_SOFTPLUS_V` joins `V_EXP_V` / `V_RECI_V` unchanged. +* `S_MAP_FP_V` is `Unit::Vector`, not `Scalar`: it holds the vector SRAM read + port for a whole row even though its destination is FP_MEM. Its `S_MAP_V_FP` + mirror is the inverse and stays `Scalar`. + +The exhaustive match is what forced this: adding an opcode to `op.rs` fails the +build until `access.rs` classifies it. That guard is the reason a re-derivation +happened at all rather than a silent omission. + +`main`'s dispatch tests moved to `pipeline_tests.rs`, so this branch's +`mod tests` in `dispatch.rs` was restored with the two guards that still mean +something after the merge. + +### And clippy finally ran + +`flake.nix` never listed `clippy` or `rustfmt` in the toolchain extensions, so +the dev shell had neither, while CI gates on both. `cargo fmt`'s error — +"not installed for the toolchain, run `rustup component add rustfmt`" — is +cargo's generic advice; there is no rustup here and the components were simply +absent from the override. I read that message as "cannot be checked locally" +rather than as a one-line flake fix, which is how a formatting failure reached +CI. + +Both are in the extensions now. On the merged tree: `cargo fmt --all -- --check` +clean, `cargo clippy --workspace --all-targets -- -D warnings` clean — **the +first time clippy has run on this work at all** — and 129 Rust tests pass, up +from 112, the difference being #114's pipeline tests plus this branch's. diff --git a/flake.nix b/flake.nix index 79d981841..4109ff1a5 100644 --- a/flake.nix +++ b/flake.nix @@ -27,7 +27,15 @@ overlays = [ rust-overlay.overlays.default ]; }; rustToolchain = pkgs.rust-bin.stable.latest.default.override { - extensions = [ "rust-src" "rust-analyzer" ]; + # clippy and rustfmt are here because CI gates on them + # (`cargo fmt --all -- --check`, `cargo clippy -- -D warnings`) and the + # dev shell could not run either. `cargo fmt` reported "not installed for + # the toolchain, run `rustup component add rustfmt`", which is cargo's + # generic advice and misleading here -- there is no rustup in this shell, + # the toolchain comes from this override, and the components were simply + # not listed. Formatting failures reached CI that a local run would have + # caught in seconds. + extensions = [ "rust-src" "rust-analyzer" "clippy" "rustfmt" ]; }; rustPlatform = pkgs.makeRustPlatform { cargo = rustToolchain; diff --git a/justfile b/justfile index 8050b1094..3abcc36fe 100644 --- a/justfile +++ b/justfile @@ -80,6 +80,55 @@ build-perf-model model batch="4" input_seq="2048" output_seq="1024": --config "$(pwd)/plena_settings.toml" \ --isa-lib "$(pwd)/analytic_models/performance/customISA_lib.json" +# Run the Mamba-2 (selective SSM) performance model. +# just build-perf-model-mamba2 [batch] [input_seq] [output_seq] +build-perf-model-mamba2 model="mamba2-2.7b" batch="4" input_seq="2048" output_seq="1024": + python3 analytic_models/performance/mamba2_model.py \ + --model {{model}} \ + --batch-size {{batch}} \ + --input-seq {{input_seq}} \ + --output-seq {{output_seq}} \ + --model-lib "$(pwd)/PLENA_Compiler/doc/Model_Lib" \ + --config "$(pwd)/plena_settings.toml" \ + --isa-lib "$(pwd)/analytic_models/performance/customISA_lib.json" + +# Same, but against a model config outside Model_Lib. +# just build-perf-model-mamba2-path /path/to/mamba2-2.7b.json 1 2048 128 +build-perf-model-mamba2-path path batch="4" input_seq="2048" output_seq="1024": + python3 analytic_models/performance/mamba2_model.py \ + --model-path {{path}} \ + --batch-size {{batch}} \ + --input-seq {{input_seq}} \ + --output-seq {{output_seq}} \ + --config "$(pwd)/plena_settings.toml" \ + --isa-lib "$(pwd)/analytic_models/performance/customISA_lib.json" + +# TTFT/TPS sweep for one model across several context lengths, JSON per point. +# just latency-sweep llama-3.1-8b 1 128 "512 2048 8192" +latency-sweep model batch="1" output_seq="128" contexts="512 2048 8192": + #!/usr/bin/env bash + set -euo pipefail + script=analytic_models/performance/llama_model.py + case "{{model}}" in + mamba2*) script=analytic_models/performance/mamba2_model.py ;; + gpt-oss*) script=analytic_models/performance/gpt_oss_model.py ;; + esac + for ctx in {{contexts}}; do + python3 "$script" \ + --model {{model}} \ + --batch-size {{batch}} \ + --input-seq "$ctx" \ + --output-seq {{output_seq}} \ + --model-lib "$(pwd)/PLENA_Compiler/doc/Model_Lib" \ + --config "$(pwd)/plena_settings.toml" \ + --isa-lib "$(pwd)/analytic_models/performance/customISA_lib.json" \ + --json --quiet + done + +# Bandwidth-model + KV-store-bugfix regression tests for the analytic perf model. +test-perf-model: + python3 analytic_models/test_perf_model_bandwidth.py + # ==================== ATen-style Operator Tests ==================== # Ensure plena.ops and PLENA_Tools/ are importable @@ -201,6 +250,34 @@ test-moe-shared-all: just test-shared-moe-deepseek-fused just test-router-policy-all +# ==================== Mamba-2 / selective SSM ==================== + +# Per-stage numerical checks of the Mamba-2 lowering against a float32 torch +# golden. Fully synthetic -- no checkpoint, no HF libs. Cases: +# dt softplus + clamp (exercises V_SOFTPLUS_V) +# cumsum a @ lower-tri ones (the prefix-scan substitute; f32 accumulate) +# decay exp(min(cs_i-cs_j,0)) (exercises S_MAP_FP_V and V_SUB_VF rorder=1) +# conv1d causal depthwise k=4 +# NOTE: the emulator's V_EXP_V is libtorch's exact exp, not the RTL's fixed-point +# model, so passing here bounds the lowering and not the silicon. See the module +# docstring of mamba2_stage_test.py. +test-mamba2-stage case="dt" *args: + python3 transactional_emulator/testbench/mamba2/mamba2_stage_test.py --case {{case}} {{args}} + +# Every Mamba-2 stage. +test-mamba2-all: + #!/usr/bin/env bash + set -euo pipefail + for case in dt cumsum decay conv1d; do + echo "=== Mamba-2 stage: $case ===" + just test-mamba2-stage "$case" + done + +# The chunked-SSD reference vs the plain recurrence. Pure torch, no emulator: +# this is what makes the chunked form usable as an intermediate golden at all. +test-mamba2-reference: + python3 -m unittest compiler.aten.tests.test_mamba2_reference compiler.aten.tests.test_mamba_stage_contract -v + # Unified model compile/emulate (use model nickname from YAML configs) # Examples: # just aten-compile smollm2 --config sliced_64x64x16_b1 diff --git a/plena_settings.toml b/plena_settings.toml index 05fddad90..6af0d116e 100644 --- a/plena_settings.toml +++ b/plena_settings.toml @@ -41,6 +41,18 @@ value = 256 [ANALYTIC.CONFIG.VECTOR_SRAM_SIZE] value = 65536 +# Accelerator clock period, in picoseconds. 1000 ps is 1 GHz. +# +# THIS IS AN ASSUMPTION, NOT A MEASUREMENT. No RTL has been synthesised, so no +# critical path has set a frequency. It is here rather than hard-coded in +# `runtime_config.rs` so that every second, microsecond and TPOT figure the +# emulator reports can be traced to a value someone chose, and changed without a +# rebuild. `runner` prints it at startup and checks it against the DRAM model's +# own clock, which for the HBM2 preset is also 1 ns -- the two being equal is a +# coincidence of that preset, not a design relationship. +[ANALYTIC.CONFIG.CLOCK_PERIOD_PS] +value = 1000 + [ANALYTIC.CONFIG.HBM_M_Prefetch_Amount] value = 2048 @@ -178,6 +190,12 @@ dc_lib_dis = 4 dc_lib_en = 8 dc_lib_dis = 20 +# V_SOFTPLUS_V: relu + abs + exp + log1p + add in one pass. Budgeted above the bare +# VECTOR_EXP_CYCLES because the unit also evaluates a logarithm. +[ANALYTIC.LATENCY.VECTOR_SOFTPLUS_CYCLES] +dc_lib_en = 2 +dc_lib_dis = 13 + [ANALYTIC.LATENCY.SCALAR_FP_LONGEST_OPERATE_CYCLES] dc_lib_en = 4 dc_lib_dis = 4 @@ -244,6 +262,18 @@ value = 4096 [TRANSACTIONAL.CONFIG.VECTOR_SRAM_SIZE] value = 4194304 +# Accelerator clock period, in picoseconds. 1000 ps is 1 GHz. +# +# THIS IS AN ASSUMPTION, NOT A MEASUREMENT. No RTL has been synthesised, so no +# critical path has set a frequency. It is here rather than hard-coded in +# `runtime_config.rs` so that every second, microsecond and TPOT figure the +# emulator reports can be traced to a value someone chose, and changed without a +# rebuild. `runner` prints it at startup and checks it against the DRAM model's +# own clock, which for the HBM2 preset is also 1 ns -- the two being equal is a +# coincidence of that preset, not a design relationship. +[TRANSACTIONAL.CONFIG.CLOCK_PERIOD_PS] +value = 1000 + [TRANSACTIONAL.CONFIG.HBM_M_Prefetch_Amount] value = 64 @@ -381,6 +411,12 @@ dc_lib_dis = 4 dc_lib_en = 8 dc_lib_dis = 20 +# V_SOFTPLUS_V: relu + abs + exp + log1p + add in one pass. Budgeted above the bare +# VECTOR_EXP_CYCLES because the unit also evaluates a logarithm. +[TRANSACTIONAL.LATENCY.VECTOR_SOFTPLUS_CYCLES] +dc_lib_en = 2 +dc_lib_dis = 13 + [TRANSACTIONAL.LATENCY.SCALAR_FP_LONGEST_OPERATE_CYCLES] dc_lib_en = 4 dc_lib_dis = 4 diff --git a/transactional_emulator/src/accelerator/access.rs b/transactional_emulator/src/accelerator/access.rs index 4f6239b32..9ce3c1f6d 100644 --- a/transactional_emulator/src/accelerator/access.rs +++ b/transactional_emulator/src/accelerator/access.rs @@ -405,8 +405,31 @@ pub(crate) fn op_access( OpAccess::new(Unit::Vector, reads, vec![vector(gp(rd), vector_tile)]) } + // V_FMA_VF is the VF family plus one thing: it reads its **destination** + // as well as its source, because `V[rd] += V[rs1] * fp[rs2]`. Grouping it + // with the arm above would under-report its vector-SRAM traffic by a row, + // which is precisely the number the FMA conversion is judged on. + op::Opcode::V_FMA_VF { + rd, + rs1, + rs2, + rmask, + } => { + let mut reads = vec![ + Gp(rd), + Gp(rs1), + Fp(rs2), + vector(gp(rs1), vector_tile), + vector(gp(rd), vector_tile), + ]; + mask_read(rmask, &mut reads); + OpAccess::new(Unit::Vector, reads, vec![vector(gp(rd), vector_tile)]) + } + // === Vector ops: one vram source, vram dest === - op::Opcode::V_EXP_V { rd, rs1, rmask } | op::Opcode::V_RECI_V { rd, rs1, rmask } => { + op::Opcode::V_EXP_V { rd, rs1, rmask } + | op::Opcode::V_RECI_V { rd, rs1, rmask } + | op::Opcode::V_SOFTPLUS_V { rd, rs1, rmask } => { let mut reads = vec![Gp(rd), Gp(rs1), vector(gp(rs1), vector_tile)]; mask_read(rmask, &mut reads); OpAccess::new(Unit::Vector, reads, vec![vector(gp(rd), vector_tile)]) @@ -507,6 +530,16 @@ pub(crate) fn op_access( vec![vector(gp(rd), vector_tile)], ), + // The mirror of S_MAP_V_FP above, with every role inverted: `rs1` is the + // VRAM row it reads, `rd` the FP_MEM base it writes. It is billed to the + // Vector unit rather than Scalar because it holds the vector SRAM read + // port for a whole row. + op::Opcode::S_MAP_FP_V { rd, rs1, imm } => OpAccess::new( + Unit::Vector, + vec![Gp(rd), Gp(rs1), vector(gp(rs1), vector_tile)], + vec![scalar_fp(gp(rd).wrapping_add(imm), vector_tile)], + ), + op::Opcode::S_ADD_INT { rd, rs1, rs2 } | op::Opcode::S_SUB_INT { rd, rs1, rs2 } | op::Opcode::S_MUL_INT { rd, rs1, rs2 } => { diff --git a/transactional_emulator/src/accelerator/dispatch.rs b/transactional_emulator/src/accelerator/dispatch.rs index 02af2e8f6..80757ed81 100644 --- a/transactional_emulator/src/accelerator/dispatch.rs +++ b/transactional_emulator/src/accelerator/dispatch.rs @@ -136,7 +136,7 @@ impl Accelerator { // reproduces serial timing exactly. if !scoreboard.is_serialize() && self.issue_async_dma(op, &access, scoreboard).await { - let finish = scoreboard.commit(&access, issue, PERIOD); + let finish = scoreboard.commit(&access, issue, *PERIOD); scoreboard.trace_op(executed_pc, op, access.unit, issue, finish); if let Some(profiler) = stage_profiler.as_deref_mut() { let elapsed_picos = (finish - issue).as_picos(); @@ -364,6 +364,27 @@ impl Accelerator { ) .await; } + // The only V-type op that reads `rd`: `V[rd] += V[rs1] * fp[rs2]`. + // `rd` is the destination *and* an operand, so it goes first -- + // swapping the first two arguments computes `V[rs1] += V[rd]*f`, + // which is finite, plausible and wrong. + op::Opcode::V_FMA_VF { + rd, + rs1, + rs2, + rmask, + } => { + let mask = self.resolve_v_mask(*rmask); + self.v_machine + .fma_scalar( + self.reg_file.read_gp(*rd), + self.reg_file.read_gp(*rs1), + self.reg_file.read_fp(*rs2).into(), + *rmask, + mask, + ) + .await; + } op::Opcode::V_MAX_VF { rd, rs1, @@ -451,6 +472,17 @@ impl Accelerator { ) .await; } + op::Opcode::V_SOFTPLUS_V { rd, rs1, rmask } => { + let mask = self.resolve_v_mask(*rmask); + self.v_machine + .softplus( + self.reg_file.read_gp(*rd), + self.reg_file.read_gp(*rs1), + *rmask, + mask, + ) + .await; + } op::Opcode::V_RECI_V { rd, rs1, rmask } => { let mask = self.resolve_v_mask(*rmask); self.v_machine @@ -567,6 +599,19 @@ impl Accelerator { .await; cycle!(*VLEN); } + op::Opcode::S_MAP_FP_V { rd, rs1, imm } => { + // Mirror of S_MAP_V_FP: VRAM row -> VLEN consecutive FP_MEM slots. + // Note the operand roles are the mirror image too: `rs1` is the + // VRAM source row and `rd` is the FP_MEM base, so that both + // instructions keep "rd names the destination memory". + let values = self + .v_machine + .vector_read_fp(self.reg_file.read_gp(*rs1)) + .await; + let start_idx = (self.reg_file.read_gp(*rd) + *imm) as usize; + self.scalar_sram.write_fp_window(start_idx, &values); + cycle!(*VLEN); + } op::Opcode::S_ADD_INT { rd, rs1, rs2 } => { self.reg_file.binop_gp(*rd, *rs1, *rs2, u32::wrapping_add); cycle!(*SCALAR_INT_BASIC_CYCLES); @@ -765,7 +810,7 @@ impl Accelerator { // Inline HBM transfers advance the clock themselves // (`after - issue`); everything else charges through // `timing::charge_cycles`. - let latency = (after - issue) + PERIOD * charged; + let latency = (after - issue) + *PERIOD * charged; let finish = scoreboard.commit(&access, issue, latency); scoreboard.trace_op(executed_pc, op, access.unit, issue, finish); profiled_elapsed_picos = Some((finish - issue).as_picos()); @@ -991,6 +1036,12 @@ fn resource_kind_for_opcode(op: &op::Opcode) -> ResourceKind { | op::Opcode::V_RECI_V { .. } | op::Opcode::V_RED_SUM { .. } | op::Opcode::V_RED_MAX { .. } + | op::Opcode::V_FMA_VF { .. } + | op::Opcode::V_SOFTPLUS_V { .. } + // S_MAP_FP_V drives the vector SRAM read port for a whole VLEN row, so it + // contends with the vector unit even though its destination is FP_MEM. Its + // mirror S_MAP_V_FP is classified Scalar for the same reason inverted. + | op::Opcode::S_MAP_FP_V { .. } | op::Opcode::V_SHFT_V { .. } => ResourceKind::Vector, op::Opcode::S_ADD_FP { .. } @@ -1026,3 +1077,59 @@ fn resource_kind_for_opcode(op: &op::Opcode) -> ResourceKind { op::Opcode::Invalid => ResourceKind::Other, } } + +#[cfg(test)] +mod tests { + use super::*; + + /// The `V_FMA_VF` execution arm must pass `rd` as the destination. + /// + /// Nothing in this crate executes a *decoded* instruction -- every + /// `fma_scalar` test calls the method directly -- so swapping the first two + /// arguments in the dispatch arm, which turns the instruction into + /// `V[rs1] += V[rd] * f`, left the whole suite green. Every other opcode has + /// the same gap; closing it properly means building an Accelerator in a test, + /// which needs Ramulator and a MatrixMachine. + /// + /// Parsing our own source is crude. `stage_profile.rs` already does it for + /// the same reason: it is the only way to tie two things together without a + /// macro, and a crude guard beats none on an argument order that is a + /// silently different instruction when wrong. + #[test] + fn the_fma_dispatch_arm_passes_rd_as_the_destination() { + let source = include_str!("dispatch.rs"); + let arm = source + .find("op::Opcode::V_FMA_VF {") + .expect("the V_FMA_VF execution arm must exist"); + let body = &source[arm..arm + 600]; + let call = body + .find(".fma_scalar(") + .expect("the arm must call fma_scalar"); + let args = &body[call..]; + let rd = args.find("read_gp(*rd)").expect("rd must be passed"); + let rs1 = args.find("read_gp(*rs1)").expect("rs1 must be passed"); + assert!( + rd < rs1, + "fma_scalar takes (vd, vs1): rd must come first, or the instruction \ + computes V[rs1] += V[rd] * f instead" + ); + assert!( + args.find("read_fp(*rs2)").is_some_and(|f| f > rs1), + "rs2 is the FP scalar and must come third" + ); + } + + /// `V_FMA_VF` belongs to the vector unit, not the scalar one. + #[test] + fn fma_is_billed_to_the_vector_unit() { + assert_eq!( + resource_kind_for_opcode(&op::Opcode::V_FMA_VF { + rd: 1, + rs1: 2, + rs2: 3, + rmask: 0, + }), + ResourceKind::Vector + ); + } +} diff --git a/transactional_emulator/src/accelerator/scalar_sram.rs b/transactional_emulator/src/accelerator/scalar_sram.rs index 4c4e1ae5e..0cec333df 100644 --- a/transactional_emulator/src/accelerator/scalar_sram.rs +++ b/transactional_emulator/src/accelerator/scalar_sram.rs @@ -45,6 +45,21 @@ impl ScalarSram { &self.fpsram[start..start + len] } + /// Bulk counterpart of `read_fp_window`, used by `S_MAP_FP_V`. + /// + /// Panics on overrun rather than truncating: an out-of-range FP_MEM base is a + /// compiler bug, and silently dropping the tail would corrupt a Mamba chunk's + /// per-row decay scalars with no diagnostic. + pub(super) fn write_fp_window(&mut self, start: usize, values: &[bf16]) { + let end = start + values.len(); + assert!( + end <= self.fpsram.len(), + "S_MAP_FP_V would write FP_MEM[{start}..{end}) past the {}-entry file", + self.fpsram.len() + ); + self.fpsram[start..end].copy_from_slice(values); + } + pub(super) fn log_debug_contents(&self) { tracing::debug!("INT SRAM Contents: \n {:?}", self.intsram); tracing::debug!("FP SRAM Contents: \n {:?}", self.fpsram); diff --git a/transactional_emulator/src/accelerator/scoreboard.rs b/transactional_emulator/src/accelerator/scoreboard.rs index bd77e3867..363e4cc38 100644 --- a/transactional_emulator/src/accelerator/scoreboard.rs +++ b/transactional_emulator/src/accelerator/scoreboard.rs @@ -227,7 +227,7 @@ impl Scoreboard { prev_issue: None, prev_finish: Instant::INIT, serialize, - issue_cost: PERIOD, + issue_cost: *PERIOD, unit_busy_until: [Instant::INIT; Unit::COUNT], gp_ready: [Instant::INIT; 16], fp_ready: [Instant::INIT; 8], @@ -463,7 +463,7 @@ mod tests { } fn cycles(n: u32) -> Duration { - PERIOD * n + *PERIOD * n } fn compute(unit: Unit, reads: Vec, writes: Vec) -> OpAccess { diff --git a/transactional_emulator/src/load_config.rs b/transactional_emulator/src/load_config.rs index 5b73badbe..4f5862c19 100644 --- a/transactional_emulator/src/load_config.rs +++ b/transactional_emulator/src/load_config.rs @@ -21,6 +21,17 @@ pub struct LatencyValue { pub dc_lib_dis: u32, } +/// Fallback cost of `V_SOFTPLUS_V` when a settings file predates the Mamba ISA +/// extension. Budgeted as exp + reciprocal-class work (`relu`, `abs`, `exp`, +/// `log1p`, `add` in one pass) rather than aliased to the bare exp latency, which +/// would understate a unit that also has to evaluate a logarithm. +fn default_softplus_cycles() -> LatencyValue { + LatencyValue { + dc_lib_en: 2, + dc_lib_dis: 13, + } +} + #[derive(Debug, Serialize, Deserialize, Clone)] pub struct FpTypeConfig { pub sign: bool, @@ -108,6 +119,19 @@ pub struct ConfigSection { pub dc_en: ConfigValue, #[serde(rename = "MAX_LOOP_INSTRUCTIONS")] pub max_loop_instructions: ConfigValueUsize, + /// Accelerator clock period in picoseconds. + /// + /// Defaulted rather than required so that a TOML written before this field + /// existed still loads. The default is the value this was hard-coded to -- + /// 1000 ps, 1 GHz -- and it is an assumption, not a measurement: no RTL has + /// been synthesised, so no critical path has set it. Every microsecond, + /// TPOT or token-per-joule figure derived here is a figure "assuming 1 GHz". + #[serde(rename = "CLOCK_PERIOD_PS", default = "default_clock_period_ps")] + pub clock_period_ps: ConfigValueUsize, +} + +fn default_clock_period_ps() -> ConfigValueUsize { + ConfigValueUsize { value: 1000 } } #[derive(Debug, Serialize, Deserialize, Clone)] @@ -150,6 +174,11 @@ pub struct LatencySection { pub vector_max_cycles: LatencyValue, #[serde(rename = "VECTOR_SUM_CYCLES")] pub vector_sum_cycles: LatencyValue, + // V_SOFTPLUS_V. `default` so a settings file predating the Mamba ISA extension + // still parses instead of failing the whole config load; the default matches the + // fallback LatencySection below. + #[serde(rename = "VECTOR_SOFTPLUS_CYCLES", default = "default_softplus_cycles")] + pub vector_softplus_cycles: LatencyValue, #[serde(rename = "SCALAR_FP_LONGEST_OPERATE_CYCLES")] pub scalar_fp_longest_operate_cycles: LatencyValue, #[serde(rename = "SCALAR_FP_BASIC_CYCLES")] @@ -181,6 +210,7 @@ impl Default for AcceleratorConfig { hbm_v_writeback_amount: ConfigValue { value: 16 }, dc_en: ConfigValue { value: 1 }, max_loop_instructions: ConfigValueUsize { value: 10000 }, + clock_period_ps: default_clock_period_ps(), }, precision: PrecisionSection { matrix_sram_type: MxDataTypeConfig { @@ -316,6 +346,7 @@ impl Default for AcceleratorConfig { dc_lib_en: 8, dc_lib_dis: 20, }, + vector_softplus_cycles: default_softplus_cycles(), scalar_fp_longest_operate_cycles: LatencyValue { dc_lib_en: 4, dc_lib_dis: 4, @@ -523,6 +554,15 @@ pub fn blen() -> u32 { // } // Latency accessor functions (automatically uses DC_EN setting from config) +/// The accelerator clock period, in picoseconds. +/// +/// Read from the active TOML rather than hard-coded, so that the assumption is +/// visible in a file and can be changed without a rebuild. Clamped to at least +/// one picosecond: a zero period would make every derived cycle count infinite. +pub fn clock_period_ps() -> u64 { + CONFIG.config.clock_period_ps.value.max(1) as u64 +} + pub fn systolic_processing_overhead() -> u32 { get_dc_lib_value(&CONFIG.latency.systolic_processing_overhead) } @@ -555,6 +595,10 @@ pub fn vector_exp_cycles() -> u32 { get_dc_lib_value(&CONFIG.latency.vector_exp_cycles) } +pub fn vector_softplus_cycles() -> u32 { + get_dc_lib_value(&CONFIG.latency.vector_softplus_cycles) +} + pub fn vector_reci_cycles() -> u32 { get_dc_lib_value(&CONFIG.latency.vector_reci_cycles) } diff --git a/transactional_emulator/src/op.rs b/transactional_emulator/src/op.rs index e8c319938..470677283 100644 --- a/transactional_emulator/src/op.rs +++ b/transactional_emulator/src/op.rs @@ -108,6 +108,19 @@ pub enum Opcode { rs2: u8, rmask: u8, }, + /// `Vector[rd] += Vector[rs1] * fp_reg` -- fused broadcast multiply-add. + /// + /// Unlike every other V-type op, `rd` is **read** as well as written. That + /// is the whole point: it collapses the `copy + multiply + add` triple that + /// a rank-1 state update or a state contraction otherwise costs, and by + /// removing the scratch row it lets a whole key sweep become one arithmetic + /// row progression -- which the compiler turns into a hardware loop. + V_FMA_VF { + rd: u8, + rs1: u8, + rs2: u8, + rmask: u8, + }, V_MAX_VF { rd: u8, rs1: u8, @@ -161,6 +174,20 @@ pub enum Opcode { rs1: u8, rmask: u8, }, + /// `V_SOFTPLUS_V rd, rs1, rmask` — elementwise `log(1 + exp(x))` over one VLEN tile. + /// + /// Mamba/Mamba-2 needs `dt = softplus(dt_raw + dt_bias)` on the critical path of a + /// multiplicative recurrence, and the ISA has no logarithm, so there is no exact + /// software lowering. Evaluated as the range-safe identity + /// `softplus(x) = relu(x) + log1p(exp(-|x|))`, which is what the RTL is expected to + /// implement: it never evaluates `exp` on a positive argument, so it cannot overflow + /// for any finite input. `rmask` follows the usual masked-vector convention (0 = whole + /// tile, otherwise the per-head bitmask in `V_MASK`). + V_SOFTPLUS_V { + rd: u8, + rs1: u8, + rmask: u8, + }, S_ADD_FP { rd: u8, rs1: u8, @@ -208,6 +235,19 @@ pub enum Opcode { rs1: u8, imm: u32, }, + /// `S_MAP_FP_V rd, rs1, imm` — the exact inverse of `S_MAP_V_FP`. + /// + /// Copies one VLEN-wide Vector SRAM row at `gp[rs1]` into VLEN consecutive FP_MEM + /// slots starting at `gp[rd] + imm`. Before this existed the only VRAM-to-scalar path + /// was `V_RED_SUM`/`V_RED_MAX`, which collapse the whole row, so extracting a single + /// lane cost a one-hot `V_MUL_VV` + `V_RED_SUM` + `S_ST_FP` triple. Mamba-2's chunked + /// scan needs one broadcast scalar per row (`cs_i`, `exp(cs_C - cs_t)`, `exp(cs_i)`), + /// which made that the dominant instruction cost of the whole kernel. + S_MAP_FP_V { + rd: u8, + rs1: u8, + imm: u32, + }, S_ADD_INT { rd: u8, @@ -524,6 +564,20 @@ impl Opcode { // 0x35..=0x37 (V_MAX_VF/V_MIN_VF/V_TOPK) are decoded with the other // masked vector ops above. 0x38 => Self::C_SET_TOPK_REG { rd }, + // Mamba / selective-SSM extensions. Encodings must stay in sync with + // PLENA_Compiler's doc/operation.svh and assembler/assembly_to_binary.py. + 0x39 => Self::V_SOFTPLUS_V { + rd, + rs1, + rmask: rs3, + }, + 0x3A => Self::S_MAP_FP_V { rd, rs1, imm: imm2 }, + 0x3B => Self::V_FMA_VF { + rd, + rs1, + rs2, + rmask: rs3, + }, _ => { tracing::error!("Unknown opcode {opcode:#x}"); Self::Invalid @@ -787,6 +841,144 @@ mod tests { assert!(matches!(Opcode::decode(0x34), Opcode::C_BREAK)); } + // ---------- Mamba / selective-SSM extensions ---------- + + #[test] + fn test_decode_v_softplus_v_reads_rmask_from_rs3() { + // rs3 (not rs2) carries rmask, matching every other masked vector op and the + // compiler's `_RMASK_VECTOR_OPS` encoding. + match Opcode::decode(rform(0x39, 1, 2, 0, 5, 0)) { + Opcode::V_SOFTPLUS_V { rd, rs1, rmask } => assert_eq!((rd, rs1, rmask), (1, 2, 5)), + other => panic!("expected V_SOFTPLUS_V, got {other:?}"), + } + } + + #[test] + fn test_decode_s_map_fp_v_imm2() { + // Same operand shape as its S_MAP_V_FP mirror: rd, rs1 and the 18-bit imm2. + match Opcode::decode(i2form(0x3A, 4, 5, 0x1F00F)) { + Opcode::S_MAP_FP_V { rd, rs1, imm } => assert_eq!((rd, rs1, imm), (4, 5, 0x1F00F)), + other => panic!("expected S_MAP_FP_V, got {other:?}"), + } + } + + /// Every opcode PLENA_Compiler declares must decode to something here. + /// + /// `decode` carries three separate comments saying encodings "must stay in + /// sync with PLENA_Compiler's doc/operation.svh". That was enforced by + /// author diligence alone -- exactly the arrangement that let the compiler's + /// FPRAM depth (1024) drift from the SystemVerilog's (512) unnoticed. The + /// submodule is checked out recursively in every CI job, so the header can + /// simply be read. + #[test] + fn every_compiler_opcode_decodes_to_something() { + let svh = include_str!("../../PLENA_Compiler/doc/operation.svh"); + + // Declared in the header but deliberately not modelled. A new entry is + // a decision to argue for in review, not a default. + const NOT_MODELLED: &[&str] = &[ + // The sentinel. Invalid is precisely what it must decode to. + "INVALID_OPCODE", + // Declared by PLENA but never implemented, in RTL or here. The + // compiler knows and refuses to emit it -- see program_ssd.py's + // ssd_chunk_cumsum docstring: it "assembles silently and then + // decodes to Invalid, so emitting it is worse than not emitting + // it". Its PREFIX_SCAN_V_ELEMENT micro-op scans within a row, which + // under any head-on-lanes layout would cumsum across heads. + "V_PS_V", + // Likewise declared and unimplemented; nothing emits it. + "C_HADAMARD_TRANSFORM", + ]; + + let mut checked = 0; + for line in svh.lines() { + let line = line.trim(); + if line.starts_with("//") { + continue; + } + let Some((name, rest)) = line.split_once('=') else { + continue; + }; + let name = name.trim(); + let Some(hex) = rest.trim().strip_prefix("6'h") else { + continue; + }; + let hex: String = hex.chars().take_while(|c| c.is_ascii_hexdigit()).collect(); + let Ok(opcode) = u32::from_str_radix(&hex, 16) else { + continue; + }; + if NOT_MODELLED.contains(&name) { + continue; + } + checked += 1; + assert!( + !matches!( + Opcode::decode(rform(opcode, 0, 0, 0, 0, 0)), + Opcode::Invalid + ), + "{name} = 6'h{opcode:02X} is declared in PLENA_Compiler's \ + doc/operation.svh but decodes to Invalid here" + ); + } + assert!( + checked > 40, + "only {checked} opcodes parsed out of the header -- the parse broke, \ + so this guard was passing vacuously" + ); + } + + #[test] + fn v_fma_vf_decodes_like_v_mul_vf() { + // Same R-type slots as V_MUL_VF; only the opcode field differs. rd is + // read as well as written, which no other V-type op does, but that is + // an execution property -- the encoding is unchanged. + match Opcode::decode(rform(0x3B, 3, 4, 2, 0, 0)) { + Opcode::V_FMA_VF { + rd, + rs1, + rs2, + rmask, + } => { + assert_eq!((rd, rs1, rs2, rmask), (3, 4, 2, 0)); + } + other => panic!("expected V_FMA_VF, got {other:?}"), + } + match Opcode::decode(rform(0x3B, 1, 2, 3, 5, 0)) { + Opcode::V_FMA_VF { rmask, .. } => assert_eq!(rmask, 5, "rs3 is the mask"), + other => panic!("expected V_FMA_VF, got {other:?}"), + } + } + + #[test] + fn test_mamba_opcodes_do_not_collide_with_existing_encodings() { + // 0x39/0x3A/0x3B were free before this work (the previous maximum was + // C_SET_TOPK_REG = 0x38). Pin that they still decode to themselves and that + // the neighbouring slots are untouched. + assert!(matches!( + Opcode::decode(rform(0x38, 1, 0, 0, 0, 0)), + Opcode::C_SET_TOPK_REG { .. } + )); + assert!(matches!( + Opcode::decode(rform(0x39, 1, 2, 0, 0, 0)), + Opcode::V_SOFTPLUS_V { .. } + )); + assert!(matches!( + Opcode::decode(rform(0x3A, 1, 2, 0, 0, 0)), + Opcode::S_MAP_FP_V { .. } + )); + // 0x3B went to V_FMA_VF, the one opcode the KDA/Mamba work adds. The + // boundary moves with it: 0x3C is the next free slot, and this assertion + // is what makes a silent collision with it fail the build. + assert!(matches!( + Opcode::decode(rform(0x3B, 1, 2, 0, 0, 0)), + Opcode::V_FMA_VF { .. } + )); + assert!(matches!( + Opcode::decode(rform(0x3C, 0, 0, 0, 0, 0)), + Opcode::Invalid + )); + } + #[test] fn test_decode_v_shft_v() { match Opcode::decode(rform(0x32, 1, 2, 3, 0, 0)) { diff --git a/transactional_emulator/src/runner.rs b/transactional_emulator/src/runner.rs index 04a129c16..93ffd7046 100644 --- a/transactional_emulator/src/runner.rs +++ b/transactional_emulator/src/runner.rs @@ -145,8 +145,10 @@ pub(crate) async fn run_from_cli() { effective_hbm_size, effective_hbm_size as f64 / (1024.0 * 1024.0 * 1024.0) ); + let dram = ramulator::Ramulator::hbm2_preset(8).unwrap(); + assert_clock_relationship(&dram); let hbm = Arc::new(memory::WithStats::new(memory::WithTiming::new( - ManuallyDrop::new(ramulator::Ramulator::hbm2_preset(8).unwrap()), + ManuallyDrop::new(dram), memory::MemoryBacked::with_capacity(effective_hbm_size), ))); @@ -320,3 +322,39 @@ pub(crate) async fn run_from_cli() { utilization ); } + +/// State the accelerator clock and its relationship to the DRAM model's, at +/// startup, instead of letting the two coincide silently. +/// +/// `stage_profile.rs` already recorded the hazard: a cycle-domain comparison +/// there "only held because the DRAM tCK happened to equal PERIOD; any preset or +/// frequency change made it fail". Equal is what the HBM2 preset gives -- 2000 +/// MBPS is a 1 ns command clock, and `CLOCK_PERIOD_PS` defaults to 1000 -- but +/// that is a property of this preset and this default, not of the design. +/// +/// The relationship required is that the DRAM period is a whole multiple of the +/// accelerator period, in either direction. Anything else means a DRAM tick and +/// an accelerator cycle do not line up on any boundary, and every cycle count +/// derived by dividing one by the other is off by a fraction nothing reports. +fn assert_clock_relationship(dram: &ramulator::Ramulator) { + let accel_ps = crate::runtime_config::PERIOD.as_picos().max(1); + let dram_ps = dram.period().as_picos().max(1); + let (hi, lo) = if dram_ps >= accel_ps { + (dram_ps, accel_ps) + } else { + (accel_ps, dram_ps) + }; + assert!( + hi % lo == 0, + "accelerator clock is {accel_ps} ps ({:.3} GHz, from CLOCK_PERIOD_PS) and the \ + DRAM model's is {dram_ps} ps; neither divides the other, so ticks never line \ + up and any cycle count derived from both is off by a fraction nothing reports", + 1e3 / accel_ps as f64, + ); + println!( + "Clock: {accel_ps} ps ({:.3} GHz) from CLOCK_PERIOD_PS -- an assumption, not a \ + synthesised frequency. DRAM model: {dram_ps} ps ({}x).", + 1e3 / accel_ps as f64, + hi / lo, + ); +} diff --git a/transactional_emulator/src/runtime_config.rs b/transactional_emulator/src/runtime_config.rs index ff3450ef1..c753422a3 100644 --- a/transactional_emulator/src/runtime_config.rs +++ b/transactional_emulator/src/runtime_config.rs @@ -5,7 +5,17 @@ use runtime::Duration; use crate::load_config::*; -pub(crate) const PERIOD: Duration = Duration::from_nanos(1); +/// The accelerator clock period, from `[.CONFIG.CLOCK_PERIOD_PS]`. +/// +/// This was a `const` of 1 ns with no stated basis, and every second, +/// microsecond and TPOT figure the emulator reports is that number's +/// consequence. It is still 1 ns by default -- nothing has measured it, because +/// no RTL has been synthesised and no critical path has set a frequency -- but +/// it now lives in a configuration file where the assumption is visible and +/// changeable, and `runner` checks its relationship to the DRAM model's own +/// clock at startup instead of leaving the two to coincide. +pub(crate) static PERIOD: LazyLock = + LazyLock::new(|| Duration::from_picos(clock_period_ps())); pub(crate) static SYSTOLIC_PROCESSING_OVERHEAD: LazyLock = LazyLock::new(systolic_processing_overhead); @@ -18,6 +28,7 @@ pub(crate) static VECTOR_MAX_CYCLES: LazyLock = LazyLock::new(vector_max_cy // to the max latency rather than add a redundant knob to every config file. pub(crate) static VECTOR_MIN_CYCLES: LazyLock = LazyLock::new(vector_max_cycles); pub(crate) static VECTOR_SUM_CYCLES: LazyLock = LazyLock::new(vector_sum_cycles); +pub(crate) static VECTOR_SOFTPLUS_CYCLES: LazyLock = LazyLock::new(vector_softplus_cycles); pub(crate) static SCALAR_FP_BASIC_CYCLES: LazyLock = LazyLock::new(scalar_fp_basic_cycles); pub(crate) static SCALAR_FP_EXP_CYCLES: LazyLock = LazyLock::new(scalar_fp_exp_cycles); pub(crate) static SCALAR_FP_SQRT_CYCLES: LazyLock = LazyLock::new(scalar_fp_sqrt_cycles); diff --git a/transactional_emulator/src/stage_profile.rs b/transactional_emulator/src/stage_profile.rs index a7de26d44..26ae6e4be 100644 --- a/transactional_emulator/src/stage_profile.rs +++ b/transactional_emulator/src/stage_profile.rs @@ -439,6 +439,36 @@ pub(crate) enum StageKind { /// difference is worth keeping: a large `Other` is a coverage problem, a /// large `NonMoe` is just an epilogue. NonMoe, + // Mamba-2 / selective-SSM stages. Mirrors MAMBA_STAGES in the compiler's + // aten/plena/program_mamba_common.py, and checked against it live by + // stage_marker_names_match_the_compiler_vocabulary, which reads both + // vocabularies out of the Python source rather than mirroring them. + MambaInProj, + MambaConv1d, + MambaDt, + MambaChunkCumsum, + MambaDecayMask, + MambaIntraChunk, + MambaStateUpdate, + MambaInterChunk, + MambaSkip, + MambaGatedNorm, + MambaOutProj, + MambaStateLoad, + MambaStateStore, + // KDA (Kimi Delta Attention) stages. Mirrors KDA_STAGES in the compiler's + // aten/plena/program_kda_common.py, checked live by + // stage_marker_names_match_the_kda_compiler_vocabulary. + KdaQkvProj, + KdaConv1d, + KdaNormalize, + KdaDecay, + KdaStateUpdate, + KdaReadout, + KdaGatedNorm, + KdaOutProj, + KdaStateLoad, + KdaStateStore, Other, } @@ -459,6 +489,29 @@ impl StageKind { StageKind::SharedExpertActivation, StageKind::SharedExpertGate, StageKind::NonMoe, + StageKind::MambaInProj, + StageKind::MambaConv1d, + StageKind::MambaDt, + StageKind::MambaChunkCumsum, + StageKind::MambaDecayMask, + StageKind::MambaIntraChunk, + StageKind::MambaStateUpdate, + StageKind::MambaInterChunk, + StageKind::MambaSkip, + StageKind::MambaGatedNorm, + StageKind::MambaOutProj, + StageKind::MambaStateLoad, + StageKind::MambaStateStore, + StageKind::KdaQkvProj, + StageKind::KdaConv1d, + StageKind::KdaNormalize, + StageKind::KdaDecay, + StageKind::KdaStateUpdate, + StageKind::KdaReadout, + StageKind::KdaGatedNorm, + StageKind::KdaOutProj, + StageKind::KdaStateLoad, + StageKind::KdaStateStore, StageKind::Other, ]; @@ -479,29 +532,42 @@ impl StageKind { StageKind::SharedExpertActivation => "shared_expert_activation", StageKind::SharedExpertGate => "shared_expert_gate", StageKind::NonMoe => "non_moe", + StageKind::MambaInProj => "mamba_in_proj", + StageKind::MambaConv1d => "mamba_conv1d", + StageKind::MambaDt => "mamba_dt", + StageKind::MambaChunkCumsum => "mamba_chunk_cumsum", + StageKind::MambaDecayMask => "mamba_decay_mask", + StageKind::MambaIntraChunk => "mamba_intra_chunk", + StageKind::MambaStateUpdate => "mamba_state_update", + StageKind::MambaInterChunk => "mamba_inter_chunk", + StageKind::MambaSkip => "mamba_skip", + StageKind::MambaGatedNorm => "mamba_gated_norm", + StageKind::MambaOutProj => "mamba_out_proj", + StageKind::MambaStateLoad => "mamba_state_load", + StageKind::MambaStateStore => "mamba_state_store", + StageKind::KdaQkvProj => "kda_qkv_proj", + StageKind::KdaConv1d => "kda_conv1d", + StageKind::KdaNormalize => "kda_normalize", + StageKind::KdaDecay => "kda_decay", + StageKind::KdaStateUpdate => "kda_state_update", + StageKind::KdaReadout => "kda_readout", + StageKind::KdaGatedNorm => "kda_gated_norm", + StageKind::KdaOutProj => "kda_out_proj", + StageKind::KdaStateLoad => "kda_state_load", + StageKind::KdaStateStore => "kda_state_store", StageKind::Other => "other", } } + /// Position in `ALL`, derived rather than hand-written: a hand-maintained + /// match here is a second source of truth for the same ordering, and every + /// per-stage counter is indexed by it, so inserting a stage in the middle + /// would silently mis-attribute cycles. fn index(self) -> usize { - match self { - StageKind::ResidualSetup => 0, - StageKind::RouterTopk => 1, - StageKind::AccumulatorInit => 2, - StageKind::Gather => 3, - StageKind::ExpertWeightAddress => 4, - StageKind::ExpertWeightPrefetch => 5, - StageKind::ExpertProjection => 6, - StageKind::ExpertActivation => 7, - StageKind::ExpertBias => 8, - StageKind::ExpertRouteWeight => 9, - StageKind::ScatterCombine => 10, - StageKind::SharedExpertProjection => 11, - StageKind::SharedExpertActivation => 12, - StageKind::SharedExpertGate => 13, - StageKind::NonMoe => 14, - StageKind::Other => 15, - } + StageKind::ALL + .iter() + .position(|stage| *stage == self) + .expect("every StageKind variant must appear in ALL") } /// Resolve the name carried by a `; @stage=` marker. @@ -517,7 +583,7 @@ impl StageKind { } /// Number of `StageKind` variants, including `Other`. -const STAGE_COUNT: usize = 16; +const STAGE_COUNT: usize = 39; /// Comment prefix carrying an explicit stage attribution, emitted by the /// compiler's `moe_stage_marker`. Kept in sync by @@ -1799,6 +1865,10 @@ M_MM gp1, gp2, gp3 /// gitlink. Normally the two agree and this is the vocabulary at the pin, /// but a dirty or deliberately-checked-out submodule is read as it stands. const COMPILER_MOE_STAGES_PATH: &str = "../PLENA_Compiler/aten/plena/program_routed_moe.py"; + /// The Mamba-2 substrate declares its own vocabulary in its own module, so the + /// guard has to read both. Same extractor, different symbol. + const COMPILER_MAMBA_STAGES_PATH: &str = "../PLENA_Compiler/aten/plena/program_mamba_common.py"; + const COMPILER_KDA_STAGES_PATH: &str = "../PLENA_Compiler/aten/plena/program_kda_common.py"; /// Recover `MOE_STAGES` from Python source using Python's own parser. /// @@ -1814,15 +1884,18 @@ import json import sys +SYMBOL = sys.argv[1] if len(sys.argv) > 1 else SYMBOL + + def _declaration(tree): for node in tree.body: if isinstance(node, ast.AnnAssign): target = node.target - if isinstance(target, ast.Name) and target.id == "MOE_STAGES" and node.value is not None: + if isinstance(target, ast.Name) and target.id == SYMBOL and node.value is not None: return node.value elif isinstance(node, ast.Assign): for target in node.targets: - if isinstance(target, ast.Name) and target.id == "MOE_STAGES": + if isinstance(target, ast.Name) and target.id == SYMBOL: return node.value return None @@ -1841,7 +1914,7 @@ def _stages(node): raise ValueError("unsupported call in MOE_STAGES: " + ast.dump(node)) value = ast.literal_eval(node) if isinstance(value, (str, bytes)): - raise ValueError("MOE_STAGES is a scalar, not a collection of stage names") + raise ValueError(f"{SYMBOL} is a scalar, not a collection of stage names") stages = set() for item in value: if not isinstance(item, str): @@ -1914,11 +1987,21 @@ else: /// node. Callers that have already checked [`python3_available`] can treat /// any `Err` as a real parse failure rather than a toolchain problem. fn try_parse_moe_stages(source: &str) -> Result, String> { + parse_stage_set(source, "MOE_STAGES") + } + + /// Run [`MOE_STAGES_EXTRACTOR`] over `source`, recovering `symbol`. + /// + /// Generic over the symbol because the vocabulary lives in two modules: the + /// routed-MoE substrate declares `MOE_STAGES` and the Mamba-2 substrate + /// declares `MAMBA_STAGES`, and both feed one `StageKind`. + fn parse_stage_set(source: &str, symbol: &str) -> Result, String> { use std::io::Write as _; let mut child = std::process::Command::new("python3") .arg("-c") .arg(MOE_STAGES_EXTRACTOR) + .arg(symbol) .stdin(std::process::Stdio::piped()) .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()) @@ -2008,7 +2091,60 @@ else: let source = std::fs::read_to_string(&path) .unwrap_or_else(|err| panic!("cannot read {}: {err}", path.display())); - Some(parse_moe_stages(&source)) + let mut stages = parse_moe_stages(&source); + + // The Mamba-2 substrate carries a second vocabulary in a second module. + // Both feed one StageKind, so the guard has to union them -- otherwise + // every Mamba variant reads as "a stage no compiler marker produces". + let mamba_path = + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(COMPILER_MAMBA_STAGES_PATH); + if mamba_path.try_exists().unwrap_or(false) { + let mamba_source = std::fs::read_to_string(&mamba_path) + .unwrap_or_else(|err| panic!("cannot read {}: {err}", mamba_path.display())); + stages.extend( + parse_stage_set(&mamba_source, "MAMBA_STAGES").unwrap_or_else(|err| { + panic!("cannot recover MAMBA_STAGES from the source:\n{err}") + }), + ); + } else { + assert!( + prerequisites_may_be_missing(), + "the Mamba stage-vocabulary guard cannot run in CI: {} is not readable", + mamba_path.display() + ); + eprintln!( + "skipping the Mamba half of the stage-vocabulary guard: {} is not readable", + mamba_path.display() + ); + } + + // ...and a third for KDA. Same reason as Mamba: one StageKind is fed by + // several compiler modules, so the union has to include every one of + // them or the extra variants read as "a stage no compiler marker + // produces". Adding the StageKind variants without extending this is + // what broke the guard when KDA landed. + let kda_path = + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(COMPILER_KDA_STAGES_PATH); + if kda_path.try_exists().unwrap_or(false) { + let kda_source = std::fs::read_to_string(&kda_path) + .unwrap_or_else(|err| panic!("cannot read {}: {err}", kda_path.display())); + stages.extend( + parse_stage_set(&kda_source, "KDA_STAGES").unwrap_or_else(|err| { + panic!("cannot recover KDA_STAGES from the source:\n{err}") + }), + ); + } else { + assert!( + prerequisites_may_be_missing(), + "the KDA stage-vocabulary guard cannot run in CI: {} is not readable", + kda_path.display() + ); + eprintln!( + "skipping the KDA half of the stage-vocabulary guard: {} is not readable", + kda_path.display() + ); + } + Some(stages) } #[test] diff --git a/transactional_emulator/src/timing.rs b/transactional_emulator/src/timing.rs index df3e236f1..370ce103a 100644 --- a/transactional_emulator/src/timing.rs +++ b/transactional_emulator/src/timing.rs @@ -59,7 +59,7 @@ pub(crate) async fn charge_cycles(cycles: u32) { match timing_mode() { TimingMode::Serial => { runtime::Executor::current() - .resolve_at(PERIOD * cycles) + .resolve_at(*PERIOD * cycles) .await; } TimingMode::Scoreboard => { @@ -95,7 +95,7 @@ mod tests { charge_cycles(3).await; }); executor.enter(Instant::ETERNITY).await; - assert_eq!(executor.now(), Instant::INIT + PERIOD * 8u32); + assert_eq!(executor.now(), Instant::INIT + *PERIOD * 8u32); } #[tokio::test] diff --git a/transactional_emulator/src/vector_machine.rs b/transactional_emulator/src/vector_machine.rs index 13a63c66f..64402f1b1 100644 --- a/transactional_emulator/src/vector_machine.rs +++ b/transactional_emulator/src/vector_machine.rs @@ -17,7 +17,7 @@ use tch::Tensor; use crate::runtime_config::{ VECTOR_ADD_CYCLES, VECTOR_EXP_CYCLES, VECTOR_MAX_CYCLES, VECTOR_MIN_CYCLES, VECTOR_MUL_CYCLES, - VECTOR_RECI_CYCLES, VECTOR_SUM_CYCLES, VLEN, + VECTOR_RECI_CYCLES, VECTOR_SOFTPLUS_CYCLES, VECTOR_SUM_CYCLES, VLEN, }; use crate::{cycle, op}; @@ -136,6 +136,42 @@ impl VectorMachine { } } + /// `Vector[vd] += Vector[vs1] * f`. + /// + /// Mirrors `mul_scalar` but reads the destination too, so the accumulate is + /// part of the instruction rather than a separate `V_ADD_VV` over a scratch + /// row. The quantisation happens once, on the sum -- which is what makes it + /// a *fused* multiply-add and not just a shorter encoding of the pair. + pub(crate) async fn fma_scalar(&self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32) { + let a = self.vram.read(vs1).await; + let d = self.vram.read(vd).await; + if rmask == 0 { + let c = + QuantTensor::quantize(d.as_tensor() + a.as_tensor() * (f as f64), d.data_type()); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } else { + // Masked-off heads keep the destination's existing value, so the + // result starts as d and not as a -- the opposite of mul_scalar, + // where the source is the base. + let result = d.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let addend = a.as_tensor().narrow(0, start, end - start) * (f as f64); + let updated = &sliced + &addend; + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, d.data_type()); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } + } + pub(crate) async fn max_scalar(&self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; if rmask == 0 { @@ -308,6 +344,53 @@ impl VectorMachine { } } + /// Elementwise `softplus(x) = log(1 + exp(x))`. + /// + /// Evaluated as `relu(x) + log1p(exp(-|x|))`. That identity is algebraically + /// exact and never feeds `exp` a positive argument, so unlike the naive + /// `log1p(exp(x))` it cannot overflow — no input clamp is needed and none is + /// applied, which matters because Mamba's `dt` feeds `exp(A*dt)` and a clamp + /// would silently flatten the decay for large `dt`. + fn softplus_tensor(x: &Tensor) -> Tensor { + x.clamp_min(0.0) + x.abs().neg().exp().log1p() + } + + pub(crate) async fn softplus(&self, vd: u32, vs1: u32, rmask: u8, mask: u32) { + let a = self.vram.read(vs1).await; + if rmask == 0 { + let c = QuantTensor::quantize(Self::softplus_tensor(a.as_tensor()), a.data_type()); + cycle!(*VECTOR_SOFTPLUS_CYCLES); + self.vram.write(vd, c).await; + } else { + let result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = Self::softplus_tensor(&sliced); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_SOFTPLUS_CYCLES); + self.vram.write(vd, c).await; + } + } + + /// Read one whole VLEN-wide VRAM row out to the scalar domain. + /// + /// The inverse direction of [`Self::vector_transfer_fp`], and the engine behind + /// `S_MAP_FP_V`. Charged the same VLEN cycles as the forward transfer. + pub(crate) async fn vector_read_fp(&self, vs1: u32) -> Vec { + let a = self.vram.read(vs1).await; + let values = + Vec::::try_from(a.as_tensor()).expect("VRAM row must be convertible to Vec"); + cycle!(*VLEN); + values.into_iter().map(bf16::from_f32).collect() + } + pub(crate) async fn reciprocal(&self, vd: u32, vs1: u32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; if rmask == 0 { @@ -473,6 +556,207 @@ mod tests { data.to_vec() } + #[tokio::test] + async fn fma_scalar_accumulates_into_the_destination() { + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(4, 4, fp_type, 4)); + let machine = VectorMachine::new(vram.clone(), 4, 2); + let ty = MxDataType::Plain(fp_type); + + vram.write( + 0, + QuantTensor::quantize(Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0]), ty), + ) + .await; + vram.write( + 4, + QuantTensor::quantize(Tensor::from_slice(&[10.0f32, 20.0, 30.0, 40.0]), ty), + ) + .await; + + machine.fma_scalar(0, 4, 0.5, 0, u32::MAX).await; + + let out = vram.read(0).await; + let src = vram.read(4).await; + *got_task.lock().unwrap() = Some(( + tensor_values(out.as_tensor()), + tensor_values(src.as_tensor()), + )); + }); + executor.enter(Instant::ETERNITY).await; + + let (dst, src) = got.lock().unwrap().take().unwrap(); + // d + a*f, not a*f: the accumulate is the instruction, not a side effect. + assert_eq!(dst, vec![6.0, 12.0, 18.0, 24.0]); + assert_eq!(src, vec![10.0, 20.0, 30.0, 40.0], "the source is read-only"); + } + + #[tokio::test] + async fn fma_scalar_leaves_masked_off_heads_holding_the_destination() { + // mul_scalar's masked path starts from the *source*; fma's must start + // from the destination, or a masked-off lane silently loses whatever it + // had accumulated so far. + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(4, 4, fp_type, 4)); + let machine = VectorMachine::new(vram.clone(), 4, 2); + let ty = MxDataType::Plain(fp_type); + + vram.write( + 0, + QuantTensor::quantize(Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0]), ty), + ) + .await; + vram.write( + 4, + QuantTensor::quantize(Tensor::from_slice(&[10.0f32, 20.0, 30.0, 40.0]), ty), + ) + .await; + + // mask_unit 2, tile 4 -> two heads. Head 0 on, head 1 off. + machine.fma_scalar(0, 4, 0.5, 1, 0b01).await; + + let out = vram.read(0).await; + *got_task.lock().unwrap() = Some(tensor_values(out.as_tensor())); + }); + executor.enter(Instant::ETERNITY).await; + + let dst = got.lock().unwrap().take().unwrap(); + assert_eq!(dst, vec![6.0, 12.0, 3.0, 4.0]); + } + + #[tokio::test] + async fn test_softplus_matches_reference_including_the_large_magnitude_tails() { + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(4, 4, fp_type, 4)); + let machine = VectorMachine::new(vram.clone(), 4, 2); + let ty = MxDataType::Plain(fp_type); + + // -100 and +100 both overflow a naive log1p(exp(x)) in bf16: exp(100) is + // inf, and exp(-100) underflows. The relu + log1p(exp(-|x|)) form must + // return ~0 and ~100 respectively rather than NaN/inf. + let input = Tensor::from_slice(&[-100.0f32, -1.0, 0.0, 100.0]); + vram.write(0, QuantTensor::quantize(input, ty)).await; + + machine.softplus(4, 0, 0, 0).await; + + let out = vram.read(4).await; + *got_task.lock().unwrap() = Some(tensor_values(out.as_tensor())); + }); + executor.enter(Instant::ETERNITY).await; + + let out = got.lock().unwrap().take().unwrap(); + // softplus(-100) = 3.7e-44 -> 0 in bf16; softplus(0) = ln 2 = 0.6931; + // softplus(-1) = 0.3133; softplus(100) = 100 to well beyond bf16 precision. + assert!( + out[0].abs() < 1e-30, + "softplus(-100) should flush to ~0, got {}", + out[0] + ); + assert!( + (out[1] - 0.3133).abs() < 0.01, + "softplus(-1), got {}", + out[1] + ); + assert!( + (out[2] - std::f32::consts::LN_2).abs() < 0.01, + "softplus(0) is ln 2, got {}", + out[2] + ); + assert!( + (out[3] - 100.0).abs() < 1.0, + "softplus(100) should pass through, got {}", + out[3] + ); + assert!( + out.iter().all(|v| v.is_finite()), + "softplus must never produce inf/NaN: {out:?}" + ); + } + + #[tokio::test] + async fn test_softplus_honors_the_per_head_mask() { + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(4, 4, fp_type, 4)); + // tile_size 4, mask_unit 2 -> two heads of two lanes each. + let machine = VectorMachine::new(vram.clone(), 4, 2); + let ty = MxDataType::Plain(fp_type); + + let input = Tensor::from_slice(&[0.0f32, 0.0, 0.0, 0.0]); + vram.write(0, QuantTensor::quantize(input.copy(), ty)).await; + vram.write(4, QuantTensor::quantize(input, ty)).await; + + // mask = 0b01 -> only head 0 (lanes 0..2) is updated. + machine.softplus(4, 0, 1, 0b01).await; + + let out = vram.read(4).await; + *got_task.lock().unwrap() = Some(tensor_values(out.as_tensor())); + }); + executor.enter(Instant::ETERNITY).await; + + let out = got.lock().unwrap().take().unwrap(); + assert!( + (out[0] - std::f32::consts::LN_2).abs() < 0.01, + "head 0 lane 0 must be softplus(0) = ln 2" + ); + assert!( + (out[1] - std::f32::consts::LN_2).abs() < 0.01, + "head 0 lane 1 must be softplus(0) = ln 2" + ); + assert_eq!(out[2], 0.0, "head 1 must be left untouched by the mask"); + assert_eq!(out[3], 0.0, "head 1 must be left untouched by the mask"); + } + + #[tokio::test] + async fn test_vector_read_fp_round_trips_with_vector_transfer_fp() { + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(4, 4, fp_type, 4)); + let machine = VectorMachine::new(vram.clone(), 4, 2); + + // FP_MEM -> VRAM (S_MAP_V_FP) then VRAM -> FP_MEM (S_MAP_FP_V) must be + // the identity: this is the contract the Mamba decay-scalar path relies on. + let source: Vec = [1.5f32, -2.25, 0.0, 7.0] + .iter() + .map(|v| bf16::from_f32(*v)) + .collect(); + machine.vector_transfer_fp(0, &source).await; + let back = machine.vector_read_fp(0).await; + + *got_task.lock().unwrap() = Some((source, back)); + }); + executor.enter(Instant::ETERNITY).await; + + let (source, back) = got.lock().unwrap().take().unwrap(); + assert_eq!( + source, back, + "S_MAP_FP_V must invert S_MAP_V_FP exactly for bf16 values" + ); + } + #[tokio::test] async fn test_vector_scalar_minmax_clamps_bf16_boundary_values() { let executor = Executor::new(); diff --git a/transactional_emulator/testbench/aten/compare/isa_analysis.py b/transactional_emulator/testbench/aten/compare/isa_analysis.py index 4f19faac7..72da14234 100644 --- a/transactional_emulator/testbench/aten/compare/isa_analysis.py +++ b/transactional_emulator/testbench/aten/compare/isa_analysis.py @@ -90,7 +90,12 @@ def instruction_cycles(self, opcode: str) -> int: return 1 if opcode.startswith("V_ADD") or opcode.startswith("V_SUB"): return self.vector_add_cycles - if opcode.startswith("V_MUL") or opcode == "V_SHFT_V": + # V_FMA_VF is a multiply-add on the multiply datapath: dispatch.rs's + # fma_scalar charges VECTOR_MUL_CYCLES, the same as mul_scalar. It does + # not start with "V_MUL", so without naming it here it fell through to + # the catch-all `return 1` below -- a 5x under-count per FMA under the + # default dc_lib_dis profile, in the one tool Phase 4 measures with. + if opcode.startswith("V_MUL") or opcode in {"V_SHFT_V", "V_FMA_VF"}: return self.vector_mul_cycles if opcode == "V_EXP_V": return self.vector_exp_cycles @@ -236,6 +241,9 @@ def analyze_asm( "M_BMM_WO", "V_ADD_VV", "V_MUL_VF", + # The recurrences run on this one; leaving it out hides the whole + # state-update cost from the per-opcode breakdown. + "V_FMA_VF", "V_EXP_V", "V_RED_MAX", "V_RED_SUM", diff --git a/transactional_emulator/testbench/kda/__init__.py b/transactional_emulator/testbench/kda/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/transactional_emulator/testbench/kda/kda_stage_test.py b/transactional_emulator/testbench/kda/kda_stage_test.py new file mode 100644 index 000000000..cf1d8062e --- /dev/null +++ b/transactional_emulator/testbench/kda/kda_stage_test.py @@ -0,0 +1,949 @@ +# ruff: noqa: E402 +"""Numerical stage tests for the KDA lowering, on the transactional emulator. + +Every KDA test until now ran on `aten/tests/isa_interpreter.py`, which is +float64 and says so. That is the right oracle for "did the lowering emit the +arithmetic I meant", and the wrong one for "does it survive the hardware's +bf16". This file closes that gap: it compiles one stage, assembles it, runs the +Rust emulator, and compares against a float32 PyTorch golden. + + covered the emitters in aten/plena/program_kda_chunk.py end to end + through assembler + emulator, V_FMA_VF as the emulator really + executes it, and the bf16 rounding applied on every vector + write. + + NOT covered the RTL's fixed-point exp model -- same caveat as the Mamba + stage tests: the emulator uses libtorch's exact exp while + PLENA_Tools' quant model is a range reduction plus a 3-term + Taylor with systematic error. KDA applies exp on the critical + path of a multiplicative recurrence, where that compounds. + +Cases: + cumprod A_t = prod_{s<=t} a_s, per key channel -- the cumulative decay + ut T = (I + tril(diag(beta) M, -1))^-1 diag(beta) -- exercises + V_FMA_VF's accumulate inside a hardware loop, which is the whole + reason the opcode exists + +Run: python3 kda_stage_test.py --case cumprod [--mlen 64] [--blen 4] +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[3] +for _compiler_root in (_REPO_ROOT / "PLENA_Compiler", _REPO_ROOT.parent / "PLENA_Compiler"): + if (_compiler_root / "aten" / "plena" / "compiler.py").exists(): + sys.path.insert(0, str(_compiler_root)) + break + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena import PlenaCompiler +from compiler.aten.plena.program_kda_chunk import kda_chunk_cols, kda_chunk_rows +from compiler.aten.plena.program_kda_prefill import ( + KdaPrefillBuffers, + kda_prefill_state_transpose_shapes, + kda_prefill_tile_shapes, +) +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.emulator_runner import run_and_assert +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + +#: Slot order of the KDA FPRAM constant block. kda_fp_constants reuses +#: MambaFPConstants, so this mirrors FP_CONST_ORDER in the Mamba stage test. +#: Pinned here because FPRAM contents are never checked at runtime. +FP_CONST_ORDER = ["zero", "one", "neg_one", "dt_min", "dt_max", "reci_group", "eps"] + + +def use_bf16_kv_precision(build_dir: Path) -> None: + """Rewrite the per-build TOML's KV classes to Plain BF16. + + The prefill path spills intermediates to HBM and reads them back as MRAM, + because a matmul's second operand can only arrive that way. Those spills go + through the ``keyvalue`` precision class, and the shipped + ``[TRANSACTIONAL.PRECISION]`` makes that Mx/e4m3 with a separate scale + stream. ``SPILLED_ACTIVATION`` sets ``set_scale=False``, so the read walks + off into the scale stream, whose ``0x7f`` bytes decode to e4m3 NaN -- which + is what this test produced before the rewrite: every output NaN. + + ``ProgramSSDMixin.require_bf16_kv_precision`` documents exactly this and is + called by nothing; the Mamba prefill path has the same exposure. The state + is a multiplicative accumulator carried across chunks, so three mantissa + bits at every boundary compounds without bound even where it does not NaN. + """ + path = Path(os.environ["PLENA_SETTINGS_TOML"]) + text = path.read_text() + for key in ("HBM_M_KV_TYPE", "HBM_V_KV_TYPE"): + head = f"[TRANSACTIONAL.PRECISION.{key}]" + start = text.index(head) + end = text.index("\n[", start + len(head)) + text = ( + text[:start] + + ( + f'{head}\nformat = "Plain"\n' + f"[TRANSACTIONAL.PRECISION.{key}.DATA_TYPE]\n" + f'type = "Fp"\nsign = true\nexponent = 8\nmantissa = 7\n' + ) + + text[end + 1 :] + ) + path.write_text(text) + + +def _mxfp8_exact(shape_, generator, *, hi: float = 2.0) -> torch.Tensor: + """Values MX-FP8 e4m3 represents exactly, so the reported error is the + kernel's and not the input staging's. Same trick as the Mamba stage test.""" + steps = int(hi * 8) + raw = torch.randint(-steps, steps + 1, shape_, generator=generator) + return raw.to(torch.float32) / 8.0 + + +def _shape(args, *, key_dim: int) -> KdaShape: + return KdaShape( + hidden_size=args.num_heads * args.mlen, + num_heads=args.num_heads, + key_dim=key_dim, + value_dim=args.mlen, + conv_kernel=4, + ) + + +def _write_comparison( + build_dir: Path, prog, var, *, rows: int, cols: int, mlen: int, atol: float = 0.0, rtol: float = 0.0 +): + """Point the golden comparison at `rows x cols` of `var`, whatever its width. + + A VRAM matrix wider than `mlen` is stored column-block-major: block `c` + starts `physical_rows * mlen` elements in, not `rows * mlen`. `check_mem`'s + stride mode already knows this and takes the block pitch as + `col_block_stride`, but only if `physical_rows` is handed to it -- it + otherwise assumes `num_batches`, which is the logical row count and is + smaller than the physical one for every padded tile. Passing it is what + makes a multi-key-block state comparable at all; without it the second block + is read from `rows * mlen` and every value past the first block is somebody + else's. + + `num_rows` covers the whole tile rather than `rows * cols / mlen`, because + the reorder has to reach the last block's live rows and those sit past the + live rows of every block before it. + """ + addr = prog.get_vram_addr(var.name) + physical_rows = var.physical_shape[0] + col_blocks = max(1, -(-cols // mlen)) + params = { + "start_row_idx": addr // mlen, + "num_rows": physical_rows * col_blocks if cols > mlen else (rows * cols) // mlen, + "num_batches": rows, + "elements_per_batch": cols, + "physical_rows": physical_rows, + "row_dim": mlen, + "use_stride_mode": cols > mlen, + "atol": atol, + "rtol": rtol, + } + with open(build_dir / "comparison_params.json", "w") as f: + json.dump(params, f, indent=2) + return addr + + +def _finish(build_dir: Path, prog, golden, input_tensors, fp_preload, order, case, args): + gen_code = prog.compile() + print(f"\nGenerated {len(gen_code.splitlines())} lines of ISA") + create_sim_env( + input_tensors, + gen_code, + {"original_output": golden}, + fp_preload, + build_dir=str(build_dir), + ) + hbm_addrs = {name: prog.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + order = sorted(order, key=lambda name: hbm_addrs[name]) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm=f"kda_{case}", + data=None, + specified_data_order=order, + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + ) + with open(build_dir / "generated_asm_code.asm", "w") as f: + f.write(gen_code) + # The prefill path spills intermediates to HBM (MRAM is writable only by + # H_PREFETCH_M, so a matmul's second operand has to arrive that way), and + # those allocations sit past everything create_mem_for_sim staged. Without + # this the emulator indexes past its HBM vector and panics rather than + # reporting a size problem. `run_emulator` reads this file. + high_water = getattr(prog, "_next_hbm_addr", 0) + if high_water: + pad = 64 * 1024 + (build_dir / "hbm_size.txt").write_text(str(high_water + pad)) + run_and_assert(build_dir, f"kda_{case}", mlen=args.mlen, blen=args.blen) + _assert_every_value_within_tolerance(build_dir, case) + + +def _assert_every_value_within_tolerance(build_dir: Path, case: str): + """Fail unless *every* value is inside the tolerance the case asked for. + + `run_and_assert` goes through `check_mem`, whose `allclose_pass` is + `match_rate >= 90.0`. That is a shared utility with many callers and this is + not the place to change what it means -- but a tenth of the values being out + of tolerance is not a pass for a kernel, and the gap is not academic: the + output scale was applied to only the first value block, leaving half of every + token 11x too large, and the case still reported PASSED at a 94.68% match + rate because the data is O(1e-3) and the tolerance is 5e-2. + + So the threshold is restated here, on the same numbers, at 100%. + """ + from transactional_emulator.testbench.emulator_runner import compare_emulator_output + + results, params = compare_emulator_output(build_dir) + rate = results.get("allclose_match_rate") + if rate is None: + raise SystemExit(f"kda_{case}: no match rate to check") + if rate < 100.0: + raise SystemExit( + f"\n[kda_{case} FAILED] {rate:.2f}% of values inside " + f"atol={params.get('atol')} rtol={params.get('rtol')}; this check " + f"wants all of them. check_mem's own bar is 90%, which is why " + f"run_and_assert did not object." + ) + print(f"[kda_{case}] every value inside tolerance ({rate:.2f}%)") + + +# ============================================================================ +# Cases +# ============================================================================ + + +def case_cumprod(args, build_dir, hw): + """A_t = prod_{s<=t} a_s, per key channel. + + A running product and not exp of a running sum: the cumulative log-decay + reaches chunk*gate_lower_bound = -80, where bf16's ulp is 0.31, so + exponentiating a stored sum costs a 17% relative error on A. See + program_kda_chunk's module docstring for the measurement. + + Drawn from exact e4m3 values in [1/2, 1] so the staging quantisation is the + identity and the remaining error is the kernel's. Unlike the Mamba stage + cases this cannot be bit-exact: the *inputs* are exact, but a product of + sixteen 3-mantissa-bit values is not representable in bf16's 8, so each of + the fifteen multiplies rounds. Observed error is one bf16 ulp (1.95e-3 + absolute, 0.6% relative), which is the floor, not a defect. The tolerance is + set a few ulp above it -- a real error in the scan is orders of magnitude + larger, as the mutation tests in test_kda_prefill.py show. + """ + mlen = args.mlen + chunk = args.chunk + key_dim = args.key_dim or mlen + shape = _shape(args, key_dim=key_dim) + rows = kda_chunk_rows(shape, mlen, chunk) + cols = kda_chunk_cols(shape, mlen) + + g = torch.Generator().manual_seed(args.seed) + # Decays are in (0, 1]; keep them exactly representable and not so small + # that a 16-fold product underflows out of the comparison's range. + a = torch.randint(4, 9, (chunk, key_dim), generator=g).to(torch.float32) / 8.0 + golden = torch.cumprod(a, dim=0) + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + prog.kda_fp_constants() + fp_preload = prog.kda_fp_constant_values() + [0.0] * 3 + + # One row per timestep, key on lanes. At key_dim > mlen the staged tensor is + # wider than a VRAM row, and load_batch lays it out column-block-major -- + # which is exactly how the scan reads it. + a_in = prog.input("A", shape=(rows, cols)) + decay = prog.load_batch(a_in, name="decay") + prev = prog.alloc("prev", rows, cols, strict=False) + + prog.kda_chunk_decay_cumprod_v0(decay=decay, prev=prev, chunk=chunk, shape=shape) + + # Rows past `chunk` are staged as 1.0, not 0: the scan does not touch them, + # but a cumulative product that leaked into them would be visible. + staged = torch.nn.functional.pad(a, (0, 0, 0, rows - chunk), value=1.0) + _write_comparison(build_dir, prog, decay, rows=chunk, cols=cols, mlen=mlen, atol=0.0, rtol=2e-2) + _finish(build_dir, prog, golden, {"A": staged}, fp_preload, ["A"], "cumprod", args) + + +def case_ut(args, build_dir, hw): + """T = (I + tril(diag(beta) M, -1))^-1 diag(beta), by forward substitution. + + The stage that exercises V_FMA_VF's accumulate inside a hardware loop -- + T[i] = beta_i (e_i - sum_{j mlen: + raise SystemExit(f"chunk {chunk} must not exceed mlen {mlen}") + + g = torch.Generator().manual_seed(args.seed) + m_dense = torch.tril(_mxfp8_exact((chunk, chunk), g, hi=0.5), -1) + beta = _mxfp8_exact((chunk,), g, hi=1.0).abs().clamp(0.125, 1.0) + golden = torch.linalg.inv(torch.eye(chunk) + torch.diag(beta) @ m_dense) @ torch.diag(beta) + golden = torch.nn.functional.pad(golden, (0, mlen - chunk)) + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + consts = prog.kda_fp_constants() + fp_preload = prog.kda_fp_constant_values() + [0.0] * 3 + + m_in = prog.input("M", shape=(chunk, mlen)) + i_in = prog.input("I", shape=(chunk, mlen)) + m_v = prog.load_batch(m_in, name="m") + i_v = prog.load_batch(i_in, name="ident") + t_out = prog.alloc("t_out", chunk, mlen, strict=False) + m_fp = prog.fp_var("m_fp", size=mlen) + beta_fp = prog.fp_var("beta_fp", size=mlen) + + prog.kda_ut_transform_v0( + m=m_v, + identity=i_v, + beta_fp=beta_fp, + t_out=t_out, + m_fp=m_fp, + consts=consts, + chunk=chunk, + shape=shape, + ) + + # beta lands in FPRAM after the constants; the preload is positional. + fp_preload = fp_preload + [0.0] * (beta_fp.address - len(fp_preload)) + beta.tolist() + + m_padded = torch.nn.functional.pad(m_dense, (0, mlen - chunk)) + identity = torch.nn.functional.pad(torch.eye(chunk), (0, mlen - chunk)) + _write_comparison(build_dir, prog, t_out, rows=chunk, cols=mlen, mlen=mlen, atol=2e-2, rtol=2e-2) + _finish( + build_dir, + prog, + golden, + {"M": m_padded, "I": identity}, + fp_preload, + ["M", "I"], + "ut", + args, + ) + + +def _sequential_chunk(q, k, v, beta, per_step_decay, state0, scale): + """The sequential gated delta rule, from reference.py::kda_step, one head. + + The oracle for prefill is not a chunked reference -- there is none; + `kda_state_engine_prefill` is a loop of steps. It is *this*: chunked prefill + must equal the recurrence it collapses. + """ + s = state0.clone() + outs = [] + for t in range(q.shape[0]): + d = s * per_step_decay[t] + e = beta[t] * (v[t] - d @ k[t]) + s = d + torch.outer(e, k[t]) + outs.append(scale * (s @ q[t])) + return torch.stack(outs), s + + +def _prefill_case(args, build_dir, hw, *, want: str, chunks: int = 1): + """``chunks`` chunks of prefill for one head, against the sequential recurrence. + + More than one matters for two reasons the single-chunk case cannot show. + The spill regions are **reused** across chunks, so a second chunk's + prefetch reads whatever the first left behind -- correct only because each + spill zero-fills the rows past its live data before storing. And the carried + state is the thing that could compound; it does not, but only running it + proves that. + """ + mlen, chunk = args.mlen, args.chunk + key_dim = args.key_dim or mlen + value_dim = args.value_dim or mlen + shape = KdaShape( + hidden_size=value_dim, + num_heads=1, + key_dim=key_dim, + value_dim=value_dim, + conv_kernel=4, + ) + use_bf16_kv_precision(build_dir) + total = chunk * chunks + if chunk > mlen: + raise SystemExit(f"chunk {chunk} must not exceed mlen {mlen}") + + g = torch.Generator().manual_seed(args.seed) + # L2-normalised, because that is what reaches the recurrence: reference.py's + # kda_step normalises q and k with rsqrt(sum + 1e-6), and the lowering does + # the same in kda_mixer_step_v0. Feeding raw projections instead made |k| ~ 5 + # and drove cond(L) to the hundreds -- an artefact of the test, not a + # property of the workload. Normalised, cond(L) is 1.0 at every chunk size. + _norm = lambda t: t * torch.rsqrt(t.square().sum(-1, keepdim=True) + 1e-6) # noqa: E731 + q = _norm(_mxfp8_exact((total, key_dim), g, hi=1.0)) + k = _norm(_mxfp8_exact((total, key_dim), g, hi=1.0)) + v = _mxfp8_exact((total, value_dim), g, hi=1.0) + beta = _mxfp8_exact((total,), g, hi=1.0).abs().clamp(0.125, 1.0) + # Per-step decays. exp(lower_bound * sigmoid(.)) with lower_bound = -5 puts + # these across (e^-5, 1); the real distribution centres near exp(-2.5), which + # makes the recurrence strongly contracting -- the reason errors do not + # compound across chunks. + a = torch.exp(shape.gate_lower_bound * torch.sigmoid(torch.randn((total, key_dim), generator=g))) + state0 = _mxfp8_exact((value_dim, key_dim), g, hi=0.5) + scale = 1.0 / key_dim**0.5 + out_ref, state_ref = _sequential_chunk(q, k, v, beta, a, state0, scale) + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + consts = prog.kda_fp_constants() + fp_preload = prog.kda_fp_constant_values() + [0.0] * 3 + + # Rows past `chunk` are POISON, not zero. Every spilled tile is prefetched + # as a whole mlen x mlen MRAM block, so the rows past its live data are read + # by the matmul whether or not anything wrote them -- and + # kda_prefill_state_tail_v0's projection contracts over exactly those rows, + # straight into the carried state. kda_prefill_spill_v0 zeroes them before + # each store; with zeros here instead, deleting that zero-fill left every + # case passing, because a fresh VRAM tile is already zero and nothing + # dirties the tail. Poison is what makes the fill observable. + POISON = 9.0 + + # Every tile at the shape the emitter demands, from the one table -- both + # the staged inputs below and the allocations further down. The old form + # allocated mlen x mlen for all of them, which is the right shape only when + # key_dim == value_dim == mlen, and the emitter's own idea of the shapes was + # written separately so nothing compared the two. + tile_shapes = kda_prefill_tile_shapes(shape, mlen, chunk) + + def _pad(t, *, poison=True, rows=None): + rows = mlen if rows is None else rows + have = t.shape[0] + if have >= rows: + return t + tail = torch.full((rows - have, t.shape[1]), POISON if poison else 0.0) + return torch.cat([t, tail], dim=0) + + def _chunkwise(t, c): + """Slice chunk `c` out of a [total, *] tensor, padded to the tile's rows.""" + return _pad(t[c * chunk : (c + 1) * chunk], rows=tile_shapes["q"][0]) + + ins = { + # Padded to mlen rows: q and k are spilled (k_hat) and projected, and a + # projection writes a whole mlen x mlen block. One staged tensor per + # chunk, because each chunk consumes its own tokens. + **{f"Q{c}": _chunkwise(q, c) for c in range(chunks)}, + **{f"K{c}": _chunkwise(k, c) for c in range(chunks)}, + # Time on lanes, zero-padded to mlen: every contraction here runs over a + # whole block, and chunk <= mlen, so the padding contributes nothing and + # keeps one uniform width across the [*, chunk] tiles. + # v_t is [value_dim, chunk]: its *lanes* past chunk are padding, and the + # projections contract over lanes, so those must stay zero. + **{ + f"VT{c}": torch.nn.functional.pad( + v[c * chunk : (c + 1) * chunk].T.contiguous(), + (0, tile_shapes["v_t"][1] - chunk, 0, tile_shapes["v_t"][0] - value_dim), + ) + for c in range(chunks) + }, + **{f"A{c}": _chunkwise(a, c) for c in range(chunks)}, + "S0": state0, + # t_mat is spilled, and it is built from `identity` -- so identity's tail + # rows become t_mat's tail rows. Poison them too. + "ID": _pad(torch.tensor(KdaPrefillBuffers.identity_values(chunk, mlen))), + # t_mat is loaded, not allocated, so its tail rows start POISONED. That + # is what makes the spill zero-fill observable: t_mat is spilled, and its + # tail feeds err_t's lanes past `chunk`, which are the contraction index + # of the state projection. With t_mat freshly allocated its tail is zero + # and multiplies every other tail away -- one accidental zero masking the + # whole class of over-read. + "TM": _pad(torch.zeros(chunk, mlen)), + "CI": _pad(torch.tensor(KdaPrefillBuffers.causal_mask_values(chunk, mlen)), poison=False), + } + v_ = {name: prog.input(name, shape=tuple(t.shape)) for name, t in ins.items()} + state_v = prog.load_batch(v_["S0"], name="state") + ident_v = prog.load_batch(v_["ID"], name="ident") + ci_v = prog.load_batch(v_["CI"], name="causal_inclusive") + + alloc = lambda n: prog.alloc(n, *tile_shapes[n], strict=False) # noqa: E731 + # q / k / v_t / decay are set per chunk in the loop below -- all four are + # consumed in place, so each chunk gets a fresh load of its own tokens. + _ph = prog.alloc("placeholder", mlen, mlen, strict=False) + buffers = KdaPrefillBuffers( + q=_ph, + k=_ph, + v_t=_ph, + decay=_ph, + k_tilde=alloc("k_tilde"), + k_end=alloc("k_end"), + gram=alloc("gram"), + readout=alloc("readout"), + t_mat=prog.load_batch(v_["TM"], name="t_mat"), + identity=ident_v, + causal_inclusive=ci_v, + err_t=alloc("err_t"), + state=state_v, + contrib=alloc("contrib"), + readout_contrib=alloc("readout_contrib"), + state_contrib=alloc("state_contrib"), + out=alloc("out"), + prev=alloc("prev"), + scale_scratch=alloc("scale_scratch"), + beta_fp=prog.fp_var("beta_fp", size=mlen), + m_fp=prog.fp_var("m_fp", size=mlen), + output_scale_fp=prog.fp_var("out_scale", size=1), + consts=consts, + ) + outs, beta_vars = [], [] + for c in range(chunks): + # Fresh q/k/v/decay per chunk: all four are consumed in place. + buffers.q = prog.load_batch(v_[f"Q{c}"], name=f"q{c}") + buffers.k = prog.load_batch(v_[f"K{c}"], name=f"k{c}") + buffers.v_t = prog.load_batch(v_[f"VT{c}"], name=f"v_t{c}") + buffers.decay = prog.load_batch(v_[f"A{c}"], name=f"decay{c}") + buffers.out = prog.alloc(f"out{c}", *tile_shapes["out"], strict=False) + buffers.beta_fp = prog.fp_var(f"beta_fp{c}", size=mlen) + beta_vars.append(buffers.beta_fp) + prog.kda_prefill_chunk_v0(buffers=buffers, chunk=chunk, shape=shape) + outs.append(buffers.out) + + out_scale = buffers.output_scale_fp + tail = max([out_scale.address + 1] + [bv.address + mlen for bv in beta_vars]) + fp = fp_preload + [0.0] * max(0, tail - len(fp_preload)) + fp[out_scale.address] = scale + for c, bv in enumerate(beta_vars): + fp[bv.address : bv.address + chunk] = beta[c * chunk : (c + 1) * chunk].tolist() + + if want == "out": + # The last chunk's output, against the same span of the reference. + target, golden, rows = outs[-1], out_ref[(chunks - 1) * chunk :], chunk + cols = tile_shapes["out"][1] + golden = torch.nn.functional.pad(golden, (0, cols - value_dim)) + else: + target, golden, rows = buffers.state, state_ref, value_dim + cols = tile_shapes["state"][1] + golden = torch.nn.functional.pad(golden, (0, cols - key_dim)) + + # Absolute, not relative: with normalised q/k the outputs are O(1e-3), so a + # relative bar on values near zero measures nothing. 5e-2 absolute is ~10x + # the observed 5e-3, and a real lowering error is orders of magnitude larger + # -- the mutation tests in test_kda_prefill.py are what cover that. + _write_comparison(build_dir, prog, target, rows=rows, cols=cols, mlen=mlen, atol=5e-2, rtol=5e-2) + _finish(build_dir, prog, golden, ins, fp, list(ins), f"prefill_{want}", args) + + +def case_prefill_out(args, build_dir, hw): + """The chunk's token outputs must equal the recurrence's, token for token.""" + _prefill_case(args, build_dir, hw, want="out") + + +def case_prefill_state(args, build_dir, hw): + """And the carried state, which is what makes chunks chain.""" + _prefill_case(args, build_dir, hw, want="state") + + +def case_state_transpose(args, build_dir, hw): + """The prefill -> decode state layout conversion. + + decode holds the state [key, value] and prefill holds it [value, key]. At + Kimi K3 key_dim == value_dim, so the shapes match and handing one to the + other is a finite plausible wrong answer. The conversion is one projection + against a staged identity. + """ + mlen = args.mlen + key_dim = args.key_dim or mlen + value_dim = args.value_dim or mlen + shape = KdaShape( + hidden_size=mlen, + num_heads=1, + key_dim=key_dim, + value_dim=value_dim, + conv_kernel=4, + ) + use_bf16_kv_precision(build_dir) + g = torch.Generator().manual_seed(args.seed) + # Distinguishable under transposition *and* under a block swap: with + # key_dim == value_dim the tile is square, so a converter that got the block + # indices the wrong way round would return something of the right shape. + state = _mxfp8_exact((value_dim, key_dim), g, hi=1.0) + golden = state.T.contiguous() + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + prog._bf16_kv_checked = True + prog.kda_fp_constants() + fp_preload = prog.kda_fp_constant_values() + [0.0] * 3 + + want = kda_prefill_state_transpose_shapes(shape, mlen) + # The identity spans the whole key axis, not `chunk` and not `mlen`: the + # transpose contracts over key, so a smaller one contracts over part of it. + ins = {"S": state, "ID": torch.eye(want["identity"][0])} + v_ = {n: prog.input(n, shape=tuple(t.shape)) for n, t in ins.items()} + state_v = prog.load_batch(v_["S"], name="state") + ident_v = prog.load_batch(v_["ID"], name="ident") + out_v = prog.alloc("state_T", *want["out"], strict=False) + + from compiler.aten.plena.program_ssd import SPILLED_ACTIVATION + + prog.kda_prefill_state_to_decode_layout_v0( + state=state_v, + identity=ident_v, + out=out_v, + shape=shape, + precision=SPILLED_ACTIVATION, + ) + _write_comparison( + build_dir, + prog, + out_v, + rows=key_dim, + cols=want["out"][1], + mlen=mlen, + atol=1e-3, + rtol=1e-3, + ) + _finish(build_dir, prog, golden, ins, fp_preload, list(ins), "state_transpose", args) + + +def _layer_reference(projected, conv_hist, conv_w, state, a_log, dt_bias, shape): + """One layer of the state engine, from the reference.""" + from compiler.aten.models.kda.reference import ( + KdaConvWeights, + KdaRecurrentState, + kda_state_engine_step, + ) + + return kda_state_engine_step( + projected, + KdaRecurrentState(state.clone(), conv_hist.clone()), + KdaConvWeights(conv_w["q"], conv_w["k"], conv_w["v"], None, None, None), + a_log, + dt_bias, + shape, + state_storage="fp32", + ) + + +def case_layer(args, build_dir, hw, *, layers: int = 1): + """The assembled KDA layer on the emulator, against kda_state_engine_step. + + Gather, three convolutions, the gates and the recurrence -- every emitter + this work added, running as one program through the assembler and the Rust + emulator rather than through the float64 ISA interpreter. + + With ``layers > 1`` each layer gets its own projection, conv history, + weights and recurrent state, and they run back to back in one program. That + is what catches a layer whose emitters collide with the previous one's -- + reused spill names, an FPRAM window written twice, a VRAM view whose name + repoints an earlier one. A single layer cannot see any of it. + """ + import torch + + from compiler.aten.models.kda.shape import KdaShape + from compiler.aten.plena.program_kda_common import ( + kda_state_row, + kda_state_rows, + kda_vector_row, + kda_vector_rows, + ) + from compiler.aten.plena.program_kda_conv import kda_conv_blocks, kda_conv_state_row + from compiler.aten.plena.program_kda_gates import ( + kda_head_blocks, + kda_key_blocks, + kda_key_row, + ) + from compiler.aten.plena.program_kda_layer import ( + kda_projection_features, + kda_projection_sections, + ) + from compiler.aten.plena.program_kda_mixer import KdaMixerBuffers + + mlen = args.mlen + heads = args.num_heads + key_dim = args.key_dim or mlen + value_dim = mlen + shape = KdaShape( + hidden_size=heads * value_dim, + num_heads=heads, + key_dim=key_dim, + value_dim=value_dim, + conv_kernel=4, + ) + key_width = shape.projection_size + value_width = heads * value_dim + kernel = shape.conv_kernel + kb = kda_key_blocks(shape, mlen) + vb = kda_vector_rows(shape, mlen) // heads + + g = torch.Generator().manual_seed(args.seed) + shared_a_log = _mxfp8_exact((heads,), g, hi=0.5) + per_layer = [] + for _ in range(layers): + per_layer.append( + dict( + projected=_mxfp8_exact((1, kda_projection_features(shape)), g, hi=1.0), + conv_hist=_mxfp8_exact((1, 2 * key_width + value_width, kernel), g, hi=0.5), + conv_w={ + n: _mxfp8_exact((w, kernel), g, hi=0.5) + for n, w in (("q", key_width), ("k", key_width), ("v", value_width)) + }, + state=_mxfp8_exact((1, heads, value_dim, key_dim), g, hi=0.5), + a_log=shared_a_log, + dt_bias=_mxfp8_exact((heads, key_dim), g, hi=0.5), + ) + ) + golden_rows = [] + for L in per_layer: + out, _ = _layer_reference( + L["projected"], + L["conv_hist"], + L["conv_w"], + L["state"], + L["a_log"], + L["dt_bias"], + shape, + ) + golden_rows.append(out[0]) + golden = torch.stack(golden_rows) + golden = torch.nn.functional.pad(golden, (0, mlen * vb - golden.shape[1])) + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + consts = prog.kda_fp_constants() + fp = prog.kda_fp_constant_values() + [0.0] * 3 + + up = lambda n: ((n + mlen - 1) // mlen) * mlen # noqa: E731 + a = lambda n, r: prog.alloc(n, up(r), mlen, strict=False) # noqa: E731 + widths = {"q": key_width, "k": key_width, "v": value_width} + offsets = {"q": 0, "k": key_width, "v": 2 * key_width} + + # One FPRAM window for every layer. FPRAM is per-layer *scratch*, not + # per-layer storage: a program that allocates it per layer overflows even + # the compiler's optimistic 1024 slots at four layers, let alone the + # hardware's 512. That is the same reuse the mixer already does per head. + ins, outs = {}, [] + for i, L in enumerate(per_layer): + # Every layer stages its own tensors; nothing is shared but the + # constants, which is what makes a collision between layers visible. + flat = L["projected"][0] + nblk = -(-flat.numel() // mlen) + proj_rows = torch.zeros(prog.blen, nblk * mlen) + for b in range(nblk): + chunk = flat[b * mlen : (b + 1) * mlen] + proj_rows[0, b * mlen : b * mlen + chunk.numel()] = chunk + ins[f"PROJ{i}"] = proj_rows + for n, w in widths.items(): + blocks = kda_conv_blocks(w, mlen) + cs = torch.zeros(blocks * kernel, mlen) + cw = torch.zeros(blocks * kernel, mlen) + for cb in range(blocks): + lo = cb * mlen + for tap in range(kernel): + row = kda_conv_state_row(w, mlen, kernel, cb, tap) + cs[row] = L["conv_hist"][0, offsets[n] + lo : offsets[n] + lo + mlen, tap] + cw[row] = L["conv_w"][n][lo : lo + mlen, tap] + ins[f"CS{i}_{n}"], ins[f"CW{i}_{n}"] = cs, cw + dtb = torch.zeros(up(heads * kb), mlen) + for h in range(heads): + for b in range(kb): + dtb[kda_key_row(shape, mlen, h, b)] = L["dt_bias"][h, b * mlen : (b + 1) * mlen] + ins[f"DTB{i}"] = dtb + st = torch.zeros(up(kda_state_rows(shape, mlen)), mlen) + for h in range(heads): + for blk in range(vb): + for key in range(key_dim): + st[kda_state_row(shape, mlen, h, blk, key)] = L["state"][0, h, blk * mlen : (blk + 1) * mlen, key] + ins[f"ST{i}"] = st + + v_ = {n: prog.input(n, shape=tuple(t.shape)) for n, t in ins.items()} + result = a("layer_out", layers * vb) + shared_fp = { + "decay": prog.fp_var("decay_and_q_hat", size=key_dim), + "k_hat": prog.fp_var("k_hat", size=key_dim), + "beta": prog.fp_var("beta", size=kda_head_blocks(shape, mlen) * mlen), + "part": prog.fp_var("part", size=kb), + "acc": prog.fp_var("acc", size=1), + "output_scale": prog.fp_var("output_scale", size=1), + "rate": prog.fp_var("rate", size=heads), + "lower_bound": prog.fp_var("lower_bound", size=1), + } + + for i in range(layers): + projected = prog.load_batch(v_[f"PROJ{i}"], name=f"projected{i}") + sections = {n: (f, c) for n, f, c in kda_projection_sections(shape, mlen)} + gathered = {n: a(f"g{i}_{n}", max(c, mlen)) for n, (_, c) in sections.items()} + conv_state = {n: prog.load_batch(v_[f"CS{i}_{n}"], name=f"cs{i}_{n}") for n in widths} + conv_weight = {n: prog.load_batch(v_[f"CW{i}_{n}"], name=f"cw{i}_{n}") for n in widths} + conv_scratch = a(f"convsc{i}", kda_conv_blocks(max(widths.values()), mlen)) + decay = shared_fp["decay"] + buffers = KdaMixerBuffers( + q=a(f"q{i}", heads * kb), + k=a(f"k{i}", heads * kb), + v=a(f"v{i}", kda_vector_rows(shape, mlen)), + gate=gathered["gate"], + dt_bias=prog.load_batch(v_[f"DTB{i}"], name=f"dtb{i}"), + beta_logit=gathered["beta"], + state=prog.load_batch(v_[f"ST{i}"], name=f"st{i}"), + out=a(f"out{i}", kda_vector_rows(shape, mlen)), + pred=a(f"pred{i}", kda_vector_rows(shape, mlen)), + err=a(f"err{i}", kda_vector_rows(shape, mlen)), + sq_scratch=a(f"sq{i}", heads * kb), + decay_fp=decay, + q_hat_fp=decay, + k_hat_fp=shared_fp["k_hat"], + beta_fp=shared_fp["beta"], + part_fp=shared_fp["part"], + acc_fp=shared_fp["acc"], + output_scale_fp=shared_fp["output_scale"], + rate_fp=shared_fp["rate"], + lower_bound_fp=shared_fp["lower_bound"], + consts=consts, + ) + prog.kda_layer_from_projected_v0( + projected=projected, + gathered=gathered, + conv_state=conv_state, + conv_weight=conv_weight, + conv_bias={}, + conv_scratch=conv_scratch, + mixer_buffers=buffers, + shape=shape, + ) + outs.append(buffers.out) + # Pack every layer's output, for every head count. This was gated on + # `heads == 1`, so under the only invocation the commit recorded + # (`--num-heads 3`) it compared the LAST layer alone -- zeroing layers + # 0..n-2's inputs left the verdict bit-identical. + for h in range(heads): + for blk in range(vb): + prog.mamba_row_copy( + result, + (i * heads + h) * vb + blk, + buffers.out, + kda_vector_row(shape, mlen, h, blk), + ) + target, rows = result, layers * heads * vb + + want = torch.zeros(rows, mlen) + for i, row in enumerate(golden_rows): + for h in range(heads): + for blk in range(vb): + lo = h * value_dim + blk * mlen + want[(i * heads + h) * vb + blk] = row[lo : lo + mlen] + golden = want + + tail = max( + shared_fp["rate"].address + heads, shared_fp["lower_bound"].address + 1, shared_fp["output_scale"].address + 1 + ) + fp = fp + [0.0] * max(0, tail - len(fp)) + fp[shared_fp["output_scale"].address] = 1.0 / key_dim**0.5 + fp[shared_fp["lower_bound"].address] = shape.gate_lower_bound + # One shared `rate` window, so every layer uses the same a_log. A real model + # reloads it per layer from HBM; there is no emitter for that yet, and what + # this case tests is that layers compose as programs. + fp[shared_fp["rate"].address : shared_fp["rate"].address + heads] = torch.exp(per_layer[0]["a_log"]).tolist() + + # atol=0. An absolute bound is what made this case vacuous: the outputs are + # O(5e-3) and atol was 2e-2, so every golden value sat inside it and writing + # zeros into the tile scored 100%. Verified -- zeroing the output, skipping + # the `v` convolution, running the mixer before the convolutions, and + # swapping the `q` and `k` gather sections all passed. + # + # The bound is relative now and set from what bf16 delivers, not from what + # makes the test green: the read-out contracts 128 terms, so the error grows + # like sqrt(128) * 2^-9 = 2.2%, and the measured mean relative error is 4.3%. + # `atol` is not 0. A pure relative bound cannot be met by the smallest + # entries: the layer's outputs run from 1e-3 to 1e-2, bf16 delivers about + # 1.8e-4 absolute on this contraction whatever the value, and on an entry of + # 1e-3 that is 18% -- so three of the sixty-four always sat outside a 12% + # bar. They were carried by `check_mem`'s 90% rule rather than by the + # tolerance, which is not a bar at all. `layer_atol` states the absolute + # floor instead, and the combined bound is checked at 100%. + # + # It stays a real check: an all-zero output scores 1.6% against this bound, + # because every golden value is far above the floor. That is the property + # the earlier `atol=2e-2` version lacked -- there, zeros scored 100%. + _write_comparison( + build_dir, + prog, + target, + rows=rows, + cols=mlen, + mlen=mlen, + atol=args.layer_atol, + rtol=args.layer_rtol, + ) + _finish(build_dir, prog, golden, ins, fp, list(ins), f"layer{layers}", args) + + +def case_layer_chain(args, build_dir, hw): + """Four layers back to back in one program.""" + case_layer(args, build_dir, hw, layers=4) + + +def case_prefill_chain_state(args, build_dir, hw): + """Three chunks in sequence, against a 3*chunk-token recurrence. + + This is the case that proves the two memory fixes. The six spill regions are + reused across chunks, so chunk 2's prefetch reads bytes chunk 1 wrote -- and + each prefetch reads a whole mlen x mlen block, more than any chunk's live + data. It is correct only because every spill zero-fills the rows past its + live data first. With a single chunk that is invisible: unallocated HBM + happens to be zero. + + It also answers whether the carried state's error compounds. It does not -- + the decay makes the recurrence strongly contracting, so old error is damped. + """ + _prefill_case(args, build_dir, hw, want="state", chunks=3) + + +def case_prefill_chain_out(args, build_dir, hw): + """The third chunk's token outputs, which depend on two chunks of carried state.""" + _prefill_case(args, build_dir, hw, want="out", chunks=3) + + +CASES = { + "cumprod": case_cumprod, + "ut": case_ut, + "prefill_out": case_prefill_out, + "prefill_state": case_prefill_state, + "prefill_chain_out": case_prefill_chain_out, + "prefill_chain_state": case_prefill_chain_state, + "state_transpose": case_state_transpose, + "layer": case_layer, + "layer_chain": case_layer_chain, +} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--case", default="cumprod", choices=sorted(CASES)) + parser.add_argument("--num-heads", type=int, default=1) + parser.add_argument("--chunk", type=int, default=16) + parser.add_argument("--key-dim", type=int, default=None) + parser.add_argument("--value-dim", type=int, default=None) + # Relative and honest. bf16 on the read-out's 128-term contraction + # gives about sqrt(128)*2^-9 = 2.2%; measured mean is 4.3%. The + # harness also passes at >= 90% of lanes inside the bound rather + # than all of them -- see PLENA_Tools/verification/check_mem.py. + parser.add_argument("--layer-rtol", type=float, default=0.12) + parser.add_argument( + "--layer-atol", + type=float, + default=2.5e-4, + help="Absolute floor for the layer cases, above bf16's ~1.8e-4 on this " + "contraction and ~40x below the smallest golden value.", + ) + args = parser.parse_args() + + build_dir = Path(__file__).parent / "build" / f"kda_{args.case}" + hw = setup_hw(args, build_dir) + + print("=" * 80) + print(f"KDA stage test: {args.case} (mlen={args.mlen}, blen={args.blen}, chunk={args.chunk})") + print("=" * 80) + + CASES[args.case](args, build_dir, hw) diff --git a/transactional_emulator/testbench/mamba2/__init__.py b/transactional_emulator/testbench/mamba2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/transactional_emulator/testbench/mamba2/mamba2_stage_test.py b/transactional_emulator/testbench/mamba2/mamba2_stage_test.py new file mode 100644 index 000000000..5182959d1 --- /dev/null +++ b/transactional_emulator/testbench/mamba2/mamba2_stage_test.py @@ -0,0 +1,350 @@ +# ruff: noqa: E402 +"""Numerical stage tests for the Mamba-2 lowering, on the transactional emulator. + +Each case compiles ONE stage of the Mamba-2 mixer to ISA, assembles it, runs the +Rust emulator, and compares the result against a float32 PyTorch golden. What is +and is not covered: + + covered the compiler emitters in aten/plena/program_mamba_common.py and + program_ssd.py, the two new opcodes (V_SOFTPLUS_V, S_MAP_FP_V) + end to end through assembler + emulator, and the bf16 rounding + the emulator applies on every vector write. + + NOT covered the RTL's fixed-point exp model. The emulator implements + V_EXP_V as libtorch's exact exp, while + PLENA_Tools/plena_quant/quant_operations/exp.py is a range + reduction plus a 3-term Taylor with ~0.5-1.5% *systematic* + relative error. Attention's softmax normalises such a bias away; + Mamba applies exp on the critical path of a multiplicative + recurrence, where it compounds. Passing here bounds the + lowering, not the silicon. + +Cases: + dt softplus(dt_raw + dt_bias) then clamp -- exercises V_SOFTPLUS_V + cumsum cs = a @ lower_triangular_ones -- the prefix-scan substitute + decay exp(min(cs_i - cs_j, 0)) * causal -- exercises S_MAP_FP_V and + V_SUB_VF's rorder=1 form + conv1d causal depthwise conv, kernel 4 + +Run: python3 mamba2_stage_test.py --case dt [--mlen 64] [--blen 4] +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import torch + +# Prefer the pinned in-repo submodule (PLENA_Simulator/PLENA_Compiler) over a +# sibling checkout: a sibling may be on a different branch and would otherwise +# silently shadow the submodule on sys.path. +_REPO_ROOT = Path(__file__).resolve().parents[3] +for _compiler_root in (_REPO_ROOT / "PLENA_Compiler", _REPO_ROOT.parent / "PLENA_Compiler"): + if (_compiler_root / "aten" / "plena" / "compiler.py").exists(): + sys.path.insert(0, str(_compiler_root)) + break + +from compiler.aten.plena import PlenaCompiler +from compiler.aten.plena.program_mamba_common import Mamba2Shape +from compiler.aten.plena.program_ssd import HOST_STAGED +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.emulator_runner import run_and_assert +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + +#: Slot order of the Mamba FPRAM constant block, matching +#: ProgramMambaCommonMixin.mamba_fp_constants. Pinned here so a reordering there +#: fails this test rather than silently feeding the kernel wrong constants -- +#: FPRAM contents are never checked at runtime. +FP_CONST_ORDER = ["zero", "one", "neg_one", "dt_min", "dt_max", "reci_group", "eps"] + + +def _mxfp8_exact(shape_, generator, *, hi: float = 2.0) -> torch.Tensor: + """Draw values that MX-FP8 e4m3 represents exactly. + + Inputs staged through HBM are quantised to e4m3 (3 mantissa bits) with a + shared e8m0 scale per block of 8. If the test drew arbitrary floats, the + reported error would be dominated by that input quantisation rather than by + the kernel, and a real lowering bug an order of magnitude smaller would hide + inside it. Multiples of 1/8 with magnitude below 2 are exact in e4m3 + (spacing is 2^-3 relative to the exponent, i.e. 0.125 across [1,2) and finer + below), so the quantisation becomes the identity and the remaining error is + the kernel's. Same trick the shared-expert MoE test uses to get atol=rtol=0. + """ + steps = int(hi * 8) + raw = torch.randint(-steps, steps + 1, shape_, generator=generator) + return raw.to(torch.float32) / 8.0 + + +def _shape(args, *, seq_len: int) -> Mamba2Shape: + mlen = args.mlen + return Mamba2Shape( + hidden_size=args.hidden_size or 2 * mlen, + num_heads=args.num_heads, + head_dim=mlen, + state_size=mlen, + n_groups=1, + conv_kernel=4, + chunk_size=mlen, + seq_len=seq_len, + time_step_min=args.dt_min, + time_step_max=args.dt_max, + ) + + +def _write_comparison(build_dir: Path, prog, var, *, rows: int, cols: int, mlen: int): + addr = prog.get_vram_addr(var.name) + params = { + "start_row_idx": addr // mlen, + "num_rows": (rows * cols) // mlen, + "num_batches": rows, + "elements_per_batch": cols, + "row_dim": mlen, + "use_stride_mode": cols > mlen, + # All four cases are bit-exact, so assert that rather than inheriting the + # default atol=rtol=0.2 at a 90% match bar -- which would let a real + # regression drift most of the way to wrong before going red, while the CI + # comment claims a hard failure. Inputs are drawn exactly representable in + # MX-FP8 precisely so this is achievable; see _mxfp8_exact. + "atol": 0.0, + "rtol": 0.0, + } + with open(build_dir / "comparison_params.json", "w") as f: + json.dump(params, f, indent=2) + return addr + + +def _finish(build_dir: Path, prog, golden, input_tensors, fp_preload, order, case, args): + gen_code = prog.compile() + print(f"\nGenerated {len(gen_code.splitlines())} lines of ISA") + create_sim_env( + input_tensors, + gen_code, + {"original_output": golden}, + fp_preload, + build_dir=str(build_dir), + ) + # Stage each tensor at the compiler's own HBM base. The allocator rounds a + # tensor's footprint up to a whole 64-byte scale row (hbm_tensor_size), so a + # tensor whose row count is not a multiple of 8 leaves a gap the contiguous + # writer would not reproduce -- every following tensor would then be read + # 32 bytes early, and the tail of the read lands in the scale stream, whose + # 0x7f bytes decode to e4m3 NaN. Same call every other testbench makes. + hbm_addrs = {name: prog.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + order = sorted(order, key=lambda name: hbm_addrs[name]) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm=f"mamba2_{case}", + data=None, + specified_data_order=order, + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + ) + with open(build_dir / "generated_asm_code.asm", "w") as f: + f.write(gen_code) + run_and_assert(build_dir, f"mamba2_{case}", mlen=args.mlen, blen=args.blen) + + +# ============================================================================ +# Cases +# ============================================================================ + + +def case_dt(args, build_dir, hw): + """dt = clamp(softplus(dt_raw + dt_bias)). Exercises V_SOFTPLUS_V.""" + mlen = args.mlen + shape = _shape(args, seq_len=mlen) + rows = args.num_heads + + g = torch.Generator().manual_seed(args.seed) + dt_raw = _mxfp8_exact((rows, mlen), g) + # Force exact powers of two into the first lanes so the softplus tails are + # still exercised: these are exact in e4m3 at any magnitude, and they are + # where a naive log1p(exp(x)) overflows (+16) or underflows (-16). + for col, val in enumerate([16.0, -16.0, 8.0, -8.0, 0.0]): + dt_raw[:, col] = val + bias_row = _mxfp8_exact((1, mlen), g, hi=1.0) + + golden = torch.nn.functional.softplus(dt_raw + (0.0 if args.no_bias else bias_row)) + golden = torch.clamp(golden, min=args.dt_min, max=args.dt_max) + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + consts = prog.mamba_fp_constants(shape) + fp_preload = prog.mamba_fp_constant_values(shape) + [0.0] * 3 + + dt_in = prog.input("DT", shape=(rows, mlen)) + # The bias is logically ONE row, but H_PREFETCH_V moves a fixed + # HBM_V_Prefetch_Amount x VLEN granule, so a 1-row HBM tensor is read a whole + # granule wide and the surplus rows come from whatever follows it in HBM. + # Staging a full granule (the row replicated) and taking a 1-row VRAM view + # keeps every byte the DMA touches defined -- what a real deployment does + # anyway when the host lays out per-head parameters. + # + # NOTE: this is *not* what produced the NaN this case used to emit. That was + # the HBM base mismatch fixed in _finish (BIAS was staged 32 B before the + # address the ISA reads it from, so the tail of the read fell in the scale + # stream and 0x7f decoded to an e4m3 NaN). With hbm_addrs passed, a plain + # 1-row BIAS also passes bit-exactly; the granule staging is kept as defence, + # not as a workaround for a hardware limit. + prefetch_rows = prog.hbm_v_prefetch_amount + bias_staged = bias_row.repeat(prefetch_rows, 1) + bias_in = prog.input("BIAS", shape=(prefetch_rows, mlen)) + dt_v = prog.load_batch(dt_in, name="dt") + bias_full = prog.load_batch(bias_in, name="dt_bias_full") + bias_v = prog.alloc_at("dt_bias", 1, mlen, prog.get_vram_addr(bias_full.name)) + + prog.mamba_dt_activation_v0(dt_v, None if args.no_bias else bias_v, consts, shape, rows=list(range(rows))) + + _write_comparison(build_dir, prog, dt_v, rows=rows, cols=mlen, mlen=mlen) + _finish(build_dir, prog, golden, {"DT": dt_raw, "BIAS": bias_staged}, fp_preload, ["DT", "BIAS"], "dt", args) + + +def case_cumsum(args, build_dir, hw): + """cs[h, t] = sum_{s<=t} a[h, s] via M_MM against lower-triangular ones.""" + mlen = args.mlen + shape = _shape(args, seq_len=mlen) + rows = args.num_heads + + torch.manual_seed(args.seed) + # a = A_h * dt is strictly negative in the real kernel; keep it so, and keep + # |cs| in a range where bf16 still resolves it (cs feeds an exponent). + g = torch.Generator().manual_seed(args.seed) + # Strictly negative (A < 0 in the real kernel) and exactly representable. + a = -(torch.randint(1, 9, (rows, mlen), generator=g).to(torch.float32) / 128.0) + golden = torch.cumsum(a, dim=1) + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + prog.mamba_fp_constants(shape) + fp_preload = prog.mamba_fp_constant_values(shape) + [0.0] * 3 + + tri = torch.tril(torch.ones(mlen, mlen)).T # U[s, t] = 1 iff s <= t + a_in = prog.input("A", shape=(rows, mlen)) + tri_in = prog.input("TRI", shape=(mlen, mlen)) + a_v = prog.load_batch(a_in, name="a_t") + cs_v = prog.alloc("cs_t", rows, mlen, strict=False) + + prog.ssd_chunk_cumsum_v0(a_v, tri_in, cs_v, shape, precision=HOST_STAGED) + + _write_comparison(build_dir, prog, cs_v, rows=rows, cols=mlen, mlen=mlen) + _finish(build_dir, prog, golden, {"A": a, "TRI": tri}, fp_preload, ["A", "TRI"], "cumsum", args) + + +def case_decay(args, build_dir, hw): + """decay[i, j] = exp(min(cs_i - cs_j, 0)) * causal[i, j] for one head. + + The stage that needs S_MAP_FP_V (the whole cs row into FPRAM in one + instruction) and V_SUB_VF's rorder=1 reverse form (cs_i - cs_j from a row of + cs_j and a scalar cs_i). + """ + mlen = args.mlen + shape = _shape(args, seq_len=mlen) + + torch.manual_seed(args.seed) + g = torch.Generator().manual_seed(args.seed) + # `cs` is a *staged* input here, so it takes the HBM MX-FP8 path and keeps + # three mantissa bits per element. A cumsum is exactly the shape e4m3 cannot + # hold -- 64 partial sums spread over a decade of dynamic range above a fixed + # step -- so drawing `a` off the natural -k/128 grid makes the reported error + # the *input* quantisation rather than the kernel's: measured 0.18 max abs + # (90% of the atol budget) while the kernel itself was within 1e-3 of + # exp(min(cs_i - cs_j, 0)) evaluated on the cs the hardware actually saw. In + # the real pipeline cs never round-trips through HBM; it comes out of + # ssd_chunk_cumsum_v0 as bf16 in VRAM. + # So draw a cumsum e4m3 represents exactly: steps of -1/8, few enough that + # every partial sum stays a multiple of 1/8 with |cs| <= 2 -- the same + # exactness argument as `_mxfp8_exact`. decay then spans exp(-2)..exp(0). + a = torch.zeros(1, mlen) + a[0, torch.randperm(mlen, generator=g)[: min(mlen, 16)]] = -0.125 + cs = torch.cumsum(a, dim=1) # [1, mlen], non-increasing + + d = cs[0][:, None] - cs[0][None, :] + golden = torch.exp(torch.clamp(d, max=0.0)) * torch.tril(torch.ones(mlen, mlen)) + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + consts = prog.mamba_fp_constants(shape) + fp_preload = prog.mamba_fp_constant_values(shape) + [0.0] * 3 + + cs_in = prog.input("CS", shape=(1, mlen)) + causal_in = prog.input("CAUSAL", shape=(mlen, mlen)) + cs_v = prog.load_batch(cs_in, name="cs_t") + causal_v = prog.load_batch(causal_in, name="causal") + decay_v = prog.alloc("decay", mlen, mlen, strict=False) + cs_fp = prog.fp_var("cs_fp", size=mlen) + + prog.ssd_decay_mask_v0(cs_v, cs_fp, decay_v, causal_v, consts, shape, head_row=0) + + _write_comparison(build_dir, prog, decay_v, rows=mlen, cols=mlen, mlen=mlen) + _finish( + build_dir, + prog, + golden, + {"CS": cs, "CAUSAL": torch.tril(torch.ones(mlen, mlen))}, + fp_preload, + ["CS", "CAUSAL"], + "decay", + args, + ) + + +def case_conv1d(args, build_dir, hw): + """Causal depthwise conv1d, kernel 4, sequence on the VRAM row axis.""" + mlen = args.mlen + seq = getattr(args, "seq_len", None) or 8 + shape = _shape(args, seq_len=seq) + k = shape.conv_kernel + + torch.manual_seed(args.seed) + g = torch.Generator().manual_seed(args.seed) + x = _mxfp8_exact((seq, mlen), g) + w = _mxfp8_exact((k, mlen), g, hi=1.0) + + golden = torch.zeros(seq, mlen) + for s in range(seq): + for j in range(k): + src = s - (k - 1) + j + if src >= 0: + golden[s] += x[src] * w[j] + + prog = PlenaCompiler(mlen=mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + prog.mamba_fp_constants(shape) + fp_preload = prog.mamba_fp_constant_values(shape) + [0.0] * 3 + + x_in = prog.input("X", shape=(seq, mlen)) + w_in = prog.input("W", shape=(k, mlen)) + x_v = prog.load_batch(x_in, name="x") + w_v = prog.load_batch(w_in, name="conv_w") + out_v = prog.alloc("conv_out", seq, mlen, strict=False) + scratch = prog.alloc("conv_scratch", 4, mlen, strict=False) + + prog.mamba_conv1d_v0(x_v, w_v, None, out_v, scratch, shape, num_rows=seq) + + _write_comparison(build_dir, prog, out_v, rows=seq, cols=mlen, mlen=mlen) + _finish(build_dir, prog, golden, {"X": x, "W": w}, fp_preload, ["X", "W"], "conv1d", args) + + +CASES = {"dt": case_dt, "cumsum": case_cumsum, "decay": case_decay, "conv1d": case_conv1d} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--case", default="dt", choices=sorted(CASES)) + parser.add_argument("--no-bias", action="store_true", help="skip dt_bias (isolates V_SOFTPLUS_V)") + parser.add_argument("--num-heads", type=int, default=4) + parser.add_argument("--dt-min", type=float, default=0.0) + parser.add_argument("--dt-max", type=float, default=float("inf")) + args = parser.parse_args() + + build_dir = Path(__file__).parent / "build" / f"mamba2_{args.case}" + hw = setup_hw(args, build_dir) + + print("=" * 80) + print(f"Mamba-2 stage test: {args.case} (mlen={args.mlen}, blen={args.blen})") + print("=" * 80) + + CASES[args.case](args, build_dir, hw) diff --git a/transactional_emulator/testbench/moe_timing/replay/test_branch_split.py b/transactional_emulator/testbench/moe_timing/replay/test_branch_split.py index 9d35b9eae..350a96f81 100644 --- a/transactional_emulator/testbench/moe_timing/replay/test_branch_split.py +++ b/transactional_emulator/testbench/moe_timing/replay/test_branch_split.py @@ -40,7 +40,39 @@ #: lm_head or the next sublayer. It is listed separately rather than folded into #: PLUMBING_STAGES because plumbing is MoE cost that both branches share, and this #: is not MoE cost. -NON_MOE_STAGES = ("non_moe",) +NON_MOE_STAGES = ( + "non_moe", + # The state-engine stages: Mamba-2's selective SSM and KDA's gated delta + # rule. They are in this bucket rather than in a branch because + # `_branch_split` measures shared-versus-routed *MoE* cost, and a recurrent + # mixer is not MoE cost at all -- it is what sits where attention would in a + # hybrid layer. Putting them in PLUMBING_STAGES would fold them into the + # per-layer MoE overhead both branches share, which is a different claim and + # a wrong one. + "mamba_in_proj", + "mamba_conv1d", + "mamba_dt", + "mamba_chunk_cumsum", + "mamba_decay_mask", + "mamba_intra_chunk", + "mamba_state_update", + "mamba_inter_chunk", + "mamba_skip", + "mamba_gated_norm", + "mamba_out_proj", + "mamba_state_load", + "mamba_state_store", + "kda_qkv_proj", + "kda_conv1d", + "kda_normalize", + "kda_decay", + "kda_state_update", + "kda_readout", + "kda_gated_norm", + "kda_out_proj", + "kda_state_load", + "kda_state_store", +) def _emulator_stage_names() -> list[str]: