From fe6ba531a26069007175159c66b2704b63f24bb0 Mon Sep 17 00:00:00 2001 From: "Qichao (Arlo) Wang" Date: Sat, 29 Aug 2026 01:39:29 +0100 Subject: [PATCH 1/6] isa: three opcodes for the static state path `main` stops at 0x38. This adds three, all ordinary fixed-function ops with their operands named in the instruction word -- no descriptor, no queue, no run-time state: 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 are about. Both recurrences used a `copy + multiply + add` triple per state row, which could not become a hardware loop: the scratch row it staged through was the *same* row every iteration, so the destination never formed an arithmetic progression and the sweep was unrolled in full. The program image therefore grew linearly with the contracted dimension -- 7,345 instructions for one KDA head at key_dim 128, 1,384 for one Mamba head at state_size 32. Folded into one FMA, each sweep is one instruction inside one hardware loop and the image is constant: 76 per (head, block) for KDA at every key_dim from 4 to 128, 54 per head for Mamba at every state_size. `_row_progression` also stops refusing a step of zero. It rejected one saying it "would cause infinite HW loop", which is not what happens -- `C_LOOP_START` takes a trip count, so a zero step is a pinned pointer and the loop terminates on the counter. Refusing it forced every sweep with a fixed accumulator onto the unrolled path. `_emit_tile_row_fma` takes `(dst, src, fp)` triples, so the three walks may progress independently. That is what lets a later commit read one operand at a stride without moving it -- and it is also the shape a post-increment operand would generalise, which the measurements price at 2.83x. Co-Authored-By: Claude Opus 5 --- asm_templates/mamba_conv1d_asm.py | 553 ++++++++++++++++++ asm_templates/preload_act.py | 7 +- assembler/assembly_to_binary.py | 32 +- assembler/parser.py | 18 +- assembler/tests/test_vector_rmask_handling.py | 69 +++ aten/plena/isa_tile_rows.py | 333 ++++++++++- aten/tests/test_v_fma_vf.py | 302 ++++++++++ doc/configuration.svh | 7 +- doc/operation.svh | 15 +- doc/plena_isa_spec.md | 103 +++- 10 files changed, 1414 insertions(+), 25 deletions(-) create mode 100644 asm_templates/mamba_conv1d_asm.py create mode 100644 aten/tests/test_v_fma_vf.py diff --git a/asm_templates/mamba_conv1d_asm.py b/asm_templates/mamba_conv1d_asm.py new file mode 100644 index 0000000..f4ef861 --- /dev/null +++ b/asm_templates/mamba_conv1d_asm.py @@ -0,0 +1,553 @@ +"""Mamba-2 mixer ASM templates: causal depthwise conv1d and the chunked SSD scan. + +Both emitters live here because they are the two Mamba-2 kernels with no existing +template to reuse; everything else in the mixer (in_proj / out_proj / RMSNorm / +SiLU) is lowered by ``projection_asm``, ``normalization_asm`` and ``silu_asm``. + +They are consumed by the *generator* pipeline (``generator/passes/code_gen.py``), +which imports them directly rather than through ``asm_templates/__init__.py``. + +VRAM layout contract +-------------------- +Every tensor here is addressed as a flat run of VLEN-wide VRAM rows with the +**sequence on the row axis**, i.e. the row holding channel-tile ``c`` at timestep +``s`` sits at ``base + (c * seq_len + s) * VLEN``. That is the layout +``doc/memory_layout.md`` calls ``[h//VLEN, b, s, VLEN]``. It is what makes the +depthwise conv1d cheap: a causal shift along ``t`` becomes a pure row-address +offset (free), and the per-channel tap weight ``w[j, :]`` is exactly a VLEN-wide +``V_MUL_VV`` operand. + +Why the conv is *not* lowered through ``im2col_asm`` +--------------------------------------------------- +``im2col_asm`` hardcodes a square KxK spatial patch, has no padding support (so +it cannot express the causal left pad), and would materialise a dense +``(M, C_in*K*K)`` GEMM against a block-diagonal weight -- wasting ``conv_dim``-fold +MACs. ``V_MUL_VF`` is equally wrong: it broadcasts one FP register across every +lane, but a depthwise tap varies *per channel*, i.e. per lane. Hence +``V_MUL_VV`` against a resident tap row. + +ISA facts these emitters respect +-------------------------------- +* ``C_LOOP_START`` takes a compile-time immediate trip count and there is no + branch, so every trip count below is a Python-side constant. +* There is no logarithm, so ``softplus`` uses the ``V_SOFTPLUS_V`` opcode; the + ``time_step_limit`` clamp reuses ``V_MAX_VF`` / ``V_MIN_VF`` (there is no + ``V_MAX_VV`` / ``V_MIN_VV``). +* There is no VRAM->VRAM move and no gather. A row is zeroed by multiplying it + by ``f0`` (hardwired 0.0) -- see ``_zero_row``. +* MRAM is write-only from HBM, so an activation-times-activation product costs a + ``H_STORE_V`` + ``H_PREFETCH_M`` round trip. The SSD scan has four such + products per chunk and pays for all four explicitly (``_spill_asm`` + + ``batched_matmul_asm``); that round trip is the dominant real cost of this + kernel on PLENA and hiding it would make the utilization report a lie. + +Fidelity note +------------- +This is the *generator* pipeline, whose purpose is structural and utilization +analysis rather than numerics (see ``docs/COMPILATION_PIPELINES.md``). The +instruction mix, the loop nest and the memory traffic are modelled faithfully; +the individual addresses are not a numerically executable Mamba-2. Places where +that gap is more than addressing are called out inline with ``; APPROX:``. +""" + +from __future__ import annotations + +import math + +from ._imm import add_large_int as _add_large_int +from ._imm import load_large_int as _load_large_int +from .batched_matmul_asm import batched_matmul_asm + +__all__ = ["mamba_conv1d_asm", "mamba_ssd_scan_asm"] + + +def _zero_row(reg: int) -> str: + """Zero the VRAM row addressed by ``gp{reg}``. + + ``f0`` is hardwired to 0.0, so ``x * 0`` is the cheapest zero the ISA offers + (there is no VRAM immediate store and no VRAM->VRAM move). The one input it + is not a true zero for is a NaN/Inf left over in the row; the generator path + never depends on the pre-state of a scratch row. + """ + return f"V_MUL_VF gp{reg}, gp{reg}, f0, 0" + + +def _vec_row_loop( + *, + body: list[str], + ptr_regs: list[int], + trip: int, + loop_reg: int, + vlen: int, + comment: str = "", +) -> list[str]: + """Wrap ``body`` in a hardware loop of ``trip`` iterations, bumping each + pointer in ``ptr_regs`` by one VRAM row per iteration. + + ``trip`` is a compile-time immediate: ``C_LOOP_START`` cannot take a + register-held trip count. + """ + if trip <= 0: + return [] + lines: list[str] = [] + if comment: + lines.append(f"; {comment} ({trip} rows)") + lines.append(f"C_LOOP_START gp{loop_reg}, {trip}") + lines.extend(body) + for r in ptr_regs: + lines.append(f"S_ADDI_INT gp{r}, gp{r}, {vlen}") + lines.append(f"C_LOOP_END gp{loop_reg}") + return lines + + +def _spill_asm( + *, + vram_base: int, + elements: int, + hbm_addr_reg: int, + vram_reg: int, + hbm_reg: int, + loop_reg: int, + vlen: int, + writeback_amount: int, + label: str, +) -> list[str]: + """VRAM -> HBM spill of ``elements`` elements, so they can come back as a + matrix operand through ``H_PREFETCH_M``. + + MRAM cannot be written from VRAM; the only path is HBM. Each ``H_STORE_V`` + moves ``HBM_V_Writeback_Amount * VLEN`` elements. + """ + per_issue = writeback_amount * vlen + issues = math.ceil(elements / per_issue) if elements > 0 else 0 + if issues == 0: + return [] + lines = [f"; spill {label}: {elements} elements -> HBM[a{hbm_addr_reg}] ({issues} x H_STORE_V)"] + lines.extend(_load_large_int(vram_reg, vram_base)) + lines.append(f"S_ADDI_INT gp{hbm_reg}, gp0, 0") + lines.append(f"C_LOOP_START gp{loop_reg}, {issues}") + lines.append(f"H_STORE_V gp{vram_reg}, gp{hbm_reg}, a{hbm_addr_reg}, 1, 0") + lines.append(f"S_ADDI_INT gp{vram_reg}, gp{vram_reg}, {per_issue}") + lines.append(f"S_ADDI_INT gp{hbm_reg}, gp{hbm_reg}, {per_issue}") + lines.append(f"C_LOOP_END gp{loop_reg}") + return lines + + +def mamba_conv1d_asm( + *, + vlen: int, + seq_len: int, + conv_dim: int, + conv_kernel: int, + alive_registers: list[int], + input_base_address: int, + output_base_address: int, + weight_base_address: int, + scratch_base_address: int, + bias_base_address: int | None = None, +) -> str: + """Causal depthwise conv1d of width ``conv_kernel`` along the sequence axis. + + ``out[c, s] = bias[c] + sum_j w[j, c] * x[c, s - (conv_kernel - 1) + j]`` + with taps that fall before ``s = 0`` dropped (the causal left pad). + + Rows are addressed per the module-level layout contract. Tap ``j`` of + channel-tile ``c`` lives at ``weight_base_address + (j * n_tiles + c) * VLEN``, + i.e. ``conv_dim`` elements apart per tap. + + The first ``conv_kernel - 1`` timesteps are unrolled because each sees a + different number of live taps; the remaining ``seq_len - conv_kernel + 1`` + share one body and run in a hardware loop. + + Args: + alive_registers: at least 8 GP registers. In order: x pointer, out + pointer, tap-0 pointer, two address temporaries, the product scratch + row, the loop counter, and the bias pointer. + """ + if conv_kernel < 1: + raise ValueError(f"conv_kernel must be >= 1, got {conv_kernel}") + if seq_len < 1: + raise ValueError(f"seq_len must be >= 1, got {seq_len}") + if len(alive_registers) < 8: + raise ValueError(f"mamba_conv1d_asm needs 8 GP registers, got {len(alive_registers)}") + + x_ptr, o_ptr, w_ptr, tmp_a, tmp_b, scr, loop_reg, bias_reg = alive_registers[:8] + + n_tiles = math.ceil(conv_dim / vlen) + k = conv_kernel + # Tap stride in elements: consecutive taps are a whole channel row apart. + tap_stride = n_tiles * vlen + + lines = [ + "; === Mamba-2 causal depthwise conv1d ===", + f"; conv_dim={conv_dim} kernel={k} seq_len={seq_len} channel_tiles={n_tiles} VLEN={vlen}", + "; Sequence is on the VRAM row axis, so the causal shift is a row-address offset", + "; and each tap is a VLEN-wide V_MUL_VV operand (V_MUL_VF would broadcast one", + "; scalar over all lanes, but a depthwise tap varies per channel = per lane).", + ] + if conv_dim % vlen: + lines.append( + f"; APPROX: conv_dim={conv_dim} is not a multiple of VLEN={vlen}; the final tile is" + ) + lines.append("; emitted full-width (lane masking via C_SET_V_MASK_REG is not modelled here).") + + lines.extend(_load_large_int(scr, scratch_base_address)) + + for c in range(n_tiles): + tile_off = c * seq_len * vlen + lines.append(f"; -- channel tile {c}/{n_tiles} --") + lines.extend(_load_large_int(w_ptr, weight_base_address + c * vlen)) + if bias_base_address is not None: + lines.extend(_load_large_int(bias_reg, bias_base_address + c * vlen)) + + # Prologue: timesteps that see the causal zero pad. + for s in range(min(k - 1, seq_len)): + first_live = (k - 1) - s + lines.append(f"; t={s}: taps 0..{first_live - 1} fall before the sequence (zero pad)") + lines.extend(_load_large_int(o_ptr, output_base_address + tile_off + s * vlen)) + for n, j in enumerate(range(first_live, k)): + src = s - (k - 1) + j + lines.extend(_load_large_int(tmp_a, input_base_address + tile_off + src * vlen)) + lines.extend(_add_large_int(tmp_b, w_ptr, j * tap_stride)) + if n == 0: + # First live tap initialises the accumulator, so no zeroing pass. + lines.append(f"V_MUL_VV gp{o_ptr}, gp{tmp_a}, gp{tmp_b}, 0") + else: + lines.append(f"V_MUL_VV gp{scr}, gp{tmp_a}, gp{tmp_b}, 0") + lines.append(f"V_ADD_VV gp{o_ptr}, gp{o_ptr}, gp{scr}, 0") + if bias_base_address is not None: + lines.append(f"V_ADD_VV gp{o_ptr}, gp{o_ptr}, gp{bias_reg}, 0") + + # Steady state: every timestep from k-1 on sees all k taps. + steady = seq_len - (k - 1) + if steady > 0: + lines.append(f"; t={k - 1}..{seq_len - 1}: all {k} taps live") + # x pointer trails the output pointer by (k-1) rows. + lines.extend(_load_large_int(x_ptr, input_base_address + tile_off)) + lines.extend(_load_large_int(o_ptr, output_base_address + tile_off + (k - 1) * vlen)) + body: list[str] = [] + for j in range(k): + body.append(f"S_ADDI_INT gp{tmp_a}, gp{x_ptr}, {j * vlen}") + body.extend(_add_large_int(tmp_b, w_ptr, j * tap_stride)) + if j == 0: + body.append(f"V_MUL_VV gp{o_ptr}, gp{tmp_a}, gp{tmp_b}, 0") + else: + body.append(f"V_MUL_VV gp{scr}, gp{tmp_a}, gp{tmp_b}, 0") + body.append(f"V_ADD_VV gp{o_ptr}, gp{o_ptr}, gp{scr}, 0") + if bias_base_address is not None: + body.append(f"V_ADD_VV gp{o_ptr}, gp{o_ptr}, gp{bias_reg}, 0") + lines.extend( + _vec_row_loop( + body=body, + ptr_regs=[x_ptr, o_ptr], + trip=steady, + loop_reg=loop_reg, + vlen=vlen, + ) + ) + + return "\n".join(lines) + "\n" + + +def mamba_ssd_scan_asm( + *, + mlen: int, + vlen: int, + blen: int, + seq_len: int, + chunk_size: int, + num_heads: int, + head_dim: int, + state_size: int, + n_groups: int, + alive_registers: list[int], + vram: dict[str, int], + act_spill_addr_reg: int, + wt_spill_addr_reg: int, + writeback_amount: int = 4, + dt_min_fp_address: int = 0, + dt_max_fp_address: int = 0, + a_decay_fp_address: int = 0, + d_skip_fp_address: int = 0, +) -> str: + """Chunked state-space-duality scan for one Mamba-2 layer. + + Emits, per chunk of ``chunk_size`` timesteps, the six stages of the SSD + recurrence: + + 1. ``dt = clamp(softplus(dt), time_step_limit)`` (``V_SOFTPLUS_V``) + 2. ``decay = exp(cumsum(dt * A))`` within the chunk (serial ``V_ADD_VV`` chain) + 3. ``G = C @ B^T`` then ``M = G * L`` (GEMM + ``V_MUL_VV``) + 4. ``Y_diag = M @ X`` (GEMM) + 5. ``S = decay * S + B^T @ (X * dt)`` (GEMM + ``V_MUL_VF``) + 6. ``Y = Y_diag + C @ S + D * X`` (GEMM + ``V_ADD_VV``) + + The four GEMMs go through :func:`batched_matmul_asm` with the head (or group) + axis as its batch axis, which is exactly the shape it was written for. Both + of their operands are activations, so each is spilled to HBM first -- see the + module docstring on why that round trip is unavoidable and is emitted rather + than elided. + + ``vram`` maps region names to VLEN-aligned element addresses and must supply: + ``x``, ``B``, ``C``, ``dt``, ``decay``, ``score``, ``state``, ``y``, + ``scratch``. + + Args: + act_spill_addr_reg / wt_spill_addr_reg: HBM address registers for the two + operand spill areas. They must differ: ``batched_matmul_asm`` + addresses both operands from offset 0 of their own register. + alive_registers: at least 8 GP registers. + """ + if len(alive_registers) < 8: + raise ValueError(f"mamba_ssd_scan_asm needs 8 GP registers, got {len(alive_registers)}") + if act_spill_addr_reg == wt_spill_addr_reg: + raise ValueError( + "act_spill_addr_reg and wt_spill_addr_reg must differ: batched_matmul_asm " + "addresses each operand from offset 0 of its own HBM address register, so " + "sharing one register would overlay the two operands." + ) + for key in ("x", "B", "C", "dt", "decay", "score", "state", "y", "scratch", "row0", "row1"): + if key not in vram: + raise KeyError(f"vram map is missing region {key!r}") + if state_size % mlen or chunk_size % mlen: + raise ValueError( + f"state_size ({state_size}) and chunk_size ({chunk_size}) must be multiples of " + f"MLEN ({mlen}): both are K dimensions of the SSD GEMMs." + ) + if num_heads % n_groups: + raise ValueError(f"num_heads ({num_heads}) must be divisible by n_groups ({n_groups})") + + p0, p1, p2, p3, tmp, scr, loop_reg, hbm_reg = alive_registers[:8] + + num_chunks = math.ceil(seq_len / chunk_size) + heads_per_group = num_heads // n_groups + group_state = n_groups * state_size + + lines = [ + "; === Mamba-2 chunked SSD scan ===", + f"; seq_len={seq_len} chunk={chunk_size} chunks={num_chunks}", + f"; heads={num_heads} head_dim={head_dim} state={state_size} groups={n_groups}" + f" (heads/group={heads_per_group})", + "; Four activation x activation GEMMs per chunk; each operand is spilled", + f"; VRAM -> HBM (a{act_spill_addr_reg} / a{wt_spill_addr_reg}) because MRAM is", + "; write-only from HBM.", + ] + # gp{scr} holds a one-row temporary for the whole kernel. batched_matmul_asm + # only clobbers p1/p2/p3, so this address survives every GEMM below. + lines.extend(_load_large_int(scr, vram["row1"])) + + def spill(vram_base: int, elements: int, addr_reg: int, label: str) -> list[str]: + return _spill_asm( + vram_base=vram_base, + elements=elements, + hbm_addr_reg=addr_reg, + vram_reg=p0, + hbm_reg=hbm_reg, + loop_reg=loop_reg, + vlen=vlen, + writeback_amount=writeback_amount, + label=label, + ) + + def gemm(b: int, m: int, k: int, n: int, result: int, label: str) -> list[str]: + out = [f"; -- GEMM {label}: ({b}, {m}, {k}) @ ({b}, {k}, {n}) --"] + out.append( + batched_matmul_asm( + mlen=mlen, + blen=blen, + b=b, + m=m, + k=k, + n=n, + alive_registers=[p1, p2, p3], + w_base_hbm_offset_reg=wt_spill_addr_reg, + # w_prefetch_amount is only used by an internal sanity assert; the + # spilled operand is contiguous so K tiles are always reachable. + w_prefetch_amount=k, + a_base_hbm_offset_reg=act_spill_addr_reg, + a_prefetch_amount=k, + result_base_address=result, + ).rstrip("\n") + ) + return out + + for ch in range(num_chunks): + rows = min(chunk_size, seq_len - ch * chunk_size) + lines.append(f"; ================ chunk {ch}/{num_chunks} (rows {ch * chunk_size}..{ch * chunk_size + rows - 1}) ================") + if rows != chunk_size: + lines.append( + f"; APPROX: tail chunk holds {rows} of {chunk_size} timesteps; the GEMMs below are" + ) + lines.append("; emitted at full chunk width (the pad rows are computed and discarded).") + + x_chunk = vram["x"] + ch * chunk_size * num_heads * head_dim + b_chunk = vram["B"] + ch * chunk_size * group_state + c_chunk = vram["C"] + ch * chunk_size * group_state + dt_chunk = vram["dt"] + ch * chunk_size * num_heads + y_chunk = vram["y"] + ch * chunk_size * num_heads * head_dim + + # ---------------- 1. dt activation ---------------- + lines.append("; @stage=mamba_dt softplus + time_step_limit clamp") + lines.append(f"S_LD_FP f1, gp0, {dt_min_fp_address}") + lines.append(f"S_LD_FP f2, gp0, {dt_max_fp_address}") + lines.extend(_load_large_int(p0, dt_chunk)) + dt_rows = max(1, math.ceil(chunk_size * num_heads / vlen)) + lines.extend( + _vec_row_loop( + body=[ + # No logarithm in the ISA: softplus is a dedicated opcode. + f"V_SOFTPLUS_V gp{p0}, gp{p0}, 0", + f"V_MAX_VF gp{p0}, gp{p0}, f1, 0", + f"V_MIN_VF gp{p0}, gp{p0}, f2, 0", + ], + ptr_regs=[p0], + trip=dt_rows, + loop_reg=loop_reg, + vlen=vlen, + comment="dt rows", + ) + ) + + # ---------------- 2. intra-chunk decay cumsum ---------------- + lines.append("; @stage=mamba_chunk_cumsum decay = exp(cumsum(dt * A)) along the chunk") + lines.append("; Serial dependent scan: V_PS_V is a within-row prefix scan, but the") + lines.append("; cumsum here runs across rows (sequence is on the row axis), so the") + lines.append("; recurrence is an explicit V_ADD_VV chain of chunk_size steps.") + lines.append(f"S_LD_FP f3, gp0, {a_decay_fp_address}") + lines.extend(_load_large_int(p1, dt_chunk)) + lines.extend(_load_large_int(p2, vram["decay"])) + lines.extend(_load_large_int(p3, vram["row0"])) + lines.append(_zero_row(p3)) # running sum accumulator + lines.extend( + _vec_row_loop( + body=[ + f"V_MUL_VF gp{scr}, gp{p1}, f3, 0", + f"V_ADD_VV gp{p3}, gp{p3}, gp{scr}, 0", + f"V_EXP_V gp{p2}, gp{p3}, 0", + ], + ptr_regs=[p1, p2], + trip=dt_rows, + loop_reg=loop_reg, + vlen=vlen, + comment="cumsum steps", + ) + ) + # Publish the per-head chunk decay scalars into FPRAM so the state update + # below can broadcast them with V_MUL_VF -- this is exactly what the + # S_MAP_FP_V (VRAM row -> VLEN FPRAM slots) opcode exists for. + lines.extend(_load_large_int(p2, vram["decay"])) + lines.extend(_load_large_int(tmp, a_decay_fp_address)) + lines.append(f"S_MAP_FP_V gp{tmp}, gp{p2}, 0") + + # ---------------- 3. intra-chunk scores ---------------- + lines.append("; @stage=mamba_intra_chunk G = C @ B^T") + lines.extend(spill(c_chunk, n_groups * chunk_size * state_size, act_spill_addr_reg, "C_chunk")) + lines.append("; B is spilled untransposed; the transpose is absorbed by H_PREFETCH_M's") + lines.append("; stride mode (element(row, col) at col*stride + row).") + lines.extend(spill(b_chunk, n_groups * chunk_size * state_size, wt_spill_addr_reg, "B_chunk")) + lines.extend( + gemm(n_groups, chunk_size, state_size, chunk_size, vram["score"], "C @ B^T") + ) + + lines.append("; @stage=mamba_decay_mask M = G * L (segment-sum decay mask)") + lines.append("; APPROX: the (chunk, chunk) causal segment-sum mask L is an outer") + lines.append("; difference of the decay vector; the ISA has no outer-product primitive,") + lines.append("; so it is modelled as a per-row V_MUL_VV against the materialised decay") + lines.append("; rows. The instruction count and traffic match; the values do not.") + lines.extend(_load_large_int(p0, vram["score"])) + lines.extend(_load_large_int(p1, vram["decay"])) + lines.extend( + _vec_row_loop( + body=[f"V_MUL_VV gp{p0}, gp{p0}, gp{p1}, 0"], + ptr_regs=[p0], + trip=max(1, num_heads * chunk_size * chunk_size // vlen), + loop_reg=loop_reg, + vlen=vlen, + comment="score rows", + ) + ) + + # ---------------- 4. intra-chunk output ---------------- + lines.append("; @stage=mamba_intra_chunk Y_diag = M @ X") + lines.extend(spill(vram["score"], num_heads * chunk_size * chunk_size, act_spill_addr_reg, "M")) + lines.extend(spill(x_chunk, num_heads * chunk_size * head_dim, wt_spill_addr_reg, "X_chunk")) + lines.extend(gemm(num_heads, chunk_size, chunk_size, head_dim, y_chunk, "M @ X")) + + # ---------------- 5. state update ---------------- + lines.append("; @stage=mamba_state_update S = decay * S + B^T @ (X * dt)") + lines.append("; dt is one scalar per (token, head) but X is head_dim wide, so the dt") + lines.append("; pointer is deliberately not advanced with X: one dt row feeds") + lines.append("; head_dim/VLEN consecutive X rows.") + lines.extend(_load_large_int(p0, vram["scratch"])) + lines.extend(_load_large_int(p1, x_chunk)) + lines.extend(_load_large_int(p2, dt_chunk)) + lines.extend( + _vec_row_loop( + body=[f"V_MUL_VV gp{p0}, gp{p1}, gp{p2}, 0"], + ptr_regs=[p0, p1], + trip=max(1, num_heads * chunk_size * head_dim // vlen), + loop_reg=loop_reg, + vlen=vlen, + comment="X * dt", + ) + ) + lines.extend(spill(b_chunk, num_heads * chunk_size * state_size, act_spill_addr_reg, "B^T")) + lines.extend(spill(vram["scratch"], num_heads * chunk_size * head_dim, wt_spill_addr_reg, "X*dt")) + lines.extend( + gemm(num_heads, state_size, chunk_size, head_dim, vram["scratch"], "B^T @ (X*dt)") + ) + lines.append("; decay the carried state, then fold in this chunk's contribution.") + lines.append(f"S_LD_FP f4, gp0, {a_decay_fp_address}") + lines.extend(_load_large_int(p0, vram["state"])) + lines.extend(_load_large_int(p1, vram["scratch"])) + lines.extend( + _vec_row_loop( + body=[ + f"V_MUL_VF gp{p0}, gp{p0}, f4, 0", + f"V_ADD_VV gp{p0}, gp{p0}, gp{p1}, 0", + ], + ptr_regs=[p0, p1], + trip=max(1, num_heads * state_size * head_dim // vlen), + loop_reg=loop_reg, + vlen=vlen, + comment="state rows", + ) + ) + + # ---------------- 6. inter-chunk contribution + skip ---------------- + lines.append("; @stage=mamba_inter_chunk Y += C @ S") + lines.extend(spill(c_chunk, num_heads * chunk_size * state_size, act_spill_addr_reg, "C_chunk")) + lines.extend(spill(vram["state"], num_heads * state_size * head_dim, wt_spill_addr_reg, "S")) + lines.extend(gemm(num_heads, chunk_size, state_size, head_dim, vram["scratch"], "C @ S")) + lines.extend(_load_large_int(p0, y_chunk)) + lines.extend(_load_large_int(p1, vram["scratch"])) + lines.extend( + _vec_row_loop( + body=[f"V_ADD_VV gp{p0}, gp{p0}, gp{p1}, 0"], + ptr_regs=[p0, p1], + trip=max(1, num_heads * chunk_size * head_dim // vlen), + loop_reg=loop_reg, + vlen=vlen, + comment="Y += Y_off", + ) + ) + + lines.append("; @stage=mamba_skip Y += D * X (per-head scalar skip)") + lines.append(f"S_LD_FP f5, gp0, {d_skip_fp_address}") + lines.extend(_load_large_int(p0, y_chunk)) + lines.extend(_load_large_int(p1, x_chunk)) + lines.extend( + _vec_row_loop( + body=[ + f"V_MUL_VF gp{scr}, gp{p1}, f5, 0", + f"V_ADD_VV gp{p0}, gp{p0}, gp{scr}, 0", + ], + ptr_regs=[p0, p1], + trip=max(1, num_heads * chunk_size * head_dim // vlen), + loop_reg=loop_reg, + vlen=vlen, + comment="skip rows", + ) + ) + + return "\n".join(lines) + "\n" diff --git a/asm_templates/preload_act.py b/asm_templates/preload_act.py index 40c590a..b56d8e0 100644 --- a/asm_templates/preload_act.py +++ b/asm_templates/preload_act.py @@ -13,7 +13,8 @@ def preload_act_asm( activation_offset_reg: int, stride_size=None, vram_stride_mult: int = 1, - storage_precision: int = 1 # bytes per element + storage_precision: int = 1, # HBM bytes per element + precision: int = 0, # H_PREFETCH_V precision class: 0 = Activation, 1 = KeyValue ) -> str: """Preload activation from HBM to VRAM. Layout: (hidden//mlen, batch, mlen).""" generated_code = "; Preload Activation Generation \n" @@ -39,7 +40,7 @@ def preload_act_asm( vram_step = elements_per_prefetch * vram_stride_mult for i in range(math.ceil(hidden_size / elements_per_prefetch)): generated_code += ( - f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, 0, 0 \n" + f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, {precision}, 0 \n" ) generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vram_step} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {elements_per_prefetch} \n" @@ -51,7 +52,7 @@ def preload_act_asm( generated_code += f"S_ADDI_INT gp{a_offset_register}, gp{a_actual_register}, 0 \n" if batch > preload_len: generated_code += f"C_LOOP_START gp{inner_loop_register}, {math.ceil(batch / preload_len)} \n" - generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_offset_register}, a{activation_offset_reg}, 1, 0 \n" + generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_offset_register}, a{activation_offset_reg}, 1, {precision} \n" generated_code += ( f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len * vram_stride_mult} \n" ) diff --git a/assembler/assembly_to_binary.py b/assembler/assembly_to_binary.py index 9cbeb35..da04fbc 100644 --- a/assembler/assembly_to_binary.py +++ b/assembler/assembly_to_binary.py @@ -14,6 +14,7 @@ "V_MUL_VV", "V_SUB_VV", "V_MUL_VF", + "V_FMA_VF", "V_EXP_V", "V_RECI_V", "V_RED_SUM", @@ -21,6 +22,7 @@ "V_MAX_VF", "V_MIN_VF", "V_TOPK", + "V_SOFTPLUS_V", } ) _IMM_RS1_RD_OPS = frozenset( @@ -32,13 +34,11 @@ "S_LD_INT", "S_ST_INT", "S_MAP_V_FP", - "V_RED_MAX", - "V_RECI_V", - "V_EXP_V", + "S_MAP_FP_V", } ) _IMM_RD_OPS = frozenset({"S_LUI_INT", "M_MV_WO", "M_BMM_WO", "M_BMV_WO"}) -_RS1_RD_OPS = frozenset({"S_MV_FP", "S_RECI_FP", "S_EXP_FP", "S_SQRT_FP", "V_EXP_V", "V_RED_SUM"}) +_RS1_RD_OPS = frozenset({"S_MV_FP", "S_RECI_FP", "S_EXP_FP", "S_SQRT_FP"}) _RD_ONLY_OPS = frozenset({"C_SET_SCALE_REG", "C_SET_STRIDE_REG", "C_SET_V_MASK_REG", "C_SET_TOPK_REG", "C_LOOP_END"}) _FUNCT_RSTRIDE_OPS = frozenset({"H_PREFETCH_M", "H_PREFETCH_V", "H_STORE_V", "V_SUB_VF"}) _RS2_RS1_RD_OPS = frozenset( @@ -105,7 +105,25 @@ def _convert_to_binary(self, instruction): # Treat omitted rmask deterministically as "mask disabled" instead of crashing on None << ... rmask = 0 - if instruction.opcode in _IMM_RS1_RD_OPS: + # _RMASK_VECTOR_OPS MUST be tested before _IMM_RS1_RD_OPS / _RS1_RD_OPS. + # V_EXP_V, V_RECI_V, V_RED_MAX and V_RED_SUM appear in more than one set, + # and the earlier branch encodes the third operand as `imm` at bit + # OPCODE_WIDTH + 2*OPERAND_WIDTH (the rs2 field) while the emulator reads + # rmask from rs3 at OPCODE_WIDTH + 3*OPERAND_WIDTH (op.rs `Opcode::decode`). + # With the old ordering a masked V_EXP_V silently executed on the whole + # tile: no diagnostic, wrong answer. rmask == 0 encodes identically under + # either ordering, so this is a no-op for every unmasked call site. + if instruction.opcode in _RMASK_VECTOR_OPS: + # 2- and 3-operand forms leave rs1/rs2 unset; the hardware reads those + # fields regardless, so encode them as 0 rather than crashing on None. + binary_instruction = ( + (rmask << (opw + 3 * ow)) + + ((rs2 or 0) << (opw + 2 * ow)) + + ((rs1 or 0) << (opw + ow)) + + ((rd or 0) << opw) + + opcode + ) + elif instruction.opcode in _IMM_RS1_RD_OPS: binary_instruction = (imm << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode elif instruction.opcode in _IMM_RD_OPS: binary_instruction = (imm << (opw + ow)) + (rd << opw) + opcode @@ -127,10 +145,6 @@ def _convert_to_binary(self, instruction): + (rd << opw) + opcode ) - elif instruction.opcode in _RMASK_VECTOR_OPS: - binary_instruction = ( - (rmask << (opw + 3 * ow)) + (rs2 << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode - ) elif instruction.opcode in _RS2_RS1_RD_OPS: binary_instruction = (rs2 << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode else: diff --git a/assembler/parser.py b/assembler/parser.py index 486c1ba..3ec7625 100644 --- a/assembler/parser.py +++ b/assembler/parser.py @@ -95,9 +95,23 @@ def __repr__(self): _REG_PREFIXES = ("gp", "f", "a") # Hoisted to module scope: these were previously re-created for every line of the # .asm (millions of times for large programs), which dominated sim_env re-parse time. -vector_masked_unary_or_reduction_ops = frozenset({"V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"}) +vector_masked_unary_or_reduction_ops = frozenset( + {"V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX", "V_SOFTPLUS_V"} +) vector_masked_binary_ops = frozenset( - {"V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF", "V_MAX_VF", "V_MIN_VF", "V_TOPK"} + { + "V_ADD_VV", + "V_ADD_VF", + "V_MUL_VV", + "V_SUB_VV", + "V_MUL_VF", + # V_FMA_VF parses like V_MUL_VF -- rd, rs1, fp2, rmask. That rd is also a + # source is invisible here; it only matters in execution. + "V_FMA_VF", + "V_MAX_VF", + "V_MIN_VF", + "V_TOPK", + } ) diff --git a/assembler/tests/test_vector_rmask_handling.py b/assembler/tests/test_vector_rmask_handling.py index 75ec606..fa24e23 100644 --- a/assembler/tests/test_vector_rmask_handling.py +++ b/assembler/tests/test_vector_rmask_handling.py @@ -39,6 +39,75 @@ def test_vector_scalar_minmax_encode_like_masked_vector_ops(self): self.assertEqual((binary >> 14) & 0xF, 3) # rs2 self.assertEqual((binary >> 18) & 0xF, 1) # rmask + def test_masked_unary_and_reduction_ops_place_rmask_in_rs3(self): + """V_EXP_V / V_RECI_V / V_RED_SUM / V_RED_MAX must carry rmask in the rs3 field. + + These four opcodes used to sit in `_IMM_RS1_RD_OPS` / `_RS1_RD_OPS` as well + as `_RMASK_VECTOR_OPS`, and the encoder tested the former first. The mask + therefore landed at bit 14 (the rs2 field) while the emulator reads it from + bit 18 (rs3, see `Opcode::decode` in transactional_emulator/src/op.rs), so a + masked V_EXP_V silently ran on the whole tile. Regression guard for that. + """ + for name in ("V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"): + with self.subTest(opcode=name): + instr = Instruction(name, 1, 2, None, 5, None, None, None) + binary = self.asm._convert_to_binary(instr) + + self.assertEqual(binary & 0x3F, self.asm.isa_definitions[name]) + self.assertEqual((binary >> 6) & 0xF, 1) # rd + self.assertEqual((binary >> 10) & 0xF, 2) # rs1 + self.assertEqual((binary >> 18) & 0xF, 5) # rmask lands in rs3 + + def test_unmasked_encoding_is_unchanged_by_the_rmask_ordering_fix(self): + """rmask == 0 must encode identically under the old and new branch order. + + This is what makes the ordering fix safe for every existing unmasked call + site: the only bits that move are the ones that were previously wrong. + """ + for name in ("V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"): + with self.subTest(opcode=name): + opcode = self.asm.isa_definitions[name] + instr = Instruction(name, 3, 7, None, 0, None, None, None) + expected = (7 << 10) + (3 << 6) + opcode + self.assertEqual(self.asm._convert_to_binary(instr), expected) + + def test_masked_three_operand_unary_round_trips_through_the_parser(self): + """`V_EXP_V gp1, gp2, 1` must reach the encoder as rmask=1, not imm=1.""" + asm_path = "/tmp/plena_test_masked_unary.asm" + with open(asm_path, "w") as f: + f.write("V_EXP_V gp1, gp2, 1\n") + + parsed = parse_asm_file(asm_path) + self.assertEqual(len(parsed), 1) + self.assertEqual(parsed[0].rmask, 1) + + binary = self.asm._convert_to_binary(parsed[0]) + self.assertEqual((binary >> 18) & 0xF, 1) + + def test_v_softplus_encodes_like_masked_unary_op(self): + instr = Instruction("V_SOFTPLUS_V", 1, 2, None, 1, None, None, None) + binary = self.asm._convert_to_binary(instr) + + self.assertEqual(binary & 0x3F, self.asm.isa_definitions["V_SOFTPLUS_V"]) + self.assertEqual((binary >> 6) & 0xF, 1) # rd + self.assertEqual((binary >> 10) & 0xF, 2) # rs1 + self.assertEqual((binary >> 18) & 0xF, 1) # rmask + + def test_s_map_fp_v_encodes_like_its_mirror_s_map_v_fp(self): + """S_MAP_FP_V is the exact inverse of S_MAP_V_FP and shares its operand shape.""" + forward = Instruction("S_MAP_V_FP", 1, 2, None, None, None, None, 4) + inverse = Instruction("S_MAP_FP_V", 1, 2, None, None, None, None, 4) + + fwd_binary = self.asm._convert_to_binary(forward) + inv_binary = self.asm._convert_to_binary(inverse) + + # Same field placement, different opcode. + self.assertEqual(fwd_binary >> 6, inv_binary >> 6) + self.assertEqual(inv_binary & 0x3F, self.asm.isa_definitions["S_MAP_FP_V"]) + self.assertEqual((inv_binary >> 6) & 0xF, 1) # rd (FPRAM base gp reg) + self.assertEqual((inv_binary >> 10) & 0xF, 2) # rs1 (VRAM row gp reg) + self.assertEqual((inv_binary >> 14) & 0x3FFFF, 4) # imm + def test_v_topk_encodes_like_masked_vector_op(self): # Non-zero rmask so the rmask-lane assertion actually exercises the field # (with rmask=0 it would pass even if the encoder dropped the lane). diff --git a/aten/plena/isa_tile_rows.py b/aten/plena/isa_tile_rows.py index 8d56f53..9a93a21 100644 --- a/aten/plena/isa_tile_rows.py +++ b/aten/plena/isa_tile_rows.py @@ -76,6 +76,28 @@ def tile_row_reci( ) -> str: return self._tile_row_single_matrix_op("tile_row_reci_asm", matrix_name, rows, tile_row_idx, tile_col_idx) + def tile_row_softplus( + self, + matrix_name: str, + rows: list[int], + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self._tile_row_single_matrix_op( + "tile_row_softplus_asm", matrix_name, rows, tile_row_idx, tile_col_idx + ) + + def tile_row_to_fpram( + self, + matrix_name: str, + row_map: list[tuple[int, int]], + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self._tile_row_single_matrix_op( + "tile_row_to_fpram_asm", matrix_name, row_map, tile_row_idx, tile_col_idx + ) + def tile_row_sub_fp( self, matrix_name: str, @@ -184,6 +206,55 @@ def tile_row_mul( src_tile_col_idx, ) + def tile_row_add_fp_broadcast( + self, + matrix_name: str, + fpram_scalar_addr: int, + rows: list[int], + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self.tile_row_add_fp_broadcast_asm( + self._tile_addr(matrix_name, tile_row_idx, tile_col_idx), fpram_scalar_addr, rows + ) + + def tile_row_max_fp_broadcast( + self, + matrix_name: str, + fpram_scalar_addr: int, + rows: list[int], + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self.tile_row_max_fp_broadcast_asm( + self._tile_addr(matrix_name, tile_row_idx, tile_col_idx), fpram_scalar_addr, rows + ) + + def tile_row_min_fp_broadcast( + self, + matrix_name: str, + fpram_scalar_addr: int, + rows: list[int], + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self.tile_row_min_fp_broadcast_asm( + self._tile_addr(matrix_name, tile_row_idx, tile_col_idx), fpram_scalar_addr, rows + ) + + def tile_row_sub_fp_broadcast( + self, + matrix_name: str, + fpram_scalar_addr: int, + rows: list[int], + reverse: bool = False, + tile_row_idx: int = 0, + tile_col_idx: int = 0, + ) -> str: + return self.tile_row_sub_fp_broadcast_asm( + self._tile_addr(matrix_name, tile_row_idx, tile_col_idx), fpram_scalar_addr, rows, reverse + ) + def tile_row_mul_fp_broadcast( self, matrix_name: str, @@ -198,6 +269,46 @@ def tile_row_mul_fp_broadcast( rows, ) + def tile_row_fma_fp_sweep( + self, + dst_matrix: str, + src_matrix: str, + fpram_base: int, + dst_rows: list[int], + src_rows: list[int], + dst_tile_row_idx: int = 0, + dst_tile_col_idx: int = 0, + src_tile_row_idx: int = 0, + src_tile_col_idx: int = 0, + ) -> str: + return self.tile_row_fma_fp_sweep_asm( + self._tile_addr(dst_matrix, dst_tile_row_idx, dst_tile_col_idx), + self._tile_addr(src_matrix, src_tile_row_idx, src_tile_col_idx), + fpram_base, + dst_rows, + src_rows, + ) + + def tile_row_fma_fp_broadcast( + self, + dst_matrix: str, + src_matrix: str, + fpram_scalar_addr: int, + dst_rows: list[int], + src_rows: list[int], + dst_tile_row_idx: int = 0, + dst_tile_col_idx: int = 0, + src_tile_row_idx: int = 0, + src_tile_col_idx: int = 0, + ) -> str: + return self.tile_row_fma_fp_broadcast_asm( + self._tile_addr(dst_matrix, dst_tile_row_idx, dst_tile_col_idx), + self._tile_addr(src_matrix, src_tile_row_idx, src_tile_col_idx), + fpram_scalar_addr, + dst_rows, + src_rows, + ) + def vram_fill_zero( self, matrix_name: str, @@ -221,8 +332,16 @@ def _arith_progression(self, values: list[int]) -> tuple[int, int, int] | None: for i in range(2, len(values)): if values[i] - values[i - 1] != step: return None - if step == 0: - return None # Constant sequence (step=0, count>1) would cause infinite HW loop + # A step of 0 is a valid progression. The old code refused it, saying it + # "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, neither of which the address step touches. The + # single-element case above 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: the prediction and the read-out walk the + # state rows against one accumulator row. Refusing it forced those onto + # the unrolled path. return (values[0], len(values), step) def _row_progression(self, rows: list[int]) -> tuple[int, int, int] | None: @@ -256,8 +375,14 @@ def _emit_tile_row_reduce( asm.instr("S_ADD_FP", fp(1), fp(0), fp(0)) asm.instr(opcode, fp(1), gp(gp_src), 0, *opcode_extra_args) asm.instr("S_ST_FP", fp(1), gp(gp_dst), 0) - asm.instr("S_ADDI_INT", gp(gp_src), gp(gp_src), row_step * self.mlen) - asm.instr("S_ADDI_INT", gp(gp_dst), gp(gp_dst), fp_step) + # A zero step is a no-op add. The FMA emitter has always skipped + # it; these three did not, and since the step-0 relaxation they + # can actually see one -- a reduction into a single FPRAM slot, + # or a broadcast constant applied down a tile. + if row_step: + asm.instr("S_ADDI_INT", gp(gp_src), gp(gp_src), row_step * self.mlen) + if fp_step: + asm.instr("S_ADDI_INT", gp(gp_dst), gp(gp_dst), fp_step) asm.instr("C_LOOP_END", gp(gp_loop)) else: for row_idx, fpram_addr in row_map: @@ -322,8 +447,10 @@ def _emit_tile_row_fp_scalar( asm.instr("C_LOOP_START", gp(gp_loop), row_count) asm.instr("S_LD_FP", fp(1), gp(gp_fp), 0) asm.instr(opcode, gp(gp_src), gp(gp_src), fp(1), *opcode_extra_args) - asm.instr("S_ADDI_INT", gp(gp_src), gp(gp_src), row_step * self.mlen) - asm.instr("S_ADDI_INT", gp(gp_fp), gp(gp_fp), fp_step) + if row_step: + asm.instr("S_ADDI_INT", gp(gp_src), gp(gp_src), row_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 row_idx, fpram_addr in row_map: @@ -337,6 +464,69 @@ def _emit_tile_row_fp_scalar( finally: self._reg.free_gp(gp_regs) + 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 all three walks are arithmetic + progressions this collapses to one hardware loop -- the case that + matters, since every recurrent sweep either steps a state tile against a + pinned accumulator or the reverse, and a pinned side is a step of 0. + + The predecessor -- ``mamba_row_copy`` then ``tile_row_mul_fp`` then + ``mamba_row_add`` -- could not loop at all: the scratch row it needed was + the *same* row every iteration, so the copy's destination was constant + and broke the progression. + """ + 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) + # A zero step is a no-op add; skipping it keeps the pinned side + # from costing an instruction per iteration. + 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 _emit_tile_row_vector_op( self, label: str, @@ -392,6 +582,50 @@ def tile_row_exp_asm(self, vram_addr: int, rows: list[int]) -> str: def tile_row_reci_asm(self, vram_addr: int, rows: list[int]) -> str: return self._emit_tile_row_unary("Reciprocal", "V_RECI_V", vram_addr, rows) + def tile_row_softplus_asm(self, vram_addr: int, rows: list[int]) -> str: + return self._emit_tile_row_unary("Softplus", "V_SOFTPLUS_V", vram_addr, rows) + + def tile_row_to_fpram_asm(self, vram_addr: int, row_map: list[tuple[int, int]]) -> str: + """Copy whole VRAM rows into FPRAM via S_MAP_FP_V, one instruction per row. + + `row_map` is [(vram_row_idx, fpram_base_addr)], and each entry moves the full + MLEN-wide row to FPRAM[fpram_base_addr : +MLEN]. This is the bulk inverse of + `S_MAP_V_FP`; the alternative (one-hot `V_MUL_VV` + `V_RED_SUM` + `S_ST_FP` + per element) costs 3 instructions per *scalar* rather than 1 per *row*. + + Note the operand order: `S_MAP_FP_V rd, rs1, imm` takes rd = FPRAM base + register and rs1 = VRAM row register, mirroring `S_MAP_V_FP` so that in both + cases `rd` names the destination memory. + """ + gp_regs = self._reg.allocate_gp(3) + gp_src, gp_dst, gp_loop = gp_regs + try: + asm = IsaBuilder().comment(f"Tile Row -> FPRAM from VRAM[{vram_addr}]") + rows = [row for row, _ in row_map] + fp_addrs = [addr_ for _, addr_ in row_map] + row_prog = self._row_progression(rows) + fp_prog = self._row_progression(fp_addrs) + + if row_prog is not None and fp_prog is not None: + row_start, row_count, row_step = row_prog + fp_start, _, fp_step = fp_prog + asm.instr("S_ADDI_INT", gp(gp_src), gp(0), vram_addr + row_start * self.mlen) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), fp_start) + asm.instr("C_LOOP_START", gp(gp_loop), row_count) + asm.instr("S_MAP_FP_V", gp(gp_dst), gp(gp_src), 0) + asm.instr("S_ADDI_INT", gp(gp_src), gp(gp_src), row_step * self.mlen) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(gp_dst), fp_step) + asm.instr("C_LOOP_END", gp(gp_loop)) + else: + for row_idx, fpram_addr in row_map: + asm.instr("S_ADDI_INT", gp(gp_src), gp(0), vram_addr + row_idx * self.mlen) + asm.instr("S_ADDI_INT", gp(gp_dst), gp(0), fpram_addr) + asm.instr("S_MAP_FP_V", gp(gp_dst), gp(gp_src), 0) + + return self._emit(asm) + finally: + self._reg.free_gp(gp_regs) + def tile_row_sub_fp_asm(self, vram_addr: int, row_map: list[tuple[int, int]]) -> str: return self._emit_tile_row_fp_scalar("Sub", "V_SUB_VF", vram_addr, row_map, opcode_extra_args=(0, 0)) @@ -416,10 +650,97 @@ def tile_row_sub_asm(self, dst_addr: int, src_addr: int, rows: list[int]) -> str def tile_row_mul_asm(self, dst_addr: int, src_addr: int, rows: list[int]) -> str: return self._emit_tile_row_vector_op("Mul", "V_MUL_VV", dst_addr, src_addr, rows) + def tile_row_fma_fp_asm( + self, + dst_addr: int, + src_addr: int, + row_map: list[tuple[int, int, int]], + ) -> str: + """``dst[d] += src[s] * FPRAM[f]``; entries applied in order.""" + return self._emit_tile_row_fma(dst_addr, src_addr, row_map) + + def tile_row_fma_fp_sweep_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``. + + Named ``_sweep`` and not ``_broadcast``: in this file ``_broadcast`` + means one slot applied to *every* row, which is the opposite of walking + a slot per row. See :meth:`tile_row_fma_fp_broadcast_asm` for that. + """ + 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) + + def tile_row_fma_fp_broadcast_asm( + self, + dst_addr: int, + src_addr: int, + fpram_scalar_addr: int, + dst_rows: list[int], + src_rows: list[int], + ) -> str: + """One FPRAM slot applied to every row pair.""" + 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_scalar_addr) for d, s in zip(dst_rows, src_rows)] + return self._emit_tile_row_fma(dst_addr, src_addr, row_map) + def tile_row_mul_fp_broadcast_asm(self, vram_addr: int, fpram_scalar_addr: int, rows: list[int]) -> str: row_map = [(r, fpram_scalar_addr) for r in rows] return self.tile_row_mul_fp_asm(vram_addr, row_map) + # The broadcast variants below apply ONE FPRAM slot to every listed row, as + # opposed to the `tile_row_*_fp` family which walks a different slot per row. + # Both shapes are needed by Mamba: the decay scalars are per-row, while the + # dt clamp bounds and the +1.0 of the sigmoid are single constants. + def tile_row_add_fp_broadcast_asm(self, vram_addr: int, fpram_scalar_addr: int, rows: list[int]) -> str: + return self.tile_row_add_fp_asm(vram_addr, [(r, fpram_scalar_addr) for r in rows]) + + def tile_row_max_fp_broadcast_asm(self, vram_addr: int, fpram_scalar_addr: int, rows: list[int]) -> str: + return self.tile_row_max_fp_asm(vram_addr, [(r, fpram_scalar_addr) for r in rows]) + + def tile_row_min_fp_broadcast_asm(self, vram_addr: int, fpram_scalar_addr: int, rows: list[int]) -> str: + return self.tile_row_min_fp_asm(vram_addr, [(r, fpram_scalar_addr) for r in rows]) + + def tile_row_sub_fp_broadcast_asm( + self, + vram_addr: int, + fpram_scalar_addr: int, + rows: list[int], + reverse: bool = False, + ) -> str: + """``row = row - c`` (reverse=False) or ``row = c - row`` (reverse=True). + + The reverse form is `V_SUB_VF`'s `rorder=1`, and it is the only way to + negate a vector on this ISA (``0.0 - x`` with the hardwired ``f0``). The + Mamba decay matrix needs it to form ``cs_i - cs_j`` from a row holding + ``cs_j`` and a scalar holding ``cs_i``. + """ + row_map = [(r, fpram_scalar_addr) for r in rows] + return self._emit_tile_row_fp_scalar( + "RSub" if reverse else "Sub", + "V_SUB_VF", + vram_addr, + row_map, + opcode_extra_args=(0, 1 if reverse else 0), + ) + def vram_fill_zero_asm( self, vram_addr: int, diff --git a/aten/tests/test_v_fma_vf.py b/aten/tests/test_v_fma_vf.py new file mode 100644 index 0000000..e9d9e26 --- /dev/null +++ b/aten/tests/test_v_fma_vf.py @@ -0,0 +1,302 @@ +"""V_FMA_VF encodes into the same R-type slots as V_MUL_VF. + +The one opcode this work adds. It differs from V_MUL_VF in exactly one bit +field -- the opcode -- because it must reach the same operand decode path in +the RTL; the difference is entirely in execution, where `rd` is read as well +as written. +""" + +import os +import sys +import tempfile +from pathlib import Path + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.assembler.assembly_to_binary import AssemblyToBinary # noqa: E402 +from compiler.assembler.parser import parse_asm_file # noqa: E402 + +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" + assert (fma ^ mul) & 0x3F == (0x3B ^ 0x12) + + +def test_v_fma_vf_carries_a_nonzero_mask(): + """rmask lands in rs3, the same slot the rest of the masked vector family + uses. A mask that silently assembled to zero would run unmasked.""" + (word,) = _assemble("V_FMA_VF gp1, gp2, f3, 5\n") + assert (word >> 18) & 0xF == 5 + + +# --------------------------------------------------------------------------- +# The emitter. The plan's static-footprint claim rests on these sweeps becoming +# hardware loops rather than unrolling. + +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.tests.isa_interpreter import Machine # noqa: E402 + +MLEN = 8 + + +def _body(asm: str) -> list[str]: + return [ + line.strip() + for line in asm.splitlines() + if line.strip() and not line.strip().startswith(";") + ] + + +def _prog(): + return PlenaCompiler(mlen=MLEN, blen=2) + + +def test_fma_sweep_over_a_progression_emits_a_hardware_loop(): + """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, and that row was the *same* row + every iteration, so the copy's destination was constant and broke the + progression. + """ + asm = _prog().tile_row_fma_fp_sweep_asm( + dst_addr=0, src_addr=8192, fpram_base=16, + dst_rows=list(range(128)), src_rows=list(range(128)), + ) + body = _body(asm) + assert sum("C_LOOP_START" in line for line in body) == 1 + assert sum("C_LOOP_END" in line for line in body) == 1 + assert sum("V_FMA_VF" in line for line 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(): + """The contraction walks the source and pins the destination. + + A pinned side is a step of 0, which `_arith_progression` used to refuse. + """ + asm = _prog().tile_row_fma_fp_sweep_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 line for line in body) == 1 + assert sum("V_FMA_VF" in line for line in body) == 1 + # The pinned pointer must not be advanced at all, not advanced by zero. + assert len(body) < 12, f"expected a compact loop, got {len(body)} instructions" + + +def test_fma_uses_no_scratch_and_no_separate_multiply(): + asm = _prog().tile_row_fma_fp_sweep_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 line for line in body) + assert not any("V_MUL_VF" in line for line in body) + + +def test_fma_rejects_mismatched_row_counts(): + with pytest.raises(ValueError, match="row counts"): + _prog().tile_row_fma_fp_sweep_asm( + dst_addr=0, src_addr=8192, fpram_base=16, + dst_rows=[0, 1], src_rows=[0], + ) + with pytest.raises(ValueError, match="row counts"): + _prog().tile_row_fma_fp_broadcast_asm( + dst_addr=0, src_addr=8192, fpram_scalar_addr=16, + dst_rows=[0, 1], src_rows=[0], + ) + + +def test_broadcast_applies_one_slot_to_every_row(): + """`_broadcast` in this file means one slot for all rows -- the opposite of + `_sweep`. Mixing them up is a silent wrong answer, so pin the difference.""" + sweep = _body(_prog().tile_row_fma_fp_sweep_asm( + dst_addr=0, src_addr=8192, fpram_base=16, + dst_rows=[0, 1, 2], src_rows=[0, 1, 2], + )) + bcast = _body(_prog().tile_row_fma_fp_broadcast_asm( + dst_addr=0, src_addr=8192, fpram_scalar_addr=16, + dst_rows=[0, 1, 2], src_rows=[0, 1, 2], + )) + # The sweep advances the FPRAM pointer; the broadcast has nothing to advance. + assert len(bcast) < len(sweep) + + +def test_the_unrolled_fallback_computes_the_same_answer_as_the_loop(): + """The `else:` branch of `_emit_tile_row_fma`, which no test reached. + + It runs whenever any of the three walks is not a progression -- a scattered + row map through the public `tile_row_fma_fp_asm`, or `ATEN_OPS_UNROLL=1`, + which routes *every* emitter through it. Swapping the two pointers there, + or replacing the whole branch with a `raise`, left the entire suite green. + + Checked four shapes: each walk non-monotonic in turn, plus a repeated + destination row, which is the read-modify-write ordering case -- the FMA + accumulates, so two entries landing on the same row must compose rather + than the later winning. + """ + scattered = [ + ("dst jumps", [(3, 0, 0), (0, 1, 1), (2, 2, 2)]), + ("src jumps", [(0, 5, 0), (1, 2, 1), (2, 7, 2)]), + ("fp jumps", [(0, 0, 5), (1, 1, 0), (2, 2, 3)]), + ("dst repeats", [(1, 0, 0), (1, 3, 2), (1, 2, 1)]), + ] + for label, row_map in scattered: + p = _prog() + dst = p.alloc("dst", MLEN, MLEN) + src = p.alloc("src", MLEN, MLEN) + weights = p.fp_var("w", size=MLEN) + dst_base = p.get_vram_layout(dst.name).vram_base_addr + src_base = p.get_vram_layout(src.name).vram_base_addr + mark = len(p.get_code()) + p.tile_row_fma_fp_asm(dst_addr=dst_base, src_addr=src_base, row_map=row_map) + code = p.get_code()[mark:] + assert "C_LOOP_START" not in code, f"{label}: expected the unrolled path" + assert code.count("V_FMA_VF") == len(row_map) + + m = Machine(vlen=MLEN, vram_words=1 << 16, fpram_words=1 << 13) + expected = [] + for r in range(MLEN): + dst_vals = [float(r * MLEN + c + 1) / 8.0 for c in range(MLEN)] + src_vals = [float(r + c) / 4.0 - 1.0 for c in range(MLEN)] + m.write_vram_row(dst_base + r * MLEN, dst_vals) + m.write_vram_row(src_base + r * MLEN, src_vals) + expected.append(list(dst_vals)) + src_rows = [[float(r + c) / 4.0 - 1.0 for c in range(MLEN)] for r in range(MLEN)] + for i in range(MLEN): + m.write_fpram(weights.address + i, [0.25 + i]) + # Apply the row map in order, exactly as the emitter does. + for d, sr, f in row_map: + for c in range(MLEN): + expected[d][c] += src_rows[sr][c] * (0.25 + f) + # The emitter takes absolute FPRAM addresses in row_map, so re-emit with + # the allocated base folded in. + p2 = _prog() + d2 = p2.alloc("dst", MLEN, MLEN) + s2 = p2.alloc("src", MLEN, MLEN) + w2 = p2.fp_var("w", size=MLEN) + mark2 = len(p2.get_code()) + p2.tile_row_fma_fp_asm( + dst_addr=p2.get_vram_layout(d2.name).vram_base_addr, + src_addr=p2.get_vram_layout(s2.name).vram_base_addr, + row_map=[(d, sr, w2.address + f) for d, sr, f in row_map], + ) + m2 = Machine(vlen=MLEN, vram_words=1 << 16, fpram_words=1 << 13) + db = p2.get_vram_layout(d2.name).vram_base_addr + sb = p2.get_vram_layout(s2.name).vram_base_addr + for r in range(MLEN): + m2.write_vram_row(db + r * MLEN, [float(r * MLEN + c + 1) / 8.0 for c in range(MLEN)]) + m2.write_vram_row(sb + r * MLEN, src_rows[r]) + for i in range(MLEN): + m2.write_fpram(w2.address + i, [0.25 + i]) + m2.run(p2.get_code()[mark2:]) + for r in range(MLEN): + got = m2.read_vram_row(db + r * MLEN, MLEN) + for c in range(MLEN): + assert abs(got[c] - expected[r][c]) < 1e-4, ( + f"{label}: row {r} lane {c}: {got[c]} != {expected[r][c]}" + ) + + +def test_the_unrolled_fallback_matches_the_looped_path_exactly(): + """Same work, two code paths. A progression run through the fallback (by + scrambling and restoring the order) must land on the same numbers.""" + rows = list(range(6)) + results = [] + for row_map in ( + [(0, r, r) for r in rows], # progression -> loop + [(0, r, r) for r in [0, 1, 2, 3, 4, 5][::-1]], # reversed -> still a + ): # progression (step -1) + p = _prog() + dst = p.alloc("dst", MLEN, MLEN) + src = p.alloc("src", MLEN, MLEN) + w = p.fp_var("w", size=MLEN) + db = p.get_vram_layout(dst.name).vram_base_addr + sb = p.get_vram_layout(src.name).vram_base_addr + mark = len(p.get_code()) + p.tile_row_fma_fp_asm( + dst_addr=db, src_addr=sb, + row_map=[(d, s, w.address + f) for d, s, f in row_map], + ) + m = Machine(vlen=MLEN, vram_words=1 << 16, fpram_words=1 << 13) + m.write_vram_row(db, [0.0] * MLEN) + for r in rows: + m.write_vram_row(sb + r * MLEN, [float(r + c) / 4.0 for c in range(MLEN)]) + m.write_fpram(w.address + r, [0.5 + r]) + m.run(p.get_code()[mark:]) + results.append(m.read_vram_row(db, MLEN)) + # Addition is commutative here up to float ordering; both orders sum the + # same six products. + for a, b in zip(*results): + assert abs(a - b) < 1e-4, f"{a} != {b}" + + +def test_a_step_zero_loop_computes_the_right_answer(): + """Allowing a step of 0 changes codegen for every emitter, so prove the + resulting loop is arithmetically right and not merely shorter. + + Contracts eight source rows into one destination row, weighted by eight + FPRAM slots -- the exact shape of KDA's prediction sweep. + """ + p = _prog() + dst = p.alloc("dst", MLEN, MLEN) + src = p.alloc("src", MLEN, MLEN) + weights = p.fp_var("w", size=MLEN) + dst_base = p.get_vram_layout(dst.name).vram_base_addr + src_base = p.get_vram_layout(src.name).vram_base_addr + mark = len(p.get_code()) + p.tile_row_fma_fp_sweep_asm( + dst_addr=dst_base, src_addr=src_base, fpram_base=weights.address, + dst_rows=[0] * MLEN, src_rows=list(range(MLEN)), + ) + code = p.get_code()[mark:] + assert "C_LOOP_START" in code, "the pinned destination must still loop" + assert code.count("V_FMA_VF") == 1, "one FMA in the loop body, not MLEN of them" + + m = Machine(vlen=MLEN, vram_words=1 << 16, fpram_words=1 << 13) + m.write_vram_row(dst_base, [0.0] * MLEN) + expected = [0.0] * MLEN + for row in range(MLEN): + values = [float(row * MLEN + col + 1) for col in range(MLEN)] + m.write_vram_row(src_base + row * MLEN, values) + m.write_fpram(weights.address + row, [0.5 + row]) + for col in range(MLEN): + expected[col] += values[col] * (0.5 + row) + m.run(code) + + got = m.read_vram_row(dst_base, MLEN) + for col in range(MLEN): + assert abs(got[col] - expected[col]) < 1e-3, ( + f"lane {col}: {got[col]} != {expected[col]}" + ) diff --git a/doc/configuration.svh b/doc/configuration.svh index 1f4b6c9..a95560d 100644 --- a/doc/configuration.svh +++ b/doc/configuration.svh @@ -8,7 +8,12 @@ import precision_pkg::*; package configuration_pkg; // Compute Unit Related parameter BLEN = 4; - parameter HLEN = 8 + // NOTE: the missing `;` here made this file invalid SystemVerilog and made + // `utils/load_config.py` (whose line-wise regex requires the trailing `;`) + // silently drop HLEN from the parsed settings dict. The declared value 8 + // disagrees with doc/plena_isa_spec.md and with plena_settings.toml (both 16); + // the value is left alone because only RTL can settle which is correct. + parameter HLEN = 8; parameter MLEN = 64; parameter VLEN = 64; parameter INST_BUFF_DEPTH = 16; diff --git a/doc/operation.svh b/doc/operation.svh index 79b3651..7c95cb4 100644 --- a/doc/operation.svh +++ b/doc/operation.svh @@ -41,7 +41,8 @@ typedef enum logic [3:0] { RECI_V_ELEMENT = 4'h5, INNER_HADAMARD_TRANSFORM = 4'h6, PREFIX_SCAN_V_ELEMENT = 4'h7, - SHIFT_V_LANES_ELEMENT = 4'h8 // renamed, use this everywhere + SHIFT_V_LANES_ELEMENT = 4'h8, // renamed, use this everywhere + SOFTPLUS_V_ELEMENT = 4'h9 // log(1+exp(x)) with large-|x| passthrough } V_ELEMENT_OP; typedef enum logic [2:0] { @@ -64,7 +65,8 @@ typedef enum logic [3:0] { ST_REG_FP = 4'hA, ST_IN_FP = 4'hB, MV_FP = 4'hC, - MAP_V_FP = 4'hD + MAP_V_FP = 4'hD, + MAP_FP_V = 4'hE // inverse of MAP_V_FP: one VLEN Vector SRAM row -> VLEN FP_MEM slots } S_FP_OP; typedef enum logic [3:0] { @@ -180,7 +182,14 @@ typedef enum logic [instruction_pkg::OPCODE_WIDTH - 1:0] { V_MAX_VF = 6'h35, V_MIN_VF = 6'h36, V_TOPK = 6'h37, - C_SET_TOPK_REG = 6'h38 + C_SET_TOPK_REG = 6'h38, + + // Mamba / selective-SSM extensions + V_SOFTPLUS_V = 6'h39, + S_MAP_FP_V = 6'h3A, + // Vector[rd] += Vector[rs1] * fp_reg. The only V-type op that reads + // rd as well as writing it. + V_FMA_VF = 6'h3B } CUSTOM_ISA_OPCODE; typedef enum logic [2:0] { diff --git a/doc/plena_isa_spec.md b/doc/plena_isa_spec.md index 89d0bdd..023b78a 100644 --- a/doc/plena_isa_spec.md +++ b/doc/plena_isa_spec.md @@ -240,12 +240,18 @@ Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs **Format:** `V_SUB_VV rd, rs1, rs2, rmask` -**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) - (Vector[gp_reg] & gp_reg)` +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) - (Vector[gp_reg] & gp_reg)` **Description:** Similar to `V_ADD_VV`, but performs element-wise subtraction. +**Note:** this line previously read `rs2 - rs1`. The hardware computes +`rs1 - rs2` -- `vector_machine.rs`'s `sub(vd, vs1, vs2)` returns +`read(vs1) - read(vs2)`, and `dispatch.rs` passes `(rd, rs1, rs2)` in that +order. Subtraction does not commute, so the old wording inverted the sign of +anything written against the spec rather than against the emulator. + ### V_SUB_VF **Format:** `V_SUB_VF rd, rs1, fp2, rmask, rorder` @@ -287,6 +293,35 @@ Similar to `V_ADD_VV`, but performs element-wise multiplication. Similar to `V_ADD_VF`, but performs element-wise multiplication. +### V_FMA_VF + +**Format:** `V_FMA_VF rd, rs1, fp2, rmask` + +**Operation:** `Vector[gp_reg] & gp_reg = (Vector[gp_reg] + Vector[gp_reg] * Broadcast(fp_reg)) & gp_reg` + +**Opcode:** `6'h3B` + +**Description:** + +Fused broadcast multiply-add. **`rd` is read as well as written** -- the only +V-type instruction for which that is true, so a register allocator must treat it +as a read-modify-write and not as a pure destination. + +The multiply and the add are fused: the result is quantised **once**, on the sum. +That is a numerical difference from the `V_MUL_VF` + `V_ADD_VV` pair it replaces, +not just a shorter encoding. + +Lanes masked off by `rmask` keep the destination's existing value. Note this is +the opposite base from `V_MUL_VF`, where a masked-off lane keeps the *source*'s +value -- there the source is what the result is built from, here it is the +destination. + +Its purpose is the state recurrences. A rank-1 update or a state contraction +otherwise costs `copy + multiply + add` through a scratch row; the fused form +removes both the copy and the scratch. Removing the scratch is the larger effect: +without it a whole key sweep is a single arithmetic row progression, which the +compiler turns into a hardware loop instead of unrolling. + ### V_MAX_VF **Format:** `V_MAX_VF rd, rs1, fp2, rmask` @@ -366,6 +401,39 @@ Fetch a (VLEN, 1) vector from the Vector SRAM using the address provided by `rs1 V_RECI_V gp2, gp1, 0 ; Vector[gp2] = 1.0 / Vector[gp1] ``` +### V_SOFTPLUS_V + +**Format:** `V_SOFTPLUS_V rd, rs1, rmask` + +**Operation:** `Vector[gp_reg] = log(1 + exp(Vector[gp_reg]))` + +**Description:** + +Element-wise softplus over a (VLEN, 1) vector. + +Exists because the ISA has **no logarithm** anywhere -- `S_FP_OP` is +{ADD, SUB, MAX, MUL, EXP, RECI, SQRT, LD, ST, MV, MAP} and `V_ELEMENT_OP` is +{ADD, SUB, MUL, EXP, RECI, HADAMARD, PREFIX_SCAN, SHIFT} -- so `log(1+exp(x))` +has no exact software lowering. A lookup table is architecturally impossible too: +there is no float-to-int conversion instruction, so FP_MEM cannot be indexed by a +data-dependent value. + +Mamba / Mamba-2 needs `dt = softplus(dt_raw + dt_bias)` on the critical path of a +multiplicative recurrence, where an approximation's error compounds across the +sequence rather than being normalised away the way attention's softmax normalises +`exp` error. + +**Implementation note:** evaluate as `relu(x) + log1p(exp(-|x|))`, not as the +naive `log1p(exp(x))`. The two are algebraically identical, but the first never +feeds `exp` a positive argument and therefore cannot overflow for any finite +input -- no clamp is needed, and none should be applied, since `dt` feeds +`exp(A*dt)` and a clamp would silently flatten the decay for large `dt`. + +**Example:** +```asm +V_SOFTPLUS_V gp2, gp1, 0 ; Vector[gp2] = softplus(Vector[gp1]) +``` + ### V_RED_SUM **Format:** `V_RED_SUM rd, rs1` @@ -551,6 +619,39 @@ Load upper immediate value into the integer register. Copy a vector of length VLEN from FP_MEM to Vector SRAM. +#### S_MAP_FP_V + +**Format:** `S_MAP_FP_V rd, rs1, imm` + +**Operation:** `FP_MEM[gp_reg + imm :+ VLEN] = Vector[gp_reg :+ VLEN]` + +**Description:** + +The exact inverse of `S_MAP_V_FP`: copy one VLEN-wide Vector SRAM row into VLEN +consecutive FP_MEM slots. Note the operand roles mirror rather than repeat -- +`rd` is the FP_MEM base and `rs1` the VRAM row, so that in both instructions `rd` +names the *destination* memory. + +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 -- 3 instructions per *scalar* +against 1 per *row* here. + +That mattered because `V_MUL_VF` can only broadcast from an FP register, while +every value a kernel wants to broadcast is computed on chip and lands in a VRAM +row. Mamba-2's chunked scan needs one such scalar per row (`cs_i` for the +segment-sum matrix, `exp(cs_C - cs_t)` for the decay-scaled X, `exp(cs_i)` for +the inter-chunk term), which made the extraction, not the matmuls, the dominant +instruction cost of the kernel. + +**Cost:** VLEN cycles, the same as `S_MAP_V_FP`. + +**Example:** +```asm +S_MAP_FP_V gp3, gp1, 0 ; FP_MEM[gp3 :+ VLEN] = Vector[gp1 :+ VLEN] +S_LD_FP f1, gp3, 5 ; then broadcast lane 5 with V_MUL_VF +``` + --- ## Memory (H-Type) Instructions From ce1bda8ab62cde2f778f264f4fd199bde3a8b199 Mon Sep 17 00:00:00 2001 From: "Qichao (Arlo) Wang" Date: Sat, 29 Aug 2026 01:39:48 +0100 Subject: [PATCH 2/6] compiler: read the machine's configuration, address the state path, and count what it emits Three pieces of substrate the recurrent lowering needs, and one bug it exposed. **FPRAM windows and tile addressing.** The recurrences want per-scalar operands -- a decay per key channel, a beta per head -- and FPRAM is 512 slots against Kimi K3's 12,288 head-scalars, so the layer streams one head at a time and reuses one window. `tile_row_to_fpram` addresses FPRAM by *position* in the row list rather than by row index, which is what makes a reused window addressable at offset 0 for every head. The `tile_row_*` family gains a `tile_col_idx`, without which a tile wider than `mlen` had its second column block silently skipped. **`instruction_stream.py` counts an emitted program three ways.** `static_count` is the program image. `dynamic_count` expands every `C_LOOP_START` by its trip count -- they differ by whatever the hardware loops multiply, and only the second prices anything; three claims on this branch were made from image counts and were wrong by 13x, 7x and 27x. `opcode_census` then says what those issued instructions *are*, which is how a quarter-of-the-layer-computes result becomes visible at all, and `self_advance_counts` splits in-loop pointer advances by whether they have a consumer to fold into, so an addressing-mode proposal can be priced against instructions that would really disappear. `dynamic_count`'s docstring claimed the `C_LOOP_END` branch issues once per trip. The code never counted it, and treating the loop boundary as zero-overhead is the right model for a trip-count loop construct -- but the two should not have disagreed. Counting it would add about 20% to these kernels and leaves every ratio unchanged. **The configuration had never been read.** `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`; 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 sufficed; together they made a silent failure indistinguishable from a working default. The reader now goes through `tomllib`, stdlib since 3.11, 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 `[MODE].active` is deliberate: `active` selects the machine the *simulator* models, not the one the compiler emits for, and they differ today. The shipped file is `analytic` (MLEN 2048) while every program here is compiled at 64 and executed by the transactional emulator, whose section declares 64. `mram_tile_capacity` consequently defaults to `MATRIX_SRAM_SIZE // mlen` -- the same arithmetic on the same two numbers as the emulator's `MatrixSram::new` -- instead of an unconditional 4. At mlen 64 that is 64, worth 2.79x on Kimi K3's input projection (13,876,267 issued against 4,967,367). At an mlen no section declares it stays 4. At mlen 2048 the ANALYTIC config gives 256/2048 = 0 tiles and the constructor raises naming both figures: a contraction split into zero-tile chunks does not terminate, and there is no substitute anyone chose. HLEN moves 64 -> 16 and BROADCAST_AMOUNT 1 -> 4 at mlen 64. Neither reaches an emitted program -- `hlen` is read only by a packed-attention precondition whose one test sets both by hand at mlen 256, `broadcast_amount` is read nowhere outside the constructor -- and the prefetch and writeback amounts, which do reach the emitters, are 4 in that section and were already 4. Co-Authored-By: Claude Opus 5 --- aten/plena/compiler.py | 155 +++++++++++++++++++--- aten/plena/instruction_stream.py | 211 ++++++++++++++++++++++++++++++ aten/plena/isa_compiler.py | 13 ++ aten/plena/memory_state.py | 16 ++- aten/plena/program_fp_tile_ops.py | 181 ++++++++++++++++++++++++- aten/plena/program_tensors.py | 15 ++- utils/load_config.py | 29 +++- 7 files changed, 587 insertions(+), 33 deletions(-) create mode 100644 aten/plena/instruction_stream.py diff --git a/aten/plena/compiler.py b/aten/plena/compiler.py index d63c6c8..2603ace 100644 --- a/aten/plena/compiler.py +++ b/aten/plena/compiler.py @@ -10,7 +10,18 @@ from compiler.aten.plena.program_fp_tile_ops import ProgramFPTileOpsMixin from compiler.aten.plena.program_moe_shared import ProgramMoeSharedMixin from compiler.aten.plena.program_routed_moe import ProgramRoutedMoeMixin +from compiler.aten.plena.program_mamba_common import ProgramMambaCommonMixin from compiler.aten.plena.program_matrix_ops import ProgramMatrixOpsMixin +from compiler.aten.plena.program_ssd import ProgramSSDMixin +from compiler.aten.plena.program_ssm_recurrent import ProgramSSMRecurrentMixin +from compiler.aten.plena.program_kda_common import ProgramKdaCommonMixin +from compiler.aten.plena.program_kda_chunk import ProgramKdaChunkMixin +from compiler.aten.plena.program_kda_conv import ProgramKdaConvMixin +from compiler.aten.plena.program_kda_gates import ProgramKdaGatesMixin +from compiler.aten.plena.program_kda_layer import ProgramKdaLayerMixin +from compiler.aten.plena.program_kda_prefill import ProgramKdaPrefillMixin +from compiler.aten.plena.program_kda_mixer import ProgramKdaMixerMixin +from compiler.aten.plena.program_kda_recurrent import ProgramKdaRecurrentMixin from compiler.aten.plena.program_tensors import ProgramTensorMixin from compiler.aten.plena.vars import FPVar, InputVar, TensorVar from compiler.utils.load_config import load_toml_config @@ -29,24 +40,100 @@ def _find_plena_settings_toml() -> Path | None: return None -def _behavior_config_value(key: str, default: int) -> int: +#: Config sections `plena_settings.toml` actually ships, in the order they are +#: searched after `BEHAVIOR`. +_MACHINE_MODES = ("TRANSACTIONAL", "ANALYTIC") + + +def _config_section(mode: str) -> dict: settings_path = _find_plena_settings_toml() if settings_path is None or not settings_path.exists(): - return default - + return {} try: - config = load_toml_config(settings_path, "CONFIG", mode="BEHAVIOR") + return load_toml_config(settings_path, "CONFIG", mode=mode) or {} except Exception: - return default + return {} - value = config.get(key, {}) + +def _section_value(section: dict, key: str) -> int | None: + value = section.get(key) if isinstance(value, dict): - value = value.get("value", default) + value = value.get("value") try: - parsed = int(value) + parsed = int(value) # type: ignore[arg-type] except (TypeError, ValueError): - return default - return parsed if parsed > 0 else default + return None + return parsed if parsed > 0 else None + + +def _config_mode_for_mlen(mlen: int) -> str | None: + """The shipped section whose ``MLEN`` is the one being compiled for. + + ``plena_settings.toml`` has no ``[BEHAVIOR]`` table -- it ships ``[MODE]``, + ``[ANALYTIC.*]`` and ``[TRANSACTIONAL.*]`` -- so a reader that asks for + ``BEHAVIOR`` finds nothing and every value silently falls back to whatever + default the caller passed. That is what this function replaces. + + Matching on ``MLEN`` rather than on ``[MODE].active`` is deliberate. + ``active`` selects which machine the *simulator* models; it does not say + what the compiler is emitting for, and the two are routinely different -- + the shipped file has ``active = "analytic"`` (MLEN 2048) while every + program on this branch is compiled at ``mlen`` 64 and executed by the + transactional emulator, whose section declares MLEN 64. Reading ``active`` + would hand a 64-wide compilation the 2048-wide machine's numbers. + + Returns ``None`` when no section matches, which is the case for the test + shapes (``mlen`` 8) and means "keep the caller's default". + """ + for mode in _MACHINE_MODES: + if _section_value(_config_section(mode), "MLEN") == mlen: + return mode + return None + + +def _behavior_config_value(key: str, default: int, mlen: int | None = None) -> int: + """``key`` from ``[BEHAVIOR]``, else from the section matching ``mlen``. + + ``BEHAVIOR`` is searched first so that a file which does define it keeps + overriding the machine sections; nothing ships one today. + """ + parsed = _section_value(_config_section("BEHAVIOR"), key) + if parsed is not None: + return parsed + if mlen is not None: + mode = _config_mode_for_mlen(mlen) + if mode is not None: + parsed = _section_value(_config_section(mode), key) + if parsed is not None: + return parsed + return default + + +def _derive_mram_tile_capacity(mlen: int) -> int | None: + """``MATRIX_SRAM_SIZE // mlen`` for the machine ``mlen`` names. + + The emulator's ``MatrixSram::new(tile_size = MLEN, depth = + MATRIX_SRAM_SIZE)`` keeps ``depth / tile_size`` tiles, so this is the same + arithmetic on the same two numbers. Returns ``None`` when no section + matches ``mlen``, and raises when one does but cannot hold a single tile -- + a compiler that split the contraction into zero-tile chunks would emit an + infinite program, so there is no sensible fallback. + """ + mode = _config_mode_for_mlen(mlen) + if mode is None: + return None + size = _section_value(_config_section(mode), "MATRIX_SRAM_SIZE") + if size is None: + return None + tiles = size // mlen + if tiles < 1: + raise ValueError( + f"[{mode}] MATRIX_SRAM_SIZE {size} holds {size / mlen:.3f} tiles of " + f"mlen {mlen}; a projection cannot be split into fewer than one " + f"tile per chunk. Pass mram_tile_capacity explicitly to override, " + f"or fix the configuration." + ) + return tiles # ============================================================================ @@ -63,6 +150,32 @@ class PlenaCompiler( # substrate (moe_expert_activation_v0, moe_materialize_route_weights_*). ProgramMoeSharedMixin, ProgramAttentionMixin, + # After ProgramMatrixOpsMixin and ProgramFPTileOpsMixin: the Mamba emitters + # are built on vram_add/vram_mul/vram_fill_zero and the tile_row_* family. + # ProgramSSDMixin and ProgramSSMRecurrentMixin must come after + # ProgramMambaCommonMixin -- both call its mamba_row_* / mamba_block_copy + # primitives and its stage-marker vocabulary. + ProgramMambaCommonMixin, + ProgramSSDMixin, + ProgramSSMRecurrentMixin, + # After ProgramSSMRecurrentMixin and ProgramMambaCommonMixin: the KDA + # emitters reuse pin_hbm_region for state residency and mamba_row_copy / + # mamba_block_copy / mamba_rsqrt_fpram for the row primitives. KDA is a + # second consumer of that substrate, not a second copy of it. + ProgramKdaCommonMixin, + # After ProgramKdaCommonMixin: the recurrence uses its layout helpers and + # stage vocabulary. + ProgramKdaRecurrentMixin, + ProgramKdaConvMixin, + ProgramKdaGatesMixin, + # After ProgramKdaGatesMixin: the chunk primitives use its key-block helpers. + ProgramKdaChunkMixin, + # Last of the chunk path: it calls the chunk primitives and the matrix ops. + ProgramKdaPrefillMixin, + # Assembles the layer around the mixer; needs the projections. + ProgramKdaLayerMixin, + # Last of the KDA mixins: it calls all the others. + ProgramKdaMixerMixin, IsaCompiler, ): """ @@ -78,7 +191,7 @@ def __init__( blen: int = 4, real_data_ratio: float = 1.125, unroll_loops: bool = False, - mram_tile_capacity: int = 4, + mram_tile_capacity: int | None = None, hbm_v_prefetch_amount: int | None = None, hbm_v_writeback_amount: int | None = None, ): @@ -88,6 +201,14 @@ def __init__( blen: Vector tile size (default 4) real_data_ratio: HBM storage ratio (MXFP8 format = 1.125) mram_tile_capacity: Number of mlen x mlen tiles that fit in MRAM. + Defaults to MATRIX_SRAM_SIZE // mlen from the config + section whose MLEN is `mlen` -- the same arithmetic + the emulator's MatrixSram does -- and to 4 when no + section matches. It was an unconditional 4 until + 2026-08-28, which is 16x below the transactional + machine's 64 and cost Kimi K3's input projection a + factor of 2.79; see + `test_the_mram_tile_default_costs_the_projection_a_factor_of_three`. hbm_v_prefetch_amount: H_PREFETCH_V transfer count. Defaults to BEHAVIOR.CONFIG.HBM_V_Prefetch_Amount in PLENA_SETTINGS_TOML / plena_settings.toml. @@ -98,6 +219,8 @@ def __init__( at ASM-gen time to eliminate C_LOOP_START/END overhead. Overridden by the ATEN_OPS_UNROLL env var ("1"=True, "0"=False). """ + if mram_tile_capacity is None: + mram_tile_capacity = _derive_mram_tile_capacity(mlen) or 4 _env_unroll = os.environ.get("ATEN_OPS_UNROLL", "") if _env_unroll == "1": unroll_loops = True @@ -111,17 +234,19 @@ def __init__( mram_tile_capacity=mram_tile_capacity, ) if hbm_v_prefetch_amount is None: - hbm_v_prefetch_amount = _behavior_config_value("HBM_V_Prefetch_Amount", 4) + hbm_v_prefetch_amount = _behavior_config_value("HBM_V_Prefetch_Amount", 4, mlen) if hbm_v_writeback_amount is None: - hbm_v_writeback_amount = _behavior_config_value("HBM_V_Writeback_Amount", 4) + hbm_v_writeback_amount = _behavior_config_value("HBM_V_Writeback_Amount", 4, mlen) if hbm_v_prefetch_amount <= 0: raise ValueError(f"hbm_v_prefetch_amount must be > 0, got {hbm_v_prefetch_amount}") if hbm_v_writeback_amount <= 0: raise ValueError(f"hbm_v_writeback_amount must be > 0, got {hbm_v_writeback_amount}") self.hbm_v_prefetch_amount = hbm_v_prefetch_amount self.hbm_v_writeback_amount = hbm_v_writeback_amount - self.hlen = _behavior_config_value("HLEN", mlen) - self.broadcast_amount = _behavior_config_value("BROADCAST_AMOUNT", max(1, mlen // max(1, self.hlen))) + self.hlen = _behavior_config_value("HLEN", mlen, mlen) + self.broadcast_amount = _behavior_config_value( + "BROADCAST_AMOUNT", max(1, mlen // max(1, self.hlen)), mlen + ) # HBM address auto-allocation self._next_hbm_addr: int = 0 diff --git a/aten/plena/instruction_stream.py b/aten/plena/instruction_stream.py new file mode 100644 index 0000000..5ace698 --- /dev/null +++ b/aten/plena/instruction_stream.py @@ -0,0 +1,211 @@ +"""Counting an emitted program two ways: as an image, and as an issue stream. + +These are different measurements and the difference is not small. A sweep that +the `V_FMA_VF` conversion collapsed into a hardware loop occupies a handful of +words in the program image and issues its body once per trip -- 192 times for a +Kimi K3 projection section, 128 times for a `key_dim` 128 recurrence. + +Which one to use follows from the question: + +* **Image** (:func:`static_count`) answers "how large is the binary" and + nothing else. It is the right instrument for the budget gates in + `test_instruction_budget.py`, whose whole purpose is to catch a sweep falling + off the hardware-loop path and unrolling -- an image failure that no numeric + test can see. +* **Issue stream** (:func:`dynamic_count`) answers "how much work is there". + Every claim about cost, about which kernel dominates, or about what an + alternative lowering would save has to be made against this one. + +Conflating them is not a hypothetical. Three claims on this branch were made +from image counts and were wrong by 13x, 7x and 27x respectively: the +projection gather's share of a layer, the conv-versus-mixer ordering, and the +`V_FMA_VF` conversion's effect on time. +""" + +from __future__ import annotations + +import re + +__all__ = [ + "static_count", + "dynamic_count", + "opcode_census", + "arithmetic_share", + "self_advance_counts", + "ARITHMETIC", +] + + +def _instructions(asm: str) -> list[str]: + return [ + line.strip() + for line in asm.splitlines() + if line.strip() and not line.strip().startswith(";") + ] + + +def static_count(asm: str) -> int: + """Instructions in the program image: one per emitted line, comments aside.""" + return len(_instructions(asm)) + + +def dynamic_count(asm: str) -> int: + """Instructions issued, with every ``C_LOOP_START`` expanded by its trip count. + + Convention: ``C_LOOP_START`` issues once, the body issues once per trip, + and ``C_LOOP_END`` is **not counted at all** -- it is treated as a + zero-overhead loop boundary that the sequencer resolves without an issue + slot, which is what a hardware loop construct with a trip-count register + normally is. Nesting is handled. + + Counting the loop end instead would add about 20% to the recurrent kernels + here. It moves any numerator and denominator together, so the ratios this + is used for do not depend on the choice -- but the absolute figures do, and + they are reported under this convention. (An earlier version of this + docstring claimed the loop end *was* counted, which the code never did.) + + Trip counts are the immediate in the ``C_LOOP_START`` word, which is what + the emitters put there; a loop whose count came from a register at runtime + would not be countable this way, and the static path does not have one. + """ + lines = _instructions(asm) + + def walk(i: int) -> tuple[int, int]: + total = 0 + while i < len(lines): + op = lines[i].split()[0].rstrip(",") + if op == "C_LOOP_START": + trips = int(re.findall(r"(-?\d+)", lines[i])[-1]) + body, i = walk(i + 1) + total += 1 + body * trips + continue + if op == "C_LOOP_END": + return total, i + 1 + total += 1 + i += 1 + return total, i + + return walk(0)[0] + + +#: Opcodes that compute. Everything else is scaffolding: pointer arithmetic, +#: loop control, and moving scalars between FPRAM and the register file. +ARITHMETIC = frozenset({ + "V_FMA_VF", "V_MUL_VV", "V_ADD_VV", "V_SUB_VV", "V_MUL_VF", "V_ADD_VF", + "V_SUB_VF", "V_MAX_VF", "V_EXP_V", "V_RECI_V", "V_SOFTPLUS_V", "V_RED_SUM", + "V_RED_MAX", "V_MOV_VF", "V_SHIFT_V", "V_CLR_V", + "M_MM", "M_TMM", "M_MM_WO", "M_BTMM", + "S_ADD_FP", "S_MUL_FP", "S_RECI_FP", "S_SQRT_FP", "S_EXP_FP", +}) + +_SELF_ADVANCE = re.compile(r"^S_ADDI_INT\s+gp(\d+)\s*,\s*gp(\d+)\s*,\s*(-?\d+)") +_LOOP_START = re.compile(r"^C_LOOP_START\s+gp(\d+)\s*,\s*(\d+)") +_GP = re.compile(r"\bgp(\d+)\b") + + +def opcode_census(asm: str) -> dict[str, int]: + """Issued instructions per opcode, hardware loops expanded. + + Same expansion as :func:`dynamic_count`, but keyed by mnemonic, which is + what turns "the kernel issues N instructions" into "and here is what they + are". The recurrent kernels spend three quarters of their issue slots on + scaffolding, and that is only visible per opcode. + """ + lines = _instructions(asm) + counts: dict[str, int] = {} + + def walk(i: int) -> tuple[dict[str, int], int]: + local: dict[str, int] = {} + while i < len(lines): + op = lines[i].split()[0].rstrip(",") + if op == "C_LOOP_START": + trips = int(re.findall(r"(-?\d+)", lines[i])[-1]) + body, i = walk(i + 1) + local[op] = local.get(op, 0) + 1 + for k, v in body.items(): + local[k] = local.get(k, 0) + v * trips + continue + if op == "C_LOOP_END": + return local, i + 1 + local[op] = local.get(op, 0) + 1 + i += 1 + return local, i + + counts, _ = walk(0) + return counts + + +def arithmetic_share(asm: str) -> float: + """Fraction of issued instructions that compute something.""" + census = opcode_census(asm) + total = sum(census.values()) + if not total: + return 0.0 + return sum(n for op, n in census.items() if op in ARITHMETIC) / total + + +def _loop_bodies(lines: list[str]) -> list[tuple[list[str], int]]: + """Every loop body paired with how many times it runs, nesting included.""" + out: list[tuple[list[str], int]] = [] + + def scan(block: list[str], trips: int) -> None: + i = 0 + while i < len(block): + m = _LOOP_START.match(block[i]) + if not m: + i += 1 + continue + reg, count = m.group(1), int(m.group(2)) + j, depth = i + 1, 1 + while j < len(block) and depth: + if _LOOP_START.match(block[j]): + depth += 1 + elif block[j].startswith(f"C_LOOP_END gp{reg}") and depth == 1: + break + elif block[j].startswith("C_LOOP_END"): + depth -= 1 + j += 1 + body = block[i + 1 : j] + out.append((body, trips * count)) + scan(body, trips * count) + i = j + 1 + + scan(lines, 1) + return out + + +def self_advance_counts(asm: str) -> tuple[int, int]: + """``(foldable, unfoldable)`` issued pointer self-advances inside loops. + + A sweep advances its pointers with ``S_ADDI_INT gpN, gpN, step`` once per + trip. A post-increment addressing mode -- the operand carrying its own + stride, the way ``_emit_tile_row_fma`` already carries independent row + progressions -- would fold those into the instruction that consumes the + pointer, and they would stop occupying an issue slot. + + "Would" only holds where the register actually has a consumer in the same + body; an advance with nothing to fold into has to stay. This reports the + split so the saving is priced against instructions that can really + disappear rather than against every advance that happens to sit in a loop. + + Loop *setup* (``S_ADDI_INT gpN, gp0, addr``) is excluded: it runs once and + survives any addressing mode. + """ + lines = _instructions(asm) + foldable = unfoldable = 0 + for body, trips in _loop_bodies(lines): + consumers: dict[str, int] = {} + for line in body: + if _SELF_ADVANCE.match(line) or line.startswith("C_LOOP"): + continue + for reg in _GP.findall(line): + consumers[reg] = consumers.get(reg, 0) + 1 + for line in body: + m = _SELF_ADVANCE.match(line) + if not m or m.group(1) != m.group(2): + continue + if consumers.get(m.group(1), 0): + foldable += trips + else: + unfoldable += trips + return foldable, unfoldable diff --git a/aten/plena/isa_compiler.py b/aten/plena/isa_compiler.py index 7d1402d..f756b4c 100644 --- a/aten/plena/isa_compiler.py +++ b/aten/plena/isa_compiler.py @@ -62,6 +62,8 @@ def load_batch( vram_object_name: str, vlen: int = 64, preload_len: int | None = None, + storage_precision: int = 1, + precision: int = 0, ) -> str: """ Load a Batch tensor from HBM to VRAM. @@ -70,6 +72,15 @@ def load_batch( logical size * real_data_ratio = 1.125. VRAM stores only the vector data (no scale), so VRAM size = logical size. + `storage_precision` is HBM bytes per element and `precision` is the + H_PREFETCH_V precision-class selector (0 = Activation, 1 = KeyValue). Both + default to the MX-FP8 activation path, which is what every existing caller + wants. They exist because `store_to_hbm` already takes the matching pair, + and a load that cannot express what the store wrote is not an inverse: the + cross-token SSM state is written KeyValue/2-byte and would otherwise be read + back as Activation/1-byte, at half the row stride and against a different + scale-section base -- garbage, silently. + Order (matters): allocate VRAM → register in symbol table → emit ISA. """ if preload_len is None: @@ -115,6 +126,8 @@ def load_batch( act_vram_offset=vram_base, activation_offset_reg=addr_reg, stride_size=w, + storage_precision=storage_precision, + precision=precision, ) self.register_allocator.free_gp(gp_regs_for_addr) diff --git a/aten/plena/memory_state.py b/aten/plena/memory_state.py index ce35cac..db99bc9 100644 --- a/aten/plena/memory_state.py +++ b/aten/plena/memory_state.py @@ -214,14 +214,26 @@ def free_fpram_object(self, name: str, strict: bool = True) -> MemoryObjectInfo self.fpram_matrices.pop(name, None) return info - def hbm_tensor_size(self, num_elements: int) -> int: + def hbm_tensor_size(self, num_elements: int, hbm_element_bytes: int | None = None) -> int: """Row-aligned HBM byte footprint of one tensor: element rows + scale rows, each padded up to hbm_row_width, matching the stager (create_mem_for_sim / plena_utils.calculate_instr_storage_offset_from_shapes). Replaces the packed `int(size * real_data_ratio)` estimate, which ignored per-tensor row padding and mislaid every following tensor whenever a region was not naturally - row-aligned (weights/K/V read zeros; common at small MLEN).""" + row-aligned (weights/K/V read zeros; common at small MLEN). + + Pass `hbm_element_bytes` for a Plain (non-MX) tensor -- a BF16 write-back, + for instance -- so the region is sized from the width actually written + rather than from the MX layout.""" row_bits = self.hbm_row_width + if hbm_element_bytes is not None and hbm_element_bytes * 8 != self.hbm_element_width: + # A Plain (non-MX) tensor: no scale stream, and a wider element than the + # MX layout assumes. Sizing such a tensor with the MX formula under-reserves + # it -- a BF16 write-back needs 2 bytes per element against the MX layout's + # 1.125 -- and the overrun silently lands on whatever tensor follows. + bytes_per_row = row_bits // 8 + total_bytes = num_elements * hbm_element_bytes + return -(-total_bytes // bytes_per_row) * bytes_per_row ew, bs, sw = self.hbm_element_width, self.hbm_block_size, self.hbm_scale_width bytes_per_row = row_bits // 8 elements_per_row = (row_bits // (ew * bs)) * bs diff --git a/aten/plena/program_fp_tile_ops.py b/aten/plena/program_fp_tile_ops.py index 61b04aa..80c3a7c 100644 --- a/aten/plena/program_fp_tile_ops.py +++ b/aten/plena/program_fp_tile_ops.py @@ -174,6 +174,36 @@ def tile_row_reci( ): super().tile_row_reci(source.name, self._default_rows(rows), tile_col_idx=tile_col_idx) + def tile_row_softplus( + self, + source: VRAMMatrixVar, + rows: Iterable[int] | None = None, + tile_col_idx: int = 0, + ): + """In-place `x = log(1 + exp(x))` over whole VRAM rows (V_SOFTPLUS_V).""" + super().tile_row_softplus(source.name, self._default_rows(rows), tile_col_idx=tile_col_idx) + + def tile_row_to_fpram( + self, + source: VRAMMatrixVar, + target_fpram_addr: int | FPVar, + rows: Iterable[int] | None = None, + target_base_offset: int = 0, + tile_col_idx: int = 0, + ): + """Copy whole VRAM rows into FPRAM (S_MAP_FP_V), MLEN scalars per row. + + Unlike `tile_row_sum`/`tile_row_max`, which reduce a row to ONE FPRAM slot, + this moves the row across intact. Row `rows[i]` lands at + `target_fpram_addr + target_base_offset + i * MLEN`. + """ + resolved_rows = self._default_rows(rows) + row_map = [ + (row, self._resolve_fpram_addr(target_fpram_addr, target_base_offset + i * self.mlen)) + for i, row in enumerate(resolved_rows) + ] + return super().tile_row_to_fpram(source.name, row_map, tile_col_idx=tile_col_idx) + def tile_row_sub_fp( self, source: VRAMMatrixVar, @@ -240,8 +270,28 @@ def tile_row_add_fp( resolved_rows = self._default_rows(rows) super().tile_row_add_fp(source.name, [(row, fp_var[row]) for row in resolved_rows], tile_col_idx=tile_col_idx) - def _tile_row_binary(self, isa_method: str, dst: VRAMMatrixVar, src: VRAMMatrixVar, rows: Iterable[int] | None): - return getattr(super(), isa_method)(dst.name, src.name, self._default_rows(rows)) + def _tile_row_binary( + self, + isa_method: str, + dst: VRAMMatrixVar, + src: VRAMMatrixVar, + rows: Iterable[int] | None, + tile_col_idx: int = 0, + ): + # One index for both operands: a binary row op is elementwise, so the + # two tiles are the same width by construction and a column block of one + # only ever pairs with the same block of the other. `vram_fill_zero` + # walks the blocks itself because it takes a single matrix; here the + # caller drives the loop, because which blocks are live differs per + # kernel (prefill's `contrib` is wider than the range any one of its + # three uses touches). + return getattr(super(), isa_method)( + dst.name, + src.name, + self._default_rows(rows), + dst_tile_col_idx=tile_col_idx, + src_tile_col_idx=tile_col_idx, + ) def _tile_row_fp_scalar( self, @@ -271,24 +321,27 @@ def tile_row_add( dst: VRAMMatrixVar, src: VRAMMatrixVar, rows: Iterable[int] | None = None, + tile_col_idx: int = 0, ): - return self._tile_row_binary("tile_row_add", dst, src, rows) + return self._tile_row_binary("tile_row_add", dst, src, rows, tile_col_idx) def tile_row_sub( self, dst: VRAMMatrixVar, src: VRAMMatrixVar, rows: Iterable[int] | None = None, + tile_col_idx: int = 0, ): - return self._tile_row_binary("tile_row_sub", dst, src, rows) + return self._tile_row_binary("tile_row_sub", dst, src, rows, tile_col_idx) def tile_row_mul( self, dst: VRAMMatrixVar, src: VRAMMatrixVar, rows: Iterable[int] | None = None, + tile_col_idx: int = 0, ): - return self._tile_row_binary("tile_row_mul", dst, src, rows) + return self._tile_row_binary("tile_row_mul", dst, src, rows, tile_col_idx) def tile_row_mul_fp_broadcast( self, @@ -297,9 +350,125 @@ def tile_row_mul_fp_broadcast( row_idx: int | None = None, rows: Iterable[int] | None = None, fpram_offset: int = 0, + tile_col_idx: int = 0, ): scalar_addr = self._resolve_fpram_addr(fpram_scalar_addr, fpram_offset) - super().tile_row_mul_fp_broadcast(source.name, scalar_addr, self._resolve_rows(row_idx=row_idx, rows=rows)) + super().tile_row_mul_fp_broadcast( + source.name, + scalar_addr, + self._resolve_rows(row_idx=row_idx, rows=rows), + tile_col_idx=tile_col_idx, + ) + + def tile_row_fma_fp_sweep( + self, + dst: VRAMMatrixVar, + src: VRAMMatrixVar, + fpram_base: int | FPVar, + dst_rows: Iterable[int], + src_rows: Iterable[int], + fpram_offset: int = 0, + ): + """``dst[d] += src[s] * FPRAM[base + i]``, one slot per row pair. + + The FMA reads ``dst`` as well as writing it, so unlike every other tile + helper here the destination must already hold what is being accumulated + into -- typically a zeroed accumulator row or the state itself. + """ + dst_rows, src_rows = list(dst_rows), list(src_rows) + base = self._resolve_fpram_addr(fpram_base, fpram_offset) + # The predecessor resolved -- and bounds-checked -- one slot per + # iteration. This resolves the base once and lets the hardware walk the + # rest, so an over-long sweep would read into whatever FPVar was + # allocated next with nothing to say so. Check the far end too. + if isinstance(fpram_base, FPVar) and dst_rows: + end = fpram_offset + len(dst_rows) - 1 + if end >= fpram_base.size: + raise ValueError( + f"FPRAM sweep of {len(dst_rows)} rows from offset " + f"{fpram_offset} reads slot {end}, past {fpram_base.name}'s " + f"{fpram_base.size}" + ) + super().tile_row_fma_fp_sweep(dst.name, src.name, base, dst_rows, src_rows) + + def tile_row_fma_fp_broadcast( + self, + dst: VRAMMatrixVar, + src: VRAMMatrixVar, + fpram_scalar_addr: int | FPVar, + dst_rows: Iterable[int], + src_rows: Iterable[int], + fpram_offset: int = 0, + ): + """``dst[d] += src[s] * FPRAM[addr]``, one slot for every row pair.""" + addr = self._resolve_fpram_addr(fpram_scalar_addr, fpram_offset) + super().tile_row_fma_fp_broadcast( + dst.name, src.name, addr, list(dst_rows), list(src_rows) + ) + + # One FPRAM slot applied to EVERY listed row, as opposed to the + # `tile_row_*_fp` family which walks a different slot per row. + def tile_row_add_fp_broadcast( + self, + source: VRAMMatrixVar, + fpram_scalar_addr: int | FPVar, + rows: Iterable[int] | None = None, + fpram_offset: int = 0, + tile_col_idx: int = 0, + ): + scalar_addr = self._resolve_fpram_addr(fpram_scalar_addr, fpram_offset) + super().tile_row_add_fp_broadcast( + source.name, scalar_addr, self._default_rows(rows), tile_col_idx=tile_col_idx + ) + + def tile_row_max_fp_broadcast( + self, + source: VRAMMatrixVar, + fpram_scalar_addr: int | FPVar, + rows: Iterable[int] | None = None, + fpram_offset: int = 0, + tile_col_idx: int = 0, + ): + scalar_addr = self._resolve_fpram_addr(fpram_scalar_addr, fpram_offset) + super().tile_row_max_fp_broadcast( + source.name, scalar_addr, self._default_rows(rows), tile_col_idx=tile_col_idx + ) + + def tile_row_min_fp_broadcast( + self, + source: VRAMMatrixVar, + fpram_scalar_addr: int | FPVar, + rows: Iterable[int] | None = None, + fpram_offset: int = 0, + tile_col_idx: int = 0, + ): + scalar_addr = self._resolve_fpram_addr(fpram_scalar_addr, fpram_offset) + super().tile_row_min_fp_broadcast( + source.name, scalar_addr, self._default_rows(rows), tile_col_idx=tile_col_idx + ) + + def tile_row_sub_fp_broadcast( + self, + source: VRAMMatrixVar, + fpram_scalar_addr: int | FPVar, + rows: Iterable[int] | None = None, + fpram_offset: int = 0, + reverse: bool = False, + tile_col_idx: int = 0, + ): + """``row -= c``, or ``row = c - row`` when ``reverse``. + + The reverse form is `V_SUB_VF`'s `rorder=1`. Mamba's decay matrix builds + ``cs_i - cs_j`` from a row of ``cs_j`` and a scalar ``cs_i`` with it. + """ + scalar_addr = self._resolve_fpram_addr(fpram_scalar_addr, fpram_offset) + super().tile_row_sub_fp_broadcast( + source.name, + scalar_addr, + self._default_rows(rows), + reverse=reverse, + tile_col_idx=tile_col_idx, + ) # ======================================================================== # FPVar operations diff --git a/aten/plena/program_tensors.py b/aten/plena/program_tensors.py index c2b6099..bb43c31 100644 --- a/aten/plena/program_tensors.py +++ b/aten/plena/program_tensors.py @@ -73,6 +73,8 @@ def load_batch( self, input_var: InputVar, name: str | None = None, + storage_precision: int = 1, + precision: int = 0, ) -> VRAMMatrixVar: """ Load tensor from HBM to VRAM (Batch type). @@ -119,6 +121,8 @@ def load_batch( vram_object_name=internal_name, vlen=self.mlen, preload_len=self.hbm_v_prefetch_amount, + storage_precision=storage_precision, + precision=precision, ) var = VRAMMatrixVar( @@ -156,14 +160,13 @@ def store( display_name = name if name is not None else f"{tensor_var.display_name}_stored" internal_name = self._scoped_name(display_name) + # Size the region from the width actually being written. store() previously + # used the MX layout unconditionally, so a BF16 write-back (2 bytes/element) + # was allocated 1.125 and overran its region by ~78% onto the next tensor. + h, w = tensor_var.physical_shape + hbm_size = self.hbm_tensor_size(h * w, hbm_element_bytes=hbm_element_bytes) if hbm_addr is None: - h, w = tensor_var.physical_shape - size = h * w - hbm_size = self.hbm_tensor_size(size) hbm_addr = self._allocate_hbm(hbm_size) - else: - h, w = tensor_var.physical_shape - hbm_size = self.hbm_tensor_size(h * w) super().store_to_hbm( tensor_name=tensor_var.name, # internal name for symbol table lookup diff --git a/utils/load_config.py b/utils/load_config.py index a1a7a8b..7aacec6 100644 --- a/utils/load_config.py +++ b/utils/load_config.py @@ -3,6 +3,11 @@ import re from pathlib import Path +try: # stdlib since 3.11, and the reason this file used to read nothing + import tomllib +except ImportError: # pragma: no cover - Python < 3.11 + tomllib = None + try: import toml except ImportError: @@ -42,9 +47,25 @@ def load_svh_settings(file_path: str | Path) -> dict[str, int]: def load_toml_config(file_path, section_to_load=None, mode="BEHAVIOR"): - if toml is None: - raise ImportError("'toml' package required for load_toml_config. Install with: pip install toml") - with open(file_path) as f: - full_toml = toml.load(f) + """``full_toml[mode][section_to_load]``, or ``{}`` if either is absent. + + Reads through `tomllib` when it is available, which since Python 3.11 is + always. It used to require the third-party `toml` package and raise + `ImportError` without it -- and `toml` is not a declared dependency and is + not installed in CI, so every caller that wrapped this in a try/except and + fell back to a default was taking the fallback unconditionally. That was + silent: the compiler's `plena_settings.toml` values had never once been + read. + """ + if tomllib is not None: + with open(file_path, "rb") as f: + full_toml = tomllib.load(f) + elif toml is not None: + with open(file_path) as f: + full_toml = toml.load(f) + else: # pragma: no cover - Python < 3.11 without `toml` + raise ImportError( + "load_toml_config needs `tomllib` (Python 3.11+) or the `toml` package" + ) mode_section = full_toml.get(mode, {}) return mode_section.get(section_to_load, {}) From 66c1439dc2109db294667563ffd7d153fc6de622 Mon Sep 17 00:00:00 2001 From: "Qichao (Arlo) Wang" Date: Sat, 29 Aug 2026 01:40:05 +0100 Subject: [PATCH 3/6] models: Mamba-2 and KDA references, and the configs they are checked against The float64 definition every lowering is compared to, plus the shape and precision types that carry a model's geometry into the compiler. `aten/models/mamba2/reference.py` is the selective SSM: the chunked scan whose `L` matrix is `exp(cs_i - cs_j)` masked causal, a 1-semiseparable decay mask that turns the chunk-internal recursion into a matmul. One scalar decay per head. `aten/models/kda/reference.py` is Kimi Delta Attention, and it is not the same shape of problem. **KDA's decay is channel-wise on the key axis** -- `exp(log_decay)[:, :, None, :]`, one value per key channel rather than one per head. That single difference is what makes KDA's decode sweeps `key_dim` lanes wide -- 6.2% of VLEN 2048 -- and it is the mechanism behind every operator-level result on this branch, including the one place a purpose-built engine would beat this lowering on architecture rather than on encoding. `KdaShape` carries hidden size, heads, key_dim, value_dim, conv kernel, chunk and the gate lower bound. `state_precision.py` refuses an MX KV type on spilled state at compile time: under MX the state decodes as e4m3 while the address stride assumes two bytes, which is a silent wrong answer rather than a fault. `doc/Model_Lib` gains four configs, fetched rather than written: kimi-linear-48b-a3b hidden 2304, 32 KDA heads, head_dim 128, 27 layers, 20 KDA and 7 full-attention nemotron-3-nano-4b hybrid_override_pattern over Mamba/attention/MLP nemotron-3-super-120b-a12b the same pattern at MoE scale mamba2-2.7b The Kimi K3 shape the measurements use -- 96 heads, key_dim and value_dim 128, 93 layers -- is a **projected** shape and not one of these. Per-head geometry matches `kimi-linear-48b-a3b` exactly (head_dim 128, conv kernel 4), so per-head numbers carry across; anything multiplied by heads or layers does not. Co-Authored-By: Claude Opus 5 --- aten/models/kda/__init__.py | 74 ++++ aten/models/kda/reference.py | 327 ++++++++++++++ aten/models/kda/shape.py | 54 +++ aten/models/kda/state_precision.py | 95 ++++ aten/models/mamba2/__init__.py | 21 + aten/models/mamba2/reference.py | 370 ++++++++++++++++ aten/native_ops.yaml | 106 +++++ aten/ops/cpu/mamba_ops.py | 98 +++++ aten/ops/plena/mamba_ops.py | 134 ++++++ aten/tests/test_kda_reference.py | 410 ++++++++++++++++++ aten/tests/test_mamba2_reference.py | 220 ++++++++++ aten/tests/test_registry_cpu_signatures.py | 43 ++ doc/Model_Lib/kimi-linear-48b-a3b.json | 86 ++++ doc/Model_Lib/mamba2-2.7b.json | 38 ++ doc/Model_Lib/nemotron-3-nano-4b.json | 54 +++ doc/Model_Lib/nemotron-3-super-120b-a12b.json | 69 +++ 16 files changed, 2199 insertions(+) create mode 100644 aten/models/kda/__init__.py create mode 100644 aten/models/kda/reference.py create mode 100644 aten/models/kda/shape.py create mode 100644 aten/models/kda/state_precision.py create mode 100644 aten/models/mamba2/__init__.py create mode 100644 aten/models/mamba2/reference.py create mode 100644 aten/ops/cpu/mamba_ops.py create mode 100644 aten/ops/plena/mamba_ops.py create mode 100644 aten/tests/test_kda_reference.py create mode 100644 aten/tests/test_mamba2_reference.py create mode 100644 aten/tests/test_registry_cpu_signatures.py create mode 100644 doc/Model_Lib/kimi-linear-48b-a3b.json create mode 100644 doc/Model_Lib/mamba2-2.7b.json create mode 100644 doc/Model_Lib/nemotron-3-nano-4b.json create mode 100644 doc/Model_Lib/nemotron-3-super-120b-a12b.json diff --git a/aten/models/kda/__init__.py b/aten/models/kda/__init__.py new file mode 100644 index 0000000..e49d041 --- /dev/null +++ b/aten/models/kda/__init__.py @@ -0,0 +1,74 @@ +"""Kimi K3 KDA (Kimi Delta Attention) CPU reference. + +Lives here, not in the Simulator, for the same reason +``aten/models/mamba2/reference.py`` does: the dependency direction across the +two repositories is Simulator -> Compiler and never the reverse, so a reference +placed here is importable from both the compiler's lowering tests and the +emulator's testbench, while one placed in the Simulator is not. + +**Re-exports are lazy on purpose.** ``aten/plena/compiler.py`` imports the KDA +lowering at module scope, and the lowering needs ``KdaShape``. Importing a +submodule runs this ``__init__`` first, so an eager ``from .reference import ...`` +here drags ``torch`` into every ``compiler.aten.plena.*`` import -- which breaks +the ``moe-stage-guard`` CI job, whose whole point is running the source-only +guards with just pytest and pyyaml installed. ``KdaShape`` therefore lives in the +torch-free ``shape`` module, and everything that genuinely needs torch resolves +through ``__getattr__`` only when someone asks for it. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +# Torch-free, so it can be imported eagerly and by the lowering. +from compiler.aten.models.kda.shape import KdaShape + +if TYPE_CHECKING: # pragma: no cover - import-time typing only + from compiler.aten.models.kda.reference import ( + KdaConvWeights, + KdaState, + KdaRecurrentState, + activate_log_decay, + kda_recurrent_sequence, + kda_state_engine_prefill, + kda_state_engine_step, + kda_step, + ) + from compiler.aten.models.kda.state_precision import ( + StateStorage, + quantize_state, + storage_bytes, + ) + +_REFERENCE_EXPORTS = frozenset( + { + "KdaConvWeights", + "KdaState", + "KdaRecurrentState", + "activate_log_decay", + "kda_recurrent_sequence", + "kda_state_engine_prefill", + "kda_state_engine_step", + "kda_step", + } +) +_PRECISION_EXPORTS = frozenset({"StateStorage", "quantize_state", "storage_bytes"}) + +__all__ = ["KdaShape", *sorted(_REFERENCE_EXPORTS | _PRECISION_EXPORTS)] + + +def __getattr__(name: str): + """Resolve the torch-backed exports on first use (PEP 562).""" + if name in _REFERENCE_EXPORTS: + from compiler.aten.models.kda import reference + + return getattr(reference, name) + if name in _PRECISION_EXPORTS: + from compiler.aten.models.kda import state_precision + + return getattr(state_precision, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted(__all__) diff --git a/aten/models/kda/reference.py b/aten/models/kda/reference.py new file mode 100644 index 0000000..60a3d71 --- /dev/null +++ b/aten/models/kda/reference.py @@ -0,0 +1,327 @@ +"""Readable FP32 CPU reference for the recurrent core of Kimi K3 KDA. + +The state layout follows FlashKDA's ``transpose_state_layout=True`` contract: +``[batch, head, value_dim, key_dim]``. This module models the recurrent KDA +core after q/k/v/decay/beta projection and short convolution. Output gating, +RMSNorm, and the output projection remain ordinary PLENA Vector/Matrix work. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +import torch.nn.functional as functional +from torch import Tensor + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.models.kda.state_precision import StateStorage, quantize_state + + +@dataclass +class KdaState: + recurrent: Tensor + + @classmethod + def zeros( + cls, + shape: KdaShape, + batch_size: int, + *, + device: torch.device | str = "cpu", + ) -> KdaState: + return cls( + torch.zeros( + batch_size, + shape.num_heads, + shape.value_dim, + shape.key_dim, + dtype=torch.float32, + device=device, + ) + ) + + def clone(self) -> KdaState: + return KdaState(self.recurrent.clone()) + + def reset_(self) -> None: + self.recurrent.zero_() + + +@dataclass +class KdaRecurrentState: + recurrent: Tensor + conv: Tensor + + @classmethod + def zeros( + cls, + shape: KdaShape, + batch_size: int, + *, + device: torch.device | str = "cpu", + ) -> KdaRecurrentState: + return cls( + recurrent=KdaState.zeros(shape, batch_size, device=device).recurrent, + conv=torch.zeros( + batch_size, + shape.num_heads * (2 * shape.key_dim + shape.value_dim), + shape.conv_kernel, + dtype=torch.float32, + device=device, + ), + ) + + +@dataclass(frozen=True) +class KdaConvWeights: + q: Tensor + k: Tensor + v: Tensor + q_bias: Tensor | None = None + k_bias: Tensor | None = None + v_bias: Tensor | None = None + + +def _require_shape(name: str, value: Tensor, expected: tuple[int, ...]) -> None: + if tuple(value.shape) != expected: + raise ValueError(f"{name} has shape {tuple(value.shape)}, expected {expected}") + + +def activate_log_decay( + gate: Tensor, + a_log: Tensor, + dt_bias: Tensor, + *, + lower_bound: float, +) -> Tensor: + """Return Kimi K3's channel-wise log decay in ``[lower_bound, 0]``.""" + if gate.ndim != 3: + raise ValueError("gate must have shape [batch, heads, key_dim]") + _, heads, key_dim = gate.shape + _require_shape("a_log", a_log, (heads,)) + _require_shape("dt_bias", dt_bias, (heads, key_dim)) + if lower_bound >= 0: + raise ValueError("lower_bound must be negative") + rate = torch.exp(a_log.float())[None, :, None] + return lower_bound * torch.sigmoid(rate * (gate.float() + dt_bias.float()[None, :, :])) + + +def kda_step( + q: Tensor, + k: Tensor, + v: Tensor, + gate: Tensor, + beta_logit: Tensor, + state: KdaState, + a_log: Tensor, + dt_bias: Tensor, + shape: KdaShape, + *, + scale: float | None = None, + state_storage: StateStorage | str = StateStorage.FP32, +) -> tuple[Tensor, KdaState]: + """Execute one recurrent KDA token with FP32 update and reduction. + + With state ``S`` stored as ``[value, key]`` the operation is:: + + D = S * exp(log_decay) + error = sigmoid(beta) * (v - D @ k) + S_new = D + outer(error, k) + output = scale * S_new @ q + + This is equivalent to the delta-rule matrix form + ``S_new = D @ (I - beta*k*k^T) + beta*v*k^T``. + """ + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("q, k, and v must have shape [batch, heads, dimension]") + batch = q.shape[0] + key_shape = (batch, shape.num_heads, shape.key_dim) + value_shape = (batch, shape.num_heads, shape.value_dim) + _require_shape("q", q, key_shape) + _require_shape("k", k, key_shape) + _require_shape("v", v, value_shape) + _require_shape("gate", gate, key_shape) + _require_shape("beta_logit", beta_logit, (batch, shape.num_heads)) + _require_shape( + "state", + state.recurrent, + (batch, shape.num_heads, shape.value_dim, shape.key_dim), + ) + + # FlashKDA's recurrent kernel applies epsilon inside rsqrt, rather than + # clamping the norm as torch.nn.functional.normalize does. + q_normalized = q.float() * torch.rsqrt(q.float().square().sum(dim=-1, keepdim=True) + 1.0e-6) + k_normalized = k.float() * torch.rsqrt(k.float().square().sum(dim=-1, keepdim=True) + 1.0e-6) + log_decay = activate_log_decay( + gate, + a_log, + dt_bias, + lower_bound=shape.gate_lower_bound, + ) + decayed = state.recurrent.float() * torch.exp(log_decay)[:, :, None, :] + prediction = torch.einsum("bhvk,bhk->bhv", decayed, k_normalized) + beta = torch.sigmoid(beta_logit.float())[:, :, None] + error = beta * (v.float() - prediction) + updated = decayed + error[:, :, :, None] * k_normalized[:, :, None, :] + output_scale = scale if scale is not None else 1.0 / math.sqrt(shape.key_dim) + output = output_scale * torch.einsum("bhvk,bhk->bhv", updated, q_normalized) + return output, KdaState(quantize_state(updated, state_storage)) + + +def _causal_conv_step( + value: Tensor, + state: Tensor, + weight: Tensor, + bias: Tensor | None, +) -> tuple[Tensor, Tensor]: + if value.ndim != 2 or state.ndim != 3: + raise ValueError("KDA conv value/state must be [batch, channels] and [batch, channels, kernel]") + batch, channels = value.shape + _require_shape("KDA conv state", state, (batch, channels, state.shape[-1])) + kernel = state.shape[-1] + if weight.shape == (channels, 1, kernel): + weight = weight[:, 0, :] + _require_shape("KDA conv weight", weight, (channels, kernel)) + updated = torch.roll(state.float(), shifts=-1, dims=-1) + updated[..., -1] = value.float() + output = (updated * weight.float().unsqueeze(0)).sum(dim=-1) + if bias is not None: + _require_shape("KDA conv bias", bias, (channels,)) + output = output + bias.float() + return functional.silu(output), updated + + +def kda_state_engine_step( + projected: Tensor, + state: KdaRecurrentState, + conv_weights: KdaConvWeights, + a_log: Tensor, + dt_bias: Tensor, + shape: KdaShape, + *, + scale: float | None = None, + state_storage: StateStorage | str = StateStorage.BF16, + conv_state_storage: StateStorage | str | None = None, +) -> tuple[Tensor, KdaRecurrentState]: + """The KDA state-engine boundary: three short convs, then the recurrence. + + Named for what it does. The boundary is a real one -- three convolutions and + a recurrence, with everything before it a projection and everything after it + a norm -- and it is the right place to compare a lowering against, however + many instructions the lowering spends crossing it. + """ + if projected.ndim != 2: + raise ValueError("projected token must have shape [batch, projection_size]") + batch = projected.shape[0] + key_width = shape.projection_size + value_width = shape.num_heads * shape.value_dim + expected = 3 * key_width + value_width + shape.num_heads + _require_shape("projected", projected, (batch, expected)) + q_raw, k_raw, v_raw, gate, beta = projected.float().split( + [key_width, key_width, value_width, key_width, shape.num_heads], dim=-1 + ) + q_state, k_state, v_state = state.conv.split([key_width, key_width, value_width], dim=1) + q, q_state = _causal_conv_step(q_raw, q_state, conv_weights.q, conv_weights.q_bias) + k, k_state = _causal_conv_step(k_raw, k_state, conv_weights.k, conv_weights.k_bias) + v, v_state = _causal_conv_step(v_raw, v_state, conv_weights.v, conv_weights.v_bias) + output, recurrent = kda_step( + q.reshape(batch, shape.num_heads, shape.key_dim), + k.reshape(batch, shape.num_heads, shape.key_dim), + v.reshape(batch, shape.num_heads, shape.value_dim), + gate.reshape(batch, shape.num_heads, shape.key_dim), + beta, + KdaState(state.recurrent), + a_log, + dt_bias, + shape, + scale=scale, + state_storage=state_storage, + ) + conv = torch.cat([q_state, k_state, v_state], dim=1) + # The descriptor carries `state_precision` and `conv_state_precision` + # independently, and the shipped Kimi configuration uses FP32 recurrent + # state with BF16 conv state. Folding both onto one parameter meant the + # combination that actually ships had no CPU reference at all. + return output.reshape(batch, -1), KdaRecurrentState( + recurrent.recurrent, + quantize_state(conv, state_storage if conv_state_storage is None else conv_state_storage), + ) + + +def kda_state_engine_prefill( + projected: Tensor, + state: KdaRecurrentState, + conv_weights: KdaConvWeights, + a_log: Tensor, + dt_bias: Tensor, + shape: KdaShape, + *, + scale: float | None = None, + state_storage: StateStorage | str = StateStorage.BF16, + conv_state_storage: StateStorage | str | None = None, +) -> tuple[Tensor, KdaRecurrentState]: + """Golden KDA prefill, defined as sequential state-engine steps.""" + if projected.ndim != 3: + raise ValueError("projected prefill input must be [batch, sequence, projection_size]") + outputs = [] + current = state + for token in projected.unbind(dim=1): + output, current = kda_state_engine_step( + token, + current, + conv_weights, + a_log, + dt_bias, + shape, + scale=scale, + state_storage=state_storage, + conv_state_storage=conv_state_storage, + ) + outputs.append(output) + return torch.stack(outputs, dim=1), current + + +def kda_recurrent_sequence( + q: Tensor, + k: Tensor, + v: Tensor, + gate: Tensor, + beta_logit: Tensor, + state: KdaState, + a_log: Tensor, + dt_bias: Tensor, + shape: KdaShape, + *, + scale: float | None = None, + state_storage: StateStorage | str = StateStorage.BF16, +) -> tuple[Tensor, KdaState]: + """Sequential golden reference for a prefill or multi-token decode trace.""" + if q.ndim != 4: + raise ValueError("sequence q must have shape [batch, tokens, heads, key_dim]") + batch, tokens, heads, key_dim = q.shape + _require_shape("q", q, (batch, tokens, shape.num_heads, shape.key_dim)) + _require_shape("k", k, (batch, tokens, heads, key_dim)) + _require_shape("v", v, (batch, tokens, shape.num_heads, shape.value_dim)) + _require_shape("gate", gate, (batch, tokens, heads, key_dim)) + _require_shape("beta_logit", beta_logit, (batch, tokens, shape.num_heads)) + + outputs = [] + current = state + for token in range(tokens): + output, current = kda_step( + q[:, token], + k[:, token], + v[:, token], + gate[:, token], + beta_logit[:, token], + current, + a_log, + dt_bias, + shape, + scale=scale, + state_storage=state_storage, + ) + outputs.append(output) + return torch.stack(outputs, dim=1), current diff --git a/aten/models/kda/shape.py b/aten/models/kda/shape.py new file mode 100644 index 0000000..da45ddd --- /dev/null +++ b/aten/models/kda/shape.py @@ -0,0 +1,54 @@ +"""KDA head geometry. + +Deliberately torch-free and in its own module. ``aten/plena/compiler.py`` +imports the KDA lowering at module scope, so anything the lowering needs at +import time is pulled into every ``compiler.aten.plena.*`` import -- including +the ones the ``moe-stage-guard`` CI job runs with only pytest and pyyaml +installed. Leaving ``KdaShape`` in ``reference.py`` next to ``import torch`` +made that job fail to collect. + +Same object for the reference and the lowering, so ``isinstance`` and dataclass +equality hold across the two. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +__all__ = ["KdaShape"] + + +@dataclass(frozen=True) +class KdaShape: + hidden_size: int + num_heads: int + key_dim: int + value_dim: int + conv_kernel: int + chunk_size: int = 16 + gate_lower_bound: float = -5.0 + + def __post_init__(self) -> None: + for name in ("hidden_size", "num_heads", "key_dim", "value_dim", "conv_kernel", "chunk_size"): + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + if self.gate_lower_bound >= 0: + raise ValueError("gate_lower_bound must be negative") + + @classmethod + def kimi_k3(cls) -> KdaShape: + return cls(7168, 96, 128, 128, 4) + + @property + def projection_size(self) -> int: + return self.num_heads * self.key_dim + + @property + def state_elements(self) -> int: + return self.num_heads * self.value_dim * self.key_dim + + @property + def conv_state_elements(self) -> int: + channels = self.num_heads * (2 * self.key_dim + self.value_dim) + return channels * self.conv_kernel + diff --git a/aten/models/kda/state_precision.py b/aten/models/kda/state_precision.py new file mode 100644 index 0000000..5f8f842 --- /dev/null +++ b/aten/models/kda/state_precision.py @@ -0,0 +1,95 @@ +"""State-storage quantizers for the Nemotron 3 Mamba CPU reference. + +All arithmetic outside :func:`quantize_state` remains FP32. ``MX8_B128`` is a +PLENA design candidate with one power-of-two scale per 128 state values. It is +deliberately named differently from OCP MXFP8, whose standard block size is 32. +""" + +from __future__ import annotations + +import math +from enum import StrEnum + +import torch +from torch import Tensor + + +class StateStorage(StrEnum): + FP32 = "fp32" + BF16 = "bf16" + FP16 = "fp16" + MX8_B128 = "mx8_b128" + + +def _round_e4m3fn(value: Tensor) -> Tensor: + """Round finite FP32 values to the E4M3FN value grid and return FP32.""" + sign = torch.sign(value) + magnitude = value.abs().clamp(max=448.0) + minimum_normal = 2.0**-6 + subnormal_step = 2.0**-9 + + subnormal = torch.round(magnitude / subnormal_step) * subnormal_step + safe = magnitude.clamp_min(minimum_normal) + exponent = torch.floor(torch.log2(safe)).clamp(min=-6, max=8) + normal_step = torch.pow(torch.tensor(2.0, device=value.device), exponent - 3) + normal = torch.round(magnitude / normal_step) * normal_step + rounded = torch.where(magnitude < minimum_normal, subnormal, normal) + return sign * rounded.clamp(max=448.0) + + +def quantize_mx8_b128(value: Tensor, *, block_size: int = 128) -> Tensor: + """Round-trip through E4M3FN with one E8M0-like scale per last-axis block.""" + if block_size <= 0: + raise ValueError("block_size must be positive") + original_shape = value.shape + if not original_shape: + raise ValueError("MX8 block quantization requires at least one dimension") + + source = value.float() + last = original_shape[-1] + blocks = math.ceil(last / block_size) + padded_last = blocks * block_size + if padded_last != last: + source = torch.nn.functional.pad(source, (0, padded_last - last)) + blocked = source.reshape(*source.shape[:-1], blocks, block_size) + maximum = blocked.abs().amax(dim=-1, keepdim=True) + + # Ceil prevents finite E4M3 values from overflowing 448. E8M0 can encode + # negative powers as well, so small state blocks still use the mantissa. + scale_exponent = torch.ceil(torch.log2((maximum / 448.0).clamp_min(2.0**-126))) + scale_exponent = scale_exponent.clamp(min=-126, max=127) + scale_exponent = torch.where(maximum == 0, torch.zeros_like(scale_exponent), scale_exponent) + scale = torch.pow(torch.tensor(2.0, device=value.device), scale_exponent) + restored = _round_e4m3fn(blocked / scale) * scale + return restored.reshape(*source.shape)[..., :last].reshape(original_shape) + + +def quantize_state( + value: Tensor, + storage: StateStorage | str, + *, + block_size: int = 128, +) -> Tensor: + """Return the FP32 value seen after writing and reading the state store.""" + storage = StateStorage(storage) + value = value.float() + if storage == StateStorage.FP32: + return value.clone() + if storage == StateStorage.BF16: + return value.to(torch.bfloat16).float() + if storage == StateStorage.FP16: + return value.to(torch.float16).float() + if storage == StateStorage.MX8_B128: + return quantize_mx8_b128(value, block_size=block_size) + raise ValueError(f"unsupported state storage: {storage}") + + +def storage_bytes(elements: int, storage: StateStorage | str, *, block_size: int = 128) -> int: + if elements < 0: + raise ValueError("elements must be non-negative") + storage = StateStorage(storage) + if storage == StateStorage.FP32: + return 4 * elements + if storage in {StateStorage.BF16, StateStorage.FP16}: + return 2 * elements + return elements + math.ceil(elements / block_size) diff --git a/aten/models/mamba2/__init__.py b/aten/models/mamba2/__init__.py new file mode 100644 index 0000000..9780a73 --- /dev/null +++ b/aten/models/mamba2/__init__.py @@ -0,0 +1,21 @@ +"""Mamba-2 reference semantics used to validate the PLENA lowering.""" + +from compiler.aten.models.mamba2.reference import ( + Mamba2Params, + Mamba2Result, + mamba2_chunked_reference, + mamba2_recurrent_reference, + mamba2_reference, + random_mamba2_params, + ssd_chunk_reference, +) + +__all__ = [ + "Mamba2Params", + "Mamba2Result", + "mamba2_chunked_reference", + "mamba2_recurrent_reference", + "mamba2_reference", + "random_mamba2_params", + "ssd_chunk_reference", +] diff --git a/aten/models/mamba2/reference.py b/aten/models/mamba2/reference.py new file mode 100644 index 0000000..d6fd98e --- /dev/null +++ b/aten/models/mamba2/reference.py @@ -0,0 +1,370 @@ +"""Mamba-2 reference semantics -- the golden the PLENA lowering is checked against. + +Three implementations, deliberately kept side by side: + +``mamba2_recurrent_reference`` + The textbook per-timestep recurrence. Slow, obviously correct, no chunking. + This is the definition; everything else is checked against it. + +``mamba2_chunked_reference`` + The SSD chunked form the PLENA prefill lowering actually implements, written + in the same order with the same intermediates. When the emulator disagrees + with the golden, comparing against *this* rather than against the recurrence + localises the failure to a stage instead of to "the layer". + +``mamba2_reference`` + The full mixer: in_proj, conv1d, dt activation, scan, gated norm, out_proj. + +The two scan implementations must agree to floating-point tolerance; +``test_mamba2_reference.py`` asserts that, which is what makes the chunked form +usable as an intermediate golden at all. + +A note on what "golden" means here +---------------------------------- +These run in float32 and are the *algorithmic* reference. They are NOT a model of +PLENA's arithmetic. In particular the emulator implements ``V_EXP_V`` as +libtorch's exact ``exp``, while the RTL model in +``PLENA_Tools/plena_quant/quant_operations/exp.py`` is a fixed-point range +reduction plus a 3-term Taylor with roughly 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. So a Mamba +accuracy result measured on the transactional emulator bounds the *lowering*, not +the silicon. Closing that gap means plumbing the hardware exp model into the +comparison, and is deliberately left as a separate, visible piece of work rather +than papered over with a loose tolerance. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass +class Mamba2Params: + """Weights and biases of one Mamba-2 mixer, in HuggingFace orientation.""" + + in_proj_weight: torch.Tensor # [in_proj_out, hidden] + conv_weight: torch.Tensor # [conv_dim, conv_kernel] + conv_bias: torch.Tensor | None # [conv_dim] + dt_bias: torch.Tensor # [num_heads] + A_log: torch.Tensor # [num_heads] + D: torch.Tensor # [num_heads] + norm_weight: torch.Tensor | None # [d_inner] + out_proj_weight: torch.Tensor # [hidden, d_inner] + + num_heads: int + head_dim: int + state_size: int + n_groups: int + conv_kernel: int + chunk_size: int + time_step_min: float = 0.0 + time_step_max: float = float("inf") + norm_eps: float = 1e-5 + + @property + def d_inner(self) -> int: + return self.num_heads * self.head_dim + + @property + def group_state(self) -> int: + return self.n_groups * self.state_size + + @property + def conv_dim(self) -> int: + return self.d_inner + 2 * self.group_state + + +@dataclass +class Mamba2Result: + """Layer output plus every intermediate a stage-level test needs.""" + + y: torch.Tensor # [batch, seq, hidden] + x: torch.Tensor # [batch, seq, num_heads, head_dim] (post conv+silu) + B: torch.Tensor # [batch, seq, n_groups, state_size] + C: torch.Tensor # [batch, seq, n_groups, state_size] + dt: torch.Tensor # [batch, seq, num_heads] (post softplus+clamp) + z: torch.Tensor # [batch, seq, d_inner] + scan_out: torch.Tensor # [batch, seq, num_heads, head_dim] + final_state: torch.Tensor # [batch, num_heads, state_size, head_dim] + + +def random_mamba2_params( + *, + hidden_size: int, + num_heads: int, + head_dim: int, + state_size: int, + n_groups: int = 1, + conv_kernel: int = 4, + chunk_size: int = 64, + use_conv_bias: bool = True, + use_norm_weight: bool = True, + seed: int = 0, + dtype: torch.dtype = torch.float32, +) -> Mamba2Params: + """Small, well-conditioned random parameters for tests. + + ``A_log`` is drawn so that ``A = -exp(A_log)`` lands in roughly ``[-4, -1]``: + the HF initialisation range, and deliberately not near zero, because + ``A -> 0`` makes the decay ~1 and hides exactly the error accumulation this + lowering has to get right. + """ + g = torch.Generator().manual_seed(seed) + d_inner = num_heads * head_dim + group_state = n_groups * state_size + in_proj_out = 2 * d_inner + 2 * group_state + num_heads + conv_dim = d_inner + 2 * group_state + + def rnd(*shape, scale=0.1): + return (torch.randn(*shape, generator=g, dtype=dtype) * scale).contiguous() + + return Mamba2Params( + in_proj_weight=rnd(in_proj_out, hidden_size), + conv_weight=rnd(conv_dim, conv_kernel, scale=0.3), + conv_bias=rnd(conv_dim, scale=0.05) if use_conv_bias else None, + dt_bias=rnd(num_heads, scale=0.5), + A_log=torch.log(torch.rand(num_heads, generator=g, dtype=dtype) * 3.0 + 1.0), + D=rnd(num_heads, scale=0.5), + norm_weight=torch.ones(d_inner, dtype=dtype) + rnd(d_inner, scale=0.05) + if use_norm_weight + else None, + out_proj_weight=rnd(hidden_size, d_inner), + num_heads=num_heads, + head_dim=head_dim, + state_size=state_size, + n_groups=n_groups, + conv_kernel=conv_kernel, + chunk_size=chunk_size, + ) + + +def _silu(t: torch.Tensor) -> torch.Tensor: + return t * torch.sigmoid(t) + + +def _causal_depthwise_conv1d( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None +) -> torch.Tensor: + """`x` is [batch, seq, channels]; `weight` is [channels, kernel]. + + Written as an explicit tap sum rather than ``F.conv1d`` so it mirrors the + lowering in ``program_mamba_common.mamba_conv1d_v0`` line for line -- the + causal left pad shows up as skipped taps in both. + """ + batch, seq, channels = x.shape + kernel = weight.shape[1] + out = torch.zeros_like(x) + for s in range(seq): + for j in range(kernel): + src = s - (kernel - 1) + j + if src < 0: + continue + out[:, s, :] += x[:, src, :] * weight[:, j] + if bias is not None: + out += bias + return out + + +def mamba2_recurrent_reference( + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + initial_state: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """The textbook selective-scan recurrence. This is the definition. + + Shapes: `x` [batch, seq, heads, head_dim]; `dt` [batch, seq, heads]; + `A` [heads]; `B`/`C` [batch, seq, groups, state]; `D` [heads]. + Returns (y [batch, seq, heads, head_dim], final_state + [batch, heads, state, head_dim]). + """ + batch, seq, heads, head_dim = x.shape + _, _, groups, state = B.shape + heads_per_group = heads // groups + + h = ( + torch.zeros(batch, heads, state, head_dim, dtype=x.dtype) + if initial_state is None + else initial_state.clone() + ) + y = torch.zeros_like(x) + + for t in range(seq): + dA = torch.exp(dt[:, t] * A) # [batch, heads] + h = h * dA[:, :, None, None] + for g in range(groups): + hs = slice(g * heads_per_group, (g + 1) * heads_per_group) + dtx = dt[:, t, hs][:, :, None] * x[:, t, hs] # [batch, heads_g, head_dim] + # rank-1 update: B[b, n] outer dtx[b, head, p] + h[:, hs] = h[:, hs] + B[:, t, g][:, None, :, None] * dtx[:, :, None, :] + y[:, t, hs] = torch.einsum("bhnp,bn->bhp", h[:, hs], C[:, t, g]) + y[:, t] = y[:, t] + D[None, :, None] * x[:, t] + return y, h + + +def ssd_chunk_reference( + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_state: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """The chunked (SSD) form, in the same order the PLENA lowering emits it. + + Each step here corresponds to one emitter in ``program_ssd.py``: + ``cs`` to ``ssd_chunk_cumsum_v0``, ``L`` to ``ssd_decay_mask_v0``, + ``Y_intra`` to ``ssd_chunk_head_v0``, the state update to + ``ssd_state_update_v0``, and ``Y_inter`` to ``ssd_inter_chunk_output_v0``. + """ + batch, seq, heads, head_dim = x.shape + _, _, groups, state = B.shape + heads_per_group = heads // groups + + h = ( + torch.zeros(batch, heads, state, head_dim, dtype=x.dtype) + if initial_state is None + else initial_state.clone() + ) + y = torch.zeros_like(x) + + for start in range(0, seq, chunk_size): + end = min(start + chunk_size, seq) + L = end - start + xc = x[:, start:end] # [b, L, h, p] + dtc = dt[:, start:end] # [b, L, h] + Bc = B[:, start:end] # [b, L, g, n] + Cc = C[:, start:end] + + a = dtc * A # [b, L, h] + cs = torch.cumsum(a, dim=1) # [b, L, h] + + # Intra-chunk: (L o G) @ X, with L[i,j] = exp(cs_i - cs_j), causal. + # Clamping to <= 0 before exp is what the lowering does and is exact on + # the causal half (cs is non-increasing because A < 0). + d_mat = cs[:, :, None, :] - cs[:, None, :, :] # [b, i, j, h] + d_mat = torch.clamp(d_mat, max=0.0) + causal = torch.tril(torch.ones(L, L, dtype=x.dtype)) + decay = torch.exp(d_mat) * causal[None, :, :, None] + + for g in range(groups): + hs = slice(g * heads_per_group, (g + 1) * heads_per_group) + G = torch.einsum("bin,bjn->bij", Cc[:, :, g], Bc[:, :, g]) # [b, i, j] + LG = G[:, :, :, None] * decay[:, :, :, hs] # [b, i, j, h] + # dt scales the *source* timestep j + LG = LG * dtc[:, None, :, hs] + y[:, start:end, hs] = torch.einsum("bijh,bjhp->bihp", LG, xc[:, :, hs]) + + # Inter-chunk: C @ h_prev, row i scaled by exp(cs_i) + y_inter = torch.einsum("bin,bhnp->bihp", Cc[:, :, g], h[:, hs]) + y[:, start:end, hs] = y[:, start:end, hs] + y_inter * torch.exp(cs[:, :, hs])[:, :, :, None] + + # State: h = h * exp(sum a) + B^T @ (X scaled by exp(cs_end - cs_j) * dt) + cs_end = cs[:, -1, hs] # [b, h] + scale = torch.exp(cs_end[:, None, :] - cs[:, :, hs]) * dtc[:, :, hs] # [b, L, h] + xd = xc[:, :, hs] * scale[:, :, :, None] + h[:, hs] = h[:, hs] * torch.exp(cs_end)[:, :, None, None] + torch.einsum( + "bjn,bjhp->bhnp", Bc[:, :, g], xd + ) + + y[:, start:end] = y[:, start:end] + D[None, None, :, None] * xc + + return y, h + + +def mamba2_reference( + hidden: torch.Tensor, + params: Mamba2Params, + *, + use_chunked: bool = True, + initial_state: torch.Tensor | None = None, +) -> Mamba2Result: + """Full Mamba-2 mixer forward pass. + + `hidden` is [batch, seq, hidden_size]. Set ``use_chunked=False`` to run the + plain recurrence instead of the SSD form. + """ + p = params + batch, seq, _ = hidden.shape + + projected = hidden @ p.in_proj_weight.T # [b, s, in_proj_out] + d_inner, gs = p.d_inner, p.group_state + z = projected[..., :d_inner] + xbc = projected[..., d_inner : d_inner + p.conv_dim] + dt_raw = projected[..., d_inner + p.conv_dim :] + + xbc = _silu(_causal_depthwise_conv1d(xbc, p.conv_weight, p.conv_bias)) + x_flat = xbc[..., :d_inner] + B = xbc[..., d_inner : d_inner + gs].reshape(batch, seq, p.n_groups, p.state_size) + C = xbc[..., d_inner + gs :].reshape(batch, seq, p.n_groups, p.state_size) + + dt = torch.nn.functional.softplus(dt_raw + p.dt_bias) + dt = torch.clamp(dt, min=p.time_step_min, max=p.time_step_max) + + x = x_flat.reshape(batch, seq, p.num_heads, p.head_dim) + A = -torch.exp(p.A_log) + + if use_chunked: + scan_out, final_state = ssd_chunk_reference( + x, dt, A, B, C, p.D, p.chunk_size, initial_state=initial_state + ) + else: + scan_out, final_state = mamba2_recurrent_reference( + x, dt, A, B, C, p.D, initial_state=initial_state + ) + + y = scan_out.reshape(batch, seq, d_inner) + + # Gated RMSNorm. + # + # The gate is applied BEFORE the variance, which is what both upstream + # implementations do: HuggingFace's `MambaRMSNormGated.forward` multiplies by + # `silu(gate)` and only then computes `hidden_states.pow(2).mean(-1)`, and + # mamba_ssm's `Mamba2` constructs `RMSNormGated(..., norm_before_gate=False)`. + # The two orders are different functions, not different roundings -- RMSNorm is + # not per-element homogeneous -- and at these shapes they differ by ~56% + # relative. Getting this backwards would be unrecoverable once real Mamba-2 + # weights are loaded, and would silently retarget every downstream + # emulator-vs-golden comparison. + # + # The reduction is over `d_inner / n_groups`, following mamba_ssm's + # `group_size=d_ssm // ngroups`. HuggingFace reduces over the whole of + # `intermediate_size`; the two agree exactly at `n_groups == 1`, which is what + # mamba2-2.7b and every other published Mamba-2 checkpoint uses. + y = y * _silu(z) + group_width = d_inner // p.n_groups + yg = y.reshape(batch, seq, p.n_groups, group_width) + rms = torch.rsqrt(yg.pow(2).mean(dim=-1, keepdim=True) + p.norm_eps) + yg = yg * rms + y = yg.reshape(batch, seq, d_inner) + if p.norm_weight is not None: + y = y * p.norm_weight + + out = y @ p.out_proj_weight.T + return Mamba2Result( + y=out, x=x, B=B, C=C, dt=dt, z=z, scan_out=scan_out, final_state=final_state + ) + + +def mamba2_chunked_reference(*args, **kwargs): + """Alias kept so callers can name the chunked scan explicitly.""" + return ssd_chunk_reference(*args, **kwargs) + + +__all__ = [ + "Mamba2Params", + "Mamba2Result", + "mamba2_chunked_reference", + "mamba2_recurrent_reference", + "mamba2_reference", + "random_mamba2_params", + "ssd_chunk_reference", +] diff --git a/aten/native_ops.yaml b/aten/native_ops.yaml index ca366ae..ae70d34 100644 --- a/aten/native_ops.yaml +++ b/aten/native_ops.yaml @@ -147,3 +147,109 @@ uses_fpram: false uses_mram: true doc: "Conv2d via im2col + linear projection: input_col @ weight_2d" + +# --------------------------------------------------------------------------- +# Mamba-2 / selective state-space model +# +# These are `composite`: each is several ISA stages, and the layer as a whole is +# deliberately NOT a single op. The chunk loop, the cross-chunk state carry and +# the transposed in_proj stay in the caller because they force HBM round trips +# (MRAM is writable only by H_PREFETCH_M) that a monolithic op would hide, taking +# the cost model with it. +# --------------------------------------------------------------------------- + +- func: "causal_conv1d(Tensor x, Tensor weight, Tensor? bias=None) -> Tensor" + category: composite + in_place: false + dispatch: + cpu: compiler.aten.ops.cpu.mamba_ops.causal_conv1d_cpu + plena: compiler.aten.ops.plena.mamba_ops.causal_conv1d_plena + plena_backend: + asm_template: null + memory_pattern: batch_only + tile_loops: manual + uses_fpram: false + uses_mram: false + doc: | + Causal depthwise conv1d over the sequence axis, kernel `conv_kernel`. + NOT routed through conv_ops.py's im2col: that hardcodes a square KxK patch, + has no padding (so it cannot express the causal left pad), and would build a + dense GEMM against a block-diagonal weight. With sequence on the VRAM row + axis the shift is a free row offset and the per-channel tap is a V_MUL_VV + operand. + +- func: "dt_activation(Tensor dt_raw, Tensor? dt_bias=None, float dt_min=0.0, float dt_max=inf) -> Tensor(dt_raw!)" + category: primitive + in_place: true + dispatch: + cpu: compiler.aten.ops.cpu.mamba_ops.dt_activation_cpu + plena: compiler.aten.ops.plena.mamba_ops.dt_activation_plena + plena_backend: + asm_template: null + memory_pattern: batch_only + tile_loops: none + uses_fpram: true + uses_mram: false + doc: | + dt = clamp(softplus(dt_raw + dt_bias), time_step_limit). Uses V_SOFTPLUS_V: + the ISA has no logarithm and no float-to-int conversion, so softplus has + neither an exact software lowering nor a table lowering, and dt feeds + exp(A*dt) on a multiplicative recurrence where an approximation compounds. + +- func: "mamba_gated_rmsnorm(Tensor y, Tensor z, Tensor? norm_weight=None) -> Tensor(y!)" + category: composite + in_place: true + dispatch: + cpu: compiler.aten.ops.cpu.mamba_ops.gated_rmsnorm_cpu + plena: compiler.aten.ops.plena.mamba_ops.gated_rmsnorm_plena + plena_backend: + asm_template: null + memory_pattern: batch_only + tile_loops: none + uses_fpram: true + uses_mram: false + doc: | + RMSNorm(y * silu(z)) -- the gate is applied BEFORE the variance, matching + HuggingFace's MambaRMSNormGated and mamba_ssm's norm_before_gate=False. The + reduction is over d_inner/n_groups rather than hidden_size, which is why + normalization_asm.py cannot be reused unchanged. + +- func: "ssd_scan(Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, Tensor D, int chunk_size) -> (Tensor y, Tensor final_state)" + category: composite + in_place: false + dispatch: + cpu: compiler.aten.ops.cpu.mamba_ops.ssd_scan_cpu + plena: compiler.aten.ops.plena.mamba_ops.ssd_scan_plena + plena_backend: + asm_template: null + memory_pattern: flash_attn + tile_loops: manual + uses_fpram: true + uses_mram: true + doc: | + Mamba-2 chunked (SSD) scan, prefill. Structurally an attention-shaped tiled + loop with a decay mask in place of the online softmax: G = C@B^T reuses + flashattn/qkt.py's M_TMM loop and Y_intra = (LoG)@X reuses pv.py's M_MM loop. + Returns (y, final_state); the caller owns the cross-chunk state carry, so the + second element is not optional bookkeeping. chunk_size has no default on + purpose: the PLENA lowering requires chunk_size == MLEN, so any literal here + would encode one build's tile size as if it were a property of Mamba-2 (whose + own config default is 256). + +- func: "ssm_recurrent_step(Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, Tensor D, Tensor state) -> (Tensor y, Tensor final_state)" + category: composite + in_place: false + dispatch: + cpu: compiler.aten.ops.cpu.mamba_ops.ssm_recurrent_step_cpu + plena: compiler.aten.ops.plena.mamba_ops.ssm_recurrent_step_plena + plena_backend: + asm_template: null + memory_pattern: batch_only + tile_loops: manual + uses_fpram: true + uses_mram: false + doc: | + One decode step: h = dA*h + dB*x, y = C.h + D*x. Returns (y, new_state). + State layout is [head][state][head_dim] with head_dim on lanes -- the only + layout under which the decay broadcast, the rank-1 input and the d_state + contraction are all existing opcodes. diff --git a/aten/ops/cpu/mamba_ops.py b/aten/ops/cpu/mamba_ops.py new file mode 100644 index 0000000..a36d337 --- /dev/null +++ b/aten/ops/cpu/mamba_ops.py @@ -0,0 +1,98 @@ +"""CPU reference implementations of the Mamba-2 ops registered in native_ops.yaml. + +Thin wrappers over ``aten/models/mamba2/reference.py`` so the registry's CPU +backend and the numerical golden are the same code -- if they were separate, a +divergence between them would be invisible until it showed up as an emulator +mismatch nobody could localise. +""" + +from __future__ import annotations + +import torch + +from compiler.aten.models.mamba2.reference import ( + _causal_depthwise_conv1d, + mamba2_recurrent_reference, + ssd_chunk_reference, +) + + +def causal_conv1d_cpu(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None): + """`x` [batch, seq, channels], `weight` [channels, kernel].""" + return _causal_depthwise_conv1d(x, weight, bias) + + +def dt_activation_cpu( + dt_raw: torch.Tensor, + dt_bias: torch.Tensor | None = None, + dt_min: float = 0.0, + dt_max: float = float("inf"), +): + dt = dt_raw if dt_bias is None else dt_raw + dt_bias + return torch.clamp(torch.nn.functional.softplus(dt), min=dt_min, max=dt_max) + + +def gated_rmsnorm_cpu( + y: torch.Tensor, + z: torch.Tensor, + norm_weight: torch.Tensor | None = None, + n_groups: int = 1, + eps: float = 1e-5, +): + """``RMSNorm(y * silu(z)) * norm_weight``, gate applied before the variance. + + This exists because pointing the registry at ``norm_ops.rms_norm_cpu`` did not + work: dispatch forwards positionally and that function's signature is + ``(input, eps, eps_offset, reci_hid_offset)``, so `z` bound to `eps` and + `sqrt(mean(y^2) + z)` went negative wherever `z` did -- NaN, silently, with the + gate and the weight both discarded. + """ + y = y * torch.nn.functional.silu(z) + *lead, d_inner = y.shape + group_width = d_inner // n_groups + yg = y.reshape(*lead, n_groups, group_width) + yg = yg * torch.rsqrt(yg.pow(2).mean(dim=-1, keepdim=True) + eps) + y = yg.reshape(*lead, d_inner) + return y if norm_weight is None else y * norm_weight + + +def ssd_scan_cpu( + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + chunk_size: int, + initial_state: torch.Tensor | None = None, +): + """Chunked SSD scan -- the form the PLENA prefill lowering implements. + + `chunk_size` is required: the PLENA lowering pins it to MLEN, so a default + here would silently encode one build's tile size, and on this CPU path it + would produce a numerically valid but differently-chunked result instead of + an error. + """ + return ssd_chunk_reference(x, dt, A, B, C, D, chunk_size, initial_state=initial_state) + + +def ssm_recurrent_step_cpu( + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + state: torch.Tensor, +): + """One decode step. `x`/`dt`/`B`/`C` carry a length-1 sequence axis.""" + return mamba2_recurrent_reference(x, dt, A, B, C, D, initial_state=state) + + +__all__ = [ + "causal_conv1d_cpu", + "gated_rmsnorm_cpu", + "dt_activation_cpu", + "ssd_scan_cpu", + "ssm_recurrent_step_cpu", +] diff --git a/aten/ops/plena/mamba_ops.py b/aten/ops/plena/mamba_ops.py new file mode 100644 index 0000000..34e29b8 --- /dev/null +++ b/aten/ops/plena/mamba_ops.py @@ -0,0 +1,134 @@ +"""PLENA backend for the Mamba-2 ops registered in native_ops.yaml. + +Each function is a thin registry adapter over the substrate mixins in +``aten/plena/`` -- the emitters live there because they are also called directly +by ``plena_frontend`` and by the emulator testbenches, which do not go through the +registry. Keeping the adapters this thin means there is exactly one lowering per +operation rather than two that can drift. + +Signature convention (shared with every other op in ``aten/ops/plena/``): the +first argument is the ``PlenaCompiler``, the rest are ``VRAMMatrixVar`` / +``InputVar`` / ``FPVar`` proxies, and the return value is a proxy for the result. +ISA is emitted eagerly as a side effect. +""" + +from __future__ import annotations + +from compiler.aten.plena.program_mamba_common import Mamba2Shape, MambaFPConstants +from compiler.aten.plena.vars import FPVar, InputVar, VRAMMatrixVar + + +def causal_conv1d_plena( + prog, + x: VRAMMatrixVar, + weight: VRAMMatrixVar, + out: VRAMMatrixVar, + scratch: VRAMMatrixVar, + shape: Mamba2Shape, + bias: VRAMMatrixVar | None = None, + num_rows: int | None = None, + history_rows: int = 0, +): + return prog.mamba_conv1d_v0( + x, weight, bias, out, scratch, shape, num_rows=num_rows, history_rows=history_rows + ) + + +def dt_activation_plena( + prog, + dt: VRAMMatrixVar, + consts: MambaFPConstants, + shape: Mamba2Shape, + dt_bias: VRAMMatrixVar | None = None, + rows: list[int] | None = None, +): + return prog.mamba_dt_activation_v0(dt, dt_bias, consts, shape, rows=rows) + + +def gated_rmsnorm_plena( + prog, + y: VRAMMatrixVar, + z: VRAMMatrixVar, + gate_scratch: VRAMMatrixVar, + sq_scratch: VRAMMatrixVar, + rms_fp: FPVar, + consts: MambaFPConstants, + shape: Mamba2Shape, + norm_weight: VRAMMatrixVar | None = None, + rows: list[int] | None = None, +): + return prog.mamba_gated_rmsnorm_v0( + y, z, norm_weight, gate_scratch, sq_scratch, rms_fp, consts, shape, rows=rows + ) + + +def ssd_scan_plena( + prog, + *, + b_chunk: InputVar, + c_chunk: VRAMMatrixVar, + x_chunk: InputVar, + decay: VRAMMatrixVar, + scores: VRAMMatrixVar, + y_out: VRAMMatrixVar, + shape: Mamba2Shape, + head_block_base: int = 0, + precision: dict | None = None, +): + """One chunk, one head, of the intra-chunk term. + + Deliberately *not* a whole-layer entry point: the chunk loop, the state + carry and the transposed in_proj are the caller's, because a single op that + hid them would also hide the HBM round trips they force (MRAM is writable + only by ``H_PREFETCH_M``) and make the resulting cost model unreadable. + """ + return prog.ssd_chunk_head_v0( + b_chunk=b_chunk, + c_chunk=c_chunk, + x_chunk=x_chunk, + decay=decay, + scores=scores, + y_out=y_out, + shape=shape, + head_block_base=head_block_base, + precision=precision, + ) + + +def ssm_recurrent_step_plena( + prog, + *, + state: VRAMMatrixVar, + x: VRAMMatrixVar, + b_fp: FPVar, + c_fp: FPVar, + da_fp: FPVar, + dt_fp: FPVar, + d_fp: FPVar, + y: VRAMMatrixVar, + scratch: VRAMMatrixVar, + shape: Mamba2Shape, + consts: MambaFPConstants, +): + return prog.ssm_decode_step_v0( + state=state, + x=x, + b_fp=b_fp, + c_fp=c_fp, + da_fp=da_fp, + dt_fp=dt_fp, + d_fp=d_fp, + y=y, + scratch=scratch, + shape=shape, + consts=consts, + ) + + +__all__ = [ + "causal_conv1d_plena", + "dt_activation_plena", + "gated_rmsnorm_plena", + "ssd_scan_plena", + "ssm_recurrent_step_plena", +] diff --git a/aten/tests/test_kda_reference.py b/aten/tests/test_kda_reference.py new file mode 100644 index 0000000..e1e64b6 --- /dev/null +++ b/aten/tests/test_kda_reference.py @@ -0,0 +1,410 @@ +from __future__ import annotations + +import math + +import pytest +import torch + +from compiler.aten.models.kda.reference import ( + KdaConvWeights, + KdaShape, + KdaState, + KdaRecurrentState, + activate_log_decay, + kda_recurrent_sequence, + kda_state_engine_prefill, + kda_state_engine_step, + kda_step, +) +from compiler.aten.models.kda.state_precision import StateStorage, quantize_state, storage_bytes + + +def _small_shape() -> KdaShape: + return KdaShape(hidden_size=8, num_heads=2, key_dim=4, value_dim=3, conv_kernel=2) + + +def test_official_kimi_k3_state_geometry() -> None: + shape = KdaShape.kimi_k3() + assert shape.projection_size == 12_288 + assert shape.state_elements == 1_572_864 + assert shape.state_elements * 2 == 3 * 1024 * 1024 + assert shape.conv_state_elements * 2 == 288 * 1024 + + +def test_log_decay_is_channelwise_and_bounded() -> None: + shape = _small_shape() + gate = torch.tensor( + [[[-100.0, -1.0, 0.0, 1.0], [100.0, 1.0, 0.0, -1.0]]], + dtype=torch.float32, + ) + decay = activate_log_decay( + gate, + torch.zeros(shape.num_heads), + torch.zeros(shape.num_heads, shape.key_dim), + lower_bound=shape.gate_lower_bound, + ) + assert torch.all(decay <= 0) + assert torch.all(decay >= shape.gate_lower_bound) + assert decay[0, 0, 0] > decay[0, 0, -1] + + +def test_log_decay_matches_its_closed_form() -> None: + """Pin the exact Kimi K3 formula, with a_log and dt_bias actually varying. + + test_log_decay_is_channelwise_and_bounded passes zeros for both, so it + holds for any function of `gate` alone -- dropping the dt_bias term or the + exp(a_log) rate leaves it green. Every other test in this file that touches + decay calls activate_log_decay as its own oracle, so nothing else can catch + an error inside it either. + """ + torch.manual_seed(23) + shape = KdaShape(hidden_size=32, num_heads=3, key_dim=5, value_dim=7, conv_kernel=4) + batch = 2 + gate = torch.randn(batch, shape.num_heads, shape.key_dim) + a_log = torch.randn(shape.num_heads) + dt_bias = torch.randn(shape.num_heads, shape.key_dim) + + actual = activate_log_decay( + gate, a_log, dt_bias, lower_bound=shape.gate_lower_bound + ) + rate = torch.exp(a_log)[None, :, None] + expected = shape.gate_lower_bound * torch.sigmoid( + rate * (gate + dt_bias[None, :, :]) + ) + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-7) + + +def test_step_matches_delta_rule_matrix_form() -> None: + torch.manual_seed(7) + shape = _small_shape() + batch = 1 + q = torch.randn(batch, shape.num_heads, shape.key_dim) + k = torch.randn_like(q) + v = torch.randn(batch, shape.num_heads, shape.value_dim) + gate = torch.randn_like(q) + beta_logit = torch.randn(batch, shape.num_heads) + initial = KdaState(torch.randn(batch, shape.num_heads, shape.value_dim, shape.key_dim)) + a_log = torch.randn(shape.num_heads) + dt_bias = torch.randn(shape.num_heads, shape.key_dim) + scale = 0.25 + + output, final = kda_step( + q, + k, + v, + gate, + beta_logit, + initial, + a_log, + dt_bias, + shape, + scale=scale, + state_storage=StateStorage.FP32, + ) + + qn = q.float() * torch.rsqrt(q.float().square().sum(dim=-1, keepdim=True) + 1.0e-6) + kn = k.float() * torch.rsqrt(k.float().square().sum(dim=-1, keepdim=True) + 1.0e-6) + log_decay = activate_log_decay(gate, a_log, dt_bias, lower_bound=shape.gate_lower_bound) + decayed = initial.recurrent * torch.exp(log_decay)[:, :, None, :] + beta = torch.sigmoid(beta_logit) + expected_states = [] + expected_outputs = [] + for head in range(shape.num_heads): + identity = torch.eye(shape.key_dim) + transition = identity - beta[0, head] * torch.outer(kn[0, head], kn[0, head]) + expected = decayed[0, head] @ transition + expected = expected + beta[0, head] * torch.outer(v[0, head], kn[0, head]) + expected_states.append(expected) + expected_outputs.append(scale * (expected @ qn[0, head])) + expected_state = torch.stack(expected_states).unsqueeze(0) + expected_output = torch.stack(expected_outputs).unsqueeze(0) + + torch.testing.assert_close(final.recurrent, expected_state) + torch.testing.assert_close(output, expected_output) + + +def test_sequence_matches_repeated_steps() -> None: + torch.manual_seed(11) + shape = _small_shape() + batch, tokens = 2, 5 + q = torch.randn(batch, tokens, shape.num_heads, shape.key_dim) + k = torch.randn_like(q) + v = torch.randn(batch, tokens, shape.num_heads, shape.value_dim) + gate = torch.randn_like(q) + beta = torch.randn(batch, tokens, shape.num_heads) + a_log = torch.randn(shape.num_heads) + dt_bias = torch.randn(shape.num_heads, shape.key_dim) + initial = KdaState.zeros(shape, batch) + + outputs, final = kda_recurrent_sequence( + q, + k, + v, + gate, + beta, + initial, + a_log, + dt_bias, + shape, + state_storage=StateStorage.FP32, + ) + assert outputs.shape == (batch, tokens, shape.num_heads, shape.value_dim) + assert final.recurrent.shape == (batch, shape.num_heads, shape.value_dim, shape.key_dim) + assert torch.isfinite(outputs).all() + + # The name of this test promises this comparison; without it the shape and + # isfinite assertions above hold even if the sequence drops the carried + # state between tokens, or emits them in reverse order. + carried = KdaState.zeros(shape, batch) + expected = [] + for token in range(tokens): + out, carried = kda_step( + q[:, token], + k[:, token], + v[:, token], + gate[:, token], + beta[:, token], + carried, + a_log, + dt_bias, + shape, + state_storage=StateStorage.FP32, + ) + expected.append(out) + torch.testing.assert_close( + outputs, torch.stack(expected, dim=1), rtol=1e-6, atol=1e-7 + ) + torch.testing.assert_close( + final.recurrent, carried.recurrent, rtol=1e-6, atol=1e-7 + ) + + +def test_prefill_includes_conv_and_matches_repeated_step() -> None: + torch.manual_seed(17) + shape = _small_shape() + batch, tokens = 2, 4 + key_width = shape.num_heads * shape.key_dim + value_width = shape.num_heads * shape.value_dim + projection_width = 3 * key_width + value_width + shape.num_heads + projected = torch.randn(batch, tokens, projection_width) + conv_weights = KdaConvWeights( + q=torch.randn(key_width, shape.conv_kernel), + k=torch.randn(key_width, shape.conv_kernel), + v=torch.randn(value_width, shape.conv_kernel), + ) + a_log = torch.randn(shape.num_heads) + dt_bias = torch.randn(shape.num_heads, shape.key_dim) + initial = KdaRecurrentState.zeros(shape, batch) + output, state = kda_state_engine_prefill( + projected, + initial, + conv_weights, + a_log, + dt_bias, + shape, + state_storage=StateStorage.BF16, + ) + expected_outputs = [] + expected_state = initial + for token in projected.unbind(1): + token_output, expected_state = kda_state_engine_step( + token, + expected_state, + conv_weights, + a_log, + dt_bias, + shape, + state_storage=StateStorage.BF16, + ) + expected_outputs.append(token_output) + torch.testing.assert_close(output, torch.stack(expected_outputs, dim=1)) + torch.testing.assert_close(state.recurrent, expected_state.recurrent) + torch.testing.assert_close(state.conv, expected_state.conv) + + +def test_default_scale_uses_key_dimension() -> None: + shape = _small_shape() + q = torch.ones(1, shape.num_heads, shape.key_dim) + k = torch.ones_like(q) + v = torch.ones(1, shape.num_heads, shape.value_dim) + gate = torch.zeros_like(q) + beta = torch.zeros(1, shape.num_heads) + state = KdaState.zeros(shape, 1) + a_log = torch.zeros(shape.num_heads) + dt_bias = torch.zeros(shape.num_heads, shape.key_dim) + + default, _ = kda_step(q, k, v, gate, beta, state, a_log, dt_bias, shape) + explicit, _ = kda_step( + q, + k, + v, + gate, + beta, + state, + a_log, + dt_bias, + shape, + scale=1.0 / math.sqrt(shape.key_dim), + ) + torch.testing.assert_close(default, explicit) + + +def test_shape_validation_rejects_wrong_state_layout() -> None: + shape = _small_shape() + q = torch.zeros(1, shape.num_heads, shape.key_dim) + v = torch.zeros(1, shape.num_heads, shape.value_dim) + wrong = KdaState(torch.zeros(1, shape.num_heads, shape.key_dim, shape.value_dim)) + with pytest.raises(ValueError, match="state has shape"): + kda_step( + q, + q, + v, + q, + torch.zeros(1, shape.num_heads), + wrong, + torch.zeros(shape.num_heads), + torch.zeros(shape.num_heads, shape.key_dim), + shape, + ) + + +def test_conv_state_precision_is_independent_of_recurrent_state_precision() -> None: + """FP32 recurrent state with BF16 conv state is what actually ships. + + The descriptor carries `state_precision` and `conv_state_precision` + separately, but this reference used to fold both onto one parameter, so the + shipped Kimi combination had no CPU reference at all. Two tokens are needed: + the first token's conv rounding only becomes observable once it is read back + as history. + """ + torch.manual_seed(23) + shape = _small_shape() + batch, tokens = 1, 2 + key_width = shape.num_heads * shape.key_dim + value_width = shape.num_heads * shape.value_dim + projected = torch.randn(batch, tokens, 3 * key_width + value_width + shape.num_heads) + conv_weights = KdaConvWeights( + q=torch.randn(key_width, shape.conv_kernel), + k=torch.randn(key_width, shape.conv_kernel), + v=torch.randn(value_width, shape.conv_kernel), + ) + a_log = torch.randn(shape.num_heads) + dt_bias = torch.randn(shape.num_heads, shape.key_dim) + initial = KdaRecurrentState.zeros(shape, batch) + + common = dict(state_storage=StateStorage.FP32) + uniform, uniform_state = kda_state_engine_prefill(projected, initial, conv_weights, a_log, dt_bias, shape, **common) + mixed, mixed_state = kda_state_engine_prefill( + projected, initial, conv_weights, a_log, dt_bias, shape, conv_state_storage=StateStorage.BF16, **common + ) + + # The recurrent state stays FP32 in both runs, so any divergence can only + # have entered through the conv history. + assert not torch.equal(uniform_state.conv, mixed_state.conv) + assert not torch.equal(uniform, mixed) + torch.testing.assert_close(uniform, mixed, rtol=2e-2, atol=2e-2) + + +def test_decode_step_matches_the_transposed_formulation(): + """``T[k, v] == S[v, k]``. The contract between this reference and the + PLENA lowering in ``aten/plena/program_kda_recurrent.py``. + + The lowering stores state transposed, with ``key`` as the row axis, because + in that orientation every step of the recurrence -- decay, predict, update, + read out -- is one broadcast-scalar operation per row. This test pins the + two orientations together so a lowering bug cannot hide behind a plausible + reimplementation of the maths. + """ + torch.manual_seed(0) + # Non-degenerate on purpose, each choice load-bearing: + # batch = 2 -- batch 1 lets a batch-0 broadcast leak pass + # key != value -- a square state makes a transposed axis numerically + # equivalent instead of a shape error, and lets the + # output scale 1/sqrt(key_dim) be confused with value_dim + # 0.05 * randn -- at unit variance ||x||^2 is O(1) and the 1e-6 epsilon + # inside rsqrt perturbs by ~2e-7, under this test's atol, + # so the FlashKDA normalisation convention below would + # not actually be pinned. At this magnitude it is. + shape = KdaShape(hidden_size=32, num_heads=3, key_dim=8, value_dim=5, conv_kernel=4) + batch = 2 + + state = KdaState( + torch.randn(batch, shape.num_heads, shape.value_dim, shape.key_dim) + ) + q = 0.05 * torch.randn(batch, shape.num_heads, shape.key_dim) + k = 0.05 * torch.randn(batch, shape.num_heads, shape.key_dim) + v = torch.randn(batch, shape.num_heads, shape.value_dim) + gate = torch.randn(batch, shape.num_heads, shape.key_dim) + beta_logit = torch.randn(batch, shape.num_heads) + a_log = torch.randn(shape.num_heads) + dt_bias = torch.randn(shape.num_heads, shape.key_dim) + + expected_out, expected_state = kda_step( + q, k, v, gate, beta_logit, state.clone(), a_log, dt_bias, shape + ) + + # --- transposed replay: T is [batch, heads, key_dim, value_dim] --------- + t = state.recurrent.transpose(-2, -1).clone() + + # Same epsilon-inside-rsqrt normalisation the reference uses; + # torch.nn.functional.normalize clamps instead and does NOT agree. + def _norm(x): + return x.float() * torch.rsqrt(x.float().square().sum(-1, keepdim=True) + 1.0e-6) + + q_n, k_n = _norm(q), _norm(k) + log_decay = activate_log_decay( + gate, a_log, dt_bias, lower_bound=shape.gate_lower_bound + ) + + # decay: one scalar per key row + t = t * torch.exp(log_decay)[:, :, :, None] + # predict: accumulate a value-length vector across key rows + pred = torch.einsum("bhkv,bhk->bhv", t, k_n) + beta = torch.sigmoid(beta_logit.float())[:, :, None] + error = beta * (v.float() - pred) + # update: rank-1, one scalar per key row + t = t + error[:, :, None, :] * k_n[:, :, :, None] + # read out: on the UPDATED state + scale = 1.0 / math.sqrt(shape.key_dim) + out = scale * torch.einsum("bhkv,bhk->bhv", t, q_n) + + torch.testing.assert_close(out, expected_out, rtol=1e-5, atol=1e-6) + torch.testing.assert_close( + t.transpose(-2, -1), expected_state.recurrent, rtol=1e-5, atol=1e-6 + ) + + +# --------------------------------------------------------------------------- +# state_precision.py coverage. +# +# These two cover a module that arrived without its tests. A second Mamba +# reference beside aten/models/mamba2/reference.py would be a drift hazard, so +# only the precision helper was kept -- and without these +# quantize_state's MX8_B128 path, _round_e4m3fn, and storage_bytes have zero +# coverage in this repo while sitting in the import graph of every KDA test. +# --------------------------------------------------------------------------- + + +def test_state_storage_round_trips_and_byte_counts() -> None: + value = torch.randn(257, generator=torch.Generator().manual_seed(23)) * 3.7 + torch.testing.assert_close(quantize_state(value, StateStorage.FP32), value) + for storage in (StateStorage.BF16, StateStorage.FP16, StateStorage.MX8_B128): + rounded = quantize_state(value, storage) + assert rounded.shape == value.shape + assert torch.isfinite(rounded).all() + assert not torch.equal(rounded, value) + assert storage_bytes(256, StateStorage.FP32) == 1024 + assert storage_bytes(256, StateStorage.BF16) == 512 + assert storage_bytes(256, StateStorage.MX8_B128) == 258 + + +def test_mx8_e4m3fn_rounding_carries_across_exponents() -> None: + value = torch.zeros(128) + value[:5] = torch.tensor([1.9375, 248.0, 432.0, 448.0, -448.0]) + restored = quantize_state(value, StateStorage.MX8_B128) + torch.testing.assert_close( + restored[:5], + torch.tensor([2.0, 256.0, 448.0, 448.0, -448.0]), + rtol=0, + atol=0, + ) diff --git a/aten/tests/test_mamba2_reference.py b/aten/tests/test_mamba2_reference.py new file mode 100644 index 0000000..f47f986 --- /dev/null +++ b/aten/tests/test_mamba2_reference.py @@ -0,0 +1,220 @@ +"""The chunked SSD reference must agree with the plain recurrence. + +This is the load-bearing test of the whole Mamba effort. The PLENA prefill +lowering implements the chunked form, so the chunked reference is what the +emulator gets compared against -- and a chunked reference that quietly disagrees +with the recurrence would make every downstream numerical result meaningless +while looking green. + +The cases below deliberately include multi-chunk sequences, a non-dividing +sequence length (so the tail chunk is exercised), multiple groups, and an +initial state, because those are the four places the chunk bookkeeping can be +wrong without a single-chunk test noticing. +""" + +from __future__ import annotations + +import unittest + +import torch + +from compiler.aten.models.mamba2.reference import ( + mamba2_recurrent_reference, + mamba2_reference, + random_mamba2_params, + ssd_chunk_reference, +) + + +def _scan_inputs(*, batch, seq, heads, head_dim, state, groups, seed=0): + g = torch.Generator().manual_seed(seed) + x = torch.randn(batch, seq, heads, head_dim, generator=g) + # dt strictly positive, as softplus guarantees downstream. + dt = torch.rand(batch, seq, heads, generator=g) * 0.2 + 0.01 + A = -torch.exp(torch.rand(heads, generator=g) * 1.0) + B = torch.randn(batch, seq, groups, state, generator=g) * 0.3 + C = torch.randn(batch, seq, groups, state, generator=g) * 0.3 + D = torch.randn(heads, generator=g) * 0.5 + return x, dt, A, B, C, D + + +class TestSSDChunkedMatchesRecurrence(unittest.TestCase): + def _assert_agrees(self, *, seq, chunk, groups=1, heads=4, initial=False, seed=0): + head_dim, state, batch = 8, 4, 2 + x, dt, A, B, C, D = _scan_inputs( + batch=batch, seq=seq, heads=heads, head_dim=head_dim, state=state, groups=groups, seed=seed + ) + init = None + if initial: + gen = torch.Generator().manual_seed(seed + 1) + init = torch.randn(batch, heads, state, head_dim, generator=gen) * 0.2 + + y_ref, h_ref = mamba2_recurrent_reference(x, dt, A, B, C, D, initial_state=init) + y_chunk, h_chunk = ssd_chunk_reference(x, dt, A, B, C, D, chunk, initial_state=init) + + torch.testing.assert_close(y_chunk, y_ref, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(h_chunk, h_ref, rtol=1e-4, atol=1e-5) + + def test_single_chunk(self): + self._assert_agrees(seq=8, chunk=8) + + def test_multiple_whole_chunks(self): + self._assert_agrees(seq=32, chunk=8) + + def test_tail_chunk_shorter_than_chunk_size(self): + # seq % chunk != 0 is the branch most likely to be live at real dims and + # least likely to be covered by a synthetic single-chunk test. + self._assert_agrees(seq=30, chunk=8) + + def test_multiple_groups(self): + self._assert_agrees(seq=24, chunk=8, groups=2, heads=4) + + def test_carries_an_initial_state(self): + # Cross-chunk state passing is the entire point of the chunked form; a + # zero initial state hides an error in exactly that term. + self._assert_agrees(seq=24, chunk=8, initial=True) + + def test_slow_decay_regime(self): + """A close to zero makes the decay ~1 and stresses accumulation. + + This is the regime that justifies an SSM (long memory) and the regime + where a chunked form that mishandles the inter-chunk decay drifts. + """ + batch, seq, heads, head_dim, state, groups, chunk = 2, 32, 4, 8, 4, 1, 8 + g = torch.Generator().manual_seed(7) + x = torch.randn(batch, seq, heads, head_dim, generator=g) + dt = torch.rand(batch, seq, heads, generator=g) * 0.02 + 0.001 + A = -torch.full((heads,), 0.05) + B = torch.randn(batch, seq, groups, state, generator=g) * 0.3 + C = torch.randn(batch, seq, groups, state, generator=g) * 0.3 + D = torch.randn(heads, generator=g) * 0.5 + + y_ref, h_ref = mamba2_recurrent_reference(x, dt, A, B, C, D) + y_chunk, h_chunk = ssd_chunk_reference(x, dt, A, B, C, D, chunk) + torch.testing.assert_close(y_chunk, y_ref, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(h_chunk, h_ref, rtol=1e-4, atol=1e-5) + + +class TestMamba2LayerReference(unittest.TestCase): + def test_full_layer_agrees_between_chunked_and_recurrent_scans(self): + params = random_mamba2_params( + hidden_size=32, num_heads=4, head_dim=8, state_size=4, n_groups=1, chunk_size=8, seed=3 + ) + hidden = torch.randn(2, 24, 32, generator=torch.Generator().manual_seed(4)) + + chunked = mamba2_reference(hidden, params, use_chunked=True) + recurrent = mamba2_reference(hidden, params, use_chunked=False) + + torch.testing.assert_close(chunked.y, recurrent.y, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(chunked.final_state, recurrent.final_state, rtol=1e-4, atol=1e-5) + + def test_intermediates_have_the_shapes_the_lowering_assumes(self): + params = random_mamba2_params( + hidden_size=32, num_heads=4, head_dim=8, state_size=4, n_groups=2, chunk_size=8, seed=5 + ) + out = mamba2_reference(torch.randn(2, 16, 32), params) + + self.assertEqual(out.y.shape, (2, 16, 32)) + self.assertEqual(out.x.shape, (2, 16, 4, 8)) + self.assertEqual(out.B.shape, (2, 16, 2, 4)) + self.assertEqual(out.C.shape, (2, 16, 2, 4)) + self.assertEqual(out.dt.shape, (2, 16, 4)) + self.assertEqual(out.final_state.shape, (2, 4, 4, 8)) + + def test_dt_is_positive_and_respects_the_clamp(self): + params = random_mamba2_params( + hidden_size=32, num_heads=4, head_dim=8, state_size=4, chunk_size=8, seed=6 + ) + params.time_step_min = 0.01 + params.time_step_max = 0.1 + out = mamba2_reference(torch.randn(2, 16, 32), params) + self.assertTrue(torch.all(out.dt >= 0.01)) + self.assertTrue(torch.all(out.dt <= 0.1)) + + def test_gate_is_applied_before_the_variance(self): + """RMSNorm(y * silu(z)), not RMSNorm(y) * silu(z). + + Both upstream implementations gate first -- HuggingFace's MambaRMSNormGated + reduces after multiplying, and mamba_ssm defaults to norm_before_gate=False. + RMSNorm is not per-element homogeneous, so the two orders are different + functions; at these shapes they differ by tens of percent, and the error is + unrecoverable once real weights are loaded. + """ + params = random_mamba2_params( + hidden_size=32, num_heads=4, head_dim=8, state_size=4, n_groups=1, chunk_size=8, seed=11 + ) + hidden = torch.randn(2, 16, 32, generator=torch.Generator().manual_seed(12)) + out = mamba2_reference(hidden, params) + + # Recompute the norm block both ways from the same scan output. + d_inner = params.d_inner + y = out.scan_out.reshape(2, 16, d_inner) + z = out.z + gate_then_norm = y * torch.nn.functional.silu(z) + gate_then_norm = gate_then_norm * torch.rsqrt( + gate_then_norm.pow(2).mean(-1, keepdim=True) + params.norm_eps + ) + gate_then_norm = gate_then_norm * params.norm_weight + expected = gate_then_norm @ params.out_proj_weight.T + + torch.testing.assert_close(out.y, expected, rtol=1e-4, atol=1e-5) + + norm_then_gate = y * torch.rsqrt(y.pow(2).mean(-1, keepdim=True) + params.norm_eps) + norm_then_gate = norm_then_gate * params.norm_weight * torch.nn.functional.silu(z) + self.assertFalse( + torch.allclose(expected, norm_then_gate @ params.out_proj_weight.T, rtol=1e-2), + "the two orders must be distinguishable, else this test proves nothing", + ) + + def test_grouped_norm_differs_from_ungrouped_when_n_groups_gt_1(self): + """n_groups > 1 must actually change the normalisation. + + Nothing else in the suite exercised the grouping, so a lowering that + normalised over the whole of d_inner would have passed everything. + """ + params = random_mamba2_params( + hidden_size=32, num_heads=4, head_dim=8, state_size=4, n_groups=2, chunk_size=8, seed=13 + ) + hidden = torch.randn(2, 16, 32, generator=torch.Generator().manual_seed(14)) + out = mamba2_reference(hidden, params) + + d_inner = params.d_inner + gated = out.scan_out.reshape(2, 16, d_inner) * torch.nn.functional.silu(out.z) + ungrouped = gated * torch.rsqrt(gated.pow(2).mean(-1, keepdim=True) + params.norm_eps) + ungrouped = (ungrouped * params.norm_weight) @ params.out_proj_weight.T + self.assertFalse( + torch.allclose(out.y, ungrouped, rtol=1e-2), + "grouped and ungrouped RMSNorm must differ at n_groups=2", + ) + + def test_full_layer_agrees_between_scans_with_multiple_groups(self): + params = random_mamba2_params( + hidden_size=32, num_heads=4, head_dim=8, state_size=4, n_groups=2, chunk_size=8, seed=15 + ) + hidden = torch.randn(2, 24, 32, generator=torch.Generator().manual_seed(16)) + chunked = mamba2_reference(hidden, params, use_chunked=True) + recurrent = mamba2_reference(hidden, params, use_chunked=False) + torch.testing.assert_close(chunked.y, recurrent.y, rtol=1e-4, atol=1e-5) + + def test_conv1d_is_causal(self): + """Perturbing a future timestep must not change an earlier output. + + Catches an off-by-one in the tap indexing, which is the single easiest + mistake to make in the conv and is invisible in an aggregate error metric. + """ + params = random_mamba2_params( + hidden_size=32, num_heads=4, head_dim=8, state_size=4, chunk_size=8, seed=8 + ) + hidden = torch.randn(1, 12, 32, generator=torch.Generator().manual_seed(9)) + base = mamba2_reference(hidden, params) + + perturbed = hidden.clone() + perturbed[:, 7:] += 5.0 + after = mamba2_reference(perturbed, params) + + torch.testing.assert_close(after.y[:, :7], base.y[:, :7], rtol=1e-5, atol=1e-6) + self.assertFalse(torch.allclose(after.y[:, 7:], base.y[:, 7:])) + + +if __name__ == "__main__": + unittest.main() diff --git a/aten/tests/test_registry_cpu_signatures.py b/aten/tests/test_registry_cpu_signatures.py new file mode 100644 index 0000000..7fcd35f --- /dev/null +++ b/aten/tests/test_registry_cpu_signatures.py @@ -0,0 +1,43 @@ +"""Every CPU backend in native_ops.yaml must accept the declared parameter list.""" +import inspect, re, unittest +from compiler.aten.ops.registry import OpRegistry + + +class TestRegistryCpuSignatures(unittest.TestCase): + def test_declared_arity_is_accepted_by_the_cpu_backend(self): + # OpRegistry.dispatch forwards *args positionally with no binding -- the + # `func:` string is used only to extract the op name -- so a CPU target whose + # real signature differs from the declaration silently binds the wrong things. + # mamba_gated_rmsnorm did exactly that: it pointed at rms_norm_cpu(input, eps, + # ...), so the gate tensor bound to `eps` and the result was NaN. + registry = OpRegistry.load() + problems = [] + # Pre-existing deviations, recorded rather than fixed so that a NEW one fails + # this test. flash_attention declares hq/hkv/h_qkv/batch_size/seq_len/kv_seq_len + # for the PLENA backend's benefit; the CPU reference genuinely needs only + # Q/K/V/scale, and every caller passes the rest by keyword, so nothing binds + # wrongly today. Removing an entry from this set is a fix; adding one needs a + # reason in the commit message. + KNOWN_DEVIATIONS = {"flash_attention"} + for name in registry.list_ops(): + schema = registry.get_op(name) + declared = re.search(r"\((.*)\)\s*->", schema.func_signature) + n_declared = 0 if not declared or not declared.group(1).strip() else len( + declared.group(1).split(",") + ) + fn = schema.resolve("cpu") + params = list(inspect.signature(fn).parameters.values()) + n_positional = sum( + 1 for p in params + if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) + ) + if n_positional < n_declared and name not in KNOWN_DEVIATIONS: + problems.append( + f"{name}: declares {n_declared} args but " + f"{fn.__module__}.{fn.__name__} takes {n_positional} positional" + ) + self.assertEqual(problems, [], "\n".join(problems)) + + +if __name__ == "__main__": + unittest.main() diff --git a/doc/Model_Lib/kimi-linear-48b-a3b.json b/doc/Model_Lib/kimi-linear-48b-a3b.json new file mode 100644 index 0000000..69e2a71 --- /dev/null +++ b/doc/Model_Lib/kimi-linear-48b-a3b.json @@ -0,0 +1,86 @@ +{ + "architectures": [ + "KimiLinearForCausalLM" + ], + "auto_map": { + "AutoConfig": "configuration_kimi.KimiLinearConfig", + "AutoModel": "modeling_kimi.KimiLinearModel", + "AutoModelForCausalLM": "modeling_kimi.KimiLinearForCausalLM" + }, + "bos_token_id": 163584, + "dtype": "bfloat16", + "eos_token_id": 163586, + "first_k_dense_replace": 1, + "head_dim": 72, + "hidden_act": "silu", + "hidden_size": 2304, + "initializer_range": 0.02, + "intermediate_size": 9216, + "kv_lora_rank": 512, + "linear_attn_config": { + "full_attn_layers": [ + 4, + 8, + 12, + 16, + 20, + 24, + 27 + ], + "head_dim": 128, + "kda_layers": [ + 1, + 2, + 3, + 5, + 6, + 7, + 9, + 10, + 11, + 13, + 14, + 15, + 17, + 18, + 19, + 21, + 22, + 23, + 25, + 26 + ], + "num_heads": 32, + "short_conv_kernel_size": 4 + }, + "mla_use_nope": true, + "model_max_length": 1048576, + "model_type": "kimi_linear", + "moe_intermediate_size": 1024, + "moe_layer_freq": 1, + "moe_renormalize": true, + "moe_router_activation_func": "sigmoid", + "num_attention_heads": 32, + "num_expert_group": 1, + "num_experts": 256, + "num_experts_per_token": 8, + "num_hidden_layers": 27, + "num_key_value_heads": 32, + "num_nextn_predict_layers": 0, + "num_shared_experts": 1, + "pad_token_id": 163839, + "q_lora_rank": null, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-05, + "rope_scaling": null, + "rope_theta": 10000.0, + "routed_scaling_factor": 2.446, + "tie_word_embeddings": false, + "topk_group": 1, + "transformers_version": "4.57.1", + "use_cache": true, + "use_grouped_topk": true, + "v_head_dim": 128, + "vocab_size": 163840 +} diff --git a/doc/Model_Lib/mamba2-2.7b.json b/doc/Model_Lib/mamba2-2.7b.json new file mode 100644 index 0000000..ca51aca --- /dev/null +++ b/doc/Model_Lib/mamba2-2.7b.json @@ -0,0 +1,38 @@ +{ + "architectures": [ + "Mamba2ForCausalLM" + ], + "bos_token_id": 0, + "chunk_size": 256, + "conv_kernel": 4, + "eos_token_id": 0, + "expand": 2, + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.1, + "layer_norm_epsilon": 1e-05, + "model_type": "mamba2", + "n_groups": 1, + "num_heads": 80, + "num_hidden_layers": 64, + "pad_token_id": 0, + "rescale_prenorm_residual": false, + "residual_in_fp32": true, + "rms_norm": true, + "state_size": 128, + "tie_word_embeddings": false, + "time_step_floor": 0.0001, + "time_step_limit": [ + 0.0, + Infinity + ], + "time_step_max": 0.1, + "time_step_min": 0.001, + "time_step_rank": "auto", + "torch_dtype": "float32", + "use_bias": false, + "use_cache": true, + "use_conv_bias": true, + "vocab_size": 50288 +} diff --git a/doc/Model_Lib/nemotron-3-nano-4b.json b/doc/Model_Lib/nemotron-3-nano-4b.json new file mode 100644 index 0000000..1186301 --- /dev/null +++ b/doc/Model_Lib/nemotron-3-nano-4b.json @@ -0,0 +1,54 @@ +{ + "architectures": [ + "NemotronHForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_nemotron_h.NemotronHConfig", + "AutoModelForCausalLM": "modeling_nemotron_h.NemotronHForCausalLM" + }, + "bos_token_id": 1, + "chunk_size": 256, + "conv_kernel": 4, + "eos_token_id": 2, + "expand": 2, + "head_dim": 128, + "hidden_dropout": 0.0, + "hidden_size": 3136, + "hybrid_override_pattern": "M-M-M-MM-M-M*-M-M*-M-M-M*-M-M-MM*-MMM-M-M-", + "initializer_range": 0.02, + "intermediate_size": 12544, + "layer_norm_epsilon": 1e-05, + "mamba_head_dim": 80, + "mamba_hidden_act": "silu", + "mamba_num_heads": 96, + "mamba_proj_bias": false, + "max_position_embeddings": 262144, + "mlp_bias": false, + "mlp_hidden_act": "relu2", + "model_type": "nemotron_h", + "n_groups": 8, + "num_attention_heads": 40, + "num_hidden_layers": 42, + "num_key_value_heads": 8, + "num_logits_to_keep": 1, + "pad_token_id": 0, + "rescale_prenorm_residual": true, + "residual_in_fp32": false, + "rms_norm_eps": 1e-05, + "sliding_window": null, + "ssm_state_size": 128, + "tie_word_embeddings": false, + "time_step_floor": 0.0001, + "time_step_max": 0.1, + "time_step_min": 0.001, + "time_step_rank": 256, + "torch_dtype": "bfloat16", + "transformers_version": "4.53.0", + "use_bias": false, + "use_cache": true, + "use_conv_bias": true, + "use_mamba_kernels": true, + "vocab_size": 131072 +} diff --git a/doc/Model_Lib/nemotron-3-super-120b-a12b.json b/doc/Model_Lib/nemotron-3-super-120b-a12b.json new file mode 100644 index 0000000..033b215 --- /dev/null +++ b/doc/Model_Lib/nemotron-3-super-120b-a12b.json @@ -0,0 +1,69 @@ +{ + "architectures": [ + "NemotronHForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_nemotron_h.NemotronHConfig", + "AutoModelForCausalLM": "modeling_nemotron_h.NemotronHForCausalLM" + }, + "bos_token_id": 1, + "chunk_size": 128, + "conv_kernel": 4, + "dtype": "bfloat16", + "eos_token_id": 2, + "expand": 2, + "head_dim": 128, + "hidden_dropout": 0.0, + "hidden_size": 4096, + "hybrid_override_pattern": "MEMEMEM*EMEMEMEM*EMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEMEM*EMEMEMEM*EMEMEMEME", + "initializer_range": 0.02, + "intermediate_size": 2688, + "layer_norm_epsilon": 1e-05, + "mamba_head_dim": 64, + "mamba_hidden_act": "silu", + "mamba_num_heads": 128, + "mamba_proj_bias": false, + "mamba_ssm_cache_dtype": "float32", + "max_position_embeddings": 262144, + "mlp_bias": false, + "mlp_hidden_act": "relu2", + "model_type": "nemotron_h", + "moe_intermediate_size": 2688, + "moe_latent_size": 1024, + "moe_shared_expert_intermediate_size": 5376, + "moe_shared_expert_overlap": false, + "mtp_hybrid_override_pattern": "*E", + "n_group": 1, + "n_groups": 8, + "n_routed_experts": 512, + "n_shared_experts": 1, + "norm_eps": 1e-05, + "norm_topk_prob": true, + "num_attention_heads": 32, + "num_experts_per_tok": 22, + "num_hidden_layers": 88, + "num_key_value_heads": 2, + "num_logits_to_keep": 1, + "num_nextn_predict_layers": 1, + "pad_token_id": 0, + "partial_rotary_factor": 1.0, + "rescale_prenorm_residual": true, + "residual_in_fp32": false, + "rope_theta": 10000, + "routed_scaling_factor": 5.0, + "sliding_window": null, + "ssm_state_size": 128, + "tie_word_embeddings": false, + "time_step_floor": 0.0001, + "time_step_max": 0.1, + "time_step_min": 0.001, + "topk_group": 1, + "transformers_version": "4.57.6", + "use_bias": false, + "use_cache": true, + "use_conv_bias": true, + "use_mamba_kernels": true, + "vocab_size": 131072 +} From 1cbe82feac0211da6643a7bda8b0ac45f3603438 Mon Sep 17 00:00:00 2001 From: "Qichao (Arlo) Wang" Date: Sat, 29 Aug 2026 01:40:05 +0100 Subject: [PATCH 4/6] mamba2: lower the selective SSM onto the static ISA Decode and the chunked scan, both as straight-line code with every address baked at assembly time. `ssm_decode_step_v0` is the recurrence `state = state * dA + B x`, then `y = C state + D x`. One `V_FMA_VF` sweep per head inside one hardware loop: 54 static instructions per head at every `state_size` from 4 to 32, against 200 / 352 / 656 / 1,384 before. Expanded, the issue stream is `11 * state_size + 33` against a pre-conversion slope of about 42 -- the conversion cut the coefficient, it did not remove the linearity. `ssd_decay_mask_v0` builds the `L` matrix -- `exp(cs_i - cs_j)` masked causal -- so the chunk-internal recursion is `Y_intra = (L . G) @ X`, a matmul the matrix machine already does. `V_SOFTPLUS_V` gives `dt` in one instruction where the previous lowering needed exp, add and log. `ssd_transposed_projection_v0` contracts on the other axis for the state-passing half. The generator learns the Mamba layer so a whole model can be parsed and lowered rather than hand-assembled, with a utilization report per stage. Every emitter takes `stage` as a **required** argument rather than defaulting it. A default is not harmless here: stage markers are sticky, so an emitter that bills itself to the wrong stage also bills everything emitted after it, and the MoE stage-attribution guard in CI is what caught that. Co-Authored-By: Claude Opus 5 --- aten/plena/program_mamba_common.py | 640 ++++++++++++++++++++++++ aten/plena/program_ssd.py | 613 +++++++++++++++++++++++ aten/plena/program_ssm_recurrent.py | 272 ++++++++++ aten/tests/test_mamba_stage_contract.py | 95 ++++ aten/tests/test_ssm_decode_step.py | 188 +++++++ generator/parser/llm_parser.py | 452 ++++++++++++++++- generator/passes/code_gen.py | 525 ++++++++++++++++++- generator/passes/utilization_report.py | 78 ++- generator/tests/test_mamba_code_gen.py | 316 ++++++++++++ generator/tests/test_mamba_parser.py | 403 +++++++++++++++ 10 files changed, 3559 insertions(+), 23 deletions(-) create mode 100644 aten/plena/program_mamba_common.py create mode 100644 aten/plena/program_ssd.py create mode 100644 aten/plena/program_ssm_recurrent.py create mode 100644 aten/tests/test_mamba_stage_contract.py create mode 100644 aten/tests/test_ssm_decode_step.py create mode 100644 generator/tests/test_mamba_code_gen.py create mode 100644 generator/tests/test_mamba_parser.py diff --git a/aten/plena/program_mamba_common.py b/aten/plena/program_mamba_common.py new file mode 100644 index 0000000..c430bf2 --- /dev/null +++ b/aten/plena/program_mamba_common.py @@ -0,0 +1,640 @@ +"""Shared Mamba / Mamba-2 (selective state-space model) program-builder helpers. + +This module holds the pieces both the prefill (`program_ssd.py`) and the decode +(`program_ssm_recurrent.py`) lowerings need: the shape contract, the FPRAM +constant block, the causal depthwise conv1d, the `dt` activation, SiLU, and the +gated RMSNorm. The two paths are deliberately separate modules because on PLENA +they are *not* the same kernel at a different sequence length -- they use +different VRAM layouts and different opcode mixes, and because `C_LOOP_START` +takes a compile-time immediate trip count and the ISA has no branch, they are two +separately compiled programs regardless. + +Scope: **Mamba-2 only**. Mamba-1's `A` is a `[d_inner, state_size]` diagonal, +which would need a `state_size`-wide vector replicated across `head_dim` rows and +applied with `V_MUL_VV`, because PLENA's only broadcast is scalar-to-vector +(`v_broadcast_en` in `doc/operation.svh` is wired solely to the `*_VF` path). +Mamba-2's `A` is one scalar per head, which is exactly what `V_MUL_VF` plus the +FPRAM row map already express. + +Hardware constraints that shaped this lowering +---------------------------------------------- +* **MRAM is write-only from HBM** (`doc/memory_layout.md`), so every + activation-times-activation product costs an `H_STORE_V` + `H_PREFETCH_M` round + trip. All such round trips here take the BF16 `keyvalue` path + (`hbm_element_bytes=2`), never the default MX-FP8 `weights` path: the SSM state + is a multiplicative accumulator carried across every chunk, and requantising it + to 3 mantissa bits at each chunk boundary compounds without bound. That is the + sharp difference from a KV cache, which is written once and never + read-modify-written, and is why e4m3 is fine there and not here. +* **`V_MUL_VF` broadcasts one FP register across all VLEN lanes.** Per-head + scalar broadcasts (the decay `dA_h`, the skip `D_h`, the per-row `exp(cs_i)`) + are only correct when one VRAM row does not straddle two heads -- i.e. when + `head_dim` is a multiple of MLEN. :meth:`Mamba2Shape.validate` enforces that + rather than letting one head's scalar silently apply to its neighbours. +* **There is no logarithm and no float-to-int conversion**, so `softplus` has no + exact software lowering and no lookup-table lowering either. It uses the + `V_SOFTPLUS_V` opcode added alongside this substrate. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass + +from compiler.aten.plena.vars import FPVar, VRAMMatrixVar + +# ============================================================================ +# Stage markers +# ============================================================================ + +#: Prefix of the explicit stage marker comment, shared with the routed-MoE +#: substrate. The emulator's stage profiler keys on this exact string. +MAMBA_STAGE_MARKER_PREFIX = "@stage=" + +#: Every Mamba stage name the emulator's ``StageKind`` understands. +#: +#: Kept in lockstep with ``transactional_emulator/src/stage_profile.rs``; the +#: cross-repo test in ``aten/tests/test_mamba_stage_contract.py`` pins it. Marker +#: emission is validated against this set so a typo fails at ASM-gen time instead +#: of silently collapsing a region into ``other``. +MAMBA_STAGES: frozenset[str] = frozenset( + { + "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", + } +) + + +def mamba_stage_marker(stage: str, detail: str = "") -> str: + """Format the explicit stage marker comment for ``stage``. + + Markers are *authoritative and sticky*: once a program contains any marker, + the emulator stops applying its legacy substring rules and attributes every + instruction to the most recent marker. A marker must therefore be emitted + whenever the stage changes, and must not be emitted mid-stage. + + The corollary is that work done inside a general-purpose helper called from a + marked region -- ``linear_projection``'s HBM weight prefetch, for instance -- + bills to the enclosing marker rather than to a stage of its own. + """ + if stage not in MAMBA_STAGES: + raise ValueError(f"unknown Mamba stage {stage!r}; expected one of {sorted(MAMBA_STAGES)}") + return f"{MAMBA_STAGE_MARKER_PREFIX}{stage}" + (f" {detail}" if detail else "") + + +# ============================================================================ +# Shape contract +# ============================================================================ + + +@dataclass(frozen=True) +class Mamba2Shape: + """Resolved Mamba-2 mixer dimensions for one layer. + + Field names mirror HuggingFace's ``Mamba2Config`` so a config JSON maps across + without a translation table. + """ + + hidden_size: int + num_heads: int + head_dim: int + state_size: int + n_groups: int + conv_kernel: int + chunk_size: int + seq_len: int + batch_size: int = 1 + #: softplus(dt) is clamped to this range, matching HF's ``time_step_limit``. + time_step_min: float = 0.0 + time_step_max: float = float("inf") + + @property + def d_inner(self) -> int: + return self.num_heads * self.head_dim + + @property + def group_state(self) -> int: + """Width of B or C: ``n_groups * state_size``.""" + return self.n_groups * self.state_size + + @property + def conv_dim(self) -> int: + """Width of the tensor the depthwise conv1d runs over: ``[x, B, C]``.""" + return self.d_inner + 2 * self.group_state + + @property + def in_proj_out(self) -> int: + """Width of the fused in_proj output: ``[z, x, B, C, dt]``.""" + return 2 * self.d_inner + 2 * self.group_state + self.num_heads + + @property + def heads_per_group(self) -> int: + return self.num_heads // self.n_groups + + @property + def num_chunks(self) -> int: + return math.ceil(self.seq_len / self.chunk_size) + + @property + def slice_offsets(self) -> dict[str, tuple[int, int]]: + """Column offset and width of each slice within the fused in_proj output.""" + z0 = 0 + x0 = z0 + self.d_inner + b0 = x0 + self.d_inner + c0 = b0 + self.group_state + dt0 = c0 + self.group_state + return { + "z": (z0, self.d_inner), + "x": (x0, self.d_inner), + "B": (b0, self.group_state), + "C": (c0, self.group_state), + "dt": (dt0, self.num_heads), + } + + def validate(self, mlen: int) -> None: + """Reject shapes this lowering would silently get wrong. + + Every check corresponds to a real failure mode, not defensive pedantry -- + see the module docstring. + """ + # head_dim must not be SMALLER than MLEN either: V_MUL_VF broadcasts one FP + # register across all VLEN lanes, so a row spanning two heads would apply one + # head's decay/skip scalar to its neighbour, silently and with no diagnostic. + if self.chunk_size != mlen: + raise ValueError( + f"chunk_size ({self.chunk_size}) must equal MLEN ({mlen}) in this lowering. " + "The tile_row_* family is per-MLEN-column-block (it defaults tile_col_idx=0 " + "and addresses rows as base + row*MLEN inside block 0), and one S_MAP_FP_V " + "moves exactly MLEN scalars, so a wider chunk would leave the columns past " + "the first block untouched and read uninitialised FPRAM for cs_i -- " + "half-processed rather than uniformly wrong, which is far harder to notice. " + "Mamba-2's default chunk_size is 256; supporting it needs the emitters to " + "loop tile_col_idx over ceil(cols/MLEN), which is deliberately not done yet." + ) + if self.head_dim != mlen: + raise ValueError( + f"head_dim ({self.head_dim}) must equal MLEN ({mlen}) in this lowering, for " + "the same per-column-block reason as chunk_size above; head_dim > MLEN would " + "silently process only the first MLEN lanes of each head." + ) + if self.state_size != mlen: + raise ValueError( + f"state_size ({self.state_size}) must equal MLEN ({mlen}) in this lowering, " + "for the same per-column-block reason as chunk_size above." + ) + if self.d_inner // self.n_groups != mlen: + raise ValueError( + f"gated RMSNorm reduces one MLEN-wide VRAM row but divides by " + f"d_inner/n_groups ({self.d_inner // self.n_groups}); those must match " + f"(MLEN={mlen}) or the normalisation constant is wrong. Accumulating " + "V_RED_SUM across group_width/MLEN rows would lift this." + ) + if self.n_groups <= 0 or self.num_heads % self.n_groups != 0: + raise ValueError( + f"num_heads ({self.num_heads}) must be divisible by n_groups ({self.n_groups}); " + "B and C are shared across the heads of a group." + ) + if self.conv_kernel < 1: + raise ValueError(f"conv_kernel must be >= 1, got {self.conv_kernel}") + if self.seq_len < 1: + raise ValueError(f"seq_len must be >= 1, got {self.seq_len}") + if self.time_step_min > self.time_step_max: + raise ValueError( + f"time_step_limit is inverted: min={self.time_step_min} > max={self.time_step_max}" + ) + if self.batch_size != 1: + raise ValueError( + f"batch_size ({self.batch_size}) must be 1: this lowering emits one sequence " + "per program. No emitter scales a row range or a state block by it, so " + "accepting it would silently produce batch-1 code. Batch by invoking the " + "program once per sequence." + ) + # FPRAM is a hardware-fixed 1024 slots (FPRAMAllocator's default, never + # overridden), and ssd_decay_mask_v0 needs chunk_size of them for one head's + # cs row -- there is no column-block loop, so the whole row must fit at once. + # With the constant block on top, that caps MLEN well below the ANALYTIC + # target of 2048. Reject it as a shape with a reason, rather than letting it + # surface as an allocator MemoryError thousands of ISA lines later. + fpram_slots = 1024 + constants = 7 + if self.chunk_size + constants > fpram_slots: + raise ValueError( + f"chunk_size ({self.chunk_size}) plus the {constants}-slot Mamba constant " + f"block exceeds the {fpram_slots}-slot FPRAM. ssd_decay_mask_v0 holds one " + "head's whole cs row in FPRAM (one S_MAP_FP_V, no column-block loop), so " + "chunk_size is bounded by FPRAM. Since validate() also requires " + "chunk_size == MLEN, this is equivalently an MLEN ceiling: the SSD prefill " + "path is not usable at the ANALYTIC target MLEN=2048." + ) + + +@dataclass(frozen=True) +class MambaFPConstants: + """The FPRAM scalars every Mamba emitter needs. + + Allocated once per layer by :meth:`ProgramMambaCommonMixin.mamba_fp_constants` + rather than per emitter, because ``FPRAMAllocator`` bump-allocates and every + re-allocation would walk further into FP_MEM. + """ + + zero: FPVar + one: FPVar + neg_one: FPVar + dt_min: FPVar + dt_max: FPVar + #: 1 / normalized_group_width, for the gated RMSNorm. + reci_group: FPVar + eps: FPVar + + def as_list(self) -> list[FPVar]: + return [self.zero, self.one, self.neg_one, self.dt_min, self.dt_max, self.reci_group, self.eps] + + +class ProgramMambaCommonMixin: + """Emitters shared by the Mamba-2 prefill and decode lowerings.""" + + # ======================================================================== + # Constants + # ======================================================================== + + def mamba_fp_constants( + self, + shape: Mamba2Shape | None = None, + *, + name_prefix: str = "mamba", + ) -> MambaFPConstants: + """Allocate (but do not initialise) the layer's FPRAM scalar block. + + ``shape`` is accepted for symmetry with :meth:`mamba_fp_constant_values` + and is unused here -- the block is seven slots regardless of geometry. + It is optional so callers with no Mamba shape (the KDA path) do not have + to fabricate one; before this it was ``Mamba2Shape.__new__(Mamba2Shape)``, + an uninitialised frozen dataclass that would ``AttributeError`` the day + this method started reading a field. + + The values come from the host through the FP_MEM preload image, the same + way the attention path seeds ``attn_scale`` / ``-inf`` / ``eps``. The + returned ``FPVar.address`` values tell the host where to write them; see + :meth:`mamba_fp_constant_values` for the matching value list. + """ + return MambaFPConstants( + zero=self.fp_var(f"{name_prefix}_zero"), + one=self.fp_var(f"{name_prefix}_one"), + neg_one=self.fp_var(f"{name_prefix}_neg_one"), + dt_min=self.fp_var(f"{name_prefix}_dt_min"), + dt_max=self.fp_var(f"{name_prefix}_dt_max"), + reci_group=self.fp_var(f"{name_prefix}_reci_group"), + eps=self.fp_var(f"{name_prefix}_eps"), + ) + + @staticmethod + def mamba_fp_constant_values( + shape: Mamba2Shape, + *, + norm_group_width: int | None = None, + eps: float = 1e-5, + ) -> list[float]: + """Host-side values matching :meth:`mamba_fp_constants`, in slot order. + + Kept next to the allocator so the two cannot drift: a mismatch here is a + silent wrong answer, because FPRAM contents are never checked at runtime. + """ + group_width = norm_group_width if norm_group_width is not None else shape.d_inner // shape.n_groups + if group_width <= 0: + raise ValueError(f"norm_group_width must be positive, got {group_width}") + dt_max = shape.time_step_max + # bf16's finite maximum; +inf would poison V_MIN_VF's clamp. + if not math.isfinite(dt_max): + dt_max = 3.3895313892515355e38 + return [0.0, 1.0, -1.0, shape.time_step_min, dt_max, 1.0 / group_width, eps] + + # ======================================================================== + # Single-row primitives + # + # The `tile_row_*` family operates on matching row indices across two + # matrices. Mamba's conv and state sweeps need *cross-row* moves (out[s] from + # x[s-3]), which `vram_add`/`vram_mul` already express through their row + # offsets. These three wrappers name the intent. + # ======================================================================== + + def mamba_row_copy(self, dst: VRAMMatrixVar, dst_row: int, src: VRAMMatrixVar, src_row: int): + """``dst[dst_row] = src[src_row]``. + + There is no VRAM->VRAM move opcode, so this is zero-then-add -- the + established idiom (``program_routed_moe.py`` uses the same one). + """ + self.vram_fill_zero(dst, rows=[dst_row]) + self.vram_add(dst, src, dst_row_offset=dst_row, src_row_offset=src_row, num_rows=1) + + def mamba_row_add(self, dst: VRAMMatrixVar, dst_row: int, src: VRAMMatrixVar, src_row: int): + """``dst[dst_row] += src[src_row]``.""" + self.vram_add(dst, src, dst_row_offset=dst_row, src_row_offset=src_row, num_rows=1) + + def mamba_row_mul(self, dst: VRAMMatrixVar, dst_row: int, src: VRAMMatrixVar, src_row: int): + """``dst[dst_row] *= src[src_row]``.""" + self.vram_mul(dst, src, dst_row_offset=dst_row, src_row_offset=src_row, num_rows=1) + + def mamba_block_copy( + self, + dst: VRAMMatrixVar, + src: VRAMMatrixVar, + *, + dst_row_offset: int = 0, + src_row_offset: int = 0, + num_rows: int | None = None, + ): + """Zero-then-add copy of a contiguous row range between two matrices.""" + rows = num_rows if num_rows is not None else src.shape[0] + self.vram_fill_zero(dst, rows=list(range(dst_row_offset, dst_row_offset + rows))) + self.vram_add( + dst, src, dst_row_offset=dst_row_offset, src_row_offset=src_row_offset, num_rows=rows + ) + + # ======================================================================== + # SiLU + # ======================================================================== + + def mamba_silu_v0( + self, + target: VRAMMatrixVar, + scratch: VRAMMatrixVar, + consts: MambaFPConstants, + *, + marker: str, + rows: Sequence[int] | None = None, + ): + """In-place ``target = target * sigmoid(target)``. + + `marker` is required and is emitted verbatim. This emitter serves two + stage vocabularies -- Mamba's and KDA's -- so there is no default that + is right for both, and the one it used to carry (``mamba_gated_norm``) + was wrong for every KDA caller. Markers are sticky, so a caller that + accepted the default billed its SiLU *and everything it emitted + afterwards* to a Mamba stage. + + Build it with ``mamba_stage_marker(stage, "silu")`` or + ``kda_stage_marker(stage, "silu")`` at the call site, which is the only + place that knows which vocabulary applies. + + Five vector ops per row after the copy, the same shape as + ``asm_templates/silu_asm.py``: negate, exp, +1, reciprocal, multiply. + `scratch` must be a distinct VRAM matrix of the same width; the sigmoid is + built there so `target` still holds `x` for the final multiply. + + The tails are safe without an explicit clamp: ``V_EXP_V`` saturates its + input at [-88, 88], so for very negative `x` the reciprocal underflows + toward zero and ``silu(x) -> 0``, which is the correct limit. + """ + resolved = list(range(target.shape[0])) if rows is None else list(rows) + self.emit_comment(marker) + self.mamba_block_copy(scratch, target, num_rows=target.shape[0]) + self.tile_row_mul_fp_broadcast(scratch, consts.neg_one, rows=resolved) + self.tile_row_exp(scratch, rows=resolved) + self.tile_row_add_fp_broadcast(scratch, consts.one, rows=resolved) + self.tile_row_reci(scratch, rows=resolved) + self.tile_row_mul(target, scratch, rows=resolved) + return target + + # ======================================================================== + # dt activation + # ======================================================================== + + def mamba_dt_activation_v0( + self, + dt: VRAMMatrixVar, + dt_bias: VRAMMatrixVar | None, + consts: MambaFPConstants, + shape: Mamba2Shape, + *, + rows: Sequence[int] | None = None, + ): + """In-place ``dt = clamp(softplus(dt + dt_bias), time_step_limit)``. + + `dt_bias` is a per-head static parameter and must already be resident in + VRAM with the same layout as `dt`. Pass None when the host has folded the + bias into the in_proj weight. + + The clamp reuses ``V_MAX_VF`` / ``V_MIN_VF``, the opcodes added for the + GPT-OSS expert clamp -- byte-for-byte the same shape. + """ + resolved = list(range(dt.shape[0])) if rows is None else list(rows) + self.emit_comment(mamba_stage_marker("mamba_dt", "softplus + clamp")) + if dt_bias is not None: + # dt_bias is ONE row broadcast across every row of dt. `tile_row_add` + # would be wrong here: it pairs matching row indices, so with a + # single-row bias every row past the first would read an unallocated + # source row and produce NaN -- silently, and only in the rows past + # the first, which is easy to miss in an aggregate error metric. + if dt_bias.shape[0] == 1: + for row in resolved: + self.vram_add(dt, dt_bias, dst_row_offset=row, src_row_offset=0, num_rows=1) + elif dt_bias.shape[0] > max(resolved): + self.tile_row_add(dt, dt_bias, rows=resolved) + else: + raise ValueError( + f"dt_bias has {dt_bias.shape[0]} rows; needs either 1 (broadcast) or more " + f"than {max(resolved)} -- tile_row_add pairs dst row r with SRC row r, not " + "with the i-th source row, so a non-zero-based `rows` needs the bias to " + "cover those same indices" + ) + self.tile_row_softplus(dt, rows=resolved) + # Skip degenerate clamps. softplus output is already strictly positive, so + # a max against 0.0 is a no-op that would only cost instructions. + if shape.time_step_min > 0.0: + self.tile_row_max_fp_broadcast(dt, consts.dt_min, rows=resolved) + if math.isfinite(shape.time_step_max): + self.tile_row_min_fp_broadcast(dt, consts.dt_max, rows=resolved) + return dt + + # ======================================================================== + # Causal depthwise conv1d + # ======================================================================== + + def mamba_conv1d_v0( + self, + x: VRAMMatrixVar, + conv_weight: VRAMMatrixVar, + conv_bias: VRAMMatrixVar | None, + out: VRAMMatrixVar, + scratch: VRAMMatrixVar, + shape: Mamba2Shape, + *, + num_rows: int | None = None, + history_rows: int = 0, + ): + """Causal depthwise conv1d of width ``conv_kernel`` along the sequence. + + Layout contract: `x` holds one VRAM row per timestep with **sequence on + the row axis** -- which is what ``doc/memory_layout.md``'s + ``[h//VLEN, b, s, VLEN]`` gives. A shift along `t` is then a pure row + offset (free), and the per-channel tap weight `w[k, :]` is exactly a + VLEN-wide ``V_MUL_VV`` operand. Depthwise conv is the one kernel this + layout is perfect for. + + Deliberately *not* routed through ``aten/ops/plena/conv_ops.py``: that + path is a dense im2col which hardcodes a square KxK spatial patch, has no + padding support (so it cannot express the causal left pad), and would + materialise a dense ``(M, C_in*K*K)`` GEMM against a block-diagonal + weight, wasting `conv_dim`-fold MACs. + + ``V_MUL_VF`` is the wrong tool for the tap weight: it broadcasts one FP + register across every lane, but the conv weight varies *per channel*, i.e. + per lane. Hence `conv_weight` as a VRAM matrix and ``V_MUL_VV``. + + `history_rows` is how many leading rows of `x` are carried-over history + from a previous call rather than new timesteps; outputs are produced for + rows ``[history_rows, history_rows + num_rows)``. With `history_rows == 0` + the first `conv_kernel - 1` outputs see the causal zero pad. + """ + k = shape.conv_kernel + rows = num_rows if num_rows is not None else x.shape[0] - history_rows + if conv_weight.shape[0] < k: + raise ValueError(f"conv_weight needs at least conv_kernel={k} rows, got {conv_weight.shape[0]}") + if history_rows + rows > x.shape[0]: + raise ValueError( + f"conv1d would read past x: history_rows={history_rows} + num_rows={rows} " + f"> x.shape[0]={x.shape[0]}" + ) + if rows > out.shape[0]: + raise ValueError(f"conv1d output needs {rows} rows, out has {out.shape[0]}") + + self.emit_comment(mamba_stage_marker("mamba_conv1d", f"k={k} rows={rows} history={history_rows}")) + for s in range(rows): + self.vram_fill_zero(out, rows=[s]) + for j in range(k): + # Tap j reads x[t - (k-1) + j]; taps before the sequence start are + # the causal zero pad and are skipped rather than masked, because + # V_MASK is HLEN-granular and cannot gate a single row anyway. + src = history_rows + s - (k - 1) + j + if src < 0: + continue + self.mamba_row_copy(scratch, 0, x, src) + self.mamba_row_mul(scratch, 0, conv_weight, j) + self.mamba_row_add(out, s, scratch, 0) + if conv_bias is not None: + self.mamba_row_add(out, s, conv_bias, 0) + return out + + # ======================================================================== + # Gated RMSNorm + # ======================================================================== + + def mamba_gated_rmsnorm_v0( + self, + y: VRAMMatrixVar, + z: VRAMMatrixVar, + norm_weight: VRAMMatrixVar | None, + gate_scratch: VRAMMatrixVar, + sq_scratch: VRAMMatrixVar, + rms_fp: FPVar, + consts: MambaFPConstants, + shape: Mamba2Shape, + *, + rows: Sequence[int] | None = None, + ): + """``y = RMSNorm(y) * silu(z)``, in place on `y`. + + Mamba-2 normalises over ``d_inner / n_groups``, not over ``hidden_size``, + so this cannot reuse ``asm_templates/normalization_asm.py`` unchanged -- + that template hardcodes its reduction bound as ``hidden_dim // vlen`` and + takes a single ``1/hidden_dim`` FPRAM slot. ``consts.reci_group`` carries + the group reciprocal instead. + + `rms_fp` must have room for one slot per row in `rows`; the caller owns it + so the allocation does not churn FPRAM once per call. + """ + resolved = list(range(y.shape[0])) if rows is None else list(rows) + if rms_fp.size < len(resolved): + raise ValueError(f"rms_fp needs {len(resolved)} slots, has {rms_fp.size}") + self.emit_comment(mamba_stage_marker("mamba_gated_norm", "rmsnorm(y * silu(z))")) + + # The gate is applied BEFORE the variance. Both upstream implementations do + # this -- HuggingFace's MambaRMSNormGated multiplies by silu(gate) and only + # then reduces, and mamba_ssm's Mamba2 uses norm_before_gate=False. The two + # orders are different functions, not different roundings. + self.mamba_block_copy(gate_scratch, z, num_rows=z.shape[0]) + self.mamba_silu_v0( + gate_scratch, sq_scratch, consts, rows=resolved, + marker=mamba_stage_marker("mamba_gated_norm", "silu"), + ) + self.tile_row_mul(y, gate_scratch, rows=resolved) + + # sum(y^2) per row -> rms_fp + self.mamba_block_copy(sq_scratch, y, num_rows=y.shape[0]) + self.tile_row_mul(sq_scratch, y, rows=resolved) + self.tile_row_sum(rms_fp, sq_scratch, rows=resolved, target_base_offset=0) + + # rms_fp = 1 / sqrt(sum/n + eps), computed in the scalar unit. There is no + # vector sqrt (S_FP_OP has SQRT, V_ELEMENT_OP does not), so per-row scalar + # work is the only option here. + self.mamba_rsqrt_fpram(rms_fp, consts, count=len(resolved)) + + self.tile_row_mul_fp(y, rms_fp, rows=resolved, fpram_base_offset=0) + if norm_weight is not None: + self.tile_row_mul(y, norm_weight, rows=resolved) + return y + + def mamba_rsqrt_fpram(self, acc: FPVar, consts: MambaFPConstants, *, count: int): + """In place ``acc[i] = 1 / sqrt(acc[i] * reci_group + eps)`` for i < count. + + Emitted as scalar FP ops because the ISA has no vector square root: + ``S_FP_OP`` carries ``SQRT_FP`` but ``V_ELEMENT_OP`` does not (see + ``doc/operation.svh``). + """ + from compiler.aten.isa_builder import IsaBuilder, fp, gp + + gp_regs = self._reg.allocate_gp(1) + # Allocate the FP registers rather than hardcoding f1-f3: RegisterAllocator + # hands them out descending (f7, f6, f5, ...), so a caller already holding + # several live FP registers would have had one silently clobbered here. + fp_regs = self.allocate_fp_reg(3) + (gp_addr,) = gp_regs + f_val, f_scale, f_eps = fp_regs + try: + asm = IsaBuilder().comment(f"Mamba rsqrt over FPRAM[{acc.address}:+{count}]") + asm.instr("S_ADDI_INT", gp(gp_addr), gp(0), consts.reci_group.address) + asm.instr("S_LD_FP", fp(f_scale), gp(gp_addr), 0) + asm.instr("S_ADDI_INT", gp(gp_addr), gp(0), consts.eps.address) + asm.instr("S_LD_FP", fp(f_eps), gp(gp_addr), 0) + for i in range(count): + asm.instr("S_ADDI_INT", gp(gp_addr), gp(0), acc.address + i) + asm.instr("S_LD_FP", fp(f_val), gp(gp_addr), 0) + asm.instr("S_MUL_FP", fp(f_val), fp(f_val), fp(f_scale)) # mean = sum / n + asm.instr("S_ADD_FP", fp(f_val), fp(f_val), fp(f_eps)) # + eps + asm.instr("S_SQRT_FP", fp(f_val), fp(f_val)) + asm.instr("S_RECI_FP", fp(f_val), fp(f_val)) + # normalization_asm.py documents a 4-instruction spacer after + # S_RECI_FP before the value is consumed. Use the same instruction it + # uses -- S_ADDI_INT gp0, gp0, 0, a write to the hardwired-zero GP + # register, i.e. an integer-pipe NOP. An S_ADD_FP would issue on the + # very scalar-FP unit whose latency the spacer exists to cover, and + # would need a register to write. + for _ in range(4): + asm.instr("S_ADDI_INT", gp(0), gp(0), 0) + asm.instr("S_ST_FP", fp(f_val), gp(gp_addr), 0) + return self._emit(asm) + finally: + self._reg.free_gp(gp_regs) + self.free_fp_reg(fp_regs) + + +__all__ = [ + "MAMBA_STAGES", + "MAMBA_STAGE_MARKER_PREFIX", + "Mamba2Shape", + "MambaFPConstants", + "ProgramMambaCommonMixin", + "mamba_stage_marker", +] diff --git a/aten/plena/program_ssd.py b/aten/plena/program_ssd.py new file mode 100644 index 0000000..584246a --- /dev/null +++ b/aten/plena/program_ssd.py @@ -0,0 +1,613 @@ +"""Mamba-2 SSD (state-space duality) chunked-scan lowering -- the prefill path. + +The chunked form of the Mamba-2 recurrence splits the sequence into chunks of +``chunk_size`` and, per chunk, computes + + G = C @ B^T [chunk, chunk] (per group) + L[i,j] = exp(cs_i - cs_j) for i >= j [chunk, chunk] (per head, causal) + Y_intra = (L o G) @ X [chunk, head_dim] (per head) + h_new = B^T @ X_d + h_prev * decay [state, head_dim] (per head) + Y_inter = C @ h_prev, row i scaled by exp(cs_i) + Y = Y_intra + Y_inter + D_h * X + +where ``a_t = A_h * dt_t`` and ``cs = cumsum_t(a)`` within the chunk. + +Why this maps well onto PLENA +----------------------------- +Three of the four matmuls drop onto the existing flash-attention templates with +no structural change: ``G = C @ B^T`` is ``asm_templates/flashattn/qkt.py``'s +``M_TMM`` loop with ``state_size`` substituted for ``head_dim``; ``Y_intra`` is +``pv.py``'s ``M_MM`` loop with ``X`` in place of ``V``; and the inter-chunk +recurrence is ``output.py``'s ``computing_o_code`` with the scalar source changed +from the softmax max-delta to the chunk segment sum -- and *simpler*, because +attention needs a different scalar per row while Mamba needs one scalar for the +whole ``[state_size, head_dim]`` block, so the ``S_LD_FP`` hoists out of the loop. +Dropping the online softmax removes about ten of the fourteen instructions in +attention's per-row inner loop. + +The one structural obstacle +--------------------------- +``M_MM`` and ``M_TMM`` both contract over the **VRAM operand's lane axis** +(``matrix_machine.rs``). Three of the four SSD products satisfy that natively, but +the state accumulation ``h[n,p] = sum_t B[t,n] * X[t,p]`` contracts over the +**row** axis of both operands, which neither expresses. There is no VRAM +transpose, no transpose instruction, and ``H_PREFETCH_M`` does not transpose +either -- ``dma.rs`` gathers contiguous runs at a stride, a strided row gather; +the sentence in ``doc/plena_isa_spec.md`` about stride mode describes a host-side +storage convention, not DMA behaviour. + +The fix is to produce ``B^T`` at its source: :meth:`ssd_transposed_projection_v0` +stages a host-transposed weight into VRAM and the layer input into MRAM, then +emits ``M_TMM``. Nothing in the ISA forbids weight-in-VRAM / activation-in-MRAM; +it is purely a convention that no existing emitter happens to break. The cost is +bounded because only the *narrow* slices need it -- ``B^T`` is +``n_groups*state_size`` wide and ``dt^T`` is ``num_heads`` wide, while the wide +``x``/``z`` slices stay in their natural orientation. + +What this path does NOT support +------------------------------- +Enforced by :meth:`Mamba2Shape.validate`, and worth stating up front because the +module reads as if it were general: + +* ``chunk_size == head_dim == state_size == MLEN`` and ``d_inner/n_groups == MLEN``. + The ``tile_row_*`` family works within ONE MLEN column block, and one + ``S_MAP_FP_V`` moves exactly MLEN scalars, so anything wider leaves the later + column blocks untouched and reads uninitialised FPRAM. +* Consequently **this path is not usable at the ANALYTIC target MLEN=2048.** + ``ssd_decay_mask_v0`` holds one head's whole ``cs`` row in FPRAM, and FPRAM is a + hardware-fixed 1024 slots, so ``chunk_size`` -- and therefore MLEN -- is capped + near 1017. Lifting it means looping ``tile_col_idx`` in every emitter and + blocking the cs row; neither is done. Note the analytic cost model in + ``analytic_models/performance/mamba2_model.py`` models the *unrestricted* shape + family, so its numbers describe a lowering this path currently rejects. +* ``hidden_size <= mram_tile_capacity * MLEN``. + :meth:`ssd_transposed_projection_v0` contracts over ``hidden_size`` and + ``vram_sub_projection_T_to`` has no K-split, so a wider hidden dimension is an + MRAM allocator overflow rather than a shape error. The module docstring's cost + argument bounds the OUTPUT dimension; the binding constraint is the contraction. +* ``batch_size == 1``. No emitter scales a row range or a state block by it. + +Numerical choices +----------------- +* **The cumsum goes through the systolic array**, not a chained ``V_ADD_VV``. + ``cs = a @ U`` against a constant lower-triangular-ones MRAM tile sums in the + f32 accumulator (all four accumulators in ``matrix_machine.rs`` are f32) and + rounds to bf16 exactly once at the ``M_MM_WO`` flush. A chained vector scan + would re-quantise to bf16 at every step, and ``cs`` feeds an *exponent*: at + ``|cs| ~ 30`` a 0.9% relative error is 0.27 absolute, i.e. a 31% error on the + decay. See :meth:`ssd_chunk_cumsum_v0`. +* **The decay matrix is never factored** as ``exp(cs_i) * exp(-cs_j)``. That form + saves about half the instructions and destroys the answer: the systolic array + computes the full LxL product including the upper triangle, where the two + factors multiply to ``e^{+60}`` and overflow bf16 *before* the causal mask is + applied. Instead ``D = cs_i - cs_j`` is formed directly, clamped to ``<= 0`` + with ``V_MIN_VF`` (free, and exact on the causal half because ``cs`` is + non-increasing), and only then exponentiated -- so every value fed to + ``V_EXP_V`` is non-positive and every result lands in ``(0, 1]``. +* **Activation round trips use the BF16 keyvalue path**, never MX-FP8. See the + module docstring of ``program_mamba_common.py``. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from compiler.aten.plena.program_mamba_common import ( + Mamba2Shape, + MambaFPConstants, + mamba_stage_marker, +) +from compiler.aten.plena.vars import FPVar, InputVar, VRAMMatrixVar +from compiler.utils.load_config import load_toml_config + +#: MRAM operand precision for a tensor the *host* staged into HBM. These arrive +#: through the ordinary MX-FP8 weight path that create_sim_env/create_mem_for_sim +#: writes, so reading them as BF16 would decode the wrong bytes -- garbage, which +#: shows up as NaN rather than as a small error. +HOST_STAGED = {"matrix_precision": "weights", "hbm_element_bytes": 1, "set_scale": True} + +#: MRAM operand precision for an *activation* the program spilled to HBM and is +#: reloading. +#: +#: This selects the ``keyvalue`` precision CLASS, not a width. It is BF16 only if +#: the build configures ``HBM_M_KV_TYPE`` / ``HBM_V_KV_TYPE`` as a Plain BF16 type +#: -- exactly the requirement ``linear_projection_bf16`` states at +#: ``program_matrix_ops.py``. Under the MX-FP8 e4m3 KV types that +#: ``plena_settings.toml`` ships, ``hbm_element_bytes=2`` only doubles the address +#: stride while the data is still decoded as 1-byte e4m3 plus a scale stream, and +#: ``set_scale=False`` leaves the scale register holding whatever the previous +#: load left. Call :meth:`ProgramSSDMixin.require_bf16_kv_precision` to make that +#: requirement a hard failure rather than silent garbage. +#: +#: Why BF16 matters here and not for a KV cache: the SSM state is a multiplicative +#: accumulator carried across every chunk, so requantising it to e4m3's 3 mantissa +#: bits at each chunk boundary compounds without bound. A KV cache is written once +#: and never read-modify-written. +SPILLED_ACTIVATION = {"matrix_precision": "keyvalue", "hbm_element_bytes": 2, "set_scale": False} + + +class ProgramSSDMixin: + """Mamba-2 chunked-scan (prefill) emitters.""" + + # ======================================================================== + # Build-level precision requirement + # ======================================================================== + + def require_bf16_kv_precision_from_active_build(self) -> None: + """Read the active TOML and check it, once per program. + + ``require_bf16_kv_precision`` puts the burden on a caller, and the + result was that **nothing ever called it** -- not the SSD emitters, not + the KDA ones, not a single test. That is the wrong division: the + ``keyvalue`` precision class is chosen by the *emitter*, in + ``SPILLED_ACTIVATION``, so the emitter is what knows the requirement + exists. A caller cannot be expected to know that passing the default + precision commits them to a build-level constraint. + + Cached on the instance: one spill-using emitter per layer would + otherwise re-read and re-parse the TOML per chunk per head. + """ + if getattr(self, "_bf16_kv_checked", False): + return + from compiler.aten.plena.compiler import _find_plena_settings_toml + + path = _find_plena_settings_toml() + if path is None or not path.exists(): + raise ValueError( + "The spilled-activation path needs the active PRECISION table to " + "verify the KV types are Plain BF16, and no plena_settings.toml " + "was found (set PLENA_SETTINGS_TOML). Refusing to proceed without " + "having checked -- under an MX KV type the spills decode as e4m3 " + "and produce a wrong answer rather than an error." + ) + self.require_bf16_kv_precision( + load_toml_config(path, "PRECISION", mode="TRANSACTIONAL") + ) + self._bf16_kv_checked = True + + def require_bf16_kv_precision(self, settings: dict | None = None) -> None: + """Fail unless the build configures the KV precision classes as Plain BF16. + + ``SPILLED_ACTIVATION`` selects the ``keyvalue`` class; the *width* comes + from the active TOML. Under the shipped MX-FP8 KV types the spilled + activations decode as e4m3 and the numerical argument in this module's + docstring does not hold. Nothing detects that at runtime -- it surfaces as + a wrong answer, not as an error -- so the check has to be explicit. + + `settings` is the parsed ``[.PRECISION]`` table; when None the check + reports that it could not verify rather than silently passing, because a + check that cannot run must not look like a check that passed. + """ + if settings is None: + raise ValueError( + "require_bf16_kv_precision needs the active PRECISION table; pass the " + "parsed [.PRECISION] section. Refusing to report success without " + "having checked -- SPILLED_ACTIVATION is silently wrong under MX-FP8 KV " + "types and produces a wrong answer rather than an error." + ) + bad = [] + for key in ("HBM_M_KV_TYPE", "HBM_V_KV_TYPE"): + node = settings.get(key, {}) + # The TOML spells this `format`, not `kind` -- see + # plena_settings.toml's [*.PRECISION.*_TYPE] tables, where a Plain + # entry is `format = "Plain"` with a `.DATA_TYPE` subtable and an Mx + # entry is `format = "Mx"` with `.ELEM` and `.SCALE`. Reading `kind` + # returned None for *every* configuration, so this rejected the + # Plain BF16 build it exists to require as well as the MX one it + # exists to refuse -- a check that could not pass, which is why + # nothing called it. + fmt = node.get("format") if isinstance(node, dict) else None + if fmt != "Plain": + bad.append(f"{key}={node}") + if bad: + raise ValueError( + "The Mamba SSD path spills activations through the keyvalue precision class " + "and requires it to be Plain BF16, but this build declares: " + + "; ".join(bad) + + ". Under an MX type the spilled operands decode as 1-byte e4m3 plus a " + "scale stream while the address stride assumes 2 bytes -- garbage, silently. " + "Either configure the KV types as Plain Fp(8,7) for this build, or pass " + "precision=HOST_STAGED and accept e4m3 for the state (see the module " + "docstring for why that compounds)." + ) + + # ======================================================================== + # Constant tiles + # ======================================================================== + + def ssd_lower_triangular_ones(self, chunk: int, *, name: str = "ssd_tri_ones") -> InputVar: + """Declare the ``[chunk, chunk]`` lower-triangular-ones constant in HBM. + + Staged as an HBM constant and prefetched, *not* materialised on chip: the + attention path's ``_build_causal_score_mask`` emits ``MLEN*MLEN`` fully + unrolled ``S_ST_FP`` instructions, which is 4096 at MLEN=64 and 4.2 + million at the target MLEN=2048. The host writes the values; see + :meth:`ssd_lower_triangular_ones_values`. + """ + return self.input(name, (chunk, chunk)) + + @staticmethod + def ssd_lower_triangular_ones_values(chunk: int) -> list[list[float]]: + """Host-side contents of :meth:`ssd_lower_triangular_ones`. + + ``U[s, t] = 1`` iff ``s <= t``. Used two ways, which is why it is one + tile and not two: as the cumsum operator (``cs[t] = sum_{s<=t} a[s]``), + and transposed-by-indexing as the causal mask on the score matrix. + """ + return [[1.0 if s <= t else 0.0 for t in range(chunk)] for s in range(chunk)] + + # ======================================================================== + # Chunk cumsum + # ======================================================================== + + def _blocks(self, n: int) -> int: + """Number of MLEN-wide tile blocks spanning `n` elements. + + The projection helpers index tensors in MLEN x MLEN *blocks*, not logical + rows -- ``vram_sub_projection_to(v, i, m, j, t, r, c)`` computes one output + tile from block row `i` of `v` and block column `j` of `m`. Getting this + wrong is a `KeyError` deep inside the layout, not a wrong answer, but the + helper keeps the arithmetic in one place. + """ + return max(1, -(-n // self.mlen)) + + def ssd_chunk_cumsum_v0( + self, + a_t: VRAMMatrixVar, + tri_ones: InputVar, + cs_t: VRAMMatrixVar, + shape: Mamba2Shape, + *, + precision: dict | None = None, + ): + """``cs_t[h, t] = sum_{s <= t} a_t[h, s]``, all heads in one pass. + + `a_t` is ``[num_heads, chunk]`` -- heads on rows, time on lanes -- which is + what the transposed ``dt`` projection produces. The sum along `t` is then a + contraction over the VRAM operand's lane axis, exactly what ``M_MM`` does, + against the lower-triangular-ones tile in MRAM. One matmul covers MLEN + heads at a time. + + This replaces the missing prefix scan, and does it *better than* the + instruction PLENA declares but does not implement. ``V_PS_V`` (0x31) would + not help even if wired up: its micro-op is ``PREFIX_SCAN_V_ELEMENT``, an + intra-row scan over a VLEN tile, so under any layout that keeps heads on + lanes it would cumsum across heads, which is meaningless. Today it also + assembles silently and then decodes to ``Invalid``, so emitting it is worse + than not emitting it. + + Routing the sum through the systolic array is also the numerically right + choice: all four accumulators in the matrix machine are f32, so the whole + chunk sums in f32 and rounds to bf16 exactly once at the ``M_MM_WO`` + flush. A chained ``V_ADD_VV`` scan would re-quantise every step, and `cs` + feeds an exponent, so its error is magnified rather than averaged. + """ + precision = HOST_STAGED if precision is None else precision + self.emit_comment( + mamba_stage_marker( + "mamba_chunk_cumsum", f"heads={shape.num_heads} chunk={shape.chunk_size}" + ) + ) + head_blocks = self._blocks(shape.num_heads) + time_blocks = self._blocks(shape.chunk_size) + for hb in range(head_blocks): + for tb in range(time_blocks): + self.vram_sub_projection_to( + a_t, + hb, + tri_ones, + tb, + cs_t, + hb, + tb, + **precision, + ) + return cs_t + + # ======================================================================== + # Decay mask + # ======================================================================== + + def ssd_decay_mask_v0( + self, + cs_t: VRAMMatrixVar, + cs_fp: FPVar, + decay: VRAMMatrixVar, + causal: VRAMMatrixVar, + consts: MambaFPConstants, + shape: Mamba2Shape, + *, + head_row: int, + ): + """Build ``decay[i, j] = exp(min(cs_i - cs_j, 0)) * causal[i, j]`` for one head. + + `cs_t` is ``[num_heads, chunk]``; row `head_row` holds this head's + cumulative sums with time on lanes. `cs_fp` receives that row via a single + ``S_MAP_FP_V`` -- the whole reason that opcode exists. Without it each + ``cs_i`` would cost a one-hot ``V_MUL_VV`` + ``V_RED_SUM`` + ``S_ST_FP`` + triple, roughly 15,000 instructions per chunk at 80 heads just to move + 5,120 numbers from the lane domain to the scalar domain -- more than the + matmuls those numbers feed. + + Then per output row `i`: copy ``cs_j`` across, reverse-subtract the scalar + ``cs_i`` to get ``cs_i - cs_j``, clamp to ``<= 0``, exponentiate, and + multiply by the causal 0/1 row. + + The clamp is what makes this safe. ``cs`` is a cumsum of ``A_h * dt`` + with ``A_h < 0``, so it is non-increasing and ``cs_i - cs_j <= 0`` + wherever ``i >= j`` -- the clamp is a no-op on the causal half and merely + stops the (about-to-be-masked) upper triangle from overflowing ``V_EXP_V``. + """ + chunk = shape.chunk_size + if cs_fp.size < chunk: + raise ValueError(f"cs_fp needs {chunk} slots for one chunk row, has {cs_fp.size}") + self.emit_comment(mamba_stage_marker("mamba_decay_mask", f"head_row={head_row}")) + + # One instruction: whole cs row -> FPRAM. + self.tile_row_to_fpram(cs_t, cs_fp, rows=[head_row]) + + for i in range(chunk): + self.mamba_row_copy(decay, i, cs_t, head_row) + # decay[i] = cs_i - cs_j (rorder=1 reverse subtract) + self.tile_row_sub_fp_broadcast(decay, cs_fp, rows=[i], fpram_offset=i, reverse=True) + self.tile_row_min_fp_broadcast(decay, consts.zero, rows=[i]) + self.tile_row_exp(decay, rows=[i]) + self.mamba_row_mul(decay, i, causal, i) + return decay + + # ======================================================================== + # Transposed projection + # ======================================================================== + + def ssd_transposed_projection_v0( + self, + weight_t_vram: VRAMMatrixVar, + x_mram: InputVar, + target: VRAMMatrixVar, + shape: Mamba2Shape, + *, + out_rows: int, + stage: str, + precision: dict | None = None, + ): + """``target[n, t] = sum_k weight_t_vram[n, k] * x_mram[t, k]`` via ``M_TMM``. + + The operand roles are inverted relative to every other emitter in the + repo: the *weight* (host-transposed, so ``[out, hidden]``) sits in VRAM + and the layer *input* sits in MRAM. That is what makes ``B^T`` and + ``dt^T`` -- with the sequence axis on lanes -- expressible at all; see the + module docstring. Nothing in the ISA forbids it: ``H_PREFETCH_V`` loads + any HBM region into VRAM and ``H_PREFETCH_M`` loads any HBM region into + MRAM. + + Use this only for the narrow slices (``B``, ``dt``). The wide ``x``/``z`` + slices should go through the ordinary ``linear_projection``. + """ + precision = SPILLED_ACTIVATION if precision is None else precision + if precision is SPILLED_ACTIVATION: + self.require_bf16_kv_precision_from_active_build() + self.emit_comment(mamba_stage_marker(stage, f"transposed, out_rows={out_rows}")) + out_blocks = self._blocks(out_rows) + time_blocks = self._blocks(shape.chunk_size) + for nb in range(out_blocks): + for tb in range(time_blocks): + self.vram_sub_projection_T_to( + weight_t_vram, + nb, + x_mram, + tb, + target, + nb, + tb, + **precision, + ) + return target + + # ======================================================================== + # One chunk, one head + # ======================================================================== + + def ssd_chunk_head_v0( + self, + *, + b_chunk: InputVar, + c_chunk: VRAMMatrixVar, + x_chunk: InputVar, + decay: VRAMMatrixVar, + scores: VRAMMatrixVar, + y_out: VRAMMatrixVar, + shape: Mamba2Shape, + head_block_base: int = 0, + precision: dict | None = None, + ): + """``Y_intra = (decay o (C @ B^T)) @ X`` for one head of one chunk. + + Two matmuls with an elementwise multiply between them, all reusing the + flash-attention shapes: + + * ``scores = C @ B^T`` -- ``M_TMM`` with `C` natural in VRAM and `B` + natural in MRAM (``qkt.py``'s form with ``state_size`` for ``head_dim``). + `B` has to be spilled to HBM and re-prefetched because MRAM is writable + only by ``H_PREFETCH_M``; ``qkt.py`` already does exactly this for `K`. + * ``scores *= decay`` -- the per-head fan-out point. ``scores`` is shared + across a group (``B`` and ``C`` are group-shared) but ``decay`` is + per-head, so this multiply is what makes the result head-specific. + * ``y_out = scores @ X`` -- ``M_MM``, ``pv.py``'s form with `X` for `V`. + """ + chunk = shape.chunk_size + precision = SPILLED_ACTIVATION if precision is None else precision + if precision is SPILLED_ACTIVATION: + self.require_bf16_kv_precision_from_active_build() + self.emit_comment(mamba_stage_marker("mamba_intra_chunk", f"chunk={chunk}")) + i_blocks = self._blocks(chunk) + j_blocks = self._blocks(chunk) + p_blocks = self._blocks(shape.head_dim) + + # scores = C @ B^T + for ib in range(i_blocks): + for jb in range(j_blocks): + self.vram_sub_projection_T_to( + c_chunk, + ib, + b_chunk, + jb, + scores, + ib, + jb, + **precision, + ) + + # scores *= decay -- the per-head fan-out point. + self.tile_row_mul(scores, decay, rows=list(range(chunk))) + + # Y_intra = scores @ X + for ib in range(i_blocks): + for pb in range(p_blocks): + self.vram_sub_projection_to( + scores, + ib, + x_chunk, + pb, + y_out, + head_block_base + ib, + pb, + **precision, + ) + return y_out + + # ======================================================================== + # Inter-chunk state + # ======================================================================== + + def ssd_state_update_v0( + self, + *, + state: VRAMMatrixVar, + b_t_chunk: VRAMMatrixVar, + x_d_chunk: InputVar, + decay_fp: FPVar, + shape: Mamba2Shape, + decay_offset: int = 0, + precision: dict | None = None, + ): + """``state = state * exp(sum_t a_t) + B^T @ X_d`` for one head. + + The decay is a single per-head scalar applied to the whole + ``[state_size, head_dim]`` block, so the ``S_LD_FP`` hoists out of the row + loop entirely -- strictly cheaper than the attention analogue in + ``output.py``, which needs a different scalar per query row. + + ``B^T @ X_d`` is the contraction that motivated the transposed projection: + `b_t_chunk` must be ``[state_size, chunk]`` with time on lanes, produced by + :meth:`ssd_transposed_projection_v0`, and `x_d_chunk` the decay-scaled + ``[chunk, head_dim]`` activation spilled to HBM. + """ + precision = SPILLED_ACTIVATION if precision is None else precision + if precision is SPILLED_ACTIVATION: + self.require_bf16_kv_precision_from_active_build() + self.emit_comment(mamba_stage_marker("mamba_state_update", f"state_size={shape.state_size}")) + n_blocks = self._blocks(shape.state_size) + p_blocks = self._blocks(shape.head_dim) + + # The projection *overwrites* its target tile (M_MM accumulates in the + # systolic array, then M_MM_WO writes and clears), so the new contribution + # has to land in scratch and be added -- writing straight into `state` + # would drop the decayed history. + contrib = self.alloc("ssd_state_contrib", shape.state_size, shape.head_dim, strict=False) + try: + for nb in range(n_blocks): + for pb in range(p_blocks): + self.vram_sub_projection_to( + b_t_chunk, + nb, + x_d_chunk, + pb, + contrib, + nb, + pb, + **precision, + ) + state_rows = list(range(state.shape[0])) + self.tile_row_mul_fp_broadcast(state, decay_fp, rows=state_rows, fpram_offset=decay_offset) + self.vram_add(state, contrib, num_rows=min(state.shape[0], contrib.shape[0])) + finally: + self.free_tensor(contrib) + return state + + def ssd_inter_chunk_output_v0( + self, + *, + c_chunk: VRAMMatrixVar, + state_prev: InputVar, + y_out: VRAMMatrixVar, + cs_fp: FPVar, + shape: Mamba2Shape, + head_row_base: int = 0, + precision: dict | None = None, + ): + """``y_out[i] += (C @ h_prev)[i] * exp(cs_i)`` for one head. + + ``M_MM`` with `C` natural in VRAM (state index on lanes) and `h_prev` + prefetched into MRAM as ``[state_size, head_dim]``. That orientation is + exactly what :meth:`ssd_state_update_v0` produces, which is why the two + compose without a transpose. + + The per-row ``exp(cs_i)`` scale comes from the FPRAM copy of the cs row, + so it costs one ``S_LD_FP`` + one ``V_MUL_VF`` per row and no extraction. + """ + chunk = shape.chunk_size + precision = SPILLED_ACTIVATION if precision is None else precision + if precision is SPILLED_ACTIVATION: + self.require_bf16_kv_precision_from_active_build() + self.emit_comment(mamba_stage_marker("mamba_inter_chunk", f"chunk={chunk}")) + i_blocks = self._blocks(chunk) + p_blocks = self._blocks(shape.head_dim) + scratch = self.alloc("ssd_inter_scratch", chunk, shape.head_dim, strict=False) + try: + for ib in range(i_blocks): + for pb in range(p_blocks): + self.vram_sub_projection_to( + c_chunk, + ib, + state_prev, + pb, + scratch, + ib, + pb, + **precision, + ) + # scale row i by exp(cs_i): the cs row is already in FPRAM, so this is + # a plain per-row scalar walk with no extraction. + self.tile_row_exp_fpram_scale(scratch, cs_fp, rows=list(range(chunk))) + self.vram_add(y_out, scratch, dst_row_offset=head_row_base, src_row_offset=0, num_rows=chunk) + finally: + self.free_tensor(scratch) + return y_out + + def tile_row_exp_fpram_scale(self, target: VRAMMatrixVar, cs_fp: FPVar, rows: Sequence[int]): + """``target[i] *= exp(cs_fp[i])``. + + The exponential is taken in the scalar unit (``S_EXP_FP``) because the + value is already a scalar; doing it in the vector unit would mean + splatting it to a row first. + """ + from compiler.aten.isa_builder import IsaBuilder, fp, gp + + resolved = list(rows) + gp_regs = self._reg.allocate_gp(2) + gp_addr, gp_row = gp_regs + try: + base = self._tile_addr(target.name) + asm = IsaBuilder().comment(f"Row scale by exp(FPRAM[{cs_fp.address}:+{len(resolved)}])") + for i, row in enumerate(resolved): + asm.instr("S_ADDI_INT", gp(gp_addr), gp(0), cs_fp.address + i) + asm.instr("S_LD_FP", fp(1), gp(gp_addr), 0) + asm.instr("S_EXP_FP", fp(1), fp(1)) + asm.instr("S_ADDI_INT", gp(gp_row), gp(0), base + row * self.mlen) + asm.instr("V_MUL_VF", gp(gp_row), gp(gp_row), fp(1), 0) + return self._emit(asm) + finally: + self._reg.free_gp(gp_regs) + + +__all__ = ["ProgramSSDMixin"] diff --git a/aten/plena/program_ssm_recurrent.py b/aten/plena/program_ssm_recurrent.py new file mode 100644 index 0000000..bb944af --- /dev/null +++ b/aten/plena/program_ssm_recurrent.py @@ -0,0 +1,272 @@ +"""Mamba-2 recurrent decode step (`seq_len == 1`) and cross-token state residency. + +At `seq_len == 1` the chunked scan collapses to the plain recurrence + + h[h, n, :] = dA_h * h[h, n, :] + B[g, n] * (dt_h * x[h, :]) + y[h, :] = sum_n C[g, n] * h[h, n, :] + D_h * x[h, :] + +which is pure elementwise work on the state plus two contractions. No matmul, no +chunking, no decay matrix. + +State layout -- the decision everything else follows from +-------------------------------------------------------- +The state is laid out ``[head][state][head_dim]``: one VRAM row per +``(head, state index)`` pair, with **head_dim on the lane axis**. Under that +layout all three terms become existing opcodes: + +* the decay ``dA_h`` is constant across a head's ``state_size`` consecutive rows, + so it is a ``V_MUL_VF`` broadcast whose FPRAM row map is an arithmetic + progression with step 0 -- one ``S_LD_FP`` per head, not per row; +* the input term is rank-1, but folding ``dt_h`` into ``x`` *first* reduces the + distinct per-row scalars from ``num_heads*head_dim*state_size`` to just + ``n_groups*state_size``, which is the difference between fitting in FPRAM and + not. It then becomes ``V_MUL_VF`` + ``V_ADD_VV`` per row; +* the ``sum_n`` contraction becomes a **cross-row** accumulation (``V_ADD_VV`` + into one output row per head) instead of an intra-row reduction. That matters: + the transposed layout would need ``V_RED_SUM`` per ``(head, head_dim)`` pair, + producing ``num_heads*head_dim`` scalars that overflow FPRAM. + +This layout requires ``VLEN <= head_dim``, which :meth:`Mamba2Shape.validate` +already enforces (as ``head_dim % mlen == 0``). + +What this path does *not* do +---------------------------- +It does not batch the per-head work through ``M_BMV``. That instruction's +broadcast runs the wrong way: it reads one VRAM row as ``broadcast_amount`` lane +groups and multiplies **every** group by one **shared** MRAM slice. Mamba-2 +decode needs the opposite -- a per-head state *matrix* contracted against a +group-shared ``C`` vector. Two further disqualifiers: the state is an activation +and MRAM is writable only by ``H_PREFETCH_M``, so any ``M_*`` form on ``h`` costs +a full HBM round trip per token; and every ``M_B*`` form silently multiplies its +result by a hardwired ``bmm_scale = 0.25`` that no opcode exposes. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from compiler.aten.plena.program_mamba_common import ( + Mamba2Shape, + MambaFPConstants, + mamba_stage_marker, +) +from compiler.aten.plena.vars import FPVar, InputVar, VRAMMatrixVar + + +class ProgramSSMRecurrentMixin: + """Mamba-2 single-token decode emitters, plus cross-token state residency.""" + + # ======================================================================== + # Cross-token persistent state + # ======================================================================== + + def pin_hbm_region(self, name: str, size: int, hbm_element_bytes: int = 2) -> int: + """Reserve an HBM range that ``_allocate_hbm`` will never hand out. + + This is the first persistent-state mechanism in the compiler: attention + takes K/V as ordinary per-block ``InputVar``s and nothing in the repo + survives across program regions (a repo-wide grep for ``kv_cache`` / + ``past_key_value`` / ``persistent`` returns no functional hits). Mamba + decode needs it because the SSM state and the conv state must be read at + the top of a step and written back at the bottom, at the *same* address, + for the next token to find. + + Returns the pinned base address. Implemented by bumping the auto-allocator + past the region rather than by adding a second allocator, so a later + ``input()``/``store()`` with no explicit address cannot collide with it. + + `hbm_element_bytes` defaults to 2 because that is what + :meth:`ssm_store_state_v0` writes; sizing the region with the MX layout + instead would under-reserve it by ~78% and the write-back would land on the + next tensor. + """ + if size <= 0: + raise ValueError(f"pinned HBM region {name!r} must have positive size, got {size}") + addr = self._allocate_hbm(self.hbm_tensor_size(size, hbm_element_bytes=hbm_element_bytes)) + if not hasattr(self, "_pinned_hbm_regions"): + self._pinned_hbm_regions: dict[str, tuple[int, int]] = {} + self._pinned_hbm_regions[name] = (addr, size) + self.emit_comment(f"pinned HBM region {name}: [{addr}, {addr + size}) elements") + return addr + + def pinned_hbm_region(self, name: str) -> tuple[int, int]: + regions = getattr(self, "_pinned_hbm_regions", {}) + if name not in regions: + raise KeyError(f"no pinned HBM region named {name!r}; call pin_hbm_region first") + return regions[name] + + def ssm_load_state_v0( + self, + name: str, + rows: int, + cols: int, + hbm_addr: int, + ) -> VRAMMatrixVar: + """Prefetch a pinned state region from HBM into VRAM. + + Declared with ``real_data_ratio=1.0`` and loaded through the BF16 path: + the SSM state is a multiplicative accumulator carried across the whole + sequence, and MX-FP8's 3 mantissa bits would compound. Quantisation error + in a decaying accumulator is amplified by ``1 / sqrt(1 - lambda^2)`` where + ``lambda = exp(A*dt)``, so at ``lambda = 0.99`` a 3.6% element error + becomes roughly 25% state error -- and ``lambda -> 1`` is exactly the + long-memory regime that justifies using an SSM at all. A KV cache has no + such amplification because it is written once and never read-modify-written. + """ + self.emit_comment(mamba_stage_marker("mamba_state_load", f"{name} [{rows},{cols}]")) + var = self.input(name, (rows, cols), hbm_addr=hbm_addr, real_data_ratio=1.0) + # Must mirror ssm_store_state_v0 exactly: same precision class (KeyValue) + # and same bytes per element (2). Until load_batch took these, it emitted + # Activation/1-byte with half the row stride and a different scale-section + # base, so the state read back was not the state written -- and the failure + # mode is a wrong answer, not an error. + return self.load_batch( + var, name=f"{name}_vram", storage_precision=2, precision=1 + ) + + def ssm_store_state_v0(self, state: VRAMMatrixVar, name: str, hbm_addr: int) -> InputVar: + """Write the updated state back over the same pinned HBM range.""" + self.emit_comment(mamba_stage_marker("mamba_state_store", name)) + return self.store( + state, + name=name, + hbm_addr=hbm_addr, + precision=1, + hbm_element_bytes=2, + real_data_ratio=1.0, + ) + + # ======================================================================== + # Decode step + # ======================================================================== + + def ssm_decode_step_v0( + self, + *, + state: VRAMMatrixVar, + x: VRAMMatrixVar, + b_fp: FPVar, + c_fp: FPVar, + da_fp: FPVar, + dt_fp: FPVar, + d_fp: FPVar, + y: VRAMMatrixVar, + scratch: VRAMMatrixVar, + shape: Mamba2Shape, + consts: MambaFPConstants, + head_rows: Sequence[int] | None = None, + ): + """One recurrent step, all heads. + + Arguments carry the per-head / per-(group, state) scalars in FPRAM: + `da_fp[h]` is ``exp(A_h * dt_h)``, `dt_fp[h]` is ``dt_h``, `d_fp[h]` is the + skip coefficient ``D_h``, and `b_fp` / `c_fp` hold ``B[g, n]`` / ``C[g, n]`` + laid out ``n``-major within each group. They get there via ``S_MAP_FP_V`` + (see :meth:`ssm_decode_scalars_to_fpram_v0`). + + `state` is ``[num_heads * state_size, head_dim]`` in the layout described + in the module docstring; `x` is ``[num_heads, head_dim]``; `y` is + ``[num_heads, head_dim]``. + """ + heads = list(range(shape.num_heads)) if head_rows is None else list(head_rows) + n_state = shape.state_size + self.emit_comment( + mamba_stage_marker("mamba_state_update", f"decode heads={len(heads)} state={n_state}") + ) + + for h in heads: + group = h // shape.heads_per_group + base = h * 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], not dt_h * 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) + + # h[h, n, :] = dA_h * h[h, n, :] + B[g, n] * xs + # + # The rank-1 update walks the state rows against a pinned scratch + # row 0, so both walks are progressions and the whole thing is one + # hardware loop. The copy/multiply/add triple it replaces could not + # loop: it needed a scratch row per step, and that row was the same + # row every iteration, which broke the progression. + state_rows = list(range(base, base + n_state)) + self.tile_row_mul_fp_broadcast(state, da_fp, rows=state_rows, fpram_offset=h) + self.tile_row_fma_fp_sweep( + state, scratch, b_fp, + dst_rows=state_rows, src_rows=[0] * n_state, + fpram_offset=group * n_state, + ) + + # y[h] = sum_n C[g, n] * h[h, n, :] -- the mirror: source walks, + # destination pinned. + self.vram_fill_zero(y, rows=[h]) + self.tile_row_fma_fp_sweep( + y, state, c_fp, + dst_rows=[h] * n_state, src_rows=state_rows, + fpram_offset=group * n_state, + ) + + # y[h] += D_h * x[h] -- one row, one slot, so a broadcast FMA. This + # is where the accumulate saves a whole scratch copy on its own. + 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 + + def ssm_decode_scalars_to_fpram_v0( + self, + *, + source: VRAMMatrixVar, + target: FPVar, + rows: Sequence[int], + stage: str, + ): + """Move computed per-head / per-state scalars from VRAM lanes into FPRAM. + + One ``S_MAP_FP_V`` per row, MLEN scalars each. Every value the decode step + broadcasts (``dA_h``, ``dt_h``, ``B[g,n]``, ``C[g,n]``) is produced on chip + by the in_proj and lands in a VRAM row, but ``V_MUL_VF`` can only broadcast + from an FP register -- and before ``S_MAP_FP_V`` the only VRAM-to-scalar + path was ``V_RED_SUM``/``V_RED_MAX``, which collapse the whole row. The + one-hot workaround costs 3 instructions per scalar; this costs 1 per row. + """ + self.emit_comment(mamba_stage_marker(stage, "vram lanes -> FPRAM")) + self.tile_row_to_fpram(source, target, rows=list(rows)) + return target + + def ssm_conv_state_roll_v0( + self, + conv_state: VRAMMatrixVar, + new_row_src: VRAMMatrixVar, + new_row_idx: int, + shape: Mamba2Shape, + ): + """Shift the rolling conv1d history by one timestep and append the new one. + + `conv_state` holds ``conv_kernel - 1`` rows of history, oldest first. The + shift is a physical copy: address immediates are baked at ASM-gen time and + there is no data-dependent addressing (no compare, no branch), so a + runtime ring pointer is not expressible. It is cheap -- ``conv_kernel - 1`` + row copies per token. + + ``V_SHFT_V`` is the wrong tool twice over: it shifts *lanes* within one + row (the channel axis here), not rows, and its direction is contradictory + between ``doc/plena_isa_spec.md`` (left) and the emulator's + ``vector_machine.rs`` (right). + """ + 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(mamba_stage_marker("mamba_conv1d", f"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 + + +__all__ = ["ProgramSSMRecurrentMixin"] diff --git a/aten/tests/test_mamba_stage_contract.py b/aten/tests/test_mamba_stage_contract.py new file mode 100644 index 0000000..a31bae0 --- /dev/null +++ b/aten/tests/test_mamba_stage_contract.py @@ -0,0 +1,95 @@ +"""The Mamba stage-marker vocabulary is a cross-repo contract; pin both ends. + +The compiler emits ``; @stage=`` comments and the emulator's ``StageKind`` +(``transactional_emulator/src/stage_profile.rs``) must know every one of them. +Markers are *authoritative and sticky*: once a program contains any marker the +emulator abandons its legacy substring rules entirely and bills every instruction +to the most recent marker. So a name the compiler emits but the emulator lacks +does not error -- it lands in ``unresolved_stage_markers`` and that region's +cycles are attributed to whatever stage preceded it. Silent, and it looks exactly +like a correct profile. + +This test owns the compiler end. The Rust test +``stage_marker_names_match_the_compiler_vocabulary`` owns the other end and +carries the same list, so a stage added on one side without the other fails a +build rather than skewing a profile. +""" + +from __future__ import annotations + +import unittest + +from compiler.aten.plena.program_mamba_common import ( + MAMBA_STAGE_MARKER_PREFIX, + MAMBA_STAGES, + mamba_stage_marker, +) + +#: Mirrors the Mamba entries of ``COMPILER_MOE_STAGES`` in +#: ``transactional_emulator/src/stage_profile.rs``. +EMULATOR_MAMBA_STAGES = frozenset( + { + "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", + } +) + + +class TestMambaStageContract(unittest.TestCase): + def test_vocabulary_matches_the_emulator(self): + missing_in_emulator = MAMBA_STAGES - EMULATOR_MAMBA_STAGES + missing_in_compiler = EMULATOR_MAMBA_STAGES - MAMBA_STAGES + self.assertEqual( + missing_in_emulator, + set(), + "compiler emits stages the emulator's StageKind does not know; they would " + "land in unresolved_stage_markers and bill to the preceding stage", + ) + self.assertEqual( + missing_in_compiler, + set(), + "emulator declares Mamba stages no compiler marker produces; those rows " + "can never be reached and will always read zero", + ) + + def test_marker_prefix_matches_the_routed_moe_substrate(self): + # Both substrates share one prefix because the emulator has one parser. + from compiler.aten.plena.program_routed_moe import MOE_STAGE_MARKER_PREFIX + + self.assertEqual(MAMBA_STAGE_MARKER_PREFIX, MOE_STAGE_MARKER_PREFIX) + + def test_marker_names_do_not_collide_with_the_moe_vocabulary(self): + from compiler.aten.plena.program_routed_moe import MOE_STAGES + + self.assertEqual( + MAMBA_STAGES & MOE_STAGES, + set(), + "a name in both vocabularies would merge Mamba and MoE cost in the profile", + ) + + def test_unknown_stage_is_rejected_at_asm_generation_time(self): + # Failing here rather than at profile-read time is the whole point: a typo + # otherwise produces a plausible-looking profile with a missing region. + with self.assertRaises(ValueError): + mamba_stage_marker("mamba_not_a_stage") + + def test_marker_format_carries_optional_detail(self): + self.assertEqual(mamba_stage_marker("mamba_dt"), "@stage=mamba_dt") + self.assertEqual( + mamba_stage_marker("mamba_dt", "softplus + clamp"), "@stage=mamba_dt softplus + clamp" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/aten/tests/test_ssm_decode_step.py b/aten/tests/test_ssm_decode_step.py new file mode 100644 index 0000000..9b0652b --- /dev/null +++ b/aten/tests/test_ssm_decode_step.py @@ -0,0 +1,188 @@ +"""`ssm_decode_step_v0` against the textbook selective-scan recurrence. + +Written when Task 8 converted this kernel onto `V_FMA_VF` and the whole suite +stayed green -- because **nothing tested it**. `grep -rl ssm_decode_step_v0` +found the emitter, its caller in `aten/ops/plena/mamba_ops.py`, and a docstring +mention in the KDA recurrence. No test. A rewrite of the state update and the +output contraction was therefore unverified. + +The oracle is `aten/tests/isa_interpreter.py`, the same one KDA's decode step +uses, so this checks the emitted instructions rather than the Python that +emitted them. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.mamba2.reference import ( # noqa: E402 + mamba2_recurrent_reference, +) +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_mamba_common import Mamba2Shape # noqa: E402 +from compiler.aten.tests.isa_interpreter import ( # noqa: E402 + Machine, + UnsupportedInstruction, +) + +MLEN = 8 + + +def _rows_up(n: int) -> int: + return ((n + MLEN - 1) // MLEN) * MLEN + + +class _Harness: + """One decode step, driven through the ISA interpreter.""" + + def __init__(self, shape: Mamba2Shape): + self.shape = shape + p = self.prog = PlenaCompiler(mlen=MLEN, blen=2) + n = shape.state_size + self.state = p.alloc("state", _rows_up(shape.num_heads * n), MLEN) + self.x = p.alloc("x", _rows_up(shape.num_heads), MLEN) + self.y = p.alloc("y", _rows_up(shape.num_heads), MLEN) + self.scratch = p.alloc("scratch", MLEN, MLEN) + self.b_fp = p.fp_var("b", size=shape.n_groups * n) + self.c_fp = p.fp_var("c", size=shape.n_groups * n) + self.da_fp = p.fp_var("da", size=shape.num_heads) + self.dt_fp = p.fp_var("dt", size=shape.num_heads) + self.d_fp = p.fp_var("d", size=shape.num_heads) + self.consts = p.mamba_fp_constants() + self.mark = len(p.get_code()) + + def _base(self, var) -> int: + return self.prog.get_vram_layout(var.name).vram_base_addr + + def emit(self) -> str: + self.prog.ssm_decode_step_v0( + state=self.state, x=self.x, b_fp=self.b_fp, c_fp=self.c_fp, + da_fp=self.da_fp, dt_fp=self.dt_fp, d_fp=self.d_fp, y=self.y, + scratch=self.scratch, shape=self.shape, consts=self.consts, + ) + return self.prog.get_code()[self.mark :] + + def run(self, *, x, state, b, c, da, dt, d): + s = self.shape + code = self.emit() + m = Machine(vlen=MLEN, vram_words=1 << 16, fpram_words=1 << 13) + for h in range(s.num_heads): + m.write_vram_row(self._base(self.x) + h * MLEN, x[h].tolist()) + # Seed y with garbage: the contraction must clear it, and a fresh + # Machine is all zeros, which would hide a missing fill. + m.write_vram_row(self._base(self.y) + h * MLEN, [7.5] * MLEN) + for n in range(s.state_size): + m.write_vram_row( + self._base(self.state) + (h * s.state_size + n) * MLEN, + state[h, n].tolist(), + ) + m.write_fpram(self.b_fp.address, b.flatten().tolist()) + m.write_fpram(self.c_fp.address, c.flatten().tolist()) + m.write_fpram(self.da_fp.address, da.tolist()) + m.write_fpram(self.dt_fp.address, dt.tolist()) + m.write_fpram(self.d_fp.address, d.tolist()) + m.write_fpram(self.consts.zero.address, self.prog.mamba_fp_constant_values(self.shape)) + m.run(code) + + y = torch.tensor([ + m.read_vram_row(self._base(self.y) + h * MLEN, s.head_dim) + for h in range(s.num_heads) + ]) + new_state = torch.tensor([ + [ + m.read_vram_row( + self._base(self.state) + (h * s.state_size + n) * MLEN, s.head_dim + ) + for n in range(s.state_size) + ] + for h in range(s.num_heads) + ]) + return y, new_state, m + + +def _case(seed: int, *, heads: int, state_size: int, groups: int = 1): + torch.manual_seed(seed) + shape = Mamba2Shape( + hidden_size=heads * MLEN, num_heads=heads, head_dim=MLEN, + state_size=state_size, n_groups=groups, conv_kernel=4, + chunk_size=16, seq_len=1, + ) + x = torch.randn(heads, MLEN) + state = torch.randn(heads, state_size, MLEN) + b = torch.randn(groups, state_size) + c = torch.randn(groups, state_size) + dt = torch.rand(heads) * 0.5 + 0.1 + a = -torch.rand(heads) - 0.5 + d = torch.randn(heads) + + # The emitter takes dA = exp(A * dt) precomputed; the reference takes A. + y_ref, state_ref = mamba2_recurrent_reference( + x=x[None, None], dt=dt[None, None], A=a, + B=b.reshape(1, 1, groups, state_size), + C=c.reshape(1, 1, groups, state_size), + D=d, initial_state=state[None], + ) + return shape, { + "x": x, "state": state, "b": b, "c": c, "d": d, "dt": dt, + "da": torch.exp(a * dt), + "y_ref": y_ref[0, 0], "state_ref": state_ref[0], + } + + +@pytest.mark.parametrize( + "seed,heads,state_size,groups", + [(1, 1, 4, 1), (2, 2, 4, 1), (3, 4, 8, 2), (4, 2, 16, 1), (5, 6, 4, 3)], +) +def test_decode_step_matches_the_reference(seed, heads, state_size, groups): + shape, c = _case(seed, heads=heads, state_size=state_size, groups=groups) + y, state, _ = _Harness(shape).run( + x=c["x"], state=c["state"], b=c["b"], c=c["c"], + da=c["da"], dt=c["dt"], d=c["d"], + ) + torch.testing.assert_close(y, c["y_ref"], rtol=5e-5, atol=5e-6) + torch.testing.assert_close(state, c["state_ref"], rtol=5e-5, atol=5e-6) + + +def test_the_output_accumulator_is_cleared(): + """`y` is reused every token, so the contraction must zero it first. The + harness seeds it with 7.5 -- on a fresh all-zero Machine a missing fill + would be invisible.""" + shape, c = _case(6, heads=2, state_size=4) + y, _, _ = _Harness(shape).run( + x=c["x"], state=c["state"], b=c["b"], c=c["c"], + da=c["da"], dt=c["dt"], d=c["d"], + ) + torch.testing.assert_close(y, c["y_ref"], rtol=5e-5, atol=5e-6) + + +def test_the_sweeps_are_hardware_loops(): + """Two FMA sweeps per head -- the rank-1 update and the output contraction + -- plus one broadcast FMA for the D skip. Each must be a single instruction + inside a loop however large state_size is.""" + for state_size in (4, 8, 16): + shape, _ = _case(7, heads=2, state_size=state_size) + code = _Harness(shape).emit() + body = [ + line.strip() + for line in code.splitlines() + if line.strip() and not line.strip().startswith(";") + ] + fmas = sum("V_FMA_VF" in line for line in body) + assert fmas == 3 * shape.num_heads, ( + f"state_size={state_size}: expected 3 FMAs per head, got {fmas}" + ) + + +def test_emits_nothing_the_oracle_cannot_model(): + shape, c = _case(8, heads=2, state_size=4) + code = _Harness(shape).emit() + try: + Machine(vlen=MLEN).run(code) + except UnsupportedInstruction as exc: # pragma: no cover - failure path + pytest.fail(f"decode step emitted an unmodelled instruction: {exc}") diff --git a/generator/parser/llm_parser.py b/generator/parser/llm_parser.py index 85a96a0..3961dc2 100644 --- a/generator/parser/llm_parser.py +++ b/generator/parser/llm_parser.py @@ -1,8 +1,14 @@ +import math from typing import Any import torch from transformers import AutoConfig, AutoModel +# Architecture strings that select the Mamba-2 (selective state-space) lowering. +# ``model_type == "mamba2"`` is the primary signal; the architecture string is +# checked too so a config that omits model_type still routes correctly. +MAMBA2_ARCHITECTURES = frozenset({"Mamba2ForCausalLM", "Mamba2Model"}) + class LLMModelParser: def __init__(self, model_name_or_path: str): @@ -26,6 +32,120 @@ def _resolve_text_config(self): return self.config.text_config return self.config + def is_mamba2(self) -> bool: + """True when this config selects the Mamba-2 selective-SSM lowering. + + Checks ``model_type`` first (the field HuggingFace's ``Mamba2Config`` + pins to ``"mamba2"``) and falls back to the ``architectures`` list so a + hand-written config without a model_type still routes correctly. + """ + if self.config is None: + self.load_model() + text_cfg = self._resolve_text_config() + if str(getattr(text_cfg, "model_type", "") or "").lower() == "mamba2": + return True + architectures = getattr(self.config, "architectures", None) or [] + return any(arch in MAMBA2_ARCHITECTURES for arch in architectures) + + def _extract_mamba2_dimensions(self) -> dict[str, Any]: + """Resolve the Mamba-2 mixer shape from a ``Mamba2Config``. + + Field names mirror HuggingFace's ``Mamba2Config`` so a config JSON maps + across without a translation table. Derived quantities: + + * ``d_inner = expand * hidden_size`` -- the mixer's inner width + * ``conv_dim = d_inner + 2 * n_groups * state_size`` + -- the depthwise conv runs over ``[x, B, C]`` + * ``in_proj_out = 2 * d_inner + 2 * n_groups * state_size + num_heads`` + -- the fused in_proj emits ``[z, x, B, C, dt]`` + """ + text_cfg = self._resolve_text_config() + + hidden_size = getattr(text_cfg, "hidden_size") + expand = getattr(text_cfg, "expand", 2) + head_dim = getattr(text_cfg, "head_dim", 64) + state_size = getattr(text_cfg, "state_size", 128) + n_groups = getattr(text_cfg, "n_groups", 1) + conv_kernel = getattr(text_cfg, "conv_kernel", 4) + chunk_size = getattr(text_cfg, "chunk_size", 256) + + d_inner = expand * hidden_size + num_heads = getattr(text_cfg, "num_heads", None) + if num_heads is None: + num_heads = d_inner // head_dim + # A mismatch here silently corrupts every downstream dimension (the dt + # slice width, the SSD batch axis, the gated-norm group count), so fail + # loudly instead of trusting whichever field happens to be wrong. + if num_heads * head_dim != d_inner: + raise ValueError( + f"Mamba-2 config is inconsistent: num_heads ({num_heads}) * head_dim " + f"({head_dim}) = {num_heads * head_dim}, but expand ({expand}) * hidden_size " + f"({hidden_size}) = {d_inner}. num_heads must equal " + "expand * hidden_size / head_dim." + ) + if num_heads % n_groups: + raise ValueError( + f"Mamba-2 config is inconsistent: num_heads ({num_heads}) is not divisible by " + f"n_groups ({n_groups}); B and C are shared across the heads of a group." + ) + + time_step_limit = getattr(text_cfg, "time_step_limit", (0.0, float("inf"))) + try: + time_step_min, time_step_max = float(time_step_limit[0]), float(time_step_limit[1]) + except (TypeError, IndexError, ValueError): + time_step_min, time_step_max = 0.0, float("inf") + + return { + "hidden_size": hidden_size, + "expand": expand, + "d_inner": d_inner, + "state_size": state_size, + "n_groups": n_groups, + "conv_kernel": conv_kernel, + "head_dim": head_dim, + "num_heads": num_heads, + "heads_per_group": num_heads // n_groups, + "chunk_size": chunk_size, + "conv_dim": d_inner + 2 * n_groups * state_size, + "in_proj_out": 2 * d_inner + 2 * n_groups * state_size + num_heads, + "use_conv_bias": getattr(text_cfg, "use_conv_bias", True), + "use_bias": getattr(text_cfg, "use_bias", False), + "rms_norm": getattr(text_cfg, "rms_norm", True), + "activation": getattr(text_cfg, "hidden_act", "silu"), + "eps": getattr(text_cfg, "layer_norm_epsilon", getattr(text_cfg, "rms_norm_eps", 1e-5)), + "time_step_min": time_step_min, + "time_step_max": time_step_max, + } + + def _mamba2_shape(self, batch_size: int, seq_len: int) -> dict[str, Any]: + """Per-layer Mamba-2 shape stamped onto every mixer node. + + Carried on each node (rather than looked up from a global) so code_gen + can derive a single consistent VRAM map from any one of them without + needing the whole graph. + """ + mamba = self._extract_mamba2_dimensions() + shape = { + k: mamba[k] + for k in ( + "hidden_size", + "d_inner", + "conv_dim", + "in_proj_out", + "state_size", + "n_groups", + "conv_kernel", + "head_dim", + "num_heads", + "heads_per_group", + "chunk_size", + ) + } + shape["seq_len"] = seq_len + shape["batch_size"] = batch_size + shape["num_chunks"] = math.ceil(seq_len / mamba["chunk_size"]) if mamba["chunk_size"] else 1 + return shape + def extract_critical_dimensions(self) -> dict[str, Any]: """Extract dimensions for attention, RMSNorm, FFN operations""" if self.config is None: @@ -49,6 +169,12 @@ def extract_critical_dimensions(self) -> dict[str, Any]: # RMSNorm dimensions dimensions["rms_norm"] = self._extract_rms_norm_dimensions() + # Mamba-2 mixer dimensions. A Mamba-2 config has no attention and no + # gated FFN, so the ``attention`` / ``ffn`` sections above are vacuous + # for it; ``mamba`` is the authoritative section. + if self.is_mamba2(): + dimensions["mamba"] = self._extract_mamba2_dimensions() + # Include vision encoder dimensions if present if hasattr(self.config, "vision_config") and self.config.vision_config is not None: vcfg = self.config.vision_config @@ -121,6 +247,12 @@ def create_symbolic_graph(self, batch_size: int = 1, seq_len: int = 512) -> dict if self.config is None: self.load_model() + # Mamba-2 replaces the whole attention + FFN sublayer pair with a single + # SSM mixer, so it gets its own builder rather than a set of branches + # threaded through the Llama-shaped one below. + if self.is_mamba2(): + return self._create_mamba2_symbolic_graph(batch_size=batch_size, seq_len=seq_len) + text_cfg = self._resolve_text_config() # Compute GQA-aware projection dimensions @@ -326,6 +458,257 @@ def create_symbolic_graph(self, batch_size: int = 1, seq_len: int = 512) -> dict return self.symbolic_graph + def _create_mamba2_symbolic_graph(self, batch_size: int = 1, seq_len: int = 512) -> dict[str, Any]: + """Build the symbolic graph for a Mamba-2 (``Mamba2ForCausalLM``) model. + + Per layer the mixer is:: + + norm -> in_proj -> conv1d -> ssd_scan -> gated_rmsnorm -> out_proj -> residual + + which is the HuggingFace ``Mamba2Mixer`` forward, one node per stage that + maps to a distinct PLENA kernel: + + * ``in_proj`` emits the fused ``[z, x, B, C, dt]`` bundle in one GEMM, + exactly as the reference implementation does -- splitting it would cost + five weight prefetches instead of one. + * ``conv1d`` is depthwise and causal over ``[x, B, C]`` only (``z`` and + ``dt`` bypass it), hence ``conv_dim`` channels rather than + ``in_proj_out``. Its SiLU is carried on the node as ``activation`` + instead of a separate node because the two always fuse. + * ``ssd_scan`` is the chunked state-space-duality recurrence. It is one + node, not a chunk-loop of nodes, because ``chunk_size`` is a lowering + parameter of a single kernel, not a graph-level structure. + * ``gated_rmsnorm`` normalises over ``d_inner / n_groups`` and multiplies + by ``silu(z)``; it cannot be folded into the plain ``normalization`` + node because that one reduces over ``hidden_size``. + + There is a single residual per layer (Mamba-2 has one sublayer, not the + attention + FFN pair), so a layer is 7 nodes rather than Llama's 6. + """ + text_cfg = self._resolve_text_config() + mamba = self._extract_mamba2_dimensions() + shape = self._mamba2_shape(batch_size=batch_size, seq_len=seq_len) + + hidden_size = mamba["hidden_size"] + d_inner = mamba["d_inner"] + conv_dim = mamba["conv_dim"] + in_proj_out = mamba["in_proj_out"] + n_groups = mamba["n_groups"] + state_size = mamba["state_size"] + num_heads = mamba["num_heads"] + eps = mamba["eps"] + + symbolic_nodes: list[dict[str, Any]] = [] + execution_order: list[str] = [] + order_counter = 0 + + def add(node: dict[str, Any]) -> None: + nonlocal order_counter + node["execution_order"] = order_counter + symbolic_nodes.append(node) + execution_order.append(node["name"]) + order_counter += 1 + + current_shape = [batch_size, seq_len, hidden_size] + + vocab_size = getattr(text_cfg, "vocab_size", None) + if vocab_size is not None: + add( + { + "name": "embed_tokens", + "operation_type": "embedding", + "operation_category": "embedding", + "input_shape": [batch_size, seq_len], + "output_shape": current_shape, + "dimensions": {"num_embeddings": vocab_size, "hidden_size": hidden_size}, + "is_data_placeholder": True, + } + ) + + # Column offsets of each slice inside the fused in_proj output, in the + # order Mamba2Mixer concatenates them. + z_off = 0 + x_off = z_off + d_inner + b_off = x_off + d_inner + c_off = b_off + n_groups * state_size + dt_off = c_off + n_groups * state_size + in_proj_slices = { + "z": [z_off, d_inner], + "x": [x_off, d_inner], + "B": [b_off, n_groups * state_size], + "C": [c_off, n_groups * state_size], + "dt": [dt_off, num_heads], + } + + num_layers = getattr(text_cfg, "num_hidden_layers", 0) + for layer_idx in range(num_layers): + add( + { + "name": f"layer_{layer_idx}_input_layernorm", + "operation_type": "normalization", + "operation_category": "normalization", + "input_shape": current_shape, + "output_shape": current_shape, + "dimensions": {"normalized_shape": hidden_size, "eps": eps}, + "is_data_placeholder": False, + } + ) + + add( + { + "name": f"layer_{layer_idx}_in_proj", + "operation_type": "projection", + "operation_category": "projection", + "input_shape": current_shape, + "output_shape": [batch_size, seq_len, in_proj_out], + "dimensions": { + "role": "mamba_in_proj", + "in_features": hidden_size, + "out_features": in_proj_out, + "use_bias": mamba["use_bias"], + "slices": in_proj_slices, + "mamba_shape": shape, + }, + "is_data_placeholder": False, + } + ) + + add( + { + "name": f"layer_{layer_idx}_conv1d", + "operation_type": "conv1d", + "operation_category": "conv1d", + "input_shape": [batch_size, seq_len, conv_dim], + "output_shape": [batch_size, seq_len, conv_dim], + "dimensions": { + "in_channels": conv_dim, + "out_channels": conv_dim, + "groups": conv_dim, # depthwise: one group per channel + "conv_dim": conv_dim, + "kernel_size": mamba["conv_kernel"], + "stride": 1, + "padding": mamba["conv_kernel"] - 1, + "causal": True, + "depthwise": True, + "use_conv_bias": mamba["use_conv_bias"], + "activation": mamba["activation"], + "seq_len": seq_len, + "mamba_shape": shape, + }, + "is_data_placeholder": False, + } + ) + + add( + { + "name": f"layer_{layer_idx}_ssd_scan", + "operation_type": "ssd_scan", + "operation_category": "ssd_scan", + "input_shape": [batch_size, seq_len, conv_dim], + "output_shape": [batch_size, seq_len, d_inner], + "dimensions": { + "d_inner": d_inner, + "num_heads": num_heads, + "head_dim": mamba["head_dim"], + "state_size": state_size, + "n_groups": n_groups, + "heads_per_group": mamba["heads_per_group"], + "chunk_size": mamba["chunk_size"], + "num_chunks": shape["num_chunks"], + "seq_len": seq_len, + "time_step_min": mamba["time_step_min"], + "time_step_max": mamba["time_step_max"], + "mamba_shape": shape, + }, + "is_data_placeholder": False, + } + ) + + add( + { + "name": f"layer_{layer_idx}_gated_rmsnorm", + "operation_type": "gated_rmsnorm", + "operation_category": "normalization", + "input_shape": [[batch_size, seq_len, d_inner], [batch_size, seq_len, d_inner]], + "output_shape": [batch_size, seq_len, d_inner], + "dimensions": { + "normalized_shape": d_inner, + # Mamba-2 normalises per group, not over the full width. + "group_size": d_inner // n_groups, + "n_groups": n_groups, + "eps": eps, + "norm_type": "gated_rms_norm", + "gate_activation": mamba["activation"], + "mamba_shape": shape, + }, + "is_data_placeholder": False, + } + ) + + add( + { + "name": f"layer_{layer_idx}_out_proj", + "operation_type": "projection", + "operation_category": "projection", + "input_shape": [batch_size, seq_len, d_inner], + "output_shape": current_shape, + "dimensions": { + "role": "mamba_out_proj", + "in_features": d_inner, + "out_features": hidden_size, + "use_bias": mamba["use_bias"], + "mamba_shape": shape, + }, + "is_data_placeholder": False, + } + ) + + add( + { + "name": f"layer_{layer_idx}_residual", + "operation_type": "elementwise_add", + "operation_category": "elementwise_add", + "input_shape": [current_shape, current_shape], + "output_shape": current_shape, + "dimensions": {"shape": [hidden_size]}, + "is_data_placeholder": False, + } + ) + + add( + { + "name": "final_layernorm", + "operation_type": "normalization", + "operation_category": "normalization", + "input_shape": current_shape, + "output_shape": current_shape, + "dimensions": {"normalized_shape": hidden_size, "eps": eps}, + "is_data_placeholder": False, + } + ) + + if vocab_size is not None: + add( + { + "name": "lm_head", + "operation_type": "lm_head", + "operation_category": "lm_head", + "input_shape": current_shape, + "output_shape": [batch_size, seq_len, vocab_size], + "dimensions": {"hidden_size": hidden_size, "vocab_size": vocab_size}, + "is_data_placeholder": False, + } + ) + + self.symbolic_graph = { + "nodes": symbolic_nodes, + "execution_order": execution_order, + "total_nodes": len(symbolic_nodes), + "architecture_family": "mamba2", + "mamba_shape": shape, + } + return self.symbolic_graph + def create_vision_symbolic_graph(self, batch_size: int = 1) -> dict | None: """Create symbolic graph for vision encoder (SigLIP/ViT style). Returns None if no vision_config present. @@ -569,23 +952,41 @@ def print_summary(self): print(f"Number of Layers: {dims['num_hidden_layers']}") print(f"Max Position Embeddings: {dims['max_position_embeddings']}") - print("\n=== Attention Dimensions ===") - att_dims = dims["attention"] - print(f"Number of Attention Heads: {att_dims['num_attention_heads']}") - print(f"Number of Key-Value Heads: {att_dims['num_key_value_heads']}") - print(f"Head Dimension: {att_dims['head_dim']}") - print(f"Key-Value Head Dimension: {att_dims['key_value_head_dim']}") - - print("\n=== FFN Dimensions ===") - ffn_dims = dims["ffn"] - print(f"Hidden Size: {ffn_dims['hidden_size']}") - print(f"Intermediate Size: {ffn_dims['intermediate_size']}") - print(f"Activation: {ffn_dims['activation']}") - - print("\n=== RMSNorm Dimensions ===") - rms_dims = dims["rms_norm"] - print(f"Normalized Shape: {rms_dims['normalized_shape']}") - print(f"Epsilon: {rms_dims['eps']}") + if "mamba" in dims: + # Mamba-2 has neither attention nor a gated FFN; printing those + # sections would just report the vacuous defaults. + m = dims["mamba"] + print("\n=== Mamba-2 Mixer Dimensions ===") + print(f"Expand: {m['expand']} -> d_inner: {m['d_inner']}") + print(f"Heads: {m['num_heads']} x head_dim {m['head_dim']}") + print(f"State Size: {m['state_size']}, Groups: {m['n_groups']}") + print(f"Conv Kernel: {m['conv_kernel']} over conv_dim {m['conv_dim']} (depthwise, causal)") + print(f"in_proj: {m['hidden_size']} -> {m['in_proj_out']} [z, x, B, C, dt]") + print(f"out_proj: {m['d_inner']} -> {m['hidden_size']}") + print(f"Chunk Size: {m['chunk_size']}") + print(f"time_step_limit: ({m['time_step_min']}, {m['time_step_max']})") + + print("\n=== Gated RMSNorm Dimensions ===") + print(f"Normalized Shape: {m['d_inner'] // m['n_groups']} (per group)") + print(f"Epsilon: {m['eps']}") + else: + print("\n=== Attention Dimensions ===") + att_dims = dims["attention"] + print(f"Number of Attention Heads: {att_dims['num_attention_heads']}") + print(f"Number of Key-Value Heads: {att_dims['num_key_value_heads']}") + print(f"Head Dimension: {att_dims['head_dim']}") + print(f"Key-Value Head Dimension: {att_dims['key_value_head_dim']}") + + print("\n=== FFN Dimensions ===") + ffn_dims = dims["ffn"] + print(f"Hidden Size: {ffn_dims['hidden_size']}") + print(f"Intermediate Size: {ffn_dims['intermediate_size']}") + print(f"Activation: {ffn_dims['activation']}") + + print("\n=== RMSNorm Dimensions ===") + rms_dims = dims["rms_norm"] + print(f"Normalized Shape: {rms_dims['normalized_shape']}") + print(f"Epsilon: {rms_dims['eps']}") # Print vision dimensions if present if "vision" in dims: @@ -656,5 +1057,22 @@ def print_symbolic_graph_details(self): ) elif category == "elementwise_add": print(f" Add: shape={dims.get('shape')}") + elif category == "projection": + print( + f" Linear ({dims.get('role', 'projection')}): " + f"{dims.get('in_features')} -> {dims.get('out_features')}" + ) + elif category == "conv1d": + print( + f" Conv1d: channels={dims.get('conv_dim')}, kernel={dims.get('kernel_size')}, " + f"depthwise={dims.get('depthwise')}, causal={dims.get('causal')}, " + f"activation={dims.get('activation')}" + ) + elif category == "ssd_scan": + print( + f" SSD: heads={dims.get('num_heads')}x{dims.get('head_dim')}, " + f"state={dims.get('state_size')}, groups={dims.get('n_groups')}, " + f"chunk={dims.get('chunk_size')} ({dims.get('num_chunks')} chunks)" + ) print() diff --git a/generator/passes/code_gen.py b/generator/passes/code_gen.py index 47ee236..a1d5210 100644 --- a/generator/passes/code_gen.py +++ b/generator/passes/code_gen.py @@ -13,6 +13,7 @@ See docs/COMPILATION_PIPELINES.md for the full architecture overview. """ +import math from pathlib import Path from typing import Any @@ -29,7 +30,14 @@ preload_addr_reg_asm, projection_asm, rms_norm_asm, + silu_asm, ) +from asm_templates._imm import load_large_int as _load_large_int + +# Imported from the module rather than the package: adding Mamba-2 to +# asm_templates/__init__.py is out of scope for this change, and a direct module +# import is equivalent for the caller. +from asm_templates.mamba_conv1d_asm import mamba_conv1d_asm, mamba_ssd_scan_asm def _load_template(template_name: str) -> str: @@ -602,6 +610,431 @@ def _generate_vision_projection_code( return code.strip() +# --------------------------------------------------------------------------- +# Mamba-2 (selective state-space) lowering +# --------------------------------------------------------------------------- + +#: HBM address registers the Mamba-2 path reuses. A Mamba-2 program has no +#: attention and no gated FFN, so the q/k/v and ffn_* registers assigned by +#: generator/scheduler/reg_assignment_lib.json are dead for it. Rather than +#: invent register indices (only a0-a7 exist), the Mamba weights and the two SSD +#: operand-spill areas are mapped onto those free slots. The names below say +#: what each one actually holds in a Mamba-2 program. +_MAMBA_ADDR_REG_ROLES = { + "token_table": "token_table_offset", + "in_proj": "q_weight_offset", + "conv1d": "k_weight_offset", + "out_proj": "v_weight_offset", + "ssd_act_spill": "ffn_gate_offset", + "ssd_wt_spill": "ffn_up_offset", +} + + +def _mamba_shape_from_graph(symbolic_graph: dict[str, Any]) -> dict[str, Any] | None: + """Return the Mamba-2 shape carried by a graph, or None for other families. + + The parser stamps it on the graph *and* on every mixer node, so a graph + assembled by hand (as tests do) still resolves. + """ + shape = symbolic_graph.get("mamba_shape") + if shape: + return shape + for node in symbolic_graph.get("nodes", []): + candidate = node.get("dimensions", {}).get("mamba_shape") + if candidate: + return candidate + return None + + +def _mamba_vram_regions(shape: dict[str, Any], vlen: int) -> dict[str, int]: + """Bump-allocate one layer's VRAM regions as VLEN-aligned element addresses. + + Every emitter for a given layer derives its addresses from this one function, + so the conv output the SSD reads is by construction the conv output the conv + wrote. Regions are laid out with the sequence on the row axis (see + ``asm_templates/mamba_conv1d_asm``'s layout contract). + + ``x``/``B``/``C``/``dt``/``z`` are column slices of the fused ``in_proj`` + output and get their own regions here. That is not a copy the ISA would have + to make: ``projection_asm`` writes its result one BLEN-wide output column + block at a time via ``M_MM_WO``, so retargeting the destination per slice + lands them in separate regions for free. The generator does not model that + retargeting, which is the one place the addresses below diverge from a + numerically executable program. + + KNOWN LIMIT: this keeps the whole sequence resident, so ``_end`` exceeds the + physical Vector SRAM (``VECTOR_SRAM_DEPTH * VLEN`` = 64 Ki elements) for any + realistic seq_len. The generator's other paths address VRAM the same + unbounded way, and the ASM still assembles because these are element + addresses rather than encoded immediates -- but a lowering meant to *run* + would have to tile the sequence, which is exactly what ``chunk_size`` exists + to enable and what this structural model does not yet do. + """ + tokens = shape["batch_size"] * shape["seq_len"] + hidden = shape["hidden_size"] + d_inner = shape["d_inner"] + conv_dim = shape["conv_dim"] + in_proj_out = shape["in_proj_out"] + state = shape["state_size"] + heads = shape["num_heads"] + head_dim = shape["head_dim"] + groups = shape["n_groups"] + chunk = shape["chunk_size"] + conv_kernel = shape["conv_kernel"] + + regions: dict[str, int] = {} + cursor = 0 + + def alloc(name: str, size: int) -> None: + nonlocal cursor + regions[name] = cursor + padded = max(size, vlen) + cursor += ((padded + vlen - 1) // vlen) * vlen + + alloc("residual", tokens * hidden) # layer input, also the residual operand + alloc("in_proj", tokens * in_proj_out) # fused [z, x, B, C, dt] + alloc("conv_w", (conv_kernel + 1) * conv_dim) # conv_kernel tap rows + bias row + alloc("conv_in", tokens * conv_dim) # the [x, B, C] slices of in_proj + alloc("conv_out", tokens * conv_dim) + alloc("z", tokens * d_inner) # gate, consumed by the gated RMSNorm + alloc("x", tokens * d_inner) + alloc("B", tokens * groups * state) + alloc("C", tokens * groups * state) + alloc("dt", tokens * heads) + alloc("decay", tokens * heads) + alloc("score", heads * chunk * chunk) + alloc("state", heads * state * head_dim) # carried across chunks + alloc("y", tokens * d_inner) + alloc( + "scratch", + max(tokens * conv_dim, tokens * d_inner, heads * chunk * chunk, heads * state * head_dim), + ) + # Single-row temporaries. Every VRAM-destination operand needs a live + # address register pointing at a real row; giving the one-row temporaries + # their own regions keeps them from aliasing the multi-row "scratch" buffer + # that the SSD GEMMs write. + alloc("row0", vlen) + alloc("row1", vlen) + alloc("out", tokens * hidden) + regions["_end"] = cursor + return regions + + +def _mamba_fp_slots(scheduler: dict[str, Any]) -> dict[str, int]: + """FP_MEM slot indices for the Mamba-2 scalar constants. + + ``generator/scheduler/mem_layout_lib.json`` has no Mamba entries, so these + are allocated immediately after the highest named slot rather than + hardcoded -- adding a slot to the JSON later shifts these instead of + colliding with them. The host must seed them the same way it seeds + ``attn_scale`` / ``eps``. + """ + fp_sram = scheduler.get("memory_layout", {}).get("fp_sram", {}) + named = [v for v in fp_sram.values() if isinstance(v, int)] + base = (max(named) + 1) if named else 6 + return { + "dt_min": base, + "dt_max": base + 1, + "a_decay": base + 2, + "d_skip": base + 3, + "one": fp_sram.get("silu_one", fp_sram.get("silu_e", 3)), + "eps": fp_sram.get("eps", 1), + "reci_group": fp_sram.get("hid_reciprocal", 2), + } + + +def _mamba_addr_reg(scheduler: dict[str, Any], role: str, default: int) -> int: + """Resolve the HBM address register a Mamba-2 tensor is addressed through.""" + hbm_addr_reg = scheduler.get("register_assignment", {}).get("hbm_addr_reg", {}) + return hbm_addr_reg.get(_MAMBA_ADDR_REG_ROLES[role], default) + + +def _prefetch_vram_asm( + *, + vram_base: int, + elements: int, + hbm_addr_reg: int, + regs: list[int], + vlen: int, + prefetch_amount: int, + label: str, +) -> str: + """H_PREFETCH_V loop bringing ``elements`` elements from HBM into VRAM.""" + per_issue = max(1, prefetch_amount) * vlen + issues = math.ceil(elements / per_issue) if elements > 0 else 0 + if issues == 0: + return "" + dst, off, loop = regs[0], regs[1], regs[2] + lines = [f"; load {label}: {elements} elements from HBM[a{hbm_addr_reg}] ({issues} x H_PREFETCH_V)"] + lines.extend(_load_large_int(dst, vram_base)) + lines.append(f"S_ADDI_INT gp{off}, gp0, 0") + lines.append(f"C_LOOP_START gp{loop}, {issues}") + lines.append(f"H_PREFETCH_V gp{dst}, gp{off}, a{hbm_addr_reg}, 1, 0") + lines.append(f"S_ADDI_INT gp{dst}, gp{dst}, {per_issue}") + lines.append(f"S_ADDI_INT gp{off}, gp{off}, {per_issue}") + lines.append(f"C_LOOP_END gp{loop}") + return "\n".join(lines) + "\n" + + +def _mamba_rowwise_asm( + *, + body: list[str], + ptr_regs: list[int], + rows: int, + loop_reg: int, + vlen: int, + comment: str, +) -> str: + """Row loop helper for the code_gen-side Mamba vector stages.""" + if rows <= 0: + return "" + lines = [f"; {comment} ({rows} rows)", f"C_LOOP_START gp{loop_reg}, {rows}"] + lines.extend(body) + for reg in ptr_regs: + lines.append(f"S_ADDI_INT gp{reg}, gp{reg}, {vlen}") + lines.append(f"C_LOOP_END gp{loop_reg}") + return "\n".join(lines) + "\n" + + +def _generate_projection_code( + node: dict[str, Any], model_info: dict[str, Any], hardware_config: dict[str, Any], scheduler: dict[str, Any] +) -> str: + """Generate assembly for a plain linear layer (Mamba-2 in_proj / out_proj). + + Reuses ``projection_asm`` unchanged: in_proj and out_proj are ordinary dense + GEMMs against a static weight, which is exactly what that template emits. + The GEMM's batch axis is the token axis (``batch_size * seq_len``), not the + model batch -- a Mamba-2 prefill projects every token in one pass. + """ + dims = node["dimensions"] + in_features = dims["in_features"] + out_features = dims["out_features"] + role = dims.get("role", "projection") + + mlen = hardware_config.get("MLEN", 64) + blen = hardware_config.get("BLEN", 4) + vlen = hardware_config.get("VLEN", 64) + vsram = scheduler["memory_layout"].get("vector_sram_addr", {}) + + shape = dims.get("mamba_shape") + if shape is not None: + regions = _mamba_vram_regions(shape, vlen) + batch = shape["batch_size"] * shape["seq_len"] + if role == "mamba_out_proj": + activation_base = regions["y"] + result_base = regions["out"] + w_reg = _mamba_addr_reg(scheduler, "out_proj", 4) + else: + activation_base = regions["residual"] + result_base = regions["in_proj"] + w_reg = _mamba_addr_reg(scheduler, "in_proj", 2) + scratch_base = regions["scratch"] + else: + regions = None + batch = model_info.get("batch", 1) + activation_base = vsram.get("block1", 0) + result_base = vsram.get("block2", 0) + scratch_base = vsram.get("k_split_scratch", vsram.get("block4", 0)) + w_reg = scheduler["register_assignment"].get("hbm_addr_reg", {}).get("q_weight_offset", 2) + + code = f""" +; Linear projection ({role}): ({batch}, {in_features}) @ ({in_features}, {out_features}) +""" + if out_features % blen: + code += ( + f"; NOTE: out_features={out_features} is not a multiple of BLEN={blen}; " + f"projection_asm emits whole BLEN column blocks, so the final " + f"{out_features % blen} column(s) are not covered.\n" + ) + code += projection_asm( + mlen=mlen, + blen=blen, + batch=batch, + hidden_size=in_features, + alive_registers=[1, 2, 3, 4, 5, 6, 7, 8], + w_base_hbm_offset_reg=w_reg, + activation_base_address=activation_base, + result_base_address=result_base, + rope_enabled=False, + out_features=out_features, + matrix_sram_size=hardware_config.get("MATRIX_SRAM_SIZE", 1024), + scratch_base_address=scratch_base, + vlen=vlen, + ) + return code.strip() + + +def _generate_conv1d_code( + node: dict[str, Any], model_info: dict[str, Any], hardware_config: dict[str, Any], scheduler: dict[str, Any] +) -> str: + """Generate assembly for the Mamba-2 causal depthwise conv1d + its SiLU. + + Deliberately not routed through ``im2col_asm``/``_no_shift``: those hardcode a + square KxK patch, have no padding (so they cannot express the causal left + pad), and would build a dense ``(M, C_in*K*K)`` GEMM against a block-diagonal + weight. ``mamba_conv1d_asm`` instead exploits the sequence-on-rows layout, + where the causal shift is a row-address offset and each tap is a VLEN-wide + ``V_MUL_VV`` operand. + """ + dims = node["dimensions"] + shape = dims["mamba_shape"] + vlen = hardware_config.get("VLEN", 64) + regions = _mamba_vram_regions(shape, vlen) + conv_dim = dims["conv_dim"] + conv_kernel = dims["kernel_size"] + tokens = shape["batch_size"] * shape["seq_len"] + use_bias = dims.get("use_conv_bias", True) + fp_slots = _mamba_fp_slots(scheduler) + + code = f""" +; === Mamba-2 depthwise causal conv1d === +; channels={conv_dim} (groups={dims["groups"]}), kernel={conv_kernel}, padding={dims["padding"]} (causal) +; Operates on the [x, B, C] slices of in_proj only; z and dt bypass the conv. +""" + # The conv weight is a static per-layer parameter: taps then (optionally) bias. + code += _prefetch_vram_asm( + vram_base=regions["conv_w"], + elements=(conv_kernel + (1 if use_bias else 0)) * conv_dim, + hbm_addr_reg=_mamba_addr_reg(scheduler, "conv1d", 3), + regs=[1, 2, 3], + vlen=vlen, + prefetch_amount=hardware_config.get("HBM_V_Prefetch_Amount", 16), + label="conv1d taps + bias", + ) + code += mamba_conv1d_asm( + vlen=vlen, + seq_len=tokens, + conv_dim=conv_dim, + conv_kernel=conv_kernel, + alive_registers=[1, 2, 3, 4, 5, 6, 7, 8], + input_base_address=regions["conv_in"], + output_base_address=regions["conv_out"], + weight_base_address=regions["conv_w"], + scratch_base_address=regions["row0"], + bias_base_address=(regions["conv_w"] + conv_kernel * conv_dim) if use_bias else None, + ) + + activation = dims.get("activation", "silu") + if activation == "silu": + code += f"\n; -- {activation} on the conv output (fused: Mamba-2 never separates them) --\n" + code += silu_asm( + const_one_fp_address=fp_slots["one"], + alive_registers=[1, 2, 3], + activation_base_address=regions["conv_out"], + # silu_asm holds its sigmoid in a single reused row, so one row is enough. + scratchpad_base_address=regions["row1"], + vlen=vlen, + batch_size=tokens, + hidden_dim=conv_dim, + ) + else: + code += f"\n; -- {activation} activation (unrecognized; no ASM emitted) --\n" + return code.strip() + + +def _generate_ssd_scan_code( + node: dict[str, Any], model_info: dict[str, Any], hardware_config: dict[str, Any], scheduler: dict[str, Any] +) -> str: + """Generate assembly for the chunked state-space-duality scan.""" + dims = node["dimensions"] + shape = dims["mamba_shape"] + vlen = hardware_config.get("VLEN", 64) + regions = _mamba_vram_regions(shape, vlen) + fp_slots = _mamba_fp_slots(scheduler) + + code = f""" +; === Mamba-2 SSD scan === +; chunk_size={dims["chunk_size"]} -> {dims["num_chunks"]} chunk(s) over seq_len={dims["seq_len"]} +; time_step_limit=({dims["time_step_min"]}, {dims["time_step_max"]}) +; The x / B / C / dt operands are the corresponding column slices of the conv +; output; they are addressed as separate VRAM regions here (see +; _mamba_vram_regions for why that costs nothing on real hardware). +""" + code += mamba_ssd_scan_asm( + mlen=hardware_config.get("MLEN", 64), + vlen=vlen, + blen=hardware_config.get("BLEN", 4), + seq_len=shape["batch_size"] * shape["seq_len"], + chunk_size=dims["chunk_size"], + num_heads=dims["num_heads"], + head_dim=dims["head_dim"], + state_size=dims["state_size"], + n_groups=dims["n_groups"], + alive_registers=[1, 2, 3, 4, 5, 6, 7, 8], + vram=regions, + act_spill_addr_reg=_mamba_addr_reg(scheduler, "ssd_act_spill", 6), + wt_spill_addr_reg=_mamba_addr_reg(scheduler, "ssd_wt_spill", 7), + writeback_amount=hardware_config.get("HBM_V_Writeback_Amount", 4), + dt_min_fp_address=fp_slots["dt_min"], + dt_max_fp_address=fp_slots["dt_max"], + a_decay_fp_address=fp_slots["a_decay"], + d_skip_fp_address=fp_slots["d_skip"], + ) + return code.strip() + + +def _generate_gated_rmsnorm_code( + node: dict[str, Any], model_info: dict[str, Any], hardware_config: dict[str, Any], scheduler: dict[str, Any] +) -> str: + """Generate assembly for Mamba-2's gated RMSNorm: ``y = RMSNorm(y) * silu(z)``. + + Cannot reuse ``_generate_normalization_code``: Mamba-2 normalises over + ``d_inner / n_groups``, not over ``hidden_size``, and the gate multiply has no + counterpart in the plain norm node. The RMSNorm body itself is + ``rms_norm_asm`` with the group width as its ``hidden_dim`` and one "batch" + row per (token, group). + """ + dims = node["dimensions"] + shape = dims["mamba_shape"] + vlen = hardware_config.get("VLEN", 64) + regions = _mamba_vram_regions(shape, vlen) + fp_slots = _mamba_fp_slots(scheduler) + + group_size = dims["group_size"] + n_groups = dims["n_groups"] + d_inner = dims["normalized_shape"] + tokens = shape["batch_size"] * shape["seq_len"] + + code = f""" +; === Mamba-2 gated RMSNorm === +; normalize over {group_size} (= d_inner {d_inner} / n_groups {n_groups}), then +; multiply by {dims["gate_activation"]}(z). 1/group_size must be seeded into the +; FP_MEM slot below -- the plain-norm path seeds 1/hidden_size there instead. +""" + code += rms_norm_asm( + _eps_offset=fp_slots["eps"], + reci_hid_offset=fp_slots["reci_group"], + alive_registers=[1, 2, 3], + activation_base_address=regions["y"], + scratchpad_base_address=regions["row0"], + vlen=vlen, + batch_size=tokens * n_groups, + hidden_dim=group_size, + ) + code += f"\n; -- gate = {dims['gate_activation']}(z), computed in place on the z slice --\n" + code += silu_asm( + const_one_fp_address=fp_slots["one"], + alive_registers=[1, 2, 3], + activation_base_address=regions["z"], + scratchpad_base_address=regions["row1"], + vlen=vlen, + batch_size=tokens, + hidden_dim=d_inner, + ) + code += "\n; -- y *= gate --\n" + code += "".join(line + "\n" for line in _load_large_int(1, regions["y"])) + code += "".join(line + "\n" for line in _load_large_int(2, regions["z"])) + code += _mamba_rowwise_asm( + body=["V_MUL_VV gp1, gp1, gp2, 0"], + ptr_regs=[1, 2], + rows=max(1, (tokens * d_inner) // vlen), + loop_reg=3, + vlen=vlen, + comment="gate multiply", + ) + return code.strip() + + def _generate_elementwise_add_code( node: dict[str, Any], model_info: dict[str, Any], hardware_config: dict[str, Any], scheduler: dict[str, Any] ) -> str: @@ -679,6 +1112,14 @@ def _generate_node_code( return header + _generate_conv2d_code(node, model_info, hardware_config, scheduler) elif operation_type == "vision_projection": return header + _generate_vision_projection_code(node, model_info, hardware_config, scheduler) + elif operation_type == "projection": + return header + _generate_projection_code(node, model_info, hardware_config, scheduler) + elif operation_type == "conv1d": + return header + _generate_conv1d_code(node, model_info, hardware_config, scheduler) + elif operation_type == "ssd_scan": + return header + _generate_ssd_scan_code(node, model_info, hardware_config, scheduler) + elif operation_type == "gated_rmsnorm": + return header + _generate_gated_rmsnorm_code(node, model_info, hardware_config, scheduler) else: raise ValueError(f"Unknown operation type: {operation_type}") @@ -716,16 +1157,74 @@ def _weight_hbm_bytes(rows: int, cols: int, hardware_config: dict) -> int: return rows * (cols // block_dim) * ((wt_block_width // 8) + (scale_width // 8)) +def _mamba_addr_reg_layout( + model_info: dict[str, Any], + shape: dict[str, Any], + hardware_config: dict[str, Any], +) -> list[tuple[str, int, int, str | None]]: + """HBM weight layout for a Mamba-2 program, as (name, rows, cols, reg_key). + + Mamba-2 has no q/k/v and no gated FFN, so the register keys those names map + to in reg_assignment_lib.json are reused here for the tensors a Mamba-2 + program actually needs (see ``_MAMBA_ADDR_REG_ROLES``). The last two entries + are not weights at all but the two SSD operand-spill areas -- their size is + the largest single operand any of the four per-chunk GEMMs spills. + """ + block_dim = hardware_config.get("block_dim", 4) + + def pad(cols: int) -> int: + # _weight_hbm_bytes requires whole MX blocks; a Mamba width that is not a + # multiple of block_dim (e.g. in_proj_out with num_heads % block_dim != 0) + # is rounded up rather than tripping its assert. + return ((cols + block_dim - 1) // block_dim) * block_dim + + vocab_size = model_info.get("vocab_size", 1024) + hidden = shape["hidden_size"] + d_inner = shape["d_inner"] + conv_dim = shape["conv_dim"] + in_proj_out = shape["in_proj_out"] + chunk = shape["chunk_size"] + state = shape["state_size"] + heads = shape["num_heads"] + head_dim = shape["head_dim"] + groups = shape["n_groups"] + + spill_elements = max( + groups * chunk * state, + heads * chunk * chunk, + heads * chunk * head_dim, + heads * chunk * state, + heads * state * head_dim, + ) + spill_rows = max(1, math.ceil(spill_elements / pad(max(head_dim, block_dim)))) + + return [ + ("token_table", vocab_size, pad(hidden), _MAMBA_ADDR_REG_ROLES["token_table"]), + ("in_proj_weight", hidden, pad(in_proj_out), _MAMBA_ADDR_REG_ROLES["in_proj"]), + # conv_kernel tap rows + one bias row, conv_dim wide. + ("conv1d_weight", shape["conv_kernel"] + 1, pad(conv_dim), _MAMBA_ADDR_REG_ROLES["conv1d"]), + ("out_proj_weight", d_inner, pad(hidden), _MAMBA_ADDR_REG_ROLES["out_proj"]), + ("ssd_act_spill", spill_rows, pad(max(head_dim, block_dim)), _MAMBA_ADDR_REG_ROLES["ssd_act_spill"]), + ("ssd_wt_spill", spill_rows, pad(max(head_dim, block_dim)), _MAMBA_ADDR_REG_ROLES["ssd_wt_spill"]), + ] + + def _generate_addr_reg_init( model_info: dict[str, Any], hardware_config: dict[str, Any], scheduler: dict[str, Any], + mamba_shape: dict[str, Any] | None = None, ) -> str: """Emit C_SET_ADDR_REG instructions for all weight HBM address registers. Computes cumulative MXFP8 byte offsets matching the HBM layout that _build_hbm_from_hf_weights writes. Must be emitted before any node that references HBM weights. + + When ``mamba_shape`` is supplied the Mamba-2 layout is used instead of the + attention + FFN one; emitting the latter for a Mamba-2 program would reserve + HBM for q/k/v/ffn tensors that do not exist and leave the registers the + Mamba nodes actually read pointing into the wrong regions. """ hidden_size = model_info.get("hidden_size", 64) intermediate_size = model_info.get("intermediate_size", 256) @@ -736,6 +1235,10 @@ def _generate_addr_reg_init( hbm_addr_reg = scheduler["register_assignment"].get("hbm_addr_reg", {}) + if mamba_shape is not None: + layout = _mamba_addr_reg_layout(model_info, mamba_shape, hardware_config) + return _emit_addr_reg_init(layout, hardware_config, scheduler) + # Weight layout: ordered list of (name, rows, cols, addr_reg_key). # addr_reg_key is None for weights without a dedicated register. layout = [ @@ -748,6 +1251,21 @@ def _generate_addr_reg_init( ("ffn_up", hidden_size, intermediate_size, "ffn_up_offset"), ("ffn_down", intermediate_size, hidden_size, "ffn_down_offset"), ] + _ = hbm_addr_reg # resolved inside _emit_addr_reg_init + return _emit_addr_reg_init(layout, hardware_config, scheduler) + + +def _emit_addr_reg_init( + layout: list[tuple[str, int, int, str | None]], + hardware_config: dict[str, Any], + scheduler: dict[str, Any], +) -> str: + """Turn an ordered HBM tensor layout into C_SET_ADDR_REG initialisation. + + Shared by the attention and Mamba-2 layouts so the two cannot drift in how + they compute cumulative MXFP8 byte offsets. + """ + hbm_addr_reg = scheduler["register_assignment"].get("hbm_addr_reg", {}) # Compute cumulative offsets offset = 0 @@ -795,8 +1313,11 @@ def code_gen_pass( # Generate program header asm_code = [_generate_program_header(model_info)] - # Initialize HBM address registers with weight byte offsets - addr_reg_init = _generate_addr_reg_init(model_info, hardware_config, scheduler) + # Initialize HBM address registers with weight byte offsets. A Mamba-2 graph + # needs a different HBM layout (no q/k/v, no FFN), detected from the graph + # itself rather than from model_info so callers need no new contract. + mamba_shape = _mamba_shape_from_graph(symbolic_graph) + addr_reg_init = _generate_addr_reg_init(model_info, hardware_config, scheduler, mamba_shape=mamba_shape) if addr_reg_init: asm_code.append(addr_reg_init) diff --git a/generator/passes/utilization_report.py b/generator/passes/utilization_report.py index 7a51a21..54835ca 100644 --- a/generator/passes/utilization_report.py +++ b/generator/passes/utilization_report.py @@ -95,6 +95,69 @@ def _report_ffn_utilization(node: dict[str, Any], model_info: dict[str, Any], m: return [overall_operation_amount, attainable_operation, theoretical_operation] +def _report_projection_utilization(node: dict[str, Any], model_info: dict[str, Any], m: int, k: int, n: int) -> list: + """Utilization of a dense linear layer (Mamba-2 in_proj / out_proj). + + Tile count follows the same shape as one term of ``_report_ffn_utilization``: + ``(out_features // m) * (in_features // k)``. + + ``attainable`` differs from the attention path's convention on purpose. That + one multiplies by ``model_info["batch_size"]`` because it models *decoding*, + where a single token leaves ``n - 1`` systolic rows idle. A Mamba-2 prefill + projects every token in the sequence, so the row count is + ``batch_size * seq_len`` and the array is fully fed whenever that exceeds + ``n`` -- capping at ``n`` is what makes attainable meet theoretical here. + """ + dims = node["dimensions"] + in_features = dims["in_features"] + out_features = dims["out_features"] + rows = model_info.get("batch_size", 1) * model_info.get("seq_len", 1) + + operation_amount = max(1, out_features // m) * max(1, in_features // k) + attainable = operation_amount * (m * k * min(rows, n)) + theoretical = operation_amount * (m * k * n) + return [operation_amount, attainable, theoretical] + + +def _report_ssd_utilization(node: dict[str, Any], model_info: dict[str, Any], m: int, k: int, n: int) -> list: + """Utilization of the chunked SSD scan's four per-chunk GEMMs. + + Per chunk, with ``Q = chunk_size``, ``P = state_size``, ``D = head_dim``: + + * ``C @ B^T`` -- (Q, P) @ (P, Q), batched over ``n_groups`` + * ``M @ X`` -- (Q, Q) @ (Q, D), batched over ``num_heads`` + * ``B^T @ (X*dt)`` -- (P, Q) @ (Q, D), batched over ``num_heads`` + * ``C @ S`` -- (Q, P) @ (P, D), batched over ``num_heads`` + + The elementwise stages (dt softplus, the decay cumsum, the decay mask, the D + skip) run on the vector unit and contribute nothing to systolic utilization, + the same way ``normalization`` and ``elementwise_add`` nodes already report + zero here. They are not free -- the cumsum in particular is a serial chain + of ``chunk_size`` dependent vector ops -- so a systolic-only report + understates the real cost of this kernel. + """ + dims = node["dimensions"] + num_heads = dims["num_heads"] + head_dim = dims["head_dim"] + state_size = dims["state_size"] + n_groups = dims["n_groups"] + chunk = dims["chunk_size"] + num_chunks = dims["num_chunks"] + + per_chunk = ( + n_groups * max(1, chunk // m) * max(1, state_size // k) # C @ B^T + + num_heads * max(1, head_dim // m) * max(1, chunk // k) # M @ X + + num_heads * max(1, head_dim // m) * max(1, chunk // k) # B^T @ (X*dt) + + num_heads * max(1, head_dim // m) * max(1, state_size // k) # C @ S + ) + operation_amount = num_chunks * per_chunk + # Rows fed per GEMM are chunk-many (or state-many for the state update); + # chunk_size >= MLEN by construction, so these tiles are always full. + attainable = operation_amount * (m * k * min(chunk, n)) + theoretical = operation_amount * (m * k * n) + return [operation_amount, attainable, theoretical] + + def _report_utilization(node: dict[str, Any], model_info: dict[str, Any], m: int, k: int, n: int) -> str: """Generate assembly code for a single symbolic graph node.""" operation_type = node["operation_type"] @@ -105,6 +168,10 @@ def _report_utilization(node: dict[str, Any], model_info: dict[str, Any], m: int return _report_flash_attn_utilization(node, model_info, 1024, m, k, n) elif operation_type == "ffn": return _report_ffn_utilization(node, model_info, m, k, n) + elif operation_type == "projection": + return _report_projection_utilization(node, model_info, m, k, n) + elif operation_type == "ssd_scan": + return _report_ssd_utilization(node, model_info, m, k, n) else: return [0, 0, 0] @@ -148,9 +215,12 @@ def analyse_overall_utilization( # Create a mapping from node names to nodes for efficient lookup node_map = {node["name"]: node for node in nodes} - overall_operations = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_attainable_flops = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_theoretical_flops = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} + # "projection" and "ssd_scan" are the Mamba-2 systolic kernels; they stay 0 + # for attention models, which never emit those node types. + accounted = ["embedding", "attention", "ffn", "projection", "ssd_scan"] + overall_operations = dict.fromkeys([*accounted, "lm_head"], 0) + overall_attainable_flops = dict.fromkeys([*accounted, "lm_head"], 0) + overall_theoretical_flops = dict.fromkeys([*accounted, "lm_head"], 0) # Generate code for each node in execution order for node_name in execution_order: @@ -158,7 +228,7 @@ def analyse_overall_utilization( node = node_map[node_name] single_op_operation = _report_utilization(node, model_info, m, k, n) operation_type = node["operation_type"] - if operation_type in ["embedding", "attention", "ffn"]: + if operation_type in accounted: overall_operations[operation_type] += single_op_operation[0] overall_attainable_flops[operation_type] += single_op_operation[1] overall_theoretical_flops[operation_type] += single_op_operation[2] diff --git a/generator/tests/test_mamba_code_gen.py b/generator/tests/test_mamba_code_gen.py new file mode 100644 index 0000000..df19c28 --- /dev/null +++ b/generator/tests/test_mamba_code_gen.py @@ -0,0 +1,316 @@ +"""Tests for code generation on Mamba-2 (selective state-space) models. + +Mirrors test_vlm_code_gen.py: the config is scaled down to hardware-compatible +dimensions so the full generator pipeline (parser -> scheduler -> code_gen -> +assembler) runs with no model download, and the correctness bar for this pipeline +is that the emitted ASM assembles cleanly. + +Shape choices, all of them load-bearing: + +* ``state_size`` and ``chunk_size`` are multiples of MLEN=64 -- both are K + dimensions of the SSD GEMMs, and ``batched_matmul_asm`` requires ``k % mlen``. +* ``num_heads`` is a multiple of BLEN=4 so ``in_proj_out`` lands on a BLEN + boundary; otherwise ``projection_asm`` (which emits whole BLEN column blocks) + would not cover the dt tail. +* ``num_heads * head_dim == expand * hidden_size``, which the parser enforces. +""" + +import os +import re +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from parser.llm_parser import LLMModelParser +from parser.hardware_parser import hardware_parser +from passes.code_gen import code_gen_pass +from scheduler import gen_scheduler + + +def _get_paths(): + compiler_root = Path(__file__).resolve().parents[2] + return { + "hw_config": str(compiler_root / "doc" / "configuration.svh"), + "precision": str(compiler_root / "doc" / "precision.svh"), + "mem_layout": str(compiler_root / "generator" / "scheduler" / "mem_layout_lib.json"), + "reg_assign": str(compiler_root / "generator" / "scheduler" / "reg_assignment_lib.json"), + "operation": str(compiler_root / "doc" / "operation.svh"), + } + + +def _make_mamba2_config(num_hidden_layers: int = 1): + return SimpleNamespace( + model_type="mamba2", + architectures=["Mamba2ForCausalLM"], + hidden_size=128, + num_hidden_layers=num_hidden_layers, + expand=2, + head_dim=64, + num_heads=4, + state_size=64, + n_groups=1, + conv_kernel=4, + chunk_size=64, + vocab_size=1024, + layer_norm_epsilon=1e-5, + time_step_limit=(0.0, float("inf")), + time_step_min=0.001, + time_step_max=0.1, + use_conv_bias=True, + use_bias=False, + rms_norm=True, + hidden_act="silu", + tie_word_embeddings=False, + ) + + +def _build(num_hidden_layers: int = 1, seq_len: int = 64): + """Return (graph, model_info, hw_config, scheduler) for a Mamba-2 decoder.""" + paths = _get_paths() + parser = LLMModelParser("mock-mamba2") + parser.config = _make_mamba2_config(num_hidden_layers) + parser.model = SimpleNamespace() + + graph = parser.create_symbolic_graph(batch_size=1, seq_len=seq_len) + hw = hardware_parser(paths["hw_config"], paths["precision"]) + dims = parser.extract_critical_dimensions() + model_config = { + "hidden_size": dims["hidden_size"], + "num_layers": dims["num_hidden_layers"], + "seq_len": seq_len, + "batch_size": 1, + "vocab_size": dims["vocab_size"], + # Mamba-2 has no FFN; d_inner is the analogous "wide" dimension the + # scheduler's VRAM block sizing keys off. + "intermediate_size": dims["mamba"]["d_inner"], + "batch": 1, + } + sched = gen_scheduler(hw, model_config, paths["mem_layout"], paths["reg_assign"]) + model_info = dict(model_config, model_name="mamba2-mock", architecture="mamba2") + return graph, model_info, hw, sched + + +def _codegen(num_hidden_layers: int = 1, seq_len: int = 64) -> str: + graph, model_info, hw, sched = _build(num_hidden_layers, seq_len) + return code_gen_pass(graph, model_info, hw, sched) + + +def _instruction_lines(asm: str) -> list[str]: + """The lines the assembler will actually encode (comments/blanks dropped).""" + out = [] + for line in asm.splitlines(): + line = line.strip() + if not line or line[0] == ";" or line.startswith("//"): + continue + out.append(line) + return out + + +class TestMamba2CodeGenStructure(unittest.TestCase): + def setUp(self): + self.asm = _codegen() + + def test_all_node_types_are_dispatched(self): + """Every Mamba-2 node type must reach an emitter, not fall through.""" + for marker in ( + "(embedding)", + "(normalization)", + "(projection)", + "(conv1d)", + "(ssd_scan)", + "(gated_rmsnorm)", + "(elementwise_add)", + "(lm_head)", + ): + self.assertIn(marker, self.asm, f"no code emitted for node type {marker}") + + def test_mixer_stages_present(self): + self.assertIn("Mamba-2 causal depthwise conv1d", self.asm) + self.assertIn("Mamba-2 chunked SSD scan", self.asm) + self.assertIn("Mamba-2 gated RMSNorm", self.asm) + for stage in ( + "@stage=mamba_dt", + "@stage=mamba_chunk_cumsum", + "@stage=mamba_intra_chunk", + "@stage=mamba_decay_mask", + "@stage=mamba_state_update", + "@stage=mamba_inter_chunk", + "@stage=mamba_skip", + ): + self.assertIn(stage, self.asm, f"SSD scan is missing stage {stage}") + + def test_conv1d_is_not_lowered_through_im2col(self): + """The depthwise causal conv must not go through the dense im2col path. + + im2col hardcodes a square KxK patch, cannot express the causal left pad, + and would build a dense GEMM against a block-diagonal weight. + """ + self.assertNotIn("im2col", self.asm.lower()) + self.assertNotIn("V_SHFT_V", self.asm) + + def test_conv1d_uses_vector_taps_not_scalar_broadcast(self): + """A depthwise tap varies per channel = per lane, so it must be a + V_MUL_VV operand rather than a V_MUL_VF scalar broadcast.""" + conv_start = self.asm.index("Mamba-2 causal depthwise conv1d") + conv_end = self.asm.index("Mamba-2 SSD scan") + conv_body = self.asm[conv_start:conv_end] + self.assertIn("V_MUL_VV", conv_body) + + def test_uses_the_mamba_specific_opcodes(self): + # softplus has no software lowering (no logarithm in the ISA). + self.assertIn("V_SOFTPLUS_V", self.asm) + # time_step_limit clamp -- V_MAX_VV / V_MIN_VV do not exist. + self.assertIn("V_MAX_VF", self.asm) + self.assertIn("V_MIN_VF", self.asm) + # per-head decay scalars: VRAM row -> FPRAM slots, for V_MUL_VF broadcast. + self.assertIn("S_MAP_FP_V", self.asm) + + def test_ssd_pays_for_the_mram_round_trip(self): + """Both SSD GEMM operands are activations, and MRAM is write-only from + HBM, so each must be spilled with H_STORE_V before H_PREFETCH_M.""" + ssd_body = self.asm[self.asm.index("Mamba-2 chunked SSD scan") :] + self.assertIn("H_STORE_V", ssd_body) + self.assertIn("H_PREFETCH_M", ssd_body) + self.assertIn("M_MM", ssd_body) + + def test_ssd_spill_operands_use_distinct_hbm_registers(self): + """Sharing one address register would overlay the two GEMM operands.""" + ssd_body = self.asm[self.asm.index("Mamba-2 chunked SSD scan") :] + stores = re.findall(r"^H_STORE_V .*?, a(\d+),", ssd_body, flags=re.MULTILINE) + self.assertTrue(stores, "no H_STORE_V found in the SSD body") + self.assertGreaterEqual(len(set(stores)), 2, "both operands spilled to the same HBM register") + + def test_no_unavailable_opcodes(self): + """Guard against inventing ISA that doesn't exist.""" + for forbidden in ("V_MAX_VV", "V_MIN_VV", "V_LOG", "V_SQRT_V", "V_GATHER", "M_MM_VV"): + self.assertNotIn(forbidden, self.asm, f"emitted non-existent opcode {forbidden}") + + def test_loop_trip_counts_are_compile_time_immediates(self): + """C_LOOP_START takes an immediate only -- there is no register form and + no branch to fall back on.""" + starts = [ln for ln in _instruction_lines(self.asm) if ln.startswith("C_LOOP_START")] + self.assertTrue(starts, "no hardware loops emitted") + for line in starts: + trip = line.split(",")[-1].strip() + self.assertTrue(trip.isdigit(), f"non-immediate loop trip count: {line}") + self.assertGreater(int(trip), 0, f"degenerate loop: {line}") + + def test_loops_are_balanced(self): + lines = _instruction_lines(self.asm) + starts = sum(1 for ln in lines if ln.startswith("C_LOOP_START")) + ends = sum(1 for ln in lines if ln.startswith("C_LOOP_END")) + self.assertEqual(starts, ends, "unbalanced C_LOOP_START / C_LOOP_END") + + def test_addr_reg_init_covers_the_mamba_tensors(self): + """The Mamba-2 HBM layout must be emitted, not the attention one.""" + header = self.asm[: self.asm.index("=== embed_tokens")] + self.assertIn("C_SET_ADDR_REG", header) + # 6 tensors: token table, in_proj, conv1d, out_proj, and the two SSD spills. + self.assertEqual(header.count("C_SET_ADDR_REG"), 6) + + +class TestMamba2CodeGenScaling(unittest.TestCase): + def test_instruction_count_scales_with_layers(self): + one = len(_instruction_lines(_codegen(num_hidden_layers=1))) + two = len(_instruction_lines(_codegen(num_hidden_layers=2))) + three = len(_instruction_lines(_codegen(num_hidden_layers=3))) + per_layer = two - one + self.assertGreater(per_layer, 0) + # Layers are emitted identically, so growth must be exactly linear. + self.assertEqual(three - two, per_layer) + + def test_instruction_count_scales_with_seq_len(self): + short = len(_instruction_lines(_codegen(seq_len=64))) + long = len(_instruction_lines(_codegen(seq_len=128))) + self.assertGreater(long, short, "doubling seq_len must cost more instructions") + + def test_chunk_count_scales_with_seq_len(self): + self.assertEqual(_codegen(seq_len=64).count("; ================ chunk"), 1) + self.assertEqual(_codegen(seq_len=128).count("; ================ chunk"), 2) + self.assertEqual(_codegen(seq_len=192).count("; ================ chunk"), 3) + + def test_conv_prologue_is_independent_of_seq_len(self): + """Only conv_kernel - 1 timesteps see the causal pad, whatever the + sequence length; the rest share one hardware-loop body.""" + for seq_len in (64, 128): + asm = _codegen(seq_len=seq_len) + self.assertEqual(asm.count("fall before the sequence (zero pad)"), 3 * (384 // 64)) + + +class TestMamba2CodeGenAssembles(unittest.TestCase): + def _assemble(self, asm_text: str, label: str) -> int: + from assembler import AssemblyToBinary + + paths = _get_paths() + asm_tool = AssemblyToBinary(paths["operation"], paths["hw_config"]) + with tempfile.NamedTemporaryFile(suffix=".asm", mode="w", delete=False) as af: + af.write(asm_text) + asm_path = af.name + mem_path = asm_path.replace(".asm", ".mem") + try: + binary = asm_tool.generate_binary(asm_path, mem_path) + self.assertGreater(os.path.getsize(mem_path), 0, f"{label}: assembled .mem is empty") + return len(binary) + except ValueError as exc: + raise AssertionError(f"{label}: assembler raised ValueError (u32 overflow): {exc}") from exc + except KeyError as exc: + raise AssertionError(f"{label}: assembler does not know opcode {exc}") from exc + finally: + for path in (asm_path, mem_path): + if os.path.exists(path): + os.unlink(path) + + def test_single_layer_assembles(self): + count = self._assemble(_codegen(num_hidden_layers=1, seq_len=64), "mamba2-1L-64") + self.assertGreater(count, 0) + + def test_multi_chunk_assembles(self): + """A sequence longer than chunk_size exercises the multi-chunk path.""" + count = self._assemble(_codegen(num_hidden_layers=1, seq_len=128), "mamba2-1L-128") + self.assertGreater(count, 0) + + def test_two_layers_assemble(self): + count = self._assemble(_codegen(num_hidden_layers=2, seq_len=64), "mamba2-2L-64") + self.assertGreater(count, 0) + + +class TestMamba2UtilizationReport(unittest.TestCase): + """The utilization pass must not report a Mamba-2 model as doing no work.""" + + def test_projection_and_ssd_are_accounted(self): + from passes.utilization_report import analyse_overall_utilization + + graph, model_info, _hw, _sched = _build(num_hidden_layers=2, seq_len=128) + report = analyse_overall_utilization(graph, model_info, 64, 64, 64) + + self.assertGreater(report["operations"]["projection"], 0, "in_proj/out_proj not accounted") + self.assertGreater(report["operations"]["ssd_scan"], 0, "SSD GEMMs not accounted") + self.assertGreater(report["theoretical_FLOPS"]["ssd_scan"], 0) + # Mamba-2 has no attention and no gated FFN. + self.assertEqual(report["operations"]["attention"], 0) + self.assertEqual(report["operations"]["ffn"], 0) + + def test_systolic_work_scales_with_layers_and_chunks(self): + from passes.utilization_report import analyse_overall_utilization + + def ops(num_layers, seq_len): + graph, model_info, _hw, _sched = _build(num_layers, seq_len) + return analyse_overall_utilization(graph, model_info, 64, 64, 64)["operations"] + + one = ops(1, 64) + two = ops(2, 64) + self.assertEqual(two["projection"], 2 * one["projection"]) + self.assertEqual(two["ssd_scan"], 2 * one["ssd_scan"]) + + # Doubling seq_len doubles the chunk count, hence the SSD GEMM count. + long = ops(1, 128) + self.assertEqual(long["ssd_scan"], 2 * one["ssd_scan"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/generator/tests/test_mamba_parser.py b/generator/tests/test_mamba_parser.py new file mode 100644 index 0000000..caa31bf --- /dev/null +++ b/generator/tests/test_mamba_parser.py @@ -0,0 +1,403 @@ +"""Tests for LLMModelParser on Mamba-2 (selective state-space) configs. + +Kept separate from test_llm_parser.py / test_vlm_parser.py the same way the VLM +tests are: Mamba-2 is a distinct architecture family with its own node sequence, +and mixing its fixtures into the attention tests would obscure which family a +regression belongs to. + +Two config sources are exercised deliberately: + +* ``doc/Model_Lib/mamba2-2.7b.json`` -- the real published shape, loaded through + HuggingFace's ``AutoConfig`` so a drift between our JSON and ``Mamba2Config`` + shows up here rather than silently downstream. +* a scaled-down ``SimpleNamespace`` -- for the structural checks, so they run + without any model download and stay readable. +""" + +import json +import os +import sys +import unittest +from pathlib import Path +from types import SimpleNamespace + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from parser.llm_parser import LLMModelParser + +MODEL_LIB_CONFIG = Path(__file__).resolve().parents[2] / "doc" / "Model_Lib" / "mamba2-2.7b.json" + +#: The per-layer node sequence Mamba-2 lowers to. One residual, not two: the +#: mixer is a single sublayer, unlike attention + FFN. +EXPECTED_LAYER_SEQUENCE = [ + "normalization", + "projection", + "conv1d", + "ssd_scan", + "gated_rmsnorm", + "projection", + "elementwise_add", +] + + +def _make_small_mamba2_config(num_hidden_layers: int = 2): + """A hardware-shaped but tiny Mamba-2 config. + + ``num_heads * head_dim == expand * hidden_size`` and + ``num_heads % n_groups == 0`` both hold, which the parser enforces. + ``num_heads`` is a multiple of BLEN so ``in_proj_out`` lands on a BLEN + boundary and the projection covers every output column. + """ + return SimpleNamespace( + model_type="mamba2", + architectures=["Mamba2ForCausalLM"], + hidden_size=128, + num_hidden_layers=num_hidden_layers, + expand=2, + head_dim=64, + num_heads=4, + state_size=64, + n_groups=1, + conv_kernel=4, + chunk_size=64, + vocab_size=1024, + layer_norm_epsilon=1e-5, + time_step_limit=(0.0, float("inf")), + time_step_min=0.001, + time_step_max=0.1, + use_conv_bias=True, + use_bias=False, + rms_norm=True, + hidden_act="silu", + tie_word_embeddings=False, + ) + + +def _make_small_parser(num_hidden_layers: int = 2) -> LLMModelParser: + parser = LLMModelParser("mock-mamba2") + parser.config = _make_small_mamba2_config(num_hidden_layers) + parser.model = SimpleNamespace() + return parser + + +def _make_model_lib_parser() -> LLMModelParser: + """Parser bound to the Model_Lib config, loaded via HuggingFace AutoConfig.""" + from transformers import AutoConfig + + parser = LLMModelParser(str(MODEL_LIB_CONFIG)) + parser.config = AutoConfig.from_pretrained(str(MODEL_LIB_CONFIG)) + parser.model = SimpleNamespace() + return parser + + +class TestMamba2ModelLibConfig(unittest.TestCase): + """The published mamba2-2.7b shape, as committed to doc/Model_Lib.""" + + def test_config_file_exists_and_is_json(self): + self.assertTrue(MODEL_LIB_CONFIG.is_file(), f"missing {MODEL_LIB_CONFIG}") + with open(MODEL_LIB_CONFIG) as f: + raw = json.load(f) + self.assertEqual(raw["model_type"], "mamba2") + self.assertEqual(raw["architectures"], ["Mamba2ForCausalLM"]) + + def test_published_dimensions(self): + with open(MODEL_LIB_CONFIG) as f: + raw = json.load(f) + self.assertEqual(raw["hidden_size"], 2560) + self.assertEqual(raw["num_hidden_layers"], 64) + self.assertEqual(raw["expand"], 2) + self.assertEqual(raw["state_size"], 128) + self.assertEqual(raw["conv_kernel"], 4) + self.assertEqual(raw["head_dim"], 64) + self.assertEqual(raw["n_groups"], 1) + self.assertEqual(raw["chunk_size"], 256) + self.assertEqual(raw["vocab_size"], 50288) + + def test_num_heads_is_consistent(self): + """num_heads must equal expand * hidden_size / head_dim. + + A mismatch would silently corrupt d_inner, the dt slice width and the + SSD batch axis, so it is pinned rather than derived. + """ + with open(MODEL_LIB_CONFIG) as f: + raw = json.load(f) + self.assertEqual(raw["num_heads"], raw["expand"] * raw["hidden_size"] // raw["head_dim"]) + self.assertEqual(raw["num_heads"], 80) + + def test_required_mamba2_fields_present(self): + with open(MODEL_LIB_CONFIG) as f: + raw = json.load(f) + for field in ( + "layer_norm_epsilon", + "time_step_limit", + "time_step_min", + "time_step_max", + "use_conv_bias", + "use_bias", + "rms_norm", + "tie_word_embeddings", + "torch_dtype", + ): + self.assertIn(field, raw, f"Model_Lib config is missing {field}") + + def test_loads_as_huggingface_mamba2config(self): + from transformers import AutoConfig, Mamba2Config + + cfg = AutoConfig.from_pretrained(str(MODEL_LIB_CONFIG)) + self.assertIsInstance(cfg, Mamba2Config) + self.assertEqual(cfg.hidden_size, 2560) + self.assertEqual(cfg.num_heads, 80) + self.assertEqual(cfg.n_groups, 1) + self.assertEqual(cfg.state_size, 128) + self.assertEqual(cfg.chunk_size, 256) + self.assertEqual(list(cfg.time_step_limit), [0.0, float("inf")]) + + +class TestMamba2Detection(unittest.TestCase): + def test_model_type_selects_mamba2(self): + self.assertTrue(_make_small_parser().is_mamba2()) + + def test_architecture_string_selects_mamba2_without_model_type(self): + parser = LLMModelParser("mock") + parser.config = SimpleNamespace( + architectures=["Mamba2ForCausalLM"], + hidden_size=128, + num_hidden_layers=1, + expand=2, + head_dim=64, + num_heads=4, + state_size=64, + n_groups=1, + conv_kernel=4, + chunk_size=64, + vocab_size=1024, + ) + parser.model = SimpleNamespace() + self.assertTrue(parser.is_mamba2()) + + def test_llama_config_is_not_mamba2(self): + parser = LLMModelParser("mock-llama") + parser.config = SimpleNamespace( + model_type="llama", + architectures=["LlamaForCausalLM"], + hidden_size=256, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=1, + intermediate_size=512, + head_dim=64, + vocab_size=1024, + ) + parser.model = SimpleNamespace() + self.assertFalse(parser.is_mamba2()) + + +class TestMamba2Dimensions(unittest.TestCase): + def test_model_lib_derived_dimensions(self): + dims = _make_model_lib_parser().extract_critical_dimensions() + self.assertIn("mamba", dims) + m = dims["mamba"] + self.assertEqual(m["d_inner"], 2 * 2560) + self.assertEqual(m["num_heads"] * m["head_dim"], m["d_inner"]) + # conv1d runs over [x, B, C] + self.assertEqual(m["conv_dim"], 5120 + 2 * 1 * 128) + # in_proj emits [z, x, B, C, dt] + self.assertEqual(m["in_proj_out"], 2 * 5120 + 2 * 1 * 128 + 80) + self.assertEqual(m["eps"], 1e-05) + self.assertEqual(m["time_step_min"], 0.0) + self.assertEqual(m["time_step_max"], float("inf")) + self.assertTrue(m["use_conv_bias"]) + self.assertFalse(m["use_bias"]) + + def test_small_config_dimensions(self): + m = _make_small_parser().extract_critical_dimensions()["mamba"] + self.assertEqual(m["d_inner"], 256) + self.assertEqual(m["conv_dim"], 256 + 2 * 64) + self.assertEqual(m["in_proj_out"], 512 + 128 + 4) + self.assertEqual(m["heads_per_group"], 4) + + def test_inconsistent_num_heads_is_rejected(self): + """A config where num_heads * head_dim != expand * hidden_size must fail + loudly, not silently produce a wrong d_inner.""" + parser = _make_small_parser() + parser.config.num_heads = 5 # 5 * 64 != 2 * 128 + with self.assertRaises(ValueError) as ctx: + parser.extract_critical_dimensions() + self.assertIn("num_heads", str(ctx.exception)) + + def test_num_heads_not_divisible_by_groups_is_rejected(self): + parser = _make_small_parser() + parser.config.n_groups = 3 # 4 heads / 3 groups + with self.assertRaises(ValueError) as ctx: + parser.extract_critical_dimensions() + self.assertIn("n_groups", str(ctx.exception)) + + +class TestMamba2SymbolicGraph(unittest.TestCase): + def test_layer_node_sequence(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=64) + types = [n["operation_type"] for n in graph["nodes"]] + self.assertEqual(types[0], "embedding") + self.assertEqual(types[1 : 1 + len(EXPECTED_LAYER_SEQUENCE)], EXPECTED_LAYER_SEQUENCE) + self.assertEqual(types[-2:], ["normalization", "lm_head"]) + + def test_layer_node_names(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=64) + names = [n["name"] for n in graph["nodes"]] + self.assertEqual( + names, + [ + "embed_tokens", + "layer_0_input_layernorm", + "layer_0_in_proj", + "layer_0_conv1d", + "layer_0_ssd_scan", + "layer_0_gated_rmsnorm", + "layer_0_out_proj", + "layer_0_residual", + "final_layernorm", + "lm_head", + ], + ) + self.assertEqual(names, graph["execution_order"]) + + def test_execution_order_is_dense_and_monotonic(self): + graph = _make_small_parser(num_hidden_layers=3).create_symbolic_graph(batch_size=1, seq_len=64) + self.assertEqual([n["execution_order"] for n in graph["nodes"]], list(range(graph["total_nodes"]))) + + def test_node_count_scales_with_layers(self): + for num_layers in (1, 2, 5): + graph = _make_small_parser(num_layers).create_symbolic_graph(batch_size=1, seq_len=64) + # embed + 7 per layer + final_norm + lm_head + self.assertEqual(graph["total_nodes"], 1 + 7 * num_layers + 1 + 1) + self.assertEqual(len(graph["execution_order"]), graph["total_nodes"]) + + def test_projection_dimensions(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=64) + by_name = {n["name"]: n for n in graph["nodes"]} + + in_proj = by_name["layer_0_in_proj"]["dimensions"] + self.assertEqual(in_proj["in_features"], 128) + self.assertEqual(in_proj["out_features"], 644) # 2*256 + 2*1*64 + 4 + # [z, x, B, C, dt] slice offsets must tile the output exactly. + offsets = in_proj["slices"] + self.assertEqual(list(offsets), ["z", "x", "B", "C", "dt"]) + cursor = 0 + for name in ("z", "x", "B", "C", "dt"): + start, width = offsets[name] + self.assertEqual(start, cursor, f"slice {name} does not abut the previous one") + cursor += width + self.assertEqual(cursor, in_proj["out_features"]) + + out_proj = by_name["layer_0_out_proj"]["dimensions"] + self.assertEqual(out_proj["in_features"], 256) + self.assertEqual(out_proj["out_features"], 128) + + def test_conv1d_dimensions(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=64) + conv = {n["name"]: n for n in graph["nodes"]}["layer_0_conv1d"]["dimensions"] + self.assertEqual(conv["conv_dim"], 384) # d_inner + 2 * n_groups * state_size + self.assertEqual(conv["in_channels"], conv["out_channels"]) + self.assertEqual(conv["groups"], conv["in_channels"], "conv1d must be depthwise") + self.assertEqual(conv["kernel_size"], 4) + self.assertEqual(conv["padding"], 3, "causal conv pads kernel_size - 1 on the left only") + self.assertTrue(conv["causal"]) + self.assertTrue(conv["depthwise"]) + self.assertEqual(conv["activation"], "silu") + + def test_ssd_scan_dimensions(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=128) + ssd = {n["name"]: n for n in graph["nodes"]}["layer_0_ssd_scan"]["dimensions"] + self.assertEqual(ssd["num_heads"], 4) + self.assertEqual(ssd["head_dim"], 64) + self.assertEqual(ssd["state_size"], 64) + self.assertEqual(ssd["n_groups"], 1) + self.assertEqual(ssd["chunk_size"], 64) + self.assertEqual(ssd["num_chunks"], 2, "128 timesteps / chunk 64") + self.assertEqual(ssd["time_step_min"], 0.0) + self.assertEqual(ssd["time_step_max"], float("inf")) + + def test_gated_rmsnorm_dimensions(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=64) + gn = {n["name"]: n for n in graph["nodes"]}["layer_0_gated_rmsnorm"]["dimensions"] + # Mamba-2 normalises per group over d_inner / n_groups, not over hidden_size. + self.assertEqual(gn["normalized_shape"], 256) + self.assertEqual(gn["group_size"], 256) + self.assertEqual(gn["norm_type"], "gated_rms_norm") + self.assertEqual(gn["gate_activation"], "silu") + + def test_shapes_thread_through_the_layer(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=64) + by_name = {n["name"]: n for n in graph["nodes"]} + self.assertEqual(by_name["layer_0_in_proj"]["output_shape"], [1, 64, 644]) + self.assertEqual(by_name["layer_0_conv1d"]["output_shape"], [1, 64, 384]) + self.assertEqual(by_name["layer_0_ssd_scan"]["output_shape"], [1, 64, 256]) + self.assertEqual(by_name["layer_0_out_proj"]["output_shape"], [1, 64, 128]) + self.assertEqual(by_name["layer_0_residual"]["output_shape"], [1, 64, 128]) + + def test_num_chunks_scales_with_seq_len(self): + for seq_len, expected_chunks in ((64, 1), (128, 2), (192, 3), (200, 4)): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph( + batch_size=1, seq_len=seq_len + ) + ssd = {n["name"]: n for n in graph["nodes"]}["layer_0_ssd_scan"]["dimensions"] + self.assertEqual(ssd["num_chunks"], expected_chunks, f"seq_len={seq_len}") + # Node count is independent of seq_len: chunking is a lowering + # parameter of one kernel, not graph structure. + self.assertEqual(graph["total_nodes"], 1 + 7 + 1 + 1) + + def test_graph_carries_family_and_shape(self): + graph = _make_small_parser(num_hidden_layers=1).create_symbolic_graph(batch_size=1, seq_len=64) + self.assertEqual(graph["architecture_family"], "mamba2") + shape = graph["mamba_shape"] + self.assertEqual(shape["seq_len"], 64) + self.assertEqual(shape["d_inner"], 256) + # Every mixer node carries the same shape so code_gen can derive one + # consistent VRAM map from any of them. + for node in graph["nodes"]: + node_shape = node["dimensions"].get("mamba_shape") + if node_shape is not None: + self.assertEqual(node_shape, shape) + + def test_model_lib_graph_shape(self): + """Full mamba2-2.7b: 64 layers, correct widths, no ASM generated.""" + graph = _make_model_lib_parser().create_symbolic_graph(batch_size=1, seq_len=256) + self.assertEqual(graph["total_nodes"], 1 + 7 * 64 + 1 + 1) + by_name = {n["name"]: n for n in graph["nodes"]} + self.assertEqual(by_name["layer_0_in_proj"]["dimensions"]["out_features"], 10576) + self.assertEqual(by_name["layer_0_conv1d"]["dimensions"]["conv_dim"], 5376) + self.assertEqual(by_name["layer_63_out_proj"]["dimensions"]["in_features"], 5120) + self.assertEqual(by_name["lm_head"]["dimensions"]["vocab_size"], 50288) + + +class TestOtherFamiliesUnaffected(unittest.TestCase): + """Guards the regression the Mamba-2 branch could most easily cause.""" + + def test_llama_graph_still_has_six_nodes_per_layer(self): + parser = LLMModelParser("mock-llama") + parser.config = SimpleNamespace( + model_type="llama", + architectures=["LlamaForCausalLM"], + hidden_size=256, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=1, + intermediate_size=512, + head_dim=64, + vocab_size=1024, + rms_norm_eps=1e-5, + hidden_act="silu", + ) + parser.model = SimpleNamespace() + graph = parser.create_symbolic_graph(batch_size=1, seq_len=32) + self.assertEqual(graph["total_nodes"], 1 + 6 * 3 + 1 + 1) + self.assertNotIn("architecture_family", graph) + types = [n["operation_type"] for n in graph["nodes"]] + self.assertEqual( + types[1:7], + ["normalization", "attention", "elementwise_add", "normalization", "ffn", "elementwise_add"], + ) + + +if __name__ == "__main__": + unittest.main() From 8cec45c61e90f66521c0925c92e33af0529aa0e6 Mon Sep 17 00:00:00 2001 From: "Qichao (Arlo) Wang" Date: Sat, 29 Aug 2026 01:40:31 +0100 Subject: [PATCH 5/6] kda: the whole layer, projection through mixer, with no state engine Gates, L2 normalisation, the causal convolution, the decode recurrence, chunked prefill and the seam that joins them -- all static, all with addresses baked at assembly time. **Decode.** `kda_decode_step_v0` is 76 static instructions per (head, block) at every key_dim from 4 to 128, against 7,345 at 128 before. Expanded, the issue stream is `16 * key_dim + 46` against a pre-conversion slope of about 57. The mixer streams one head at a time so FPRAM stays bounded. **Prefill past one block.** `key_dim > mlen` used to be refused outright. Lifted by moving the key axis from folded rows onto lanes: the systolic array contracts a VRAM operand's *lanes*, and `vram_sub_projection_asm_impl` already walks column blocks at stride `physical_rows * mlen`, so **zero projection calls changed** in the seven products this touches. `kda_prefill_tile_shapes` becomes the single source of truth for tile geometry -- the divergence between the emitter's idea of the tiles and the caller's is how a spill came to prefetch past its allocation. One real defect fell out. `out = scale * (...)` ended in a single `tile_row_mul_fp_broadcast`, which scales **column block 0 only**, so at `value_dim` 128 half of every token was left `sqrt(key_dim)` = 11.3x too large. The harness reported PASSED at 94.68% of values inside tolerance, because `check_mem`'s bar is 90%. Fixed by looping value blocks; the stage test now restates the bar at 100% on the same numbers, and four mutations that survived the 90% rule die under it. **The projection seam.** Everything past the projection wants a dense tile whose row is a feature block; the projection produces `[tokens, 3*key_width + value_width + heads]`. Those are the same bytes -- column block `c` of any VRAM tile is at `base + c * physical_rows * mlen` -- so `vram_tall_view` and `vram_column_block_view` name one by registering a VRAM object at the computed address and **emit nothing**. The gather that follows is 14 static instructions per section, 70 for a Kimi K3 layer. Its cost is not 70: the loop body runs once per feature block, 192 of them for q, k, v and gate, so it **issues 4,650 instructions against 492,681 for the layer, 0.94%**. It is not caused by packing. `M_MM_WO` writes a `blen x blen` sub-tile and the writeback loops 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 weights are arranged. Five separate projections were built and counted against one packed one: identical `M_MM` count (21,560), 0.12% more instructions for the separate form, the same gather. What removes it is not moving the data. `kda_conv_step_v0` takes `x_new_row_base` and `x_new_row_stride`, and the sweeps underneath collapse any arithmetic progression into one hardware loop whose step is an `S_ADDI_INT` immediate -- a step of `blen` costs what a step of 1 costs. Feeding it the tall view directly removes **4,041 of the 4,650** issued instructions for nothing. Gate and beta reach a different consumer and still gather, at 1,176 issued, 0.24%; removing those needs either `(dst, src)` row pairs in a row-op emitter every kernel shares, or padding `dt_bias` from 192 rows to 768 at 6.8 MiB more HBM traffic per token. Both trades are bad and the test records the price. **No state engine.** `test_no_state_engine.py` pins all three opcodes by slot. Nothing here reads a descriptor, holds a queue, or keeps residency. The recurrence runs on the ordinary vector unit, which is what lets a consumer read an operand at a stride instead of needing the data rearranged for it. Co-Authored-By: Claude Opus 5 --- aten/plena/program_kda_chunk.py | 298 +++++++ aten/plena/program_kda_common.py | 640 +++++++++++++++ aten/plena/program_kda_conv.py | 192 +++++ aten/plena/program_kda_gates.py | 248 ++++++ aten/plena/program_kda_layer.py | 402 ++++++++++ aten/plena/program_kda_mixer.py | 238 ++++++ aten/plena/program_kda_prefill.py | 820 ++++++++++++++++++++ aten/plena/program_kda_recurrent.py | 290 +++++++ aten/tests/test_kda_conv.py | 246 ++++++ aten/tests/test_kda_decode_step.py | 619 +++++++++++++++ aten/tests/test_kda_gates.py | 392 ++++++++++ aten/tests/test_kda_layer.py | 942 +++++++++++++++++++++++ aten/tests/test_kda_lowering.py | 461 +++++++++++ aten/tests/test_kda_mixer.py | 471 ++++++++++++ aten/tests/test_kda_normalize.py | 301 ++++++++ aten/tests/test_kda_prefill.py | 336 ++++++++ aten/tests/test_kda_prefill_structure.py | 558 ++++++++++++++ aten/tests/test_kda_stage_contract.py | 112 +++ aten/tests/test_spill_precision_guard.py | 99 +++ 19 files changed, 7665 insertions(+) create mode 100644 aten/plena/program_kda_chunk.py create mode 100644 aten/plena/program_kda_common.py create mode 100644 aten/plena/program_kda_conv.py create mode 100644 aten/plena/program_kda_gates.py create mode 100644 aten/plena/program_kda_layer.py create mode 100644 aten/plena/program_kda_mixer.py create mode 100644 aten/plena/program_kda_prefill.py create mode 100644 aten/plena/program_kda_recurrent.py create mode 100644 aten/tests/test_kda_conv.py create mode 100644 aten/tests/test_kda_decode_step.py create mode 100644 aten/tests/test_kda_gates.py create mode 100644 aten/tests/test_kda_layer.py create mode 100644 aten/tests/test_kda_lowering.py create mode 100644 aten/tests/test_kda_mixer.py create mode 100644 aten/tests/test_kda_normalize.py create mode 100644 aten/tests/test_kda_prefill.py create mode 100644 aten/tests/test_kda_prefill_structure.py create mode 100644 aten/tests/test_kda_stage_contract.py create mode 100644 aten/tests/test_spill_precision_guard.py diff --git a/aten/plena/program_kda_chunk.py b/aten/plena/program_kda_chunk.py new file mode 100644 index 0000000..b8e0426 --- /dev/null +++ b/aten/plena/program_kda_chunk.py @@ -0,0 +1,298 @@ +"""KDA chunk primitives: the cumulative decay, and the UT transform. + +Chunked prefill collapses a chunk of ``C`` tokens into four products against the +incoming state, instead of ``C`` sequential rank-1 updates. For the gated delta +rule the collapse is not free: within a chunk the update is a product of rank-1 +projectors, and turning that into a single matrix needs the WY / UT transform, +whose core is the inverse of a ``[C, C]`` lower-triangular matrix. + +The algebra, verified against ``kda_step`` run sequentially +----------------------------------------------------------- +With state ``S[value, key]``, per-step decay ``a_t[key]`` and cumulative decay +``A_t[key] = prod_{s<=t} a_s[key]``:: + + k_hat_t = k_t * A_t bounded by |k|, since A <= 1 + q_hat_t = q_t * A_t + M[t, s] = sum_key k_t k_s A_t/A_s for s < t + N[t, s] = sum_key q_t k_s A_t/A_s for s <= t + T = (I + tril(diag(beta) M, -1))^-1 diag(beta) + E = T (V - K_hat S_0^T) [C, value] + out_t = scale * (S_0 q_hat_t + sum_{s<=t} N[t,s] E_s) + S_C = A_C * S_0 + sum_s E_s outer (k_s * A_C/A_s) + +**KDA's decay is channel-wise on the key axis**, not one scalar per timestep -- +``reference.py`` applies ``exp(log_decay)[:, :, None, :]``. The textbook chunking +assumes a scalar decay and gives ``M = tril(diag(beta) K K^T, -1)`` with no decay +weighting at all; against the sequential reference that form is wrong by 1.8e-1, +which is an error and not a rounding. The ``A_t/A_s`` weighting is not optional. + +Two numerical constraints, both measured +---------------------------------------- +**The cumulative decay is a running product, not ``exp`` of a running sum.** +``A_t = exp(c_t)`` where ``c_t`` is the cumulative log-decay, and it is tempting +to follow ``program_ssd.py``'s cumsum -- a matmul against lower-triangular ones, +accumulated in f32. But ``c`` reaches ``chunk * gate_lower_bound = -80``, and at +that magnitude bf16's ulp is 0.31, so storing ``c`` and *then* exponentiating +costs a 17% relative error on ``A``. Multiplying the per-step decays instead +keeps every intermediate in ``[e^-5, 1]``, where bf16's error is relative rather +than magnified. End-to-end against the sequential reference, chunk 16, key 128: + + route output error state error + cumsum in f32 8.1e-05 9.5e-04 (the algebra's own floor) + cumsum in bf16 3.7e-04 1.1e-02 + cumprod in bf16 5.9e-05 2.4e-03 + +**``chunk * |gate_lower_bound|`` must stay under 88.** ``M`` and ``N`` are formed +as a matmul of ``k * A`` against ``k / A``, and ``1/A`` reaches +``exp(chunk * |lower_bound|)``. bf16's maximum is 3.39e38, whose log is 88.72, so +at Kimi K3's lower bound of -5 the last chunk size that works is **17**. Chunk 18 +overflows to inf. That is why this file asserts the product rather than the chunk +size: raising either without the other is what breaks it. + +The reciprocal is only ever used inside a *difference* -- every place ``1/A_s`` +appears it is multiplied by some ``A_t`` with ``t > s``, so the product is +bounded by 1. Forming the two halves separately is what costs the range, and it +is done that way because it is the only form the matrix engine can contract. The +cancellation this causes is severe in ``M``'s smallest entries (up to 315% +relative) but does not reach the output: those entries contribute nothing, and +end-to-end the matmul form and an exactly-computed one agree to 6e-5. + +Where the key axis lives: on lanes +---------------------------------- +A ``[chunk, key_dim]`` tile here is ``chunk`` rows of ``key_dim`` lanes, spread +across ``kda_key_blocks`` **column blocks**. It is not the row-folded +``[chunk * key_blocks, mlen]`` layout the decode path uses for its state. + +The choice is forced by the projections, and it is what lifts ``key_dim > mlen``. +Five of prefill's seven products contract over the key axis, and the systolic +array contracts a VRAM operand's *lanes*: ``vram_sub_projection_asm_impl`` walks +the operand's column blocks at stride ``physical_rows * mlen`` and accumulates +inside the array, so a key axis on lanes contracts across as many blocks as it +needs with no extra instruction and no explicit sum. With the key block folded +into rows instead, one projection can only ever see ``mlen`` of the key axis, and +a gram of the folded tile computes the ``kb != kb'`` cross terms as well -- half +the MACs thrown away, plus a strided gather to recover the diagonal. + +The row ops pay for it, and the price is small: what used to be one sweep over +``chunk * key_blocks`` rows is ``key_blocks`` sweeps over ``chunk`` rows, passing +``tile_col_idx``. Same instruction count to within the loop headers, and rows +inside a column block are still consecutive, so each sweep is still one hardware +loop. +""" + +from __future__ import annotations + +# NOTE: these emitters bill to the existing ``kda_decay`` and +# ``kda_state_update`` stages. Prefill deserves its own stage names so Phase 4 +# can separate its cost from decode's, but KDA_STAGES is guarded bidirectionally +# against the emulator's stage_profile.rs, so adding one means editing both +# repos together. Done with the prefill layer, not here. + +import math + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena.program_kda_common import kda_stage_marker +from compiler.aten.plena.program_kda_gates import kda_key_blocks +from compiler.aten.plena.vars import FPVar, VRAMMatrixVar + +__all__ = [ + "ProgramKdaChunkMixin", + "kda_chunk_cols", + "kda_chunk_rows", + "kda_max_chunk_for", + "kda_round_up", +] + +#: ``ln`` of bf16's largest finite value. ``1/A`` over a chunk reaches +#: ``exp(chunk * |gate_lower_bound|)``, and this is where that stops being +#: representable. +_BF16_MAX_LOG = math.log(3.3895313892515355e38) + + +def kda_max_chunk_for(gate_lower_bound: float) -> int: + """Largest chunk whose ``1/A`` still fits bf16. 17 at Kimi K3's -5.""" + if gate_lower_bound >= 0: + raise ValueError(f"gate_lower_bound must be negative, got {gate_lower_bound}") + return int(_BF16_MAX_LOG // abs(gate_lower_bound)) + + +def kda_round_up(n: int, mlen: int) -> int: + """``n`` rounded up to a whole number of ``mlen``-sized blocks, minimum one. + + Every prefill tile is allocated at block granularity in both axes. Rows, + because a projection writes a whole ``mlen x mlen`` block and a spill + prefetches one; columns, because ``_target_tile_addr`` places column block + ``c`` at ``c * physical_rows * mlen``, so a tile whose declared width is not + a multiple of ``mlen`` reserves less than the blocks it will be addressed by + and the last one runs off the end. + """ + if mlen <= 0: + raise ValueError(f"mlen must be positive, got {mlen}") + return max(1, -(-n // mlen)) * mlen + + +def kda_chunk_rows(shape: KdaShape, mlen: int, chunk: int) -> int: + """Rows of a ``[chunk, key_dim]`` prefill tile: one per timestep. + + The key axis is on **lanes**, not folded into rows -- see the module + docstring. ``shape`` is taken for the key-block validation it carries and + for symmetry with the rest of the family. + """ + kda_key_blocks(shape, mlen) + return kda_round_up(chunk, mlen) + + +def kda_chunk_cols(shape: KdaShape, mlen: int) -> int: + """Columns of a ``[chunk, key_dim]`` prefill tile: the whole key axis.""" + return kda_key_blocks(shape, mlen) * mlen + + +class ProgramKdaChunkMixin: + """Chunk-level primitives for KDA prefill.""" + + def kda_chunk_check_range(self, chunk: int, shape: KdaShape) -> None: + """Refuse a chunk whose reciprocal decay would overflow bf16. + + Asserted rather than commented because the failure is silent: ``1/A`` + becomes ``inf``, ``M`` becomes ``nan``, and the triangular solve returns + ``nan`` for the whole chunk. + """ + if chunk < 1: + raise ValueError(f"chunk must be positive, got {chunk}") + span = chunk * abs(shape.gate_lower_bound) + if span >= _BF16_MAX_LOG: + raise ValueError( + f"chunk {chunk} at gate_lower_bound {shape.gate_lower_bound} gives a " + f"reciprocal decay of exp({span:.1f}), past bf16's exp({_BF16_MAX_LOG:.1f}). " + f"The largest chunk that fits is " + f"{kda_max_chunk_for(shape.gate_lower_bound)}" + ) + + def kda_chunk_decay_cumprod_v0( + self, + *, + decay: VRAMMatrixVar, + prev: VRAMMatrixVar, + chunk: int, + shape: KdaShape, + ) -> VRAMMatrixVar: + """``decay[t] <- prod_{s<=t} decay[s]``, in place, per key channel. + + ``decay`` arrives holding the **per-step** decay ``exp(log_decay[t, key])`` + -- which is exactly what :meth:`kda_decay_scalars_v0` leaves in ``gate``. + It leaves holding the cumulative decay ``A_t``. + + A running product and not ``exp`` of a running sum; the module docstring + has the measurement. ``prev`` is scratch of the same shape: the row ops + take one row list for both operands, so the previous timestep has to be + staged at the *current* timestep's row before the multiply. + """ + self.kda_chunk_check_range(chunk, shape) + blocks = kda_key_blocks(shape, self.mlen) + rows = kda_chunk_rows(shape, self.mlen, chunk) + cols = kda_chunk_cols(shape, self.mlen) + for name, tile in (("decay", decay), ("prev", prev)): + if tile.shape[1] != cols: + raise ValueError( + f"{name} is {tile.shape[1]} columns wide; a [chunk, key_dim] " + f"tile carries the key axis on lanes and must be exactly " + f"key_dim ({cols}) wide" + ) + if tile.shape[0] < rows: + raise ValueError(f"{name} needs {rows} rows, has {tile.shape[0]}") + if prev.name == decay.name: + raise ValueError("prev must be distinct from decay") + + self.emit_comment(kda_stage_marker("kda_decay", f"chunk cumprod chunk={chunk}")) + for t in range(1, chunk): + # One copy covers every key block: `mamba_block_copy` goes through + # `vram_fill_zero` and `vram_add`, both of which walk the column + # blocks themselves. Only the multiply has to be driven per block, + # because a binary row op pairs one block with one block. + self.mamba_block_copy( + prev, decay, + dst_row_offset=t, src_row_offset=t - 1, num_rows=1, + ) + for kb in range(blocks): + self.tile_row_mul(decay, prev, rows=[t], tile_col_idx=kb) + return decay + + def kda_ut_transform_v0( + self, + *, + m: VRAMMatrixVar, + identity: VRAMMatrixVar, + beta_fp: FPVar, + t_out: VRAMMatrixVar, + m_fp: FPVar, + consts, + chunk: int, + shape: KdaShape, + beta_offset: int = 0, + ) -> VRAMMatrixVar: + """``T = (I + tril(diag(beta) M, -1))^-1 diag(beta)``, by forward substitution. + + ``m`` is the ``[chunk, chunk]`` decay-weighted key gram, consumed in place + and left negated. ``identity`` is a host-staged ``[chunk, chunk]`` identity + -- staged rather than built, for the reason ``ssd_lower_triangular_ones`` + gives: materialising a constant tile on chip costs one ``S_ST_FP`` per + element. ``m_fp`` needs one padded row (``mlen`` slots). + + Solving ``L T = diag(beta)`` row by row, with ``L = I + tril(diag(beta)M, -1)`` + lower unitriangular:: + + T[i] = beta_i * (e_i - sum_{j self.mlen: + raise ValueError( + f"chunk {chunk} exceeds mlen {self.mlen}; the [chunk, chunk] tiles " + f"here are one row per timestep and must fit a single block" + ) + if m_fp.size < self.mlen: + raise ValueError( + f"m_fp holds {m_fp.size} slots but S_MAP_FP_V writes a whole row " + f"({self.mlen}); size it to mlen, not to chunk" + ) + if beta_fp.size < beta_offset + chunk: + raise ValueError( + f"beta_fp needs {beta_offset + chunk} slots, has {beta_fp.size}" + ) + for name, tile in (("t_out", t_out), ("identity", identity)): + if tile.name == m.name: + raise ValueError(f"{name} must not alias m") + + self.emit_comment(kda_stage_marker("kda_state_update", f"UT transform chunk={chunk}")) + rows = list(range(chunk)) + # The sweep accumulates, so negate once rather than subtracting per row. + self.tile_row_mul_fp_broadcast(m, consts.neg_one, rows=rows) + + for i in rows: + self.mamba_row_copy(t_out, i, identity, i) + if i: + # Row i of the negated M into FPRAM, then one loop over j < i. + self.tile_row_to_fpram(m, m_fp, rows=[i]) + self.tile_row_fma_fp_sweep( + t_out, t_out, m_fp, + dst_rows=[i] * i, src_rows=list(range(i)), + ) + self.tile_row_mul_fp_broadcast( + t_out, beta_fp, rows=[i], fpram_offset=beta_offset + i + ) + return t_out diff --git a/aten/plena/program_kda_common.py b/aten/plena/program_kda_common.py new file mode 100644 index 0000000..8fb7ba5 --- /dev/null +++ b/aten/plena/program_kda_common.py @@ -0,0 +1,640 @@ +"""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]`` in +``aten/models/kda/reference.py``. This is deliberate, and it is the decision the +whole KDA lowering rests on. + +In the reference orientation the per-key decay is strided across rows and both +contractions (``prediction = S @ k`` and ``output = S @ q``) run *within* a row, +so each costs a ``V_MUL_VV`` plus a ``V_RED_SUM`` per row and lands its result in +an FP register that then has to be gathered back into a vector. + +Transposed, ``key`` becomes the row axis and all four steps of the recurrence +reduce to "one row times one scalar":: + + 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 +``_row_progression`` recognises and turns into a hardware loop rather than an +unrolled block. ``aten/tests/test_kda_reference.py``'s +``test_decode_step_matches_the_transposed_formulation`` pins the two orientations +together so the lowering cannot be validated against a plausible-but-different +reimplementation of the maths. + +Mamba-2 already stores state the same way -- ``state_size`` as the row axis, see +``program_ssm_recurrent.py`` -- so the two algorithms share the emitters rather +than each getting their own. + +**Prefill does not use this layout.** ``program_kda_prefill`` holds the state +``[value, key]``, because that is what makes its seven chunk-level products land +on the projection primitives without an explicit transpose. At Kimi K3 +``key_dim == value_dim == 128``, so the two shapes match and passing one path's +state to the other is a finite, plausible, wrong answer rather than an error. +The conversion is ``kda_prefill_state_to_decode_layout_v0``, and it belongs at +the boundary -- once per layer, not inside either path. + +Column blocks are folded into the row index +------------------------------------------- +A VRAM row is ``mlen`` elements wide, and Kimi K3's ``value_dim`` is 128 against +a default ``mlen`` of 64. Storing state as ``[key, value]`` would put one key's +row across two column blocks, and the helper family that would have to walk them +is inconsistent about it: ``tile_row_mul_fp`` takes a ``tile_col_idx``, but +``tile_row_sum``, ``tile_row_mul``, ``tile_row_sub``, ``vram_fill_zero``, +``mamba_row_copy`` and ``mamba_row_add`` do not. Mixing "walks every block" with +"silently block 0 only" is how a state ends up correct in its first 64 lanes and +stale in the rest -- no error, plausible numbers. + +So the block index is folded into the row instead: +``row = (head * blocks + block) * key_dim + key``. Every row is then exactly one +block wide, no helper needs a column argument, and fixing ``(head, block)`` +leaves the keys at consecutive rows -- a unit-stride progression, so each sweep +is still one hardware loop. + +When ``value_dim == mlen`` this is byte-identical to the un-blocked layout, so it +is a generalisation rather than a second scheme to maintain. ``Mamba2Shape`` +takes the other route for the same problem -- it *rejects* shapes wider than +``mlen`` (``program_mamba_common.py:175-188``) and its error text names the +helper fix as "deliberately not done yet". Folding the block into the row gets +the same safety without touching emitters that Mamba and attention are using. + +Precision -- and what this module can and cannot promise +-------------------------------------------------------- +State travels through the ``keyvalue`` precision class (``precision=1``), the +same one ``ssm_load_state_v0`` uses. **That class selects a name, not a width.** +The width comes from the active ``[.PRECISION]`` table, and the shipped +one declares ``HBM_V_KV_TYPE`` as ``format = "Mx"`` with e4m3 elements -- +1 byte plus a scale stream, not 2 and not 4. ``storage_precision`` only feeds +the compiler's own address arithmetic; it cannot change what the DMA decodes. + +So this module **requires** Plain BF16 KV types rather than asserting them. +Call :meth:`kda_require_state_precision_v0` with the parsed precision table +before lowering a layer. Under an MX KV type the state decodes as e4m3 while +the address stride assumes 2 bytes: garbage, silently, with no runtime error. +``f5eb36a`` found and fixed exactly this for the Mamba SSD path +(``require_bf16_kv_precision``); KDA reuses that guard rather than repeating +the mistake. + +Why BF16 and not something narrower: 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`` -- exactly the long-memory regime where the amplification +is worst. e4m3's 3 mantissa bits compound badly there. FP32 would be better +still, but there is no FP32 path through the KV precision class, so BF16 is +the widest this ISA offers and the numerical claim is scoped to it. + +Load and store must agree on both ``storage_precision`` and +``hbm_element_bytes``: a mismatch changes the row stride and the scale-section +base, so the state read back is not the state written -- again a wrong answer +rather than an error. Both are 2 here, and +:meth:`kda_pin_state_v0` reserves at the same width so the write-back cannot +overrun onto the next tensor. +""" + +from __future__ import annotations + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena.vars import FPVar, InputVar, VRAMMatrixVar + +__all__ = [ + "KDA_STAGES", + "KdaShape", + "ProgramKdaCommonMixin", + "kda_blocks", + "kda_conv_blocks", + "kda_conv_channels", + "kda_conv_state_row", + "kda_head_base", + "kda_stage_marker", + "kda_state_row", + "kda_state_rows", + "kda_vector_row", + "kda_vector_rows", +] + +KDA_STAGE_MARKER_PREFIX = "@stage=" + +#: Every KDA stage name the emulator's ``StageKind`` understands. +#: +#: Kept in lockstep with ``transactional_emulator/src/stage_profile.rs``; the +#: cross-repo test in ``aten/tests/test_kda_stage_contract.py`` pins it. Marker +#: emission is validated against this set so a typo fails at ASM-gen time rather +#: than silently collapsing a region into the preceding stage. +KDA_STAGES: frozenset[str] = frozenset( + { + "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 kda_stage_marker(stage: str, detail: str = "") -> str: + """Format the explicit stage marker comment for ``stage``. + + The ``@`` is load-bearing. ``stage_profile.rs``'s ``extract_stage_tag`` + matches on ``"@stage="``; a bare ``stage=`` comment is not a marker at all, + so every instruction under it bills to the *previous* marker -- silently, + and the resulting profile looks perfectly plausible. + + Markers are authoritative and sticky, exactly as in ``mamba_stage_marker``: + once a program contains any marker the emulator stops applying its legacy + substring rules. Emit one whenever the stage changes and never mid-stage. + """ + if stage not in KDA_STAGES: + raise ValueError(f"unknown KDA stage {stage!r}; expected one of {sorted(KDA_STAGES)}") + return f"{KDA_STAGE_MARKER_PREFIX}{stage}" + (f" {detail}" if detail else "") + + +def kda_blocks(shape: KdaShape, mlen: int) -> int: + """Column blocks one value vector spans. + + A VRAM row is ``mlen`` elements wide, so a ``value_dim``-wide vector needs + ``value_dim / mlen`` of them. Kimi K3 is ``value_dim = 128`` against a + default ``mlen = 64``, so 2. + """ + if mlen <= 0: + raise ValueError(f"mlen must be positive, got {mlen}") + if shape.value_dim % mlen: + raise ValueError( + f"value_dim ({shape.value_dim}) must be a multiple of mlen ({mlen}); a " + f"partial trailing block would leave lanes past value_dim holding " + f"whatever was there before, and they would still be summed" + ) + return shape.value_dim // mlen + + +def kda_state_rows(shape: KdaShape, mlen: int) -> int: + """Rows of the flattened state tile: one per ``(head, block, key)``.""" + return shape.num_heads * kda_blocks(shape, mlen) * shape.key_dim + + +def kda_state_row(shape: KdaShape, mlen: int, head: int, block: int, key: int) -> int: + """Row holding ``block`` of head ``head``'s state at key ``key``. + + ``(head * blocks + block) * key_dim + key`` -- head outermost, then column + block, then key. Fixing ``(head, block)`` therefore gives keys + ``0..key_dim-1`` at consecutive rows, a unit-stride arithmetic progression, + which is what ``_row_progression`` needs to emit a hardware loop instead of + an unrolled block. + """ + blocks = kda_blocks(shape, mlen) + if not 0 <= head < shape.num_heads: + raise ValueError(f"head {head} out of range for {shape.num_heads} heads") + if not 0 <= block < blocks: + raise ValueError(f"block {block} out of range for {blocks} blocks") + if not 0 <= key < shape.key_dim: + raise ValueError(f"key {key} out of range for key_dim {shape.key_dim}") + return (head * blocks + block) * shape.key_dim + key + + +def kda_head_base(shape: KdaShape, mlen: int, head: int, block: int = 0) -> int: + """First state row of ``head``'s ``block``.""" + return kda_state_row(shape, mlen, head, block, 0) + + +def kda_vector_rows(shape: KdaShape, mlen: int) -> int: + """Rows of a flattened ``[num_heads, value_dim]`` tile. + + ``q``, ``v``, the output, and the two accumulators are all value-width, so + they are blocked the same way the state is -- one row per + ``(head, block)``. Keeping every value-width tile on the same convention is + the point: a mixture is how a sweep ends up reading block 0 of one tile + against block 1 of another. + """ + return shape.num_heads * kda_blocks(shape, mlen) + + +def kda_vector_row(shape: KdaShape, mlen: int, head: int, block: int) -> int: + """Row holding ``block`` of head ``head``'s value-width vector.""" + blocks = kda_blocks(shape, mlen) + if not 0 <= head < shape.num_heads: + raise ValueError(f"head {head} out of range for {shape.num_heads} heads") + if not 0 <= block < blocks: + raise ValueError(f"block {block} out of range for {blocks} blocks") + return head * blocks + block + + +def kda_conv_blocks(channels: int, mlen: int) -> int: + """Column blocks a ``channels``-wide vector spans. + + KDA's conv channel counts are ``num_heads * key_dim`` for q and k and + ``num_heads * value_dim`` for v -- 12,288 and 12,288 for Kimi K3. Everything + the conv does is elementwise across channels, so the blocks are independent + and this is a plain loop bound, not a correctness question the way the + key-axis reduction is. + """ + if mlen <= 0: + raise ValueError(f"mlen must be positive, got {mlen}") + if channels % mlen: + raise ValueError( + f"channels ({channels}) must be a multiple of mlen ({mlen}); a partial " + f"trailing block would convolve lanes that hold nothing" + ) + return channels // mlen + + +def kda_conv_state_row(channels: int, mlen: int, kernel: int, block: int, tap: int) -> int: + """Row holding tap ``tap`` of channel block ``block``. + + ``block * kernel + tap``, so one block's taps are consecutive and the + history shift is a run of adjacent row copies -- the same reason the state + layout puts keys consecutive within a ``(head, block)``. + + ``tap`` runs oldest (0) to newest (``kernel - 1``), matching the reference's + ``torch.roll(state, shifts=-1, dims=-1)`` followed by writing the new value + at ``[..., -1]``. + """ + blocks = kda_conv_blocks(channels, mlen) + if not 0 <= block < blocks: + raise ValueError(f"block {block} out of range for {blocks} blocks") + if not 0 <= tap < kernel: + raise ValueError(f"tap {tap} out of range for kernel {kernel}") + return block * kernel + tap + + +def kda_conv_channels(shape: KdaShape) -> int: + """Channels in the q/k/v short-convolution history. + + ``2 * key_dim + value_dim`` per head: q and k are key-width, v is + value-width. Matches ``KdaShape.conv_state_elements`` in the reference. + """ + return shape.num_heads * (2 * shape.key_dim + shape.value_dim) + + +class ProgramKdaCommonMixin: + """State residency and data movement for KDA. + + Requires ``ProgramSSMRecurrentMixin`` (``pin_hbm_region``), + ``ProgramMambaCommonMixin`` (``mamba_row_copy``, ``mamba_rsqrt_fpram``) and + ``ProgramTensorMixin`` (``input``, ``load_batch``, ``store``). + """ + + # -- residency --------------------------------------------------------- + + def kda_require_state_precision_v0(self, settings: dict | None = None) -> None: + """Fail unless this build configures the KV precision classes as Plain BF16. + + Delegates to ``require_bf16_kv_precision``; see the module docstring for + why a check and not an assertion. Passing ``settings=None`` raises rather + than passing, because a check that could not run must not look like a + check that passed. + """ + self.require_bf16_kv_precision(settings) + + def kda_pin_state_v0(self, name: str, shape: KdaShape) -> int: + """Reserve the pinned HBM range for one head-group's recurrent state. + + Sized from the flattened tile, ``state_rows x mlen``. That is the same + element count as ``num_heads * key_dim * value_dim`` -- folding the block + into the row redistributes elements, it does not add any. + """ + return self.pin_hbm_region( + name, + kda_state_rows(shape, self.mlen) * self.mlen, + hbm_element_bytes=2, + ) + + def kda_load_state_v0(self, name: str, shape: KdaShape, hbm_addr: int) -> VRAMMatrixVar: + """Prefetch the pinned FP32 state tile into VRAM, transposed. + + Returns a ``[num_heads * blocks * key_dim, mlen]`` tile. Use + :func:`kda_state_row` to index it; head ``h``'s block ``c`` occupies + ``key_dim`` consecutive rows starting at + ``(h * blocks + c) * key_dim``. + """ + rows = kda_state_rows(shape, self.mlen) + self.emit_comment(kda_stage_marker("kda_state_load", f"{name} [{rows},{self.mlen}]")) + var = self.input( + name, + (rows, self.mlen), + hbm_addr=hbm_addr, + real_data_ratio=1.0, + ) + return self.load_batch( + var, name=f"{name}_vram", storage_precision=2, precision=1 + ) + + def kda_store_state_v0( + self, state: VRAMMatrixVar, name: str, hbm_addr: int + ) -> InputVar: + """Write the updated state back over the same pinned HBM range. + + Mirrors :meth:`kda_load_state_v0` exactly -- same precision class, + same bytes per element. See the module docstring for why that matters. + """ + self.emit_comment(kda_stage_marker("kda_state_store", name)) + return self.store( + state, + name=name, + hbm_addr=hbm_addr, + precision=1, + hbm_element_bytes=2, + real_data_ratio=1.0, + ) + + # -- short convolution history ---------------------------------------- + + def kda_conv_state_roll_v0( + self, + conv_state: VRAMMatrixVar, + new_row_src: VRAMMatrixVar, + new_row_idx: int, + shape: KdaShape, + ) -> VRAMMatrixVar: + """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, + not a ring pointer. Address immediates are baked at ASM-gen time and + there is no data-dependent addressing, so a runtime read pointer is not + expressible. It costs ``conv_kernel - 1`` row copies per token. + """ + 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(kda_stage_marker("kda_conv1d", f"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 + + # -- normalisation ----------------------------------------------------- + + def kda_l2_normalize_v0( + self, + vec: VRAMMatrixVar, + rows: list[int], + sq_scratch: VRAMMatrixVar, + acc_fp: FPVar, + consts, + ) -> VRAMMatrixVar: + """In-place L2 normalisation of the listed VRAM rows. + + ``x <- x * rsqrt(sum(x^2) + eps)``. + + The epsilon goes **inside** the rsqrt, matching FlashKDA's recurrent + kernel and the CPU reference (``reference.py``: ``rsqrt(sum + 1e-6)``). + ``torch.nn.functional.normalize`` clamps the norm instead and does not + agree; see the D3 note in the execution log for how easily that + difference hides under a loose tolerance. + + ``consts.reci_group`` must be ``1.0``: ``mamba_rsqrt_fpram`` computes + ``1 / sqrt(acc * reci_group + eps)``, and L2 is ``sqrt(sum)``, not + ``sqrt(mean)``. Passing a gated-RMSNorm constants block here, whose + ``reci_group`` is ``1 / group_width``, silently normalises by the wrong + factor -- and produces plausible finite numbers while doing it. Callers + get a constants block from :meth:`kda_fp_constants`. + + Follows ``mamba_gated_rmsnorm_v0`` step for step: block-copy, square by + multiplying the copy against the source, ``V_RED_SUM`` into FPRAM, + scalar rsqrt, then a per-row broadcast multiply. There is no vector + square root -- ``S_FP_OP`` carries ``SQRT_FP``, ``V_ELEMENT_OP`` does + not -- so the rsqrt is unavoidably scalar, one FP op per row. + + **Each row is normalised by its own norm.** That is right for a vector + that fits one row, and wrong for one that does not. + + This matters because KDA normalises ``q`` and ``k``, which are + *key*-width, and Kimi K3 is ``key_dim = 128`` against ``mlen = 64``. The + row folding that solved ``value_dim`` does **not** apply here: the L2 + norm contracts over the key axis, so the two halves of one ``q`` share a + single norm. Splitting them into two rows and calling this once would + normalise each half by its own partial norm -- finite, plausible, wrong. + + Nor can the sums simply be accumulated: ``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. + + The layer lowering therefore has to reduce each block into its own slot + and combine them with scalar ``S_ADD_FP`` before the rsqrt -- ``blocks-1`` + extra scalar adds per vector. Until that exists, callers must pass + vectors that fit one row. + """ + if not rows: + return vec + if sq_scratch.name == vec.name: + # mamba_block_copy is zero-then-add, so dst is src zeroes the vector + # and every later step operates on zeros: sum 0, rsqrt(1e-6) = 1000, + # vec *= 1000 on zeros. Silently returns an all-zero vector. + raise ValueError("kda_l2_normalize_v0 needs a scratch tile distinct from vec") + if acc_fp.size < len(rows): + raise ValueError( + f"acc_fp holds {acc_fp.size} slots but {len(rows)} rows are being " + f"normalised; tile_row_sum writes one slot per row" + ) + if vec.shape[1] != self.mlen: + # Every tile in the KDA path is exactly one column block wide (see + # the module docstring). A wider one would be normalised by a + # partial norm, because tile_row_mul / tile_row_sum / tile_row_mul_fp + # all default to tile_col_idx=0 while vram_fill_zero walks them all. + raise ValueError( + f"kda_l2_normalize_v0 needs a tile exactly mlen ({self.mlen}) wide, " + f"got {vec.shape[1]}" + ) + self.emit_comment(kda_stage_marker("kda_normalize", f"rows={len(rows)}")) + self.mamba_block_copy(sq_scratch, vec, num_rows=vec.shape[0]) + self.tile_row_mul(sq_scratch, vec, rows=rows) + self.tile_row_sum(acc_fp, sq_scratch, rows=rows, target_base_offset=0) + self.mamba_rsqrt_fpram(acc_fp, consts, count=len(rows)) + self.tile_row_mul_fp(vec, acc_fp, rows=rows, fpram_base_offset=0) + return vec + + def _kda_scale_rows( + self, + tile: VRAMMatrixVar, + fp: FPVar, + rows: list[int], + fpram_offset: int, + ) -> None: + """``tile[rows[i]] *= FPRAM[fp + fpram_offset + i]``. + + Wraps a trap in ``_fpram_row_map``: with more than one row it walks + ``base_offset + i``, but with exactly one row it ignores ``base_offset`` + and uses ``single_offset`` instead. Passing ``fpram_base_offset`` for a + single row therefore silently reads slot 0 -- which for the per-key + scalars here means every row scaled by head 0's first key. + """ + if len(rows) == 1: + self.tile_row_mul_fp(tile, fp, rows=rows, fpram_offset=fpram_offset) + else: + self.tile_row_mul_fp(tile, fp, rows=rows, fpram_base_offset=fpram_offset) + + def _kda_row_sum(self, target_fp, source, rows: list[int], offset: int) -> None: + """``target_fp[offset + i] = sum(source[rows[i]])``. + + Wraps the same ``_fpram_row_map`` trap :meth:`_kda_scale_rows` above documents: + with more than one row it walks ``base_offset + i``, but with exactly one + it ignores ``base_offset`` and uses ``single_offset``. Passing the wrong + one sends every call to slot 0, which for a per-block reduction means + only the last block survives -- finite, plausible, wrong, and invisible + at ``vectors > 1``. + """ + if len(rows) == 1: + self.tile_row_sum(target_fp, source, rows=rows, target_offset=offset) + else: + self.tile_row_sum(target_fp, source, rows=rows, target_base_offset=offset) + + def kda_l2_normalize_blocked_v0( + self, + vec: VRAMMatrixVar, + *, + vectors: int, + blocks: int, + sq_scratch: VRAMMatrixVar, + part_fp: FPVar, + acc_fp: FPVar, + consts, + first_row: int = 0, + ) -> VRAMMatrixVar: + """L2-normalise ``vectors`` vectors that each span ``blocks`` rows. + + ``vec`` is ``[vectors * blocks, mlen]`` with row ``v * blocks + c``. + + This is the cross-block case :meth:`kda_l2_normalize_v0` cannot do. KDA + normalises ``q`` and ``k``, which are *key*-width, and Kimi K3 is + ``key_dim = 128`` against ``mlen = 64`` -- but the norm **contracts over + the key axis**, so the two halves of one ``q`` share a single norm. + Normalising each row on its own would divide each half by its own partial + norm: finite, plausible, wrong. + + The sums cannot simply be accumulated 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, even though ``V_RED_SUM`` itself accumulates + into its FP register. + + So each block is reduced into its own slot and the slots are folded with + scalar ``S_ADD_FP``. The partial sums are laid out **block-major**, + ``part_fp[c * vectors + v]``, precisely so the fold is ``blocks - 1`` + contiguous elementwise FPRAM adds of length ``vectors`` rather than + ``vectors`` separate reductions. + + ``part_fp`` needs ``vectors * blocks`` slots, ``acc_fp`` needs + ``vectors``. ``consts.reci_group`` must be 1.0 -- see + :meth:`kda_fp_constant_values`. + """ + if vectors < 1 or blocks < 1: + raise ValueError(f"vectors and blocks must be positive, got {vectors}, {blocks}") + if vec.shape[1] != self.mlen: + raise ValueError( + f"vec is {vec.shape[1]} columns wide; must be exactly mlen ({self.mlen})" + ) + if vec.shape[0] < first_row + vectors * blocks: + raise ValueError( + f"vec needs {first_row + vectors * blocks} rows for {vectors} vectors " + f"x {blocks} blocks starting at {first_row}, has {vec.shape[0]}" + ) + if sq_scratch.name == vec.name: + raise ValueError("sq_scratch must be distinct from vec") + if sq_scratch.shape[0] < first_row + vectors * blocks or sq_scratch.shape[1] != self.mlen: + raise ValueError("sq_scratch must match vec's shape") + if part_fp.size < vectors * blocks: + raise ValueError( + f"part_fp holds {part_fp.size} slots, needs {vectors * blocks}" + ) + if acc_fp.size < vectors: + raise ValueError(f"acc_fp holds {acc_fp.size} slots, needs {vectors}") + + self.emit_comment( + kda_stage_marker("kda_normalize", f"vectors={vectors} blocks={blocks}") + ) + # `first_row` lets the mixer normalise one head's slice of a tile that + # holds every head, rather than copying it out first. + rows = list(range(first_row, first_row + vectors * blocks)) + # Copy only the live rows. VRAM row counts must be a multiple of mlen + # (memory_state.py:309), so vec.shape[0] is the padded height -- at + # mlen=64 with two live rows that is 64 rows copied for 2 rows of work, + # about 380 wasted dynamic instructions. Only `rows` is reduced below, + # so the padding never contributes. + # Offset both sides so the copy is `vectors * blocks` rows wherever the + # slice sits. Copying from row 0 up to the last live row instead made + # the cost O(first_row), and the mixer walks first_row across every + # head -- at Kimi K3 that is 18,624 rows moved for 384 rows of live + # data. The rows stay at their original indices in sq_scratch so the + # multiply and the reductions below can share one row list. + self.mamba_block_copy( + sq_scratch, vec, + dst_row_offset=first_row, src_row_offset=first_row, + num_rows=vectors * blocks, + ) + self.tile_row_mul(sq_scratch, vec, rows=rows) + + # Block c of every vector reduces into part_fp[c * vectors ..]. Rows for + # fixed c step by `blocks` and slots step by 1 -- both progressions, so + # each block is one hardware loop. + for c in range(blocks): + self._kda_row_sum( + part_fp, + sq_scratch, + [first_row + v * blocks + c for v in range(vectors)], + c * vectors, + ) + + # Fold the per-block sums: acc[0..vectors) = sum_c part[c*vectors ..]. + # Explicit addresses: the FPVar-level helpers take whole vars, and these + # operate on a slice of part_fp. + self.fpvar_copy_asm(part_fp.address, acc_fp.address, vectors) + for c in range(1, blocks): + self.fpvar_add_asm( + acc_fp.address, part_fp.address + c * vectors, acc_fp.address, vectors + ) + + self.mamba_rsqrt_fpram(acc_fp, consts, count=vectors) + + # Every row of vector v scales by the single acc[v]. + for c in range(blocks): + self._kda_scale_rows( + vec, acc_fp, [first_row + v * blocks + c for v in range(vectors)], 0 + ) + return vec + + def kda_fp_constants(self, name_prefix: str = "kda"): + """Allocate the KDA layer's FPRAM scalar block. + + Reuses ``MambaFPConstants`` -- same seven slots, and a second allocator + would bump-walk FP_MEM twice. ``mamba_fp_constants`` takes no shape, so + nothing about Mamba's geometry comes along. + + ``dt_min`` / ``dt_max`` are unused by KDA (they clamp Mamba's ``dt``); + they stay allocated so slot order matches what ``mamba_rsqrt_fpram`` and + the other shared emitters index into. + """ + return self.mamba_fp_constants(name_prefix=name_prefix) + + @staticmethod + def kda_fp_constant_values(eps: float = 1.0e-6) -> list[float]: + """Host-side values matching :meth:`kda_fp_constants`, in slot order. + + Two deliberate differences from ``mamba_fp_constant_values``: + + ``reci_group = 1.0``. ``mamba_rsqrt_fpram`` computes + ``1 / sqrt(acc * reci_group + eps)``. Mamba wants ``sqrt(mean)`` so it + passes ``1 / group_width``; KDA's q/k normalisation is the L2 norm + ``sqrt(sum)``, so the factor must be 1. + + ``eps = 1e-6``, not ``1e-5``, matching FlashKDA's recurrent kernel and + the CPU reference (``reference.py``: ``rsqrt(sum + 1.0e-6)``). The + difference is small but it is a real convention, and + ``test_kda_reference.py`` pins it. + + ``dt_min`` / ``dt_max`` are zero and bf16-max: unused by KDA, but a + NaN or +inf parked in a live FPRAM slot would poison any emitter that + later reached for it. + """ + return [0.0, 1.0, -1.0, 0.0, 3.3895313892515355e38, 1.0, eps] diff --git a/aten/plena/program_kda_conv.py b/aten/plena/program_kda_conv.py new file mode 100644 index 0000000..e6bb43e --- /dev/null +++ b/aten/plena/program_kda_conv.py @@ -0,0 +1,192 @@ +"""KDA's three short causal convolutions, one decode step at a time. + +The reference (``aten/models/kda/reference.py::_causal_conv_step``) is:: + + updated = roll(state, -1, dim=-1) # drop the oldest tap + updated[..., -1] = value # append this token + output = sum_t updated[..., t] * weight[..., t] + output = silu(output + bias) + +It is **depthwise**: every channel carries its own ``kernel``-long tap vector. +So the tap weight varies per *lane*, which is why ``V_MUL_VF`` is the wrong tool +(it broadcasts one FP register across every lane) and the weight lives in VRAM +as a matrix for ``V_MUL_VV``. ``mamba_conv1d_v0`` makes the same argument at +``program_mamba_common.py:481-484``; this module is the same kernel with two +differences that make it its own emitter rather than a call into that one: + +* Mamba convolves a *sequence* held in ``x`` with the taps at row offsets. + KDA decode has one new token and a **carried** conv state, so the shift is a + real row copy every step rather than a free change of offset. +* KDA's channel counts are ``num_heads * key_dim`` and ``num_heads * value_dim`` + -- 12,288 apiece for Kimi K3 -- so everything is blocked by ``mlen``, and the + block index has to reach the inner loop. ``mamba_conv1d_v0`` addresses rows + from 0 and has nowhere to put it. + +Adding a base-row argument to ``mamba_conv1d_v0`` would have avoided the +duplication, but that emitter is on Mamba's decode and prefill paths and this +work has no business changing their signatures. The inner loop below is +deliberately the same shape so the two stay legible together. + +Layout, all tiles exactly ``mlen`` wide: + +=============== ===================================== ========================== +tile shape row index +=============== ===================================== ========================== +``conv_state`` ``[channel_blocks * kernel, mlen]`` ``kda_conv_state_row`` +``weight`` ``[channel_blocks * kernel, mlen]`` same +``bias`` ``[channel_blocks, mlen]`` channel block +``x_new`` ``[channel_blocks, mlen]`` channel block +``out`` ``[channel_blocks, mlen]`` channel block +=============== ===================================== ========================== + +Taps run oldest (0) to newest (``kernel - 1``), matching the reference's +``roll(-1)`` then write-at-``[-1]``. +""" + +from __future__ import annotations + +from compiler.aten.plena.program_kda_common import ( + kda_conv_blocks, + kda_conv_state_row, + kda_stage_marker, +) +from compiler.aten.plena.vars import VRAMMatrixVar + +__all__ = ["ProgramKdaConvMixin"] + + +class ProgramKdaConvMixin: + """One causal-conv decode step, blocked by channel. + + Requires ``ProgramMambaCommonMixin`` (``mamba_row_copy``, ``mamba_row_mul``, + ``mamba_row_add``, ``mamba_silu_v0``) and ``ProgramFPTileOpsMixin`` + (``vram_fill_zero``). + """ + + def kda_conv_step_v0( + self, + *, + x_new: VRAMMatrixVar, + conv_state: VRAMMatrixVar, + weight: VRAMMatrixVar, + bias: VRAMMatrixVar | None, + out: VRAMMatrixVar, + scratch: VRAMMatrixVar, + consts, + channels: int, + kernel: int, + apply_silu: bool = True, + x_new_row_base: int = 0, + x_new_row_stride: int = 1, + ) -> VRAMMatrixVar: + """Advance the history by one token and convolve. + + ``conv_state`` is updated in place: this is the carried state, and the + next token reads what this call leaves behind. + + ``x_new_row_base`` and ``x_new_row_stride`` say where block ``cb`` of + this token lives in ``x_new``: at row ``base + cb * stride``. The + default 0/1 is a dense tile, one feature block per row. + + A stride exists because the projection cannot produce a dense tile. + ``M_MM_WO`` writes a ``blen x blen`` sub-tile and the writeback loops + 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`` whatever the weights look like. The + convolution wants one token's blocks as consecutive rows. The mismatch + is exactly ``blen`` and it is a property of the matrix writeback, not + of how the weight matrices are packed, so no rearrangement of the + projection removes it. + + Reading across it costs nothing. ``mamba_row_copy`` takes a row index, + and the sweeps underneath collapse any arithmetic progression into one + hardware loop where the step is an ``S_ADDI_INT`` immediate -- a step + of ``blen`` is the same instruction as a step of 1. So passing the + projection tile with ``x_new_row_stride=blen`` does what a gather into + a dense tile does, for nothing. + """ + blocks = kda_conv_blocks(channels, self.mlen) + if kernel < 1: + raise ValueError(f"kernel must be at least 1, got {kernel}") + for name, tile in ( + ("x_new", x_new), ("conv_state", conv_state), ("weight", weight), + ("out", out), ("scratch", scratch), + ): + if tile.shape[1] != self.mlen: + raise ValueError( + f"{name} is {tile.shape[1]} columns wide; every KDA tile must be " + f"exactly mlen ({self.mlen})" + ) + needed = blocks * kernel + if conv_state.shape[0] < needed: + raise ValueError( + f"conv_state needs {needed} rows for {blocks} blocks x kernel {kernel}, " + f"has {conv_state.shape[0]}" + ) + if weight.shape[0] < needed: + raise ValueError(f"weight needs {needed} rows, has {weight.shape[0]}") + if x_new_row_stride < 1: + raise ValueError(f"x_new_row_stride must be at least 1, got {x_new_row_stride}") + if x_new_row_base < 0: + raise ValueError(f"x_new_row_base must not be negative, got {x_new_row_base}") + x_new_last = x_new_row_base + (blocks - 1) * x_new_row_stride + if out.shape[0] < blocks: + raise ValueError(f"out needs {blocks} rows, has {out.shape[0]}") + if x_new.shape[0] <= x_new_last: + raise ValueError( + f"x_new needs row {x_new_last} for block {blocks - 1} " + f"(base {x_new_row_base}, stride {x_new_row_stride}), " + f"has {x_new.shape[0]} rows" + ) + aliased = {conv_state.name, out.name, x_new.name, weight.name} + if bias is not None: + aliased.add(bias.name) + if scratch.name in aliased: + # mamba_row_copy is zero-then-add, so a scratch aliased onto any + # operand wipes it mid-loop. + raise ValueError(f"scratch must be distinct from {sorted(aliased)}") + + self.emit_comment( + kda_stage_marker("kda_conv1d", f"channels={channels} k={kernel} blocks={blocks}") + ) + for cb in range(blocks): + # Shift the history one tap older, then append this token. A real + # copy, not a ring pointer: address immediates are baked at ASM-gen + # time and there is no data-dependent addressing. + for tap in range(kernel - 1): + self.mamba_row_copy( + conv_state, + kda_conv_state_row(channels, self.mlen, kernel, cb, tap), + conv_state, + kda_conv_state_row(channels, self.mlen, kernel, cb, tap + 1), + ) + self.mamba_row_copy( + conv_state, + kda_conv_state_row(channels, self.mlen, kernel, cb, kernel - 1), + x_new, + x_new_row_base + cb * x_new_row_stride, + ) + + # out[cb] = sum_t state[cb, t] * weight[cb, t] + self.vram_fill_zero(out, rows=[cb]) + for tap in range(kernel): + row = kda_conv_state_row(channels, self.mlen, kernel, cb, tap) + self.mamba_row_copy(scratch, 0, conv_state, row) + self.mamba_row_mul(scratch, 0, weight, row) + self.mamba_row_add(out, cb, scratch, 0) + if bias is not None: + self.mamba_row_add(out, cb, bias, cb) + + if apply_silu: + # `marker` is required, and this is why: mamba_silu_v0 serves both + # stage vocabularies, and a Mamba default would bill KDA's SiLU -- + # and, because markers are sticky, everything emitted after it -- to + # a Mamba stage. + self.mamba_silu_v0( + out, + scratch, + consts, + rows=list(range(blocks)), + marker=kda_stage_marker("kda_conv1d", "silu"), + ) + return out diff --git a/aten/plena/program_kda_gates.py b/aten/plena/program_kda_gates.py new file mode 100644 index 0000000..42e5f2e --- /dev/null +++ b/aten/plena/program_kda_gates.py @@ -0,0 +1,248 @@ +"""KDA's decay and beta scalars, and their route into FPRAM. + +The recurrence consumes ``decay[h, k]``, ``beta[h]``, ``q_hat[h, k]`` and +``k_hat[h, k]`` as **FPRAM scalars**, because ``V_MUL_VF`` broadcasts one FP +register across every lane and each of those varies per row of the state tile. +They are produced on chip as VRAM rows, so each has to cross with +``S_MAP_FP_V`` -- one instruction per row, versus the one-hot +``V_MUL_VV`` + ``V_RED_SUM`` + ``S_ST_FP`` triple per *scalar* that was the only +route before it existed. + +From ``aten/models/kda/reference.py::activate_log_decay`` and ``kda_step``:: + + rate = exp(a_log[h]) # host constant + log_decay = gate_lower_bound * sigmoid(rate * (gate + dt_bias)) + decay = exp(log_decay) + beta = sigmoid(beta_logit) + +``rate`` is a function of a weight alone, so the host precomputes it and it +arrives as an FPRAM constant rather than costing a vector ``exp`` per token. +``gate_lower_bound`` likewise -- it lives on ``KdaShape`` and is a build +constant. + +Sigmoid is done in place with four vector ops and no scratch:: + + x <- -x V_MUL_VF by consts.neg_one + x <- exp(x) V_EXP_V + x <- x + 1 V_ADD_VF by consts.one + x <- 1/x V_RECI_V + +``mamba_silu_v0`` runs the same sequence but needs a scratch tile because it +keeps ``x`` alive for a final multiply. Nothing here does, so the scratch and +its copy are dropped. + +Layout. ``gate``, ``dt_bias`` and the decay output are **key**-width, blocked +one row per ``(head, key block)``; ``beta_logit`` is one value per head, blocked +by head. Both use their own row helpers rather than +:func:`kda_vector_row`, which blocks *value*-width tiles -- for Kimi K3 the two +happen to have the same block count, which is exactly the kind of coincidence +that hides a mix-up. + +The FPRAM landing addresses fall out of the row order. ``tile_row_to_fpram`` +puts row ``i`` at ``base + i * mlen``, so iterating ``(head, key block)`` in +order lands head ``h``'s block ``kb`` at ``h * key_dim + kb * mlen`` -- which is +exactly the ``key``-major-within-head layout the recurrence indexes. +""" + +from __future__ import annotations + +from typing import Sequence + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena.program_kda_common import kda_stage_marker +from compiler.aten.plena.vars import FPVar, VRAMMatrixVar + +__all__ = [ + "ProgramKdaGatesMixin", + "kda_head_blocks", + "kda_key_blocks", + "kda_key_row", +] + + +def kda_key_blocks(shape: KdaShape, mlen: int) -> int: + """Column blocks a key-width vector spans.""" + if mlen <= 0: + raise ValueError(f"mlen must be positive, got {mlen}") + if shape.key_dim % mlen: + raise ValueError( + f"key_dim ({shape.key_dim}) must be a multiple of mlen ({mlen}); a " + f"partial trailing block would carry lanes that hold nothing into " + f"the FPRAM scalars" + ) + return shape.key_dim // mlen + + +def kda_key_row(shape: KdaShape, mlen: int, head: int, block: int) -> int: + """Row holding block ``block`` of head ``head``'s key-width vector. + + ``head * key_blocks + block``: head outermost, so a single + ``tile_row_to_fpram`` over rows ``0..num_heads*key_blocks-1`` lands the + scalars key-major within each head, matching what the recurrence indexes. + """ + blocks = kda_key_blocks(shape, mlen) + if not 0 <= head < shape.num_heads: + raise ValueError(f"head {head} out of range for {shape.num_heads} heads") + if not 0 <= block < blocks: + raise ValueError(f"block {block} out of range for {blocks} key blocks") + return head * blocks + block + + +def kda_head_blocks(shape: KdaShape, mlen: int) -> int: + """Rows needed to hold one value per head. + + Unlike the key and value widths this is allowed a partial trailing block -- + Kimi K3 has 96 heads against ``mlen`` 64. The lanes past ``num_heads`` are + never indexed, because the recurrence reads ``beta_fp[h]`` for + ``h < num_heads``. + """ + if mlen <= 0: + raise ValueError(f"mlen must be positive, got {mlen}") + return (shape.num_heads + mlen - 1) // mlen + + +class ProgramKdaGatesMixin: + """Decay and beta, from projected activations to FPRAM scalars.""" + + def _kda_sigmoid_inplace(self, tile: VRAMMatrixVar, rows: list[int], consts) -> None: + """``tile[rows] <- 1 / (1 + exp(-tile[rows]))``, no scratch. + + ``V_EXP_V`` saturates its input to ``[-88, 88]`` rather than + overflowing, so the tails need no explicit clamp: a large positive ``x`` + gives ``exp(-x) -> 0`` and sigmoid ``-> 1``; a large negative gives + ``exp(-x) -> e^88`` and sigmoid ``-> 0``. Neither produces inf or NaN. + """ + self.tile_row_mul_fp_broadcast(tile, consts.neg_one, rows=rows) + self.tile_row_exp(tile, rows=rows) + self.tile_row_add_fp_broadcast(tile, consts.one, rows=rows) + self.tile_row_reci(tile, rows=rows) + + def kda_decay_scalars_v0( + self, + *, + gate: VRAMMatrixVar, + dt_bias: VRAMMatrixVar, + rate_fp: FPVar, + lower_bound_fp: FPVar, + decay_fp: FPVar, + consts, + shape: KdaShape, + heads: Sequence[int] | None = None, + ) -> FPVar: + """``decay[h,k] = exp(lower_bound * sigmoid(rate[h] * (gate + dt_bias)))``. + + ``gate`` is consumed in place and left holding the decay. ``rate_fp`` + carries ``exp(a_log[h])``, precomputed by the host because it depends + only on a weight. ``lower_bound_fp`` carries ``shape.gate_lower_bound``. + + ``heads`` restricts the work to a subset, which is how the mixer keeps + FPRAM bounded. Producing every head's scalars at once needs + ``num_heads * key_dim`` slots -- 12,288 for Kimi K3 against an FPRAM of + 512 -- so the layer streams one head at a time and reuses the window. + With a single head the scalars land at ``decay_fp[0 .. key_dim)``, + because ``tile_row_to_fpram`` addresses by position in ``rows`` rather + than by row index. + """ + blocks = kda_key_blocks(shape, self.mlen) + selected = list(range(shape.num_heads)) if heads is None else list(heads) + rows = [kda_key_row(shape, self.mlen, h, b) for h in selected for b in range(blocks)] + for name, tile in (("gate", gate), ("dt_bias", dt_bias)): + if tile.shape[1] != self.mlen: + raise ValueError( + f"{name} is {tile.shape[1]} columns wide; must be exactly " + f"mlen ({self.mlen})" + ) + # max(rows) + 1, not len(rows): with a head subset the rows are a + # slice out of the middle, so the count is smaller than the highest + # row number. Checking the count accepts a tile that the emitter + # then reads and writes past the end of -- silently, into whatever + # VRAM object follows. + needed = max(rows) + 1 + if tile.shape[0] < needed: + raise ValueError( + f"{name} needs {needed} rows to reach head {max(selected)}, " + f"has {tile.shape[0]}" + ) + if rate_fp.size < max(selected) + 1: + raise ValueError( + f"rate_fp holds {rate_fp.size} slots but head {max(selected)} is " + f"selected; it is indexed by head number, not by position" + ) + for h in selected: + if not 0 <= h < shape.num_heads: + raise ValueError(f"head {h} out of range for {shape.num_heads} heads") + # Sized from what S_MAP_FP_V writes, which here happens to equal + # num_heads * key_dim exactly because key_dim is a whole number of + # blocks. Stated in terms of the write so it stays correct if that + # changes. + written = len(rows) * self.mlen + if decay_fp.size < written: + raise ValueError( + f"decay_fp holds {decay_fp.size} slots but S_MAP_FP_V writes " + f"{written} ({len(rows)} rows x mlen {self.mlen})" + ) + + self.emit_comment(kda_stage_marker("kda_decay", f"heads={len(selected)}")) + self.tile_row_add(gate, dt_bias, rows=rows) + + # rate is per head, and rate_fp is indexed by head number -- so the + # offset is the head, not its position in `selected`. For a fixed key + # block a contiguous head range gives a row progression and a slot + # progression; a scattered one falls back to unrolled, still correct. + for kb in range(blocks): + for h in selected: + self._kda_scale_rows( + gate, rate_fp, [kda_key_row(shape, self.mlen, h, kb)], h + ) + + self._kda_sigmoid_inplace(gate, rows, consts) + self.tile_row_mul_fp_broadcast(gate, lower_bound_fp, rows=rows) + self.tile_row_exp(gate, rows=rows) + + # Row i -> decay_fp[i * mlen], and i = h*key_blocks + kb, so head h's + # block kb lands at h*key_dim + kb*mlen: key-major within head. + self.tile_row_to_fpram(gate, decay_fp, rows=rows) + return decay_fp + + def kda_beta_scalars_v0( + self, + *, + beta_logit: VRAMMatrixVar, + beta_fp: FPVar, + consts, + shape: KdaShape, + ) -> FPVar: + """``beta[h] = sigmoid(beta_logit[h])``, consumed in place. + + ``beta_logit`` holds one value per head, so its last row may be partly + empty -- Kimi K3 is 96 heads against ``mlen`` 64. Sigmoid of whatever + those lanes hold is harmless to *read*: they are never indexed. It is + not harmless to *write*, which is why ``beta_fp`` must be sized to the + padded row count; see the check below. + """ + if beta_logit.shape[1] != self.mlen: + raise ValueError( + f"beta_logit is {beta_logit.shape[1]} columns wide; must be exactly " + f"mlen ({self.mlen})" + ) + rows = list(range(kda_head_blocks(shape, self.mlen))) + if beta_logit.shape[0] < len(rows): + raise ValueError(f"beta_logit needs {len(rows)} rows") + # S_MAP_FP_V moves a whole mlen-wide row, so the write is rows*mlen + # slots, not num_heads. Checking num_heads lets the tail of the last row + # land in the next FPVar: at 96 heads and mlen 64 that is 32 slots of + # sigmoid(padding) written outside the allocation. `_resolve_fpram_addr` + # only bounds-checks the base offset, and the emulator's scalar SRAM only + # knows the whole file, so nothing downstream catches it either. + written = len(rows) * self.mlen + if beta_fp.size < written: + raise ValueError( + f"beta_fp holds {beta_fp.size} slots but S_MAP_FP_V writes " + f"{written} ({len(rows)} rows x mlen {self.mlen}); size it to the " + f"padded row count, not to num_heads" + ) + + self.emit_comment(kda_stage_marker("kda_decay", f"beta heads={shape.num_heads}")) + self._kda_sigmoid_inplace(beta_logit, rows, consts) + self.tile_row_to_fpram(beta_logit, beta_fp, rows=rows) + return beta_fp diff --git a/aten/plena/program_kda_layer.py b/aten/plena/program_kda_layer.py new file mode 100644 index 0000000..7c89aa9 --- /dev/null +++ b/aten/plena/program_kda_layer.py @@ -0,0 +1,402 @@ +"""A whole KDA layer: the input projection, the convolutions, the mixer, the +output projection. + +`program_kda_mixer` stops at the state-engine boundary, which is the right place +for it -- the projections are ordinary Matrix work. This file is what puts them +back, because a whole-model instruction count needs the whole layer. + +The reshape that is not a reshape +--------------------------------- +The input projection produces one tile:: + + projected = hidden @ W_in [tokens, 3*key_width + value_width + heads] + +and the mixer wants ``q``, ``k``, ``gate`` as ``[.., mlen]`` tiles per +``(head, key block)``. Those look like different layouts, and a naive conversion +is a scatter of one row per column block -- 700-odd of them per layer at Kimi +K3's shape, each a full block operation. + +They are the same bytes. A VRAM tile's column block ``c`` sits at +``base + c * physical_rows * mlen`` (``memory.py``'s ``col_block_base``), and +rows *within* a block are linear. So the projection tile already **is** a dense +``[blocks * physical_rows, mlen]`` tile: feature block ``c``'s token ``t`` is at +row ``c * stride + t``, with ``stride = physical_rows``. + +:meth:`vram_tall_view` names that whole tile and +:meth:`vram_column_block_view` names a single block; both register a VRAM object +at the computed address rather than allocating, so **neither emits an +instruction**. The stride is ``physical_rows * mlen`` and not ``mlen * mlen`` -- +those coincide only when a tile happens to be exactly ``mlen`` rows tall, and a +projection allocated ``strict=False`` is padded to ``blen`` rows instead. + +One gather, and the emitters are as they were +--------------------------------------------- +Everything on the decode path past the projection wants a **dense** tile whose +row is a feature block. :meth:`kda_gather_projection_v0` moves one section into +place with a single ``V_FMA_VF`` sweep over the tall view -- both the source rows +(step ``stride``) and the destination rows (step 1) are arithmetic progressions, +so it is one hardware loop. + +**Fourteen static instructions per section, independent of how many blocks it +spans.** All five sections of a Kimi K3 layer come to **70**, against 53,757 for +the three convolutions and 39,526 for the mixer. + +Those are *image* sizes, and they are the wrong number to price the gather with. +The loop body runs once per feature block -- 192 of them for q, k, v and gate -- +so the gather **issues 4,650 instructions against 492,681 for the layer, 0.94%**. +The static count is what the ``V_FMA_VF`` conversion bought: a program image that +does not grow with ``key_dim``. The dynamic count is what the gather costs, and +it is the only one of the two that may be compared against another way of +getting the sections into place. ``test_instruction_budget.py`` pins both. + +Why the gather exists, and how to not need it +--------------------------------------------- +Not because the sections share a tile. ``M_MM_WO`` writes a ``blen x blen`` +sub-tile and the writeback loops 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``, whatever the weights look +like. The consumers want one token's blocks as consecutive rows, and the +mismatch is exactly ``blen`` -- a property of the matrix writeback, not of the +packing. + +So splitting the packed projection into five, each writing straight to its +consumer, does not help, and measurement agrees: identical ``M_MM`` count, +0.12% more instructions for the extra setup, the same 4,650-instruction gather. +``test_separate_projections_do_not_help`` records the numbers. + +What does help is not moving the data at all. Every consumer takes explicit row +indices and ``_row_progression`` accepts any constant step, so reading the +projection at stride ``blen`` collapses into the same single hardware loop -- +the step is an ``S_ADDI_INT`` immediate, and ``blen`` costs what ``1`` costs. +``kda_conv_step_v0`` takes ``x_new_row_base`` and ``x_new_row_stride`` for +this; feeding it the tall view directly removes **4,041 of the 4,650** issued +instructions, the q/k/v share, for nothing. The gate and beta sections reach a +different consumer and still gather. + +An earlier version of this file claimed the layer could not be assembled without +first giving ``kda_conv_step_v0`` and ``kda_l2_normalize_blocked_v0`` a +column-block index. Both halves of that were wrong. The conv reads ``x_new`` at a +single row index rather than walking blocks as rows; and the normalisation never +sees the projection at all, because the mixer's ``q`` and ``k`` are the +convolutions' output, which the compiler allocates dense -- +``test_kda_mixer.py``'s conv/mixer seam test already pins that. Threading a +column-block index through the Var layer was measured at +25% on the +normalisation, and it would have lost a cross-block hardware loop, to solve a +problem that is not on the path. +""" + +from __future__ import annotations + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena.program_kda_common import kda_stage_marker +from compiler.aten.plena.program_kda_gates import kda_key_blocks +from compiler.aten.plena.vars import InputVar, VRAMMatrixVar + +__all__ = [ + "ProgramKdaLayerMixin", + "kda_projection_features", + "kda_projection_sections", + "kda_projection_width", +] + + +def kda_projection_features(shape: KdaShape) -> int: + """Logical width of the input projection, matching ``kda_state_engine_step``. + + ``3 * key_width + value_width + num_heads``: q, k and gate are key-width, + v is value-width, beta is one per head. + """ + key_width = shape.projection_size + return 3 * key_width + shape.num_heads * shape.value_dim + shape.num_heads + + +def kda_projection_sections(shape: KdaShape, mlen: int) -> list[tuple[str, int, int]]: + """``(name, first column block, block count)`` for each projected section. + + Sections are **block-aligned**, in the order ``kda_state_engine_step`` + splits them. Only ``beta`` is not naturally a whole block -- it is one value + per head, 96 at Kimi K3 against an ``mlen`` of 64 -- so it is padded to two, + the same padding ``kda_head_blocks`` already applies. The lanes past + ``num_heads`` are never indexed. + + Alignment is what makes the split free: a section that began mid-block could + not be named as a view, and would need a scatter instead. + """ + key_width = shape.projection_size + value_width = shape.num_heads * shape.value_dim + out, block = [], 0 + for name, width in ( + ("q", key_width), ("k", key_width), ("v", value_width), + ("gate", key_width), ("beta", shape.num_heads), + ): + count = max(1, -(-width // mlen)) + out.append((name, block, count)) + block += count + return out + + +def kda_projection_width(shape: KdaShape, mlen: int) -> int: + """Padded width of the input projection. + + Larger than :func:`kda_projection_features` by the padding on ``beta``. The + **host must materialise ``W_in`` at this width**, with the padding columns + zero -- the lowering reads whole blocks, so a narrower weight would leave + the tail of beta's block holding whatever preceded it in HBM. + """ + name, block, count = kda_projection_sections(shape, mlen)[-1] + return (block + count) * mlen + + +class ProgramKdaLayerMixin: + """Assembles a KDA layer around the mixer.""" + + def vram_column_block_view( + self, + tile: VRAMMatrixVar, + block: int, + *, + name: str, + rows: int | None = None, + ) -> VRAMMatrixVar: + """Name column block ``block`` of ``tile`` as its own ``[rows, mlen]`` tile. + + A view over the parent's bytes -- writing through one is visible through + the other -- and it **emits nothing**. + + The address comes from :meth:`_tile_addr`, the same function every other + consumer uses. Column blocks are strided by ``physical_rows * mlen`` + (``memory.py``'s ``col_block_base``), not by ``mlen * mlen``; those + coincide only when a tile happens to be exactly ``mlen`` rows tall, + which was the one shape the first version of this was tested at. A + projection allocated ``strict=False`` is padded to ``blen`` rows, not + ``mlen``, so a real decode projection was off by ``mlen / blen``. + """ + mlen = self.mlen + # The layout is authoritative, not the var: for a view they disagree, + # the var carrying the rows asked for and the layout the parent's. + layout = self.get_vram_layout(tile.name) + phys_rows, phys_cols = layout.physical_shape + # ceil, though `alloc` pads physical_cols to a multiple of mlen either + # way -- strict requires it, non-strict rounds up. Kept as ceil so a + # partial trailing block would stay addressable if that changes. + blocks = -(-phys_cols // mlen) + if not 0 <= block < blocks: + raise ValueError( + f"column block {block} out of range for {tile.name}, which has " + f"{blocks} ({phys_cols} columns at mlen {mlen})" + ) + rows = phys_rows if rows is None else rows + if rows > phys_rows: + raise ValueError(f"view of {rows} rows exceeds {tile.name}'s {phys_rows}") + return self._named_view( + name, rows, mlen, self._tile_addr(tile.name, 0, block), (phys_rows, mlen) + ) + + def vram_tall_view( + self, tile: VRAMMatrixVar, *, name: str + ) -> tuple[VRAMMatrixVar, int]: + """The whole of ``tile`` as one dense ``[blocks * physical_rows, mlen]`` tile. + + Returns ``(view, stride)`` where ``stride`` is ``physical_rows`` -- the + row distance between consecutive column blocks. + + Column block ``c`` sits at ``base + c * physical_rows * mlen`` and rows + *within* a block are linear, so the tile's bytes already are a dense + single-column-block tile. Feature block ``c``'s token ``t`` is at row + ``c * stride + t``. That is what lets a gather be one hardware loop + rather than a copy per block: both the source rows (step ``stride``) and + the destination rows (step 1) are arithmetic progressions. + + Emits nothing. + """ + mlen = self.mlen + layout = self.get_vram_layout(tile.name) + phys_rows, phys_cols = layout.physical_shape + blocks = -(-phys_cols // mlen) + return ( + self._named_view( + name, blocks * phys_rows, mlen, layout.vram_base_addr, + (blocks * phys_rows, mlen), + ), + phys_rows, + ) + + def _named_view( + self, name: str, rows: int, cols: int, addr: int, + physical_shape: tuple[int, int], + ) -> VRAMMatrixVar: + """`alloc_at` with a duplicate-name guard. + + ``register_vram_matrix`` overwrites without complaint and a + ``VRAMMatrixVar`` resolves its address by name at emit time, so a + duplicate silently repoints every view already handed out under that + name. With a default prefix and a layer that runs 93 times that is a + certainty rather than a risk. + """ + internal = self._scoped_name(name) + if internal in self.vram_matrices: + raise ValueError( + f"a VRAM object named {internal!r} already exists; give this view " + f"a distinct name (views resolve by name at emit time, so a " + f"duplicate repoints the earlier one)" + ) + return self.alloc_at(name, rows, cols, addr, physical_shape=physical_shape) + + def kda_gather_projection_v0( + self, + *, + projected: VRAMMatrixVar, + dst: VRAMMatrixVar, + section: str, + shape: KdaShape, + consts, + token: int = 0, + name: str = "kda_tall", + ) -> VRAMMatrixVar: + """Gather one section of ``projected`` into ``dst``'s first rows. + + The section's feature blocks live at rows ``(first + c) * stride + token`` + of the tall view; ``dst`` wants them at rows ``0 .. count-1``. Both are + arithmetic progressions, so ``_emit_tile_row_fma`` collapses the whole + gather into **one hardware loop** -- 14 static instructions per section, + independent of how many blocks it spans. All five sections of a Kimi K3 + layer come to about 70 in the image; expanded, the loop bodies issue + 4,650 instructions, 0.94% of the layer's 492,681. See the module + docstring on why only the second number prices the gather. + + The multiply is by ``1.0``: the FMA accumulates, so ``dst`` is zeroed + first and the sweep adds the source in. There is no VRAM-to-VRAM move + opcode, and this is cheaper than the copy idiom because the copy's + destination would be a different row each iteration through + ``mamba_row_copy``, which does not form a progression. + """ + sections = {n: (first, count) for n, first, count in + kda_projection_sections(shape, self.mlen)} + if section not in sections: + raise ValueError(f"unknown section {section!r}; have {sorted(sections)}") + first, count = sections[section] + if dst.shape[0] < count: + raise ValueError(f"dst needs {count} rows for {section}, has {dst.shape[0]}") + + tall, stride = self.vram_tall_view(projected, name=f"{name}_{section}") + if token >= stride: + raise ValueError( + f"token {token} is past the projection's {stride} physical rows" + ) + self.emit_comment( + kda_stage_marker("kda_qkv_proj", f"gather {section} blocks={count}") + ) + rows = list(range(count)) + self.vram_fill_zero(dst, rows=rows) + self.tile_row_fma_fp_broadcast( + dst, tall, consts.one, + dst_rows=rows, + src_rows=[(first + c) * stride + token for c in range(count)], + ) + return dst + + def kda_split_projection_v0( + self, + *, + projected: VRAMMatrixVar, + shape: KdaShape, + rows: int, + prefix: str = "kda", + ) -> dict[str, list[VRAMMatrixVar]]: + """Views onto ``projected`` for each of q, k, v, gate and beta. + + Returns lists of per-column-block views, in the order + ``kda_state_engine_step`` splits them: q, k, v, gate, beta. Every entry + is an alias into ``projected`` -- nothing is copied. + """ + mlen = self.mlen + # The invariant the split rests on is per *head*, not per section: + # feature `h*key_dim + b*mlen` must start on a block boundary. A section + # width can be a clean multiple while key_dim is not -- 4 heads of + # key_dim 6 gives key_width 24, divisible by mlen 8, with heads 1, 2 and + # 3 all starting mid-block. kda_key_blocks already raises exactly this. + kda_key_blocks(shape, mlen) + value_width = shape.num_heads * shape.value_dim + if value_width % mlen: + raise ValueError( + f"value_width ({value_width}) must be a multiple of mlen " + f"({mlen}); a partial block cannot be named as a view" + ) + out: dict[str, list[VRAMMatrixVar]] = {} + for name, first, count in kda_projection_sections(shape, mlen): + out[name] = [ + self.vram_column_block_view( + projected, first + i, name=f"{prefix}_{name}{i}", rows=rows + ) + for i in range(count) + ] + return out + + def kda_layer_from_projected_v0( + self, + *, + projected: VRAMMatrixVar, + gathered: dict[str, VRAMMatrixVar], + conv_state: dict[str, VRAMMatrixVar], + conv_weight: dict[str, VRAMMatrixVar], + conv_bias: dict[str, VRAMMatrixVar | None], + conv_scratch: VRAMMatrixVar, + mixer_buffers, + shape: KdaShape, + token: int = 0, + ) -> VRAMMatrixVar: + """One KDA layer for one decode token, from the projection to ``out``. + + Safe to call repeatedly in one program: the views it takes are named + after ``projected``, so stacking layers needs nothing from the caller + beyond giving each its own projection tile. + + Exactly ``kda_state_engine_step``'s boundary: takes ``projected``, + returns the mixer's output. The input and output projections are + ordinary Matrix work and stay with the caller -- the same line + ``program_kda_mixer`` draws. + + ``gathered`` supplies a dense destination per section; ``conv_state`` is + carried across tokens and updated in place. + + The order follows the reference: gather, convolve q, k and v, then the + gates, then the recurrence. All three convolutions must finish before + the mixer starts, because the mixer normalises ``q`` and ``k`` in place + and they are the convolutions' output. + """ + for section in ("q", "k", "v", "gate", "beta"): + if section not in gathered: + raise ValueError(f"gathered is missing {section!r}") + self.kda_gather_projection_v0( + projected=projected, dst=gathered[section], section=section, + shape=shape, consts=mixer_buffers.consts, token=token, + # From the projection's own name, which is unique per layer. + # A fixed prefix collides on the second layer -- the guard in + # `_named_view` catches it, but only after the caller has + # already written the layer. + name=f"{projected.display_name}_tall_t{token}", + ) + + widths = { + "q": shape.projection_size, + "k": shape.projection_size, + "v": shape.num_heads * shape.value_dim, + } + outs = {"q": mixer_buffers.q, "k": mixer_buffers.k, "v": mixer_buffers.v} + for section, channels in widths.items(): + self.kda_conv_step_v0( + x_new=gathered[section], conv_state=conv_state[section], + weight=conv_weight[section], bias=conv_bias.get(section), + out=outs[section], scratch=conv_scratch, + consts=mixer_buffers.consts, channels=channels, + kernel=shape.conv_kernel, + ) + + mixer_buffers.gate = gathered["gate"] + mixer_buffers.beta_logit = gathered["beta"] + self.kda_beta_scalars_v0( + beta_logit=mixer_buffers.beta_logit, beta_fp=mixer_buffers.beta_fp, + consts=mixer_buffers.consts, shape=shape, + ) + self.kda_mixer_step_v0(buffers=mixer_buffers, shape=shape) + return mixer_buffers.out diff --git a/aten/plena/program_kda_mixer.py b/aten/plena/program_kda_mixer.py new file mode 100644 index 0000000..e22a20f --- /dev/null +++ b/aten/plena/program_kda_mixer.py @@ -0,0 +1,238 @@ +"""The KDA mixer for one decode token: conv, gates, recurrence. + +This is the boundary `aten/models/kda/reference.py::kda_state_engine_step` +defines -- everything between the input projection and the output projection. +The projections themselves are ordinary Matrix work and are not part of it. + + q_raw, k_raw, v_raw, gate, beta = split(projected) + q, k, v = causal_conv(q_raw), causal_conv(k_raw), causal_conv(v_raw) + output = kda_step(q, k, v, gate, beta, state, a_log, dt_bias) + +Why it streams per head +----------------------- +FPRAM is 512 slots (``FP_SRAM_DEPTH``, ``PLENA_RTL/src/definitions/configuration.svh``). +The recurrence reads ``decay[h,k]``, ``q_hat[h,k]`` and ``k_hat[h,k]`` as FPRAM +scalars, because ``V_MUL_VF`` broadcasts one FP register across every lane while +each of those varies per row of the state tile. Materialising all of them costs +``3 * num_heads * key_dim`` slots -- 36,960 for Kimi K3, seventy-two times the +file. + +So the scalars are produced one head at a time into a window that is reused +before the next head. The window is ``2 * key_dim``, not ``3``: ``decay`` and +``q_hat`` are never live at the same time, because decay is dead once the +predict half has run and ``q_hat`` is not read until the read-out. That is why +the recurrence is called as its two halves here rather than through +``kda_decode_step_v0``, and why a caller may pass the same ``FPVar`` as both +``decay_fp`` and ``q_hat_fp``. At Kimi K3 the difference is 620 slots against +492 -- over the 512-slot file, versus inside it. The convolutions do **not** stream: they touch no FPRAM, they are +elementwise across channels, and running them once for all heads keeps their +per-channel-block hardware loops intact. + +Ordering constraint +------------------- +``q`` and ``k`` are normalised **in place**, and the normalisation of head ``h`` +must happen after the conv has written every head. That is why the three convs +run to completion before the per-head loop starts, rather than being folded into +it: a per-head conv would have to re-roll the shared conv history, and the +history is one tile covering all channels. +""" + +from __future__ import annotations + +from typing import Sequence + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena.program_kda_common import kda_blocks, kda_stage_marker +from compiler.aten.plena.program_kda_gates import ( + kda_head_blocks, + kda_key_blocks, + kda_key_row, +) +from compiler.aten.plena.vars import FPVar, VRAMMatrixVar + +__all__ = ["KdaMixerBuffers", "ProgramKdaMixerMixin"] + + +class KdaMixerBuffers: + """The tiles and FPRAM windows one mixer step needs. + + Grouped into an object because there are fourteen of them and a positional + mistake between two same-shaped tiles is a silent wrong answer. + """ + + def __init__( + self, + *, + q: VRAMMatrixVar, + k: VRAMMatrixVar, + v: VRAMMatrixVar, + gate: VRAMMatrixVar, + dt_bias: VRAMMatrixVar, + beta_logit: VRAMMatrixVar, + state: VRAMMatrixVar, + out: VRAMMatrixVar, + pred: VRAMMatrixVar, + err: VRAMMatrixVar, + sq_scratch: VRAMMatrixVar, + decay_fp: FPVar, + q_hat_fp: FPVar, + k_hat_fp: FPVar, + beta_fp: FPVar, + part_fp: FPVar, + acc_fp: FPVar, + output_scale_fp: FPVar, + rate_fp: FPVar, + lower_bound_fp: FPVar, + consts, + ): + self.q, self.k, self.v = q, k, v + self.gate, self.dt_bias = gate, dt_bias + self.beta_logit = beta_logit + self.state, self.out = state, out + self.pred, self.err = pred, err + self.sq_scratch = sq_scratch + self.decay_fp, self.q_hat_fp, self.k_hat_fp = decay_fp, q_hat_fp, k_hat_fp + self.beta_fp, self.part_fp, self.acc_fp = beta_fp, part_fp, acc_fp + self.output_scale_fp = output_scale_fp + self.rate_fp, self.lower_bound_fp = rate_fp, lower_bound_fp + self.consts = consts + + +class ProgramKdaMixerMixin: + """Assembles the conv, gate and recurrence emitters into one decode step.""" + + def kda_mixer_fpram_slots(self, shape: KdaShape) -> dict[str, int]: + """Every FPRAM slot one mixer step needs, itemised, plus ``"total"``. + + Reported so a caller can check it against ``FP_SRAM_DEPTH`` (512) before + lowering rather than after allocation fails -- note that the compiler's + own :class:`FPRAMAllocator` defaults to 1024 and so will *not* catch an + overflow. The two numbers disagree; the SystemVerilog is authoritative. + + Counting only the per-head window used to hide the overflow. ``beta_fp`` + is not one slot: ``S_MAP_FP_V`` moves whole rows, so it costs + ``ceil(num_heads / mlen) * mlen`` -- 128 for Kimi K3, not 1. ``rate_fp`` + is indexed by head number and so costs ``num_heads``. With those counted + honestly the old ``3 * key_dim`` window came to 620 slots at Kimi K3, + over the file by 108. Sharing one window between ``decay`` and ``q_hat`` + (see :meth:`kda_decode_predict_v0`) brings it to 492. + """ + key_blocks = kda_key_blocks(shape, self.mlen) + items = { + # decay and q_hat share one window; k_hat needs its own, because it + # is read by both halves of the recurrence. + "k_hat": shape.key_dim, + "decay_or_q_hat": shape.key_dim, + "beta": kda_head_blocks(shape, self.mlen) * self.mlen, + "rate": shape.num_heads, + "part": key_blocks, + "acc": 1, + "output_scale": 1, + "lower_bound": 1, + "consts": len(self.kda_fp_constant_values()), + } + items["total"] = sum(items.values()) + return items + + def kda_mixer_step_v0( + self, + *, + buffers: KdaMixerBuffers, + shape: KdaShape, + head_rows: Sequence[int] | None = None, + ) -> VRAMMatrixVar: + """One KDA decode token, from post-conv activations to ``out``. + + ``buffers.q`` / ``.k`` hold the *convolved* q and k, key-width and + blocked by :func:`kda_key_row`; ``.v`` is value-width and blocked by + :func:`kda_vector_row`. ``.gate`` is key-width and is consumed in place. + Call :meth:`kda_conv_step_v0` three times first to produce them -- + see the module docstring for why that is not folded in here. + + **Call :meth:`kda_beta_scalars_v0` once before this**, not per call. + It is not folded in because it consumes ``beta_logit`` in place and + sigmoid is not idempotent: a second application moves the value by ~0.43. + Folding it in made ``head_rows`` silently wrong -- lowering three heads + in three calls gave head 0 the right answer and the rest + ``sigmoid(sigmoid(...))``. + + In-place consumption, and what it means for a second token: + + * ``gate`` becomes the decay. **Not idempotent** (~0.90 off), so reseed it. + * ``beta_logit`` becomes beta. **Not idempotent** (~0.43 off), reseed it. + * ``q`` and ``k`` are normalised. Idempotent -- a unit vector normalises + to itself -- but they must be reseeded anyway because the next token's + projections differ. + * ``state`` is updated in place, and that is the point: it carries. + """ + heads = list(range(shape.num_heads)) if head_rows is None else list(head_rows) + key_blocks = kda_key_blocks(shape, self.mlen) + b = buffers + + for name, want in ( + ("q_hat_fp", shape.key_dim), + ("k_hat_fp", shape.key_dim), + ("decay_fp", shape.key_dim), + ): + fp = getattr(b, name) + if fp.size < want: + raise ValueError( + f"{name} holds {fp.size} slots but the per-head window needs " + f"{want}; it is reused every head, so size it to key_dim and " + f"not to num_heads * key_dim" + ) + if b.part_fp.size < key_blocks: + raise ValueError(f"part_fp needs {key_blocks} slots, has {b.part_fp.size}") + if b.acc_fp.size < 1: + raise ValueError("acc_fp needs at least one slot") + + for h in heads: + first = kda_key_row(shape, self.mlen, h, 0) + rows = [kda_key_row(shape, self.mlen, h, kb) for kb in range(key_blocks)] + + # k first: the predict half reads k_hat and decay, and q_hat may be + # the same storage as decay (see the module docstring), so q cannot + # be materialised until the predict half is done with decay. + self.emit_comment(kda_stage_marker("kda_normalize", f"k head={h}")) + self.kda_l2_normalize_blocked_v0( + b.k, vectors=1, blocks=key_blocks, sq_scratch=b.sq_scratch, + part_fp=b.part_fp, acc_fp=b.acc_fp, consts=b.consts, + first_row=first, + ) + # Row `first + i` lands at offset `i * mlen`, so the window is + # key-major and starts at 0 -- which is what both halves index with + # fp_base = 0 for a single head. + self.tile_row_to_fpram(b.k, b.k_hat_fp, rows=rows) + + self.kda_decay_scalars_v0( + gate=b.gate, dt_bias=b.dt_bias, + rate_fp=b.rate_fp, lower_bound_fp=b.lower_bound_fp, + decay_fp=b.decay_fp, consts=b.consts, shape=shape, heads=[h], + ) + self.kda_decode_predict_v0( + state=b.state, + k_fp=b.k_hat_fp, decay_fp=b.decay_fp, beta_fp=b.beta_fp, + v=b.v, pred=b.pred, err=b.err, shape=shape, head_rows=[h], + # The scalar window is reused every head, so this head's decay + # and k_hat sit at offset 0 -- not at h * key_dim. beta_fp is + # the exception: one slot per head, produced once before this + # call, and indexed by head number directly. + fp_head_stride=0, + ) + + # decay is dead now, so q_hat may overwrite it. + self.emit_comment(kda_stage_marker("kda_normalize", f"q head={h}")) + self.kda_l2_normalize_blocked_v0( + b.q, vectors=1, blocks=key_blocks, sq_scratch=b.sq_scratch, + part_fp=b.part_fp, acc_fp=b.acc_fp, consts=b.consts, + first_row=first, + ) + self.tile_row_to_fpram(b.q, b.q_hat_fp, rows=rows) + + self.kda_decode_update_v0( + state=b.state, + k_fp=b.k_hat_fp, q_fp=b.q_hat_fp, o=b.out, err=b.err, + shape=shape, output_scale_fp=b.output_scale_fp, + head_rows=[h], fp_head_stride=0, + ) + return b.out diff --git a/aten/plena/program_kda_prefill.py b/aten/plena/program_kda_prefill.py new file mode 100644 index 0000000..1ea43da --- /dev/null +++ b/aten/plena/program_kda_prefill.py @@ -0,0 +1,820 @@ +"""KDA chunked prefill: one chunk of ``C`` tokens against the incoming state. + +The algebra is derived and measured in ``program_kda_chunk``'s module docstring; +this file composes it. Per head, with ``S_0[value, key]`` the incoming state and +``A_t[key]`` the cumulative decay from :meth:`kda_chunk_decay_cumprod_v0`:: + + k_hat = k * A q_hat = q * A k_tilde = k / A + M = tril(k_hat @ k_tilde^T, -1) + N = tril(q_hat @ k_tilde^T) diagonal included + T = (I + tril(diag(beta) M, -1))^-1 diag(beta) + W = V - k_hat @ S_0^T + E = T @ W + out = scale * (q_hat @ S_0^T + N @ E) + S_C = A_C * S_0 + E^T @ (k * A_C / A) + +Seven matrix products. Every one 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, which is the arrangement ``ssd_chunk_head_v0`` already +uses for Mamba's ``B``: MRAM is writable only by ``H_PREFETCH_M``. + +``q`` and ``k`` must arrive L2-normalised +---------------------------------------- +This is a hard precondition, not a nicety. ``reference.py::kda_step`` normalises +both with ``rsqrt(sum + 1e-6)`` before the recurrence sees them, and +``kda_mixer_step_v0`` does the same. Feeding raw projections instead makes the +substitution's matrix ``L = I + tril(diag(beta) M, -1)`` catastrophically +ill-conditioned, and the whole chunk is then worthless: + + inputs cond(L) f32 error bf16 error + normalised, decay e^-5..1 1.0 - 1.1 2e-08 4e-04 + raw projections, |k| ~ 5 1e6 - 6e10 3e+21 2e+25 + +Chunk size +---------- +:meth:`kda_chunk_check_range` bounds it at 17, where ``1/A`` stops fitting bf16. +Within that, **the error does not depend on the chunk size**, measured against +the sequential recurrence at three seeds each, worst case: + + chunk cond(L) f32 bf16 out bf16 state + 4 1.0 9e-09 1.5e-04 2.5e-03 + 8 1.0 1.1e-08 5.0e-04 2.0e-03 + 16 1.0 - 1.1 2.2e-08 3.8e-04 3.8e-03 + +Nor does it **compound across chunks**. Eight chunks chained against an +equal-length sequential reference, worst state error per chunk: + + chunk 8: 1.7e-3 3.9e-3 1.4e-3 5.2e-3 1.7e-3 2.2e-3 5.2e-3 3.1e-3 + chunk 16: 2.3e-3 2.9e-3 2.4e-3 3.8e-3 3.2e-3 2.9e-3 3.9e-3 2.0e-3 + +Flat, because the decay makes the recurrence strongly contracting -- old error is +damped, not accumulated. An earlier version of this file claimed the opposite, +recommending chunk 8 on the strength of a conditioning table that ran from 3 to +759. That table was measured with **unnormalised** ``q``/``k`` and unnaturally +mild decays, so it described the top row of the table above and not this kernel. +Chunk 16 stands, and it is the size that makes the instruction count worth it. + +Only the read-out is masked +--------------------------- +``N`` is masked inclusively -- token ``t`` reads the state after its own update, +and letting it see ``s > t`` would be a causality violation worth 36% error. + +``M`` is **not masked at all**. The substitution reads only ``M[i, j]`` for +``j < i``, so a mask on it is a second defence against the same mistake -- and a +second defence that nothing distinguishes from the first is worse than none: +with both present, swapping ``M``'s mask for the inclusive one, dropping the +``-1`` from it, or extending the substitution's sweep to ``j <= i`` each left the +emulator's answer **bit-identical**. Only applying two of them together produced +an error, of seven orders of magnitude. The one defence that remains is pinned by +``test_ut_transform_ignores_m_on_and_above_the_diagonal``, which feeds a polluted +diagonal and demands an identical result. + +``M``'s upper triangle is left holding large values -- ``A_t/A_s`` for ``t < s`` +is a product of reciprocal decays -- but they are finite by construction, which +is what :meth:`kda_chunk_check_range` really guarantees, and they are never read. + +""" + +from __future__ import annotations + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena.program_kda_chunk import kda_round_up +from compiler.aten.plena.program_kda_common import kda_stage_marker +from compiler.aten.plena.vars import FPVar, InputVar, VRAMMatrixVar + +__all__ = [ + "KdaPrefillBuffers", + "ProgramKdaPrefillMixin", + "kda_prefill_state_transpose_shapes", + "kda_prefill_tile_shapes", +] + + +def kda_prefill_tile_shapes( + shape: KdaShape, mlen: int, chunk: int +) -> dict[str, tuple[int, int]]: + """Required ``(rows, cols)`` of every tile :class:`KdaPrefillBuffers` holds. + + One source of truth, consulted by the emitter's validation *and* by whoever + allocates. It exists because the two used to be written separately: the + emitter carried a hand-kept list of row counts and every caller allocated + ``mlen x mlen`` regardless, and the gap between them is exactly how a spill + came to prefetch seven kilobytes past its allocation. + + Both axes are block-rounded; :func:`kda_round_up` says why. The width is what + encodes each tile's role, so it is worth reading off: + + * ``key_dim`` wide -- the key axis is on lanes and gets contracted by a + projection: ``q``, ``k``, ``decay``, ``k_tilde``, ``k_end``, ``prev``, + and the state and its two scratch tiles. + * ``value_dim`` wide -- value is the output column axis: ``out`` and the + read-out's accumulator. + * ``chunk`` wide -- time is on lanes, which is what makes ``E^T @ k_end`` + a contraction the array can do without an explicit transpose: ``v_t``, + ``err_t``, and the four ``[chunk, chunk]`` tiles. + """ + r = lambda n: kda_round_up(n, mlen) # noqa: E731 + chunk_rows, chunk_cols = r(chunk), r(chunk) + key_cols = r(shape.key_dim) + value_rows, value_cols = r(shape.value_dim), r(shape.value_dim) + return { + "q": (chunk_rows, key_cols), + "k": (chunk_rows, key_cols), + "decay": (chunk_rows, key_cols), + "k_tilde": (chunk_rows, key_cols), + "k_end": (chunk_rows, key_cols), + "prev": (chunk_rows, key_cols), + "state": (value_rows, key_cols), + "scale_scratch": (value_rows, key_cols), + "state_contrib": (value_rows, key_cols), + "v_t": (value_rows, chunk_cols), + "err_t": (value_rows, chunk_cols), + "contrib": (value_rows, chunk_cols), + "out": (chunk_rows, value_cols), + "readout_contrib": (chunk_rows, value_cols), + "gram": (chunk_rows, chunk_cols), + "readout": (chunk_rows, chunk_cols), + "t_mat": (chunk_rows, chunk_cols), + "identity": (chunk_rows, chunk_cols), + "causal_inclusive": (chunk_rows, chunk_cols), + } + + +def kda_prefill_state_transpose_shapes( + shape: KdaShape, mlen: int +) -> dict[str, tuple[int, int]]: + """Required ``(rows, cols)`` for :meth:`kda_prefill_state_to_decode_layout_v0`. + + Separate from :func:`kda_prefill_tile_shapes` because the transpose runs once + per layer rather than once per chunk, and because its ``identity`` is a + different tile from the chunk loop's: that one is ``[chunk, chunk]``, this + one has to span the whole key axis. + """ + r = lambda n: kda_round_up(n, mlen) # noqa: E731 + return { + "state": (r(shape.value_dim), r(shape.key_dim)), + "identity": (r(shape.key_dim), r(shape.key_dim)), + "out": (r(shape.key_dim), r(shape.value_dim)), + } + + +class KdaPrefillBuffers: + """Tiles and FPRAM windows one chunk of prefill needs, for one head. + + Grouped for the same reason ``KdaMixerBuffers`` is: a positional mistake + between two same-shaped ``[chunk, chunk]`` tiles is a silent wrong answer. + :func:`kda_prefill_tile_shapes` gives every tile's required shape, and the + emitter checks each one against it. + + Three accumulators, not one + --------------------------- + ``contrib``, ``readout_contrib`` and ``state_contrib`` are all + projection targets that get combined into something else, and a single + scratch tile used to serve all three. It cannot any more: they are + ``[value_dim, chunk]``, ``[chunk, value_dim]`` and ``[value_dim, key_dim]``, + and once those widths differ ``vram_add`` refuses the mismatch -- which is + the good outcome. Sizing one tile to the maximum of the three would have + passed the shape check by coincidence whenever ``key_dim >= chunk``, and + silently added the wrong lanes when it did not. + + Layout, and why two of these are transposed + ------------------------------------------ + ``q``, ``k`` and ``decay`` are ``[chunk, key_dim]``: one row per timestep, + the key axis on **lanes** across ``kda_key_blocks`` column blocks. See + ``program_kda_chunk``'s module docstring for why the key axis is there and + not folded into rows -- it is what lets ``key_dim`` exceed ``mlen``. + ``v_t`` and ``err_t`` are ``[value_dim, chunk]``, **time on lanes**. + + That is not a preference. The final state needs ``E^T @ k_end``, a + contraction over time on *both* operands, and the systolic array contracts a + VRAM operand's lanes against an MRAM operand's rows. Carrying the error + transposed makes every one of the seven products fall onto + ``vram_sub_projection_to`` or its ``_T_to`` sibling with no explicit + transpose anywhere -- the same move ``ssd_state_update_v0`` makes when it + demands ``b_t_chunk`` with time on lanes. + """ + + def __init__( + self, + *, + q: VRAMMatrixVar, + k: VRAMMatrixVar, + v_t: VRAMMatrixVar, + decay: VRAMMatrixVar, + k_tilde: VRAMMatrixVar, + k_end: VRAMMatrixVar, + gram: VRAMMatrixVar, + readout: VRAMMatrixVar, + t_mat: VRAMMatrixVar, + identity: VRAMMatrixVar, + causal_inclusive: VRAMMatrixVar, + err_t: VRAMMatrixVar, + state: VRAMMatrixVar, + contrib: VRAMMatrixVar, + readout_contrib: VRAMMatrixVar, + state_contrib: VRAMMatrixVar, + out: VRAMMatrixVar, + prev: VRAMMatrixVar, + scale_scratch: VRAMMatrixVar, + beta_fp: FPVar, + m_fp: FPVar, + output_scale_fp: FPVar, + consts, + ): + self.q, self.k, self.v_t = q, k, v_t + self.decay, self.k_tilde, self.k_end = decay, k_tilde, k_end + self.gram, self.readout = gram, readout + self.t_mat, self.identity = t_mat, identity + self.causal_inclusive = causal_inclusive + self.err_t = err_t + self.state, self.contrib, self.out = state, contrib, out + self.readout_contrib, self.state_contrib = readout_contrib, state_contrib + self.prev, self.scale_scratch = prev, scale_scratch + self.beta_fp, self.m_fp = beta_fp, m_fp + self.output_scale_fp = output_scale_fp + self.consts = consts + + @staticmethod + def causal_mask_values(chunk: int, mlen: int) -> list[list[float]]: + """Host contents of ``causal_inclusive``: 1 iff ``s <= t``. + + Inclusive because token ``t`` reads the state *after* its own update. + There is deliberately no strict variant: the only other place one was + used was on the gram feeding the substitution, which reads only + ``j < i`` on its own. Staged rather than built on chip for the reason + ``ssd_lower_triangular_ones`` gives: materialising a constant tile costs + one ``S_ST_FP`` per element. + """ + return [ + [1.0 if (s < chunk and s <= t) else 0.0 for s in range(mlen)] + for t in range(chunk) + ] + + @staticmethod + def identity_values(chunk: int, mlen: int) -> list[list[float]]: + return [[1.0 if s == t else 0.0 for s in range(mlen)] for t in range(chunk)] + + @staticmethod + def state_transpose_identity_values(key_dim: int) -> list[list[float]]: + """Host contents of the identity :meth:`kda_prefill_state_to_decode_layout_v0` projects against. + + A different tile from :meth:`identity_values`, and bigger: the transpose + contracts over the whole key axis, so this one is ``key_dim x key_dim`` + while the UT transform's is ``chunk x chunk``. Handing over the smaller + one gives a finite, wrong answer whenever ``key_dim > chunk``, which is + every real shape -- hence the separate name. + """ + return [[1.0 if s == t else 0.0 for s in range(key_dim)] for t in range(key_dim)] + + +class ProgramKdaPrefillMixin: + """One chunk of KDA prefill, for one head. + + Requires ``ProgramKdaChunkMixin`` (the cumulative decay and the UT + transform), ``ProgramMatrixOpsMixin`` (the projections) and the tile-row + family. + """ + + def _kda_blocks(self, n: int) -> int: + """MLEN-wide tile blocks spanning ``n`` elements. + + Same arithmetic as ``ProgramSsdMixin._blocks``; duplicated rather than + inherited so the KDA path does not drag in the whole SSD mixin. + """ + return max(1, -(-n // self.mlen)) + + def kda_prefill_state_to_decode_layout_v0( + self, + *, + state: VRAMMatrixVar, + identity: VRAMMatrixVar, + out: VRAMMatrixVar, + shape: KdaShape, + precision: dict, + ) -> VRAMMatrixVar: + """Transpose the carried state from prefill's layout into decode's. + + **The two paths hold the state transposed relative to each other**, and + at Kimi K3 ``key_dim == value_dim == 128``, so the shapes match and + handing one to the other produces a finite, plausible, wrong answer + rather than an error. Both layouts are deliberate: + + * decode is ``[key, value]`` (:func:`kda_state_row`) so that each sweep + is one arithmetic row progression and becomes a hardware loop -- see + ``program_kda_common``'s module docstring. Mamba-2 stores state the + same way, which is why they share emitters. + * prefill is ``[value, key]`` because that is what makes all seven of its + products land on the projection primitives without an explicit + transpose. + + So the conversion belongs at the boundary, once per layer, not inside + either path. It is a projection against a staged identity: + ``out[i][j] = sum_k I[i][k] * state[j][k] = state[j][i]``. + + With more than one block the block indices swap as well as the elements, + and both loops here do exactly that. Writing the result block index pair + as ``(ib, jb)`` with ``ib`` over **key** blocks and ``jb`` over **value** + blocks:: + + out[ib][jb][r][c] = sum_k I[ib*mlen+r][k] * state[jb*mlen+c][k] + = state[jb*mlen+c][ib*mlen+r] + + so the element transpose and the block transpose fall out of the same + pair of loops, provided ``identity`` spans the whole key axis. It is a + separate, larger tile from the chunk loop's ``[chunk, chunk]`` identity: + :meth:`KdaPrefillBuffers.state_transpose_identity_values` builds it and + :func:`kda_prefill_state_transpose_shapes` gives all three shapes. + + The result lands directly in decode's layout with no further move. + Decode addresses its state at row ``block * key_dim + key`` (see + :func:`kda_state_row`), and a VRAM matrix is column-block-major with + column block ``jb`` at ``jb * physical_rows * mlen`` -- so with ``out`` + exactly ``key_dim`` rows tall the two coincide. That is a real + constraint, not an observation, and it is asserted below: pad ``out`` to + ``mlen`` rows when ``key_dim < mlen`` and the value blocks would land at + the wrong stride. + """ + want = kda_prefill_state_transpose_shapes(shape, self.mlen) + for name, tile in (("state", state), ("identity", identity), ("out", out)): + rows, cols = want[name] + if tile.shape[0] < rows or tile.shape[1] != cols: + raise ValueError( + f"{name} must be at least {rows} rows and exactly {cols} " + f"columns for key_dim {shape.key_dim} x value_dim " + f"{shape.value_dim} at mlen {self.mlen}; got {tile.shape}" + ) + if out.name == state.name: + raise ValueError("out must not alias state -- a projection overwrites it") + + key_blocks = self._kda_blocks(shape.key_dim) + val_blocks = self._kda_blocks(shape.value_dim) + if val_blocks > 1 and out.physical_shape[0] != shape.key_dim: + raise ValueError( + f"out is {out.physical_shape[0]} rows tall but decode indexes its " + f"state at block * key_dim + key, so with {val_blocks} value " + f"blocks it must be exactly key_dim ({shape.key_dim}) rows -- " + f"otherwise the value blocks land at the wrong stride" + ) + + self.emit_comment(kda_stage_marker("kda_state_store", "prefill -> decode layout")) + state_hbm = self.kda_prefill_spill_v0( + tile=state, name="kda_state_transpose_spill", + live_rows=shape.value_dim, precision=precision, + ) + for ib in range(key_blocks): + for jb in range(val_blocks): + self.vram_sub_projection_T_to( + identity, ib, state_hbm, jb, out, ib, jb, **precision + ) + return out + + def kda_prefill_broadcast_row_v0( + self, + *, + dst: VRAMMatrixVar, + src: VRAMMatrixVar, + src_row: int, + count: int, + ) -> VRAMMatrixVar: + """Copy one row of ``src`` into the first ``count`` rows of ``dst``, by doubling. + + The row ops take one row list for both operands, so broadcasting a + single row across a tile means a copy per destination row -- 64 of them + at ``mlen`` 64, and 128 at Kimi K3's ``value_dim``. Copying the + already-copied region instead doubles the filled span each time, so it + is ``ceil(log2(count)) + 1`` copies: 7 instead of 64. + + Every copy covers all of the tile's column blocks, so the whole key axis + is broadcast at once: ``mamba_block_copy`` goes through + ``vram_fill_zero`` and ``vram_add``, and both walk the blocks + themselves. ``dst`` and ``src`` must therefore be the same width -- + ``vram_add`` asserts it -- which for both callers here means both are + ``key_dim`` wide. + """ + if count < 1: + raise ValueError(f"count must be positive, got {count}") + if dst.shape[1] != src.shape[1]: + raise ValueError( + f"broadcast needs matching widths: dst {dst.shape[1]}, src " + f"{src.shape[1]}. Every copy spans all column blocks" + ) + self.mamba_block_copy( + dst, src, dst_row_offset=0, src_row_offset=src_row, num_rows=1 + ) + filled = 1 + while filled < count: + take = min(filled, count - filled) + self.mamba_block_copy( + dst, dst, dst_row_offset=filled, src_row_offset=0, num_rows=take + ) + filled += take + return dst + + def kda_prefill_spill_v0( + self, + *, + tile: VRAMMatrixVar, + name: str, + live_rows: int, + precision: dict, + ) -> InputVar: + """Spill ``tile`` to HBM for use as an MRAM operand, safely and reusably. + + Two things this does that a bare :meth:`store` does not. + + **It zero-fills the rows past ``live_rows``.** ``store`` sizes the HBM + region from the tile's real height, but ``load_sub_matrix_col`` / + ``load_sub_matrix_row`` prefetch a whole ``mlen x mlen`` block + unconditionally -- ``k_block_count`` selects whole blocks and cannot trim + a partial one. At ``chunk`` 8 and ``mlen`` 64 that is 1,024 bytes written + and 8,192 read: the prefetch runs 7,168 bytes past the allocation, which + at 1,024-byte spacing is straight through the next seven spills. + + For the four ``_T_to`` products the over-read rows land in output columns + past ``chunk`` and are never used. But + :meth:`kda_prefill_state_tail_v0` uses the plain ``_to`` form, which + contracts over the MRAM operand's **rows** -- so garbage there is summed + directly into the carried state. It has been correct only because + ``k_end`` happens to be allocated last, with nothing after it. Adding a + spill, or reordering them, would have broken it silently. Requiring the + tile to be ``mlen`` rows tall and zeroing the tail makes the prefetch + read exactly what was written. + + **It reuses one HBM region per name.** Each call used to allocate fresh: + six regions per chunk per head, never reclaimed, about 20 KB per chunk -- + 10 MB per head over a 4k-token prefill. Reuse would have been unsafe + before, because the over-read would then have found the *previous* + chunk's real data instead of unallocated zeros. The two defects had to be + fixed together, and this is where. + """ + if tile.shape[0] < self.mlen: + raise ValueError( + f"{name} spills into an mlen x mlen MRAM block, so its tile must " + f"be {self.mlen} rows tall; got {tile.shape[0]}. The rows past " + f"live_rows are zeroed here so the prefetch reads what was written" + ) + if live_rows < tile.shape[0]: + # To the tile's own height, not to `mlen`. A tile taller than one + # block -- the state at `value_dim` 128 against `mlen` 64 -- has dead + # rows past `mlen` too, and bounding the fill at `mlen` made the + # range empty and zeroed nothing at all. + self.vram_fill_zero(tile, rows=list(range(live_rows, tile.shape[0]))) + regions = getattr(self, "_kda_spill_regions", None) + if regions is None: + regions = self._kda_spill_regions = {} + addr = regions.get(name) + var = self.store( + tile, name=name, hbm_addr=addr, precision=1, + hbm_element_bytes=precision["hbm_element_bytes"], + ) + regions[name] = self.get_hbm_layout(var.name).hbm_base_addr + return var + + def kda_prefill_reciprocal_decay_v0( + self, + *, + decay: VRAMMatrixVar, + k: VRAMMatrixVar, + k_tilde: VRAMMatrixVar, + chunk: int, + shape: KdaShape, + ) -> VRAMMatrixVar: + """``k_tilde = k / A``, the one unbounded quantity in the chunk. + + ``1/A`` reaches ``exp(chunk * |gate_lower_bound|)``, which is why + :meth:`kda_chunk_check_range` refuses a chunk past 17 at Kimi K3's -5. + It is formed here and consumed immediately by the two grams; nothing + downstream keeps it. + + Note this is a *reciprocal of the cumulative decay*, not a separate + reverse scan: ``V_RECI_V`` on a value in ``[e^-80, 1]`` is exact to + bf16's relative precision, whereas a reverse product would accumulate a + second set of roundings. + """ + self.kda_chunk_check_range(chunk, shape) + rows = list(range(chunk)) + self.emit_comment(kda_stage_marker("kda_decay", f"reciprocal chunk={chunk}")) + # One copy for every column block; the two row ops once per block. + self.mamba_block_copy(k_tilde, decay, num_rows=chunk) + for kb in range(self._kda_blocks(shape.key_dim)): + self.tile_row_reci(k_tilde, rows=rows, tile_col_idx=kb) + self.tile_row_mul(k_tilde, k, rows=rows, tile_col_idx=kb) + return k_tilde + + def kda_prefill_gram_v0( + self, + *, + left: VRAMMatrixVar, + right_hbm: InputVar, + out: VRAMMatrixVar, + chunk: int, + shape: KdaShape, + precision: dict, + mask: VRAMMatrixVar | None = None, + ) -> VRAMMatrixVar: + """``out = mask * (left @ right_hbm^T)``, a ``[chunk, chunk]`` gram. + + ``left`` is ``k_hat`` or ``q_hat`` in VRAM; ``right_hbm`` is ``k_tilde`` + spilled. When a mask is given it is applied after the matmul rather than + folded into the operands, because the systolic array has no triangular + mode -- the same arrangement the attention path uses for its causal mask. + + ``mask=None`` is for the gram feeding the substitution, which reads only + ``M[i, j]`` for ``j < i``. Masking it too was a second defence that made + the first one untestable; see the module docstring. + """ + i_blocks = self._kda_blocks(chunk) + j_blocks = self._kda_blocks(chunk) + self.emit_comment(kda_stage_marker("kda_state_update", f"gram chunk={chunk}")) + for ib in range(i_blocks): + for jb in range(j_blocks): + self.vram_sub_projection_T_to( + left, ib, right_hbm, jb, out, ib, jb, **precision + ) + if mask is not None: + self.tile_row_mul(out, mask, rows=list(range(chunk))) + return out + + def kda_prefill_chunk_v0( + self, + *, + buffers: KdaPrefillBuffers, + chunk: int, + shape: KdaShape, + precision: dict | None = None, + ) -> tuple[VRAMMatrixVar, VRAMMatrixVar]: + """One chunk of prefill for one head: ``(out, state)``. + + ``buffers.state`` is read as ``S_0`` and left holding ``S_C``, so chunks + chain by calling this again. ``q``, ``k`` and ``decay`` are consumed in + place: ``decay`` becomes the cumulative decay and then ``A_C/A``, and + ``q``/``k`` become ``q_hat``/``k_hat``. Reseed all three per chunk. + + ``precision`` defaults to the spilled-activation class, which must be + **Plain BF16**, and this method checks that itself against the active + TOML. Under the shipped MX-FP8 KV types the spills decode as e4m3 and the + read walks into the scale stream, whose ``0x7f`` bytes are e4m3 NaN -- + every output ``nan``. Even where it does not NaN, the state is a + multiplicative accumulator carried across every chunk, so three mantissa + bits at each boundary compounds without bound. + """ + from compiler.aten.plena.program_ssd import SPILLED_ACTIVATION + + precision = SPILLED_ACTIVATION if precision is None else precision + if precision is SPILLED_ACTIVATION: + # The emitter picks the keyvalue class, so the emitter checks it. + # Leaving this to the caller is what left it unchecked everywhere. + self.require_bf16_kv_precision_from_active_build() + self.kda_chunk_check_range(chunk, shape) + b = buffers + key_blocks = self._kda_blocks(shape.key_dim) + val_blocks = self._kda_blocks(shape.value_dim) + t_blocks = self._kda_blocks(chunk) + chunk_rows = list(range(chunk)) + + if chunk > self.mlen: + raise ValueError( + f"chunk {chunk} exceeds mlen {self.mlen}; the [chunk, chunk] tiles " + f"are one row per timestep and must fit a single block" + ) + # Every tile against the one table. Rows must be at least what the table + # asks (a caller is free to over-allocate); the width must match exactly, + # because a projection derives the length of its contraction from the + # operand's column-block count -- too narrow contracts over part of the + # axis, too wide contracts against lanes that hold nothing, and both are + # finite wrong answers rather than errors. + want = kda_prefill_tile_shapes(shape, self.mlen, chunk) + for name in want: + tile = getattr(b, name) + rows, cols = want[name] + if tile.shape[0] < rows: + raise ValueError( + f"{name} needs {rows} rows, has {tile.shape[0]}" + ) + if tile.shape[1] != cols: + raise ValueError( + f"{name} is {tile.shape[1]} columns wide; it must be exactly " + f"{cols} at key_dim {shape.key_dim}, value_dim " + f"{shape.value_dim}, chunk {chunk}, mlen {self.mlen}" + ) + for left, right in ( + ("prev", "decay"), ("out", "readout_contrib"), + ("state", "state_contrib"), ("state", "scale_scratch"), + ("v_t", "contrib"), ("k_tilde", "k_end"), + ): + if getattr(b, left).name == getattr(b, right).name: + raise ValueError( + f"{left} and {right} must be distinct tiles; one is read " + f"while the other is written" + ) + + self.emit_comment( + kda_stage_marker("kda_state_update", f"prefill chunk={chunk}") + ) + + # -- 1. A_t = prod_{s<=t} a_s, per key channel ------------------------ + self.kda_chunk_decay_cumprod_v0( + decay=b.decay, prev=b.prev, chunk=chunk, shape=shape + ) + + # -- 2. k_tilde = k / A, the one unbounded quantity ------------------- + self.kda_prefill_reciprocal_decay_v0( + decay=b.decay, k=b.k, k_tilde=b.k_tilde, chunk=chunk, shape=shape + ) + k_tilde_hbm = self.kda_prefill_spill_v0( + tile=b.k_tilde, name="kda_k_tilde_spill", + live_rows=chunk, precision=precision, + ) + + # -- 3. k_end = k_tilde * A_C, while k is still the raw key ---------- + # k_end[t] = k[t] * A_C / A_t, and k_tilde already carries the 1/A_t. + # A_C <= A_t for every t in the chunk (decays are <= 1), so the product + # is bounded by |k| even though k_tilde on its own is not. + self.kda_prefill_final_decay_ratio_v0( + k_tilde=b.k_tilde, decay=b.decay, k_end=b.k_end, + chunk=chunk, shape=shape, + ) + + # -- 4. k_hat and q_hat, in place ------------------------------------- + # Once per key block: a binary row op pairs one column block of the + # destination with the same block of the source. + for kb in range(key_blocks): + self.tile_row_mul(b.k, b.decay, rows=chunk_rows, tile_col_idx=kb) + self.tile_row_mul(b.q, b.decay, rows=chunk_rows, tile_col_idx=kb) + + # -- 5. the two grams, masked --------------------------------------- + # No mask on this one. See the module docstring: the substitution reads + # only j < i, so masking M was a redundant second defence -- and it made + # the first one impossible to test, because breaking either alone left + # the answer bit-identical. + self.kda_prefill_gram_v0( + left=b.k, right_hbm=k_tilde_hbm, + out=b.gram, chunk=chunk, shape=shape, precision=precision, + ) + # This one *is* load-bearing: N[t, s] for s > t would let token t read a + # future token's error. Swapping it for the strict mask is a 36% error. + self.kda_prefill_gram_v0( + left=b.q, right_hbm=k_tilde_hbm, mask=b.causal_inclusive, + out=b.readout, chunk=chunk, shape=shape, precision=precision, + ) + + # -- 6. T = (I + tril(diag(beta) M, -1))^-1 diag(beta) ---------------- + self.kda_ut_transform_v0( + m=b.gram, identity=b.identity, beta_fp=b.beta_fp, t_out=b.t_mat, + m_fp=b.m_fp, consts=b.consts, chunk=chunk, shape=shape, + ) + + # -- 7. W^T = v^T - S_0 @ k_hat^T ------------------------------------- + # Carried transposed: (S_0 @ k_hat^T)[v, t] = sum_key S_0[v,key] k_hat[t,key], + # which is `_T_to` with the state natural in VRAM and k_hat in MRAM. + k_hat_hbm = self.kda_prefill_spill_v0( + tile=b.k, name="kda_k_hat_spill", + live_rows=chunk, precision=precision, + ) + for vb in range(val_blocks): + for tb in range(t_blocks): + self.vram_sub_projection_T_to( + b.state, vb, k_hat_hbm, tb, b.contrib, vb, tb, **precision + ) + self.tile_row_sub(b.v_t, b.contrib, rows=list(range(shape.value_dim))) + + # -- 8. E^T = W^T @ T^T ----------------------------------------------- + # (W^T T^T)[v, t] = sum_s W^T[v,s] T[t,s] -- again `_T_to`, with T in MRAM. + t_hbm = self.kda_prefill_spill_v0( + tile=b.t_mat, name="kda_t_spill", + live_rows=chunk, precision=precision, + ) + for vb in range(val_blocks): + for tb in range(t_blocks): + self.vram_sub_projection_T_to( + b.v_t, vb, t_hbm, tb, b.err_t, vb, tb, **precision + ) + + # -- 9. out = scale * (q_hat @ S_0^T + N @ E) ------------------------- + state_hbm = self.kda_prefill_spill_v0( + tile=b.state, name="kda_state_spill", + live_rows=shape.value_dim, precision=precision, + ) + for tb in range(t_blocks): + for vb in range(val_blocks): + self.vram_sub_projection_T_to( + b.q, tb, state_hbm, vb, b.out, tb, vb, **precision + ) + err_hbm = self.kda_prefill_spill_v0( + tile=b.err_t, name="kda_err_spill", + live_rows=shape.value_dim, precision=precision, + ) + # (N @ E)[t, v] = sum_s N[t,s] E^T[v,s] -- it lands in a scratch tile + # because a projection overwrites its target (M_MM accumulates in the + # array, M_MM_WO writes and clears), so it cannot add into `out`. + # `readout_contrib` and not `contrib`: this one is [chunk, value_dim], + # and step 7's is [value_dim, chunk]. + for tb in range(t_blocks): + for vb in range(val_blocks): + self.vram_sub_projection_T_to( + b.readout, tb, err_hbm, vb, b.readout_contrib, tb, vb, **precision + ) + self.vram_add(b.out, b.readout_contrib, num_rows=chunk) + self.emit_comment(kda_stage_marker("kda_readout", f"prefill chunk={chunk}")) + # Once per value block. `out` is `[chunk, value_dim]`, and a row op + # reaches one column block per call -- so at value_dim 128 against mlen + # 64 a single call left the upper half of every token unscaled, i.e. + # `sqrt(key_dim)` = 11.3x too large. It read as a rounding problem in the + # aggregate (94.7% of values still inside a 5e-2 tolerance on data of + # order 1e-3) and only separating the blocks showed it. + for vb in range(val_blocks): + self.tile_row_mul_fp_broadcast( + b.out, b.output_scale_fp, rows=chunk_rows, tile_col_idx=vb + ) + + # -- 10. S_C = A_C * S_0 + E^T @ (k * A_C / A) ------------------------ + # k_end is bounded by |k|: A_C/A_s is a product of decays over (s, C], + # all <= 1. Forming k * A_C / A_s rather than (k/A_s) * A_C keeps it that + # way -- the second grouping goes through 1/A_s, which is what bounds the + # chunk size in the first place. + self.kda_prefill_state_tail_v0( + k_end=b.k_end, decay=b.decay, err_t=b.err_t, state=b.state, + state_contrib=b.state_contrib, scale_scratch=b.scale_scratch, + chunk=chunk, shape=shape, precision=precision, + ) + return b.out, b.state + + def kda_prefill_final_decay_ratio_v0( + self, + *, + k_tilde: VRAMMatrixVar, + decay: VRAMMatrixVar, + k_end: VRAMMatrixVar, + chunk: int, + shape: KdaShape, + ) -> VRAMMatrixVar: + """``k_end[t] = k[t] * A_C / A_t``, from ``k_tilde = k / A``. + + The final state contracts ``E^T`` against this. Written as + ``k_tilde * A_C`` rather than ``(k * A_C) / A`` because the first form + reuses a tile that already exists, and both are the same number: the + ratio ``A_C / A_t`` is a product of decays over ``(t, C]``, every one at + most 1, so ``k_end`` is bounded by ``|k|`` however small ``A_t`` gets. + """ + blocks = self._kda_blocks(shape.key_dim) + rows = list(range(chunk)) + self.emit_comment( + kda_stage_marker("kda_state_store", f"final decay ratio chunk={chunk}") + ) + # A_C is the last timestep's row, broadcast down every timestep. One row + # per timestep now, spanning every key block, so no destination pitch. + self.kda_prefill_broadcast_row_v0( + dst=k_end, src=decay, src_row=chunk - 1, count=chunk + ) + for kb in range(blocks): + self.tile_row_mul(k_end, k_tilde, rows=rows, tile_col_idx=kb) + return k_end + + def kda_prefill_state_tail_v0( + self, + *, + k_end: VRAMMatrixVar, + decay: VRAMMatrixVar, + err_t: VRAMMatrixVar, + state: VRAMMatrixVar, + state_contrib: VRAMMatrixVar, + scale_scratch: VRAMMatrixVar, + chunk: int, + shape: KdaShape, + precision: dict, + ) -> VRAMMatrixVar: + """``state = A_C * state + E^T @ k_end``, closing the chunk. + + ``A_C`` is per key channel, and the state's rows are *values*, so the + decay is applied along the lane axis of every row -- one + :meth:`tile_row_mul` against the broadcast row, not a per-row scalar. + That only works because the state is held ``[value, key]``; under the + reference's ``[value, key]`` it is the same thing, and under a + ``[key, value]`` layout it would be a per-row scalar instead. + """ + key_blocks = self._kda_blocks(shape.key_dim) + val_blocks = self._kda_blocks(shape.value_dim) + self.emit_comment(kda_stage_marker("kda_state_store", f"prefill chunk={chunk}")) + + k_end_hbm = self.kda_prefill_spill_v0( + tile=k_end, name="kda_k_end_spill", + live_rows=chunk, precision=precision, + ) + for vb in range(val_blocks): + for kb in range(key_blocks): + self.vram_sub_projection_to( + err_t, vb, k_end_hbm, kb, state_contrib, vb, kb, **precision + ) + + # state *= A_C. The row ops take one row list for both operands, so + # A_C -- a single row -- is broadcast into a tile shaped like the state + # first. Same staging the cumulative-decay scan needs, for the same + # reason. + self.kda_prefill_broadcast_row_v0( + dst=scale_scratch, src=decay, src_row=chunk - 1, count=shape.value_dim + ) + value_rows = list(range(shape.value_dim)) + for kb in range(key_blocks): + self.tile_row_mul(state, scale_scratch, rows=value_rows, tile_col_idx=kb) + self.vram_add(state, state_contrib, num_rows=shape.value_dim) + return state diff --git a/aten/plena/program_kda_recurrent.py b/aten/plena/program_kda_recurrent.py new file mode 100644 index 0000000..4536daa --- /dev/null +++ b/aten/plena/program_kda_recurrent.py @@ -0,0 +1,290 @@ +"""The KDA recurrent decode step: two sweeps over the transposed state tile. + +The maths, per head, with state held transposed as ``T[key, value]`` (see +``program_kda_common``):: + + T[k, :] = decay[k] * T[k, :] for each k + pred[:] = sum_k k_hat[k] * T[k, :] + err[:] = beta * (v[:] - pred[:]) + T[k, :] += err[:] * k_hat[k] for each k + out[:] = output_scale * sum_k q_hat[k] * T[k, :] + +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 does +``updated = state + err*k; state = updated; reduced += updated*q`` inside one +loop, so update and read-out cannot be reordered into separate passes over the +state. Updating every key row and then reducing over all of them *is* legal, +because each key's update touches only that key's row; that equivalence is what +lets both be single row sweeps. + +State is read twice and written once. That is inherent to the delta rule, not an +artefact of this lowering. + +Fused form +---------- +Each sweep is one ``V_FMA_VF`` inside one hardware loop. This module was first +written with only instructions that already existed -- the ``copy + multiply + +add`` triple, via ``mamba_row_copy`` / ``tile_row_mul_fp`` / ``mamba_row_add`` +-- to establish that the static ISA expresses KDA with no new opcode. It does, +and that is now demonstrated rather than claimed, so the triple has been +retired: it staged every row through a scratch row, and because that was the +*same* row each iteration the destination never formed a progression, which +forced the whole sweep to unroll. Static cost went from 7,345 instructions at +``key_dim`` 128 to 76, and stopped depending on ``key_dim`` at all. + +Column blocks +------------- +``value_dim`` may exceed ``mlen``; the block index is folded into the row +(``program_kda_common``'s module docstring explains why). Every tile here is +therefore exactly one column block wide, and the recurrence runs once per +block. The two sweeps are per ``(head, block)``, and within one the keys sit at +consecutive rows -- so each stays a single hardware loop. + +Blocks are independent for the decay, update and read-out, because those are +elementwise along ``value``. The prediction is not: it contracts over ``key`` +for a fixed ``value`` lane, so block ``c`` of ``pred`` only ever sees block +``c`` of the state. That is why the whole recurrence nests block outside key +rather than the other way round. +""" + +from __future__ import annotations + +from typing import Sequence + +from compiler.aten.models.kda.shape import KdaShape +from compiler.aten.plena.program_kda_common import ( + kda_blocks, + kda_stage_marker, + kda_state_row, + kda_vector_row, +) +from compiler.aten.plena.vars import FPVar, VRAMMatrixVar + +__all__ = ["ProgramKdaRecurrentMixin"] + + +class ProgramKdaRecurrentMixin: + """The KDA decode recurrence. + + Requires ``ProgramKdaCommonMixin`` (stage markers, layout helpers), + ``ProgramMambaCommonMixin`` (``mamba_row_copy``, ``mamba_row_add``) and + ``ProgramFPTileOpsMixin`` (``tile_row_mul_fp``, ``vram_fill_zero``). + """ + + def kda_decode_step_v0( + self, + *, + state: VRAMMatrixVar, + q_fp: FPVar, + k_fp: FPVar, + decay_fp: FPVar, + beta_fp: FPVar, + v: VRAMMatrixVar, + o: VRAMMatrixVar, + pred: VRAMMatrixVar, + err: VRAMMatrixVar, + shape: KdaShape, + output_scale_fp: FPVar, + head_rows: Sequence[int] | None = None, + fp_head_stride: int | None = None, + ) -> VRAMMatrixVar: + """One recurrent KDA token, all heads. + + ``q_fp`` / ``k_fp`` / ``decay_fp`` are FPRAM arrays holding the + already-normalised ``q_hat[h, k]``, ``k_hat[h, k]`` and ``decay[h, k]``, + laid out ``k``-major within each head -- so head ``h``'s scalars start at + offset ``h * key_dim``. ``beta_fp[h]`` and ``output_scale_fp[0]`` are + single scalars. They get to FPRAM through ``S_MAP_FP_V``; see + ``ssm_decode_scalars_to_fpram_v0``. + + ``state`` is ``[num_heads * blocks * key_dim, mlen]``, indexed by + :func:`kda_state_row`. ``v``, ``o``, ``pred`` and ``err`` are + ``[num_heads * blocks, mlen]``, indexed by :func:`kda_vector_row`. + Every tile is exactly one column block wide -- the block index lives in + the row. There is no scratch tile: the FMA accumulates in place, which + is most of what it bought. Before Task 8 each sweep staged a row through + scratch, and because that was the *same* row every iteration the + destination never formed a progression, so the sweep could not become a + hardware loop. + + A convenience wrapper over :meth:`kda_decode_predict_v0` and + :meth:`kda_decode_update_v0`. Call those two directly when ``q_fp`` and + ``decay_fp`` must share one FPRAM window; see the predict half. + """ + heads, stride = self._kda_decode_setup( + shape=shape, tiles=(("state", state), ("v", v), ("o", o), + ("pred", pred), ("err", err)), + state=state, pred=pred, err=err, + head_rows=head_rows, fp_head_stride=fp_head_stride, + ) + self.kda_decode_predict_v0( + state=state, k_fp=k_fp, decay_fp=decay_fp, + beta_fp=beta_fp, v=v, pred=pred, err=err, shape=shape, + head_rows=heads, fp_head_stride=stride, + ) + self.kda_decode_update_v0( + state=state, k_fp=k_fp, q_fp=q_fp, o=o, err=err, + shape=shape, output_scale_fp=output_scale_fp, + head_rows=heads, fp_head_stride=stride, + ) + return o + + def _kda_decode_setup( + self, *, shape, tiles, state, pred, err, head_rows, fp_head_stride + ): + """Shared validation for both halves; returns (heads, stride).""" + for name, tile in tiles: + if tile.shape[1] != self.mlen: + raise ValueError( + f"{name} is {tile.shape[1]} columns wide; every KDA tile must be " + f"exactly mlen ({self.mlen}), with the column block folded into " + f"the row index -- see program_kda_common's module docstring" + ) + for name, tile in (("pred", pred), ("err", err)): + if tile is not None and tile.name == state.name: + raise ValueError(f"{name} must not alias state") + + heads = list(range(shape.num_heads)) if head_rows is None else list(head_rows) + # How far apart consecutive heads' scalars sit in FPRAM. The default, + # key_dim, is the layout where every head's decay / q_hat / k_hat are + # materialised at once. Pass 0 when the caller streams one head at a + # time into a reused window -- FPRAM is 512 slots and the all-at-once + # layout wants 3 * num_heads * key_dim, which is 36,960 for Kimi K3. + stride = shape.key_dim if fp_head_stride is None else fp_head_stride + if stride < 0: + raise ValueError(f"fp_head_stride must not be negative, got {stride}") + if stride == 0 and len(heads) > 1: + raise ValueError( + "fp_head_stride=0 means every head reads the same FPRAM window, " + f"so only one head may be lowered per call; got {len(heads)}" + ) + return heads, stride + + def kda_decode_predict_v0( + self, + *, + state: VRAMMatrixVar, + k_fp: FPVar, + decay_fp: FPVar, + beta_fp: FPVar, + v: VRAMMatrixVar, + pred: VRAMMatrixVar, + err: VRAMMatrixVar, + shape: KdaShape, + head_rows: Sequence[int] | None = None, + fp_head_stride: int | None = None, + ) -> VRAMMatrixVar: + """First half: decay the state, predict, and form the error. + + Reads ``decay_fp`` and ``k_fp``; never touches ``q_fp``. Leaves ``err`` + holding ``beta * (v - pred)`` and ``state`` holding the decayed state. + + Split from the second half so a caller that is short of FPRAM can let + ``q_fp`` and ``decay_fp`` be the *same* window: decay is dead once every + block of this head has been predicted, and ``q_hat`` is not read until + the read-out. That halves the per-head window from ``3 * key_dim`` to + ``2 * key_dim``, which is what makes Kimi K3 fit in 512 slots at all. + """ + heads, stride = self._kda_decode_setup( + shape=shape, tiles=(("state", state), ("v", v), ("pred", pred), + ("err", err)), + state=state, pred=pred, err=err, + head_rows=head_rows, fp_head_stride=fp_head_stride, + ) + blocks = kda_blocks(shape, self.mlen) + k_dim = shape.key_dim + + for h in heads: + fp_base = h * stride + for c in range(blocks): + first = kda_state_row(shape, self.mlen, h, c, 0) + key_rows = list(range(first, first + k_dim)) + acc = kda_vector_row(shape, self.mlen, h, c) + + # decay is per (head, key), so this walks a different FPRAM + # slot per row -- the tile_row_mul_fp family, not _broadcast. + self.emit_comment(kda_stage_marker("kda_decay", f"head={h} block={c}")) + self._kda_scale_rows(state, decay_fp, key_rows, fp_base) + + self.emit_comment( + kda_stage_marker("kda_state_update", f"predict head={h} block={c}") + ) + # pred[acc] += state[row] * k_hat[i], contracting over key. + # The destination is pinned and the source walks, so both are + # progressions and the whole contraction is one hardware loop. + self.vram_fill_zero(pred, rows=[acc]) + self.tile_row_fma_fp_sweep( + pred, state, k_fp, + dst_rows=[acc] * len(key_rows), src_rows=key_rows, + fpram_offset=fp_base, + ) + + # err = beta_h * (v - pred), per block. beta_fp is one slot per + # head and produced for every head at once, so it is indexed by + # head number and not through fp_base. + self.emit_comment( + kda_stage_marker("kda_state_update", f"error head={h} block={c}") + ) + self.mamba_row_copy(err, acc, v, acc) + self.tile_row_sub(err, pred, rows=[acc]) + self._kda_scale_rows(err, beta_fp, [acc], h) + return err + + def kda_decode_update_v0( + self, + *, + state: VRAMMatrixVar, + k_fp: FPVar, + q_fp: FPVar, + o: VRAMMatrixVar, + err: VRAMMatrixVar, + shape: KdaShape, + output_scale_fp: FPVar, + head_rows: Sequence[int] | None = None, + fp_head_stride: int | None = None, + ) -> VRAMMatrixVar: + """Second half: the rank-1 update, then the read-out. + + Reads ``k_fp`` and ``q_fp``; never touches ``decay_fp``. Must run after + :meth:`kda_decode_predict_v0` on the same head, and must see the ``err`` + it produced. + """ + heads, stride = self._kda_decode_setup( + shape=shape, tiles=(("state", state), ("o", o), ("err", err)), + state=state, pred=None, err=err, + head_rows=head_rows, fp_head_stride=fp_head_stride, + ) + blocks = kda_blocks(shape, self.mlen) + k_dim = shape.key_dim + + for h in heads: + fp_base = h * stride + for c in range(blocks): + first = kda_state_row(shape, self.mlen, h, c, 0) + key_rows = list(range(first, first + k_dim)) + acc = kda_vector_row(shape, self.mlen, h, c) + + # Every key's update touches only that key's row, so updating + # all of them and then reducing is identical to the reference's + # interleaved `updated = ...; reduced += updated*q`. + self.emit_comment( + kda_stage_marker("kda_state_update", f"rank-1 head={h} block={c}") + ) + # state[row] += err[acc] * k_hat[i] -- the mirror of the + # prediction: the source is pinned and the destination walks. + self.tile_row_fma_fp_sweep( + state, err, k_fp, + dst_rows=key_rows, src_rows=[acc] * len(key_rows), + fpram_offset=fp_base, + ) + + self.emit_comment(kda_stage_marker("kda_readout", f"head={h} block={c}")) + self.vram_fill_zero(o, rows=[acc]) + self.tile_row_fma_fp_sweep( + o, state, q_fp, + dst_rows=[acc] * len(key_rows), src_rows=key_rows, + fpram_offset=fp_base, + ) + self._kda_scale_rows(o, output_scale_fp, [acc], 0) + return o diff --git a/aten/tests/test_kda_conv.py b/aten/tests/test_kda_conv.py new file mode 100644 index 0000000..d6e1fda --- /dev/null +++ b/aten/tests/test_kda_conv.py @@ -0,0 +1,246 @@ +"""KDA's causal conv step, executed against the CPU reference. + +Same approach as `test_kda_decode_step.py`: the emitted assembly runs on +`aten/tests/isa_interpreter.py`, so these compare numbers rather than asserting +which opcodes appear. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda import reference as kda_ref # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_kda_common import ( # noqa: E402 + kda_conv_blocks, + kda_conv_state_row, +) +from compiler.aten.tests.isa_interpreter import Machine # noqa: E402 + +MLEN = 8 +#: exp() saturates well before this; the reference's silu does not, so keeping +#: inputs modest is about comparing the same function, not about hiding a bug. +SCALE = 0.6 + + +def _rows_up(n: int) -> int: + return ((n + MLEN - 1) // MLEN) * MLEN + + +class _ConvHarness: + def __init__(self, channels: int, kernel: int): + self.channels, self.kernel = channels, kernel + self.blocks = kda_conv_blocks(channels, MLEN) + p = self.prog = PlenaCompiler(mlen=MLEN, blen=2) + + rows = _rows_up(self.blocks * kernel) + self.conv_state = p.alloc("conv_state", rows, MLEN) + self.weight = p.alloc("weight", rows, MLEN) + blk = _rows_up(self.blocks) + self.bias = p.alloc("bias", blk, MLEN) + self.x_new = p.alloc("x_new", blk, MLEN) + self.out = p.alloc("out", blk, MLEN) + self.scratch = p.alloc("scratch", MLEN, MLEN) + self.consts = p.kda_fp_constants() + self.mark = len(p.get_code()) + + def _base(self, var) -> int: + return self.prog.get_vram_layout(var.name).vram_base_addr + + def emit(self, *, use_bias: bool = True, apply_silu: bool = True) -> str: + self.prog.kda_conv_step_v0( + x_new=self.x_new, conv_state=self.conv_state, weight=self.weight, + bias=self.bias if use_bias else None, out=self.out, scratch=self.scratch, + consts=self.consts, channels=self.channels, kernel=self.kernel, + apply_silu=apply_silu, + ) + return self.prog.get_code()[self.mark :] + + def _wide(self, m, tile, vec) -> None: + """Scatter a channels-wide vector across its column blocks.""" + for cb in range(self.blocks): + m.write_vram_row( + self._base(tile) + cb * MLEN, vec[cb * MLEN : (cb + 1) * MLEN].tolist() + ) + + def run(self, *, state, weight, bias, x_new, **kw): + """state [channels, kernel]; weight [channels, kernel]; bias/x_new [channels].""" + code = self.emit(**kw) + m = Machine(vlen=MLEN) + for cb in range(self.blocks): + lanes = slice(cb * MLEN, (cb + 1) * MLEN) + for t in range(self.kernel): + row = kda_conv_state_row(self.channels, MLEN, self.kernel, cb, t) + m.write_vram_row( + self._base(self.conv_state) + row * MLEN, state[lanes, t].tolist() + ) + m.write_vram_row( + self._base(self.weight) + row * MLEN, weight[lanes, t].tolist() + ) + self._wide(m, self.bias, bias) + self._wide(m, self.x_new, x_new) + # `out` is a persistent tile: on token n+1 it holds token n's + # activations, and q/k/v reuse it. Seeding it with garbage makes the + # emitter's fill-zero load-bearing -- a fresh Machine zeroes VRAM, so + # without this an accumulator that is never cleared is indistinguishable + # from one that happens to start at zero. + for cb in range(self.out.shape[0]): + m.write_vram_row(self._base(self.out) + cb * MLEN, [7.5] * MLEN) + m.write_fpram( + self.consts.zero.address, PlenaCompiler.kda_fp_constant_values() + ) + m.run(code) + + out = torch.tensor( + [ + x + for cb in range(self.blocks) + for x in m.read_vram_row(self._base(self.out) + cb * MLEN, MLEN) + ], + dtype=torch.float32, + ) + new_state = torch.tensor( + [ + [ + m.read_vram_row( + self._base(self.conv_state) + + kda_conv_state_row(self.channels, MLEN, self.kernel, cb, t) * MLEN, + MLEN, + ) + for t in range(self.kernel) + ] + for cb in range(self.blocks) + ], + dtype=torch.float32, + ) # [blocks, kernel, mlen] + new_state = new_state.permute(0, 2, 1).reshape(self.channels, self.kernel) + return out, new_state, m + + +def _case(seed: int, channels: int, kernel: int): + torch.manual_seed(seed) + state = SCALE * torch.randn(channels, kernel) + weight = SCALE * torch.randn(channels, kernel) + bias = SCALE * torch.randn(channels) + x_new = SCALE * torch.randn(channels) + out, new = kda_ref._causal_conv_step( + x_new[None, :], state[None, :, :], weight, bias + ) + return dict( + state=state, weight=weight, bias=bias, x_new=x_new, + expected_out=out[0], expected_state=new[0], + ) + + +@pytest.mark.parametrize( + "seed,channels,kernel", + [ + (1, MLEN, 4), # one channel block, the Kimi kernel width + (2, MLEN * 2, 4), # two blocks: the case Kimi actually needs + (3, MLEN * 3, 2), # three blocks, short kernel + (4, MLEN, 1), # kernel 1: no history, pure pointwise + ], +) +def test_conv_step_matches_the_reference(seed, channels, kernel): + c = _case(seed, channels, kernel) + h = _ConvHarness(channels, kernel) + out, new_state, _ = h.run( + state=c["state"], weight=c["weight"], bias=c["bias"], x_new=c["x_new"] + ) + torch.testing.assert_close(out, c["expected_out"], rtol=2e-5, atol=2e-6) + torch.testing.assert_close(new_state, c["expected_state"], rtol=1e-5, atol=1e-6) + + +def test_history_shifts_rather_than_being_overwritten(): + """The oldest tap must fall off and the rest move down by one. A roll in the + wrong direction, or an append that clobbers a live tap, still produces + finite output of the right shape.""" + kernel = 4 + c = _case(7, MLEN, kernel) + h = _ConvHarness(MLEN, kernel) + _, new_state, _ = h.run( + state=c["state"], weight=c["weight"], bias=c["bias"], x_new=c["x_new"] + ) + # taps 0..k-2 of the new state are taps 1..k-1 of the old + torch.testing.assert_close(new_state[:, :-1], c["state"][:, 1:], rtol=0, atol=0) + # the newest tap is this token + torch.testing.assert_close(new_state[:, -1], c["x_new"], rtol=1e-6, atol=1e-7) + + +def test_bias_is_optional_and_actually_applied(): + c = _case(8, MLEN, 4) + with_bias, _, _ = _ConvHarness(MLEN, 4).run( + state=c["state"], weight=c["weight"], bias=c["bias"], x_new=c["x_new"] + ) + without, _, _ = _ConvHarness(MLEN, 4).run( + state=c["state"], weight=c["weight"], bias=c["bias"], x_new=c["x_new"], + use_bias=False, + ) + assert not torch.allclose(with_bias, without, rtol=1e-3, atol=1e-3) + + ref_no_bias, _ = kda_ref._causal_conv_step( + c["x_new"][None, :], c["state"][None, :, :], c["weight"], None + ) + torch.testing.assert_close(without, ref_no_bias[0], rtol=2e-5, atol=2e-6) + + +def test_channel_blocks_do_not_leak_into_each_other(): + """Every channel block has its own history and its own taps.""" + channels, kernel = MLEN * 3, 3 + c = _case(9, channels, kernel) + h = _ConvHarness(channels, kernel) + out, _, _ = h.run( + state=c["state"], weight=c["weight"], bias=c["bias"], x_new=c["x_new"] + ) + torch.testing.assert_close(out, c["expected_out"], rtol=2e-5, atol=2e-6) + + # perturb only block 1's input; blocks 0 and 2 must not move + perturbed = c["x_new"].clone() + perturbed[MLEN : 2 * MLEN] += 1.0 + out2, _, _ = _ConvHarness(channels, kernel).run( + state=c["state"], weight=c["weight"], bias=c["bias"], x_new=perturbed + ) + torch.testing.assert_close(out2[:MLEN], out[:MLEN], rtol=0, atol=0) + torch.testing.assert_close(out2[2 * MLEN :], out[2 * MLEN :], rtol=0, atol=0) + assert not torch.allclose(out2[MLEN : 2 * MLEN], out[MLEN : 2 * MLEN]) + + +def test_rejects_shapes_it_would_get_wrong(): + p = PlenaCompiler(mlen=MLEN, blen=2) + consts = p.kda_fp_constants() + ok = dict( + conv_state=p.alloc("cs", MLEN, MLEN), weight=p.alloc("w", MLEN, MLEN), + bias=None, out=p.alloc("o", MLEN, MLEN), scratch=p.alloc("sc", MLEN, MLEN), + consts=consts, channels=MLEN, kernel=4, + ) + with pytest.raises(ValueError, match="multiple of mlen"): + p.kda_conv_step_v0(x_new=p.alloc("x1", MLEN, MLEN), **{**ok, "channels": MLEN + 3}) + with pytest.raises(ValueError, match="exactly mlen"): + p.kda_conv_step_v0(x_new=p.alloc("x2", MLEN, MLEN * 2), **ok) + with pytest.raises(ValueError, match="conv_state needs"): + p.kda_conv_step_v0( + x_new=p.alloc("x3", MLEN, MLEN), **{**ok, "channels": MLEN * 4, "kernel": 4} + ) + # scratch aliased onto ANY operand: mamba_row_copy is zero-then-add, so it + # wipes whatever it aliases mid-loop. + for victim in ("conv_state", "weight", "out"): + with pytest.raises(ValueError, match="distinct"): + p.kda_conv_step_v0( + x_new=p.alloc(f"xa_{victim}", MLEN, MLEN), + **{**ok, "scratch": ok[victim]}, + ) + biased = {**ok, "bias": p.alloc("b1", MLEN, MLEN)} + with pytest.raises(ValueError, match="distinct"): + p.kda_conv_step_v0( + x_new=p.alloc("xb", MLEN, MLEN), + **{**biased, "scratch": biased["bias"]}, + ) + x_alias = p.alloc("xc", MLEN, MLEN) + with pytest.raises(ValueError, match="distinct"): + p.kda_conv_step_v0(x_new=x_alias, **{**ok, "scratch": x_alias}) diff --git a/aten/tests/test_kda_decode_step.py b/aten/tests/test_kda_decode_step.py new file mode 100644 index 0000000..3d44f47 --- /dev/null +++ b/aten/tests/test_kda_decode_step.py @@ -0,0 +1,619 @@ +"""The KDA decode step, executed and compared against the CPU reference. + +`aten/tests/isa_interpreter.py` runs the emitted assembly, so these are +numerical tests of the lowering rather than assertions about which opcodes +appear. See that module's docstring for what it does and does not model. + +Phase 0: the lowering uses only instructions that already exist. `V_FMA_VF` +arrives in Task 7 and Task 8 converts this kernel onto it; the counts recorded +by `test_records_the_phase0_instruction_count` are the baseline that change is +measured against. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.reference import KdaState, activate_log_decay, kda_step # noqa: E402 +from compiler.aten.models.kda.shape import KdaShape # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_kda_common import ( # noqa: E402 + kda_blocks, + kda_stage_marker, + kda_state_row, + kda_state_rows, + kda_vector_row, + kda_vector_rows, +) +from compiler.aten.tests.isa_interpreter import Machine, UnsupportedInstruction # noqa: E402 + +#: What the recurrence emitted before Task 8. Named rather than inlined so the +#: assertion below reads as "one opcode more than this" and not as a magic list. +_PHASE0_OPCODES = frozenset( + { + "S_ADDI_INT", + "S_LD_FP", + "S_ST_FP", + "S_ADD_FP", + "V_MUL_VF", + "V_ADD_VF", + "V_SUB_VF", + "V_ADD_VV", + "V_SUB_VV", + "V_MUL_VV", + "V_RED_SUM", + "C_LOOP_START", + "C_LOOP_END", + } +) + +MLEN = 8 + + +def _rows_up(n: int) -> int: + """VRAM row counts must be a multiple of mlen.""" + return ((n + MLEN - 1) // MLEN) * MLEN + + +class _Harness: + """Lower one decode step, seed VRAM/FPRAM from tensors, execute, read back.""" + + def __init__(self, shape: KdaShape): + self.shape = shape + self.prog = PlenaCompiler(mlen=MLEN, blen=2) + p, s = self.prog, shape + + # Every tile is exactly mlen wide; the column block lives in the row. + self.state = p.alloc("state", _rows_up(kda_state_rows(s, MLEN)), MLEN) + vec_rows = _rows_up(kda_vector_rows(s, MLEN)) + self.v = p.alloc("v", vec_rows, MLEN) + self.o = p.alloc("o", vec_rows, MLEN) + self.pred = p.alloc("pred", vec_rows, MLEN) + self.err = p.alloc("err", vec_rows, MLEN) + + n_key = s.num_heads * s.key_dim + self.q_fp = p.fp_var("q_hat", size=n_key) + self.k_fp = p.fp_var("k_hat", size=n_key) + self.decay_fp = p.fp_var("decay", size=n_key) + self.beta_fp = p.fp_var("beta", size=s.num_heads) + self.scale_fp = p.fp_var("out_scale", size=1) + + self.mark = len(p.get_code()) + + def emit(self, **kw): + self.prog.kda_decode_step_v0( + state=self.state, q_fp=self.q_fp, k_fp=self.k_fp, decay_fp=self.decay_fp, + beta_fp=self.beta_fp, v=self.v, o=self.o, pred=self.pred, + err=self.err, shape=self.shape, output_scale_fp=self.scale_fp, **kw, + ) + return self.prog.get_code()[self.mark :] + + def _base(self, var) -> int: + return self.prog.get_vram_layout(var.name).vram_base_addr + + def seed_state(self, m, state) -> None: + """state is [heads, key, value]; scatter it into the flattened tile.""" + s = self.shape + for h in range(s.num_heads): + for c in range(kda_blocks(s, MLEN)): + lanes = slice(c * MLEN, (c + 1) * MLEN) + for j in range(s.key_dim): + m.write_vram_row( + self._base(self.state) + kda_state_row(s, MLEN, h, c, j) * MLEN, + state[h, j, lanes].tolist(), + ) + + def seed_token(self, m, *, v, q_hat, k_hat, decay, beta, output_scale) -> None: + """Per-token inputs. Deliberately separate from seed_state: a decode + loop rewrites these every token while the state stays on chip.""" + s = self.shape + for h in range(s.num_heads): + for c in range(kda_blocks(s, MLEN)): + m.write_vram_row( + self._base(self.v) + kda_vector_row(s, MLEN, h, c) * MLEN, + v[h, c * MLEN : (c + 1) * MLEN].tolist(), + ) + m.write_fpram(self.q_fp.address, q_hat.reshape(-1).tolist()) + m.write_fpram(self.k_fp.address, k_hat.reshape(-1).tolist()) + m.write_fpram(self.decay_fp.address, decay.reshape(-1).tolist()) + m.write_fpram(self.beta_fp.address, beta.tolist()) + m.write_fpram(self.scale_fp.address, [output_scale]) + + def read_vector(self, m, tile): + """Gather a flattened [heads*blocks, mlen] tile back to [heads, value].""" + s = self.shape + return torch.tensor( + [ + [ + x + for c in range(kda_blocks(s, MLEN)) + for x in m.read_vram_row( + self._base(tile) + kda_vector_row(s, MLEN, h, c) * MLEN, MLEN + ) + ] + for h in range(s.num_heads) + ], + dtype=torch.float32, + ) + + def run(self, *, state, q_hat, k_hat, decay, beta, v, output_scale): + """state [heads,key,value]; q/k/decay [heads,key]; beta [heads]; v [heads,value].""" + s = self.shape + code = self.emit() + m = Machine(vlen=MLEN) + + # Seed the accumulators with garbage. A fresh Machine is all zeros, so + # without this 'the program clears it' and 'it started at zero' are + # indistinguishable -- deleting either vram_fill_zero used to leave + # every test here green but one. + for _acc in (self.pred, self.o): + for _r in range(_acc.shape[0]): + m.write_vram_row(self._base(_acc) + _r * MLEN, [7.5] * MLEN) + + blocks = kda_blocks(s, MLEN) + self.seed_state(m, state) + self.seed_token(m, v=v, q_hat=q_hat, k_hat=k_hat, decay=decay, + beta=beta, output_scale=output_scale) + m.run(code) + + out = torch.tensor( + [ + [ + x + for c in range(blocks) + for x in m.read_vram_row( + self._base(self.o) + kda_vector_row(s, MLEN, h, c) * MLEN, MLEN + ) + ] + for h in range(s.num_heads) + ], + dtype=torch.float32, + ) + new_state = torch.tensor( + [ + [ + [ + x + for c in range(blocks) + for x in m.read_vram_row( + self._base(self.state) + + kda_state_row(s, MLEN, h, c, j) * MLEN, + MLEN, + ) + ] + for j in range(s.key_dim) + ] + for h in range(s.num_heads) + ], + dtype=torch.float32, + ) + return out, new_state, m + + +def _case(seed: int, *, num_heads: int, key_dim: int, value_dim: int = MLEN): + """One random decode step: the shape, the reference result, and the FPRAM + scalars the lowering expects to be handed already normalised.""" + torch.manual_seed(seed) + shape = KdaShape( + hidden_size=num_heads * value_dim, + num_heads=num_heads, + key_dim=key_dim, + value_dim=value_dim, + conv_kernel=4, + ) + state = torch.randn(1, num_heads, value_dim, key_dim) + q = 0.5 * torch.randn(1, num_heads, key_dim) + k = 0.5 * torch.randn(1, num_heads, key_dim) + v = torch.randn(1, num_heads, value_dim) + gate = torch.randn(1, num_heads, key_dim) + beta_logit = torch.randn(1, num_heads) + a_log = torch.randn(num_heads) + dt_bias = torch.randn(num_heads, key_dim) + + out, new = kda_step( + q, k, v, gate, beta_logit, KdaState(state.clone()), a_log, dt_bias, shape + ) + + def _norm(x): + # Must match reference.py exactly: epsilon inside rsqrt, not F.normalize. + return x.float() * torch.rsqrt(x.float().square().sum(-1, keepdim=True) + 1.0e-6) + + log_decay = activate_log_decay( + gate, a_log, dt_bias, lower_bound=shape.gate_lower_bound + ) + return shape, { + # inputs to the lowering, all batch-0 + "state_T": state[0].transpose(-2, -1).contiguous(), # [heads, key, value] + "q_hat": _norm(q)[0], + "k_hat": _norm(k)[0], + "decay": torch.exp(log_decay)[0], + "beta": torch.sigmoid(beta_logit.float())[0], + "v": v[0], + "output_scale": 1.0 / (shape.key_dim**0.5), + # expected results + "out": out[0], + "state_out_T": new.recurrent[0].transpose(-2, -1).contiguous(), + } + + +def _run(shape, ref): + h = _Harness(shape) + out, new_state, machine = h.run( + state=ref["state_T"], + q_hat=ref["q_hat"], + k_hat=ref["k_hat"], + decay=ref["decay"], + beta=ref["beta"], + v=ref["v"], + output_scale=ref["output_scale"], + ) + return out, new_state, machine + + +# --------------------------------------------------------------------------- +# Numerical equivalence with the CPU reference +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "seed,heads,key_dim", + [ + (1, 1, 2), # smallest thing that can be wrong + (2, 1, 5), # key_dim != value_dim: a transposed axis is not a shape error + (3, 3, 4), # multiple heads: catches per-head FPRAM offset bugs + (4, 2, 8), # key_dim == value_dim == mlen, the square case + ], +) +def test_decode_step_matches_the_reference(seed, heads, key_dim): + shape, ref = _case(seed, num_heads=heads, key_dim=key_dim) + out, new_state, _ = _run(shape, ref) + torch.testing.assert_close(out, ref["out"], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(new_state, ref["state_out_T"], rtol=1e-5, atol=1e-6) + + +def test_read_out_uses_the_updated_state_not_the_decayed_one(): + """The delta rule reads the state *after* the rank-1 update. + + Reducing over the decayed-but-not-updated state is a plausible reordering + that changes the answer, and the output alone would still look reasonable. + """ + shape, ref = _case(7, num_heads=2, key_dim=4) + out, _, _ = _run(shape, ref) + + stale = torch.einsum( + "hkv,hk->hv", ref["state_T"] * ref["decay"][:, :, None], ref["q_hat"] + ) * ref["output_scale"] + assert not torch.allclose(out, stale, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(out, ref["out"], rtol=1e-5, atol=1e-6) + + +def test_heads_do_not_leak_into_each_other(): + """Head h must touch only its own state rows and its own FPRAM slice. + Running one head at a time must reproduce the all-heads run.""" + shape, ref = _case(5, num_heads=3, key_dim=4) + together, _, _ = _run(shape, ref) + + for h in range(shape.num_heads): + harness = _Harness(shape) + code = harness.emit(head_rows=[h]) + m = Machine(vlen=MLEN) + harness.seed_state(m, ref["state_T"]) + harness.seed_token( + m, v=ref["v"], q_hat=ref["q_hat"], k_hat=ref["k_hat"], + decay=ref["decay"], beta=ref["beta"], output_scale=ref["output_scale"], + ) + m.run(code) + + torch.testing.assert_close( + harness.read_vector(m, harness.o)[h], together[h], rtol=1e-5, atol=1e-6 + ) + + # every other head's state untouched, bit for bit + for other in range(shape.num_heads): + if other == h: + continue + for c in range(kda_blocks(shape, MLEN)): + for j in range(shape.key_dim): + got = m.read_vram_row( + harness._base(harness.state) + + kda_state_row(shape, MLEN, other, c, j) * MLEN, + MLEN, + ) + torch.testing.assert_close( + torch.tensor(got), + ref["state_T"][other, j, c * MLEN : (c + 1) * MLEN], + rtol=0, atol=0, + ) + + +# --------------------------------------------------------------------------- +# Guards and baseline +# --------------------------------------------------------------------------- + + +def test_rejects_a_value_dim_that_is_not_a_whole_number_of_blocks(): + """A partial trailing block would leave lanes past value_dim holding + whatever was there before -- and they would still be summed.""" + shape = KdaShape( + hidden_size=64, num_heads=1, key_dim=2, value_dim=MLEN + 3, conv_kernel=4 + ) + with pytest.raises(ValueError, match="multiple of mlen"): + _Harness(shape).emit() + + +def test_rejects_a_tile_that_is_not_exactly_one_block_wide(): + """Every KDA tile carries the column block in its row index. A wider tile + would be swept by helpers that silently cover only block 0.""" + shape, _ = _case(1, num_heads=1, key_dim=2) + h = _Harness(shape) + wide = h.prog.alloc("wide", MLEN, MLEN * 2) + with pytest.raises(ValueError, match="exactly mlen"): + h.prog.kda_decode_step_v0( + state=h.state, q_fp=h.q_fp, k_fp=h.k_fp, + decay_fp=h.decay_fp, beta_fp=h.beta_fp, v=h.v, o=h.o, pred=wide, + err=h.err, shape=shape, output_scale_fp=h.scale_fp, + ) + + +def test_rejects_pred_aliased_to_state(): + """Was `scratch` aliased to state. Task 8 removed the scratch tile -- the + FMA accumulates in place -- so the aliasing that still matters is `pred` + and `err`, which the recurrence writes while reading state.""" + shape, _ = _case(1, num_heads=1, key_dim=2) + h = _Harness(shape) + with pytest.raises(ValueError, match="must not alias state"): + h.prog.kda_decode_step_v0( + state=h.state, q_fp=h.q_fp, k_fp=h.k_fp, + decay_fp=h.decay_fp, beta_fp=h.beta_fp, v=h.v, o=h.o, pred=h.state, + err=h.err, shape=shape, output_scale_fp=h.scale_fp, + ) + + +def test_rejects_a_negative_fp_head_stride(): + shape, _ = _case(1, num_heads=2, key_dim=2) + h = _Harness(shape) + with pytest.raises(ValueError, match="must not be negative"): + h.prog.kda_decode_step_v0( + state=h.state, q_fp=h.q_fp, k_fp=h.k_fp, + decay_fp=h.decay_fp, beta_fp=h.beta_fp, v=h.v, o=h.o, pred=h.pred, + err=h.err, shape=shape, output_scale_fp=h.scale_fp, + fp_head_stride=-1, + ) + + +def test_rejects_more_than_one_head_when_the_window_is_reused(): + """fp_head_stride=0 means every head reads the same FPRAM slots. Two heads + would then silently share one head's decay, q_hat and k_hat -- a finite, + plausible, wrong answer, which is exactly what this guard exists to stop.""" + shape, _ = _case(1, num_heads=2, key_dim=2) + h = _Harness(shape) + with pytest.raises(ValueError, match="only one head may be lowered"): + h.prog.kda_decode_step_v0( + state=h.state, q_fp=h.q_fp, k_fp=h.k_fp, + decay_fp=h.decay_fp, beta_fp=h.beta_fp, v=h.v, o=h.o, pred=h.pred, + err=h.err, shape=shape, output_scale_fp=h.scale_fp, + head_rows=[0, 1], fp_head_stride=0, + ) + + +def test_the_two_halves_emit_the_same_program_as_the_whole_step(): + """kda_decode_step_v0 is a wrapper. If the halves ever drift from it, the + mixer -- which calls the halves so it can share one FPRAM window between + decay and q_hat -- silently stops matching what these tests cover.""" + shape, _ = _case(3, num_heads=2, key_dim=2) + + whole = _Harness(shape) + whole.prog.kda_decode_step_v0( + state=whole.state, q_fp=whole.q_fp, + k_fp=whole.k_fp, decay_fp=whole.decay_fp, beta_fp=whole.beta_fp, + v=whole.v, o=whole.o, pred=whole.pred, err=whole.err, shape=shape, + output_scale_fp=whole.scale_fp, + ) + + halves = _Harness(shape) + halves.prog.kda_decode_predict_v0( + state=halves.state, k_fp=halves.k_fp, + decay_fp=halves.decay_fp, beta_fp=halves.beta_fp, v=halves.v, + pred=halves.pred, err=halves.err, shape=shape, + ) + halves.prog.kda_decode_update_v0( + state=halves.state, k_fp=halves.k_fp, + q_fp=halves.q_fp, o=halves.o, err=halves.err, shape=shape, + output_scale_fp=halves.scale_fp, + ) + + assert whole.prog.get_code()[whole.mark :] == halves.prog.get_code()[halves.mark :] + + +def test_emits_exactly_one_opcode_beyond_the_phase0_set(): + """This asserted `V_FMA_VF not in code` through Phase 0, when the point was + that KDA needs no new opcode. That claim was demonstrated and now it is + spent: Task 8 converts the sweeps onto the one instruction Phase 1 adds. + + What still has to hold is that it is exactly *one*. The interpreter raises + on anything it does not model, so a lowering that reaches for a second new + opcode fails here rather than being silently half-checked. + """ + shape, ref = _case(1, num_heads=2, key_dim=4) + code = _Harness(shape).emit() + assert "V_FMA_VF" in code, "the sweeps are supposed to be fused now" + ops = { + line.strip().replace(",", " ").split()[0] + for line in code.splitlines() + if line.strip() and not line.strip().startswith(";") + } + assert ops - _PHASE0_OPCODES == {"V_FMA_VF"}, ( + f"beyond Phase 0's set the decode step may emit only V_FMA_VF, " + f"got {sorted(ops - _PHASE0_OPCODES)}" + ) + try: + Machine(vlen=MLEN).run(code) + except UnsupportedInstruction as exc: # pragma: no cover - failure path + pytest.fail(f"decode step emitted an unmodelled instruction: {exc}") + + +def test_the_sweeps_are_hardware_loops_not_unrolled(): + """The static-footprint claim. Each (head, block) contributes three sweeps + -- predict, rank-1, read-out -- and each must be one FMA inside one loop, + however many keys it walks. + + Without this, a change that broke the row progression would stay green: the + numbers would still be right, and the program would just be 64x larger. + """ + for key_dim in (4, 8, 16): + shape, _ = _case(2, num_heads=1, key_dim=key_dim) + code = _Harness(shape).emit() + body = [ + line.strip() + for line in code.splitlines() + if line.strip() and not line.strip().startswith(";") + ] + blocks = kda_blocks(shape, MLEN) + assert sum("V_FMA_VF" in line for line in body) == 3 * blocks, ( + f"key_dim={key_dim}: expected one FMA per sweep, got " + f"{sum('V_FMA_VF' in line for line in body)}" + ) + + +def test_records_the_phase0_instruction_count(): + """Baseline for Task 8. Not a budget -- the number this prints, against the + one Task 8 prints, is the measured case for spending opcode 0x3B.""" + shape, ref = _case(3, num_heads=1, key_dim=8) + code = _Harness(shape).emit() + static = len([l for l in code.splitlines() if l.strip() and not l.strip().startswith(";")]) + _, _, machine = _run(shape, ref) + print( + f"PHASE0_KDA_DECODE heads={shape.num_heads} key_dim={shape.key_dim} " + f"static={static} dynamic={machine.executed}" + ) + assert static > 0 and machine.executed > 0 + + +def test_two_tokens_in_one_program_match_two_reference_steps(): + """The accumulators must be zeroed, and only a second token can prove it. + + Every other test runs one step on a freshly-zeroed Machine, so an + accumulator that is never cleared is indistinguishable from one that happens + to start at zero. A real decode loop reuses the same `pred` and `o` tiles + every token: without the fill-zero, token 1 accumulates on top of token 0's + residue. Deleting either `vram_fill_zero` left the rest of this file green. + + This also covers the state being updated in place on chip and carried into + the next step, which is the whole point of a recurrent kernel. + """ + torch.manual_seed(31) + heads, key_dim, value_dim = 2, 4, MLEN * 2 + shape = KdaShape( + hidden_size=heads * value_dim, num_heads=heads, key_dim=key_dim, + value_dim=value_dim, conv_kernel=4, + ) + + state = torch.randn(1, heads, value_dim, key_dim) + tokens = [ + dict( + q=0.5 * torch.randn(1, heads, key_dim), + k=0.5 * torch.randn(1, heads, key_dim), + v=torch.randn(1, heads, value_dim), + gate=torch.randn(1, heads, key_dim), + beta_logit=torch.randn(1, heads), + ) + for _ in range(2) + ] + a_log = torch.randn(heads) + dt_bias = torch.randn(heads, key_dim) + + def _norm(x): + return x.float() * torch.rsqrt(x.float().square().sum(-1, keepdim=True) + 1.0e-6) + + # reference: two sequential steps carrying state + carried = KdaState(state.clone()) + expected_out = [] + for t in tokens: + out, carried = kda_step( + t["q"], t["k"], t["v"], t["gate"], t["beta_logit"], carried, + a_log, dt_bias, shape, + ) + expected_out.append(out[0]) + + # lowering: one program, the step emitted twice, same tiles reused + h = _Harness(shape) + programs = [] + for _ in range(2): + start = len(h.prog.get_code()) + h.emit() + programs.append(h.prog.get_code()[start:]) + + m = Machine(vlen=MLEN) + h.seed_state(m, state[0].transpose(-2, -1).contiguous()) + + # Token 1's scalars are written between the halves -- what a real decode + # loop does through S_MAP_FP_V. VRAM is never reset: `state` carries, and + # `pred` / `o` must be cleared by the program itself. + actual_out = [] + for t, program in zip(tokens, programs): + log_decay = activate_log_decay( + t["gate"], a_log, dt_bias, lower_bound=shape.gate_lower_bound + ) + h.seed_token( + m, + v=t["v"][0], + q_hat=_norm(t["q"])[0], + k_hat=_norm(t["k"])[0], + decay=torch.exp(log_decay)[0], + beta=torch.sigmoid(t["beta_logit"].float())[0], + output_scale=1.0 / (key_dim**0.5), + ) + m.run(program) + actual_out.append(h.read_vector(m, h.o)) + + torch.testing.assert_close(actual_out[0], expected_out[0], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(actual_out[1], expected_out[1], rtol=1e-5, atol=1e-6) + + +# --------------------------------------------------------------------------- +# Multiple column blocks -- the case the flattened layout exists for +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "seed,heads,key_dim,blocks", + [ + (11, 1, 3, 2), # smallest multi-block case + (12, 2, 5, 2), # several heads, key_dim != value_dim + (13, 1, 4, 3), # three blocks: catches an off-by-one in the block stride + ], +) +def test_decode_step_matches_the_reference_across_column_blocks(seed, heads, key_dim, blocks): + """value_dim > mlen. Kimi K3 is value_dim=128 against a default mlen=64. + + With the old [key, value] layout this was impossible: the helper family is + split between "walks every column block" and "silently block 0 only", so the + state came out correct in its first mlen lanes and stale in the rest. + """ + shape, ref = _case(seed, num_heads=heads, key_dim=key_dim, value_dim=MLEN * blocks) + assert kda_blocks(shape, MLEN) == blocks + out, new_state, _ = _run(shape, ref) + torch.testing.assert_close(out, ref["out"], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(new_state, ref["state_out_T"], rtol=1e-5, atol=1e-6) + + +def test_every_lane_past_the_first_block_is_actually_updated(): + """The failure the flattened layout prevents, asserted directly: lanes in + block 1 must differ from their input, not merely be finite.""" + shape, ref = _case(14, num_heads=1, key_dim=4, value_dim=MLEN * 2) + _, new_state, _ = _run(shape, ref) + before = ref["state_T"] + tail_before = before[:, :, MLEN:] + tail_after = new_state[:, :, MLEN:] + assert not torch.allclose(tail_after, tail_before, rtol=1e-4, atol=1e-5), ( + "lanes past the first column block were never touched" + ) + torch.testing.assert_close( + tail_after, ref["state_out_T"][:, :, MLEN:], rtol=1e-5, atol=1e-6 + ) diff --git a/aten/tests/test_kda_gates.py b/aten/tests/test_kda_gates.py new file mode 100644 index 0000000..a73f59a --- /dev/null +++ b/aten/tests/test_kda_gates.py @@ -0,0 +1,392 @@ +"""KDA's decay and beta scalars, executed against the CPU reference. + +`activate_log_decay` is the oracle for decay; `beta = sigmoid(beta_logit)` comes +from `kda_step`. Both end up in FPRAM, so these tests read the FP memory back +rather than a VRAM tile. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.reference import activate_log_decay # noqa: E402 +from compiler.aten.models.kda.shape import KdaShape # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_kda_gates import ( # noqa: E402 + kda_head_blocks, + kda_key_blocks, + kda_key_row, +) +from compiler.aten.tests.isa_interpreter import Machine # noqa: E402 + +MLEN = 8 + + +def _rows_up(n: int) -> int: + return ((n + MLEN - 1) // MLEN) * MLEN + + +def _shape(num_heads: int, key_dim: int) -> KdaShape: + return KdaShape( + hidden_size=num_heads * MLEN, + num_heads=num_heads, + key_dim=key_dim, + value_dim=MLEN, + conv_kernel=4, + ) + + +class _GateHarness: + def __init__(self, shape: KdaShape): + self.shape = shape + p = self.prog = PlenaCompiler(mlen=MLEN, blen=2) + kb = kda_key_blocks(shape, MLEN) + rows = _rows_up(shape.num_heads * kb) + self.gate = p.alloc("gate", rows, MLEN) + self.dt_bias = p.alloc("dt_bias", rows, MLEN) + self.beta_logit = p.alloc("beta_logit", _rows_up(kda_head_blocks(shape, MLEN)), MLEN) + self.rate = p.fp_var("rate", size=shape.num_heads) + self.lower = p.fp_var("lower", size=1) + self.decay = p.fp_var("decay", size=shape.num_heads * shape.key_dim) + # Sized to exactly what S_MAP_FP_V writes -- rows x mlen, not num_heads. + # The sentinel after it turns an over-write into a visible failure + # instead of a silent clobber of whatever came next. + self.beta = p.fp_var("beta", size=kda_head_blocks(shape, MLEN) * MLEN) + self.sentinel = p.fp_var("sentinel", size=MLEN) + self.consts = p.kda_fp_constants() + self.mark = len(p.get_code()) + + def _base(self, var) -> int: + return self.prog.get_vram_layout(var.name).vram_base_addr + + def run(self, *, gate, dt_bias, a_log, beta_logit): + """gate/dt_bias [heads, key_dim]; a_log [heads]; beta_logit [heads].""" + s = self.shape + self.prog.kda_decay_scalars_v0( + gate=self.gate, dt_bias=self.dt_bias, rate_fp=self.rate, + lower_bound_fp=self.lower, decay_fp=self.decay, consts=self.consts, + shape=s, + ) + self.prog.kda_beta_scalars_v0( + beta_logit=self.beta_logit, beta_fp=self.beta, consts=self.consts, shape=s, + ) + code = self.prog.get_code()[self.mark :] + + m = Machine(vlen=MLEN) + kb = kda_key_blocks(s, MLEN) + for h in range(s.num_heads): + for b in range(kb): + lanes = slice(b * MLEN, (b + 1) * MLEN) + row = kda_key_row(s, MLEN, h, b) + m.write_vram_row(self._base(self.gate) + row * MLEN, gate[h, lanes].tolist()) + m.write_vram_row( + self._base(self.dt_bias) + row * MLEN, dt_bias[h, lanes].tolist() + ) + for hb in range(kda_head_blocks(s, MLEN)): + chunk = beta_logit[hb * MLEN : (hb + 1) * MLEN].tolist() + m.write_vram_row( + self._base(self.beta_logit) + hb * MLEN, chunk + [0.0] * (MLEN - len(chunk)) + ) + + m.write_fpram(self.consts.zero.address, PlenaCompiler.kda_fp_constant_values()) + m.write_fpram(self.sentinel.address, [-99.0] * MLEN) + # rate = exp(a_log), a host constant: it depends only on a weight. + m.write_fpram(self.rate.address, torch.exp(a_log).tolist()) + m.write_fpram(self.lower.address, [s.gate_lower_bound]) + m.run(code) + + decay = torch.tensor( + [ + self.decay_slot(m, h, k) + for h in range(s.num_heads) + for k in range(s.key_dim) + ], + dtype=torch.float32, + ).reshape(s.num_heads, s.key_dim) + beta = torch.tensor( + [m.fpram[self.beta.address + h] for h in range(s.num_heads)], + dtype=torch.float32, + ) + assert all( + m.fpram[self.sentinel.address + i] == -99.0 for i in range(MLEN) + ), "an FPRAM write ran past its allocation into the next one" + return decay, beta, m + + def decay_slot(self, m, head: int, key: int) -> float: + """key-major within head, which is what the recurrence indexes.""" + return m.fpram[self.decay.address + head * self.shape.key_dim + key] + + +def _case(seed: int, num_heads: int, key_dim: int): + torch.manual_seed(seed) + shape = _shape(num_heads, key_dim) + gate = 0.8 * torch.randn(num_heads, key_dim) + dt_bias = 0.8 * torch.randn(num_heads, key_dim) + a_log = 0.5 * torch.randn(num_heads) + beta_logit = torch.randn(num_heads) + + log_decay = activate_log_decay( + gate[None, :, :], a_log, dt_bias, lower_bound=shape.gate_lower_bound + ) + return shape, dict( + gate=gate, dt_bias=dt_bias, a_log=a_log, beta_logit=beta_logit, + expected_decay=torch.exp(log_decay)[0], + expected_beta=torch.sigmoid(beta_logit), + ) + + +@pytest.mark.parametrize( + "seed,heads,key_dim", + [ + (1, 1, MLEN), # one head, one key block + (2, 3, MLEN), # several heads + (3, 1, MLEN * 2), # one head spanning two key blocks + (4, 4, MLEN * 2), # several heads, two key blocks -- Kimi's shape + (5, 2, MLEN * 3), # three key blocks + ], +) +def test_decay_and_beta_match_the_reference(seed, heads, key_dim): + shape, c = _case(seed, heads, key_dim) + decay, beta, _ = _GateHarness(shape).run( + gate=c["gate"], dt_bias=c["dt_bias"], a_log=c["a_log"], + beta_logit=c["beta_logit"], + ) + torch.testing.assert_close(decay, c["expected_decay"], rtol=3e-5, atol=3e-6) + torch.testing.assert_close(beta, c["expected_beta"], rtol=2e-5, atol=2e-6) + + +def test_heads_with_a_partial_beta_row_are_handled(): + """beta_logit holds one value per head, so its last row is partly empty + whenever num_heads is not a multiple of mlen -- Kimi K3 is 96 against 64.""" + heads = MLEN + 3 + shape, c = _case(6, heads, MLEN) + assert kda_head_blocks(shape, MLEN) == 2 + _, beta, _ = _GateHarness(shape).run( + gate=c["gate"], dt_bias=c["dt_bias"], a_log=c["a_log"], + beta_logit=c["beta_logit"], + ) + torch.testing.assert_close(beta, c["expected_beta"], rtol=2e-5, atol=2e-6) + + +def test_decay_uses_each_head_own_rate(): + """`rate = exp(a_log[h])` is per head. Broadcasting one head's rate to all + of them still yields plausible decays in (0, 1).""" + shape, c = _case(7, 3, MLEN * 2) + decay, _, _ = _GateHarness(shape).run( + gate=c["gate"], dt_bias=c["dt_bias"], a_log=c["a_log"], + beta_logit=c["beta_logit"], + ) + torch.testing.assert_close(decay, c["expected_decay"], rtol=3e-5, atol=3e-6) + + # every head's rate genuinely differs, so a shared rate would be visible + same_rate = torch.full_like(c["a_log"], c["a_log"][0].item()) + wrong = torch.exp( + activate_log_decay( + c["gate"][None, :, :], same_rate, c["dt_bias"], + lower_bound=shape.gate_lower_bound, + ) + )[0] + assert not torch.allclose(decay, wrong, rtol=1e-3, atol=1e-3) + + +def test_decay_lands_key_major_within_each_head(): + """The recurrence indexes decay_fp[h * key_dim + k]. tile_row_to_fpram puts + row i at base + i*mlen, so the row ordering is what makes that true.""" + shape, c = _case(8, 3, MLEN * 2) + h = _GateHarness(shape) + decay, _, m = h.run( + gate=c["gate"], dt_bias=c["dt_bias"], a_log=c["a_log"], + beta_logit=c["beta_logit"], + ) + for head in range(shape.num_heads): + for key in range(shape.key_dim): + assert h.decay_slot(m, head, key) == pytest.approx( + c["expected_decay"][head, key].item(), rel=3e-5, abs=3e-6 + ) + + +def test_decay_is_bounded_by_the_gate_lower_bound(): + """log_decay is in [lower_bound, 0], so decay is in [exp(lower_bound), 1].""" + shape, c = _case(9, 3, MLEN * 2) + decay, _, _ = _GateHarness(shape).run( + gate=c["gate"], dt_bias=c["dt_bias"], a_log=c["a_log"], + beta_logit=c["beta_logit"], + ) + import math + + assert decay.min().item() >= math.exp(shape.gate_lower_bound) - 1e-5 + assert decay.max().item() <= 1.0 + 1e-5 + + +def test_rejects_shapes_and_undersized_fpram(): + p = PlenaCompiler(mlen=MLEN, blen=2) + shape = _shape(2, MLEN) + consts = p.kda_fp_constants() + ok = dict( + dt_bias=p.alloc("db", MLEN, MLEN), rate_fp=p.fp_var("r", size=2), + lower_bound_fp=p.fp_var("l", size=1), consts=consts, shape=shape, + ) + with pytest.raises(ValueError, match="exactly mlen"): + p.kda_decay_scalars_v0( + gate=p.alloc("g1", MLEN, MLEN * 2), + decay_fp=p.fp_var("d1", size=2 * MLEN), **ok, + ) + with pytest.raises(ValueError, match="decay_fp"): + p.kda_decay_scalars_v0( + gate=p.alloc("g2", MLEN, MLEN), decay_fp=p.fp_var("d2", size=3), **ok, + ) + with pytest.raises(ValueError, match="rate_fp"): + p.kda_decay_scalars_v0( + gate=p.alloc("g3", MLEN, MLEN), decay_fp=p.fp_var("d3", size=2 * MLEN), + **{**ok, "rate_fp": p.fp_var("r2", size=1)}, + ) + with pytest.raises(ValueError, match="multiple of mlen"): + kda_key_blocks(_shape(1, MLEN + 1), MLEN) + + +def test_beta_fp_must_be_sized_for_the_full_row_write(): + """S_MAP_FP_V moves a whole mlen row, so sizing beta_fp to num_heads lets + the tail land in the next allocation. At Kimi K3's 96 heads against mlen 64 + that is 32 slots of sigmoid(padding) written outside the buffer.""" + heads = MLEN + 3 + shape = _shape(heads, MLEN) + p = PlenaCompiler(mlen=MLEN, blen=2) + consts = p.kda_fp_constants() + beta_logit = p.alloc("bl", _rows_up(kda_head_blocks(shape, MLEN)), MLEN) + with pytest.raises(ValueError, match="padded row count"): + p.kda_beta_scalars_v0( + beta_logit=beta_logit, + beta_fp=p.fp_var("small", size=heads), + consts=consts, + shape=shape, + ) + # exactly the padded size is accepted + p.kda_beta_scalars_v0( + beta_logit=beta_logit, + beta_fp=p.fp_var("right", size=kda_head_blocks(shape, MLEN) * MLEN), + consts=consts, + shape=shape, + ) + + +# --------------------------------------------------------------------------- +# Per-head streaming. +# +# FPRAM is 512 slots in RTL. Producing every head's decay at once needs +# num_heads * key_dim -- 12,288 for Kimi K3 -- so the mixer streams one head at +# a time and reuses the window. These pin that the streamed result is identical +# to the all-at-once one. +# --------------------------------------------------------------------------- + + +class _PerHeadHarness(_GateHarness): + """Same tiles, but decay_fp sized for one head rather than all of them.""" + + def __init__(self, shape: KdaShape): + super().__init__(shape) + self.one_head_decay = self.prog.fp_var("one_head_decay", size=shape.key_dim) + self.mark = len(self.prog.get_code()) + + def run_head(self, head: int, *, gate, dt_bias, a_log): + s = self.shape + self.prog.kda_decay_scalars_v0( + gate=self.gate, dt_bias=self.dt_bias, rate_fp=self.rate, + lower_bound_fp=self.lower, decay_fp=self.one_head_decay, + consts=self.consts, shape=s, heads=[head], + ) + code = self.prog.get_code()[self.mark :] + self.mark = len(self.prog.get_code()) + + m = Machine(vlen=MLEN) + kb = kda_key_blocks(s, MLEN) + for h in range(s.num_heads): + for b in range(kb): + lanes = slice(b * MLEN, (b + 1) * MLEN) + row = kda_key_row(s, MLEN, h, b) + m.write_vram_row(self._base(self.gate) + row * MLEN, gate[h, lanes].tolist()) + m.write_vram_row( + self._base(self.dt_bias) + row * MLEN, dt_bias[h, lanes].tolist() + ) + m.write_fpram(self.consts.zero.address, PlenaCompiler.kda_fp_constant_values()) + m.write_fpram(self.rate.address, torch.exp(a_log).tolist()) + m.write_fpram(self.lower.address, [s.gate_lower_bound]) + m.run(code) + return torch.tensor( + [m.fpram[self.one_head_decay.address + k] for k in range(s.key_dim)], + dtype=torch.float32, + ) + + +@pytest.mark.parametrize("heads,key_dim", [(3, MLEN), (4, MLEN * 2), (2, MLEN * 3)]) +def test_one_head_at_a_time_matches_all_at_once(heads, key_dim): + """The streamed window holds key_dim slots instead of num_heads*key_dim, + and the scalars land at offset 0 because tile_row_to_fpram addresses by + position in `rows`, not by row index.""" + shape, c = _case(21, heads, key_dim) + for h in range(heads): + got = _PerHeadHarness(shape).run_head( + h, gate=c["gate"], dt_bias=c["dt_bias"], a_log=c["a_log"] + ) + torch.testing.assert_close( + got, c["expected_decay"][h], rtol=3e-5, atol=3e-6 + ) + + +def test_streaming_uses_the_selected_head_rate_not_its_position(): + """rate_fp is indexed by head number. Using the position within `heads` + would silently give every streamed head rate[0].""" + shape, c = _case(22, 4, MLEN * 2) + last = shape.num_heads - 1 + got = _PerHeadHarness(shape).run_head( + last, gate=c["gate"], dt_bias=c["dt_bias"], a_log=c["a_log"] + ) + torch.testing.assert_close(got, c["expected_decay"][last], rtol=3e-5, atol=3e-6) + + # head 0's rate must give a different answer, or the test proves nothing + with_rate0 = torch.exp( + activate_log_decay( + c["gate"][None, last : last + 1, :], + c["a_log"][:1], + c["dt_bias"][last : last + 1], + lower_bound=shape.gate_lower_bound, + ) + )[0, 0] + assert not torch.allclose(got, with_rate0, rtol=1e-3, atol=1e-3) + + +def test_streaming_rejects_a_head_out_of_range(): + shape, _ = _case(23, 2, MLEN) + h = _PerHeadHarness(shape) + with pytest.raises(ValueError, match="out of range"): + h.prog.kda_decay_scalars_v0( + gate=h.gate, dt_bias=h.dt_bias, rate_fp=h.rate, lower_bound_fp=h.lower, + decay_fp=h.one_head_decay, consts=h.consts, shape=shape, heads=[2], + ) + + +def test_a_head_subset_bounds_the_tile_by_row_number_not_by_count(): + """With heads=[h] the rows are a slice out of the middle, so the row count + is smaller than the highest row number. Checking the count accepts a tile + the emitter then reads and writes past the end of -- into whatever VRAM + object was allocated next, with no error anywhere in the stack.""" + shape = KdaShape( + hidden_size=32, num_heads=16, key_dim=MLEN, value_dim=MLEN, conv_kernel=4 + ) + p = PlenaCompiler(mlen=MLEN, blen=2) + # Room for one head's worth of rows, but head 12 lives at row 12. + gate = p.alloc("gate", MLEN, MLEN) + dt_bias = p.alloc("dt_bias", MLEN, MLEN) + with pytest.raises(ValueError, match="to reach head 12"): + p.kda_decay_scalars_v0( + gate=gate, dt_bias=dt_bias, + rate_fp=p.fp_var("rate", size=shape.num_heads), + lower_bound_fp=p.fp_var("lb", size=1), + decay_fp=p.fp_var("decay", size=MLEN), + consts=p.kda_fp_constants(), shape=shape, heads=[12], + ) diff --git a/aten/tests/test_kda_layer.py b/aten/tests/test_kda_layer.py new file mode 100644 index 0000000..88c0b6c --- /dev/null +++ b/aten/tests/test_kda_layer.py @@ -0,0 +1,942 @@ +"""The KDA layer's projection split, which is a naming operation and not a copy. + +The input projection produces one tile of `3*key_width + value_width + heads` +columns; the mixer wants per-`(head, key block)` tiles of `mlen`. Those look +like different layouts and a naive conversion is a scatter of one row per +column block -- around 700 per layer at Kimi K3's shape, each a full block +operation. + +They are the same bytes. Column block `c` of any VRAM tile sits at +`base + c * mlen * mlen`, so block `c` of `[tokens, N]` *is* the `[tokens, mlen]` +tile at that address. `vram_column_block_view` names one by registering a VRAM +object at the computed address rather than allocating, and emits nothing. + +These tests pin that it emits nothing, that the alias is real, and that the +section boundaries match the order `kda_state_engine_step` splits in -- an +off-by-one-block there hands the mixer `k` where it expects `q`, which is a +finite plausible wrong answer. +""" + +from __future__ import annotations + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.shape import KdaShape # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.instruction_stream import ( # noqa: E402 + dynamic_count, + static_count, +) +from compiler.aten.plena.program_kda_conv import kda_conv_blocks # noqa: E402 +from compiler.aten.plena.program_kda_layer import ( # noqa: E402 + kda_projection_features, + kda_projection_sections, + kda_projection_width, +) +from compiler.aten.tests.isa_interpreter import Machine # noqa: E402 + +MLEN = 8 + + +def _shape(heads=4, key_dim=MLEN, value_dim=MLEN) -> KdaShape: + return KdaShape( + hidden_size=heads * value_dim, num_heads=heads, key_dim=key_dim, + value_dim=value_dim, conv_kernel=4, + ) + + +def _static(asm: str) -> int: + return len([ln for ln in asm.splitlines() if ln.strip() and not ln.strip().startswith(";")]) + + +def test_the_split_emits_nothing(): + """The whole point. A scatter would be ~700 block operations per layer at + Kimi K3; naming a view is free.""" + shape = _shape() + p = PlenaCompiler(mlen=MLEN, blen=2) + projected = p.alloc("projected", MLEN, kda_projection_width(shape, MLEN)) + mark = len(p.get_code()) + p.kda_split_projection_v0(projected=projected, shape=shape, rows=4) + assert _static(p.get_code()[mark:]) == 0 + + +def test_the_sections_follow_the_reference_split_order(): + """`kda_state_engine_step` splits [key, key, value, key, heads] in that + order. Getting it wrong hands the mixer k where it expects q.""" + shape = _shape(heads=4, key_dim=MLEN * 2, value_dim=MLEN) + sections = kda_projection_sections(shape, MLEN) + assert [name for name, _, _ in sections] == ["q", "k", "v", "gate", "beta"] + key_blocks = shape.projection_size // MLEN + value_blocks = shape.num_heads * shape.value_dim // MLEN + assert [count for _, _, count in sections] == [ + key_blocks, key_blocks, value_blocks, key_blocks, 1 + ] + # Contiguous, no gaps and no overlaps. + expected_first = 0 + for _, first, count in sections: + assert first == expected_first + expected_first += count + + +def test_beta_is_padded_to_a_block_and_the_width_says_so(): + """beta is one value per head -- 96 at Kimi K3 against mlen 64 -- so it is + the one section that is not naturally a whole block. The host has to + materialise W_in at the padded width with the padding columns zero; a + narrower weight leaves the tail of beta's block holding whatever preceded + it in HBM.""" + kimi = KdaShape( + hidden_size=7168, num_heads=96, key_dim=128, value_dim=128, conv_kernel=4 + ) + logical = kda_projection_features(kimi) + padded = kda_projection_width(kimi, 64) + assert padded > logical, "beta's 96 heads do not fill a 64-lane block" + assert padded % 64 == 0 + # Exactly the beta padding, nothing else. + assert padded - logical == 64 * 2 - 96 + + +@pytest.mark.parametrize( + "parent_rows,strict", + [ + (MLEN, True), # the only shape the first version tested + (MLEN * 2, True), # two row blocks -- the stride is physical_rows, not mlen + (MLEN * 4, True), # four + (2, False), # padded: alloc pads rows to blen, not to mlen + ], +) +def test_a_view_lands_on_its_parents_column_block(parent_rows, strict): + """The view's address must equal the parent's own block address. + + Column blocks are strided by `physical_rows * mlen`, not `mlen * mlen`; + those coincide only when a tile is exactly mlen rows tall, which is the one + shape this was first tested at. Both the right formula and a wrong one + passed, because the test seeded and read at the view's *own* address -- it + was self-consistent with any address the view invented. + """ + shape = _shape() + p = PlenaCompiler(mlen=MLEN, blen=2) + width = kda_projection_width(shape, MLEN) + projected = p.alloc("projected", parent_rows, width, strict=strict) + n_blocks = p.get_vram_layout(projected.name).physical_shape[1] // MLEN + for block in range(n_blocks): + view = p.vram_column_block_view(projected, block, name=f"v{block}", rows=1) + assert p.get_vram_layout(view.name).vram_base_addr == p._tile_addr( + projected.name, 0, block + ), f"view of block {block} is not where the parent's block {block} is" + + +def test_a_view_aliases_its_parent(): + """Written through the PARENT's address, read through the view. + + Seeding at the view's own address would pass for any address at all; the + point is that the two agree. + """ + shape = _shape() + p = PlenaCompiler(mlen=MLEN, blen=2) + width = kda_projection_width(shape, MLEN) + # Two row blocks, so the stride actually discriminates. + projected = p.alloc("projected", MLEN * 2, width) + views = p.kda_split_projection_v0(projected=projected, shape=shape, rows=MLEN) + block = kda_projection_sections(shape, MLEN)[2][1] + 1 # v's second block + target = views["v"][1] + + parent_addr = p._tile_addr(projected.name, 0, block) + assert p.get_vram_layout(target.name).vram_base_addr == parent_addr + + consts = p.kda_fp_constants() + two = p.fp_var("two", size=1) + mark = len(p.get_code()) + p.tile_row_mul_fp_broadcast(target, two, rows=[0, 1]) + code = p.get_code()[mark:] + + m = Machine(vlen=MLEN, vram_words=1 << 18, fpram_words=1 << 13) + m.write_fpram(consts.zero.address, p.kda_fp_constant_values()) + m.write_fpram(two.address, [2.0]) + seed = [float(i + 1) for i in range(MLEN)] + # Through the parent's address, not the view's. + m.write_vram_row(parent_addr, seed) + m.write_vram_row(parent_addr + MLEN, seed) + # A neighbouring block must not move -- a wrong stride scaled someone else's + # bytes and left these alone. + neighbour = p._tile_addr(projected.name, 0, block + 1) + m.write_vram_row(neighbour, seed) + m.run(code) + + assert m.read_vram_row(parent_addr, MLEN) == [v * 2 for v in seed] + assert m.read_vram_row(neighbour, MLEN) == seed, "the view touched another block" + + +def test_a_view_is_bounded_by_its_parent(): + shape = _shape() + p = PlenaCompiler(mlen=MLEN, blen=2) + projected = p.alloc("projected", MLEN, kda_projection_width(shape, MLEN)) + n_blocks = projected.physical_shape[1] // MLEN + with pytest.raises(ValueError, match="out of range"): + p.vram_column_block_view(projected, n_blocks, name="past_the_end") + phys_rows = p.get_vram_layout(projected.name).physical_shape[0] + # One past, not far past: `rows=MLEN*2` cleared the boundary by so much that + # a permissive off-by-one bound survived. + with pytest.raises(ValueError, match="exceeds"): + p.vram_column_block_view(projected, 0, name="too_tall", rows=phys_rows + 1) + + +def test_every_per_head_tile_is_a_block_boundary(): + """The split is only free because each `(head, key block)` starts on one. + At Kimi K3 that is 96 heads x 2 blocks for each of q, k and gate.""" + kimi = KdaShape( + hidden_size=7168, num_heads=96, key_dim=128, value_dim=128, conv_kernel=4 + ) + mlen = 64 + for head in range(kimi.num_heads): + for block in range(kimi.key_dim // mlen): + feature = head * kimi.key_dim + block * mlen + assert feature % mlen == 0, ( + f"head {head} block {block} starts at feature {feature}, " + f"which is not a block boundary" + ) + + +def test_a_duplicate_view_name_is_refused(): + """`register_vram_matrix` overwrites without complaint and views resolve by + name at emit time, so a duplicate silently repoints every view already + handed out. `prefix` defaults to "kda", so two splits in one program collide + -- and the layer this is for runs 93 times.""" + shape = _shape() + p = PlenaCompiler(mlen=MLEN, blen=2) + width = kda_projection_width(shape, MLEN) + a = p.alloc("proj_a", MLEN, width) + b = p.alloc("proj_b", MLEN, width) + p.kda_split_projection_v0(projected=a, shape=shape, rows=1) + with pytest.raises(ValueError, match="already exists"): + p.kda_split_projection_v0(projected=b, shape=shape, rows=1) + # A distinct prefix is the way through. + p.kda_split_projection_v0(projected=b, shape=shape, rows=1, prefix="layer1") + + +def test_the_split_needs_key_dim_on_a_block_boundary_not_just_key_width(): + """The invariant is per head. 4 heads of key_dim 6 gives key_width 24, + divisible by mlen 8 -- while heads 1, 2 and 3 all start mid-block.""" + bad = KdaShape(hidden_size=32, num_heads=4, key_dim=6, value_dim=MLEN, conv_kernel=4) + p = PlenaCompiler(mlen=MLEN, blen=2) + assert (bad.num_heads * bad.key_dim) % MLEN == 0, "key_width alone looks fine" + projected = p.alloc("projected", MLEN, MLEN * 8) + with pytest.raises(ValueError, match="multiple of mlen"): + p.kda_split_projection_v0(projected=projected, shape=bad, rows=1) + + +def test_a_view_carries_its_parents_physical_shape(): + """So a view of a view addresses correctly. + + Built by hand, a view's var carried the rows asked for while its layout + carried the parent's -- and reading the var gave a nested view the wrong + bounds and, past one row block, the wrong address. Going through `alloc_at` + with an explicit physical_shape makes the two agree by construction. + """ + p = PlenaCompiler(mlen=MLEN, blen=2) + parent = p.alloc("parent", MLEN * 2, MLEN * 4) + view = p.vram_column_block_view(parent, 1, name="v1", rows=MLEN) + assert view.physical_shape == p.get_vram_layout(view.name).physical_shape + + nested = p.vram_column_block_view(view, 0, name="v2", rows=MLEN * 2) + assert p.get_vram_layout(nested.name).vram_base_addr == p._tile_addr( + parent.name, 0, 1 + ) + + +def test_the_tall_view_makes_a_projection_one_dense_tile(): + """Column block `c` is at `base + c*physical_rows*mlen` and rows within a + block are linear, so the whole tile already *is* a dense + `[blocks*physical_rows, mlen]` tile. Feature block `c`'s token `t` is at row + `c*stride + t` -- which is what lets a gather be one hardware loop.""" + p = PlenaCompiler(mlen=MLEN, blen=2) + for rows, strict in ((MLEN, True), (MLEN * 2, True), (2, False)): + parent = p.alloc(f"p{rows}{strict}", rows, MLEN * 4, strict=strict) + tall, stride = p.vram_tall_view(parent, name=f"t{rows}{strict}") + phys_rows = p.get_vram_layout(parent.name).physical_shape[0] + assert stride == phys_rows + assert tall.shape[0] == 4 * phys_rows + for c in range(4): + # Row c*stride of the tall view is block c's row 0. + want = p._tile_addr(parent.name, 0, c) + got = p.get_vram_layout(tall.name).vram_base_addr + c * stride * MLEN + assert got == want, f"block {c}: tall view says {got}, parent says {want}" + + +def test_the_gather_is_one_hardware_loop_per_section(): + """14 static instructions per section, independent of block count. A copy + per block would be one loop scaffold each -- 10 instructions per block, so + 7,674 for a Kimi K3 layer against about 70.""" + shape = _shape(heads=8, key_dim=MLEN * 2) + p = PlenaCompiler(mlen=MLEN, blen=2) + consts = p.kda_fp_constants() + projected = p.alloc("projected", MLEN, kda_projection_width(shape, MLEN)) + counts = {} + for section, (first, count) in { + n: (f, c) for n, f, c in kda_projection_sections(shape, MLEN) + }.items(): + dst = p.alloc(f"dst_{section}", max(count, MLEN), MLEN, strict=False) + mark = len(p.get_code()) + p.kda_gather_projection_v0( + projected=projected, dst=dst, section=section, shape=shape, + consts=consts, name=f"tall_{section}", + ) + code = p.get_code()[mark:] + counts[section] = (count, _static(code), code.count("V_FMA_VF")) + for section, (blocks, static, fmas) in counts.items(): + assert fmas == 1, f"{section} spans {blocks} blocks but emitted {fmas} FMAs" + assert static < 30, f"{section}: {static} instructions for {blocks} blocks" + # q spans 16 blocks and beta 1. They differ by the loop's two step + # increments, which a single-block sweep skips -- not by block count. + assert counts["q"][1] - counts["beta"][1] <= 2, counts + assert counts["q"][0] == 16 and counts["beta"][0] == 1, counts + + +def test_the_gather_cost_does_not_grow_with_head_count(): + """The property, stated directly: 96 heads must cost what 2 heads cost.""" + costs = {} + for heads in (2, 8, 96): + shape = _shape(heads=heads, key_dim=MLEN * 2) + p = PlenaCompiler(mlen=MLEN, blen=2) + consts = p.kda_fp_constants() + projected = p.alloc("projected", MLEN, kda_projection_width(shape, MLEN)) + first, count = {n: (f, c) for n, f, c in + kda_projection_sections(shape, MLEN)}["q"] + dst = p.alloc("dst", max(count, MLEN), MLEN, strict=False) + mark = len(p.get_code()) + p.kda_gather_projection_v0( + projected=projected, dst=dst, section="q", shape=shape, consts=consts, + ) + costs[heads] = _static(p.get_code()[mark:]) + assert len(set(costs.values())) == 1, f"gather cost grows with heads: {costs}" + + +def test_the_gather_moves_the_right_features(): + """Section boundaries and the tall view's stride, checked against values. + + An off-by-one block here hands the mixer `k` where it expects `q`. + """ + import torch + + shape = _shape(heads=2, key_dim=MLEN * 2) + p = PlenaCompiler(mlen=MLEN, blen=2) + consts = p.kda_fp_constants() + width = kda_projection_width(shape, MLEN) + projected = p.alloc("projected", MLEN, width) + sections = {n: (f, c) for n, f, c in kda_projection_sections(shape, MLEN)} + dsts = { + n: p.alloc(f"dst_{n}", max(c, MLEN), MLEN, strict=False) + for n, (_, c) in sections.items() + } + mark = len(p.get_code()) + for section in sections: + p.kda_gather_projection_v0( + projected=projected, dst=dsts[section], section=section, + shape=shape, consts=consts, name=f"tall_{section}", + ) + code = p.get_code()[mark:] + + m = Machine(vlen=MLEN, vram_words=1 << 18, fpram_words=1 << 13) + m.write_fpram(consts.zero.address, p.kda_fp_constant_values()) + n_blocks = p.get_vram_layout(projected.name).physical_shape[1] // MLEN + # Block b's token 0 is marked with b, so a misrouted gather is obvious. + for b in range(n_blocks): + m.write_vram_row(p._tile_addr(projected.name, 0, b), [float(b)] * MLEN) + m.run(code) + + for section, (first, count) in sections.items(): + base = p.get_vram_layout(dsts[section].name).vram_base_addr + got = [m.read_vram_row(base + r * MLEN, 1)[0] for r in range(count)] + want = [float(first + c) for c in range(count)] + assert got == want, f"{section}: gathered blocks {got}, expected {want}" + + +# --------------------------------------------------------------------------- +# The assembled layer, against the reference's own boundary. + + +def _layer_case(seed: int, heads: int, key_dim: int, value_dim: int, kernel: int = 4): + """Everything `kda_state_engine_step` needs, plus the compiled equivalent.""" + import torch + + from compiler.aten.models.kda.reference import ( + KdaConvWeights, + KdaRecurrentState, + kda_state_engine_step, + ) + from compiler.aten.plena.program_kda_common import ( + 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_mixer import KdaMixerBuffers + + torch.manual_seed(seed) + shape = KdaShape( + hidden_size=heads * value_dim, num_heads=heads, key_dim=key_dim, + value_dim=value_dim, conv_kernel=kernel, + ) + key_width = shape.projection_size + value_width = heads * value_dim + projected = torch.randn(1, kda_projection_features(shape)) * 0.5 + conv_w = { + n: torch.randn(w, kernel) * 0.3 + for n, w in (("q", key_width), ("k", key_width), ("v", value_width)) + } + # `kernel` taps, not kernel-1: both the reference and the emitter roll the + # window and overwrite the last with the new token, so the state carries the + # full window. + conv_hist = torch.randn(1, key_width * 2 + value_width, kernel) * 0.3 + state = KdaRecurrentState( + torch.randn(1, heads, value_dim, key_dim) * 0.1, conv_hist.clone() + ) + a_log = torch.randn(heads) * 0.5 + dt_bias = torch.randn(heads, key_dim) * 0.5 + out_ref, state_ref = kda_state_engine_step( + projected, state, + KdaConvWeights(conv_w["q"], conv_w["k"], conv_w["v"], None, None, None), + a_log, dt_bias, shape, state_storage="fp32", + ) + return dict( + shape=shape, projected=projected, conv_w=conv_w, conv_hist=conv_hist, + state=state, a_log=a_log, dt_bias=dt_bias, + out_ref=out_ref, state_ref=state_ref, + key_width=key_width, value_width=value_width, + ) + + +def test_the_assembled_layer_matches_kda_state_engine_step(): + """Gather, three convolutions, the gates and the recurrence, against the + reference's own boundary. + + This is what an earlier version of this file said could not be built without + changing two emitters. It needed neither: one gather over a tall view, and + everything downstream is as it was. + """ + import torch + + 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_mixer import KdaMixerBuffers + + c = _layer_case(3, heads=3, key_dim=MLEN * 2, value_dim=MLEN) + shape = c["shape"] + p = PlenaCompiler(mlen=MLEN, blen=2) + consts = p.kda_fp_constants() + kb = kda_key_blocks(shape, MLEN) + up = lambda n: ((n + MLEN - 1) // MLEN) * MLEN # noqa: E731 + a = lambda n, r: p.alloc(n, up(r), MLEN, strict=False) # noqa: E731 + + # The projection output, allocated at blen rows -- the whole point of the + # stride being physical_rows rather than mlen. + projected = p.alloc("projected", p.blen, kda_projection_width(shape, MLEN), + strict=False) + sections = {n: (f, cnt) for n, f, cnt in kda_projection_sections(shape, MLEN)} + gathered = { + n: a(f"g_{n}", max(cnt, MLEN)) for n, (_, cnt) in sections.items() + } + widths = {"q": c["key_width"], "k": c["key_width"], "v": c["value_width"]} + conv_state = {n: a(f"cs_{n}", kda_conv_blocks(w, MLEN) * shape.conv_kernel) + for n, w in widths.items()} + conv_weight = {n: a(f"cw_{n}", kda_conv_blocks(w, MLEN) * shape.conv_kernel) + for n, w in widths.items()} + conv_scratch = a("conv_scratch", kda_conv_blocks(max(widths.values()), MLEN)) + + decay = p.fp_var("decay_and_q_hat", 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=gathered["gate"], dt_bias=a("dt_bias", shape.num_heads * kb), + beta_logit=gathered["beta"], + state=a("state", kda_state_rows(shape, MLEN)), + out=a("out", kda_vector_rows(shape, MLEN)), + pred=a("pred", kda_vector_rows(shape, MLEN)), + err=a("err", kda_vector_rows(shape, MLEN)), + sq_scratch=a("sq_scratch", shape.num_heads * kb), + decay_fp=decay, q_hat_fp=decay, + k_hat_fp=p.fp_var("k_hat", size=shape.key_dim), + beta_fp=p.fp_var("beta", size=kda_head_blocks(shape, MLEN) * MLEN), + part_fp=p.fp_var("part", size=kb), acc_fp=p.fp_var("acc", size=1), + output_scale_fp=p.fp_var("output_scale", size=1), + rate_fp=p.fp_var("rate", size=shape.num_heads), + lower_bound_fp=p.fp_var("lower_bound", size=1), consts=consts, + ) + mark = len(p.get_code()) + p.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, + ) + code = p.get_code()[mark:] + + m = Machine(vlen=MLEN, vram_words=1 << 20, fpram_words=1 << 14) + m.write_fpram(consts.zero.address, p.kda_fp_constant_values()) + m.write_fpram(buffers.output_scale_fp.address, [1.0 / shape.key_dim ** 0.5]) + m.write_fpram(buffers.rate_fp.address, torch.exp(c["a_log"]).tolist()) + m.write_fpram(buffers.lower_bound_fp.address, [shape.gate_lower_bound]) + + # The projection output: feature f at column block f//mlen, row 0. + flat = c["projected"][0] + for blk in range(-(-flat.numel() // MLEN)): + chunk = flat[blk * MLEN : (blk + 1) * MLEN].tolist() + chunk += [0.0] * (MLEN - len(chunk)) + m.write_vram_row(p._tile_addr(projected.name, 0, blk), chunk) + + # Conv history and weights, blocked by kda_conv_state_row. + hist = c["conv_hist"][0] + offsets = {"q": 0, "k": c["key_width"], "v": 2 * c["key_width"]} + for name, width in widths.items(): + blocks = kda_conv_blocks(width, MLEN) + cs_base = p.get_vram_layout(conv_state[name].name).vram_base_addr + cw_base = p.get_vram_layout(conv_weight[name].name).vram_base_addr + for cb in range(blocks): + lo = cb * MLEN + for tap in range(shape.conv_kernel): + row = kda_conv_state_row(width, MLEN, shape.conv_kernel, cb, tap) + vals = hist[offsets[name] + lo : offsets[name] + lo + MLEN, tap] + m.write_vram_row(cs_base + row * MLEN, vals.tolist()) + m.write_vram_row( + cw_base + row * MLEN, + c["conv_w"][name][lo : lo + MLEN, tap].tolist(), + ) + + # dt_bias, and the incoming recurrent state, transposed as decode holds it. + dt_base = p.get_vram_layout(buffers.dt_bias.name).vram_base_addr + for h in range(shape.num_heads): + for b in range(kb): + m.write_vram_row( + dt_base + kda_key_row(shape, MLEN, h, b) * MLEN, + c["dt_bias"][h, b * MLEN : (b + 1) * MLEN].tolist(), + ) + st_base = p.get_vram_layout(buffers.state.name).vram_base_addr + blocks_v = kda_vector_rows(shape, MLEN) // shape.num_heads + for h in range(shape.num_heads): + for blk in range(blocks_v): + for key in range(shape.key_dim): + row = kda_state_row(shape, MLEN, h, blk, key) + vals = c["state"].recurrent[0, h, blk * MLEN : (blk + 1) * MLEN, key] + m.write_vram_row(st_base + row * MLEN, vals.tolist()) + m.run(code) + + out_base = p.get_vram_layout(buffers.out.name).vram_base_addr + got = torch.tensor([ + sum((m.read_vram_row( + out_base + kda_vector_row(shape, MLEN, h, blk) * MLEN, MLEN) + for blk in range(blocks_v)), []) + for h in range(shape.num_heads) + ]).reshape(1, -1) + torch.testing.assert_close(got, c["out_ref"], rtol=5e-4, atol=5e-5) + + +# --------------------------------------------------------------------------- +# Layer composition at depth: the HBM address space, and FPRAM reuse. + + +def _kimi_layers(layers: int): + """Stack `layers` Kimi K3 KDA layers and return the program.""" + from compiler.aten.plena.program_kda_common import kda_state_rows, kda_vector_rows + from compiler.aten.plena.program_kda_conv import kda_conv_blocks + from compiler.aten.plena.program_kda_gates import kda_head_blocks, kda_key_blocks + from compiler.aten.plena.program_kda_mixer import KdaMixerBuffers + + mlen = 64 + shape = KdaShape.kimi_k3() + key_width = shape.projection_size + value_width = shape.num_heads * shape.value_dim + kernel = shape.conv_kernel + kb = kda_key_blocks(shape, mlen) + up = lambda n: ((n + mlen - 1) // mlen) * mlen # noqa: E731 + counts = {n: c for n, _, c in kda_projection_sections(shape, mlen)} + + p = PlenaCompiler(mlen=mlen, blen=4) + # Bypasses kda_require_state_precision_v0, which a real lowering must call + # -- under an MX KV type the state decodes as e4m3 while the address stride + # assumes 2 bytes. Acceptable for a structural measurement; it does mean the + # program measured here could not be executed as configured. + p._bf16_kv_checked = True + consts = p.kda_fp_constants() + a = lambda n, r: p.alloc(n, up(r), mlen, strict=False) # noqa: E731 + # One FPRAM window for the whole model. It is scratch, not storage. + shared = { + k: p.fp_var(k, size=s) + for k, s in ( + ("decay", shape.key_dim), ("k_hat", shape.key_dim), + ("beta", kda_head_blocks(shape, mlen) * mlen), ("part", kb), + ("acc", 1), ("os", 1), ("rate", shape.num_heads), ("lb", 1), + ) + } + widths = {"q": key_width, "k": key_width, "v": value_width} + for i in range(layers): + projected = p.load_batch( + p.input(f"proj{i}", (p.blen, kda_projection_width(shape, mlen))), + name=f"projv{i}", + ) + gathered = {n: a(f"g{i}_{n}", max(c, mlen)) for n, c in counts.items()} + conv_state = { + n: p.load_batch( + p.input(f"cs{i}_{n}", (kda_conv_blocks(w, mlen) * kernel, mlen)), + name=f"csv{i}_{n}", + ) + for n, w in widths.items() + } + conv_weight = { + n: p.load_batch( + p.input(f"cw{i}_{n}", (kda_conv_blocks(w, mlen) * kernel, mlen)), + name=f"cwv{i}_{n}", + ) + for n, w in widths.items() + } + # The state, dt_bias and the projection come from HBM, exactly as + # `kda_stage_test.py::case_layer` stages them. Allocating them in VRAM + # instead made the HBM figure count only the conv tensors -- one seventh + # of a layer's real traffic -- and the two tests then disagreed about + # where a layer's state lives. + state_v = p.load_batch( + p.input(f"st{i}", (up(kda_state_rows(shape, mlen)), mlen)), name=f"stv{i}" + ) + dtb_v = p.load_batch( + p.input(f"dtb{i}", (up(shape.num_heads * kb), mlen)), name=f"dtbv{i}" + ) + buffers = KdaMixerBuffers( + q=a(f"q{i}", shape.num_heads * kb), k=a(f"k{i}", shape.num_heads * kb), + v=a(f"v{i}", kda_vector_rows(shape, mlen)), gate=gathered["gate"], + dt_bias=dtb_v, beta_logit=gathered["beta"], + state=state_v, + out=a(f"o{i}", kda_vector_rows(shape, mlen)), + pred=a(f"pr{i}", kda_vector_rows(shape, mlen)), + err=a(f"er{i}", kda_vector_rows(shape, mlen)), + sq_scratch=a(f"sq{i}", shape.num_heads * kb), + decay_fp=shared["decay"], q_hat_fp=shared["decay"], + k_hat_fp=shared["k_hat"], beta_fp=shared["beta"], + part_fp=shared["part"], acc_fp=shared["acc"], + output_scale_fp=shared["os"], rate_fp=shared["rate"], + lower_bound_fp=shared["lb"], consts=consts, + ) + p.kda_layer_from_projected_v0( + projected=projected, gathered=gathered, conv_state=conv_state, + conv_weight=conv_weight, conv_bias={}, + conv_scratch=a(f"csc{i}", kda_conv_blocks(max(widths.values()), mlen)), + mixer_buffers=buffers, shape=shape, + ) + return p + + +def test_layers_compose_without_name_collisions(): + """Stacking layers must need nothing from the caller. + + Every view the layer takes is named after its projection tile, which is + unique per layer. With a fixed prefix the second layer collided -- and + because `register_vram_matrix` overwrites silently and views resolve by name + at emit time, without the duplicate guard it would have repointed layer 1's + views at layer 2's projection instead of raising. + """ + _kimi_layers(4) + + +#: Measured 2026-08-26. Everything a layer stages -- conv history and weights, +#: the recurrent state, dt_bias and the projection tile -- excluding only the +#: projection *weights*, which any implementation needs and which are +#: model-inherent. +#: +#: Gated because the emulator preloads HBM from a flat file starting at offset 0, +#: so the allocation has to span every address the program touches. A layout that +#: spreads its regions over a wide address space needs an allocation and a file +#: as large as that span, whatever the live data comes to, and stops being +#: executable long before it stops being describable. Emulator HBM_SIZE is +#: 16 GiB. +#: +#: An earlier version of this counted only the conv tensors, because the helper +#: put the state and the projection in VRAM while the emulator test staged them +#: from HBM. That reported 324 KiB per layer against a real 2,282. +KIMI_HBM_PER_LAYER_MAX = 2560 * 1024 +KIMI_LAYERS_FULL = 93 + + +def test_the_hbm_footprint_does_not_repeat_the_descriptor_blowup(): + one = _kimi_layers(1)._next_hbm_addr + two = _kimi_layers(2)._next_hbm_addr + per_layer = two - one + full = one + (KIMI_LAYERS_FULL - 1) * per_layer + print( + f"KIMI_K3_HBM per_layer={per_layer} full_{KIMI_LAYERS_FULL}={full} " + f"({full / (1 << 30):.4f} GiB)" + ) + assert per_layer <= KIMI_HBM_PER_LAYER_MAX, ( + f"{per_layer} bytes per layer exceeds {KIMI_HBM_PER_LAYER_MAX}" + ) + assert full < (1 << 30), ( + f"{full / (1 << 30):.2f} GiB at {KIMI_LAYERS_FULL} layers, against a 1 GiB " + f"bar. The emulator is configured for 16 GiB and preloads it from a flat " + f"file, so the whole span has to be allocated and written -- a layout that " + f"crosses the configured size stops being executable at all" + ) + + +def test_fpram_fits_the_hardware_file_at_full_depth(): + """The whole model's FPRAM must fit `FP_SRAM_DEPTH`, which is **512**. + + The compiler's own `FPRAMAllocator` defaults to 1024 and so will not catch + an overflow -- that disagreement with the SystemVerilog is why this asserts + the hardware number directly. + + Sharing the window across layers is the *caller's* decision, and this helper + makes it; no KDA emitter allocates FPRAM at all (`grep fp_var + aten/plena/program_kda_*.py` is empty, and a mutation that adds a per-call + `fp_var` turns this red). So what is pinned here is two things: that the + emitters stay caller-allocated, and that one shared window is enough for a + 93-layer model. A caller that allocates per layer overflows 1024 by the + fourth layer and 512 by the second; that is a mistake this test cannot make + on their behalf, but the number it asserts is the one they have to meet. + """ + used = {n: _kimi_layers(n).fpram_allocator.next_free for n in (1, 2, 4)} + assert len(set(used.values())) == 1, ( + f"FPRAM use grows with depth: {used}. It is scratch, not storage." + ) + full = _kimi_layers(KIMI_LAYERS_FULL).fpram_allocator.next_free + print(f"KIMI_K3_FPRAM full_{KIMI_LAYERS_FULL}={full}") + assert full == used[1], "93 layers must use exactly what one layer uses" + assert full <= 512, ( + f"{full} FPRAM slots against FP_SRAM_DEPTH 512 " + f"(PLENA_RTL/src/definitions/configuration.svh)" + ) + + +# --------------------------------------------------------------------------- +# The gather is not forced by packing, and it is not forced at all. + + +def _conv_reading_projection(direct: bool, mlen: int = 64): + """Kimi K3's three convolutions, fed by gather-then-conv or by stride.""" + shape = KdaShape.kimi_k3() + sections = {n: (f, c) for n, f, c in kda_projection_sections(shape, mlen)} + widths = { + "q": shape.projection_size, "k": shape.projection_size, + "v": shape.num_heads * shape.value_dim, + } + p = PlenaCompiler(mlen=mlen, blen=4) + consts = p.kda_fp_constants() + projected = p.alloc( + "proj", p.blen, kda_projection_width(shape, mlen), strict=False + ) + up = lambda n: ((n + mlen - 1) // mlen) * mlen # noqa: E731 + a = lambda n, r: p.alloc(n, up(r), mlen, strict=False) # noqa: E731 + mark = len(p.get_code()) + for name, width in widths.items(): + first, count = sections[name] + blocks = kda_conv_blocks(width, mlen) + common = dict( + conv_state=a(f"cs_{name}", blocks * shape.conv_kernel), + weight=a(f"w_{name}", blocks * shape.conv_kernel), bias=None, + out=a(f"o_{name}", blocks), scratch=a(f"sc_{name}", blocks), + consts=consts, channels=width, kernel=shape.conv_kernel, + ) + if direct: + tall, stride = p.vram_tall_view(projected, name=f"tall_{name}") + p.kda_conv_step_v0( + x_new=tall, x_new_row_base=first * stride, + x_new_row_stride=stride, **common, + ) + else: + dst = a(f"g_{name}", max(count, mlen)) + p.kda_gather_projection_v0( + projected=projected, dst=dst, section=name, shape=shape, + consts=consts, name=f"t_{name}", + ) + p.kda_conv_step_v0(x_new=dst, **common) + return p.get_code()[mark:] + + +def test_reading_the_projection_at_a_stride_removes_the_gather(): + """The gather is overhead, and the consumer can simply read across it. + + The projection cannot produce a dense tile. `M_MM_WO` writes a + `blen x blen` sub-tile and the writeback loops cover `mlen / blen` column + groups, so the smallest thing it can lay down is `blen` token-rows by + `mlen` lanes: column block `c` is at row `c * blen`. The convolution wants + one token's blocks as consecutive rows, and the mismatch is exactly `blen`. + + That mismatch is a property of the matrix writeback, not of how the weight + matrices are packed -- which is why splitting the packed projection into + five separate ones changes nothing (`test_separate_projections_do_not_help` + measures that). + + Reading across it is free. `mamba_row_copy` takes a row index and the + sweeps underneath collapse any arithmetic progression into one hardware + loop whose step is an `S_ADDI_INT` immediate, so a step of `blen` costs + what a step of 1 costs. Feeding the convolution the projection tile with + `x_new_row_stride=blen` therefore does the gather's job for nothing. + """ + gathered = _conv_reading_projection(direct=False) + direct = _conv_reading_projection(direct=True) + saved = dynamic_count(gathered) - dynamic_count(direct) + print( + f"KDA_CONV_FEED gather+conv static={static_count(gathered)} " + f"dynamic={dynamic_count(gathered)} | stride static={static_count(direct)} " + f"dynamic={dynamic_count(direct)} | saved={saved}" + ) + assert saved > 3_500, ( + f"reading at a stride saved only {saved} issued instructions; the three " + f"q/k/v gathers were 3,474 plus their zero-fills when measured" + ) + assert dynamic_count(direct) < dynamic_count(gathered) + assert static_count(direct) < static_count(gathered) + + +def test_the_stride_reads_exactly_the_rows_the_gather_would_have(): + """Equivalence, without running anything: the source rows must match. + + `kda_gather_projection_v0` copies tall row `(first + c) * stride + token` + into dense row `c`, and the convolution then reads dense row `cb`. Reading + the tall view at `base = first * stride`, step `stride`, visits + `(first + cb) * stride` -- the same rows, for token 0. The emitted + `VRAM Matrix Add` comments carry the row ranges, so this is checkable on + the assembly rather than by argument. + """ + mlen, blen, kernel = 8, 2, 2 + channels = 4 * mlen + first, stride = 3, blen + p = PlenaCompiler(mlen=mlen, blen=blen) + consts = p.kda_fp_constants() + a = lambda n, r: p.alloc(n, r, mlen, strict=False) # noqa: E731 + blocks = kda_conv_blocks(channels, mlen) + tall = a("tall", (first + blocks) * stride) + common = dict( + conv_state=a("cs", blocks * kernel), weight=a("w", blocks * kernel), + bias=None, out=a("o", blocks), scratch=a("sc", blocks), consts=consts, + channels=channels, kernel=kernel, apply_silu=False, + ) + mark = len(p.get_code()) + p.kda_conv_step_v0( + x_new=tall, x_new_row_base=first * stride, x_new_row_stride=stride, + **common, + ) + reads = [ + ln for ln in p.get_code()[mark:].splitlines() if "+= tall[" in ln + ] + want = [f"tall[{(first + cb) * stride}:{(first + cb) * stride + 1}]" + for cb in range(blocks)] + got = [ln.split("+= ")[1].split(" ")[0] for ln in reads] + print(f"KDA_CONV_FEED source rows {got}") + assert got == want, f"read {got}, the gather would have read {want}" + + +def test_separate_projections_do_not_help(): + """Splitting the packed projection into five changes neither cost. + + The proposal is that the gather exists because q, k, v, gate and beta share + one output tile, and that projecting them separately -- each writing + straight to where its consumer wants it -- would remove it. Measured at + Kimi K3, `mlen` 64, `blen` 4, with the projection lowered through + `linear_projection` both ways: + + packed matmul 1,090,417 static / 13,876,267 issued, 21,560 M_MM + separate matmul 1,107,395 static / 13,893,245 issued, 21,560 M_MM + gather 70 static / 4,650 issued, both + + Identical matmul work, 0.12% more instructions for the separate form from + the extra per-projection setup, and the same gather. It cannot help, + because the gather is forced by the `blen`-row writeback granularity rather + than by the sharing -- see `kda_conv_step_v0`'s docstring. + + This test pins the part that is cheap to recompute: the section count and + total block count are unchanged by splitting, which is what makes the two + matmul costs equal. The full build is 14M instructions and too slow to run + in CI; the numbers above were measured on 2026-08-28. + """ + mlen = 64 + shape = KdaShape.kimi_k3() + sections = kda_projection_sections(shape, mlen) + packed_blocks = sum(c for _, _, c in sections) + separate_blocks = sum( + max(1, -(-w // mlen)) for w in ( + shape.projection_size, shape.projection_size, + shape.num_heads * shape.value_dim, shape.projection_size, + shape.num_heads, + ) + ) + print(f"KDA_PROJ packed_blocks={packed_blocks} separate_blocks={separate_blocks}") + assert packed_blocks == separate_blocks, ( + f"packing changed the output block count ({packed_blocks} vs " + f"{separate_blocks}); the two lowerings would then not be equal-cost " + f"and the measurement above would need redoing" + ) + + +def test_what_the_gate_and_beta_gathers_would_cost_to_remove(): + """The remaining 1,176 instructions, and why they stay. + + `kda_conv_step_v0` reading at a stride removed q, k and v. Gate and beta + reach `kda_decay_scalars_v0` and `kda_beta_scalars_v0` instead, and those + two are not symmetric: + + * **beta** is only single-tile work -- `_kda_sigmoid_inplace` then + `tile_row_to_fpram`, both taking an explicit row list, and + `tile_row_to_fpram` addresses FPRAM by *position* in that list rather + than by row index. A stride would drop straight in. It is worth 18 + instructions. + * **gate** is not. `kda_decay_scalars_v0` opens with + `tile_row_add(gate, dt_bias, rows=rows)`, 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. It is worth 1,158. + + Together **1,176 issued instructions, 0.24% of the state-engine path**, for + a change to a row-op emitter that every kernel on the branch shares. The + other way round -- laying `dt_bias` out at the same stride so one row list + still serves both -- needs no ISA change but pads it from 192 rows to 768, + which is 74 KiB more per layer and 6.8 MiB more HBM traffic per token + across 93 layers, to save 1,176 instructions. Both trades are bad. + + So they stay, and this test records the price rather than the intention. + If the row-op emitter ever grows `(dst, src)` pairs for another reason, + gate becomes free and this should be revisited. + """ + mlen = 64 + shape = KdaShape.kimi_k3() + p = PlenaCompiler(mlen=mlen, blen=4) + consts = p.kda_fp_constants() + projected = p.alloc( + "proj", p.blen, kda_projection_width(shape, mlen), strict=False + ) + remaining = 0 + for name, _first, count in kda_projection_sections(shape, mlen): + if name not in ("gate", "beta"): + continue + dst = p.alloc(f"d_{name}", max(count, mlen), mlen, strict=False) + mark = len(p.get_code()) + p.kda_gather_projection_v0( + projected=projected, dst=dst, section=name, shape=shape, + consts=consts, name=f"t_{name}", + ) + remaining += dynamic_count(p.get_code()[mark:]) + share = remaining / 492_681 + print(f"KDA_GATHER_REMAINING dynamic={remaining} share={share:.3%}") + assert remaining < 1_500, ( + f"the gate and beta gathers are {remaining} issued instructions; they " + f"were 1,176, and a rise means a section grew or the sweep unrolled" + ) + assert share < 0.005, ( + f"at {share:.3%} of the state-engine path the trade above should be " + f"re-examined; it was 0.24% when the decision was made" + ) diff --git a/aten/tests/test_kda_lowering.py b/aten/tests/test_kda_lowering.py new file mode 100644 index 0000000..0420651 --- /dev/null +++ b/aten/tests/test_kda_lowering.py @@ -0,0 +1,461 @@ +"""KDA lowering: state layout, data movement, and normalisation. + +Numerical validation against the CPU reference lives in the emulator testbench; +what these tests can check without an emulator is that the emitted ISA has the +shape the design requires -- the transposed state tile, load/store agreeing on +precision, and normalisation using the ISA's scalar sqrt rather than a vector op +that does not exist. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.reference import KdaShape # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_kda_common import ( # noqa: E402 + kda_blocks, + kda_conv_blocks, + kda_conv_state_row, + kda_conv_channels, + kda_head_base, + kda_state_row, + kda_state_rows, + kda_vector_row, + kda_vector_rows, +) + +MLEN = 8 + +#: HBM bytes per element for recurrent state. BF16, because the state travels +#: through the keyvalue precision class and Plain BF16 is the widest that class +#: offers; see program_kda_common's module docstring. +STATE_BYTES_PER_ELEMENT = 2 + + +def _vram_base(prog, var) -> int: + """Byte address of a VRAM tile, via the compiler's own symbol table.""" + return prog.get_vram_layout(var.name).vram_base_addr + + +def _prog() -> PlenaCompiler: + return PlenaCompiler(mlen=MLEN, blen=2) + + +def _shape(**kw) -> KdaShape: + base = dict( + hidden_size=MLEN * 4, + num_heads=2, + key_dim=4, + value_dim=MLEN, + conv_kernel=4, + ) + base.update(kw) + return KdaShape(**base) + + +def _body(code: str) -> list[str]: + return [ + line.strip() + for line in code.splitlines() + if line.strip() and not line.strip().startswith((";", "//")) + ] + + +def _trace(code: str) -> list[tuple[str, list[str], dict[str, int]]]: + """Emitted instructions with GP registers resolved to their last constant. + + Opcode-presence assertions are close to worthless here: `V_MUL_VF` appears + because `vram_fill_zero` zeroes a tile, `V_RED_SUM` appears as long as any + reduction survives. What distinguishes a correct emitter from a broken one + is *which addresses* the vector ops read and write, and those arrive through + `S_ADDI_INT gpN, gp0, `. This resolves that much of the dataflow -- + enough to pin operands, not enough to be a simulator. + """ + regs: dict[str, int] = {"gp0": 0} + out: list[tuple[str, list[str], dict[str, int]]] = [] + for line in _body(code): + parts = line.replace(",", " ").split() + op, args = parts[0], parts[1:] + if op == "S_ADDI_INT" and len(args) == 3: + dst, src, imm = args + if src in regs and imm.lstrip("-").isdigit(): + regs[dst] = regs[src] + int(imm) + else: + regs.pop(dst, None) + out.append((op, args, dict(regs))) + return out + + +def _operands(trace, opcode: str) -> list[tuple[int, ...]]: + """Resolved addresses of every occurrence of `opcode`, in order.""" + found = [] + for op, args, regs in trace: + if op != opcode: + continue + vals = tuple(regs[a] for a in args if a in regs) + found.append(vals) + return found + + +# --------------------------------------------------------------------------- +# Layout helpers +# --------------------------------------------------------------------------- + + +def test_state_rows_covers_every_head_block_and_key(): + shape = _shape(num_heads=3, key_dim=5, value_dim=MLEN) + assert kda_blocks(shape, MLEN) == 1 + assert kda_state_rows(shape, MLEN) == 15 + assert [kda_head_base(shape, MLEN, h) for h in range(3)] == [0, 5, 10] + + wide = _shape(num_heads=3, key_dim=5, value_dim=MLEN * 2) + assert kda_blocks(wide, MLEN) == 2 + assert kda_state_rows(wide, MLEN) == 30 + + +def test_flattened_layout_is_a_bijection(): + """Every (head, block, key) gets its own row, and together they tile the + whole state exactly -- no overlap, no gap. Overlap is one head's sweep + corrupting another's; a gap is a row nothing ever writes.""" + shape = _shape(num_heads=4, key_dim=6, value_dim=MLEN * 3) + blocks = kda_blocks(shape, MLEN) + rows = [ + kda_state_row(shape, MLEN, h, c, j) + for h in range(shape.num_heads) + for c in range(blocks) + for j in range(shape.key_dim) + ] + assert len(rows) == len(set(rows)), "two (head, block, key) share a row" + assert set(rows) == set(range(kda_state_rows(shape, MLEN))) + + +def test_row_ordering_is_head_then_block_then_key(): + """Pin the ordering, not merely that it is *some* bijection. + + A bijection test accepts a head/block swap, and so does a unit-stride test. + That matters because this is the HBM byte layout of the pinned state tensor + (kda_load_state_v0 maps row r to offset r*mlen), so once a prefill kernel or + a checkpoint loader writes the initial state it is a cross-component + contract. At value_dim == mlen the swap is the identity, which is why this + has to run at blocks > 1. + """ + shape = _shape(num_heads=3, key_dim=4, value_dim=MLEN * 2) + blocks = kda_blocks(shape, MLEN) + expected = [ + (h * blocks + c) * shape.key_dim + for h in range(shape.num_heads) + for c in range(blocks) + ] + assert [ + kda_state_row(shape, MLEN, h, c, 0) + for h in range(shape.num_heads) + for c in range(blocks) + ] == expected + + # kda_head_base must honour its block argument, not silently return block 0 + assert [ + kda_head_base(shape, MLEN, h, c) + for h in range(shape.num_heads) + for c in range(blocks) + ] == expected + + +def test_keys_are_consecutive_within_a_head_and_block(): + """This is what makes each sweep a single hardware loop: fixing + (head, block) must leave the keys at unit stride.""" + shape = _shape(num_heads=3, key_dim=6, value_dim=MLEN * 2) + for h in range(shape.num_heads): + for c in range(kda_blocks(shape, MLEN)): + rows = [kda_state_row(shape, MLEN, h, c, j) for j in range(shape.key_dim)] + steps = {b - a for a, b in zip(rows, rows[1:])} + assert steps == {1}, f"head {h} block {c} is not unit stride: {rows}" + + +def test_vector_rows_block_the_same_way_as_the_state(): + """q, v, the output and the accumulators are all value-width, so they must + be blocked on the same convention -- a mixture reads block 0 of one tile + against block 1 of another.""" + shape = _shape(num_heads=3, key_dim=4, value_dim=MLEN * 2) + blocks = kda_blocks(shape, MLEN) + assert kda_vector_rows(shape, MLEN) == shape.num_heads * blocks + rows = [ + kda_vector_row(shape, MLEN, h, c) + for h in range(shape.num_heads) + for c in range(blocks) + ] + assert rows == sorted(rows) and len(set(rows)) == len(rows) + + +def test_layout_helpers_reject_out_of_range(): + shape = _shape(num_heads=2, key_dim=3, value_dim=MLEN * 2) + for bad in (dict(head=2), dict(head=-1), dict(block=2), dict(key=3)): + args = dict(head=0, block=0, key=0) | bad + with pytest.raises(ValueError, match="out of range"): + kda_state_row(shape, MLEN, **args) + + +def test_value_dim_must_be_a_whole_number_of_blocks(): + """A partial trailing block leaves lanes past value_dim holding stale data + that the reduction would still sum.""" + shape = _shape(num_heads=1, key_dim=2, value_dim=MLEN + 3) + with pytest.raises(ValueError, match="multiple of mlen"): + kda_blocks(shape, MLEN) + + +def test_conv_channels_matches_the_reference_geometry(): + """2*key + value per head. Cross-checked against KdaShape's own accounting + so the lowering and the reference cannot disagree about conv state size.""" + shape = _shape(num_heads=3, key_dim=5, value_dim=MLEN) + assert kda_conv_channels(shape) == 3 * (2 * 5 + MLEN) + assert kda_conv_channels(shape) * shape.conv_kernel == shape.conv_state_elements + + +# --------------------------------------------------------------------------- +# State residency +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("blocks", [1, 2, 3]) +def test_state_loads_in_the_transposed_layout(blocks): + prog, shape = _prog(), _shape(value_dim=MLEN * blocks) + addr = prog.kda_pin_state_v0("kda_state", shape) + state = prog.kda_load_state_v0("kda_state", shape, addr) + assert state.shape == (kda_state_rows(shape, MLEN), MLEN) + + +def _scale_reg_immediates(code: str) -> list[int]: + """Value written to C_SET_SCALE_REG, which encodes rows*cols*bytes_per_elem.""" + return [regs[args[0]] for op, args, regs in _trace(code) if op == "C_SET_SCALE_REG"] + + +@pytest.mark.parametrize("blocks", [1, 2, 3]) +def test_state_store_mirrors_the_load_precision(blocks): + """A load/store precision mismatch changes the row stride, so the state read + back is not the state written -- a wrong answer, not an error. + + Asserting `stored.hbm_addr == addr` would be a tautology: `store()` builds + the InputVar from the address it was handed. The precision has to be read + out of the emitted ISA, and C_SET_SCALE_REG carries it: its immediate is + rows * value_dim * bytes_per_element on both sides. + """ + prog, shape = _prog(), _shape(value_dim=MLEN * blocks) + addr = prog.kda_pin_state_v0("kda_state", shape) + + mark = len(prog.get_code()) + state = prog.kda_load_state_v0("kda_state", shape, addr) + load_code = prog.get_code()[mark:] + + mark = len(prog.get_code()) + stored = prog.kda_store_state_v0(state, "kda_state", addr) + store_code = prog.get_code()[mark:] + + assert stored.hbm_addr == addr + expected = kda_state_rows(shape, MLEN) * MLEN * STATE_BYTES_PER_ELEMENT + assert _scale_reg_immediates(load_code) == [expected] + assert _scale_reg_immediates(store_code) == [expected] + + # The precision-class operand of the transfer itself: 1 == KeyValue on both. + (prefetch,) = [l for l in _body(load_code) if l.startswith("H_PREFETCH_V")] + (store_v,) = [l for l in _body(store_code) if l.startswith("H_STORE_V")] + assert prefetch.replace(",", " ").split()[-1] == "1" + assert store_v.replace(",", " ").split()[-1] == "1" + + +@pytest.mark.parametrize("blocks", [1, 2, 3]) +def test_state_round_trips_to_the_same_address(blocks): + """Decode reads the state at the top of a step and writes it back at the + bottom; if the two addresses differ the next token reads stale state.""" + prog, shape = _prog(), _shape(value_dim=MLEN * blocks) + addr = prog.kda_pin_state_v0("kda_state", shape) + state = prog.kda_load_state_v0("kda_state", shape, addr) + prog.kda_store_state_v0(state, "kda_state", addr) + reloaded = prog.kda_load_state_v0("kda_state", shape, addr) + assert reloaded.shape == state.shape + base, size = prog.pinned_hbm_region("kda_state") + assert base == addr + assert size == kda_state_rows(shape, MLEN) * MLEN + + +@pytest.mark.parametrize("blocks", [1, 2, 3]) +def test_pinned_state_region_does_not_collide_with_later_tensors(blocks): + """pin_hbm_region must keep the auto-allocator off the state range, or a + later input()/store() silently overwrites the carried state.""" + prog, shape = _prog(), _shape(value_dim=MLEN * blocks) + addr = prog.kda_pin_state_v0("kda_state", shape) + _base, size_elements = prog.pinned_hbm_region("kda_state") + # pinned_hbm_region reports ELEMENTS; hbm_addr is BYTES. Comparing the two + # directly is how an under-reserved region passes this test. + size_bytes = prog.hbm_tensor_size( + size_elements, hbm_element_bytes=STATE_BYTES_PER_ELEMENT + ) + other = prog.input("unrelated", (MLEN, MLEN)) + assert other.hbm_addr >= addr + size_bytes + + +# --------------------------------------------------------------------------- +# Convolution history +# --------------------------------------------------------------------------- + + +def test_conv_roll_shifts_history_and_appends(): + prog, shape = _prog(), _shape(conv_kernel=4) + # VRAM row counts must be a multiple of mlen (memory_state.py:309), so the + # conv_kernel-1 rows of history live in the first rows of an mlen-row tile. + conv = prog.alloc("conv", MLEN, MLEN) + new = prog.alloc("new", MLEN, MLEN) + mark = len(prog.get_code()) + prog.kda_conv_state_roll_v0(conv, new, 0, shape) + code = prog.get_code()[mark:] + assert _body(code), "roll emitted nothing" + + # Counting V_ADD_VV alone cannot tell a correct roll from one that shifts the + # wrong way or appends from the wrong tile -- all three emit 3 copies. Pin + # the (dst, src) row addresses instead. + conv_base = _vram_base(prog, conv) + new_base = _vram_base(prog, new) + history = shape.conv_kernel - 1 + expected = [ + (conv_base + i * MLEN, conv_base + (i + 1) * MLEN) for i in range(history - 1) + ] + [(conv_base + (history - 1) * MLEN, new_base)] + assert _operands(_trace(code), "V_ADD_VV") == [ + (dst, dst, src) for dst, src in expected + ] + + +def test_conv_roll_is_a_noop_for_kernel_one(): + prog, shape = _prog(), _shape(conv_kernel=1) + conv = prog.alloc("conv", MLEN, MLEN) + new = prog.alloc("new", MLEN, MLEN) + before = prog.get_code() + prog.kda_conv_state_roll_v0(conv, new, 0, shape) + assert prog.get_code() == before + + +def test_conv_roll_rejects_an_undersized_history(): + # conv_kernel=10 needs 9 rows of history; the tile below holds mlen=8. + prog, shape = _prog(), _shape(conv_kernel=10) + conv = prog.alloc("conv", MLEN, MLEN) + new = prog.alloc("new", MLEN, MLEN) + with pytest.raises(ValueError, match="conv_state needs"): + prog.kda_conv_state_roll_v0(conv, new, 0, shape) + + +# --------------------------------------------------------------------------- +# Normalisation +# --------------------------------------------------------------------------- + + +def test_l2_normalize_uses_the_scalar_sqrt_path(): + """There is no vector square root -- S_FP_OP carries SQRT_FP, V_ELEMENT_OP + does not -- so the reciprocal norm must come out of the scalar unit.""" + prog, shape = _prog(), _shape() + vec = prog.alloc("qk", MLEN, MLEN) + scratch = prog.alloc("sq", MLEN, MLEN) + consts = prog.kda_fp_constants() + acc = prog.fp_var("kda_norm_acc", size=2) + rows = [0, 1] + mark = len(prog.get_code()) + prog.kda_l2_normalize_v0(vec, rows, scratch, acc, consts) + code = prog.get_code()[mark:] + trace = _trace(code) + emitted = _body(code) + + vec_base = _vram_base(prog, vec) + scratch_base = _vram_base(prog, scratch) + + # Scalar rsqrt, one pair per row: there is no vector square root. + assert sum("S_SQRT_FP" in line for line in emitted) == len(rows) + assert sum("S_RECI_FP" in line for line in emitted) == len(rows) + + # The block copy that stages the vector into the scratch. mamba_block_copy + # is zero-then-add, so its V_ADD_VV carries the direction; swapping src and + # dst zeroes the vector and leaves every other operand below unchanged -- + # the one mutation that survived the first round of this test. + copies = _operands(trace, "V_ADD_VV") + assert copies, "no block copy emitted" + assert all( + dst == scratch_base and src == vec_base for dst, _, src in copies + ), f"block copy runs the wrong way: {copies}" + + # The square must write the scratch and read the vector -- not the reverse, + # which would zero the vector, and not vector-times-vector in place. + muls = _operands(trace, "V_MUL_VV") + assert muls, "no elementwise square emitted" + assert all(dst == scratch_base and src == vec_base for dst, _, src in muls) + + # The reduction must read the squared scratch, not the raw vector. + reds = _operands(trace, "V_RED_SUM") + assert reds, "no reduction emitted" + assert all(src == scratch_base for *_, src in reds) + + # The final scaling must write the vector itself. + scale_targets = {dst for dst, *_ in _operands(trace, "V_MUL_VF")} + assert vec_base in scale_targets + + +def test_kda_constant_values_use_l2_not_rms_and_the_flashkda_epsilon(): + """reci_group must be 1.0: mamba_rsqrt_fpram computes + 1/sqrt(acc*reci_group + eps), and L2 is sqrt(sum), not sqrt(mean). + eps must be 1e-6 to match the CPU reference's rsqrt(sum + 1e-6).""" + values = PlenaCompiler.kda_fp_constant_values() + zero, one, neg_one, dt_min, dt_max, reci_group, eps = values + assert (zero, one, neg_one) == (0.0, 1.0, -1.0) + assert reci_group == 1.0 + assert eps == pytest.approx(1.0e-6) + assert torch.isfinite(torch.tensor([dt_min, dt_max])).all() + + +def test_kda_constants_allocate_seven_distinct_slots(): + prog = _prog() + consts = prog.kda_fp_constants() + addresses = [v.address for v in consts.as_list()] + assert len(set(addresses)) == 7 + assert len(PlenaCompiler.kda_fp_constant_values()) == 7 + + +# --------------------------------------------------------------------------- +# Conv state layout +# --------------------------------------------------------------------------- + + +def test_conv_state_rows_are_block_major_with_consecutive_taps(): + """Pin the formula, not merely that it is a bijection. + + The emitter, the test scatter and the test gather all funnel through this + helper, so any bijection cancels out and the numerics stay green. But the + layout claim -- one block's taps consecutive, so the history shift is a run + of adjacent row copies -- rests on this exact ordering. + """ + channels, kernel = MLEN * 3, 4 + blocks = kda_conv_blocks(channels, MLEN) + assert blocks == 3 + rows = [ + kda_conv_state_row(channels, MLEN, kernel, cb, t) + for cb in range(blocks) + for t in range(kernel) + ] + assert rows == list(range(blocks * kernel)) + + for cb in range(blocks): + taps = [kda_conv_state_row(channels, MLEN, kernel, cb, t) for t in range(kernel)] + assert taps == list(range(cb * kernel, (cb + 1) * kernel)) + + +def test_conv_state_row_rejects_out_of_range(): + channels, kernel = MLEN * 2, 4 + with pytest.raises(ValueError, match="out of range"): + kda_conv_state_row(channels, MLEN, kernel, 2, 0) + with pytest.raises(ValueError, match="out of range"): + kda_conv_state_row(channels, MLEN, kernel, 0, kernel) + + +def test_conv_blocks_rejects_a_partial_trailing_block(): + with pytest.raises(ValueError, match="multiple of mlen"): + kda_conv_blocks(MLEN + 1, MLEN) diff --git a/aten/tests/test_kda_mixer.py b/aten/tests/test_kda_mixer.py new file mode 100644 index 0000000..25b0f68 --- /dev/null +++ b/aten/tests/test_kda_mixer.py @@ -0,0 +1,471 @@ +"""The whole KDA mixer for one decode token, against `kda_state_engine_step`. + +This is the first test that runs conv, gates and recurrence together, so it is +also the first that can catch a mismatch *between* them -- a layout convention +that each side reads consistently on its own but differently from the other. + +The mixer streams per head because FPRAM is 512 slots and materialising every +head's decay / q_hat / k_hat costs 3 * num_heads * key_dim. These tests check +the streamed result equals the reference, which computes all heads at once. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.reference import ( # noqa: E402 + KdaConvWeights, + KdaState, + KdaRecurrentState, + kda_state_engine_step, +) +from compiler.aten.models.kda.shape import KdaShape # noqa: E402 +from compiler.aten.models.kda.state_precision import StateStorage # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_kda_common import ( # noqa: E402 + kda_blocks, + kda_state_row, + kda_state_rows, + kda_vector_row, + kda_vector_rows, +) +from compiler.aten.plena.program_kda_gates import ( # noqa: E402 + kda_head_blocks, + kda_key_blocks, + kda_key_row, +) +from compiler.aten.plena.program_kda_mixer import KdaMixerBuffers # noqa: E402 +from compiler.aten.tests.isa_interpreter import Machine # noqa: E402 + +MLEN = 8 + + +def _rows_up(n: int) -> int: + return ((n + MLEN - 1) // MLEN) * MLEN + + +class _MixerHarness: + """Post-conv activations in, mixer output and updated state out. + + The convolutions are covered by test_kda_conv.py; feeding their outputs in + directly keeps this test on the seam between the gate and recurrence + layouts, which is what has not been exercised before. + """ + + def __init__(self, shape: KdaShape): + self.shape = shape + p = self.prog = PlenaCompiler(mlen=MLEN, blen=2) + kb = kda_key_blocks(shape, MLEN) + vb = kda_blocks(shape, MLEN) + + key_rows = _rows_up(shape.num_heads * kb) + vec_rows = _rows_up(kda_vector_rows(shape, MLEN)) + self.q = p.alloc("q", key_rows, MLEN) + self.k = p.alloc("k", key_rows, MLEN) + self.gate = p.alloc("gate", key_rows, MLEN) + self.dt_bias = p.alloc("dt_bias", key_rows, MLEN) + self.v = p.alloc("v", vec_rows, MLEN) + self.out = p.alloc("out", vec_rows, MLEN) + self.pred = p.alloc("pred", vec_rows, MLEN) + self.err = p.alloc("err", vec_rows, MLEN) + self.state = p.alloc("state", _rows_up(kda_state_rows(shape, MLEN)), MLEN) + self.beta_logit = p.alloc( + "beta_logit", _rows_up(kda_head_blocks(shape, MLEN)), MLEN + ) + self.sq_scratch = p.alloc("sq_scratch", key_rows, MLEN) + + # The per-head window, reused every head. decay and q_hat are one + # allocation because they are never live at the same time -- that is + # what brings Kimi K3 from 620 slots to 492, under FP_SRAM_DEPTH. + self.k_hat = p.fp_var("k_hat", size=shape.key_dim) + self.decay = p.fp_var("decay_and_q_hat", size=shape.key_dim) + self.q_hat = self.decay + self.beta = p.fp_var("beta", size=kda_head_blocks(shape, MLEN) * MLEN) + self.part = p.fp_var("part", size=kb) + self.acc = p.fp_var("acc", size=1) + self.scale = p.fp_var("scale", size=1) + self.rate = p.fp_var("rate", size=shape.num_heads) + self.lower = p.fp_var("lower", size=1) + self.consts = p.kda_fp_constants() + self.vb, self.kb = vb, kb + self.mark = len(p.get_code()) + + def fp_vars(self): + """Every FPRAM allocation this harness makes, for the budget test.""" + return [ + self.k_hat, self.decay, self.beta, self.part, self.acc, self.scale, + self.rate, self.lower, *self.consts.as_list(), + ] + + def _base(self, var) -> int: + return self.prog.get_vram_layout(var.name).vram_base_addr + + def buffers(self) -> KdaMixerBuffers: + return KdaMixerBuffers( + q=self.q, k=self.k, v=self.v, gate=self.gate, dt_bias=self.dt_bias, + beta_logit=self.beta_logit, state=self.state, out=self.out, + pred=self.pred, err=self.err, sq_scratch=self.sq_scratch, decay_fp=self.decay, q_hat_fp=self.q_hat, + k_hat_fp=self.k_hat, beta_fp=self.beta, part_fp=self.part, + acc_fp=self.acc, output_scale_fp=self.scale, rate_fp=self.rate, + lower_bound_fp=self.lower, consts=self.consts, + ) + + def _seed_key_tile(self, m, tile, values) -> None: + """values is [heads, key_dim].""" + s = self.shape + for h in range(s.num_heads): + for b in range(self.kb): + m.write_vram_row( + self._base(tile) + kda_key_row(s, MLEN, h, b) * MLEN, + values[h, b * MLEN : (b + 1) * MLEN].tolist(), + ) + + def _seed_value_tile(self, m, tile, values) -> None: + """values is [heads, value_dim].""" + s = self.shape + for h in range(s.num_heads): + for c in range(self.vb): + m.write_vram_row( + self._base(tile) + kda_vector_row(s, MLEN, h, c) * MLEN, + values[h, c * MLEN : (c + 1) * MLEN].tolist(), + ) + + def run(self, *, q, k, v, gate, dt_bias, beta_logit, a_log, state, output_scale, + per_head=False, head_rows_override=None): + """``per_head`` lowers one head per ``kda_mixer_step_v0`` call. + + That is the shape a whole-layer loop takes, and it must give bit-for-bit + the same program as one call over every head -- which it only does + because beta is produced once, outside. It used to be folded into the + mixer, and the second call then computed ``sigmoid(sigmoid(beta))``. + """ + s = self.shape + b = self.buffers() + self.prog.kda_beta_scalars_v0( + beta_logit=b.beta_logit, beta_fp=b.beta_fp, consts=b.consts, shape=s + ) + if head_rows_override is not None: + self.prog.kda_mixer_step_v0( + buffers=b, shape=s, head_rows=head_rows_override + ) + elif per_head: + for h in range(s.num_heads): + self.prog.kda_mixer_step_v0(buffers=b, shape=s, head_rows=[h]) + else: + self.prog.kda_mixer_step_v0(buffers=b, shape=s) + code = self.prog.get_code()[self.mark :] + + m = Machine(vlen=MLEN, vram_words=1 << 16, fpram_words=1 << 13) + self._seed_key_tile(m, self.q, q) + self._seed_key_tile(m, self.k, k) + self._seed_key_tile(m, self.gate, gate) + self._seed_key_tile(m, self.dt_bias, dt_bias) + self._seed_value_tile(m, self.v, v) + for h in range(s.num_heads): + for c in range(self.vb): + for j in range(s.key_dim): + m.write_vram_row( + self._base(self.state) + kda_state_row(s, MLEN, h, c, j) * MLEN, + state[h, j, c * MLEN : (c + 1) * MLEN].tolist(), + ) + for hb in range(kda_head_blocks(s, MLEN)): + chunk = beta_logit[hb * MLEN : (hb + 1) * MLEN].tolist() + m.write_vram_row( + self._base(self.beta_logit) + hb * MLEN, + chunk + [0.0] * (MLEN - len(chunk)), + ) + + m.write_fpram(self.consts.zero.address, PlenaCompiler.kda_fp_constant_values()) + m.write_fpram(self.rate.address, torch.exp(a_log).tolist()) + m.write_fpram(self.lower.address, [s.gate_lower_bound]) + m.write_fpram(self.scale.address, [output_scale]) + m.run(code) + + out = torch.tensor( + [ + [ + val + for c in range(self.vb) + for val in m.read_vram_row( + self._base(self.out) + kda_vector_row(s, MLEN, h, c) * MLEN, MLEN + ) + ] + for h in range(s.num_heads) + ], + dtype=torch.float32, + ) + new_state = torch.tensor( + [ + [ + [ + val + for c in range(self.vb) + for val in m.read_vram_row( + self._base(self.state) + + kda_state_row(s, MLEN, h, c, j) * MLEN, + MLEN, + ) + ] + for j in range(s.key_dim) + ] + for h in range(s.num_heads) + ], + dtype=torch.float32, + ) + return out, new_state, m + + +def _case(seed: int, heads: int, key_dim: int, value_dim: int): + """Post-conv activations plus the reference result for one token.""" + torch.manual_seed(seed) + shape = KdaShape( + hidden_size=heads * value_dim, num_heads=heads, key_dim=key_dim, + value_dim=value_dim, conv_kernel=4, + ) + q = 0.5 * torch.randn(heads, key_dim) + k = 0.5 * torch.randn(heads, key_dim) + v = torch.randn(heads, value_dim) + gate = 0.8 * torch.randn(heads, key_dim) + dt_bias = 0.8 * torch.randn(heads, key_dim) + beta_logit = torch.randn(heads) + a_log = 0.5 * torch.randn(heads) + state = torch.randn(heads, value_dim, key_dim) + + # kda_step is the post-conv half of kda_state_engine_step. + from compiler.aten.models.kda.reference import kda_step + + out, new = kda_step( + q[None], k[None], v[None], gate[None], beta_logit[None], + KdaState(state[None].clone()), a_log, dt_bias, shape, + ) + output_scale = 1.0 / (key_dim**0.5) + return shape, dict( + q=q, k=k, v=v, gate=gate, dt_bias=dt_bias, beta_logit=beta_logit, + a_log=a_log, output_scale=output_scale, + state_in_T=state.transpose(-2, -1).contiguous(), # [heads, key, value] + expected_out=out[0], + expected_state_T=new.recurrent[0].transpose(-2, -1).contiguous(), + ) + + +@pytest.mark.parametrize( + "seed,heads,key_dim,value_dim", + [ + (1, 1, MLEN, MLEN), # smallest complete mixer + (2, 3, MLEN, MLEN), # several heads streaming + (3, 2, MLEN * 2, MLEN), # key spans two blocks + (4, 2, MLEN, MLEN * 2), # value spans two blocks + (5, 3, MLEN * 2, MLEN * 2), # both -- Kimi K3's shape + ], +) +def test_mixer_matches_the_reference(seed, heads, key_dim, value_dim): + shape, c = _case(seed, heads, key_dim, value_dim) + out, new_state, _ = _MixerHarness(shape).run( + q=c["q"], k=c["k"], v=c["v"], gate=c["gate"], dt_bias=c["dt_bias"], + beta_logit=c["beta_logit"], a_log=c["a_log"], state=c["state_in_T"], + output_scale=c["output_scale"], + ) + torch.testing.assert_close(out, c["expected_out"], rtol=5e-5, atol=5e-6) + torch.testing.assert_close(new_state, c["expected_state_T"], rtol=5e-5, atol=5e-6) + + +def test_the_scalar_window_really_is_reused(): + """The whole point of streaming. If the emitter had kept the all-at-once + layout it would index decay_fp[h * key_dim], which for h > 0 runs past a + key_dim-sized window -- so this shape passing at all is the evidence.""" + shape, c = _case(6, 4, MLEN * 2, MLEN) + h = _MixerHarness(shape) + assert h.decay.size == shape.key_dim + assert h.q_hat.size == shape.key_dim + items = h.prog.kda_mixer_fpram_slots(shape) + assert items["k_hat"] == items["decay_or_q_hat"] == shape.key_dim + + out, _, _ = h.run( + q=c["q"], k=c["k"], v=c["v"], gate=c["gate"], dt_bias=c["dt_bias"], + beta_logit=c["beta_logit"], a_log=c["a_log"], state=c["state_in_T"], + output_scale=c["output_scale"], + ) + torch.testing.assert_close(out, c["expected_out"], rtol=5e-5, atol=5e-6) + + +def test_fpram_stays_within_the_hardware_file(): + """FP_SRAM_DEPTH is 512, and the whole allocation has to fit -- not just + the per-head window. + + This test asserted `3 * key_dim + 1 <= 512` and passed while the real + footprint was 620. Two things were missing: beta costs a padded row per + S_MAP_FP_V (128 at Kimi K3, not 1) and rate is indexed by head number + (96). Counting those, `3 * key_dim + beta` alone is 512 -- so the old + layout could not fit however the rest was arranged. + """ + kimi = KdaShape( + hidden_size=7168, num_heads=96, key_dim=128, value_dim=128, conv_kernel=4 + ) + p = PlenaCompiler(mlen=64, blen=4) + items = p.kda_mixer_fpram_slots(kimi) + + assert items["beta"] == 128, "S_MAP_FP_V writes whole rows, not one slot" + assert items["rate"] == 96, "rate_fp is indexed by head number" + assert items["total"] == 492 + assert items["total"] <= 512, ( + f"the mixer needs {items['total']} FPRAM slots against FP_SRAM_DEPTH 512: " + f"{items}" + ) + + # The layout this replaced, so a regression to it is loud rather than quiet. + unshared = items["total"] + kimi.key_dim + assert unshared > 512, ( + "a separate decay and q_hat window is what overflowed; if this ever fits, " + "sharing them is no longer load-bearing" + ) + all_at_once = 3 * kimi.num_heads * kimi.key_dim + kimi.num_heads + assert all_at_once > 512 * 70, ( + "if this ever fits, the streaming was unnecessary -- recheck the shape" + ) + + +def test_fpram_accounting_matches_what_is_actually_allocated(): + """The itemised count must equal the allocator's high-water mark. + + Otherwise it is a comment, not a budget: the compiler's FPRAMAllocator + defaults to 1024 slots, so nothing else in the stack would catch a mixer + that overflows the real 512. + """ + shape, _ = _case(11, 3, MLEN, MLEN) + h = _MixerHarness(shape) + items = h.prog.kda_mixer_fpram_slots(shape) + + used = [(v.address, v.size) for v in h.fp_vars()] + high_water = max(a + n for a, n in used) + assert high_water <= items["total"], ( + f"the harness allocates {high_water} slots but the budget claims " + f"{items['total']}: {items}" + ) + + +def test_decay_and_q_hat_may_share_one_window(): + """The two are never live at once, and the mixer must still be correct when + a caller exploits that -- which is the only reason Kimi K3 fits.""" + shape, c = _case(12, 2, MLEN * 2, MLEN) + h = _MixerHarness(shape) + assert h.q_hat is h.decay, "the harness is meant to exercise the shared window" + out, state_out, _ = h.run( + q=c["q"], k=c["k"], v=c["v"], gate=c["gate"], dt_bias=c["dt_bias"], + beta_logit=c["beta_logit"], a_log=c["a_log"], state=c["state_in_T"], + output_scale=c["output_scale"], + ) + torch.testing.assert_close(out, c["expected_out"], rtol=5e-5, atol=5e-6) + torch.testing.assert_close(state_out, c["expected_state_T"], rtol=5e-5, atol=5e-6) + + +def test_heads_can_be_lowered_individually(): + """head_rows selects a subset, which is what a per-head outer loop uses. + + Three calls of one head each must equal one call of three heads. This + failed before beta moved out of the mixer: head 0 was right and heads 1 + and 2 were off by 1.2e-2, because each call re-applied sigmoid to + beta_logit in place. + """ + shape, c = _case(8, 3, MLEN, MLEN) + args = dict( + q=c["q"], k=c["k"], v=c["v"], gate=c["gate"], dt_bias=c["dt_bias"], + beta_logit=c["beta_logit"], a_log=c["a_log"], state=c["state_in_T"], + output_scale=c["output_scale"], + ) + split, state_split, _ = _MixerHarness(shape).run(per_head=True, **args) + torch.testing.assert_close(split, c["expected_out"], rtol=5e-5, atol=5e-6) + torch.testing.assert_close(state_split, c["expected_state_T"], rtol=5e-5, atol=5e-6) + + together, state_together, _ = _MixerHarness(shape).run(**args) + torch.testing.assert_close(split, together, rtol=0, atol=0) + torch.testing.assert_close(state_split, state_together, rtol=0, atol=0) + + +def test_head_rows_actually_restricts_the_work(): + """Lowering head 0 alone must leave the other heads' state untouched. + + Without this, an implementation that ignores head_rows and does every head + passes every other test in this file. + """ + shape, c = _case(9, 3, MLEN, MLEN) + h = _MixerHarness(shape) + _, state_out, _ = h.run(head_rows_override=[0], **dict( + q=c["q"], k=c["k"], v=c["v"], gate=c["gate"], dt_bias=c["dt_bias"], + beta_logit=c["beta_logit"], a_log=c["a_log"], state=c["state_in_T"], + output_scale=c["output_scale"], + )) + torch.testing.assert_close( + state_out[0], c["expected_state_T"][0], rtol=5e-5, atol=5e-6 + ) + for head in range(1, shape.num_heads): + torch.testing.assert_close( + state_out[head], c["state_in_T"][head], rtol=0, atol=0 + ) + + +# --------------------------------------------------------------------------- +# The conv/mixer seam. +# +# kda_conv_step_v0 writes its output blocked by *channel*: block cb covers +# channels [cb*mlen, (cb+1)*mlen). kda_mixer_step_v0 reads q and k blocked by +# *key*: kda_key_row(h, b) = h*key_blocks + b. Those are two different helpers +# in two different modules, and the mixer feeds directly from the conv, so if +# they disagree every KDA result is wrong in a way neither module's own tests +# can see. +# --------------------------------------------------------------------------- + + +def test_conv_channel_blocks_and_mixer_key_blocks_are_the_same_rows(): + """For a key-width tensor the channel index is h*key_dim + k, so channel + block cb is head cb//key_blocks's key block cb%key_blocks -- which is + exactly kda_key_row. This holds because key_dim is required to be a whole + number of blocks; it would not if a head's keys straddled a block.""" + from compiler.aten.plena.program_kda_common import kda_conv_blocks + + for heads, key_dim in [(1, MLEN), (3, MLEN), (2, MLEN * 2), (5, MLEN * 3)]: + shape = KdaShape( + hidden_size=heads * MLEN, num_heads=heads, key_dim=key_dim, + value_dim=MLEN, conv_kernel=4, + ) + channels = heads * key_dim + key_blocks = kda_key_blocks(shape, MLEN) + assert kda_conv_blocks(channels, MLEN) == heads * key_blocks + + for h in range(heads): + for b in range(key_blocks): + # the channel block holding head h's key block b + first_channel = h * key_dim + b * MLEN + channel_block = first_channel // MLEN + assert channel_block == kda_key_row(shape, MLEN, h, b), ( + f"conv writes head {h} key block {b} to row {channel_block} " + f"but the mixer reads it from {kda_key_row(shape, MLEN, h, b)}" + ) + + +def test_conv_channel_blocks_and_mixer_value_blocks_are_the_same_rows(): + """Same seam for v, which is value-width: channel index is h*value_dim + j, + so channel block cb is kda_vector_row(h, cb%value_blocks).""" + from compiler.aten.plena.program_kda_common import kda_conv_blocks + + for heads, value_dim in [(1, MLEN), (3, MLEN), (2, MLEN * 2), (4, MLEN * 3)]: + shape = KdaShape( + hidden_size=heads * value_dim, num_heads=heads, key_dim=MLEN, + value_dim=value_dim, conv_kernel=4, + ) + channels = heads * value_dim + value_blocks = kda_blocks(shape, MLEN) + assert kda_conv_blocks(channels, MLEN) == heads * value_blocks + + for h in range(heads): + for c in range(value_blocks): + channel_block = (h * value_dim + c * MLEN) // MLEN + assert channel_block == kda_vector_row(shape, MLEN, h, c), ( + f"conv writes head {h} value block {c} to row {channel_block} " + f"but the mixer reads it from {kda_vector_row(shape, MLEN, h, c)}" + ) diff --git a/aten/tests/test_kda_normalize.py b/aten/tests/test_kda_normalize.py new file mode 100644 index 0000000..fd6754b --- /dev/null +++ b/aten/tests/test_kda_normalize.py @@ -0,0 +1,301 @@ +"""KDA's L2 normalisation of q and k, executed against the reference. + +`q` and `k` are key-width, and Kimi K3 is `key_dim = 128` against a default +`mlen = 64`. The norm contracts over the key axis, so the two halves of one `q` +share a single norm -- this is the one place in KDA where the column-block +folding used everywhere else does *not* apply, and normalising each row on its +own would divide each half by its own partial norm. + +The reference is `reference.py`'s `x * rsqrt(sum(x^2) + 1e-6)`, matching +FlashKDA's recurrent kernel. `torch.nn.functional.normalize` clamps instead and +does not agree. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.tests.isa_interpreter import Machine # noqa: E402 + +MLEN = 8 +EPS = 1.0e-6 + + +def _rows_up(n: int) -> int: + return ((n + MLEN - 1) // MLEN) * MLEN + + +def _reference(x: torch.Tensor) -> torch.Tensor: + """[vectors, width] -> normalised, exactly as reference.py does it.""" + return x * torch.rsqrt(x.square().sum(-1, keepdim=True) + EPS) + + +class _NormHarness: + def __init__(self, vectors: int, blocks: int): + self.vectors, self.blocks = vectors, blocks + p = self.prog = PlenaCompiler(mlen=MLEN, blen=2) + rows = _rows_up(vectors * blocks) + self.vec = p.alloc("vec", rows, MLEN) + self.sq = p.alloc("sq", rows, MLEN) + self.part = p.fp_var("part", size=vectors * blocks) + self.acc = p.fp_var("acc", size=vectors) + self.consts = p.kda_fp_constants() + self.mark = len(p.get_code()) + + def _base(self, var) -> int: + return self.prog.get_vram_layout(var.name).vram_base_addr + + def run(self, x: torch.Tensor): + """x is [vectors, blocks * mlen].""" + self.prog.kda_l2_normalize_blocked_v0( + self.vec, vectors=self.vectors, blocks=self.blocks, + sq_scratch=self.sq, part_fp=self.part, acc_fp=self.acc, + consts=self.consts, + ) + code = self.prog.get_code()[self.mark :] + + m = Machine(vlen=MLEN) + for v in range(self.vectors): + for c in range(self.blocks): + m.write_vram_row( + self._base(self.vec) + (v * self.blocks + c) * MLEN, + x[v, c * MLEN : (c + 1) * MLEN].tolist(), + ) + m.write_fpram(self.consts.zero.address, PlenaCompiler.kda_fp_constant_values()) + m.run(code) + + return torch.tensor( + [ + [ + val + for c in range(self.blocks) + for val in m.read_vram_row( + self._base(self.vec) + (v * self.blocks + c) * MLEN, MLEN + ) + ] + for v in range(self.vectors) + ], + dtype=torch.float32, + ), m + + +@pytest.mark.parametrize( + "seed,vectors,blocks", + [ + (1, 1, 1), # degenerate: one vector, one block + (2, 3, 1), # several vectors, still one block + (3, 1, 2), # THE case: one vector spanning two blocks + (4, 4, 2), # several vectors, two blocks -- Kimi K3's shape + (5, 2, 3), # three blocks + ], +) +def test_blocked_normalize_matches_the_reference(seed, vectors, blocks): + torch.manual_seed(seed) + x = 0.7 * torch.randn(vectors, blocks * MLEN) + got, _ = _NormHarness(vectors, blocks).run(x) + torch.testing.assert_close(got, _reference(x), rtol=2e-5, atol=2e-6) + + +def test_both_halves_share_one_norm(): + """The whole point. If each row were normalised by its own partial norm the + result would still be finite and unit-ish per half -- this pins the *joint* + norm by comparing against a per-half normalisation and requiring they differ. + """ + torch.manual_seed(11) + x = 0.7 * torch.randn(1, 2 * MLEN) + # make the halves very unequal, so a per-half norm is clearly wrong + x[0, MLEN:] *= 6.0 + + got, _ = _NormHarness(1, 2).run(x) + torch.testing.assert_close(got, _reference(x), rtol=2e-5, atol=2e-6) + + per_half = torch.cat( + [_reference(x[:, :MLEN]), _reference(x[:, MLEN:])], dim=-1 + ) + assert not torch.allclose(got, per_half, rtol=1e-2, atol=1e-2), ( + "each half was normalised by its own norm" + ) + + # and the joint norm really is 1 (up to eps) + assert got.square().sum().item() == pytest.approx(1.0, abs=2e-4) + + +def test_vectors_do_not_share_a_norm(): + """Each vector gets its own norm; a fold that summed across vectors would + give every one the same scale.""" + torch.manual_seed(12) + x = 0.7 * torch.randn(3, 2 * MLEN) + x[1] *= 5.0 + got, _ = _NormHarness(3, 2).run(x) + for v in range(3): + assert got[v].square().sum().item() == pytest.approx(1.0, abs=2e-4) + + +def test_rejects_undersized_fpram_and_aliased_scratch(): + p = PlenaCompiler(mlen=MLEN, blen=2) + vec = p.alloc("v", MLEN, MLEN) + sq = p.alloc("s", MLEN, MLEN) + consts = p.kda_fp_constants() + ok = dict(vectors=2, blocks=2, sq_scratch=sq, consts=consts) + + with pytest.raises(ValueError, match="part_fp"): + p.kda_l2_normalize_blocked_v0( + vec, part_fp=p.fp_var("p1", size=3), acc_fp=p.fp_var("a1", size=2), **ok + ) + with pytest.raises(ValueError, match="acc_fp"): + p.kda_l2_normalize_blocked_v0( + vec, part_fp=p.fp_var("p2", size=4), acc_fp=p.fp_var("a2", size=1), **ok + ) + with pytest.raises(ValueError, match="distinct"): + p.kda_l2_normalize_blocked_v0( + vec, part_fp=p.fp_var("p3", size=4), acc_fp=p.fp_var("a3", size=2), + **{**ok, "sq_scratch": vec}, + ) + + +# --------------------------------------------------------------------------- +# first_row: normalising one head's slice of a tile that holds every head. +# +# The mixer streams per head to keep FPRAM bounded, and q/k arrive as one tile +# covering all heads. Copying a head's rows out first would cost key_blocks row +# copies per head per tensor; first_row addresses them in place instead. +# --------------------------------------------------------------------------- + + +class _SliceHarness: + def __init__(self, heads: int, blocks: int): + self.heads, self.blocks = heads, blocks + p = self.prog = PlenaCompiler(mlen=MLEN, blen=2) + rows = _rows_up(heads * blocks) + self.vec = p.alloc("vec", rows, MLEN) + self.sq = p.alloc("sq", rows, MLEN) + self.part = p.fp_var("part", size=blocks) + self.acc = p.fp_var("acc", size=1) + self.consts = p.kda_fp_constants() + self.mark = len(p.get_code()) + + def _base(self, var) -> int: + return self.prog.get_vram_layout(var.name).vram_base_addr + + def run(self, x: torch.Tensor, head: int): + """x is [heads, blocks*mlen]; normalise only `head`'s rows in place.""" + self.prog.kda_l2_normalize_blocked_v0( + self.vec, vectors=1, blocks=self.blocks, sq_scratch=self.sq, + part_fp=self.part, acc_fp=self.acc, consts=self.consts, + first_row=head * self.blocks, + ) + code = self.prog.get_code()[self.mark :] + + m = Machine(vlen=MLEN) + for h in range(self.heads): + for c in range(self.blocks): + m.write_vram_row( + self._base(self.vec) + (h * self.blocks + c) * MLEN, + x[h, c * MLEN : (c + 1) * MLEN].tolist(), + ) + m.write_fpram(self.consts.zero.address, PlenaCompiler.kda_fp_constant_values()) + m.run(code) + + return torch.tensor( + [ + [ + val + for c in range(self.blocks) + for val in m.read_vram_row( + self._base(self.vec) + (h * self.blocks + c) * MLEN, MLEN + ) + ] + for h in range(self.heads) + ], + dtype=torch.float32, + ) + + +@pytest.mark.parametrize("heads,blocks", [(3, 1), (3, 2), (2, 3)]) +def test_first_row_normalises_only_that_head(heads, blocks): + torch.manual_seed(31) + x = 0.7 * torch.randn(heads, blocks * MLEN) + for head in range(heads): + got = _SliceHarness(heads, blocks).run(x, head) + # the selected head is normalised + torch.testing.assert_close( + got[head], _reference(x[head : head + 1])[0], rtol=2e-5, atol=2e-6 + ) + # every other head is untouched, bit for bit + for other in range(heads): + if other != head: + torch.testing.assert_close(got[other], x[other], rtol=0, atol=0) + + +def test_first_row_uses_its_own_slice_not_row_zero(): + """A first_row that is accepted but ignored would normalise head 0 every + time. Making the heads' magnitudes differ makes that visible.""" + torch.manual_seed(32) + heads, blocks = 3, 2 + x = 0.7 * torch.randn(heads, blocks * MLEN) + x[0] *= 8.0 + got = _SliceHarness(heads, blocks).run(x, 2) + torch.testing.assert_close(got[0], x[0], rtol=0, atol=0) + assert got[2].square().sum().item() == pytest.approx(1.0, abs=2e-4) + + +def test_sq_scratch_must_reach_past_first_row(): + """`vec`'s own bound fires first when both are short, so a too-small + sq_scratch needs a vec that is big enough. Without this the sq_scratch + check could be written against `vectors * blocks` alone and every test + still passed -- while the mixer, which walks first_row across heads, wrote + the squares of the last head past the end of the scratch.""" + p = PlenaCompiler(mlen=MLEN, blen=2) + vec = p.alloc("v", MLEN * 2, MLEN) + sq = p.alloc("s", MLEN, MLEN) + with pytest.raises(ValueError, match="sq_scratch must match"): + p.kda_l2_normalize_blocked_v0( + vec, vectors=1, blocks=2, sq_scratch=sq, + part_fp=p.fp_var("p", size=2), acc_fp=p.fp_var("a", size=1), + consts=p.kda_fp_constants(), first_row=MLEN, + ) + + +def test_the_squares_copy_is_only_as_long_as_the_slice(): + """Copying from row 0 up to the last live row makes the cost O(first_row), + and the mixer moves first_row across every head -- quadratic in head count. + Two slices of the same width must cost the same however far in they sit.""" + counts = {} + for first_row in (0, MLEN * 4, MLEN * 32): + p = PlenaCompiler(mlen=MLEN, blen=2) + vec = p.alloc("v", MLEN * 40, MLEN) + sq = p.alloc("s", MLEN * 40, MLEN) + consts = p.kda_fp_constants() + mark = len(p.get_code()) + p.kda_l2_normalize_blocked_v0( + vec, vectors=1, blocks=2, sq_scratch=sq, + part_fp=p.fp_var(f"p{first_row}", size=2), + acc_fp=p.fp_var(f"a{first_row}", size=1), + consts=consts, first_row=first_row, + ) + counts[first_row] = len(p.get_code()) - mark + + # A row of block_copy costs about 80 instructions, so a copy that started + # at row 0 would put 256 * 80 between the two ends. What is left is the + # wider immediate for a further-out address, which is a constant per call. + spread = counts[MLEN * 32] - counts[0] + assert spread < 100, f"cost still grows with first_row: {counts}" + + +def test_first_row_bounds_are_checked(): + p = PlenaCompiler(mlen=MLEN, blen=2) + vec = p.alloc("v", MLEN, MLEN) + sq = p.alloc("s", MLEN, MLEN) + with pytest.raises(ValueError, match="rows"): + p.kda_l2_normalize_blocked_v0( + vec, vectors=1, blocks=2, sq_scratch=sq, + part_fp=p.fp_var("p", size=2), acc_fp=p.fp_var("a", size=1), + consts=p.kda_fp_constants(), first_row=MLEN, + ) diff --git a/aten/tests/test_kda_prefill.py b/aten/tests/test_kda_prefill.py new file mode 100644 index 0000000..d570dda --- /dev/null +++ b/aten/tests/test_kda_prefill.py @@ -0,0 +1,336 @@ +"""KDA chunk primitives: the cumulative decay and the UT transform. + +Both run through the ISA interpreter, so what is checked is the emitted +instructions rather than the Python that emitted them. + +The chunked algebra these serve was verified against `kda_step` run +sequentially before any of it was lowered; `program_kda_chunk`'s module +docstring carries the derivation and the measurements. Two results from that +work are pinned here as tests, because they are the ones that are easy to +"simplify" back into a wrong answer: + +* the decay weighting in `M` is not optional -- dropping it is wrong by 1.8e-1, +* the cumulative decay is a running product, not `exp` of a running sum. +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.shape import KdaShape # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_kda_chunk import ( # noqa: E402 + kda_chunk_cols, + kda_chunk_rows, + kda_max_chunk_for, +) +from compiler.aten.plena.program_kda_gates import kda_key_blocks # noqa: E402 +from compiler.aten.tests.isa_interpreter import ( # noqa: E402 + Machine, + UnsupportedInstruction, +) + +MLEN = 8 + + +def _rows_up(n: int, mlen: int = MLEN) -> int: + return ((n + mlen - 1) // mlen) * mlen + + +def _static(asm: str) -> int: + return len([ln for ln in asm.splitlines() if ln.strip() and not ln.strip().startswith(";")]) + + +def _shape(key_dim: int, mlen: int = MLEN) -> KdaShape: + return KdaShape( + hidden_size=8, num_heads=1, key_dim=key_dim, value_dim=mlen, conv_kernel=4 + ) + + +# --------------------------------------------------------------------------- +# The cumulative decay. + + +def _run_cumprod(seed: int, key_dim: int, chunk: int): + torch.manual_seed(seed) + shape = _shape(key_dim) + p = PlenaCompiler(mlen=MLEN, blen=2) + blocks = kda_key_blocks(shape, MLEN) + rows = kda_chunk_rows(shape, MLEN, chunk) + cols = kda_chunk_cols(shape, MLEN) + decay = p.alloc("decay", rows, cols) + prev = p.alloc("prev", rows, cols) + mark = len(p.get_code()) + p.kda_chunk_decay_cumprod_v0(decay=decay, prev=prev, chunk=chunk, shape=shape) + code = p.get_code()[mark:] + + # Per-step decays, the range kda_decay_scalars_v0 produces: exp of something + # in [gate_lower_bound, 0], so every factor is in [e^-5, 1]. + d = torch.exp(shape.gate_lower_bound * torch.sigmoid(torch.randn(chunk, key_dim))) + m = Machine(vlen=MLEN, vram_words=1 << 16, fpram_words=1 << 13) + # The key axis is on lanes: block `b` of timestep `t` is row `t` of column + # block `b`, which the tile addressing places at `b * rows * MLEN + t * MLEN`. + addr = lambda b, t: p.get_vram_tile_addr(decay.name, 0, b) + t * MLEN # noqa: E731 + for t in range(chunk): + for b in range(blocks): + m.write_vram_row(addr(b, t), d[t, b * MLEN : (b + 1) * MLEN].tolist()) + m.run(code) + got = torch.tensor([ + sum((m.read_vram_row(addr(b, t), MLEN) for b in range(blocks)), []) + for t in range(chunk) + ]) + return got, torch.cumprod(d, dim=0), code + + +@pytest.mark.parametrize( + "seed,key_dim,chunk", + [(1, MLEN, 2), (2, MLEN, 4), (3, MLEN, 16), (4, MLEN * 2, 16), (5, MLEN * 4, 8)], +) +def test_cumulative_decay_matches_a_running_product(seed, key_dim, chunk): + got, want, _ = _run_cumprod(seed, key_dim, chunk) + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-7) + + +def test_cumulative_decay_spans_every_key_block(): + """key_dim > mlen puts the key axis across column blocks. A scan that ran + only over block 0 would leave the upper half holding the per-step decay -- + finite, plausible, wrong.""" + got, want, _ = _run_cumprod(6, MLEN * 2, 8) + torch.testing.assert_close(got, want, rtol=1e-5, atol=1e-7) + # The two halves decay independently; if they had been chained the second + # would be the product of both. + assert got.shape[1] == MLEN * 2 + + +def test_cumulative_decay_is_not_a_plain_copy(): + """The scan must actually multiply. Guards against a version that emits the + staging copy and drops the multiply, which leaves decay holding the + per-step values -- correct at t=0 and wrong everywhere after.""" + got, want, _ = _run_cumprod(7, MLEN, 8) + assert (got[-1] < got[0]).all(), "a cumulative product of factors < 1 must decrease" + + +# --------------------------------------------------------------------------- +# The UT transform. + + +def _run_ut(seed: int, chunk: int, mlen: int = MLEN, *, scale: float = 0.3): + torch.manual_seed(seed) + shape = _shape(mlen, mlen) + p = PlenaCompiler(mlen=mlen, blen=2) + up = _rows_up(chunk, mlen) + m_tile = p.alloc("m", up, mlen) + identity = p.alloc("ident", up, mlen) + t_out = p.alloc("t_out", up, mlen) + m_fp = p.fp_var("m_fp", size=mlen) + beta_fp = p.fp_var("beta", size=up) + consts = p.kda_fp_constants() + mark = len(p.get_code()) + p.kda_ut_transform_v0( + m=m_tile, identity=identity, beta_fp=beta_fp, t_out=t_out, + m_fp=m_fp, consts=consts, chunk=chunk, shape=shape, + ) + code = p.get_code()[mark:] + + m_dense = torch.tril(torch.randn(chunk, chunk) * scale, -1) + beta = torch.sigmoid(torch.randn(chunk)) + want = torch.linalg.inv( + torch.eye(chunk) + torch.diag(beta) @ m_dense + ) @ torch.diag(beta) + + mac = Machine(vlen=mlen, vram_words=1 << 18, fpram_words=1 << 13) + mb = p.get_vram_layout(m_tile.name).vram_base_addr + ib = p.get_vram_layout(identity.name).vram_base_addr + tb = p.get_vram_layout(t_out.name).vram_base_addr + pad = [0.0] * (mlen - chunk) + for i in range(chunk): + mac.write_vram_row(mb + i * mlen, m_dense[i].tolist() + pad) + mac.write_vram_row( + ib + i * mlen, [1.0 if j == i else 0.0 for j in range(chunk)] + pad + ) + # Garbage, so "the program wrote it" is distinguishable from "it was + # already zero" -- a fresh Machine is all zeros. + mac.write_vram_row(tb + i * mlen, [9.5] * mlen) + mac.write_fpram(beta_fp.address, beta.tolist()) + mac.write_fpram(consts.zero.address, p.kda_fp_constant_values()) + mac.run(code) + got = torch.tensor([mac.read_vram_row(tb + i * mlen, chunk) for i in range(chunk)]) + return got, want, code, mac + + +@pytest.mark.parametrize("seed,chunk,mlen", [(1, 2, 8), (2, 4, 8), (3, 8, 8), (4, 16, 16)]) +def test_ut_transform_matches_a_dense_inverse(seed, chunk, mlen): + got, want, _, _ = _run_ut(seed, chunk, mlen) + torch.testing.assert_close(got, want, rtol=5e-5, atol=5e-6) + + +def test_ut_transform_ignores_m_on_and_above_the_diagonal(): + """The substitution reads only M[i, j] for j < i. + + Every other test feeds a strictly lower-triangular M, so the emitter + could read the diagonal and still pass -- extending its sweep from + range(i) to range(i + 1) left the whole file green. In the real + pipeline M comes from a gram matrix whose diagonal is |k|^2, near 1, + so a mask that is off by one row would be a real wrong answer rather than a + harmless zero. Pin the contract instead of the accident. + """ + torch.manual_seed(13) + chunk, mlen = 8, 8 + shape = _shape(mlen, mlen) + strict = torch.tril(torch.randn(chunk, chunk) * 0.3, -1) + # Same strict lower triangle, but the diagonal and above are loud. + polluted = strict + torch.triu(torch.randn(chunk, chunk) * 5.0 + 3.0) + + outs = [] + for m_dense in (strict, polluted): + p = PlenaCompiler(mlen=mlen, blen=2) + up = _rows_up(chunk, mlen) + m_tile = p.alloc("m", up, mlen) + identity = p.alloc("ident", up, mlen) + t_out = p.alloc("t_out", up, mlen) + m_fp = p.fp_var("m_fp", size=mlen) + beta_fp = p.fp_var("beta", size=up) + consts = p.kda_fp_constants() + mark = len(p.get_code()) + p.kda_ut_transform_v0( + m=m_tile, identity=identity, beta_fp=beta_fp, t_out=t_out, + m_fp=m_fp, consts=consts, chunk=chunk, shape=shape, + ) + code = p.get_code()[mark:] + + torch.manual_seed(13) + beta = torch.sigmoid(torch.randn(chunk)) + mac = Machine(vlen=mlen, vram_words=1 << 18, fpram_words=1 << 13) + mb = p.get_vram_layout(m_tile.name).vram_base_addr + ib = p.get_vram_layout(identity.name).vram_base_addr + tb = p.get_vram_layout(t_out.name).vram_base_addr + pad = [0.0] * (mlen - chunk) + for i in range(chunk): + mac.write_vram_row(mb + i * mlen, m_dense[i].tolist() + pad) + mac.write_vram_row( + ib + i * mlen, [1.0 if j == i else 0.0 for j in range(chunk)] + pad + ) + mac.write_vram_row(tb + i * mlen, [9.5] * mlen) + mac.write_fpram(beta_fp.address, beta.tolist()) + mac.write_fpram(consts.zero.address, p.kda_fp_constant_values()) + mac.run(code) + outs.append(torch.tensor([mac.read_vram_row(tb + i * mlen, chunk) for i in range(chunk)])) + + torch.testing.assert_close(outs[0], outs[1], rtol=0, atol=0) + + +def test_ut_transform_is_lower_triangular(): + """Forward substitution must never write above the diagonal.""" + got, _, _, _ = _run_ut(5, 8) + assert torch.triu(got, diagonal=1).abs().max() < 1e-6 + + +def test_ut_transform_diagonal_is_beta(): + """L is unitriangular, so T's diagonal is exactly beta. A missing identity + row or a misapplied beta shows up here first.""" + torch.manual_seed(9) + got, want, _, _ = _run_ut(9, 8) + torch.testing.assert_close(torch.diagonal(got), torch.diagonal(want), rtol=5e-5, atol=5e-6) + + +def test_ut_transform_overwrites_its_destination(): + """t_out is seeded with 9.5. Anything left over means a row was not written.""" + got, _, _, _ = _run_ut(10, 8) + assert (got.abs() < 5.0).all(), f"stale seed value survived: {got}" + + +def test_ut_transform_is_one_hardware_loop_per_row(): + """The substitution's inner sum walks j < i with the destination pinned, so + it is a progression and must not unroll. Without this a change that broke + the progression would still be numerically right and quadratically larger.""" + counts = {} + for chunk in (4, 8, 16): + _, _, code, _ = _run_ut(11, chunk, 16) + counts[chunk] = code.count("V_FMA_VF") + # One FMA per row that has predecessors, whatever the row length. + assert counts == {4: 3, 8: 7, 16: 15}, counts + + +def test_ut_transform_emits_nothing_the_oracle_cannot_model(): + _, _, code, _ = _run_ut(12, 8) + try: + Machine(vlen=MLEN).run(code) + except UnsupportedInstruction as exc: # pragma: no cover - failure path + pytest.fail(f"UT transform emitted an unmodelled instruction: {exc}") + + +# --------------------------------------------------------------------------- +# The chunk-size bound. + + +def test_rejects_a_chunk_whose_reciprocal_decay_overflows_bf16(): + """M is formed as a matmul of k*A against k/A, and 1/A reaches + exp(chunk * |gate_lower_bound|). bf16 tops out at exp(88.7), so at Kimi K3's + -5 the last chunk that works is 17. The failure is silent -- inf, then nan + through the whole solve -- so it is refused rather than warned about.""" + p = PlenaCompiler(mlen=64, blen=2) + shape = KdaShape( + hidden_size=8, num_heads=1, key_dim=64, value_dim=64, conv_kernel=4 + ) + assert shape.gate_lower_bound == -5.0 + for ok in (1, 16, 17): + p.kda_chunk_check_range(ok, shape) + for bad in (18, 32, 64): + with pytest.raises(ValueError, match="past bf16"): + p.kda_chunk_check_range(bad, shape) + + +def test_the_chunk_bound_scales_with_the_gate_lower_bound(): + assert kda_max_chunk_for(-5.0) == 17 + assert kda_max_chunk_for(-2.5) == 35 + assert kda_max_chunk_for(-10.0) == 8 + with pytest.raises(ValueError, match="must be negative"): + kda_max_chunk_for(0.0) + + +def test_rejects_a_chunk_wider_than_mlen(): + """The [chunk, chunk] tiles here are one row per timestep, so the row must + hold the whole chunk in a single block.""" + p = PlenaCompiler(mlen=MLEN, blen=2) + shape = _shape(MLEN) + up = _rows_up(MLEN * 2) + with pytest.raises(ValueError, match="exceeds mlen"): + p.kda_ut_transform_v0( + m=p.alloc("m", up, MLEN), identity=p.alloc("i", up, MLEN), + beta_fp=p.fp_var("b", size=up), t_out=p.alloc("t", up, MLEN), + m_fp=p.fp_var("mf", size=MLEN), consts=p.kda_fp_constants(), + chunk=MLEN * 2, shape=shape, + ) + + +def test_rejects_an_m_fp_sized_to_chunk_rather_than_mlen(): + """S_MAP_FP_V moves a whole row, so m_fp must be mlen slots. Sizing it to + chunk writes past the allocation -- the same trap beta_fp had.""" + p = PlenaCompiler(mlen=MLEN, blen=2) + shape = _shape(MLEN) + with pytest.raises(ValueError, match="whole row"): + p.kda_ut_transform_v0( + m=p.alloc("m", MLEN, MLEN), identity=p.alloc("i", MLEN, MLEN), + beta_fp=p.fp_var("b", size=MLEN), t_out=p.alloc("t", MLEN, MLEN), + m_fp=p.fp_var("mf", size=4), consts=p.kda_fp_constants(), + chunk=4, shape=shape, + ) + + +def test_rejects_t_out_aliased_to_m(): + p = PlenaCompiler(mlen=MLEN, blen=2) + shape = _shape(MLEN) + m = p.alloc("m", MLEN, MLEN) + with pytest.raises(ValueError, match="must not alias m"): + p.kda_ut_transform_v0( + m=m, identity=p.alloc("i", MLEN, MLEN), + beta_fp=p.fp_var("b", size=MLEN), t_out=m, + m_fp=p.fp_var("mf", size=MLEN), consts=p.kda_fp_constants(), + chunk=4, shape=shape, + ) diff --git a/aten/tests/test_kda_prefill_structure.py b/aten/tests/test_kda_prefill_structure.py new file mode 100644 index 0000000..779cfd1 --- /dev/null +++ b/aten/tests/test_kda_prefill_structure.py @@ -0,0 +1,558 @@ +"""Structural coverage for `kda_prefill_chunk_v0`. + +The prefill layer had **no pytest coverage at all**. Its only caller anywhere +was `transactional_emulator/testbench/kda/kda_stage_test.py`, a hand-run script +in no CI job, and the ISA interpreter models no matmul so numerical coverage in +pytest is not available. + +What *is* available is the emitted assembly. Every projection announces its +operands and tile indices in a comment: + + ; VRAM Sub Projection T To: k[0][:] @ kda_k_tilde_spill[0][:]^T -> gram[0][0] + ; VRAM Sub Projection To: err_t[0][:] @ kda_k_end_spill[:][0] -> contrib[0][0] + +so operand order, the `_to` versus `_T_to` choice, the block indices, and the +spill/prefetch pairing can all be pinned without executing anything. Those are +exactly the mistakes that produce a finite, plausible, wrong answer -- a reviewer +found that swapping two block indices in one projection passed bit-identically, +because at the only shape ever tested both were 0. + +Numerical correctness stays with the emulator stage test. This file is what +fails in CI when the composition changes. +""" + +from __future__ import annotations + +import os +import re +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.shape import KdaShape # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.instruction_stream import ( # noqa: E402 + dynamic_count, + static_count, +) +from compiler.aten.plena.program_kda_prefill import ( # noqa: E402 + KdaPrefillBuffers, + kda_prefill_state_transpose_shapes, + kda_prefill_tile_shapes, +) + +MLEN = 64 + +#: `; VRAM Sub Projection [T ]To: A[i][:] @ B[...] -> C[r][c]` +_PROJ = re.compile( + r";\s*VRAM Sub Projection (T )?To:\s*" + r"(?P\w+)\[(?P\d+)\]\[:\]\s*@\s*" + r"(?P\w+)(?P\[[^\]]*\]\[[^\]]*\])\^?T?\s*->\s*" + r"(?P\w+)\[(?P\d+)\]\[(?P\d+)\]" +) +_STORE = re.compile(r";\s*Store (?P\w+) from VRAM to HBM") + + +def _build( + chunk: int = 16, + mlen: int = MLEN, + key_dim: int | None = None, + value_dim: int | None = None, + emit: bool = True, +): + """Compile one chunk, allocating every tile at the shape the emitter demands. + + Allocation goes through `kda_prefill_tile_shapes` rather than a local + `mlen x mlen` for everything, which is what it used to be. That divergence + is how a spill came to prefetch past its allocation: the emitter's idea of + the tiles and the caller's were written separately and only overlapped at + the one shape anyone ran. + """ + shape = KdaShape( + hidden_size=mlen, num_heads=1, + key_dim=key_dim or mlen, value_dim=value_dim or mlen, conv_kernel=4, + ) + p = PlenaCompiler(mlen=mlen, blen=4) + # The build-level KV precision check reads the active TOML and the shipped + # one is MX; that path has its own tests in test_spill_precision_guard.py. + p._bf16_kv_checked = True + want = kda_prefill_tile_shapes(shape, mlen, chunk) + tiles = {n: p.alloc(n, r, c, strict=False) for n, (r, c) in want.items()} + buffers = KdaPrefillBuffers( + **tiles, + beta_fp=p.fp_var("beta", size=mlen), m_fp=p.fp_var("m", size=mlen), + output_scale_fp=p.fp_var("os", size=1), consts=p.kda_fp_constants(), + ) + mark = len(p.get_code()) + if emit: + p.kda_prefill_chunk_v0(buffers=buffers, chunk=chunk, shape=shape) + return p, buffers, p.get_code()[mark:] + + +def _projections(code: str): + out = [] + for line in code.splitlines(): + m = _PROJ.search(line) + if m: + out.append({ + "transposed": bool(m.group(1)), + "vram": m.group("vram"), "vi": int(m.group("vi")), + "mram": m.group("mram"), "mslot": m.group("mslot"), + "target": m.group("target"), + "tr": int(m.group("tr")), "tc": int(m.group("tc")), + }) + return out + + +#: The seven products, in emission order. Each entry is +#: (vram operand, mram operand, target, transposed). +#: +#: `transposed` is `_T_to`, which contracts the MRAM operand's *lanes*; the +#: plain form contracts its *rows*. The last one is the only plain form, and it +#: is the one that reads past a spill's live rows -- see +#: `kda_prefill_spill_v0` for why that matters. +EXPECTED = [ + ("k", "kda_k_tilde_spill", "gram", True), # M = k_hat @ k_tilde^T + ("q", "kda_k_tilde_spill", "readout", True), # N = q_hat @ k_tilde^T + ("state", "kda_k_hat_spill", "contrib", True), # S_0 @ k_hat^T + ("v_t", "kda_t_spill", "err_t", True), # E^T = W^T @ T^T + ("q", "kda_state_spill", "out", True), # q_hat @ S_0^T + ("readout", "kda_err_spill", "readout_contrib", True), # N @ E + ("err_t", "kda_k_end_spill", "state_contrib", False), # E^T @ k_end +] + + +def test_the_seven_products_and_their_operand_order(): + """Operand order, target, and `_to` versus `_T_to` for every product. + + Reversing any product's operands, or swapping the two forms, is a finite + plausible wrong answer that only the emulator would otherwise catch. + """ + _, _, code = _build() + got = [(p["vram"], p["mram"], p["target"], p["transposed"]) for p in _projections(code)] + assert got == EXPECTED, f"\ngot {got}\nexpected {EXPECTED}" + + +def test_the_state_tail_product_contracts_over_mram_rows(): + """`E^T @ k_end` must be the plain form, and it must be the only one. + + `_T_to` contracts the MRAM operand's lanes; the plain form contracts its + rows. Here the contraction is over *time*, which is `k_end`'s row axis, so + `_T_to` would contract over key instead -- shapes match, answer does not. + """ + _, _, code = _build() + plain = [p for p in _projections(code) if not p["transposed"]] + assert len(plain) == 1, f"expected exactly one plain projection, got {plain}" + assert plain[0]["mram"] == "kda_k_end_spill" + + +def test_every_spill_is_stored_before_the_projection_that_reads_it(): + """A prefetch of a region written later reads the previous chunk's data -- + silently, because the regions are reused.""" + _, _, code = _build() + order, seen = [], set() + for line in code.splitlines(): + m = _STORE.search(line) + if m: + order.append(m.group("tile")) + p = _PROJ.search(line) + if p and p.group("mram").startswith("kda_"): + seen.add(p.group("mram")) + spilled = {f"kda_{t}_spill" for t in order} | { + "kda_k_tilde_spill" if "k_tilde" in order else "", + "kda_k_hat_spill" if "k" in order else "", + "kda_t_spill" if "t_mat" in order else "", + "kda_state_spill" if "state" in order else "", + "kda_err_spill" if "err_t" in order else "", + "kda_k_end_spill" if "k_end" in order else "", + } + assert p.group("mram") in spilled, ( + f"{p.group('mram')} is read before it is stored; stores so far: {order}" + ) + assert len(seen) == 6, f"expected six distinct spills, saw {sorted(seen)}" + + +def test_each_spill_zeroes_its_tail_before_storing(): + """The prefetch pulls a whole mlen x mlen block, more than any chunk's live + data. Each spill owns its tail so the read is defined regardless of what the + tile held before -- the regions are reused across chunks. + + This is not observable numerically today: 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. That is precisely the accident this + pins, so that it stays an invariant rather than a coincidence. + """ + _, _, code = _build() + lines = code.splitlines() + stores = [(i, _STORE.search(ln).group("tile")) + for i, ln in enumerate(lines) if _STORE.search(ln)] + assert len(stores) == 6, [t for _, t in stores] + # `state` and `err_t` are value-width, so at value_dim == mlen they have no + # tail and correctly emit no fill. The other four are chunk- or key-width. + needs_tail = {"k_tilde", "k", "t_mat", "k_end"} + for i, tile in stores: + window = "\n".join(lines[max(0, i - 40) : i]) + has_fill = "VRAM Fill Zero" in window + if tile in needs_tail: + assert has_fill, f"{tile} is spilled without zeroing its tail first" + + +def test_the_broadcast_is_logarithmic_not_linear(): + """Copying one row into N used to be N copies -- 64 at mlen 64, 128 at Kimi's + value_dim. Doubling the filled span makes it ceil(log2(N)) + 1. + + Asserted by scaling rather than by a threshold: the state's decay is + broadcast across `value_dim` rows, so a linear form grows with `mlen` and a + logarithmic one barely does. A fixed number would only pin today's shape. + """ + counts = {} + for mlen in (16, 32, 64): + _, _, code = _build(chunk=8, mlen=mlen) + counts[mlen] = code.count("VRAM Fill Zero") + growth = counts[64] - counts[16] + assert growth <= 8, ( + f"fill count grows with mlen ({counts}); a linear broadcast would add " + f"one per row, i.e. 48 across this range" + ) + + +@pytest.mark.parametrize("key_dim,value_dim", [(8, 8), (16, 8), (8, 16), (16, 16), (32, 16)]) +def test_the_seven_products_survive_every_block_count(key_dim, value_dim): + """The composition must not change when an axis stops fitting one block. + + `key_dim > mlen` was refused outright until the key axis moved onto lanes; + Kimi K3 is 128 against mlen 64 on both axes. What multi-block changes is how + many times each product is emitted, not which products there are or which + way round their operands go -- so the distinct product list is invariant and + is checked as one. + """ + _, _, code = _build(chunk=4, mlen=8, key_dim=key_dim, value_dim=value_dim) + got = [(x["vram"], x["mram"], x["target"], x["transposed"]) for x in _projections(code)] + seen, distinct = set(), [] + for entry in got: + if entry not in seen: + seen.add(entry) + distinct.append(entry) + assert distinct == EXPECTED, f"\ngot {distinct}\nexpected {EXPECTED}" + + +def test_multi_block_axes_emit_one_projection_per_block_pair(): + """Each product is emitted once per (row block, column block) of its target. + + Pins the block *count*, which is what the invariant product list above + cannot see. `state @ k_hat^T` writes `[value_dim, chunk]`, so it runs once + per value block; `q_hat @ S_0^T` writes `[chunk, value_dim]`, once per value + block too; `E^T @ k_end` writes `[value_dim, key_dim]`, so value blocks + times key blocks. Getting these wrong leaves part of the target holding + whatever it held before -- zeros on a fresh tile, so silent. + """ + _, _, code = _build(chunk=4, mlen=8, key_dim=16, value_dim=16) + counts = {} + for x in _projections(code): + counts[(x["vram"], x["mram"])] = counts.get((x["vram"], x["mram"]), 0) + 1 + # key_blocks = value_blocks = 2, t_blocks = 1 at chunk 4. + assert counts[("k", "kda_k_tilde_spill")] == 1 # gram is [chunk, chunk] + assert counts[("q", "kda_k_tilde_spill")] == 1 + assert counts[("state", "kda_k_hat_spill")] == 2 # value blocks + assert counts[("v_t", "kda_t_spill")] == 2 # value blocks + assert counts[("q", "kda_state_spill")] == 2 # value blocks + assert counts[("readout", "kda_err_spill")] == 2 # value blocks + assert counts[("err_t", "kda_k_end_spill")] == 4 # value x key blocks + + +def test_multi_block_state_update_covers_every_target_block(): + """`E^T @ k_end` must write all four (value, key) blocks of the state + accumulator, each exactly once. A loop that ran only the diagonal, or that + reused one target index, leaves half the state stale.""" + _, _, code = _build(chunk=4, mlen=8, key_dim=16, value_dim=16) + tail = [(x["tr"], x["tc"]) for x in _projections(code) + if x["mram"] == "kda_k_end_spill"] + assert sorted(tail) == [(0, 0), (0, 1), (1, 0), (1, 1)], tail + + +def test_refuses_a_tile_whose_width_does_not_match_its_axis(): + """The width is what tells a projection how far its contraction runs, so a + tile one block too narrow contracts over half the key axis and returns a + finite wrong answer. Checked exactly, not as a lower bound.""" + p, buffers, _ = _build(chunk=4, mlen=8, key_dim=16, value_dim=8, emit=False) + shape = KdaShape(hidden_size=8, num_heads=1, key_dim=16, value_dim=8, conv_kernel=4) + buffers.k_tilde = p.alloc("narrow", 8, 8, strict=False) + with pytest.raises(ValueError, match="columns wide"): + p.kda_prefill_chunk_v0(buffers=buffers, chunk=4, shape=shape) + + +def test_refuses_a_chunk_wider_than_mlen(): + """At mlen 64 the bf16 range check fires first -- the largest chunk it + allows is 17 -- so this needs a small mlen to reach the width check at all.""" + with pytest.raises(ValueError, match="exceeds mlen"): + _build(chunk=16, mlen=8, key_dim=8) + + +def test_refuses_a_tile_that_is_not_mlen_rows(): + """Every projection writes a whole mlen x mlen block and every spill is + prefetched as one, so a tile that is either must be mlen rows tall however + few rows carry data. Sizing them to `chunk` is what let the prefetch run + past the allocation.""" + p, buffers, _ = _build(emit=False) + buffers.k_end = p.alloc("short", 16, MLEN, strict=False) + with pytest.raises(ValueError, match="rows"): + p.kda_prefill_chunk_v0( + buffers=buffers, chunk=16, + shape=KdaShape( + hidden_size=MLEN, num_heads=1, key_dim=MLEN, + value_dim=MLEN, conv_kernel=4, + ), + ) + + +#: Measured 2026-08-26, mlen 64, one head, keyed by (chunk, key_dim, value_dim). +#: Set ~10% above measured; raising one needs a line in the plan saying why. +#: +#: The single-block row is what it was before the key axis moved onto lanes +#: (1476 against a 1500 measurement at chunk 16), so the layout change is free at +#: the shape that already worked. The multi-block rows are new: they are what +#: Kimi K3 actually costs, 2426 at its 128 x 128 head against 1476 for a 64 x 64 +#: one -- 1.6x the instructions for 4x the state. +PREFILL_STATIC_MAX = { + (4, 64, 64): 900, (4, 128, 64): 1370, (4, 128, 128): 1620, + (8, 64, 64): 1150, (8, 128, 64): 1710, (8, 128, 128): 1980, + (16, 64, 64): 1650, (16, 128, 64): 2380, (16, 128, 128): 2670, +} + + +@pytest.mark.parametrize("shape_key", sorted(PREFILL_STATIC_MAX)) +def test_static_instruction_budget(shape_key): + chunk, key_dim, value_dim = shape_key + _, _, code = _build(chunk=chunk, key_dim=key_dim, value_dim=value_dim) + static = static_count(code) + print(f"KDA_PREFILL chunk={chunk} key={key_dim} value={value_dim} static={static}") + assert static <= PREFILL_STATIC_MAX[shape_key], ( + f"{static} static instructions at chunk {chunk}, key_dim {key_dim}, " + f"value_dim {value_dim} exceeds {PREFILL_STATIC_MAX[shape_key]}" + ) + + +#: Issued instructions, measured 2026-08-27 at `mlen` 64, one head. Set ~15% +#: above. These are the numbers that price the kernel; `PREFILL_STATIC_MAX` +#: above prices the binary and the two do not scale together. +PREFILL_DYNAMIC_MAX = { + (4, 64, 64): 21_800, (4, 128, 64): 29_800, (4, 128, 128): 50_500, + (8, 64, 64): 22_200, (8, 128, 64): 30_400, (8, 128, 128): 51_000, + (16, 64, 64): 23_200, (16, 128, 64): 31_700, (16, 128, 128): 52_400, +} + + +@pytest.mark.parametrize("shape_key", sorted(PREFILL_DYNAMIC_MAX)) +def test_dynamic_instruction_budget(shape_key): + chunk, key_dim, value_dim = shape_key + _, _, code = _build(chunk=chunk, key_dim=key_dim, value_dim=value_dim) + dynamic = dynamic_count(code) + print( + f"KDA_PREFILL chunk={chunk} key={key_dim} value={value_dim} " + f"dynamic={dynamic} static={static_count(code)}" + ) + assert dynamic <= PREFILL_DYNAMIC_MAX[shape_key], ( + f"{dynamic} issued instructions at chunk {chunk}, key_dim {key_dim}, " + f"value_dim {value_dim} exceeds {PREFILL_DYNAMIC_MAX[shape_key]}" + ) + + +def test_the_head_scales_worse_than_the_image_says(): + """`128 x 128` against `64 x 64`: 1.6x by image, 2.3x by issue stream. + + The measurement note in `static_path_measurements.md` read the image ratio + as the cost of quadrupling the state -- "1.6x the instructions for 4x the + state". The kernel is sublinear in the state either way, which was the + point, but by 2.3x rather than 1.6x. Pinned so the two ratios are always + reported together. + """ + ratios = {} + for chunk in (4, 8, 16): + small = _build(chunk=chunk, key_dim=64, value_dim=64)[2] + large = _build(chunk=chunk, key_dim=128, value_dim=128)[2] + ratios[chunk] = ( + static_count(large) / static_count(small), + dynamic_count(large) / dynamic_count(small), + ) + print( + f"KDA_PREFILL chunk={chunk} 128x128/64x64 " + f"static={ratios[chunk][0]:.2f}x dynamic={ratios[chunk][1]:.2f}x" + ) + for chunk, (image, work) in ratios.items(): + assert image < work, ( + f"at chunk {chunk} the image ratio {image:.2f} is not below the " + f"work ratio {work:.2f}; the image used to understate the scaling " + f"and if it stops doing so the reason should be written down" + ) + assert work < 4.0, ( + f"at chunk {chunk} the kernel costs {work:.2f}x for 4x the state; " + f"above 4.0 it is no longer sublinear and the layout move regressed" + ) + + +def test_chunk_is_nearly_free_in_issued_instructions(): + """Quadrupling `chunk` costs 6% of the issue stream and 80% of the image. + + A consequence of prefill being spill-dominated: the per-chunk cost is the + fixed traffic of filling and storing `mlen`-sized tiles, which does not + move with how many tokens the chunk holds. It is worth pinning because the + `chunk` cap of 17 -- set by bf16 range on `1/A`, not by the lowering -- is + therefore more expensive than the image suggests. Every token that cannot + join a chunk pays the fixed cost again. + + Measured at `128 x 128`, `mlen` 64: chunk 4 -> 16 moves the image from + 1,470 to 2,426 and the issue stream from 43,851 to 45,515. + """ + small = _build(chunk=4, key_dim=128, value_dim=128)[2] + large = _build(chunk=16, key_dim=128, value_dim=128)[2] + image = static_count(large) / static_count(small) + work = dynamic_count(large) / dynamic_count(small) + print(f"KDA_PREFILL chunk 4->16 static={image:.2f}x dynamic={work:.3f}x") + assert work < 1.15, ( + f"chunk 4 -> 16 cost {work:.3f}x of the issue stream; it was 1.04x when " + f"measured, and a rise means the chunk dimension started paying for " + f"itself somewhere it did not before" + ) + assert image > 1.4, ( + f"the image grew only {image:.2f}x; it was 1.65x, and a fall means a " + f"sweep that was unrolled over `chunk` became a loop -- good news, but " + f"it changes what this test is documenting" + ) + + +# --------------------------------------------------------------------------- +# The prefill -> decode state layout boundary. + + +def test_the_two_state_layouts_are_transposes_of_each_other(): + """Documents the mismatch by executing it. + + decode indexes the state with `kda_state_row` -- one row per key, value on + lanes -- and prefill holds it one row per value, key on lanes. At Kimi K3 + `key_dim == value_dim == 128`, so the shapes match and handing prefill's + state straight to decode gives a finite, plausible, wrong answer. + + Both layouts are deliberate: decode's makes each sweep a row progression and + therefore a hardware loop, and prefill's makes all seven of its products + land on the projection primitives without an explicit transpose. So the + conversion belongs at the boundary, and + `kda_prefill_state_to_decode_layout_v0` is it. + """ + import torch + + from compiler.aten.plena.program_kda_common import kda_state_row + + key_dim = value_dim = MLEN + shape = KdaShape( + hidden_size=value_dim, num_heads=1, key_dim=key_dim, + value_dim=value_dim, conv_kernel=4, + ) + # A state whose (value, key) entry is distinguishable from (key, value). + dense = torch.arange(value_dim * key_dim).reshape(value_dim, key_dim).float() + + # decode reads one row per key with value on lanes, so what it sees is the + # transpose of what prefill leaves behind. + assert kda_state_row(shape, MLEN, 0, 0, 2) == 2, "one row per key" + as_decode_sees_it = torch.stack([dense[:, k] for k in range(key_dim)]) + torch.testing.assert_close(as_decode_sees_it, dense.T) + assert not torch.allclose(as_decode_sees_it, dense), ( + "if these ever coincide the mismatch is gone and the converter is dead code" + ) + + +def _transpose(mlen: int = MLEN, key_dim: int | None = None, value_dim: int | None = None): + """Compile the layout converter alone, tiles sized from its own shape table.""" + from compiler.aten.plena.program_ssd import SPILLED_ACTIVATION + + shape = KdaShape( + hidden_size=mlen, num_heads=1, + key_dim=key_dim or mlen, value_dim=value_dim or mlen, conv_kernel=4, + ) + p = PlenaCompiler(mlen=mlen, blen=4) + p._bf16_kv_checked = True + want = kda_prefill_state_transpose_shapes(shape, mlen) + tiles = { + "state": p.alloc("state", *want["state"], strict=False), + "identity": p.alloc("t_identity", *want["identity"], strict=False), + "out": p.alloc("state_T", *want["out"], strict=False), + } + mark = len(p.get_code()) + p.kda_prefill_state_to_decode_layout_v0( + shape=shape, precision=SPILLED_ACTIVATION, **tiles + ) + return p, tiles, shape, _projections(p.get_code()[mark:]) + + +def test_the_layout_converter_projects_the_state_against_the_identity(): + """A transpose is `out[i][j] = sum_k I[i][k] * state[j][k]`, so every + projection is the `_T_to` form with the identity as the VRAM operand and the + spilled state as MRAM.""" + _, _, _, projections = _transpose() + assert len(projections) == 1, projections + only = projections[0] + assert only["transposed"], "a transpose is the _T_to form" + assert only["vram"] == "t_identity" + assert only["mram"] == "kda_state_transpose_spill" + assert only["target"] == "state_T" + + +def test_the_layout_converter_transposes_the_block_indices_too(): + """With more than one block the block indices swap as well as the elements. + + `out[ib][jb] = I[ib][:] @ state[jb][:]^T` with `ib` over key blocks and `jb` + over value blocks, so a converter that emitted `(jb, ib)` -- or that only + ran the diagonal -- would move the elements correctly inside each block and + put the blocks in the wrong place. At `key_dim == value_dim` the tile is + square, so nothing about the shapes would object. + """ + _, _, _, projections = _transpose(mlen=8, key_dim=16, value_dim=16) + pairs = sorted((x["tr"], x["tc"]) for x in projections) + assert pairs == [(0, 0), (0, 1), (1, 0), (1, 1)], pairs + # The MRAM operand is indexed by the *value* block while the target row is + # the *key* block, which is the swap itself. `_T_to` prints its operand as + # `name[row][:]`, so the first bracket is the value block and it must equal + # the target's column index, not its row index. + for x in projections: + value_block = int(x["mslot"].split("]")[0].lstrip("[")) + assert value_block == x["tc"], ( + f"out[{x['tr']}][{x['tc']}] reads state block {value_block}; " + f"it must read the value block that becomes its column" + ) + + +def test_the_layout_converter_refuses_a_chunk_width_identity(): + """Its identity spans the whole key axis, not `chunk`. The UT transform's + `[chunk, chunk]` identity is a different tile, and handing that one over + contracts against part of the key axis for a finite wrong answer.""" + from compiler.aten.plena.program_ssd import SPILLED_ACTIVATION + + shape = KdaShape(hidden_size=8, num_heads=1, key_dim=16, value_dim=16, conv_kernel=4) + p = PlenaCompiler(mlen=8, blen=4) + p._bf16_kv_checked = True + with pytest.raises(ValueError, match="columns"): + p.kda_prefill_state_to_decode_layout_v0( + state=p.alloc("state", 16, 16, strict=False), + identity=p.alloc("chunk_identity", 8, 8, strict=False), + out=p.alloc("state_T", 16, 16, strict=False), + shape=shape, precision=SPILLED_ACTIVATION, + ) + + +def test_the_layout_converter_refuses_to_alias_its_output(): + """A projection overwrites its target, so transposing in place would read + half-written data.""" + from compiler.aten.plena.program_ssd import SPILLED_ACTIVATION + + shape = KdaShape( + hidden_size=MLEN, num_heads=1, key_dim=MLEN, value_dim=MLEN, conv_kernel=4 + ) + p = PlenaCompiler(mlen=MLEN, blen=4) + p._bf16_kv_checked = True + state = p.alloc("state", MLEN, MLEN, strict=False) + with pytest.raises(ValueError, match="must not alias"): + p.kda_prefill_state_to_decode_layout_v0( + state=state, identity=p.alloc("t_identity", MLEN, MLEN, strict=False), + out=state, shape=shape, precision=SPILLED_ACTIVATION, + ) diff --git a/aten/tests/test_kda_stage_contract.py b/aten/tests/test_kda_stage_contract.py new file mode 100644 index 0000000..551729c --- /dev/null +++ b/aten/tests/test_kda_stage_contract.py @@ -0,0 +1,112 @@ +"""The KDA stage-marker vocabulary is a cross-repo contract; pin both ends. + +Same contract, same failure mode, as ``test_mamba_stage_contract.py``: the +compiler emits ``; @stage=`` comments and the emulator's ``StageKind`` +(``transactional_emulator/src/stage_profile.rs``) must know every one of them. +Markers are authoritative and sticky -- once a program contains any marker the +emulator abandons its legacy substring rules and bills every instruction to the +most recent marker. A name the compiler emits but the emulator lacks does not +error; it lands in ``unresolved_stage_markers`` and that region's cycles go to +whatever stage preceded it. + +There is a sharper version of the same bug, and this file exists because the +KDA emitters shipped with it: a comment reading ``stage=kda_normalize``, with no +``@``, is not a marker at all. ``extract_stage_tag`` matches on ``"@stage="``, +so such a comment is invisible and the whole region bills to the previous +marker -- with nothing in ``unresolved_stage_markers`` to show for it. That is +why every KDA emitter routes through ``kda_stage_marker`` rather than +formatting its own string. +""" + +from __future__ import annotations + +import re +import unittest +from pathlib import Path + +from compiler.aten.plena.program_kda_common import ( + KDA_STAGE_MARKER_PREFIX, + KDA_STAGES, + kda_stage_marker, +) + +#: Mirrors the KDA entries of ``StageKind`` in +#: ``transactional_emulator/src/stage_profile.rs``. +EMULATOR_KDA_STAGES = frozenset( + { + "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", + } +) + +_PLENA = Path(__file__).resolve().parents[1] / "plena" +#: Every module that emits KDA stage markers. The grep guards below are only +#: worth anything if they cover all of them -- program_kda_conv.py was added +#: after this file and was not scanned until it was listed here. +_KDA_EMITTERS = sorted(_PLENA.glob("program_kda_*.py")) + + +class KdaStageContractTest(unittest.TestCase): + def test_compiler_and_emulator_vocabularies_agree(self): + self.assertEqual(KDA_STAGES, EMULATOR_KDA_STAGES) + + def test_marker_prefix_is_the_one_the_emulator_matches(self): + """A bare ``stage=`` comment is silently not a marker.""" + self.assertEqual(KDA_STAGE_MARKER_PREFIX, "@stage=") + self.assertTrue(kda_stage_marker("kda_normalize").startswith("@stage=")) + + def test_marker_rejects_an_unknown_stage(self): + with self.assertRaises(ValueError): + kda_stage_marker("kda_not_a_stage") + + def test_detail_does_not_leak_into_the_tag(self): + """``extract_stage_tag`` takes the first whitespace-delimited token, so a + detail string must never run into the name.""" + marker = kda_stage_marker("kda_state_load", "weights [16,8]") + tag = marker[len(KDA_STAGE_MARKER_PREFIX) :].split()[0] + self.assertEqual(tag, "kda_state_load") + + def test_the_guard_scans_every_kda_emitter(self): + """A module added later is not covered until it is listed. The grep + guards below pass vacuously on files they never read.""" + names = {p.name for p in _KDA_EMITTERS} + self.assertIn("program_kda_common.py", names) + self.assertIn("program_kda_conv.py", names) + self.assertIn("program_kda_recurrent.py", names) + + def test_no_emitter_formats_a_marker_by_hand(self): + """Every ``stage=`` in the KDA emitters must come from kda_stage_marker. + + The shipped bug was four hand-written f-strings missing the ``@``. A + grep guard is crude, but it is the only thing that catches the fifth one. + """ + source = "\n".join(p.read_text() for p in _KDA_EMITTERS) + offenders = [ + line.strip() + for line in source.splitlines() + if "stage=" in line + and "kda_stage_marker" not in line + and "KDA_STAGE_MARKER_PREFIX" not in line + and not line.lstrip().startswith(("#", "*", '"', "'")) + and "``" not in line + ] + self.assertEqual(offenders, [], f"hand-formatted stage markers: {offenders}") + + def test_every_emitted_marker_name_is_in_the_vocabulary(self): + """Catches a marker whose name was typo'd inside a kda_stage_marker call.""" + source = "\n".join(p.read_text() for p in _KDA_EMITTERS) + names = set(re.findall(r'kda_stage_marker\(\s*"([^"]+)"', source)) + self.assertTrue(names, "no markers found; this guard would pass vacuously") + self.assertLessEqual(names, KDA_STAGES) + + +if __name__ == "__main__": + unittest.main() diff --git a/aten/tests/test_spill_precision_guard.py b/aten/tests/test_spill_precision_guard.py new file mode 100644 index 0000000..2774a68 --- /dev/null +++ b/aten/tests/test_spill_precision_guard.py @@ -0,0 +1,99 @@ +"""The spilled-activation path must refuse an MX KV configuration. + +`SPILLED_ACTIVATION` selects the `keyvalue` precision class with +`set_scale=False`. Under the shipped `[TRANSACTIONAL.PRECISION]`, that class is +Mx/e4m3 with a separate scale stream, so the read walks into the scale stream +whose `0x7f` bytes decode to e4m3 NaN. On the emulator every output of KDA's +chunked prefill came back `nan`. + +`require_bf16_kv_precision` was written for exactly this and had two problems: +it read `kind` where the TOML spells `format`, so it rejected *every* +configuration including the Plain BF16 one it exists to require; and it put the +burden on a caller, with the result that nothing in the repo ever called it -- +not the four SSD emitters, not the KDA one, not a test. + +Both are fixed, and both directions are pinned here. The emitters now check the +active build themselves, because the emitter is what chooses the precision class. +""" + +from __future__ import annotations + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.plena import PlenaCompiler # noqa: E402 + +_PLAIN_BF16 = { + "format": "Plain", + "DATA_TYPE": {"type": "Fp", "sign": True, "exponent": 8, "mantissa": 7}, +} +_MX_E4M3 = { + "format": "Mx", + "block": 8, + "ELEM": {"type": "Fp", "sign": True, "exponent": 4, "mantissa": 3}, + "SCALE": {"type": "Fp", "sign": False, "exponent": 8, "mantissa": 0}, +} + + +def _settings(node: dict) -> dict: + return {"HBM_M_KV_TYPE": node, "HBM_V_KV_TYPE": node} + + +def test_accepts_plain_bf16(): + """The check has to be able to pass. It could not before: reading `kind` + where the TOML says `format` returned None for every configuration, so a + correctly configured build was rejected too -- which is why calling it + anywhere would have failed the build.""" + PlenaCompiler(mlen=64, blen=4).require_bf16_kv_precision(_settings(_PLAIN_BF16)) + + +def test_rejects_the_shipped_mx_kv_types(): + with pytest.raises(ValueError, match="Plain BF16"): + PlenaCompiler(mlen=64, blen=4).require_bf16_kv_precision(_settings(_MX_E4M3)) + + +def test_rejects_a_half_configured_build(): + """Only one of the two classes converted is still wrong, and is the easier + mistake to make by hand.""" + for node_m, node_v in ((_PLAIN_BF16, _MX_E4M3), (_MX_E4M3, _PLAIN_BF16)): + with pytest.raises(ValueError, match="Plain BF16"): + PlenaCompiler(mlen=64, blen=4).require_bf16_kv_precision( + {"HBM_M_KV_TYPE": node_m, "HBM_V_KV_TYPE": node_v} + ) + + +def test_refuses_to_report_success_without_the_table(): + """A check that cannot run must not look like a check that passed.""" + with pytest.raises(ValueError, match="needs the active PRECISION table"): + PlenaCompiler(mlen=64, blen=4).require_bf16_kv_precision(None) + + +def test_the_active_build_check_is_cached(): + """One spill-using emitter per layer would otherwise re-read and re-parse + the TOML per chunk per head.""" + prog = PlenaCompiler(mlen=64, blen=4) + prog._bf16_kv_checked = True + # With the flag set it must not touch the filesystem at all -- if it did, + # the shipped MX types would raise. + prog.require_bf16_kv_precision_from_active_build() + + +def test_the_shipped_settings_would_be_refused(): + """The repository's own plena_settings.toml configures MX KV types, so any + build that spills without overriding them is refused rather than silently + producing NaN. If this ever starts passing, the shipped defaults changed and + the workaround in kda_stage_test.py can go.""" + import tomllib + from pathlib import Path + + root = Path(__file__).resolve().parents[3] + toml = root / "plena_settings.toml" + if not toml.exists(): # pragma: no cover - depends on checkout layout + pytest.skip("plena_settings.toml not found from this checkout") + precision = tomllib.loads(toml.read_text())["TRANSACTIONAL"]["PRECISION"] + with pytest.raises(ValueError, match="Plain BF16"): + PlenaCompiler(mlen=64, blen=4).require_bf16_kv_precision(precision) From 27119a91ad5b9ae77c9a1e3eb082ece434c93ba2 Mon Sep 17 00:00:00 2001 From: "Qichao (Arlo) Wang" Date: Sat, 29 Aug 2026 01:40:54 +0100 Subject: [PATCH 6/6] tests: an ISA interpreter, and budgets that price the work rather than the image **The interpreter** executes emitted assembly in float64, so a lowering can be checked against its reference without the Rust emulator. It is how the mutation testing behind these guards was done: change one thing in an emitter, confirm a test dies. **The budgets** are the only guard against a change that is numerically correct but quietly unrolls a sweep that used to be a hardware loop -- unrolling is what makes a whole-model image tens of megabytes, and no numeric test can see it. They are image budgets and are right to be. What is not right is reading an image count as a cost, which this branch did four times: | 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 | 97x | **3.6x** | | Kimi K3 head vs 64x64, for 4x state | 1.6x | **2.3x** | The third was the headline result. "The cost went from linear in the contracted dimension to constant, a factor of 97" describes the image; a hardware loop's body still runs once per trip and a recurrence must touch every state row. Expanded, both kernels are still exactly linear and fit every measured point with no residual: KDA: issued = 16 * key_dim + 46 174 / 302 / 1,070 / 2,094 Mamba: issued = 11 * state_size + 33 77 / 121 / 209 / 385 against pre-conversion slopes of about 57 and 42. The conversion cut the coefficient 3.6x and stopped the image growing -- two real results, one of them 27x the other. The second matters most for where effort goes: the convolutions are the larger half of the image and an eighth of the work, because a four-tap FIR is wide in the image and shallow in work while a recurrence in a hardware loop is the reverse. **What the issue slots are spent on.** Per opcode, the layer's 492,681 issued instructions are 124,428 arithmetic and 368,253 scaffolding -- 25.3%. Every `V_FMA_VF` sits in a body of five, of which one computes: an FPRAM scalar load, the FMA, two pointer advances, and a loop end that is not counted. That 25% is the figure a dedicated recurrence engine beats by construction, so it is worth knowing how much of the gap is encoding rather than architecture. A post-increment operand folds `S_ADDI_INT gpN, gpN, step` into the instruction that consumes the pointer -- but only where the register has a consumer in the same body, so `self_advance_counts` prices it against those. **Every one of them qualifies: 215,634 foldable, 0 not.** now 492,681 issued 25.3% arithmetic P1 post-inc on the operands 277,047 issued 44.9% 1.78x P2 P1 + the FPRAM scalar too 174,129 issued 71.5% 2.83x An addressing field on instructions that already exist, against a dedicated recurrence datapath with its own banked state SRAM plus a layout engine to feed it. Neither is synthesised, so neither has an area number. What it does not touch: KDA's channel-wise decay makes the operand `key_dim` lanes wide, 6.2% of VLEN 2048 -- a width mismatch, not an encoding one, and the two multiply on a wide machine. New gates: dynamic shares for the gather and the conv/mixer ordering, the two linear fits, `PREFILL_DYNAMIC_MAX` over nine shapes, the arithmetic share, the foldable-advance split with `unfoldable == 0` asserted, and the input projection -- 389,717 static and 4,967,367 issued for one Kimi K3 layer, which had no budget at all because the projection was never in the lowering. That row is never added to or ranked against the state-engine path: an `M_MM` is a 64x64x64 matmul and a `V_FMA_VF` is a 64-lane vector op. One more fell out of the prefill table: chunk 4 -> 16 costs 1.65x of the image and **1.04x of the issue stream**, because prefill is spill-dominated and the fixed per-chunk traffic does not move with how many tokens ride along. That makes the bf16 `chunk` cap of 17 costlier than it looked. The config-derived tests skip where there is no config: `plena_settings.toml` lives in the simulator repository and this one is a submodule of it, so this repo's own CI has no file to read -- which is exactly why nothing here caught the reader being broken. The numeric validation comes from the simulator's `aten quick` job, which runs `linear` and `flash-attention` through the emulator at mlen 64. CI also gains clippy and rustfmt, recursive submodules for the `include_str!` guard, and 23 Mamba/KDA stages added to the stage classification. Co-Authored-By: Claude Opus 5 --- .github/workflows/ci.yml | 90 +++ aten/tests/isa_interpreter.py | 374 +++++++++++ aten/tests/test_instruction_budget.py | 879 ++++++++++++++++++++++++++ aten/tests/test_isa_interpreter.py | 607 ++++++++++++++++++ aten/tests/test_kda_decode_step.py | 6 + aten/tests/test_no_state_engine.py | 179 ++++++ aten/tests/test_real_checkpoint.py | 200 ++++++ docs/ARCHITECTURE.md | 22 +- docs/COMPILATION_PIPELINES.md | 36 ++ 9 files changed, 2392 insertions(+), 1 deletion(-) create mode 100644 aten/tests/isa_interpreter.py create mode 100644 aten/tests/test_instruction_budget.py create mode 100644 aten/tests/test_isa_interpreter.py create mode 100644 aten/tests/test_no_state_engine.py create mode 100644 aten/tests/test_real_checkpoint.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3732abc..611b9f0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -58,6 +58,16 @@ jobs: - name: Run the stage-terminator guard run: PYTHONPATH=. python3 -m pytest aten/tests/test_moe_stage_terminator.py -v + # The Mamba stage vocabulary is the other half of a cross-repo contract: a + # name the compiler emits but the emulator's StageKind lacks does not error, + # it bills to the preceding stage. Source-only, so it belongs in this job. + - name: Run the Mamba stage-marker contract guard + run: PYTHONPATH=. python3 -m pytest aten/tests/test_mamba_stage_contract.py -v + # Same contract for KDA. Source-only like the rest of this job, which is + # also why the KDA lowering must stay importable without torch. + - name: Run the KDA stage-marker contract guard + run: PYTHONPATH=. python3 -m pytest aten/tests/test_kda_stage_contract.py -v + unit-tests: name: Generator unit tests runs-on: ubuntu-latest @@ -71,6 +81,9 @@ jobs: run: | pip install torch --index-url https://download.pytorch.org/whl/cpu pip install transformers + # The generator tests below are scripts, but every step after them + # invokes pytest. This job never installed it. + pip install pytest - name: Run parser tests run: | PYTHONPATH=. python3 generator/tests/test_llm_parser.py @@ -83,6 +96,83 @@ jobs: - name: Run VLM codegen tests run: | PYTHONPATH=. python3 generator/tests/test_vlm_code_gen.py + # The chunked SSD reference vs the plain recurrence. Pure torch, no + # emulator: if these two disagree, every downstream Mamba numerical result + # is meaningless while still looking green. + - name: Run Mamba-2 reference cross-check + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_mamba2_reference.py + # The KDA reference, and the contract pinning its [value, key] orientation + # to the transposed [key, value] layout the PLENA lowering uses. If those + # two ever drift the lowering is validated against the wrong maths. + - name: Run KDA reference cross-check + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_reference.py + # The KDA lowering: transposed state tile, load/store precision agreement, + # conv history roll, and L2 normalisation on the scalar sqrt path. + - name: Run KDA lowering guards + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_lowering.py + # The decode step, executed against the CPU reference by the small ISA + # interpreter in aten/tests/isa_interpreter.py. This is the check that the + # lowering computes the right numbers, not merely plausible opcodes. + - name: Run the KDA decode step numerics + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_decode_step.py + # The oracle those numerics depend on. Anchored separately because the KDA + # lowering only emits a subset of the forms it models -- mutating the + # interpreter's V_MUL_VF source handling left the decode tests green. + - name: Run the ISA interpreter unit tests + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_isa_interpreter.py + # KDA's three short causal convolutions, against _causal_conv_step. + - name: Run the KDA causal conv numerics + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_conv.py + # q and k are key-width, and the L2 norm contracts over the key axis, so + # this is the one KDA reduction the column-block folding cannot split. + - name: Run the KDA normalisation numerics + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_normalize.py + # decay and beta, and the S_MAP_FP_V route into FPRAM the recurrence reads. + - name: Run the KDA gate-scalar numerics + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_gates.py + # The whole mixer: conv output through gates and recurrence to the token + # output, per-head streamed, plus the conv/mixer layout seam. + - name: Run the KDA mixer numerics + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_mixer.py + PYTHONPATH=. python3 -m pytest -q aten/tests/test_v_fma_vf.py + PYTHONPATH=. python3 -m pytest -q aten/tests/test_ssm_decode_step.py + # Chunk primitives for prefill: the cumulative decay and the UT transform. + - name: Run the KDA chunk primitives + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_prefill.py + PYTHONPATH=. python3 -m pytest -q aten/tests/test_spill_precision_guard.py + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_prefill_structure.py + PYTHONPATH=. python3 -m pytest -q aten/tests/test_kda_layer.py + # The design's central claim: one new opcode, and no descriptor machinery. + - name: Run the static-path guard + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_no_state_engine.py + # Real published weights. Skips when the checkpoint is not cached, so it + # never requires network -- fetch once with + # snapshot_download("AntonV/mamba2-130m-hf"). + - name: Run the real-checkpoint binding + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_real_checkpoint.py + # The static-footprint gate. A numerically-correct change that unrolls a + # sweep passes every other test and makes the program image enormous. + - name: Run the instruction budgets + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_instruction_budget.py + # Belongs here, not in moe-stage-guard: the guard resolves every registered + # CPU backend, and resolving imports the module, and every CPU backend + # imports torch. moe-stage-guard deliberately installs only pytest+pyyaml. + - name: Run the registry CPU-signature guard + run: | + PYTHONPATH=. python3 -m pytest -q aten/tests/test_registry_cpu_signatures.py codegen-smoke: name: Codegen + assembler smoke diff --git a/aten/tests/isa_interpreter.py b/aten/tests/isa_interpreter.py new file mode 100644 index 0000000..db85505 --- /dev/null +++ b/aten/tests/isa_interpreter.py @@ -0,0 +1,374 @@ +"""A tiny interpreter for the ISA subset the KDA decode lowering emits. + +Why this exists +--------------- +The transactional emulator is the authority on ISA semantics, but running it +needs ``libramulator`` built, which is a C++ dependency the compiler repo's own +test suite does not have. Without something in between, a lowering change can +only be checked for *shape* -- "an S_SQRT_FP appears somewhere" -- and the review +of Task 2 showed exactly how far that gets you: thirteen mutations that change +the arithmetic left every such assertion green. + +So this executes the emitted assembly against the CPU reference. It is a test +oracle, deliberately narrow: + +* Seven opcodes, listed in ``_OPS``. Anything else raises rather than being + skipped, so a lowering that starts emitting something new fails loudly here + instead of being silently half-checked. +* No timing, no precision model, no memory hierarchy. Values are float64, + while the emulator's FP registers are bf16 and every VRAM write is quantised. + "Validated numerically" here means the *arithmetic and addressing* are right, + not that the result holds at hardware precision. +* Batch is not modelled, because the lowering has no batch concept. + +It is a second implementation of ISA semantics and could therefore drift from +the emulator. That is a real risk and the reason it is scoped to seven opcodes +whose meaning is one line each in ``doc/plena_isa_spec.md``. The whole-model +emulator run is what finally cross-checks it; until then this catches the class +of bug that actually occurs in lowering work -- wrong row, wrong FPRAM slot, +wrong order. + +Semantics, from ``doc/plena_isa_spec.md`` and +``transactional_emulator/src/accelerator/loop_state.rs``: + +=================================== ====================================== +``S_ADDI_INT rd, rs1, imm`` ``gp[rd] = gp[rs1] + imm`` +``S_LD_FP fd, rs1, imm`` ``f[fd] = FPRAM[gp[rs1] + imm]`` +``V_MUL_VF rd, rs1, f2, rmask`` ``V[rd] = V[rs1] * f[f2]`` +``V_ADD_VF rd, rs1, f2, rmask`` ``V[rd] = V[rs1] + f[f2]`` +``V_SUB_VF rd, rs1, f2, rmask, ord`` ``ord=0: V[rs1] - f[f2]``; ``ord=1: f[f2] - V[rs1]`` +``V_ADD_VV rd, rs1, rs2, rmask`` ``V[rd] = V[rs1] + V[rs2]`` +``V_SUB_VV rd, rs1, rs2, rmask`` ``V[rd] = V[rs1] - V[rs2]`` +``V_MUL_VV rd, rs1, rs2, rmask`` ``V[rd] = V[rs1] * V[rs2]`` +``V_EXP_V rd, rs1, rmask`` ``V[rd] = exp(clamp(V[rs1], -88, 88))`` +``V_RECI_V rd, rs1, rmask`` ``V[rd] = 1 / V[rs1]``; ``1/0 -> inf`` +``V_RED_SUM fd, rs1, ...`` ``f[fd] += sum(V[rs1])`` -- **accumulates** +``S_ST_FP fd, rs1, imm`` ``FPRAM[gp[rs1] + imm] = f[fd]`` (fd is the source) +``S_MAP_FP_V rd, rs1, imm`` ``FPRAM[gp[rd]+imm ..+vlen] = V[gp[rs1]]`` +``S_ADD_FP fd, fs1, fs2`` ``f[fd] = f[fs1] + f[fs2]`` +``S_SUB_FP fd, fs1, fs2`` ``f[fd] = f[fs1] - f[fs2]`` +``S_MUL_FP fd, fs1, fs2`` ``f[fd] = f[fs1] * f[fs2]`` +``S_RECI_FP fd, fs1`` ``f[fd] = 1 / f[fs1]`` +``S_SQRT_FP fd, fs1`` ``f[fd] = sqrt(f[fs1])`` + +A write to ``f0`` by any of those is **discarded**, matching +``dispatch.rs:415-424``. +``C_LOOP_START rd, imm`` ``gp[rd] = imm``; remember this pc +``C_LOOP_END rd`` if ``gp[rd] > 1``: ``gp[rd] -= 1``, jump + to ``start_pc + 1``; else fall through +=================================== ====================================== + +``V_SUB_VV`` is ``rs1 - rs2``. Take that from the emulator +(``vector_machine.rs::sub``), not from the operation line in +``doc/plena_isa_spec.md``, which read ``rs2 - rs1`` until this was found -- +"correcting" this interpreter to match that wording inverts the sign of KDA's +error term. + +``gp0`` and ``f0`` are **ordinary writable registers that initialise to zero**, +exactly as in the emulator (``accelerator/registers.rs`` reads and writes them +through the same plain arrays as any other). They are not hardwired. What keeps +them zero is ``aten/plena/registers.py``'s allocator, which hands out gp1+/f1+ +and never gives either away -- so the zeroing idiom ``V_MUL_VF gpX, gpX, f0`` +rests on a compiler convention, not on hardware. + +Modelling that faithfully rather than hardwiring zero here is deliberate: a +hardwired oracle would keep passing an emitter that started writing ``gp0``, +which on real hardware would corrupt every address computed from it. +``test_isa_interpreter.py`` and ``test_kda_decode_step.py`` guard the +convention directly instead. + +Vector operands are element addresses; each op touches ``vlen`` consecutive +elements. ``rmask`` is ignored: every KDA emission passes 0 (all lanes), and +asserting that is cheaper than modelling the mask register. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field + +__all__ = ["Machine", "UnsupportedInstruction"] + +_OPS = frozenset( + { + "S_ADDI_INT", + "S_LD_FP", + "V_MUL_VF", + "V_FMA_VF", + "V_ADD_VF", + "V_SUB_VF", + "V_ADD_VV", + "V_SUB_VV", + "V_MUL_VV", + "V_EXP_V", + "V_RECI_V", + "V_RED_SUM", + "S_ST_FP", + "S_MAP_FP_V", + "S_ADD_FP", + "S_SUB_FP", + "S_MUL_FP", + "S_RECI_FP", + "S_SQRT_FP", + "C_LOOP_START", + "C_LOOP_END", + } +) + +#: Scalar FP ops whose destination is an FP register. ``dispatch.rs:415-424`` +#: makes every one of these a no-op when ``rd == 0``, so ``f0`` cannot be +#: clobbered by them -- which is what the ``V_MUL_VF gpX, gpX, f0`` zeroing +#: idiom rests on. ``S_LD_FP`` is deliberately absent from that guard list, and +#: therefore from this one. +_FP_DEST_OPS = frozenset( + {"S_ADD_FP", "S_SUB_FP", "S_MUL_FP", "S_RECI_FP", "S_SQRT_FP"} +) + +#: ``exp`` saturates rather than overflowing, matching +#: ``vector_machine.rs::exp``, which clamps to keep bf16 from going infinite. +_EXP_CLAMP = 88.0 + +#: Guards against a runaway loop turning a failing test into a hang. +_MAX_STEPS = 5_000_000 + + +class UnsupportedInstruction(RuntimeError): + """Raised for any opcode outside `_OPS`, rather than skipping it.""" + + +def _reg(token: str, prefix: str) -> int: + if not token.startswith(prefix): + raise UnsupportedInstruction(f"expected {prefix} register, got {token!r}") + return int(token[len(prefix) :]) + + +@dataclass +class Machine: + vlen: int + vram_words: int = 1 << 16 + fpram_words: int = 1 << 12 + vram: list[float] = field(init=False) + fpram: list[float] = field(init=False) + gp: list[int] = field(init=False) + fp: list[float] = field(init=False) + executed: int = field(init=False, default=0) + + def __post_init__(self) -> None: + self.vram = [0.0] * self.vram_words + self.fpram = [0.0] * self.fpram_words + self.gp = [0] * 16 + self.fp = [0.0] * 8 + + # -- memory helpers --------------------------------------------------- + + def write_vram_row(self, addr: int, values) -> None: + vals = list(values) + if len(vals) > self.vlen: + raise ValueError(f"row is {len(vals)} wide, vlen is {self.vlen}") + self.vram[addr : addr + len(vals)] = [float(x) for x in vals] + + def read_vram_row(self, addr: int, count: int | None = None) -> list[float]: + n = self.vlen if count is None else count + return list(self.vram[addr : addr + n]) + + def write_fpram(self, addr: int, values) -> None: + vals = [float(x) for x in values] + self.fpram[addr : addr + len(vals)] = vals + + # -- execution -------------------------------------------------------- + + def run(self, code: str) -> "Machine": + program = _parse(code) + loop_starts: dict[int, int] = {} + pc = 0 + steps = 0 + while pc < len(program): + steps += 1 + if steps > _MAX_STEPS: + raise RuntimeError(f"exceeded {_MAX_STEPS} steps; runaway loop at pc {pc}") + op, args = program[pc] + if op == "S_ADDI_INT": + rd, rs1, imm = _reg(args[0], "gp"), _reg(args[1], "gp"), int(args[2]) + self._set_gp(rd, self._gp(rs1) + imm) + elif op == "S_LD_FP": + fd, rs1, imm = _reg(args[0], "f"), _reg(args[1], "gp"), int(args[2]) + self.fp[fd] = self.fpram[self._gp(rs1) + imm] + elif op == "V_RED_SUM": + # The masked reduce_sum in vector_machine.rs:404-418 is a + # materially different computation -- it broadcasts each head's + # sum back over its slice before summing -- so refuse rather + # than model it wrongly. + # Every emitter writes zeros here; rather than pin which slot + # carries rmask, refuse any non-zero trailing operand. + if any(a != "0" for a in args[2:]): + raise UnsupportedInstruction( + f"V_RED_SUM with non-zero mask operands {args[2:]}; " + f"only the unmasked form is modelled" + ) + fd, rs1 = _reg(args[0], "f"), _reg(args[1], "gp") + src = self._gp(rs1) + if src % self.vlen: + raise UnsupportedInstruction( + f"V_RED_SUM address {src} is not a multiple of vlen {self.vlen}" + ) + # Accumulates: dispatch.rs seeds reduce_sum with the current f[rd]. + self._set_fp(fd, self.fp[fd] + sum(self.vram[src : src + self.vlen])) + elif op.startswith("V_"): + self._vector(op, args) + elif op == "S_MAP_FP_V": + # Mirror of S_MAP_V_FP, and the operand roles mirror too: rs1 is + # the VRAM source row, rd the FP_MEM base, so that in both + # instructions "rd names the destination memory". + rd, rs1, imm = _reg(args[0], "gp"), _reg(args[1], "gp"), int(args[2]) + src = self._gp(rs1) + if src % self.vlen: + raise UnsupportedInstruction( + f"S_MAP_FP_V source {src} is not a multiple of vlen {self.vlen}" + ) + start = self._gp(rd) + imm + # Python slice assignment past the end *extends* the list rather + # than failing, so an out-of-bounds write would silently land at + # the wrong address and stay green. The emulator asserts on the + # FP SRAM bound (scalar_sram.rs), so refuse here too. + if start < 0 or start + self.vlen > len(self.fpram): + raise UnsupportedInstruction( + f"S_MAP_FP_V writes FPRAM[{start}:{start + self.vlen}] past " + f"the modelled {len(self.fpram)} slots" + ) + self.fpram[start : start + self.vlen] = self.vram[src : src + self.vlen] + elif op == "S_ST_FP": + fd, rs1, imm = _reg(args[0], "f"), _reg(args[1], "gp"), int(args[2]) + self.fpram[self._gp(rs1) + imm] = self.fp[fd] + elif op in _FP_DEST_OPS: + self._scalar_fp(op, args) + elif op == "C_LOOP_START": + rd, imm = _reg(args[0], "gp"), int(args[1]) + if imm <= 0: + raise UnsupportedInstruction(f"C_LOOP_START with count {imm}") + self._set_gp(rd, imm) + loop_starts[rd] = pc + elif op == "C_LOOP_END": + rd = _reg(args[0], "gp") + if rd not in loop_starts: + raise UnsupportedInstruction(f"C_LOOP_END gp{rd} with no matching start") + if self._gp(rd) > 1: + self._set_gp(rd, self._gp(rd) - 1) + pc = loop_starts[rd] + 1 + self.executed += 1 + continue + else: # pragma: no cover - _parse already rejects these + raise UnsupportedInstruction(op) + self.executed += 1 + pc += 1 + return self + + def _scalar_fp(self, op: str, args: list[str]) -> None: + fd = _reg(args[0], "f") + a = self.fp[_reg(args[1], "f")] + if op == "S_RECI_FP": + # bf16::ONE / a. Python's `a == 0.0` is true for -0.0 too, so the + # sign has to come from copysign or -0.0 would yield +inf. + out = math.copysign(math.inf, a) if a == 0.0 else 1.0 / a + elif op == "S_SQRT_FP": + out = math.sqrt(a) if a >= 0.0 else math.nan + else: + b = self.fp[_reg(args[2], "f")] + out = {"S_ADD_FP": a + b, "S_SUB_FP": a - b, "S_MUL_FP": a * b}[op] + self._set_fp(fd, out) + + def _set_fp(self, index: int, value: float) -> None: + # dispatch.rs discards these writes when rd == 0. + if index != 0: + self.fp[index] = value + + def _vector(self, op: str, args: list[str]) -> None: + # V_SUB_VF carries rorder after rmask; every other V-type ends on rmask. + rmask = args[3] if op == "V_SUB_VF" else args[-1] + if rmask != "0": + raise UnsupportedInstruction(f"{op} with rmask={rmask}; only 0 is modelled") + dst = self._gp(_reg(args[0], "gp")) + src = self._gp(_reg(args[1], "gp")) + # The emulator's vector SRAM asserts VLEN alignment (lib/sram addr_to_cell); + # an unaligned vector address panics there rather than reading across a + # row boundary. Model that, so a misaligned emission fails here too. + for name, addr in (("dst", dst), ("src", src)): + if addr % self.vlen: + raise UnsupportedInstruction( + f"{op} {name} address {addr} is not a multiple of vlen {self.vlen}" + ) + n = self.vlen + if op in ("V_EXP_V", "V_RECI_V"): + if op == "V_EXP_V": + out = [ + math.exp(min(max(self.vram[src + i], -_EXP_CLAMP), _EXP_CLAMP)) + for i in range(n) + ] + else: + # vector_machine.rs uses tensor.reciprocal(), which yields + # +/-inf on zero rather than trapping. Model that: a lowering + # that reciprocates an unused lane must behave the same here. + out = [ + ( + math.copysign(math.inf, self.vram[src + i]) + if self.vram[src + i] == 0.0 + else 1.0 / self.vram[src + i] + ) + for i in range(n) + ] + elif op == "V_FMA_VF": + # The only V-type op that reads its destination. Modelled as one + # expression rather than a multiply followed by an add, matching + # vector_machine.rs::fma_scalar, which quantises once on the sum. + f = self.fp[_reg(args[2], "f")] + out = [self.vram[dst + i] + self.vram[src + i] * f for i in range(n)] + elif op in ("V_MUL_VF", "V_ADD_VF", "V_SUB_VF"): + f = self.fp[_reg(args[2], "f")] + if op == "V_MUL_VF": + out = [self.vram[src + i] * f for i in range(n)] + elif op == "V_ADD_VF": + out = [self.vram[src + i] + f for i in range(n)] + else: + reverse = len(args) > 4 and args[4] == "1" + out = [ + (f - self.vram[src + i]) if reverse else (self.vram[src + i] - f) + for i in range(n) + ] + else: + src2 = self._gp(_reg(args[2], "gp")) + if src2 % self.vlen: + raise UnsupportedInstruction( + f"{op} src2 address {src2} is not a multiple of vlen {self.vlen}" + ) + if op == "V_ADD_VV": + out = [self.vram[src + i] + self.vram[src2 + i] for i in range(n)] + elif op == "V_SUB_VV": + out = [self.vram[src + i] - self.vram[src2 + i] for i in range(n)] + elif op == "V_MUL_VV": + out = [self.vram[src + i] * self.vram[src2 + i] for i in range(n)] + else: # pragma: no cover - _parse rejects anything else + raise UnsupportedInstruction(op) + self.vram[dst : dst + n] = out + + def _gp(self, index: int) -> int: + return self.gp[index] + + def _set_gp(self, index: int, value: int) -> None: + self.gp[index] = value + + +def _parse(code: str) -> list[tuple[str, list[str]]]: + program: list[tuple[str, list[str]]] = [] + for raw in code.splitlines(): + line = raw.split(";", 1)[0].split("//", 1)[0].strip() + if not line: + continue + parts = line.replace(",", " ").split() + op, args = parts[0], parts[1:] + if op not in _OPS: + raise UnsupportedInstruction( + f"{op!r} is outside the modelled subset {sorted(_OPS)}. " + f"The lowering emitted something this oracle does not understand; " + f"extend it deliberately rather than letting the check pass." + ) + program.append((op, args)) + return program diff --git a/aten/tests/test_instruction_budget.py b/aten/tests/test_instruction_budget.py new file mode 100644 index 0000000..47b56c2 --- /dev/null +++ b/aten/tests/test_instruction_budget.py @@ -0,0 +1,879 @@ +"""Static instruction-count gates for the recurrent kernels. + +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, and no numeric test can see it. + +What these gates assert is that the **image** is flat in the contracted +dimension. Before Task 8 both kernels emitted a `copy + multiply + add` triple +per state row, which could not become a hardware loop -- the scratch row it +staged through was the *same* row every iteration, so the destination never +formed an arithmetic progression -- and the image therefore grew linearly: + + KDA decode, one head: key_dim 8 -> 505 Mamba decode, one head: + key_dim 16 -> 961 state 4 -> 200 + key_dim 64 -> 3,697 state 8 -> 352 + key_dim128 -> 7,345 state 16 -> 656 + state 32 -> 1,384 + +After the conversion onto `V_FMA_VF` each sweep is one instruction inside one +hardware loop, so the image is constant: **76** per (head, block) for KDA at +every key_dim from 4 to 128, and **54** per head for Mamba at every state_size. + +**That constancy is a property of the image and of nothing else, and the +`97x` this file used to quote for it was a comparison of an unrolled program's +size against a looped program's size.** The work did not become constant and +could not have: a recurrence has to touch every state row, and no encoding +changes how many rows there are. Measured with the loops expanded (`_dynamic`, +`mlen` 8, one head), both kernels are still exactly linear -- + + KDA: 16 * key_dim + 46 (174, 302, 1,070, 2,094 at 8/16/64/128) + Mamba: 11 * state_size + 33 (77, 121, 209, 385 at 4/8/16/32) + +-- against slopes of about 57 and 38 before. What the conversion bought is a +**3.6x smaller coefficient** on a line that stays a line, plus an image that +stops growing. Both are real; only the second is a factor of 97. +`test_the_conversion_reduced_the_slope_not_the_linearity` pins the fits. + +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 needs a line in +docs/superpowers/plans/2026-08-25-static-mamba-kda.md saying why. +""" + +from __future__ import annotations + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.kda.shape import KdaShape # noqa: E402 +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_kda_common import kda_blocks # noqa: E402 +from compiler.aten.plena.instruction_stream import ( # noqa: E402 + ARITHMETIC, + arithmetic_share, + dynamic_count as _dynamic, + opcode_census, + self_advance_counts, + static_count as _static, +) +from compiler.aten.plena.program_mamba_common import Mamba2Shape # noqa: E402 + +MLEN = 8 + +#: Measured 2026-08-26 after Task 8: 76 per (head, block), flat in key_dim. +KDA_DECODE_STATIC_INSTR_PER_UNIT_MAX = 85 + +#: Measured 2026-08-26 after Task 8: 54 per head, flat in state_size. +MAMBA_DECODE_STATIC_INSTR_PER_HEAD_MAX = 60 + + +def _rows_up(n: int) -> int: + return ((n + MLEN - 1) // MLEN) * MLEN + + +def _kda_decode_asm(num_heads: int, key_dim: int) -> tuple[str, int]: + shape = KdaShape( + hidden_size=num_heads * MLEN, num_heads=num_heads, key_dim=key_dim, + value_dim=MLEN, conv_kernel=4, + ) + p = PlenaCompiler(mlen=MLEN, blen=2) + blocks = kda_blocks(shape, MLEN) + state = p.alloc("state", _rows_up(num_heads * blocks * key_dim), MLEN) + vec = lambda name: p.alloc(name, _rows_up(num_heads * blocks), MLEN) # noqa: E731 + v, o, pred, err = vec("v"), vec("o"), vec("pred"), vec("err") + fp = lambda name, n: p.fp_var(name, size=n) # noqa: E731 + mark = len(p.get_code()) + p.kda_decode_step_v0( + state=state, + q_fp=fp("q", num_heads * key_dim), k_fp=fp("k", num_heads * key_dim), + decay_fp=fp("d", num_heads * key_dim), beta_fp=fp("b", _rows_up(num_heads)), + v=v, o=o, pred=pred, err=err, shape=shape, + output_scale_fp=fp("s", 1), + ) + return p.get_code()[mark:], num_heads * blocks + + +def _kda_decode(num_heads: int, key_dim: int) -> tuple[int, int]: + asm, units = _kda_decode_asm(num_heads, key_dim) + return _static(asm), units + + +def _mamba_decode_asm(num_heads: int, state_size: int) -> str: + shape = Mamba2Shape( + hidden_size=num_heads * MLEN, num_heads=num_heads, head_dim=MLEN, + state_size=state_size, n_groups=1, conv_kernel=4, chunk_size=16, seq_len=1, + ) + p = PlenaCompiler(mlen=MLEN, blen=2) + state = p.alloc("state", _rows_up(num_heads * state_size), MLEN) + x = p.alloc("x", _rows_up(num_heads), MLEN) + y = p.alloc("y", _rows_up(num_heads), MLEN) + scratch = p.alloc("scratch", MLEN, MLEN) + mark = len(p.get_code()) + p.ssm_decode_step_v0( + state=state, x=x, + b_fp=p.fp_var("b", size=state_size), c_fp=p.fp_var("c", size=state_size), + da_fp=p.fp_var("da", size=num_heads), dt_fp=p.fp_var("dt", size=num_heads), + d_fp=p.fp_var("d", size=num_heads), y=y, scratch=scratch, + shape=shape, consts=p.mamba_fp_constants(), + ) + return p.get_code()[mark:] + + +def _mamba_decode(num_heads: int, state_size: int) -> int: + return _static(_mamba_decode_asm(num_heads, state_size)) + + +@pytest.mark.parametrize("num_heads,key_dim", [(1, 4), (1, 8), (2, 8), (4, 16), (1, 64)]) +def test_kda_decode_step_static_instruction_count(num_heads, key_dim): + static, units = _kda_decode(num_heads, key_dim) + per_unit = static // units + print( + f"KDA_DECODE heads={num_heads} key_dim={key_dim} " + f"static={static} per_unit={per_unit}" + ) + assert per_unit <= KDA_DECODE_STATIC_INSTR_PER_UNIT_MAX, ( + f"{per_unit} static instructions per (head, block) exceeds " + f"{KDA_DECODE_STATIC_INSTR_PER_UNIT_MAX}; a sweep probably fell off the " + f"hardware-loop path" + ) + + +@pytest.mark.parametrize("num_heads,state_size", [(1, 4), (2, 8), (4, 16), (8, 16)]) +def test_mamba_decode_step_static_instruction_count(num_heads, state_size): + static = _mamba_decode(num_heads, state_size) + per_head = static // num_heads + print( + f"MAMBA_DECODE heads={num_heads} state={state_size} " + f"static={static} per_head={per_head}" + ) + 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 the " + f"hardware-loop path" + ) + + +def test_the_cost_is_flat_in_the_contracted_dimension(): + """The claim the budgets encode, asserted directly. + + A budget alone would still pass if the cost grew and someone raised it. + This fails on *any* growth, which is what the hardware loop actually buys. + """ + kda = {key_dim: _kda_decode(1, key_dim)[0] for key_dim in (4, 8, 16, 64)} + assert len(set(kda.values())) == 1, ( + f"KDA decode cost varies with key_dim: {kda}; before Task 8 it was " + f"505 / 961 / 3,697 at key_dim 8 / 16 / 64" + ) + + mamba = {n: _mamba_decode(1, n) for n in (4, 8, 16, 32)} + assert len(set(mamba.values())) == 1, ( + f"Mamba decode cost varies with state_size: {mamba}; before Task 8 it " + f"was 200 / 352 / 656 at state_size 4 / 8 / 16" + ) + + +# =========================================================================== +# Whole-model scale. +# +# The static footprint of a model's recurrent path, which is the part this work +# changed. The projections, MoE blocks, norms and embeddings are shared code this +# work did not touch, and no whole-model program has been compiled here -- so what +# is gated is one layer's state-engine path times the layer count, and that is +# stated rather than presented as a whole-model total. +# +# The layer count and head count below are a projected K3 shape, not a released +# model: doc/Model_Lib/kimi-linear-48b-a3b.json is hidden 2304 with 32 KDA heads +# over 27 layers. The per-head geometry is the same (head_dim 128, conv kernel 4), +# so the per-head numbers carry across; anything multiplied by heads or layers +# does not. +# +# The layer *does* assemble end to end: gather, three convolutions, the gates and +# the recurrence, matching `kda_state_engine_step` -- see +# `test_kda_layer.py::test_the_assembled_layer_matches_kda_state_engine_step`. +# An earlier version of this comment claimed it could not, on the grounds that +# two emitters would first need a column-block index. That was wrong twice over; +# `program_kda_layer`'s module docstring has the correction. +# +# The convolutions are the larger half of the *image* and were ungated until this +# was measured. They are not the larger half of the work -- see +# `test_the_image_and_the_work_rank_the_kernels_differently`. + +#: Kimi K3, from `KdaShape.kimi_k3()`: 96 heads, key_dim and value_dim 128. +KIMI_LAYERS = 93 + +#: Measured 2026-08-26 at mlen 64. Set ~10% above; raising it needs a line in +#: docs/superpowers/plans/2026-08-25-static-mamba-kda.md saying why. +KIMI_MIXER_STATIC_PER_LAYER_MAX = 43_500 + + +def _kimi_mixer_asm() -> str: + 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 + + mlen = 64 + shape = KdaShape.kimi_k3() + p = PlenaCompiler(mlen=mlen, blen=4) + kb = kda_key_blocks(shape, mlen) + up = lambda n: ((n + mlen - 1) // mlen) * mlen # noqa: E731 + a = lambda n, r: p.alloc(n, up(r), mlen, strict=False) # noqa: E731 + decay = p.fp_var("decay_and_q_hat", 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("gate", shape.num_heads * kb), + dt_bias=a("dt_bias", shape.num_heads * kb), + beta_logit=a("beta_logit", kda_head_blocks(shape, mlen)), + state=a("state", kda_state_rows(shape, mlen)), + out=a("out", kda_vector_rows(shape, mlen)), + pred=a("pred", kda_vector_rows(shape, mlen)), + err=a("err", kda_vector_rows(shape, mlen)), + sq_scratch=a("sq_scratch", shape.num_heads * kb), + decay_fp=decay, q_hat_fp=decay, # one window; see kda_decode_predict_v0 + k_hat_fp=p.fp_var("k_hat", size=shape.key_dim), + beta_fp=p.fp_var("beta", size=kda_head_blocks(shape, mlen) * mlen), + part_fp=p.fp_var("part", size=kb), acc_fp=p.fp_var("acc", size=1), + output_scale_fp=p.fp_var("output_scale", size=1), + rate_fp=p.fp_var("rate", size=shape.num_heads), + lower_bound_fp=p.fp_var("lower_bound", size=1), + consts=p.kda_fp_constants(), + ) + mark = len(p.get_code()) + p.kda_beta_scalars_v0( + beta_logit=buffers.beta_logit, beta_fp=buffers.beta_fp, + consts=buffers.consts, shape=shape, + ) + p.kda_mixer_step_v0(buffers=buffers, shape=shape) + return p.get_code()[mark:] + + +def _kimi_mixer_static() -> int: + return _static(_kimi_mixer_asm()) + + +def test_kimi_k3_mixer_static_instruction_count(): + per_layer = _kimi_mixer_static() + print( + f"KIMI_K3_MIXER per_layer={per_layer} layers={KIMI_LAYERS} " + f"total={per_layer * KIMI_LAYERS}" + ) + assert per_layer <= KIMI_MIXER_STATIC_PER_LAYER_MAX, ( + f"{per_layer} static instructions per layer exceeds " + f"{KIMI_MIXER_STATIC_PER_LAYER_MAX}; a sweep probably fell off the " + f"hardware-loop path" + ) + + +def _fit_line(points: list[tuple[int, int]]) -> tuple[float, float]: + """Exact `(slope, intercept)` for points that lie on a line, else raise.""" + (x0, y0), (x1, y1) = points[0], points[1] + slope = (y1 - y0) / (x1 - x0) + intercept = y0 - slope * x0 + for x, y in points: + if abs(slope * x + intercept - y) > 1e-9: + raise AssertionError(f"not collinear at x={x}: {y} vs {slope * x + intercept}") + return slope, intercept + + +def test_the_conversion_reduced_the_slope_not_the_linearity(): + """What `V_FMA_VF` actually bought, measured on the issue stream. + + This file quoted `7,345 -> 76, a factor of 97` and the phrase "the cost + went from linear in the contracted dimension to constant". Both describe + the program image. The image did go constant, and 97 is the right number + for it -- but it is the ratio between an unrolled program's *size* and a + looped program's *size*, and nothing about a hardware loop makes its body + run fewer times. + + Expanded, both kernels are still exactly linear in the contracted + dimension, and have to be: a recurrence touches every state row and the + encoding cannot change how many rows there are. What changed is the + coefficient -- the `copy + multiply + add` triple became one FMA, so the + per-row cost fell by about 3.6x. That is the honest figure for the + conversion's effect on time, and it is 27x smaller than the one this file + used to report. + + The budgets above stay image budgets, which is correct: they exist to + catch a sweep falling off the hardware-loop path, and that is an image + failure needing an image measurement. + """ + kda = [(kd, _dynamic(_kda_decode_asm(1, kd)[0])) for kd in (8, 16, 64, 128)] + mamba = [(ss, _dynamic(_mamba_decode_asm(1, ss))) for ss in (4, 8, 16, 32)] + kda_slope, kda_c = _fit_line(kda) + mamba_slope, mamba_c = _fit_line(mamba) + print( + f"KDA_DECODE dynamic = {kda_slope:.0f}*key_dim + {kda_c:.0f} {kda}\n" + f"MAMBA_DECODE dynamic = {mamba_slope:.0f}*state_size + {mamba_c:.0f} {mamba}" + ) + + # The image is flat over the same range -- that is the budgeted property. + assert len({_static(_kda_decode_asm(1, kd)[0]) for kd, _ in kda}) == 1 + assert len({_static(_mamba_decode_asm(1, ss)) for ss, _ in mamba}) == 1 + + # The work is not. Slopes measured 2026-08-27; bounds are ~40% either side. + assert 10 <= kda_slope <= 24, f"KDA slope {kda_slope} left its measured band" + assert 7 <= mamba_slope <= 16, f"Mamba slope {mamba_slope} left its measured band" + + # Pre-conversion slopes, from the table in this module's docstring: + # KDA (7,345-505)/(128-8) = 57.0, Mamba (1,384-200)/(32-4) = 42.3. + assert (57.0 / kda_slope) < 10, ( + f"the conversion cut the KDA slope by {57.0 / kda_slope:.1f}x; if this " + f"ever approaches the image's 97x, the two measurements have been " + f"conflated again" + ) + assert (42.3 / mamba_slope) < 10, ( + f"the conversion cut the Mamba slope by {42.3 / mamba_slope:.1f}x" + ) + + +def test_the_recurrence_is_the_bulk_of_the_layers_work(): + """The recurrence dominates the layer, and the image says the opposite. + + The old version of this test multiplied a per-unit *image* count by the + unit count and compared it against the layer's *image*, concluding that + without the `V_FMA_VF` conversion the recurrence would be 97% of the + layer. Two image numbers divided cannot say what fraction of a layer + anything is. + + Measured on the issue stream, the mixer -- which is the recurrence plus + the gates around it -- is 87% of the layer against the three + convolutions' 12%, while by image the convolutions are the larger half. + So the conclusion the old test reached happens to survive; its arithmetic + did not. + """ + conv = sum(_dynamic(a) for a in _kimi_conv_asm()) + mixer = _dynamic(_kimi_mixer_asm()) + layer = conv + mixer + _dynamic(_kimi_gather_asm()) + print(f"KIMI_K3 mixer share of issued layer = {mixer / layer:.1%}") + assert mixer / layer > 0.7, ( + f"the mixer is {mixer / layer:.1%} of the issued layer; it was 87% when " + f"measured, and a large fall means work moved out of the recurrence" + ) + + +#: Measured 2026-08-26 at mlen 64, Kimi K3. The three convolutions, 58% of the +#: layer's *image* -- larger than the mixer, and previously ungated. Their share +#: of the issued stream is 12%; the two orderings are opposite. +KIMI_CONV_STATIC_PER_LAYER_MAX = 59_000 + +#: The whole state-engine path: gather + conv x3 + gates + recurrence. +KIMI_LAYER_STATIC_MAX = 103_000 + + +def _kimi_conv_asm() -> list[str]: + from compiler.aten.plena.program_kda_conv import kda_conv_blocks + + mlen = 64 + shape = KdaShape.kimi_k3() + out: list[str] = [] + for channels in ( + shape.projection_size, shape.projection_size, + shape.num_heads * shape.value_dim, + ): + p = PlenaCompiler(mlen=mlen, blen=4) + consts = p.kda_fp_constants() + blocks = kda_conv_blocks(channels, mlen) + up = lambda n: ((n + mlen - 1) // mlen) * mlen # noqa: E731 + a = lambda n, r: p.alloc(n, up(r), mlen, strict=False) # noqa: E731 + mark = len(p.get_code()) + p.kda_conv_step_v0( + x_new=a("x", blocks), conv_state=a("cs", blocks * shape.conv_kernel), + weight=a("w", blocks * shape.conv_kernel), bias=None, + out=a("o", blocks), scratch=a("sc", blocks), consts=consts, + channels=channels, kernel=shape.conv_kernel, + ) + out.append(p.get_code()[mark:]) + return out + + +def _kimi_conv_static() -> int: + return sum(_static(a) for a in _kimi_conv_asm()) + + +def test_kimi_k3_conv_static_instruction_count(): + """The convolutions, the larger half of the layer's program image. + + They were ungated while the mixer was, which had the budget covering 42% of + the thing it was meant to protect. + + "Larger half" is a statement about the image and nothing else. The + convolutions issue 59,409 instructions against the mixer's 428,622, so by + work they are the smaller *eighth*. Both budgets are image budgets and are + right to be -- they guard against unrolling -- but neither ranks the + kernels. + """ + total = _kimi_conv_static() + print(f"KIMI_K3_CONV q+k+v static={total}") + assert total <= KIMI_CONV_STATIC_PER_LAYER_MAX, ( + f"{total} static instructions for the three convolutions exceeds " + f"{KIMI_CONV_STATIC_PER_LAYER_MAX}" + ) + + +def test_kimi_k3_layer_static_instruction_count(): + """Gather + conv x3 + gates + recurrence, and the 93-layer total.""" + gather = 70 # five sections, one hardware loop each; see test_kda_layer + layer = _kimi_conv_static() + _kimi_mixer_static() + gather + print( + f"KIMI_K3_LAYER static={layer} layers={KIMI_LAYERS} " + f"total={layer * KIMI_LAYERS}" + ) + assert layer <= KIMI_LAYER_STATIC_MAX, ( + f"{layer} static instructions per layer exceeds {KIMI_LAYER_STATIC_MAX}" + ) + + +def _kimi_gather_asm() -> str: + """The five-section projection gather, at Kimi K3's real block counts.""" + from compiler.aten.plena.program_kda_layer import ( + kda_projection_sections, + kda_projection_width, + ) + + mlen = 64 + shape = KdaShape.kimi_k3() + p = PlenaCompiler(mlen=mlen, blen=4) + consts = p.kda_fp_constants() + projected = p.alloc("proj", 64, kda_projection_width(shape, mlen), strict=False) + mark = len(p.get_code()) + for name, _first, count in kda_projection_sections(shape, mlen): + dst = p.alloc(f"dst_{name}", max(count, mlen), mlen, strict=False) + p.kda_gather_projection_v0( + projected=projected, dst=dst, section=name, shape=shape, + consts=consts, name=f"tall_{name}", + ) + return p.get_code()[mark:] + + +#: Measured 2026-08-27 at mlen 64, Kimi K3: 4,650 executed against 492,681 for +#: the layer, 0.94%. Set ~50% above; the ratio, not the absolute, is the claim. +KIMI_GATHER_DYNAMIC_SHARE_MAX = 0.015 + + +def test_the_gather_is_one_percent_of_the_layer_when_the_loops_are_expanded(): + """The projection split cost, measured as issued instructions. + + This test previously divided **static** counts -- 70 against 93,353 -- and + reported 0.07%. That number was real but it answered the wrong question: + the gather is one hardware loop per section and its body runs once per + feature block, 192 of them for q/k/v/gate. Expanded, the gather issues + 4,650 instructions against 492,681 for the layer, and the honest figure is + **0.94%** -- thirteen times the static one. + + Both numbers matter and they are not interchangeable. The static count is + what the `V_FMA_VF` conversion bought (a program image that does not grow + with `key_dim`) and it belongs in the budget gates above. The dynamic count + is what the gather *costs*, and it is the only one that may be compared + against an alternative way of getting the sections into place. + + 0.94% is still small enough that no amount of making the gather faster + changes the layer, which is the conclusion the old test was reaching for. + It is not small enough to call the gather free, which is what it said. + """ + gather = _dynamic(_kimi_gather_asm()) + layer = ( + sum(_dynamic(a) for a in _kimi_conv_asm()) + + _dynamic(_kimi_mixer_asm()) + + gather + ) + share = gather / layer + print( + f"KIMI_K3_GATHER dynamic={gather} layer_dynamic={layer} " + f"share={share:.4%} (static share was {70 / 93_353:.4%})" + ) + assert share <= KIMI_GATHER_DYNAMIC_SHARE_MAX, ( + f"the gather is {share:.3%} of the issued layer, over the " + f"{KIMI_GATHER_DYNAMIC_SHARE_MAX:.1%} budget" + ) + assert share > 10 * (70 / 93_353), ( + "the dynamic share should be an order of magnitude above the static " + "one; if it is not, the gather stopped being a hardware loop" + ) + + +def test_the_image_and_the_work_rank_the_kernels_differently(): + """The convolutions are the larger half of the image and an eighth of the work. + + Pinned because the inversion is severe enough to misdirect effort, and it + already did: the comments in this file called the convolutions "the larger + half of the layer" on the strength of 53,757 against 39,526. Expanded, the + convolutions issue 59,409 and the mixer 428,622 -- the mixer is 87% of the + layer and the convolutions 12%, the opposite order. + + The reason is structural rather than incidental. The convolutions are a + four-tap FIR over feature blocks: the image carries one instruction per tap + per block because the taps are separate weights, and almost nothing loops. + The mixer is the recurrence, which the `V_FMA_VF` conversion collapsed into + hardware loops whose bodies run once per key block per head -- a small + image over a large trip count, which is exactly the shape the image cannot + see. + + So this file's budgets, all of which are image budgets, must never be read + as a cost ranking. They gate one failure mode -- a sweep falling off the + hardware-loop path and unrolling -- and they are the right instrument for + it. Anything about where the layer's time goes has to come from `_dynamic`. + """ + conv = sum(_dynamic(a) for a in _kimi_conv_asm()) + mixer = _dynamic(_kimi_mixer_asm()) + gather = _dynamic(_kimi_gather_asm()) + layer = conv + mixer + gather + print( + f"KIMI_K3_LAYER dynamic conv={conv} ({conv / layer:.1%}) " + f"mixer={mixer} ({mixer / layer:.1%}) gather={gather} " + f"({gather / layer:.1%}) layer={layer}" + ) + assert _kimi_conv_static() > _kimi_mixer_static(), ( + "the convolutions should still be the larger half of the image" + ) + assert mixer > 3 * conv, ( + f"the mixer should dominate the issued stream ({mixer} vs {conv}); if it " + f"no longer does, either a recurrence sweep unrolled or a convolution " + f"grew a loop, and the budgets above cannot tell you which" + ) + + +# --------------------------------------------------------------------------- +# The input projection, which no budget covered until 2026-08-28. + +#: Measured at Kimi K3, `mlen` 64, `blen` 4, `mram_tile_capacity` 64 -- the +#: capacity the shipped transactional config implies. Set ~10% above. +KIMI_PROJECTION_STATIC_MAX = 430_000 +KIMI_PROJECTION_DYNAMIC_MAX = 5_470_000 + +#: `PlenaCompiler`'s unconditional default until 2026-08-28. +LEGACY_MRAM_TILES = 4 + + +def _shipped_config_mode(mlen: int) -> str | None: + """The config section declaring `mlen`, or None if the file is not here. + + `plena_settings.toml` lives in the **simulator** repository; this one is a + submodule of it. A checkout of the compiler alone -- which is what its own + CI does -- has no such file, `_find_plena_settings_toml` returns None, and + every config-derived value falls back to its hardcoded default. + + That split is worth naming rather than papering over. It means the compiler + emits different programs standalone than it does inside the simulator + checkout, and it is why the reader could be broken in two ways at once for + as long as it was: the job that would have caught it never had a file to + read. + """ + from compiler.aten.plena.compiler import _config_mode_for_mlen, _find_plena_settings_toml + + path = _find_plena_settings_toml() + # `_find_plena_settings_toml` returns PLENA_SETTINGS_TOML unchecked, so a + # stale or deliberately bogus env var gives a path that is not there. + if path is None or not path.exists(): + return None + return _config_mode_for_mlen(mlen) + +#: `MatrixSram::new(tile_size = MLEN, depth = MATRIX_SRAM_SIZE)` keeps +#: `depth / tile_size` tiles, so the shipped transactional config -- MLEN 64, +#: MATRIX_SRAM_SIZE 4096 -- is 64. This is now what `PlenaCompiler(mlen=64)` +#: derives. +CONFIGURED_MRAM_TILES = 64 + + +def _kimi_projection_asm(mram_tile_capacity: int) -> str: + """`X @ W_in` for one Kimi K3 KDA layer, through `linear_projection`.""" + from compiler.aten.plena.program_kda_layer import kda_projection_width + + mlen, blen = 64, 4 + shape = KdaShape.kimi_k3() + p = PlenaCompiler(mlen=mlen, blen=blen, mram_tile_capacity=mram_tile_capacity) + x = p.load_batch(p.input("x", (blen, shape.hidden_size)), name="xv") + weight = p.input("W_in", (shape.hidden_size, kda_projection_width(shape, mlen))) + mark = len(p.get_code()) + p.linear_projection(x, weight, name="projected") + return p.get_code()[mark:] + + +def test_kimi_k3_projection_instruction_budget(): + """The largest kernel in the layer, and the last one to get a gate. + + The projection was never in the KDA lowering -- `test_kda_layer.py` loads + its result rather than computing it -- so nothing here covered it. Built, + it is 389,717 static and 4,967,367 issued instructions for one layer, + against 93,353 and 492,681 for the whole state-engine path. + + Those two pairs must **not** be added, ranked, or turned into a + percentage. An `M_MM` is a `64 x 64 x 64` matmul and a `V_FMA_VF` is a + 64-lane vector operation; counting them in the same unit says nothing + about time. The gate exists because a kernel this size should not be + outside every budget, not because the comparison means anything. + """ + asm = _kimi_projection_asm(CONFIGURED_MRAM_TILES) + static, dynamic = _static(asm), _dynamic(asm) + print(f"KIMI_K3_PROJECTION static={static} dynamic={dynamic}") + assert static <= KIMI_PROJECTION_STATIC_MAX, ( + f"{static} static instructions exceeds {KIMI_PROJECTION_STATIC_MAX}" + ) + assert dynamic <= KIMI_PROJECTION_DYNAMIC_MAX, ( + f"{dynamic} issued instructions exceeds {KIMI_PROJECTION_DYNAMIC_MAX}" + ) + + +def test_the_mram_tile_capacity_now_comes_from_the_configuration(): + """It was an unconditional 4, and the machine it compiles for has 64. + + `PlenaCompiler` took `mram_tile_capacity: int = 4` and nothing derived it + from `MATRIX_SRAM_SIZE`. The emulator's `MatrixSram::new` keeps + `depth / tile_size` tiles, so the shipped transactional config -- MLEN 64 + over MATRIX_SRAM_SIZE 4096 -- is 64 tiles, sixteen times the default. + + Kimi K3's projection contracts over 112 k-tiles and `linear_projection` + re-streams the weights once per capacity-sized chunk, so the gap was + expensive: + + cap 4 (32 KiB) 1,090,417 static 13,876,267 issued + cap 16 (128 KiB) 524,467 static 6,680,617 issued + cap 64 (512 KiB) 389,717 static 4,967,367 issued + cap 112 (896 KiB) 362,767 static 4,624,717 issued + + The default now derives from the config section whose `MLEN` is the one + being compiled for -- not from `[MODE].active`, which selects the machine + the *simulator* models and is `analytic` (MLEN 2048) while every program + here is compiled at 64 and run by the transactional emulator. A shape no + section declares, such as the `mlen` 8 used by most tests in this file, + keeps the old default and is unaffected. + """ + # An `mlen` no section declares, and an explicit override, hold with or + # without the file -- so they are checked before the skip. + assert PlenaCompiler(mlen=MLEN, blen=2).mram_tile_capacity == LEGACY_MRAM_TILES + assert PlenaCompiler(mlen=64, blen=4, mram_tile_capacity=7).mram_tile_capacity == 7 + if _shipped_config_mode(64) is None: + pytest.skip("no plena_settings.toml declaring MLEN 64 on this checkout") + assert PlenaCompiler(mlen=64, blen=4).mram_tile_capacity == CONFIGURED_MRAM_TILES + + default = _dynamic(_kimi_projection_asm(LEGACY_MRAM_TILES)) + configured = _dynamic(_kimi_projection_asm(CONFIGURED_MRAM_TILES)) + ratio = default / configured + print( + f"KIMI_K3_PROJECTION cap={LEGACY_MRAM_TILES} dynamic={default} " + f"cap={CONFIGURED_MRAM_TILES} dynamic={configured} ratio={ratio:.2f}" + ) + assert 2.2 <= ratio <= 3.5, ( + f"the old default cost {ratio:.2f}x the configured capacity; it was " + f"2.79x when measured, and a move means either the config or the " + f"K-split strategy changed" + ) + + +def test_the_analytic_config_cannot_hold_one_matrix_tile(): + """`256 / 2048 = 0`, and a zero-tile capacity has no sane fallback. + + `ANALYTIC` declares MLEN 2048 and MATRIX_SRAM_SIZE 256, so the emulator's + own arithmetic gives it no whole matrix tile at all. Nothing reads it + today -- that mode does not drive the transactional emulator -- but the + derivation must not quietly substitute a number nobody chose, because a + contraction split into zero-tile chunks does not terminate. It raises, and + names both figures. + """ + # An explicit capacity overrides whatever the config says, so the mode + # stays compilable either way. True with or without the file. + assert PlenaCompiler(mlen=2048, blen=4, mram_tile_capacity=1).mram_tile_capacity == 1 + if _shipped_config_mode(2048) is None: + pytest.skip("no plena_settings.toml declaring MLEN 2048 on this checkout") + with pytest.raises(ValueError, match="MATRIX_SRAM_SIZE 256"): + PlenaCompiler(mlen=2048, blen=4) + + +def test_the_configuration_is_actually_read(): + """Two independent reasons it never was, both fixed here. + + `_behavior_config_value` asked `load_toml_config` for `[BEHAVIOR].CONFIG` + and fell back to a caller-supplied default on any exception. It took that + fallback every single time, for two reasons at once: + + 1. `plena_settings.toml` has no `[BEHAVIOR]` 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 the import failed there too -- and the + `except Exception` around the call turned "cannot read the config" into + "the default is correct". + + So `PlenaCompiler`'s docstring promised HLEN, BROADCAST_AMOUNT and the + prefetch amounts came from the TOML, and none of them ever did. The reader + now goes through `tomllib`, which is stdlib from 3.11, and searches + `BEHAVIOR` then the section matching `mlen`. + + At `mlen` 64 that moves HLEN from 64 to the transactional config's 16 and + BROADCAST_AMOUNT from 1 to 4. Neither changes 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, which do reach + the emitters, are 4 in that section and were already 4. + """ + from compiler.aten.plena.compiler import ( # noqa: F401 + _config_mode_for_mlen, + _config_section, + _find_plena_settings_toml, + ) + + settings = _find_plena_settings_toml() + if settings is None or not settings.exists(): + pytest.skip( + "no plena_settings.toml on this checkout -- it lives in the " + "simulator repository, and the compiler's own CI checks out only " + "this one, which is why nothing here ever caught the reader" + ) + + assert _config_section("TRANSACTIONAL"), ( + "the transactional CONFIG table came back empty; the reader is not " + "reading the file" + ) + assert _config_section("BEHAVIOR") == {}, ( + "a [BEHAVIOR] table now exists; it takes priority over the machine " + "sections and this test's expectations below need rechecking" + ) + assert _config_mode_for_mlen(64) == "TRANSACTIONAL" + assert _config_mode_for_mlen(2048) == "ANALYTIC" + assert _config_mode_for_mlen(MLEN) is None, ( + f"mlen {MLEN} now matches a declared machine; the tests in this file " + f"would start reading its numbers instead of the defaults" + ) + + p = PlenaCompiler(mlen=64, blen=4) + print( + f"CONFIG_AT_MLEN_64 hlen={p.hlen} broadcast={p.broadcast_amount} " + f"prefetch={p.hbm_v_prefetch_amount} writeback={p.hbm_v_writeback_amount} " + f"mram_tiles={p.mram_tile_capacity}" + ) + assert (p.hlen, p.broadcast_amount) == (16, 4), ( + "these are the transactional config's values; before the fix they were " + "64 and 1, which is what the hardcoded defaults gave" + ) + assert (p.hbm_v_prefetch_amount, p.hbm_v_writeback_amount) == (4, 4), ( + "these reach the emitters, so a change here moves every load_batch on " + "the branch and needs its own decision" + ) + + +# --------------------------------------------------------------------------- +# What the issue slots are actually spent on, and what an addressing mode +# would buy. Measured 2026-08-29. + +#: A quarter of the issued layer computes; the rest is scaffolding. +KIMI_ARITHMETIC_SHARE_MIN = 0.22 + +#: `S_ADDI_INT gpN, gpN, step` inside a loop, folded into the instruction that +#: consumes the pointer. Every one of them has such a consumer today. +KIMI_FOLDABLE_ADVANCES_MIN = 200_000 + + +def _kimi_layer_parts() -> dict[str, str]: + return { + "conv": "".join(_kimi_conv_asm()), + "mixer": _kimi_mixer_asm(), + "gather": _kimi_gather_asm(), + } + + +def test_three_quarters_of_the_issued_layer_is_not_arithmetic(): + """Where the issue slots go, per opcode. + + `dynamic_count` says a Kimi K3 layer issues 492,681 instructions. This says + what they are: **124,428 compute and 368,253 do not**. The scaffolding is + pointer arithmetic (215,634 self-advances inside loops), scalar loads from + FPRAM (102,918), and loop setup. + + It is the shape of the recurrent sweep that does it. Every `V_FMA_VF` sits + in a body of five: + + 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 <- free, see `dynamic_count` + + One in four. This is the number to beat, and it is what a dedicated + recurrence engine beats by construction -- one instruction covering a whole + head has no instruction overhead at all. + """ + parts = _kimi_layer_parts() + issued = sum(_dynamic(a) for a in parts.values()) + work = sum( + sum(v for op, v in opcode_census(a).items() if op in ARITHMETIC) + for a in parts.values() + ) + share = work / issued + for name, asm in parts.items(): + print(f"KIMI_K3_{name.upper()} arithmetic_share={arithmetic_share(asm):.1%}") + print(f"KIMI_K3_LAYER issued={issued} arithmetic={work} share={share:.1%}") + assert share >= KIMI_ARITHMETIC_SHARE_MIN, ( + f"only {share:.1%} of the issued layer computes; it was 25.3% when " + f"measured, and a fall means scaffolding grew" + ) + assert share < 0.5, ( + "if arithmetic ever exceeds half the issue stream, the sweeps have " + "changed shape and the addressing-mode case below needs re-pricing" + ) + + +def test_every_pointer_advance_could_fold_into_a_post_increment(): + """Pricing an addressing mode instead of an engine. + + A post-increment operand -- the pointer carrying its own stride, the way + `_emit_tile_row_fma` already carries independent row progressions -- would + fold `S_ADDI_INT gpN, gpN, step` into the instruction that consumes the + pointer. That only works where the register has a consumer in the same loop + body, so the saving is priced against those and not against every advance + that happens to sit in a loop. + + **Every one of them qualifies.** 215,634 foldable, 0 not, across the + convolutions, the mixer and the gather. 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. + + Two proposals, measured on the Kimi K3 layer: + + now 492,681 issued 25.3% arithmetic + P1 post-inc on the 277,047 issued 44.9% arithmetic 1.78x + vector operands + P2 P1 + the FPRAM 174,129 issued 71.5% arithmetic 2.83x + scalar auto-advancing too + + P2 is an addressing field on instructions that already exist. The + alternative it is competing with is a dedicated recurrence datapath with + its own banked state SRAM plus a layout engine to feed it, and the two are + not close in cost. + + What neither proposal touches: KDA's decay is channel-wise on the key axis, + so the operand is `key_dim` lanes wide -- 6.2% of `VLEN` 2048. That is a + *width* mismatch, not an encoding one, and no addressing mode reaches it. + On a machine whose width matches the operand the two inefficiencies do not + compound; on a wide one they do, and that is the one place a purpose-built + engine wins on architecture rather than on encoding. + """ + parts = _kimi_layer_parts() + issued = sum(_dynamic(a) for a in parts.values()) + foldable = unfoldable = 0 + fp_loads = 0 + for asm in parts.values(): + f, u = self_advance_counts(asm) + foldable += f + unfoldable += u + fp_loads += opcode_census(asm).get("S_LD_FP", 0) + p1, p2 = issued - foldable, issued - foldable - fp_loads + print( + f"KIMI_K3_LAYER issued={issued} foldable={foldable} unfoldable={unfoldable} " + f"fp_loads={fp_loads} P1={p1} ({issued / p1:.2f}x) P2={p2} ({issued / p2:.2f}x)" + ) + assert unfoldable == 0, ( + f"{unfoldable} pointer advances have no consumer in their loop body and " + f"could not fold; the P1 figure above would need re-deriving" + ) + assert foldable >= KIMI_FOLDABLE_ADVANCES_MIN, ( + f"{foldable} foldable advances, below the {KIMI_FOLDABLE_ADVANCES_MIN} " + f"measured; the sweeps changed shape" + ) + assert issued / p1 > 1.6, f"P1 is worth {issued / p1:.2f}x, was 1.78x" + assert issued / p2 > 2.4, f"P2 is worth {issued / p2:.2f}x, was 2.83x" diff --git a/aten/tests/test_isa_interpreter.py b/aten/tests/test_isa_interpreter.py new file mode 100644 index 0000000..36d062f --- /dev/null +++ b/aten/tests/test_isa_interpreter.py @@ -0,0 +1,607 @@ +"""Unit tests for the ISA oracle in `aten/tests/isa_interpreter.py`. + +The interpreter is what makes `test_kda_decode_step.py` a numerical check +rather than a shape check, so it needs anchoring independently of what the KDA +lowering happens to emit. Mutating it and running only the KDA tests is not +enough: making `V_MUL_VF` read its destination instead of its source left those +green, because the lowering only ever emits the in-place form `V_MUL_VF gpX, +gpX, fY`. Every opcode is therefore exercised here in a form the lowering does +not currently produce. + +Semantics are taken from the emulator, not from `doc/plena_isa_spec.md` -- +`V_SUB_VV` is the case where the two disagreed. +""" + +from __future__ import annotations + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.tests.isa_interpreter import Machine, UnsupportedInstruction # noqa: E402 + +VLEN = 4 + + +def _machine() -> Machine: + return Machine(vlen=VLEN, vram_words=256, fpram_words=64) + + +# --------------------------------------------------------------------------- +# Vector ops -- out of place, so dst and src cannot be confused +# --------------------------------------------------------------------------- + + +def test_v_mul_vf_reads_its_source_not_its_destination(): + m = _machine() + m.write_vram_row(0, [9.0, 9.0, 9.0, 9.0]) # destination, must be overwritten + m.write_vram_row(4, [1.0, 2.0, 3.0, 4.0]) # source + m.write_fpram(0, [2.5]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + V_MUL_VF gp1, gp2, f1, 0 + """ + ) + assert m.read_vram_row(0) == [2.5, 5.0, 7.5, 10.0] + assert m.read_vram_row(4) == [1.0, 2.0, 3.0, 4.0], "source must be untouched" + + +def test_v_sub_vv_is_rs1_minus_rs2(): + """The emulator computes rs1 - rs2. doc/plena_isa_spec.md said rs2 - rs1 + until this was found; subtraction does not commute, and KDA's error term + flips sign under the wrong one.""" + m = _machine() + m.write_vram_row(4, [10.0, 10.0, 10.0, 10.0]) + m.write_vram_row(8, [1.0, 2.0, 3.0, 4.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + S_ADDI_INT gp3, gp0, 8 + V_SUB_VV gp1, gp2, gp3, 0 + """ + ) + assert m.read_vram_row(0) == [9.0, 8.0, 7.0, 6.0] + + +def test_v_add_vv_writes_a_third_destination(): + m = _machine() + m.write_vram_row(4, [1.0, 2.0, 3.0, 4.0]) + m.write_vram_row(8, [10.0, 20.0, 30.0, 40.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + S_ADDI_INT gp3, gp0, 8 + V_ADD_VV gp1, gp2, gp3, 0 + """ + ) + assert m.read_vram_row(0) == [11.0, 22.0, 33.0, 44.0] + assert m.read_vram_row(4) == [1.0, 2.0, 3.0, 4.0] + + +# --------------------------------------------------------------------------- +# Scalar and control +# --------------------------------------------------------------------------- + + +def test_s_addi_int_accumulates(): + m = _machine() + m.run( + """ + S_ADDI_INT gp1, gp0, 7 + S_ADDI_INT gp1, gp1, 5 + """ + ) + assert m.gp[1] == 12 + + +def test_gp0_and_f0_are_writable_registers_not_hardwired_zero(): + """Faithful to the emulator, where both are plain array slots. Their zero + value is a compiler convention (`RegisterAllocator` never hands them out), + so an emitter that wrote one would break address arithmetic on real + hardware -- and a hardwired oracle would keep passing it.""" + m = _machine() + m.run( + """ + S_ADDI_INT gp0, gp0, 99 + S_ADDI_INT gp1, gp0, 1 + """ + ) + assert m.gp[0] == 99 + assert m.gp[1] == 100 + + +def test_s_ld_fp_addresses_through_the_base_register(): + m = _machine() + m.write_fpram(0, [10.0, 11.0, 12.0, 13.0]) + m.write_vram_row(0, [1.0, 1.0, 1.0, 1.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 2 + S_LD_FP f1, gp2, 1 + V_MUL_VF gp1, gp1, f1, 0 + """ + ) + # base 2 + imm 1 -> slot 3 + assert m.read_vram_row(0) == [13.0, 13.0, 13.0, 13.0] + + +def test_loop_runs_exactly_the_declared_count(): + m = _machine() + m.write_vram_row(0, [1.0, 1.0, 1.0, 1.0]) + m.write_fpram(0, [2.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + C_LOOP_START gp2, 3 + V_MUL_VF gp1, gp1, f1, 0 + C_LOOP_END gp2 + """ + ) + assert m.read_vram_row(0) == [8.0, 8.0, 8.0, 8.0], "2^3, so exactly 3 iterations" + + +def test_loop_body_advances_a_pointer_each_iteration(): + """The pattern every tile_row_* sweep emits: the loop body bumps the + address register, so an off-by-one in the trip count shows up as a row that + was skipped or written twice.""" + m = _machine() + for row in range(3): + m.write_vram_row(row * VLEN, [1.0] * VLEN) + m.write_vram_row(3 * VLEN, [5.0] * VLEN) # must be left alone + m.write_fpram(0, [3.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + C_LOOP_START gp2, 3 + V_MUL_VF gp1, gp1, f1, 0 + S_ADDI_INT gp1, gp1, 4 + C_LOOP_END gp2 + """ + ) + for row in range(3): + assert m.read_vram_row(row * VLEN) == [3.0] * VLEN + assert m.read_vram_row(3 * VLEN) == [5.0] * VLEN + + +# --------------------------------------------------------------------------- +# Failure modes: the oracle must refuse rather than guess +# --------------------------------------------------------------------------- + + +def test_unmodelled_opcode_raises(): + # V_FMA_VF used to be the example here; it is modelled now, so this needs an + # opcode that genuinely is not. V_MAX_VF exists in the ISA and no KDA or + # Mamba lowering emits it. + with pytest.raises(UnsupportedInstruction, match="outside the modelled subset"): + _machine().run("V_MAX_VF gp1, gp2, f1, 0") + + +def test_fma_accumulates_into_its_destination(): + """The oracle has to model the accumulate, or a lowering that used FMA where + it meant multiply would agree with it and disagree with the hardware.""" + m = _machine() + m.write_vram_row(0, [1.0, 2.0, 3.0, 4.0]) + m.write_vram_row(4, [10.0, 20.0, 30.0, 40.0]) + m.write_fpram(0, [0.5]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + V_FMA_VF gp1, gp2, f1, 0 + """ + ) + assert m.read_vram_row(0, 4) == [6.0, 12.0, 18.0, 24.0] + assert m.read_vram_row(4, 4) == [10.0, 20.0, 30.0, 40.0], "source is read-only" + + +def test_fma_with_the_destination_as_its_own_source(): + """`x += x * f` is `x * (1 + f)`. Reading and writing the same row must not + see a partially-updated value part-way through the lane sweep.""" + m = _machine() + m.write_vram_row(0, [1.0, 2.0, 3.0, 4.0]) + m.write_fpram(0, [2.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + V_FMA_VF gp1, gp1, f1, 0 + """ + ) + assert m.read_vram_row(0, 4) == [3.0, 6.0, 9.0, 12.0] + + +def test_nonzero_rmask_raises(): + with pytest.raises(UnsupportedInstruction, match="rmask"): + _machine().run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + V_ADD_VV gp1, gp2, gp2, 1 + """ + ) + + +def test_misaligned_vector_address_raises(): + """The emulator's vector SRAM asserts VLEN alignment; reading across a row + boundary is a panic there, not a silent wrap.""" + with pytest.raises(UnsupportedInstruction, match="not a multiple of vlen"): + _machine().run( + """ + S_ADDI_INT gp1, gp0, 3 + S_ADDI_INT gp2, gp0, 0 + V_ADD_VV gp1, gp2, gp2, 0 + """ + ) + + +def test_comments_and_stage_markers_are_ignored(): + m = _machine() + m.write_vram_row(0, [1.0, 1.0, 1.0, 1.0]) + m.run( + """ + ; @stage=kda_decay head=0 + S_ADDI_INT gp1, gp0, 0 + // trailing style comment + V_ADD_VV gp1, gp1, gp1, 0 + """ + ) + assert m.read_vram_row(0) == [2.0, 2.0, 2.0, 2.0] + + +def test_runaway_loop_is_bounded(): + """A loop whose counter is rewritten inside the body would otherwise hang + the suite instead of failing it.""" + with pytest.raises(RuntimeError, match="runaway loop"): + _machine().run( + """ + C_LOOP_START gp2, 2 + S_ADDI_INT gp2, gp0, 2 + C_LOOP_END gp2 + """ + ) + + +# --------------------------------------------------------------------------- +# Ops added for the causal conv. Each is exercised in a form the lowering does +# not emit, for the same reason as V_MUL_VF above. +# --------------------------------------------------------------------------- + + +def test_v_mul_vv_is_elementwise_and_out_of_place(): + m = _machine() + m.write_vram_row(0, [9.0] * VLEN) + m.write_vram_row(4, [1.0, 2.0, 3.0, 4.0]) + m.write_vram_row(8, [10.0, 100.0, 1000.0, 0.5]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + S_ADDI_INT gp3, gp0, 8 + V_MUL_VV gp1, gp2, gp3, 0 + """ + ) + assert m.read_vram_row(0) == [10.0, 200.0, 3000.0, 2.0] + assert m.read_vram_row(4) == [1.0, 2.0, 3.0, 4.0] + + +def test_v_add_vf_broadcasts_the_scalar(): + m = _machine() + m.write_vram_row(4, [1.0, 2.0, 3.0, 4.0]) + m.write_fpram(0, [0.5]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + V_ADD_VF gp1, gp2, f1, 0 + """ + ) + assert m.read_vram_row(0) == [1.5, 2.5, 3.5, 4.5] + + +def test_v_sub_vf_honours_rorder(): + """rorder=1 reverses the operands. The negate idiom `V_SUB_VF rd, rs1, f0, 0, 1` + depends on it, and getting it backwards flips the sign of everything after.""" + m = _machine() + m.write_vram_row(4, [1.0, 2.0, 3.0, 4.0]) + m.write_fpram(0, [10.0]) + prologue = """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + """ + m.run(prologue + "V_SUB_VF gp1, gp2, f1, 0, 0\n") + assert m.read_vram_row(0) == [-9.0, -8.0, -7.0, -6.0] + + m2 = _machine() + m2.write_vram_row(4, [1.0, 2.0, 3.0, 4.0]) + m2.write_fpram(0, [10.0]) + m2.run(prologue + "V_SUB_VF gp1, gp2, f1, 0, 1\n") + assert m2.read_vram_row(0) == [9.0, 8.0, 7.0, 6.0] + + +def test_v_exp_v_saturates_rather_than_overflowing(): + """vector_machine.rs clamps to [-88, 88] before exp, so bf16 cannot go + infinite. An unclamped model would return inf and any comparison against it + would then pass or fail for the wrong reason.""" + import math + + m = _machine() + m.write_vram_row(4, [0.0, 1.0, 200.0, -200.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + V_EXP_V gp1, gp2, 0 + """ + ) + got = m.read_vram_row(0) + assert got[0] == pytest.approx(1.0) + assert got[1] == pytest.approx(math.e) + assert got[2] == pytest.approx(math.exp(88.0)) + assert got[3] == pytest.approx(math.exp(-88.0)) + assert all(math.isfinite(x) for x in got) + + +def test_v_reci_v_on_zero_returns_inf_like_the_emulator(): + """vector_machine.rs uses tensor.reciprocal(), which yields inf rather than + trapping. A model that raises would crash the oracle on a lowering the + hardware runs fine -- e.g. one that reciprocates an unused lane.""" + import math + + m = _machine() + m.write_vram_row(4, [0.0, 2.0, 0.0, -4.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + V_RECI_V gp1, gp2, 0 + """ + ) + got = m.read_vram_row(0) + assert math.isinf(got[0]) and math.isinf(got[2]) + assert got[1] == 0.5 and got[3] == -0.25 + + +def test_v_reci_v_is_elementwise_reciprocal(): + m = _machine() + m.write_vram_row(4, [1.0, 2.0, 4.0, -0.5]) + m.run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + V_RECI_V gp1, gp2, 0 + """ + ) + assert m.read_vram_row(0) == [1.0, 0.5, 0.25, -2.0] + + +# --------------------------------------------------------------------------- +# Scalar FP and the reduction. +# +# These carry the two semantics this oracle took from the emulator rather than +# from doc/plena_isa_spec.md, and neither was reachable through the KDA +# lowerings: every V_RED_SUM the emitters produce is preceded by +# `S_ADD_FP f1, f0, f0`, so f1 always enters at zero and accumulate-vs-overwrite +# never shows; and no emitter ever names f0 as a destination, because +# RegisterAllocator hands out f1..f7. Both are pinned here directly. +# --------------------------------------------------------------------------- + + +def _fp_prog(body: str) -> str: + return """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 4 + """ + body + + +def test_v_red_sum_accumulates_into_its_fp_register(): + """dispatch.rs seeds reduce_sum with the current f[rd], so a second + reduction adds to the first. The compiler relies on the opposite -- it + clears f1 before every V_RED_SUM -- so nothing in the lowering exercises + this, and an oracle that overwrote would go unnoticed.""" + m = _machine() + m.write_vram_row(0, [1.0, 2.0, 3.0, 4.0]) # sums to 10 + m.write_vram_row(4, [10.0, 20.0, 30.0, 40.0]) # sums to 100 + m.run( + _fp_prog( + """ + V_RED_SUM f1, gp1, 0, 0 + V_RED_SUM f1, gp2, 0, 0 + S_ADDI_INT gp3, gp0, 0 + S_ST_FP f1, gp3, 0 + """ + ) + ) + assert m.fpram[0] == 110.0, "second V_RED_SUM must add, not replace" + + +@pytest.mark.parametrize("form", ["V_RED_SUM f1, gp1, 1, 0", "V_RED_SUM f1, gp1, 0, 1"]) +def test_v_red_sum_refuses_a_mask_it_does_not_model(form): + """The masked reduce_sum in vector_machine.rs is a materially different + computation -- it broadcasts each head's sum back over its slice before + summing. Refuse rather than model it wrongly, whichever operand slot the + assembler puts the mask in.""" + with pytest.raises(UnsupportedInstruction, match="mask"): + _machine().run(_fp_prog(form + "\n")) + + +def test_scalar_fp_writes_to_f0_are_discarded(): + """dispatch.rs:417-424 makes every scalar FP op a no-op when rd == 0. That + is what keeps f0 zero for the `V_MUL_VF gpX, gpX, f0` zeroing idiom -- + RegisterAllocator never handing out f0 is the other half.""" + m = _machine() + m.write_fpram(0, [7.0, 3.0]) + m.write_vram_row(0, [5.0, 5.0, 5.0, 5.0]) + m.run( + """ + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + S_LD_FP f2, gp3, 1 + S_ADD_FP f0, f1, f2 + S_ADDI_INT gp1, gp0, 0 + V_MUL_VF gp1, gp1, f0, 0 + """ + ) + assert m.fp[0] == 0.0, "S_ADD_FP with rd=0 must be discarded" + assert m.read_vram_row(0) == [0.0] * VLEN, "f0 must still zero the row" + + +def test_s_ld_fp_is_not_guarded_against_f0(): + """S_LD_FP is deliberately absent from dispatch.rs's rd==0 guard list, so + modelling it as guarded would diverge.""" + m = _machine() + m.write_fpram(0, [9.0]) + m.run( + """ + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f0, gp3, 0 + """ + ) + assert m.fp[0] == 9.0 + + +def test_s_st_fp_source_and_address_roles(): + """S_ST_FP { rd, rs1, imm } stores f[rd] to FPRAM[gp[rs1] + imm] -- rd is + the source, which is the reverse of the usual destination-first reading.""" + m = _machine() + m.write_fpram(0, [42.0]) + m.run( + """ + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + S_ADDI_INT gp4, gp0, 5 + S_ST_FP f1, gp4, 2 + """ + ) + assert m.fpram[7] == 42.0, "address is gp[rs1] + imm = 5 + 2" + assert m.fpram[5] == 0.0 and m.fpram[0] == 42.0 + + +def test_scalar_fp_arithmetic(): + m = _machine() + m.write_fpram(0, [10.0, 4.0]) + m.run( + """ + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + S_LD_FP f2, gp3, 1 + S_ADD_FP f3, f1, f2 + S_SUB_FP f4, f1, f2 + S_MUL_FP f5, f1, f2 + S_ADDI_INT gp4, gp0, 10 + S_ST_FP f3, gp4, 0 + S_ST_FP f4, gp4, 1 + S_ST_FP f5, gp4, 2 + """ + ) + assert m.fpram[10] == 14.0 + assert m.fpram[11] == 6.0, "S_SUB_FP is fs1 - fs2, not the reverse" + assert m.fpram[12] == 40.0 + + +def test_s_sqrt_and_s_reci_edge_cases(): + """sqrt of a negative is NaN (f32::sqrt), reciprocal of zero is inf + (bf16::ONE / 0), and -0.0 gives -inf rather than +inf.""" + import math + + m = _machine() + m.write_fpram(0, [9.0, -4.0, 0.0, -0.0]) + m.run( + """ + S_ADDI_INT gp3, gp0, 0 + S_LD_FP f1, gp3, 0 + S_SQRT_FP f2, f1 + S_LD_FP f3, gp3, 1 + S_SQRT_FP f4, f3 + S_LD_FP f5, gp3, 2 + S_RECI_FP f6, f5 + S_LD_FP f7, gp3, 3 + S_RECI_FP f1, f7 + S_ADDI_INT gp4, gp0, 10 + S_ST_FP f2, gp4, 0 + S_ST_FP f4, gp4, 1 + S_ST_FP f6, gp4, 2 + S_ST_FP f1, gp4, 3 + """ + ) + assert m.fpram[10] == 3.0 + assert math.isnan(m.fpram[11]) + assert m.fpram[12] == math.inf + assert m.fpram[13] == -math.inf + + +def test_v_red_sum_rejects_a_misaligned_address(): + with pytest.raises(UnsupportedInstruction, match="not a multiple of vlen"): + _machine().run( + """ + S_ADDI_INT gp1, gp0, 3 + V_RED_SUM f1, gp1, 0, 0 + """ + ) + + +def test_s_map_fp_v_moves_a_vram_row_into_fpram(): + """Operand roles mirror S_MAP_V_FP: rs1 is the VRAM source row and rd the + FP_MEM base, so that in both instructions rd names the destination memory. + Reading them the usual destination-first way gets it backwards.""" + m = _machine() + m.write_vram_row(8, [1.0, 2.0, 3.0, 4.0]) + m.run( + """ + S_ADDI_INT gp1, gp0, 5 + S_ADDI_INT gp2, gp0, 8 + S_MAP_FP_V gp1, gp2, 2 + """ + ) + # destination is gp[rd] + imm = 5 + 2 + assert m.fpram[7:11] == [1.0, 2.0, 3.0, 4.0] + assert m.fpram[5] == 0.0 and m.fpram[6] == 0.0 + assert m.read_vram_row(8) == [1.0, 2.0, 3.0, 4.0], "source untouched" + + +def test_s_map_fp_v_rejects_a_misaligned_source(): + with pytest.raises(UnsupportedInstruction, match="not a multiple of vlen"): + _machine().run( + """ + S_ADDI_INT gp1, gp0, 0 + S_ADDI_INT gp2, gp0, 3 + S_MAP_FP_V gp1, gp2, 0 + """ + ) + + +def test_s_map_fp_v_refuses_to_write_past_the_modelled_fpram(): + """Python slice assignment past the end *extends* the list rather than + failing, so an out-of-bounds write would silently land at the wrong address + and stay green. The emulator asserts on the FP SRAM bound.""" + m = Machine(vlen=VLEN, vram_words=64, fpram_words=16) + m.write_vram_row(0, [1.0, 2.0, 3.0, 4.0]) + with pytest.raises(UnsupportedInstruction, match="past the modelled"): + m.run( + """ + S_ADDI_INT gp1, gp0, 14 + S_ADDI_INT gp2, gp0, 0 + S_MAP_FP_V gp1, gp2, 0 + """ + ) + assert len(m.fpram) == 16, "the FPRAM must not have grown" diff --git a/aten/tests/test_kda_decode_step.py b/aten/tests/test_kda_decode_step.py index 3d44f47..9beb465 100644 --- a/aten/tests/test_kda_decode_step.py +++ b/aten/tests/test_kda_decode_step.py @@ -439,6 +439,12 @@ def test_emits_exactly_one_opcode_beyond_the_phase0_set(): What still has to hold is that it is exactly *one*. The interpreter raises on anything it does not model, so a lowering that reaches for a second new opcode fails here rather than being silently half-checked. + + "One" is a statement about **this kernel**, not about the branch. The + branch adds three opcodes -- `V_SOFTPLUS_V` 0x39, `S_MAP_FP_V` 0x3A and + `V_FMA_VF` 0x3B -- and `test_no_state_engine.py` pins that count. The + decode step reaches for only the third of them; the other two belong to + Mamba's `dt` and to the FPRAM window fill. """ shape, ref = _case(1, num_heads=2, key_dim=4) code = _Harness(shape).emit() diff --git a/aten/tests/test_no_state_engine.py b/aten/tests/test_no_state_engine.py new file mode 100644 index 0000000..b3309b3 --- /dev/null +++ b/aten/tests/test_no_state_engine.py @@ -0,0 +1,179 @@ +"""The static path carries no descriptor machinery, and must not acquire any. + +PLENA is a **statically scheduled** accelerator: the compiler decides every +address at compile time. A descriptor fetched from HBM at runtime re-decides +something the compiler already decided, and a residency cache defers a decision +the compiler is not allowed to defer. Neither belongs on this path, and this file +fails if either reappears. + +What the recurrent kernels need instead is three ordinary instructions -- +`V_SOFTPLUS_V` (0x39), `S_MAP_FP_V` (0x3A) and `V_FMA_VF` (0x3B) -- each a +fixed-function ALU or move op with its operands named in the instruction word. +Everything else -- the state layout, the per-head streaming, the chunked +prefill -- is compiler work. + +There is nothing to delete: this branch was cut from `a4b3e7de` on main, which +never had any of it. So this file is not a cleanup, it is the guard that keeps +the property. A grep is a check that ran once; a test is a check that keeps +running. +""" + +from __future__ import annotations + +import os +import re +import subprocess +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +#: The vocabulary of a run-time state engine: two module names and two opcode +#: names that only exist if the machine fetches its own work descriptors. None of +#: them may appear outside the documents that explain why they do not. +_FORBIDDEN = ("X_STATE", "L_SCATTER_M", "StateDescriptor", "state_engine") + +#: Paths whose *purpose* is to record the decision not to build one. +_ALLOWED_PREFIXES = ( + "docs/superpowers/", # the plan and the progress log + "doc/static_path_measurements.md", + "aten/tests/test_no_state_engine.py", +) + +_SEARCH_SUFFIXES = (".rs", ".py", ".svh", ".json", ".toml") + + +def _repo_root() -> Path: + """The Simulator checkout, which contains the Compiler submodule.""" + here = Path(__file__).resolve() + for parent in here.parents: + if (parent / "transactional_emulator").is_dir(): + return parent + pytest.skip("not inside a Simulator checkout") + + +def _tracked_files(root: Path) -> list[Path]: + out: list[Path] = [] + for repo in (root, root / "PLENA_Compiler"): + if not (repo / ".git").exists(): + continue + listing = subprocess.run( + ["git", "-C", str(repo), "ls-files"], + capture_output=True, text=True, check=True, + ).stdout.splitlines() + out.extend(repo / name for name in listing) + return out + + +def test_no_descriptor_machinery_anywhere(): + """No opcode, module or type from the descriptor design. + + `kda_state_engine_step` and `KdaRecurrentState` in the CPU reference are *function + and dataclass names*, not the instruction -- they are the boundary the + reference defines, and the static lowering targets exactly that boundary. + The pattern below is word-anchored so it does not catch them. + """ + root = _repo_root() + pattern = re.compile(r"\b(" + "|".join(_FORBIDDEN) + r")\b") + hits: list[str] = [] + for path in _tracked_files(root): + rel = str(path.relative_to(root)) if path.is_relative_to(root) else str(path) + rel = rel.replace("PLENA_Compiler/", "", 1) + if rel.startswith(_ALLOWED_PREFIXES) or path.suffix not in _SEARCH_SUFFIXES: + continue + try: + text = path.read_text(errors="ignore") + except OSError: # pragma: no cover - unreadable file + continue + for lineno, line in enumerate(text.splitlines(), 1): + if pattern.search(line): + hits.append(f"{rel}:{lineno}: {line.strip()[:90]}") + assert not hits, ( + "descriptor machinery has come back:\n " + "\n ".join(hits[:20]) + ) + + +def test_the_state_engine_directories_do_not_exist(): + root = _repo_root() + for gone in ( + root / "transactional_emulator" / "src" / "state_engine", + root / "PLENA_Compiler" / "aten" / "state", + root / "PLENA_Compiler" / "spec", + ): + assert not gone.exists(), f"{gone} exists; the static path does not need it" + + +def test_the_descriptor_opcodes_are_still_free(): + """0x3D and 0x3F stay undefined. + + They are the slots a descriptor-driven state instruction and a + descriptor-driven scatter would take. The compiler's opcode table must not + define them and the emulator asserts 0x3F decodes to Invalid, so filling + either needs a written argument rather than a free slot. + """ + root = _repo_root() + svh = (root / "PLENA_Compiler" / "doc" / "operation.svh").read_text() + defined = { + int(m.group(2), 16): m.group(1) + for m in re.finditer(r"(\w+)\s*=\s*6'h([0-9A-Fa-f]+)", svh) + } + # 0x3D and 0x3F are the slots a descriptor-driven state instruction and a + # descriptor-driven scatter would naturally take. Held free deliberately: + # a free slot is not an argument for filling it. + for opcode in (0x3D, 0x3F): + assert opcode not in defined, ( + f"0x{opcode:02X} is defined as {defined[opcode]}. This path carries no " + f"descriptor machinery; adding an opcode here needs an argument in " + f"docs/superpowers/plans/, not just a free slot" + ) + for opcode, name in ((0x39, "V_SOFTPLUS_V"), (0x3A, "S_MAP_FP_V"), (0x3B, "V_FMA_VF")): + assert defined.get(opcode) == name, ( + f"{name} should hold 0x{opcode:02X} -- these three are what this work " + f"adds to the ISA, and every one is an ordinary fixed-function op" + ) + assert max(defined) == 0x3B, ( + f"the highest opcode is 0x{max(defined):02X}; adding another needs an " + f"argument, not just a free slot" + ) + + +#: What `origin/main` stops at -- `C_SET_TOPK_REG`. Everything past it is this +#: branch's, and there are three. +MAIN_LAST_OPCODE = 0x38 + +#: `V_SOFTPLUS_V`, `S_MAP_FP_V`, `V_FMA_VF`. +OPCODES_ADDED_HERE = {0x39, 0x3A, 0x3B} + + +def test_exactly_three_opcodes_were_added(): + """The whole design rests on the count staying at three, and on what they are. + + This test used to be called `test_exactly_one_opcode_was_added` and asserted + `{0x3B}`, with a comment claiming `0x35..0x3A` predate this work. Half of + that was wrong: `origin/main`'s `operation.svh` stops at `0x38` + (`C_SET_TOPK_REG`), and `V_SOFTPLUS_V` at `0x39` and `S_MAP_FP_V` at `0x3A` + are both added here. The branch adds **three**. + + The old form did not merely misdescribe itself -- it looked only at + `c > 0x3A`, so a fourth opcode landing on `0x39` or `0x3A` would have gone + straight through. It passed while asserting something false about a range + it had excluded. + + What the count is guarding is the claim that a static recurrence needs no + new *mechanism*: all three are ordinary fixed-function ops with their + operands in the instruction word. A fourth that reads a descriptor, holds a + queue, or keeps residency would break that claim however few of them there + were, which is what the other tests in this file check. + """ + root = _repo_root() + svh = (root / "PLENA_Compiler" / "doc" / "operation.svh").read_text() + codes = {int(m.group(1), 16) + for m in re.finditer(r"\w+\s*=\s*6'h([0-9A-Fa-f]+)", svh)} + added_by_this_work = {c for c in codes if c > MAIN_LAST_OPCODE} + assert added_by_this_work == OPCODES_ADDED_HERE, ( + f"opcodes past {hex(MAIN_LAST_OPCODE)}: " + f"{sorted(hex(c) for c in added_by_this_work)}; this work adds exactly " + f"three -- V_SOFTPLUS_V 0x39, S_MAP_FP_V 0x3A, V_FMA_VF 0x3B" + ) diff --git a/aten/tests/test_real_checkpoint.py b/aten/tests/test_real_checkpoint.py new file mode 100644 index 0000000..6cf4877 --- /dev/null +++ b/aten/tests/test_real_checkpoint.py @@ -0,0 +1,200 @@ +"""Real published weights, not `randn`. + +Synthetic weights make every lowering test pass on values that no trained model +would produce, so a defect that only appears on realistic structure never +surfaces. This binds a real one: layer 0 of `AntonV/mamba2-130m-hf`, a published +Mamba-2 with 24 layers, `state_size` 128, 24 heads of `head_dim` 64. + +What a real checkpoint tests that synthetic weights cannot: + +* **Structure.** Layer 0's `A` is not a spread: 23 of its 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. +* **Layout.** The projection packs `z`, `x`, `B`, `C` and `dt` into one 3352-wide + tensor in a fixed order; `conv1d.weight` is `[channels, 1, kernel]`, not + `[channels, kernel]`. Getting either wrong is a plausible wrong answer that + synthetic weights, generated to whatever shape the test assumed, cannot catch. + +Skipped when the checkpoint is not cached: this must not require network in CI. +Fetch it once with + + python -c "from huggingface_hub import snapshot_download as d; d('AntonV/mamba2-130m-hf')" +""" + +from __future__ import annotations + +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from compiler.aten.models.mamba2.reference import ( # noqa: E402 + mamba2_recurrent_reference, +) +from compiler.aten.plena import PlenaCompiler # noqa: E402 +from compiler.aten.plena.program_mamba_common import Mamba2Shape # noqa: E402 +from compiler.aten.tests.isa_interpreter import Machine # noqa: E402 + +REPO = "AntonV/mamba2-130m-hf" +MLEN = 64 + + +def _checkpoint(): + """Layer 0's mixer weights and the model config, or skip.""" + try: + from huggingface_hub import snapshot_download + from safetensors import safe_open + except ImportError: # pragma: no cover - depends on the environment + pytest.skip("huggingface_hub / safetensors not installed") + import json + from pathlib import Path + + try: + path = Path(snapshot_download(REPO, local_files_only=True)) + except Exception: # pragma: no cover - not cached + pytest.skip( + f"{REPO} is not cached; fetch it with snapshot_download('{REPO}')" + ) + config = json.loads((path / "config.json").read_text()) + weights = {} + with safe_open(path / "model.safetensors", framework="pt") as f: + for key in f.keys(): + if key.startswith("backbone.layers.0.mixer."): + weights[key.split("mixer.")[1]] = f.get_tensor(key).float() + return config, weights + + +def test_the_real_decay_is_clustered_with_an_outlier(): + """The structure real weights have and `randn` does not. + + Layer 0's `A` is not a spread: 23 of its 24 heads sit between -0.73 and + -0.27, and one is at -5.06, about ten times the rest. A tight cluster plus a + single outlier is what training produces and what a symmetric draw does not, + and it is the head most likely to expose a decay that underflows. + + Asserted from the checkpoint rather than from a threshold I picked: an + earlier version of this demanded `min < -10` on the strength of nothing. + """ + _, w = _checkpoint() + a = -torch.exp(w["A_log"]) + outlier = a.min() + rest = a[a != outlier] + assert outlier < 5 * rest.min(), ( + f"the outlier at {outlier:.2f} is no longer far from the cluster " + f"({rest.min():.2f}..{rest.max():.2f})" + ) + assert rest.std() < 0.15, ( + f"the other heads are no longer clustered (std {rest.std():.3f})" + ) + + +def test_mamba2_decode_step_on_published_weights(): + """`ssm_decode_step_v0` against the textbook recurrence, on layer 0's real + `A`, `D` and `dt_bias`, with the projection's own split order. + + The emitter takes `dA = exp(A * dt)` precomputed, because it depends only on + a weight and a scalar; the reference takes raw `A` and computes it. Feeding + both from the same checkpoint is the point. + """ + config, w = _checkpoint() + heads = config["num_heads"] + head_dim = config["head_dim"] + state = config["state_size"] + groups = config["n_groups"] + assert head_dim == MLEN, f"this test assumes head_dim == mlen, got {head_dim}" + + torch.manual_seed(0) + # Activations, since B, C, dt and x come from the projection and not from a + # weight. The weights under test are A_log, D and dt_bias. + x = torch.randn(heads, head_dim) * 0.5 + b = torch.randn(groups, state) * 0.5 + c = torch.randn(groups, state) * 0.5 + dt = torch.nn.functional.softplus(torch.randn(heads) * 0.5 + w["dt_bias"]) + a = -torch.exp(w["A_log"]) + d = w["D"] + + y_ref, state_ref = mamba2_recurrent_reference( + x=x[None, None], dt=dt[None, None], A=a, + B=b.reshape(1, 1, groups, state), C=c.reshape(1, 1, groups, state), + D=d, initial_state=torch.zeros(1, heads, state, head_dim), + ) + + shape = Mamba2Shape( + hidden_size=config["hidden_size"], num_heads=heads, head_dim=head_dim, + state_size=state, n_groups=groups, conv_kernel=config["conv_kernel"], + chunk_size=config["chunk_size"], seq_len=1, + ) + p = PlenaCompiler(mlen=MLEN, blen=4) + up = lambda n: ((n + MLEN - 1) // MLEN) * MLEN # noqa: E731 + st = p.alloc("state", up(heads * state), MLEN) + xv = p.alloc("x", up(heads), MLEN) + yv = p.alloc("y", up(heads), MLEN) + scratch = p.alloc("scratch", MLEN, MLEN) + fps = { + "b": p.fp_var("b", size=groups * state), + "c": p.fp_var("c", size=groups * state), + "da": p.fp_var("da", size=heads), + "dt": p.fp_var("dt", size=heads), + "d": p.fp_var("d", size=heads), + } + consts = p.mamba_fp_constants() + mark = len(p.get_code()) + p.ssm_decode_step_v0( + state=st, x=xv, b_fp=fps["b"], c_fp=fps["c"], da_fp=fps["da"], + dt_fp=fps["dt"], d_fp=fps["d"], y=yv, scratch=scratch, + shape=shape, consts=consts, + ) + code = p.get_code()[mark:] + + m = Machine(vlen=MLEN, vram_words=1 << 20, fpram_words=1 << 14) + base = lambda v: p.get_vram_layout(v.name).vram_base_addr # noqa: E731 + for h in range(heads): + m.write_vram_row(base(xv) + h * MLEN, x[h].tolist()) + m.write_vram_row(base(yv) + h * MLEN, [7.5] * MLEN) # garbage: must clear + for n in range(state): + m.write_vram_row(base(st) + (h * state + n) * MLEN, [0.0] * MLEN) + m.write_fpram(fps["b"].address, b.flatten().tolist()) + m.write_fpram(fps["c"].address, c.flatten().tolist()) + m.write_fpram(fps["da"].address, torch.exp(a * dt).tolist()) + m.write_fpram(fps["dt"].address, dt.tolist()) + m.write_fpram(fps["d"].address, d.tolist()) + m.write_fpram(consts.zero.address, p.mamba_fp_constant_values(shape)) + m.run(code) + + y = torch.tensor([m.read_vram_row(base(yv) + h * MLEN, head_dim) + for h in range(heads)]) + torch.testing.assert_close(y, y_ref[0, 0], rtol=1e-4, atol=1e-5) + + got_state = torch.tensor([ + [m.read_vram_row(base(st) + (h * state + n) * MLEN, head_dim) + for n in range(state)] + for h in range(heads) + ]) + torch.testing.assert_close(got_state, state_ref[0], rtol=1e-4, atol=1e-5) + + +def test_the_checkpoints_projection_layout_is_what_the_lowering_assumes(): + """`in_proj` packs `z`, `x`, `B`, `C`, `dt` in one tensor, in that order. + + The widths have to add up, and `conv1d.weight` is `[channels, 1, kernel]` + rather than `[channels, kernel]` -- the reference's `_causal_conv_step` + accepts both, and a lowering that assumed the wrong one would silently + transpose a 4-tap filter. + """ + config, w = _checkpoint() + heads, head_dim = config["num_heads"], config["head_dim"] + state, groups = config["state_size"], config["n_groups"] + inner = heads * head_dim + expected = 2 * inner + 2 * groups * state + heads + assert w["in_proj.weight"].shape == (expected, config["hidden_size"]), ( + f"in_proj is {tuple(w['in_proj.weight'].shape)}; z|x|B|C|dt should be " + f"{expected} wide" + ) + conv_channels = inner + 2 * groups * state + assert w["conv1d.weight"].shape == (conv_channels, 1, config["conv_kernel"]) + assert w["D"].shape == (heads,) and w["A_log"].shape == (heads,) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 518d75f..596c0f4 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -31,9 +31,16 @@ PLENA_Compiler/ | | |-- compiler.py # PlenaCompiler composition class | | |-- memory_state.py # Tensor/input/FP memory state | | |-- program_*.py # High-level program operations +| | | program_routed_moe.py / program_moe_shared.py # MoE substrate +| | | program_mamba_common.py # Mamba-2 shapes, conv1d, dt, gated norm +| | | program_ssd.py # Mamba-2 chunked-scan prefill +| | | program_ssm_recurrent.py # Mamba-2 decode + persistent HBM state | | +-- isa_*.py # Low-level ISA emitters +| |-- models/ # Per-architecture reference semantics +| | |-- gpt_oss/ # Routed-MoE golden + checkpoint loading +| | +-- mamba2/ # Recurrent + chunked SSD golden | +-- ops/ # Registered ATen op implementations -| |-- registry.py # Op dispatch registry +| |-- registry.py # Op dispatch registry (native_ops.yaml) | |-- plena/ # PLENA backend wrappers | +-- cpu/ # CPU reference implementations | @@ -107,5 +114,18 @@ HF config.json -> emulator smoke / utilization analysis ``` +## Architectures + +| Family | Compiler substrate | Reference | Emulator tests | +|--------|--------------------|-----------|----------------| +| Dense decoder (Llama, Qwen) | `asm_templates/` + `aten/plena/program_attention.py` | `aten/ops/cpu/` | `testbench/aten/` | +| Routed MoE (GPT-OSS, DeepSeek, Qwen2-MoE) | `aten/plena/program_routed_moe.py`, `program_moe_shared.py` | `aten/models/gpt_oss/` | `testbench/routed_moe/` | +| Mamba-2 / selective SSM | `aten/plena/program_mamba_common.py`, `program_ssd.py`, `program_ssm_recurrent.py` | `aten/models/mamba2/` | `testbench/mamba2/` | + +Mamba-2 uses two ISA extensions added with it: `V_SOFTPLUS_V` (0x39) and +`S_MAP_FP_V` (0x3A). Scope is Mamba-2 only -- Mamba-1's `A` is a +`[d_inner, state_size]` diagonal, which PLENA's scalar-only broadcast cannot +express cheaply. See the module docstrings for the full reasoning. + See `docs/COMPILATION_PIPELINES.md` for detailed comparison, known gaps, and guidance on when to use which pipeline. diff --git a/docs/COMPILATION_PIPELINES.md b/docs/COMPILATION_PIPELINES.md index 9e666ea..01a6559 100644 --- a/docs/COMPILATION_PIPELINES.md +++ b/docs/COMPILATION_PIPELINES.md @@ -167,6 +167,7 @@ call into these. | `store_act_asm.py` | Activation store-back to HBM | | `reset_reg_asm.py` | Register reset helpers | | `_imm.py`, `_k_split.py` | Large immediate + K-split utilities | +| `mamba_conv1d_asm.py` | Mamba-2 causal depthwise conv1d + SSD scan (generator path) | ### `assembler/` -- ASM text to binary @@ -217,3 +218,38 @@ call into these. 5. **MXFP8 quantization**: Generator tests that do load weights (`_build_hbm_from_hf_weights`) apply quantization, but the standard codegen path has no weight awareness at all. + +6. **Mamba-2 SSD values are approximate on this path.** The generator emits + structurally valid, correctly-costed ISA for the chunked scan, but three + pieces are marked `; APPROX:` inline: the `(chunk, chunk)` segment-sum mask is + applied per row because no outer-product primitive exists, tail chunks are + emitted at full chunk width, and the `dt` pointer is not advanced with `X`. + Instruction count and memory traffic are right; the numbers are not. Use the + ATen path (`aten/plena/program_ssd.py`, verified bit-exact against + `aten/models/mamba2/`) for anything numerical. + +7. **The ATen Mamba-2 path is MLEN-bound, and not usable at the target config.** + `Mamba2Shape.validate` requires `chunk_size == head_dim == state_size == MLEN` + because `tile_row_*` works within one MLEN column block and one `S_MAP_FP_V` + moves exactly MLEN scalars. `ssd_decay_mask_v0` additionally holds a whole `cs` + row in the 1024-slot FPRAM, capping MLEN near 1017 -- so the SSD prefill path + does not compile at the ANALYTIC `MLEN=VLEN=2048` that `plena_settings.toml` + declares active. The analytic cost model + (`analytic_models/performance/mamba2_model.py`) models the unrestricted shape + family instead, so **the analytic numbers and the compilable shapes are + currently disjoint**. Lifting this means looping `tile_col_idx` in every + emitter and blocking the cs row. + +8. **The Mamba spill path needs a Plain BF16 KV precision, which no shipped + config provides.** `SPILLED_ACTIVATION` selects the `keyvalue` precision + *class*; the width comes from the active TOML, where both KV types are MX-FP8. + Under that config `hbm_element_bytes=2` only doubles the address stride while + the data still decodes as e4m3 plus a scale stream. Call + `ProgramSSDMixin.require_bf16_kv_precision` to turn that into a hard failure, + or configure the build's KV types as Plain Fp(8,7). This is the same + requirement `linear_projection_bf16` already states. + +9. **VRAM residency is not modelled for Mamba.** `_mamba_vram_regions` keeps the + whole sequence resident, which exceeds `VECTOR_SRAM_DEPTH * VLEN` even at + seq 64. The other generator paths address VRAM the same unbounded way, so this + is not new, but a runnable Mamba lowering needs sequence tiling.