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/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/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/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/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/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_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/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/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/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_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..9beb465 --- /dev/null +++ b/aten/tests/test_kda_decode_step.py @@ -0,0 +1,625 @@ +"""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. + + "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() + 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_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_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_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_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_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/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/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) 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/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/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 +} 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 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. 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() 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, {})