From 49171dd648cdf9d3499b467fbf94ec75f17391d1 Mon Sep 17 00:00:00 2001 From: Dineth Hettiarachchi Date: Sun, 25 Oct 2026 13:57:39 +0530 Subject: [PATCH 1/2] docs: Phase 5 design (overlap, NAIVE schedule, tile operations, perf baselines) --- docs/ARCHITECTURE.md | 55 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index abc533b..64a941f 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -258,3 +258,58 @@ N = 32, non-causal (`traffic.py`): numbers in 9.3 are the Phase 4 baseline that Phase 5 will be measured against. - NAIVE mode is Phase 5. - The HBM model has no bandwidth limit beyond one word per cycle, and writes are not latency-modelled. + +## 10. Phase 5: overlap, NAIVE schedule, perf baselines (design, written before RTL) + +### 10.1 Overlap (`cfg_overlap_en`, FLASH only) + +- K and V are double-buffered in `tiny_attn_top` (`kbuf[2]`, `vbuf[2]`, 4 words each). `cur` selects the buffer in use. +- A prefetch sequencer, separate from the main FSM, starts when the group loop of K/V tile `kj` begins and `kj` is not the + block's last tile. It loads K then V of tile `kj+1` into buffer `!cur` over the one DMA engine. +- The main FSM, at the end of tile `kj`'s group loop, waits for the prefetch to finish (`C_WAITPF`, zero cycles if it has), + swaps `cur`, then runs the weight load (4 cycles, serial: PE weights are stationary) and the next group loop. +- Not overlapped: the Q load, the first K/V tile of each block, and the O store. +- Only one DMA burst is in flight at a time, so the prefetch and the main FSM never issue DMA together. The protocol + guarantees the prefetch is done before any block-boundary DMA. + +### 10.2 NAIVE schedule (`cfg_mode = NAIVE`) + +Same arithmetic as FLASH, with the normalization deferred as in the golden model (`fixed.naive`). Blocks are always +4 rows. Off-chip layout: Z tile `(qi,kj)` at `Z_BASE + 64·qi + 8·kj` (8 words, 2 int16 per word); P tile at `P_BASE + …`. +HBM grows to 2048 words (`HBM_AW = 11`). + +| Pass | Per tile pair `(qi, kj)` | Block-level | +|---|---|---| +| 1 | load K tile, weight load, `OP_NZ`: MM1 + scale/mask → Z rows out; DMA write Z tile | Q block loaded once | +| 2 | DMA read Z tile → `OP_NM`: row max into `m` | `m` is final after the last kj | +| 3 | DMA read Z tile → `OP_NP`: `P = exp2(Z − m)`, `l += ΣP`; DMA write P tile | | +| 4 | DMA read P tile, load V tile, weight load, `OP_NPV`: MM2 with P, O += P·V with α = 1 | normalize, store O | + +Masking, the `MASK` sentinel, and the causal diagonal are the same as the golden model. Pass 2 and pass 3 visit the +same tiles, so pass 2 finishes a block before pass 3 starts on it. + +Traffic follows `traffic.py` NAIVE term by term: pass 1 reads `Q` once per block and `K` per tile; passes 2 and 3 each +read `Z` per tile; pass 3 writes `P` per tile; pass 4 reads `P` and `V` per tile; O is written once per block. + +NAIVE has no overlap in Phase 5. + +### 10.3 Tile-core operations + +`tile_start` carries an operation: + +| Op | Input | Output | State effect | +|---|---|---|---| +| `OP_FLASH` | Q (streamed), K/V weights | O rows (normalized at block end) | m, l, α, O (online) | +| `OP_NZ` | Q (streamed), K weights | Z rows on the output port | none | +| `OP_NM` | Z tile (`tin`) | none | m = max(m, rowmax) | +| `OP_NP` | Z tile (`tin`) | P rows on the output port | l += ΣP (α = 1 by construction) | +| `OP_NPV` | P tile (`tin`), V weights | none | O += P·V (α forced to 1) | + +The output port (`out_*`) carries normalized O, Z or P rows; they go to the staging buffer, and the DMA writes them out. + +### 10.4 Performance baselines + +Phase 5 pins measured busy cycles for a grid of `(N, causal, qrows, L, overlap, mode)` in `sim/perf_baseline.json`. +The FLASH non-overlapped case stays exact against the closed-form model from Phase 4. The overlapped and NAIVE cases +are regression baselines measured from the RTL. Each one is checked by `sim/test_perf.py`; a change fails the test and +must be justified in review. From 8bec8d6d921d67f18b79b15dfaa0a102fb861758 Mon Sep 17 00:00:00 2001 From: Dineth Hettiarachchi Date: Sun, 25 Oct 2026 14:17:17 +0530 Subject: [PATCH 2/2] feat: Phase 5 NAIVE schedule, K/V overlap, tile operations, perf baselines naive_ctrl runs the four NAIVE passes block by block (Z and P tiles round-trip through HBM, HBM grows to 2048 words). flash_ctrl gains double-buffered K/V prefetch (cfg_overlap). The tile core gains OP_NZ/OP_NM/OP_NP/OP_NPV. E2E: bit-exact O for FLASH and NAIVE, with and without overlap, traffic and tile counts exact, FLASH sequential cycles exact against the model, overlap and NAIVE cycles pinned in sim/perf_baseline.json. --- CONTEXT.md | 10 +- README.md | 2 +- docs/ARCHITECTURE.md | 26 +++++ docs/PLAN.md | 2 +- rtl/attn_tile.sv | 115 ++++++++++++++++----- rtl/flash_ctrl.sv | 81 +++++++++++++-- rtl/naive_ctrl.sv | 220 +++++++++++++++++++++++++++++++++++++++++ rtl/tiny_attn_pkg.sv | 12 ++- rtl/tiny_attn_top.sv | 154 ++++++++++++++++++++--------- sim/lint.sh | 2 +- sim/perf_baseline.json | 80 +++++++++++++++ sim/test_e2e.py | 116 ++++++++++++++++++---- 12 files changed, 712 insertions(+), 108 deletions(-) create mode 100644 rtl/naive_ctrl.sv create mode 100644 sim/perf_baseline.json diff --git a/CONTEXT.md b/CONTEXT.md index e71f95b..423b34b 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -28,7 +28,15 @@ On a fresh machine delete the symlinks and run `.tools/setup.sh` and `.tools/ems - Phase 4 (memory + FLASH controller + top): DONE. `sram_bank, hbm_model, dma, flash_ctrl, perf_counters, tiny_attn_top` and the 32-row attn_tile. End to end: bit-exact O, HBM bytes, tiles, K/V loads and busy cycles all equal the golden / closed-form models for every N, causal, cfg_qrows and latency tested (`SEEDS=200 bash sim/run_e2e.sh`). -- Next: Phase 5 (cfg_overlap_en ping-pong K/V prefetch, NAIVE schedule, latency sweep, perf baseline pins). +- Phase 5 (overlap, NAIVE, perf pins): DONE. `naive_ctrl` (4 passes, block-by-block), FLASH double-buffered K/V prefetch + (`cfg_overlap`), tile operations OP_NZ / OP_NM / OP_NP / OP_NPV, HBM 2048 words. Full e2e (200 seeds) and unit suites + pass; perf baseline in `sim/perf_baseline.json` (re-record with `PERF_RECORD=1 COCOTB_TEST_FILTER=perf_grid bash sim/run_e2e.sh`). +- Next: Phase 6 (WASM harness: dual models, bindings, state schema, Node test against the golden fixtures). + +## Gotchas found in Phase 5 +- The NAIVE tile operations reuse the online-softmax unit: OP_NM uses its m path, OP_NP uses its P and l path with alpha = 1 + (exp2(0) = 32768), OP_NPV uses o_accumulator with alpha forced to 32768. Golden equality holds, but that identity is the thing to re-check first if NAIVE ever breaks. +- A wrong alpha in the NAIVE O update fails the e2e test (checked by injecting it). ## Gotchas found in Phase 4 - Disk: the root filesystem was 100% full during Phase 4. The Bash tool writes its output to /tmp, so with no space every diff --git a/README.md b/README.md index 35593ea..bfc4ee7 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ Fourth project in the series after [tiny-gpu](https://github.com/deaneeth/tiny-g [tiny-tpu](https://github.com/deaneeth/tiny-tpu) and tiny-cpu. Same rule as the others: nothing on screen is faked; every value comes from a `dbg_*` port on the compiled RTL. -**Status:** Phases 0-4 done: toolchain spike, bit-exact golden models and numerics study, the RTL arithmetic blocks, the tile core, and the full FLASH engine (HBM model, DMA, controller) matching the golden model bit for bit, with byte counts and cycle counts pinned. See [`docs/PLAN.md`](docs/PLAN.md) for the roadmap. +**Status:** Phases 0-5 done: the FLASH engine, the NAIVE baseline and K/V overlap are bit-exact against the golden model, with byte counts, tile counts and cycle counts verified (FLASH sequential exactly against a closed-form model; overlap and NAIVE against pinned baselines). See [`docs/PLAN.md`](docs/PLAN.md) for the roadmap. ## Quick start diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 64a941f..5d0222e 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -313,3 +313,29 @@ Phase 5 pins measured busy cycles for a grid of `(N, causal, qrows, L, overlap, The FLASH non-overlapped case stays exact against the closed-form model from Phase 4. The overlapped and NAIVE cases are regression baselines measured from the RTL. Each one is checked by `sim/test_perf.py`; a change fails the test and must be justified in review. + +### 10.5 Measured results (as built) + +From `sim/perf_baseline.json` (N = 32, non-causal, L = 10): + +| Configuration | Busy cycles | vs sequential FLASH | +|---|---|---| +| FLASH, `cfg_qrows = 4`, no overlap | 4113 | 1.00 | +| FLASH, `cfg_qrows = 4`, overlap | 2937 | 0.71 | +| FLASH, `cfg_qrows = 16`, no overlap | 2133 | 0.52 | +| FLASH, `cfg_qrows = 16`, overlap | 1657 | 0.40 | +| NAIVE | 10961 | 2.66 | + +Overlap saves about 29% at `cfg_qrows = 4` (the K/V loads are most of the time there) and about 22% at `cfg_qrows = 16`. +Causal runs follow the same pattern (2433 → 1845 at `cfg_qrows = 4`). NAIVE is 2.7× slower than FLASH at this size. Phase 5 does not break that gap down by pass; the pass +breakdown is a candidate for Phase 9 profiling. + +The FLASH sequential case equals the cycle model in `sim/golden/schedule.py` exactly. The overlap and NAIVE cases are +pinned baselines; `test_e2e.py::perf_grid` also checks that overlap is never slower than sequential for the same +configuration. + +Deviations from the first plan, recorded here so they are not forgotten: +- HBM grew from 256 to 2048 words to hold the NAIVE Z and P tiles. +- NAIVE runs block by block, not pass by pass over all blocks. Values and traffic totals are identical (the + per-row computations are independent), but the sequencing differs from `fixed.naive`'s loop order. +- Overlap applies to FLASH only. NAIVE is sequential. diff --git a/docs/PLAN.md b/docs/PLAN.md index 1836403..e0bb212 100644 --- a/docs/PLAN.md +++ b/docs/PLAN.md @@ -438,7 +438,7 @@ The estimates assume one person part-time, at tiny-cpu's pace. - Write the controller timing diagram in `ARCHITECTURE.md` **before** coding it. - **Exit:** full attention is bit-exact for 200 seeds × N × causal; HBM bytes equal `traffic.py`. -### Phase 5: NAIVE mode + overlap + perf pinning (1 week) +### Phase 5: NAIVE mode + overlap + perf pinning (1 week) ✅ DONE - NAIVE schedule (Z and P written to and read back from HBM), `cfg_overlap_en` ping-pong prefetch, and the `cfg_hbm_latency` sweep. - **Exit:** the whole random matrix (§5 layer 3) is green; `perf_baseline.json` is pinned; the accuracy test is green. *At this point the engine is fully verified hardware and worth publishing on its own.* diff --git a/rtl/attn_tile.sv b/rtl/attn_tile.sv index cba8ed8..7cff75b 100644 --- a/rtl/attn_tile.sv +++ b/rtl/attn_tile.sv @@ -1,14 +1,17 @@ -// Tile core for one Q block of up to 32 rows (8 row groups of 4), driven by flash_ctrl. +// Tile core for one Q block of up to 32 rows (8 row groups of 4), driven by flash_ctrl or naive_ctrl. // -// Per K/V tile: kv_load (4 cycles) puts Kᵀ into MM1's weights and V into MM2's weights. Then every unmasked -// row group of the block runs once through the pipeline with tile_start(tile_grp): -// MM1 (A = Q rows, W = Kᵀ) -> scale/mask -> online softmax -> MM2 (A = P rows, W = V) -> O update. -// Row state (m, l, alpha, O) is kept per row in 32-entry arrays; the row index is {group, row-in-group}. -// After the last K/V tile, blk_norm normalizes the first blk_nrows rows and emits them on out_*. +// Operations (tile_op, see tiny_attn_pkg): +// OP_FLASH Per K/V tile: kv_load (4 cycles) puts Kᵀ into MM1 and V into MM2. Then each unmasked row group runs +// MM1 (A = Q rows, W = Kᵀ) -> scale/mask -> online softmax -> MM2 (A = P, W = V) -> O update. +// Row state (m, l, alpha, O) is per row, index {group, row-in-group}. After the last tile, blk_norm. +// OP_NZ MM1 + scale/mask only: each Z row goes to the output port. No state change. +// OP_NM tin (4 Z rows) -> m = max(m, rowmax) per row. +// OP_NP tin (4 Z rows) -> P = exp2(Z - m) per row (alpha = 1 by construction), l += sum(P); P rows to the output port. +// OP_NPV tin (4 P rows, 9 bit in 16-bit lanes) through MM2 with the loaded V weights; O += P.V with alpha = 1. // -// Timing of one group tile (c0 = first cycle after tile_start): Q rows enter MM1 at c0..c0+3, S rows leave MM1 -// at c0+7..c0+10, the softmax stage runs on each, P rows enter MM2 at c0+8..c0+11, PV rows leave MM2 at -// c0+15..c0+18, tile_done at c0+19. tile_start to tile_done is 20 cycles (see sim/test_attn_tile.py). +// Timing of an OP_FLASH group tile (c0 = first cycle after tile_start): Q rows enter MM1 at c0..c0+3, S rows leave +// MM1 at c0+7..c0+10, P rows enter MM2 at c0+8..c0+11, PV rows leave at c0+15..c0+18, tile_done at c0+19. +// tile_start to tile_done is 20 cycles (sim/test_attn_tile.py). /* verilator lint_off PINCONNECTEMPTY */ // mm_engine.busy is unused: tile_busy comes from the FSM module attn_tile import tiny_attn_pkg::*; @@ -28,8 +31,10 @@ module attn_tile input logic [4*Q_W-1:0] v_row, // row r of V input logic tile_start, + input logic [2:0] tile_op, input logic [2:0] tile_grp, // row group within the block input logic tile_diag, // this group's diagonal K/V tile (causal mask inside) + input logic [4*4*16-1:0] tin, // NAIVE input tile: 4 rows x 4 lanes x 16 bit, row-major output logic tile_busy, output logic tile_done, @@ -37,8 +42,8 @@ module attn_tile input logic [5:0] blk_nrows, // rows to normalize (4 per group in use) output logic blk_done, output logic out_valid, - output logic [4:0] out_idx, // row within the block - output logic [4*OUT_W-1:0] out_row, // Q8.8 V-LSB values, lane c in bits [c*16 +: 16] + output logic [4:0] out_idx, // row within the block (or tile row for NZ / NP) + output logic [4*OUT_W-1:0] out_row, // Q8.8 V-LSB values (O), or Z / P lanes // ---- debug: real registers and wires. Tile values are for the group of the most recent tile. output logic [2:0] dbg_cur_grp, @@ -57,7 +62,7 @@ module attn_tile output logic [16*P_W-1:0] dbg_mm2_a, output logic [16*PV_W-1:0] dbg_mm2_psum ); - typedef enum logic [1:0] {S_IDLE, S_RUN, S_NORM} state_t; + typedef enum logic [1:0] {S_IDLE, S_RUN, S_ROWS, S_NORM} state_t; state_t state; // ---- per-row state (row = {group, row-in-group}) @@ -67,12 +72,13 @@ module attn_tile logic [E_W-1:0] alpha_q [32]; logic [4*O_W-1:0] o_q [32]; - // ---- per-tile values (4 rows of the current group) + // ---- per-tile values (4 rows of the current group), FLASH only logic [4*S_W-1:0] s_t [4]; logic [4*Z_W-1:0] z_t [4]; logic [4*P_W-1:0] p_t [4]; logic [2:0] cur_grp; + logic [2:0] op_q; logic diag_q, streaming, p_valid; logic [1:0] qcnt, srow, orow; logic [4*P_W-1:0] p_reg; @@ -83,13 +89,33 @@ module attn_tile assign sidx = {cur_grp, srow}; assign oidx = {cur_grp, orow}; + logic is_flash, is_nz, is_nm, is_np, is_npv, mm1_run; + assign is_flash = (op_q == OP_FLASH); + assign is_nz = (op_q == OP_NZ); + assign is_nm = (op_q == OP_NM); + assign is_np = (op_q == OP_NP); + assign is_npv = (op_q == OP_NPV); + assign mm1_run = is_flash || is_nz; + + // ---- tin rows (NAIVE): row i = tin[i*64 +: 64], lane j = bits [j*16 +: 16] + logic [63:0] tin_row; + assign tin_row = tin[{srow, 6'b0} +: 64]; + + // NPV streams P rows from tin: low 9 bits of each lane + logic [4*P_W-1:0] npv_row; + logic [7:0] npv_base; + assign npv_base = {qcnt, 6'b0}; + for (genvar j = 0; j < 4; j++) begin : g_npv + assign npv_row[j*P_W +: P_W] = tin[npv_base + 8'(j*16) +: P_W]; + end + // ---- MM1: S = Q . Kᵀ logic mm1_c_valid; logic [4*S_W-1:0] mm1_c_row; mm_engine #(.N(4), .A_W(Q_W), .A_SIGNED(1'b1), .W_W(Q_W), .ACC_W(S_W)) u_mm1 ( .clk(clk), .rst_n(rst_n), .w_row_valid(kv_load_valid), .w_row_idx(kv_load_idx), .w_row(kt_row), - .a_valid(streaming), .a_row(q_buf[{cur_grp, qcnt}]), + .a_valid(streaming && mm1_run), .a_row(q_buf[{cur_grp, qcnt}]), .c_valid(mm1_c_valid), .c_row(mm1_c_row), .busy(), .dbg_w(dbg_mm1_w), .dbg_a(dbg_mm1_a), .dbg_psum(dbg_mm1_psum) ); @@ -105,30 +131,40 @@ module attn_tile ); end - // ---- online softmax for the row leaving MM1 + // ---- online softmax: FLASH takes Z from MM1; NM and NP take Z from tin + logic [4*Z_W-1:0] z_src; logic signed [Z_W-1:0] m_new; logic [E_W-1:0] alpha_new; logic [4*P_W-1:0] p_row; logic [L_W-1:0] l_new; + assign z_src = (is_nm || is_np) ? tin_row : z_row; online_softmax u_softmax ( - .z(z_row), .m_old(m_q[sidx]), .l_old(l_q[sidx]), + .z(z_src), .m_old(m_q[sidx]), .l_old(l_q[sidx]), .m_new(m_new), .alpha(alpha_new), .p(p_row), .l_new(l_new) ); - // ---- MM2: PV = P . V + // P row as 16-bit lanes (NAIVE stores P in HBM as 16-bit lanes; the value is at most 256) + logic [4*OUT_W-1:0] p_lanes; + for (genvar j = 0; j < 4; j++) begin : g_plane + assign p_lanes[j*OUT_W +: OUT_W] = OUT_W'(p_row[j*P_W +: P_W]); + end + + // ---- MM2: PV = P . V (FLASH: P from the softmax; NPV: P from tin) logic mm2_c_valid; logic [4*PV_W-1:0] mm2_c_row; mm_engine #(.N(4), .A_W(P_W), .A_SIGNED(1'b0), .W_W(Q_W), .ACC_W(PV_W)) u_mm2 ( .clk(clk), .rst_n(rst_n), .w_row_valid(kv_load_valid), .w_row_idx(kv_load_idx), .w_row(v_row), - .a_valid(p_valid), .a_row(p_reg), + .a_valid(is_npv ? streaming : p_valid), .a_row(is_npv ? npv_row : p_reg), .c_valid(mm2_c_valid), .c_row(mm2_c_row), .busy(), .dbg_w(dbg_mm2_w), .dbg_a(dbg_mm2_a), .dbg_psum(dbg_mm2_psum) ); - // ---- O update for the row leaving MM2 (alpha of that row was stored by the softmax stage) + // ---- O update for the row leaving MM2. NPV has alpha = 1 (UQ1.15 1.0 = 32768). logic [4*O_W-1:0] o_new; - o_accumulator u_oacc (.o_old(o_q[oidx]), .alpha(alpha_q[oidx]), .pv(mm2_c_row), .o_new(o_new)); + logic [E_W-1:0] alpha_o; + assign alpha_o = is_npv ? E_W'(32768) : alpha_q[oidx]; + o_accumulator u_oacc (.o_old(o_q[oidx]), .alpha(alpha_o), .pv(mm2_c_row), .o_new(o_new)); // ---- normalization: one row per cycle, shared reciprocal logic [R_W-1:0] r_n; @@ -142,7 +178,7 @@ module attn_tile ); end - assign tile_busy = (state != S_IDLE); + assign tile_busy = (state != S_IDLE); assign dbg_cur_grp = cur_grp; assign dbg_srow = srow; @@ -159,6 +195,7 @@ module attn_tile always_ff @(posedge clk or negedge rst_n) begin if (!rst_n) begin state <= S_IDLE; + op_q <= '0; streaming <= 1'b0; p_valid <= 1'b0; p_reg <= '0; @@ -190,7 +227,7 @@ module attn_tile tile_done <= 1'b0; blk_done <= 1'b0; out_valid <= 1'b0; - p_valid <= mm1_c_valid; + p_valid <= mm1_c_valid && is_flash; if (q_load_valid) q_buf[q_load_idx] <= q_load_row; @@ -205,13 +242,14 @@ module attn_tile case (state) S_IDLE: begin if (tile_start) begin - state <= S_RUN; + op_q <= tile_op; cur_grp <= tile_grp; diag_q <= tile_diag; - streaming <= 1'b1; + streaming <= (tile_op != OP_NM) && (tile_op != OP_NP); qcnt <= '0; srow <= '0; orow <= '0; + state <= ((tile_op == OP_NM) || (tile_op == OP_NP)) ? S_ROWS : S_RUN; end else if (blk_norm) begin state <= S_NORM; nrows_q <= blk_nrows; @@ -225,7 +263,7 @@ module attn_tile if (qcnt == 2'd3) streaming <= 1'b0; end - if (mm1_c_valid) begin + if (mm1_c_valid && is_flash) begin s_t[srow] <= mm1_c_row; z_t[srow] <= z_row; p_t[srow] <= p_row; @@ -236,6 +274,17 @@ module attn_tile srow <= srow + 2'd1; end + if (mm1_c_valid && is_nz) begin + out_valid <= 1'b1; + out_idx <= {3'b0, srow}; + out_row <= z_row; + srow <= srow + 2'd1; + if (srow == 2'd3) begin + tile_done <= 1'b1; + state <= S_IDLE; + end + end + if (mm2_c_valid) begin o_q[oidx] <= o_new; orow <= orow + 2'd1; @@ -246,6 +295,22 @@ module attn_tile end end + S_ROWS: begin + // one tile row per cycle: NM updates m, NP updates l and emits P + if (is_nm) m_q[sidx] <= m_new; + if (is_np) begin + l_q[sidx] <= l_new; + out_valid <= 1'b1; + out_idx <= {3'b0, srow}; + out_row <= p_lanes; + end + srow <= srow + 2'd1; + if (srow == 2'd3) begin + tile_done <= 1'b1; + state <= S_IDLE; + end + end + S_NORM: begin out_valid <= 1'b1; out_idx <= nrow; diff --git a/rtl/flash_ctrl.sv b/rtl/flash_ctrl.sv index f1484dc..ade10de 100644 --- a/rtl/flash_ctrl.sv +++ b/rtl/flash_ctrl.sv @@ -1,9 +1,15 @@ -// FLASH-2 controller (no overlap; cfg_overlap_en is Phase 5). Q blocks of G = cfg_qrows/4 row groups, outer loop. +// FLASH-2 controller. Q blocks of G = cfg_qrows/4 row groups, outer loop. // Per block: clear state, load Q rows, then for each K/V tile kj up to the block's last tile: load K and V (DMA), // load Kᵀ/V into the engines (4 cycles), run every unmasked row group of the block through the tile core. // After the last K/V tile: normalize the block's rows, store O (DMA write), next block. // -// The state sequence is modelled cycle-exactly in sim/golden/schedule.py; tests/test_e2e.py asserts it. +// cfg_overlap = 1 (double buffering): while the groups of K/V tile kj run, a separate prefetch sequencer loads K and V +// of tile kj+1 into the other buffer over the same DMA engine. At the end of tile kj the main FSM waits for the +// prefetch if it is not finished, then swaps buffers and loads the weights. The prefetch never overlaps a block-boundary +// DMA: the last tile of a block has no prefetch, and the main FSM waits for the prefetch before any other DMA. +// +// The overlap-off state sequence is modelled cycle-exactly in sim/golden/schedule.py; the overlap-on numbers are pinned +// baselines (sim/perf_baseline.json). module flash_ctrl import tiny_attn_pkg::*; ( @@ -13,6 +19,7 @@ module flash_ctrl input logic [3:0] cfg_nb, // K/V tiles = sequence length / 4 (1..8) input logic [3:0] cfg_gg, // row groups per Q block = cfg_qrows / 4 (1, 2, 4, 8) input logic cfg_causal, + input logic cfg_overlap, output logic busy, output logic done, // one-cycle pulse @@ -23,13 +30,16 @@ module flash_ctrl output logic [HBM_AW-1:0] dma_addr, output logic [6:0] dma_len, output logic [1:0] dma_dst, // read destination: 0 = Q, 1 = K, 2 = V + output logic dma_buf, // K/V buffer for the read words input logic dma_done, // tile core output logic blk_clear, output logic kv_load_valid, output logic [1:0] kv_load_idx, + output logic kv_sel, // K/V buffer used by the weight load output logic tile_start, + output logic [2:0] tile_op, output logic [2:0] tile_grp, output logic tile_diag, input logic tile_done, @@ -50,25 +60,30 @@ module flash_ctrl C_LDV, C_LDV_W, C_WLOAD, C_TILE, C_TILE_W, + C_WAITPF, C_NORM, C_NORM_W, C_STO, C_STO_W, C_DONE } cstate_t; cstate_t st; + typedef enum logic [2:0] {PF_IDLE, PF_K_ISSUE, PF_K_W, PF_V_ISSUE, PF_V_W, PF_DONE} pstate_t; + pstate_t pf; + logic [3:0] gb; // first row group of the current block logic [3:0] bg; // row groups in the current block logic [3:0] last; // last K/V tile index of the current block logic [3:0] kj; // current K/V tile logic [3:0] gi; // row group within the block being tiled logic [1:0] wcnt; // weight-load row + logic cur; // K/V buffer of the current tile (overlap mode; always 0 otherwise) logic [3:0] nb, gg; assign nb = cfg_nb; assign gg = cfg_gg; logic [3:0] rem, bg_c, last_c, g_glob; - logic at_end, skip, last_kj; + logic at_end, skip, last_kj, pf_active, pf_k, pf_v; always_comb begin rem = nb - gb; bg_c = (rem < gg) ? rem : gg; @@ -77,23 +92,36 @@ module flash_ctrl at_end = (gi >= bg); skip = cfg_causal && (kj > g_glob); last_kj = (kj == last); + pf_k = (pf == PF_K_ISSUE) || (pf == PF_K_W); + pf_v = (pf == PF_V_ISSUE) || (pf == PF_V_W); + pf_active = pf_k || pf_v; end assign busy = (st != C_IDLE); assign done = (st == C_DONE); assign dbg_state = st; - assign dma_start = (st == C_LDQ) || (st == C_LDK) || (st == C_LDV) || (st == C_STO); - assign dma_dir = (st == C_STO); - assign dma_dst = ((st == C_LDK) || (st == C_LDK_W)) ? 2'd1 : - ((st == C_LDV) || (st == C_LDV_W)) ? 2'd2 : 2'd0; + assign dma_start = (st == C_LDQ) || (st == C_LDK) || (st == C_LDV) || (st == C_STO) || + (pf == PF_K_ISSUE) || (pf == PF_V_ISSUE); + assign dma_dir = (st == C_STO) || (st == C_STO_W); + assign dma_dst = ((st == C_LDK) || (st == C_LDK_W) || pf_k) ? 2'd1 : + ((st == C_LDV) || (st == C_LDV_W) || pf_v) ? 2'd2 : 2'd0; + assign dma_buf = pf_active ? ~cur : cur; + + // prefetch tile index is kj + 1; the main FSM does not advance kj until the prefetch is consumed + logic [3:0] kj_next; + assign kj_next = kj + 4'd1; always_comb begin case (st) C_LDK, C_LDK_W: dma_addr = HBM_AW'(K_BASE) + HBM_AW'({kj, 2'b00}); C_LDV, C_LDV_W: dma_addr = HBM_AW'(V_BASE) + HBM_AW'({kj, 2'b00}); C_STO, C_STO_W: dma_addr = HBM_AW'(O_BASE) + HBM_AW'({gb, 3'b000}); - default: dma_addr = HBM_AW'({gb, 2'b00}); // Q + default: begin + if (pf_k) dma_addr = HBM_AW'(K_BASE) + HBM_AW'({kj_next, 2'b00}); + else if (pf_v) dma_addr = HBM_AW'(V_BASE) + HBM_AW'({kj_next, 2'b00}); + else dma_addr = HBM_AW'({gb, 2'b00}); // Q + end endcase end @@ -108,24 +136,38 @@ module flash_ctrl assign blk_clear = (st == C_BLK); assign kv_load_valid = (st == C_WLOAD); assign kv_load_idx = wcnt; + assign kv_sel = cur; assign tile_start = (st == C_TILE) && !at_end && !skip; + assign tile_op = OP_FLASH; assign tile_grp = gi[2:0]; assign tile_diag = cfg_causal && (kj == g_glob); assign blk_norm = (st == C_NORM); assign blk_nrows = 6'({bg, 2'b00}); - assign ev_kv_load = (st == C_LDK); + assign ev_kv_load = (st == C_LDK) || (pf == PF_K_ISSUE); assign ev_tile = tile_start; always_ff @(posedge clk or negedge rst_n) begin if (!rst_n) begin st <= C_IDLE; + pf <= PF_IDLE; gb <= '0; bg <= '0; last <= '0; kj <= '0; gi <= '0; wcnt <= '0; + cur <= 1'b0; end else begin + // ---- prefetch sequencer (runs alongside the main FSM) + case (pf) + PF_K_ISSUE: pf <= PF_K_W; + PF_K_W: if (dma_done) pf <= PF_V_ISSUE; + PF_V_ISSUE: pf <= PF_V_W; + PF_V_W: if (dma_done) pf <= PF_DONE; + default: ; + endcase + + // ---- main FSM case (st) C_IDLE: begin if (start) begin @@ -138,6 +180,7 @@ module flash_ctrl last <= last_c; kj <= '0; gi <= '0; + cur <= 1'b0; st <= C_LDQ; end C_LDQ: st <= C_LDQ_W; @@ -156,6 +199,7 @@ module flash_ctrl if (wcnt == 2'd3) begin gi <= '0; st <= C_TILE; + if (cfg_overlap && (kj != last)) pf <= PF_K_ISSUE; end end C_TILE: begin @@ -163,9 +207,17 @@ module flash_ctrl gi <= '0; if (last_kj) begin st <= C_NORM; - end else begin + end else if (!cfg_overlap) begin kj <= kj + 4'd1; st <= C_LDK; + end else if (pf == PF_DONE) begin + kj <= kj + 4'd1; + cur <= ~cur; + pf <= PF_IDLE; + wcnt <= '0; + st <= C_WLOAD; + end else begin + st <= C_WAITPF; end end else if (skip) begin gi <= gi + 4'd1; @@ -179,6 +231,15 @@ module flash_ctrl st <= C_TILE; end end + C_WAITPF: begin + if (pf == PF_DONE) begin + kj <= kj + 4'd1; + cur <= ~cur; + pf <= PF_IDLE; + wcnt <= '0; + st <= C_WLOAD; + end + end C_NORM: st <= C_NORM_W; C_NORM_W: if (blk_done) st <= C_STO; C_STO: st <= C_STO_W; diff --git a/rtl/naive_ctrl.sv b/rtl/naive_ctrl.sv new file mode 100644 index 0000000..24c2f6f --- /dev/null +++ b/rtl/naive_ctrl.sv @@ -0,0 +1,220 @@ +// NAIVE (3-pass safe softmax, FlashAttention paper Algorithm 0) controller. Blocks are 4 rows (one row group). +// Per block qi, the four passes run in order over the block's K/V tiles kj = 0 .. last (last = qi when causal): +// pass 1 load K tile, weight load, OP_NZ: Z tile -> staging -> DMA write to Z[qi][kj] +// pass 2 DMA read Z[qi][kj] -> tin, OP_NM: row max +// pass 3 DMA read Z[qi][kj] -> tin, OP_NP: P tile -> staging -> DMA write to P[qi][kj], l accumulates +// pass 4 DMA read P[qi][kj] -> tin, load V tile, weight load, OP_NPV: O += P.V +// then normalize the block and store O. Per-row results do not depend on the order of tiles, so this block-by-block +// order produces the same values and the same traffic totals as fixed.naive (which runs pass by pass over all blocks). +// +// No K/V overlap in NAIVE (Phase 5 applies it to FLASH only). Buffer 0 is used for every K/V load. +module naive_ctrl + import tiny_attn_pkg::*; +( + input logic clk, + input logic rst_n, + input logic start, // one-cycle pulse while idle + input logic [3:0] cfg_nb, // K/V tiles = sequence length / 4 (1..8) + input logic cfg_causal, + + output logic busy, + output logic done, + + // DMA + output logic dma_start, + output logic dma_dir, // 0: read, 1: write + output logic [HBM_AW-1:0] dma_addr, + output logic [6:0] dma_len, + output logic [1:0] dma_dst, // read destination: 0 = Q, 1 = K, 2 = V, 3 = tin (Z or P tile) + output logic dma_buf, + input logic dma_done, + + // tile core + output logic blk_clear, + output logic kv_load_valid, + output logic [1:0] kv_load_idx, + output logic kv_sel, + output logic tile_start, + output logic [2:0] tile_op, + output logic [2:0] tile_grp, + output logic tile_diag, + input logic tile_done, + output logic blk_norm, + output logic [5:0] blk_nrows, + input logic blk_done, + + output logic ev_kv_load, // K tile loads (pass 1) + output logic ev_tile, // tile pairs (pass 1) + + output logic [4:0] dbg_state +); + typedef enum logic [4:0] { + N_IDLE, N_BLK, + N_LDQ, N_LDQ_W, + N_ISSUE, N_WAIT, + N_WLOAD, + N_TILE, N_TILE_W, + N_STORE, N_STORE_W, + N_NEXT, + N_NORM, N_NORM_W, + N_STO, N_STO_W, + N_DONE + } nstate_t; + nstate_t st; + + logic [2:0] qi, kj, p; // block, K/V tile, pass (1..4) + logic step; // load step inside a pass (pass 4 loads P then V) + logic [1:0] wcnt; + + logic [3:0] last; + assign last = cfg_causal ? {1'b0, qi} : (cfg_nb - 4'd1); + logic last_kj; + assign last_kj = ({1'b0, kj} == last); + + logic [HBM_AW-1:0] z_addr, p_addr; + assign z_addr = HBM_AW'(Z_BASE) + HBM_AW'({qi, kj, 3'b000}); + assign p_addr = HBM_AW'(P_BASE) + HBM_AW'({qi, kj, 3'b000}); + + assign busy = (st != N_IDLE); + assign done = (st == N_DONE); + assign dbg_state = st; + assign dma_buf = 1'b0; + + // DMA start and descriptor by state + assign dma_start = (st == N_LDQ) || (st == N_ISSUE) || (st == N_STORE) || (st == N_STO); + assign dma_dir = (st == N_STORE) || (st == N_STO); + + always_comb begin + case (st) + N_LDQ, N_LDQ_W: begin + dma_addr = HBM_AW'(Q_BASE) + HBM_AW'({qi, 2'b00}); + dma_len = 7'd4; + dma_dst = 2'd0; + end + N_STORE, N_STORE_W: begin + dma_addr = (p == 3'd1) ? z_addr : p_addr; + dma_len = 7'd8; + dma_dst = 2'd0; + end + N_STO, N_STO_W: begin + dma_addr = HBM_AW'(O_BASE) + HBM_AW'({qi, 3'b000}); + dma_len = 7'd8; + dma_dst = 2'd0; + end + default: begin // N_ISSUE, N_WAIT: loads of the current pass + if (p == 3'd1) begin + dma_addr = HBM_AW'(K_BASE) + HBM_AW'({kj, 2'b00}); + dma_len = 7'd4; + dma_dst = 2'd1; + end else if (p == 3'd4 && step == 1'b1) begin + dma_addr = HBM_AW'(V_BASE) + HBM_AW'({kj, 2'b00}); + dma_len = 7'd4; + dma_dst = 2'd2; + end else if (p == 3'd4) begin + dma_addr = p_addr; + dma_len = 7'd8; + dma_dst = 2'd3; + end else begin // passes 2 and 3: Z tile + dma_addr = z_addr; + dma_len = 7'd8; + dma_dst = 2'd3; + end + end + endcase + end + + assign blk_clear = (st == N_BLK); + assign kv_load_valid = (st == N_WLOAD); + assign kv_load_idx = wcnt; + assign kv_sel = 1'b0; + assign tile_start = (st == N_TILE); + assign tile_grp = 3'd0; + always_comb begin + case (p) + 3'd1: tile_op = OP_NZ; + 3'd2: tile_op = OP_NM; + 3'd3: tile_op = OP_NP; + default: tile_op = OP_NPV; + endcase + end + assign tile_diag = (p == 3'd1) && cfg_causal && (kj == qi); + assign blk_norm = (st == N_NORM); + assign blk_nrows = 6'd4; + assign ev_kv_load = (st == N_ISSUE) && (p == 3'd1); + assign ev_tile = (st == N_TILE) && (p == 3'd1); + + always_ff @(posedge clk or negedge rst_n) begin + if (!rst_n) begin + st <= N_IDLE; + qi <= '0; + kj <= '0; + p <= 3'd1; + step <= 1'b0; + wcnt <= '0; + end else begin + case (st) + N_IDLE: if (start) begin + qi <= '0; + st <= N_BLK; + end + N_BLK: begin + p <= 3'd1; + kj <= '0; + step <= 1'b0; + st <= N_LDQ; + end + N_LDQ: st <= N_LDQ_W; + N_LDQ_W: if (dma_done) st <= N_ISSUE; + N_ISSUE: st <= N_WAIT; + N_WAIT: begin + if (dma_done) begin + if (p == 3'd4 && step == 1'b0) begin + step <= 1'b1; + st <= N_ISSUE; + end else begin + step <= 1'b0; + wcnt <= '0; + st <= (p == 3'd1 || p == 3'd4) ? N_WLOAD : N_TILE; + end + end + end + N_WLOAD: begin + wcnt <= wcnt + 2'd1; + if (wcnt == 2'd3) st <= N_TILE; + end + N_TILE: st <= N_TILE_W; + N_TILE_W: if (tile_done) st <= (p == 3'd1 || p == 3'd3) ? N_STORE : N_NEXT; + N_STORE: st <= N_STORE_W; + N_STORE_W: if (dma_done) st <= N_NEXT; + N_NEXT: begin + if (!last_kj) begin + kj <= kj + 3'd1; + st <= N_ISSUE; + end else if (p == 3'd4) begin + st <= N_NORM; + end else begin + p <= p + 3'd1; + kj <= '0; + step <= 1'b0; + st <= N_ISSUE; + end + end + N_NORM: st <= N_NORM_W; + N_NORM_W: if (blk_done) st <= N_STO; + N_STO: st <= N_STO_W; + N_STO_W: begin + if (dma_done) begin + if (qi == cfg_nb[2:0] - 3'd1) begin + st <= N_DONE; + end else begin + qi <= qi + 3'd1; + st <= N_BLK; + end + end + end + N_DONE: st <= N_IDLE; + default: st <= N_IDLE; + endcase + end + end +endmodule diff --git a/rtl/tiny_attn_pkg.sv b/rtl/tiny_attn_pkg.sv index 4d21bd1..48ff4a5 100644 --- a/rtl/tiny_attn_pkg.sv +++ b/rtl/tiny_attn_pkg.sv @@ -25,10 +25,20 @@ package tiny_attn_pkg; localparam logic signed [Z_W-1:0] Z_MASK = 16'sh8000; // masked logit; unmasked Z saturates to +-32767 localparam logic signed [Z_W-1:0] Z_MAX = 16'sh7fff; localparam logic signed [Z_W-1:0] Z_MIN = 16'sh8001; + // Tile-core operations (tile_op) + localparam logic [2:0] OP_FLASH = 3'd0; // online softmax + O (FlashAttention) + localparam logic [2:0] OP_NZ = 3'd1; // NAIVE pass 1: Z rows out + localparam logic [2:0] OP_NM = 3'd2; // NAIVE pass 2: row max from Z + localparam logic [2:0] OP_NP = 3'd3; // NAIVE pass 3: P rows out, l + localparam logic [2:0] OP_NPV = 3'd4; // NAIVE pass 4: O += P.V + // HBM layout (word addresses). Q, K, V: one word per row (4 int8). O: two words per row (lanes 0-1, 2-3). - localparam int HBM_AW = 8; // 256 words + // Z and P tiles (NAIVE): 8 words each (4 rows x 4 lanes x 16 bit), tile (qi,kj) at base + 64*qi + 8*kj. + localparam int HBM_AW = 11; // 2048 words localparam int Q_BASE = 0; localparam int K_BASE = 32; localparam int V_BASE = 64; localparam int O_BASE = 96; + localparam int Z_BASE = 256; + localparam int P_BASE = 768; endpackage diff --git a/rtl/tiny_attn_top.sv b/rtl/tiny_attn_top.sv index dec1efe..1bf89c1 100644 --- a/rtl/tiny_attn_top.sv +++ b/rtl/tiny_attn_top.sv @@ -1,7 +1,7 @@ -// TinyAttention engine, Phase 4: FLASH-2 schedule over an off-chip (HBM) model. +// TinyAttention engine: FLASH-2 or NAIVE schedule over an off-chip (HBM) model, optionally with K/V overlap. // -// Host protocol: write Q, K, V rows into HBM with the load port (one word per row: lane c in bits [c*8 +: 8]), -// set the config, pulse start, wait for done, read O back (two words per row, lanes 0-1 then 2-3, int16). +// Host protocol: write Q, K, V rows into HBM with the load port (one word per row: lane c in bits [c*8 +: 8]). +// Set the config, pulse start, wait for done, read O back (two words per row, lanes 0-1 then 2-3, int16). // Load and readback only while idle. Counters are cleared by start and hold their values after done. module tiny_attn_top import tiny_attn_pkg::*; @@ -13,6 +13,8 @@ module tiny_attn_top input logic [3:0] cfg_nb, // K/V tiles: sequence length / 4 (1..8) input logic [3:0] cfg_gg, // row groups per Q block: cfg_qrows / 4 (1, 2, 4, 8, <= cfg_nb) input logic cfg_causal, + input logic cfg_naive, // 0: FLASH, 1: NAIVE + input logic cfg_overlap, // FLASH only: double-buffered K/V prefetch input logic [6:0] cfg_hbm_latency, // cycles before the first word of a read burst input logic [SCALE_W-1:0] cfg_scale, // from the host's quantization @@ -54,63 +56,106 @@ module tiny_attn_top output logic [16*P_W-1:0] dbg_mm2_a, output logic [16*PV_W-1:0] dbg_mm2_psum ); - // ---- controller <-> DMA - logic dma_start, dma_dir, dma_done; + // ---- controllers (only the one selected by cfg_naive is started) + logic f_busy, f_done, n_busy, n_done; + logic f_dma_start, f_dma_dir, f_dma_buf, n_dma_start, n_dma_dir, n_dma_buf; + logic [HBM_AW-1:0] f_dma_addr, n_dma_addr; + logic [6:0] f_dma_len, n_dma_len; + logic [1:0] f_dma_dst, n_dma_dst; + logic f_blk_clear, n_blk_clear, f_kv_valid, n_kv_valid, f_kv_sel, n_kv_sel; + logic [1:0] f_kv_idx, n_kv_idx; + logic f_tile_start, n_tile_start, f_tile_diag, n_tile_diag, f_blk_norm, n_blk_norm; + logic [2:0] f_tile_op, n_tile_op, f_tile_grp, n_tile_grp; + logic [5:0] f_blk_nrows, n_blk_nrows; + logic f_ev_kv, n_ev_kv, f_ev_tile, n_ev_tile; + logic [4:0] f_dbg, n_dbg; + + // ---- shared DMA and tile core + logic dma_start, dma_dir, dma_buf, dma_done; logic [HBM_AW-1:0] dma_addr; logic [6:0] dma_len; logic [1:0] dma_dst; - - // ---- DMA <-> HBM and stream - logic dma_rd_en, dma_wr_en; - logic [HBM_AW-1:0] dma_rd_addr, dma_wr_addr; - logic [31:0] hbm_rd_word, dma_wr_word; - logic wd_valid; - logic [5:0] wd_idx; - logic [31:0] wd_data, src_word; - logic [4*OUT_W-1:0] stg_rdata; - - // ---- controller <-> tile core - logic blk_clear, kv_load_valid, tile_start, tile_diag, tile_done, blk_norm, blk_done; + logic blk_clear, kv_load_valid, kv_sel, tile_start, tile_diag, tile_done, blk_norm, blk_done; logic [1:0] kv_load_idx; - logic [2:0] tile_grp; + logic [2:0] tile_op, tile_grp; logic [5:0] blk_nrows; - logic out_valid; - logic [4:0] out_idx; - logic [4*OUT_W-1:0] out_row; - logic ev_kv_load, ev_tile; - // ---- DMA read words are routed by destination + assign dma_start = cfg_naive ? n_dma_start : f_dma_start; + assign dma_dir = cfg_naive ? n_dma_dir : f_dma_dir; + assign dma_buf = cfg_naive ? n_dma_buf : f_dma_buf; + assign dma_addr = cfg_naive ? n_dma_addr : f_dma_addr; + assign dma_len = cfg_naive ? n_dma_len : f_dma_len; + assign dma_dst = cfg_naive ? n_dma_dst : f_dma_dst; + assign blk_clear = cfg_naive ? n_blk_clear : f_blk_clear; + assign kv_load_valid = cfg_naive ? n_kv_valid : f_kv_valid; + assign kv_load_idx = cfg_naive ? n_kv_idx : f_kv_idx; + assign kv_sel = cfg_naive ? n_kv_sel : f_kv_sel; + assign tile_start = cfg_naive ? n_tile_start : f_tile_start; + assign tile_op = cfg_naive ? n_tile_op : f_tile_op; + assign tile_grp = cfg_naive ? n_tile_grp : f_tile_grp; + assign tile_diag = cfg_naive ? n_tile_diag : f_tile_diag; + assign blk_norm = cfg_naive ? n_blk_norm : f_blk_norm; + assign blk_nrows = cfg_naive ? n_blk_nrows : f_blk_nrows; + assign ev_kv_load = cfg_naive ? n_ev_kv : f_ev_kv; + assign ev_tile = cfg_naive ? n_ev_tile : f_ev_tile; + assign busy = cfg_naive ? n_busy : f_busy; + assign done = cfg_naive ? n_done : f_done; + assign dbg_ctrl_state = cfg_naive ? n_dbg : f_dbg; + + // ---- K/V double buffers and the NAIVE tile input (Z or P tile: 8 words, 4 rows x 4 lanes x 16 bit) + logic [31:0] kreg [2][4]; + logic [31:0] vreg [2][4]; + logic [31:0] tin_words [8]; + logic q_load_valid; logic [4:0] q_load_idx; logic [4*Q_W-1:0] q_load_row; - logic [31:0] kreg [4]; - logic [31:0] vreg [4]; + logic [5:0] wd_idx; + logic wd_valid; + logic [31:0] wd_data, src_word; + logic [4*OUT_W-1:0] stg_rdata; + // DMA read words, routed by destination assign q_load_valid = wd_valid && (dma_dst == 2'd0); assign q_load_idx = wd_idx[4:0]; assign q_load_row = wd_data; always_ff @(posedge clk or negedge rst_n) begin if (!rst_n) begin - for (int i = 0; i < 4; i++) begin - kreg[i] <= '0; - vreg[i] <= '0; + for (int b = 0; b < 2; b++) begin + for (int i = 0; i < 4; i++) begin + kreg[b][i] <= '0; + vreg[b][i] <= '0; + end end + for (int w = 0; w < 8; w++) tin_words[w] <= '0; end else if (wd_valid) begin - if (dma_dst == 2'd1) kreg[wd_idx[1:0]] <= wd_data; - if (dma_dst == 2'd2) vreg[wd_idx[1:0]] <= wd_data; + if (dma_dst == 2'd1) kreg[dma_buf][wd_idx[1:0]] <= wd_data; + if (dma_dst == 2'd2) vreg[dma_buf][wd_idx[1:0]] <= wd_data; + if (dma_dst == 2'd3) tin_words[wd_idx[2:0]] <= wd_data; end end - // Kᵀ row r: element i is K[i][r], i.e. byte r of kreg[i]. V row r is vreg[r]. + // Kᵀ row r: element i is K[i][r], i.e. byte r of kreg[buf][i]. V row r is vreg[buf][r]. logic [4*Q_W-1:0] kt_row, v_row; + logic [31:0] kt_src [4]; for (genvar i = 0; i < 4; i++) begin : g_kt - assign kt_row[i*Q_W +: Q_W] = kreg[i][{kv_load_idx, 3'b000} +: Q_W]; + assign kt_src[i] = kreg[kv_sel][i]; + assign kt_row[i*Q_W +: Q_W] = kt_src[i][{kv_load_idx, 3'b000} +: Q_W]; + end + assign v_row = {vreg[kv_sel][kv_load_idx]}; + + logic [4*4*16-1:0] tin; + for (genvar w = 0; w < 8; w++) begin : g_tin + assign tin[w*32 +: 32] = tin_words[w]; end - assign v_row = vreg[kv_load_idx]; - // ---- output staging: one 64-bit row per block row; DMA write reads it as two 32-bit words + logic out_valid; + logic [4:0] out_idx; + logic [4*OUT_W-1:0] out_row; + + // ---- output staging: one 64-bit row per output row; DMA writes read it as two 32-bit words sram_bank #(.DEPTH(32), .W(4*OUT_W)) u_stg ( .clk(clk), .we(out_valid), .waddr(out_idx), .wdata(out_row), .raddr(wd_idx[5:1]), .rdata(stg_rdata) @@ -118,6 +163,10 @@ module tiny_attn_top assign src_word = wd_idx[0] ? stg_rdata[63:32] : stg_rdata[31:0]; // ---- HBM + logic dma_rd_en, dma_wr_en; + logic [HBM_AW-1:0] dma_rd_addr, dma_wr_addr; + logic [31:0] hbm_rd_word, dma_wr_word; + hbm_model #(.ADDR_W(HBM_AW)) u_hbm ( .clk(clk), .rst_n(rst_n), .clr(start), .rd_en(dma_rd_en), .rd_addr(dma_rd_addr), .rd_word(hbm_rd_word), @@ -136,18 +185,29 @@ module tiny_attn_top .wd_valid(wd_valid), .wd_idx(wd_idx), .wd_data(wd_data), .src_word(src_word) ); - // ---- control - flash_ctrl u_ctrl ( - .clk(clk), .rst_n(rst_n), .start(start), - .cfg_nb(cfg_nb), .cfg_gg(cfg_gg), .cfg_causal(cfg_causal), - .busy(busy), .done(done), - .dma_start(dma_start), .dma_dir(dma_dir), .dma_addr(dma_addr), .dma_len(dma_len), - .dma_dst(dma_dst), .dma_done(dma_done), - .blk_clear(blk_clear), .kv_load_valid(kv_load_valid), .kv_load_idx(kv_load_idx), - .tile_start(tile_start), .tile_grp(tile_grp), .tile_diag(tile_diag), .tile_done(tile_done), - .blk_norm(blk_norm), .blk_nrows(blk_nrows), .blk_done(blk_done), - .ev_kv_load(ev_kv_load), .ev_tile(ev_tile), - .dbg_state(dbg_ctrl_state) + // ---- controllers + flash_ctrl u_fctrl ( + .clk(clk), .rst_n(rst_n), .start(start && !cfg_naive), + .cfg_nb(cfg_nb), .cfg_gg(cfg_gg), .cfg_causal(cfg_causal), .cfg_overlap(cfg_overlap), + .busy(f_busy), .done(f_done), + .dma_start(f_dma_start), .dma_dir(f_dma_dir), .dma_addr(f_dma_addr), .dma_len(f_dma_len), + .dma_dst(f_dma_dst), .dma_buf(f_dma_buf), .dma_done(dma_done), + .blk_clear(f_blk_clear), .kv_load_valid(f_kv_valid), .kv_load_idx(f_kv_idx), .kv_sel(f_kv_sel), + .tile_start(f_tile_start), .tile_op(f_tile_op), .tile_grp(f_tile_grp), .tile_diag(f_tile_diag), + .tile_done(tile_done), .blk_norm(f_blk_norm), .blk_nrows(f_blk_nrows), .blk_done(blk_done), + .ev_kv_load(f_ev_kv), .ev_tile(f_ev_tile), .dbg_state(f_dbg) + ); + + naive_ctrl u_nctrl ( + .clk(clk), .rst_n(rst_n), .start(start && cfg_naive), + .cfg_nb(cfg_nb), .cfg_causal(cfg_causal), + .busy(n_busy), .done(n_done), + .dma_start(n_dma_start), .dma_dir(n_dma_dir), .dma_addr(n_dma_addr), .dma_len(n_dma_len), + .dma_dst(n_dma_dst), .dma_buf(n_dma_buf), .dma_done(dma_done), + .blk_clear(n_blk_clear), .kv_load_valid(n_kv_valid), .kv_load_idx(n_kv_idx), .kv_sel(n_kv_sel), + .tile_start(n_tile_start), .tile_op(n_tile_op), .tile_grp(n_tile_grp), .tile_diag(n_tile_diag), + .tile_done(tile_done), .blk_norm(n_blk_norm), .blk_nrows(n_blk_nrows), .blk_done(blk_done), + .ev_kv_load(n_ev_kv), .ev_tile(n_ev_tile), .dbg_state(n_dbg) ); // ---- tile core @@ -156,7 +216,7 @@ module tiny_attn_top .blk_clear(blk_clear), .q_load_valid(q_load_valid), .q_load_idx(q_load_idx), .q_load_row(q_load_row), .kv_load_valid(kv_load_valid), .kv_load_idx(kv_load_idx), .kt_row(kt_row), .v_row(v_row), - .tile_start(tile_start), .tile_grp(tile_grp), .tile_diag(tile_diag), + .tile_start(tile_start), .tile_op(tile_op), .tile_grp(tile_grp), .tile_diag(tile_diag), .tin(tin), .tile_busy(dbg_tile_busy), .tile_done(tile_done), .blk_norm(blk_norm), .blk_nrows(blk_nrows), .blk_done(blk_done), .out_valid(out_valid), .out_idx(out_idx), .out_row(out_row), diff --git a/sim/lint.sh b/sim/lint.sh index 184f91e..391a020 100644 --- a/sim/lint.sh +++ b/sim/lint.sh @@ -5,7 +5,7 @@ set -euo pipefail cd "$(dirname "$0")/../rtl" PKGS="lut_pkg.sv tiny_attn_pkg.sv" ALL="$PKGS $(ls *.sv | grep -v '^lut_pkg.sv$' | grep -v '^tiny_attn_pkg.sv$' | grep -v '^spike' | tr '\n' ' ')" -for top in pe systolic_array mm_engine row_max scale_mask exp2_unit recip_unit online_softmax o_accumulator normalize attn_tile dma hbm_model sram_bank flash_ctrl perf_counters tiny_attn_top; do +for top in pe systolic_array mm_engine row_max scale_mask exp2_unit recip_unit online_softmax o_accumulator normalize attn_tile dma hbm_model sram_bank flash_ctrl naive_ctrl perf_counters tiny_attn_top; do echo "lint: $top" verilator --lint-only -Wall --top-module "$top" $ALL done diff --git a/sim/perf_baseline.json b/sim/perf_baseline.json new file mode 100644 index 0000000..29e405d --- /dev/null +++ b/sim/perf_baseline.json @@ -0,0 +1,80 @@ +{ + "flash/N16/c0/q16/o0/L0": 485, + "flash/N16/c0/q16/o0/L10": 575, + "flash/N16/c0/q16/o0/L40": 845, + "flash/N16/c0/q16/o1/L0": 443, + "flash/N16/c0/q16/o1/L10": 473, + "flash/N16/c0/q16/o1/L40": 593, + "flash/N16/c0/q4/o0/L0": 737, + "flash/N16/c0/q4/o0/L10": 1097, + "flash/N16/c0/q4/o0/L40": 2177, + "flash/N16/c0/q4/o1/L0": 569, + "flash/N16/c0/q4/o1/L10": 845, + "flash/N16/c0/q4/o1/L40": 1925, + "flash/N16/c1/q16/o0/L0": 365, + "flash/N16/c1/q16/o0/L10": 455, + "flash/N16/c1/q16/o0/L40": 725, + "flash/N16/c1/q16/o1/L0": 323, + "flash/N16/c1/q16/o1/L10": 353, + "flash/N16/c1/q16/o1/L40": 533, + "flash/N16/c1/q4/o0/L0": 497, + "flash/N16/c1/q4/o0/L10": 737, + "flash/N16/c1/q4/o0/L40": 1457, + "flash/N16/c1/q4/o1/L0": 413, + "flash/N16/c1/q4/o1/L10": 611, + "flash/N16/c1/q4/o1/L40": 1331, + "flash/N32/c0/q16/o0/L0": 1793, + "flash/N32/c0/q16/o0/L10": 2133, + "flash/N32/c0/q16/o0/L40": 3153, + "flash/N32/c0/q16/o1/L0": 1597, + "flash/N32/c0/q16/o1/L10": 1657, + "flash/N32/c0/q16/o1/L40": 1977, + "flash/N32/c0/q4/o0/L0": 2753, + "flash/N32/c0/q4/o0/L10": 4113, + "flash/N32/c0/q4/o0/L40": 8193, + "flash/N32/c0/q4/o1/L0": 1969, + "flash/N32/c0/q4/o1/L10": 2937, + "flash/N32/c0/q4/o1/L40": 7017, + "flash/N32/c1/q16/o0/L0": 1141, + "flash/N32/c1/q16/o0/L10": 1401, + "flash/N32/c1/q16/o0/L40": 2181, + "flash/N32/c1/q16/o1/L0": 1001, + "flash/N32/c1/q16/o1/L10": 1061, + "flash/N32/c1/q16/o1/L40": 1461, + "flash/N32/c1/q4/o0/L0": 1633, + "flash/N32/c1/q4/o0/L10": 2433, + "flash/N32/c1/q4/o0/L40": 4833, + "flash/N32/c1/q4/o1/L0": 1241, + "flash/N32/c1/q4/o1/L10": 1845, + "flash/N32/c1/q4/o1/L40": 4245, + "flash/N8/c0/q4/o0/L0": 209, + "flash/N8/c0/q4/o0/L10": 309, + "flash/N8/c0/q4/o0/L40": 609, + "flash/N8/c0/q4/o1/L0": 181, + "flash/N8/c0/q4/o1/L10": 267, + "flash/N8/c0/q4/o1/L40": 567, + "flash/N8/c1/q4/o0/L0": 169, + "flash/N8/c1/q4/o0/L10": 249, + "flash/N8/c1/q4/o0/L40": 489, + "flash/N8/c1/q4/o1/L0": 155, + "flash/N8/c1/q4/o1/L10": 228, + "flash/N8/c1/q4/o1/L40": 468, + "naive/N16/c0/L0": 1969, + "naive/N16/c0/L10": 2809, + "naive/N16/c0/L40": 5329, + "naive/N16/c1/L0": 1267, + "naive/N16/c1/L10": 1807, + "naive/N16/c1/L40": 3427, + "naive/N32/c0/L0": 7681, + "naive/N32/c0/L10": 10961, + "naive/N32/c0/L40": 20801, + "naive/N32/c1/L0": 4405, + "naive/N32/c1/L10": 6285, + "naive/N32/c1/L40": 11925, + "naive/N8/c0/L0": 517, + "naive/N8/c0/L10": 737, + "naive/N8/c0/L40": 1397, + "naive/N8/c1/L0": 400, + "naive/N8/c1/L10": 570, + "naive/N8/c1/L40": 1080 +} diff --git a/sim/test_e2e.py b/sim/test_e2e.py index a62ef45..09ab674 100644 --- a/sim/test_e2e.py +++ b/sim/test_e2e.py @@ -1,13 +1,16 @@ -"""End-to-end FLASH through tiny_attn_top, against the golden model and the closed forms. +"""End-to-end engine test through tiny_attn_top, against the golden model and the closed forms. Per run, checked exactly: - * O (every value, int16 Q8.8) == fixed.flash(qblock=cfg_qrows) (bit-exact) - * HBM read bytes / write bytes == traffic.py (FLASH) - * tiles computed, K/V tile loads == traffic.py / kv_tile_loads - * busy cycles == sim/golden/schedule.py (the controller's state sequence) - -SEEDS (env, default 200): each seed runs every sequence length N in {4..32} x causal {off, on}, with random -cfg_qrows (<= N), cfg_hbm_latency in 0..64 and random inputs. qrows_sweep covers every cfg_qrows value. + * O (every value, int16 Q8.8) == fixed.flash (FLASH) or fixed.naive (NAIVE), bit for bit + * HBM read / write bytes == traffic.py for the mode + * tiles computed == traffic.py tile pairs + * K/V tile loads == kv_tile_loads (FLASH) or tile pairs (NAIVE: one K load per pair) + * busy cycles == sim/golden/schedule.py, FLASH without overlap only (the other cases are + pinned baselines in sim/perf_baseline.json, see test_perf.py) + +SEEDS (env, default 200): every seed runs every N in {4..32} x causal {off, on}, with random mode (FLASH or NAIVE), +random overlap (FLASH), random cfg_qrows (<= N), cfg_hbm_latency in 0..64 and random inputs. +qrows_sweep covers every cfg_qrows value and both overlap settings. """ import os @@ -27,7 +30,7 @@ async def idle(dut): - for sig in ("start", "ld_en", "cfg_causal"): + for sig in ("start", "ld_en", "cfg_causal", "cfg_naive", "cfg_overlap"): getattr(dut, sig).value = 0 dut.ld_addr.value = 0 dut.ld_word.value = 0 @@ -53,10 +56,14 @@ async def host_load(dut, qz, n): dut.ld_en.value = 0 -async def run_one(dut, qz, n, causal, qrows, latency): +async def run_one(dut, qz, n, causal, qrows, latency, naive=False, overlap=False): + mode = "NAIVE" if naive else "FLASH" + qblock = 4 if naive else qrows dut.cfg_nb.value = n // 4 - dut.cfg_gg.value = qrows // 4 + dut.cfg_gg.value = qblock // 4 dut.cfg_causal.value = int(causal) + dut.cfg_naive.value = int(naive) + dut.cfg_overlap.value = int(overlap and not naive) dut.cfg_hbm_latency.value = latency dut.cfg_scale.value = qz.cfg_scale await host_load(dut, qz, n) @@ -73,15 +80,19 @@ async def run_one(dut, qz, n, causal, qrows, latency): raise AssertionError("engine never finished") await next_cycle(dut) # the DONE cycle is counted; read the counters after it - ctx = f"N={n} causal={causal} qrows={qrows} L={latency}" - t = traffic(n, causal, "FLASH", qblock=qrows) + ctx = f"{mode} N={n} causal={causal} qrows={qblock} L={latency} overlap={int(overlap and not naive)}" + t = traffic(n, causal, mode, qblock=qblock) assert int(dut.hbm_rd_bytes.value) == t["rd"], f"{ctx}: read bytes {int(dut.hbm_rd_bytes.value)} != {t['rd']}" assert int(dut.hbm_wr_bytes.value) == t["wr"], f"{ctx}: write bytes {int(dut.hbm_wr_bytes.value)} != {t['wr']}" assert int(dut.tiles.value) == t["tiles"], f"{ctx}: tiles {int(dut.tiles.value)} != {t['tiles']}" - kv = kv_tile_loads(n, causal, qrows) + if naive: + kv = t["tiles"] # one K tile load per tile pair (pass 1) + else: + kv = kv_tile_loads(n, causal, qrows) assert int(dut.kv_loads.value) == kv, f"{ctx}: K/V loads {int(dut.kv_loads.value)} != {kv}" - want_cycles = sched_cycles(n, causal, qrows, latency) - assert int(dut.cycles.value) == want_cycles, f"{ctx}: busy cycles {int(dut.cycles.value)} != model {want_cycles}" + if not naive and not overlap: + want_cycles = sched_cycles(n, causal, qrows, latency) + assert int(dut.cycles.value) == want_cycles, f"{ctx}: busy cycles {int(dut.cycles.value)} != model {want_cycles}" got = np.zeros((n, 4), dtype=np.int64) for i in range(n): @@ -93,12 +104,13 @@ async def run_one(dut, qz, n, causal, qrows, latency): v = (w >> (16 * lane)) & 0xFFFF got[i, 2 * half + lane] = v - 65536 if v >= 32768 else v await next_cycle(dut) - want = fixed.flash(qz, n, causal, qblock=qrows).out.astype(np.int64) + want = (fixed.naive(qz, n, causal) if naive else fixed.flash(qz, n, causal, qblock=qrows)).out.astype(np.int64) if not np.array_equal(got, want): bad = np.argwhere(got != want)[0] raise AssertionError( f"{ctx}: O differs at {tuple(int(b) for b in bad)}: expected {want[tuple(bad)]} got {got[tuple(bad)]}" ) + return int(dut.cycles.value) async def setup(dut): @@ -109,7 +121,7 @@ async def setup(dut): @cocotb.test() async def qrows_sweep(dut): - """Every cfg_qrows value x every N x causal, two seeds each. The output must not depend on qrows.""" + """Every cfg_qrows x N x causal, FLASH with and without overlap, plus NAIVE. Two seeds each.""" await setup(dut) rng = random.Random(51) for _ in range(2): @@ -117,17 +129,79 @@ async def qrows_sweep(dut): for causal in (False, True): qz = fixed.quantize(*make_inputs(rng, n, 2.5)) for qrows in [q for q in (4, 8, 16, 32) if q <= n]: - await run_one(dut, qz, n, causal, qrows, rng.randint(0, 64)) + for overlap in (False, True): + await run_one(dut, qz, n, causal, qrows, rng.randint(0, 64), overlap=overlap) + await run_one(dut, qz, n, causal, 4, rng.randint(0, 64), naive=True) @cocotb.test() async def random_seeds(dut): - """SEEDS seeds x every N x causal, random cfg_qrows (<= N), latency 0..64 and input amplitude.""" + """SEEDS seeds x every N x causal, random mode, overlap, cfg_qrows (<= N), latency 0..64 and amplitude.""" await setup(dut) rng = random.Random(52) for _ in range(SEEDS): for n in NS: for causal in (False, True): + naive = rng.random() < 0.35 qrows = rng.choice([q for q in (4, 8, 16, 32) if q <= n]) + overlap = rng.random() < 0.5 qz = fixed.quantize(*make_inputs(rng, n, rng.choice([0.5, 1.5, 2.5, 3.5]))) - await run_one(dut, qz, n, causal, qrows, rng.randint(0, 64)) + await run_one(dut, qz, n, causal, qrows, rng.randint(0, 64), naive=naive, overlap=overlap) + + +def _parse_flash_key(key): + fields = key.split("/") + n = int(fields[1][1:]) + causal = bool(int(fields[2][1:])) + qrows = int(fields[3][1:]) + latency = int(fields[5][1:]) + return n, causal, qrows, latency + + +PERF_BASELINE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "perf_baseline.json") + + +@cocotb.test() +async def perf_grid(dut): + """Busy cycles over a grid of N x causal x mode x overlap x latency. Pinned in sim/perf_baseline.json. + + PERF_RECORD=1 rewrites the baseline. Otherwise every value must equal the baseline, and overlap must never be + slower than the sequential schedule for the same configuration. + """ + import json + + await setup(dut) + rng = random.Random(53) + measured = {} + for n in (8, 16, 32): + for causal in (False, True): + qz = fixed.quantize(*make_inputs(rng, n, 2.5)) + for qrows in (4, 16): + if qrows > n: + continue + for overlap in (False, True): + for latency in (0, 10, 40): + key = f"flash/N{n}/c{int(causal)}/q{qrows}/o{int(overlap)}/L{latency}" + measured[key] = await run_one(dut, qz, n, causal, qrows, latency, overlap=overlap) + for latency in (0, 10, 40): + key = f"naive/N{n}/c{int(causal)}/L{latency}" + measured[key] = await run_one(dut, qz, n, causal, 4, latency, naive=True) + + for key, cyc in measured.items(): + if key.startswith("flash/") and "/o0/" in key: + n, causal, qrows, latency = _parse_flash_key(key) + assert cyc == sched_cycles(n, causal, qrows, latency), f"{key}: {cyc} != schedule model" + if "/o1/" in key: + seq = measured[key.replace("/o1/", "/o0/")] + assert cyc <= seq, f"{key}: overlap {cyc} slower than sequential {seq}" + + if os.environ.get("PERF_RECORD"): + with open(PERF_BASELINE, "w") as f: + json.dump(dict(sorted(measured.items())), f, indent=1) + f.write("\n") + return + with open(PERF_BASELINE) as f: + baseline = json.load(f) + assert set(baseline) == set(measured), "perf grid changed: re-record with PERF_RECORD=1" + diffs = {k: (baseline[k], measured[k]) for k in measured if baseline[k] != measured[k]} + assert not diffs, f"busy cycles changed (baseline, measured): {dict(list(diffs.items())[:8])}"