Skip to main content

logicaffeine_compile/spatial/decode_engine/rtl/
attend.rs

1//! The attention unit's vector path as Verilog: one query head over its sequence's cached rows, exactly as
2//! [`crate::spatial::decode_engine::numerics::attend_vector`] computes it.
3//!
4//! Three passes over the positions. The keys stream in, `HD/16` beats of sixteen elements a position: each beat's
5//! products with the head's queries accumulate in sixteen lanes (element `i` in lane `i mod 16`), the lanes fold by the
6//! pairwise tree, the dot times the scale is the position's score, kept, and the scores' maximum is taken as the
7//! softmax takes it (a NaN never larger, a largest zero `+0`). Then each `exp(s − max)` is kept and summed in sixteen
8//! lanes by position, the sum folded by the tree and its reciprocal taken. Then the values stream in, `VD/16` beats a
9//! position: each position's probability `e·(1/Σ)` times its row adds into lane `position mod 16` of every element,
10//! and the lanes fold by the tree into the head's output, sixteen elements a beat.
11
12/// The unit's two paths, with the f32 units, the functions, the lane's one rounding and the block RAM they are built
13/// from.
14pub fn attend_verilog() -> String {
15    format!(
16        "{}\n{}\n{}\n{}\n{}",
17        super::ram::ram_verilog(&[32]),
18        super::vector::vector_verilog(),
19        ROUND_VERILOG,
20        ATTEND_VERILOG,
21        ATTEND_INT8_VERILOG
22    )
23}
24
25/// The lane's rounding, `v·2^least` to f32 once, without the lane around it.
26const ROUND_VERILOG: &str = r#"
27module logos_round88(input signed [87:0] v, input signed [11:0] least, output reg [31:0] f);
28  reg [87:0] mag;
29  reg        neg;
30  integer    lead, i, s;
31  reg [87:0] kept, rem, half;
32  reg [24:0] m;
33  reg signed [13:0] unb;
34  always @* begin
35    neg = v[87];
36    mag = neg ? -v : v;
37    lead = -1;
38    for (i = 0; i < 88; i = i + 1) if (mag[i]) lead = i;
39    f = 32'd0;
40    if (lead >= 0) begin
41      if (lead <= 23) begin
42        m = mag[24:0] << (23 - lead);
43        unb = least + $signed(lead[13:0]);
44      end else begin
45        s = lead - 23;
46        kept = mag >> s;
47        rem = mag & ((88'd1 << s) - 88'd1);
48        half = 88'd1 << (s - 1);
49        m = kept[24:0] + ((rem > half || (rem == half && kept[0])) ? 25'd1 : 25'd0);
50        unb = least + $signed(lead[13:0]);
51        if (m[24]) begin m = m >> 1; unb = unb + 14'sd1; end
52      end
53      if (unb > 14'sd127) f = {neg, 8'hff, 23'd0};
54      else if (unb < -14'sd126) f = {neg, 31'd0};
55      else f = {neg, unb[7:0] + 8'd127, m[22:0]};
56    end
57  end
58endmodule
59"#;
60
61const ATTEND_VERILOG: &str = r#"
62// One query head of HD elements over up to P_MAX positions of VD-element values.
63module logos_attend_vector #(parameter HD = 64, VD = 64, P_MAX = 64)
64  (input clk, input rst,
65   input start, input [15:0] positions, input [32*HD-1:0] q, input [31:0] scale, output busy,
66   input k_valid, input [511:0] k_beat, output k_ready,
67   input v_valid, input [511:0] v_beat, output v_ready,
68   output reg o_valid, output reg [511:0] o_beat, output reg o_last);
69  localparam [15:0] KB = HD / 16, VB = VD / 16;
70  localparam IDLE = 3'd0, KEY = 3'd1, FOLD = 3'd2, EXP = 3'd3, INV = 3'd4, VALUE = 3'd5, OUT = 3'd6;
71  reg [2:0] state;
72  reg [15:0] j, b, n;
73  reg [31:0] acc [0:16*VD-1];
74  reg [511:0] lanes;
75  reg [31:0] m, sc, inv;
76  reg [32*HD-1:0] qh;
77  assign busy = (state != IDLE);
78  assign k_ready = (state == KEY);
79  assign v_ready = (state == VALUE);
80
81  // A key beat's products into its sixteen lanes.
82  wire [511:0] q_beat = qh[512*b +: 512];
83  wire [511:0] k_next;
84  genvar l, t;
85  generate
86    for (l = 0; l < 16; l = l + 1) begin : kl
87      wire [31:0] p;
88      logos_fmul mk(.a(q_beat[32*l +: 32]), .b(k_beat[32*l +: 32]), .y(p));
89      logos_fadd ak(.a(b == 16'd0 ? 32'd0 : lanes[32*l +: 32]), .b(p), .y(k_next[32*l +: 32]));
90    end
91  endgenerate
92  // The fold of a key's lanes: its dot, its score, whether it is the largest yet.
93  wire [31:0] dot, score;
94  logos_tree16 tk(.lanes(lanes), .sum(dot));
95  logos_fmul ms(.a(dot), .b(sc), .y(score));
96  wire score_nan = (score[30:23] == 8'hff) && (score[22:0] != 23'd0);
97  wire m_lt;
98  logos_flt cmp(.a(m), .b(score), .y(m_lt));
99  // The largest score, a largest zero +0.
100  wire [31:0] mx = (m[30:0] == 31'd0) ? 32'd0 : m;
101  // One position's exp(s - max), its score and exp read back from block RAM.
102  wire [31:0] score_j, ex_j;
103  wire [31:0] diff, e_j;
104  logos_fsub sd(.a(score_j), .b(mx), .y(diff));
105  logos_exp xe(.x_in(diff), .y(e_j));
106  wire [31:0] e_lane, e_next;
107  assign e_lane = lanes[32*j[3:0] +: 32];
108  logos_fadd ae(.a(e_lane), .b(e_j), .y(e_next));
109  // The sum of the exps and its reciprocal.
110  wire [31:0] total, total_inv;
111  logos_tree16 te(.lanes(lanes), .sum(total));
112  logos_recip rr(.x_in(total), .y(total_inv));
113  // A value beat: the position's probability times each element into lane j mod 16 of its accumulator.
114  wire [31:0] prob;
115  logos_fmul mp(.a(ex_j), .b(inv), .y(prob));
116  wire [511:0] v_next;
117  generate
118    for (l = 0; l < 16; l = l + 1) begin : vl
119      wire [31:0] pv;
120      logos_fmul mv(.a(prob), .b(v_beat[32*l +: 32]), .y(pv));
121      logos_fadd av(.a(acc[VD*j[3:0] + 16*b + l]), .b(pv), .y(v_next[32*l +: 32]));
122    end
123  endgenerate
124  // An output beat: each of its sixteen elements' sixteen position-lanes folded by the tree.
125  wire [511:0] out_beat;
126  generate
127    for (l = 0; l < 16; l = l + 1) begin : ol
128      wire [511:0] col;
129      for (t = 0; t < 16; t = t + 1) begin : ot
130        assign col[32*t +: 32] = acc[VD*t + 16*b + l];
131      end
132      logos_tree16 to(.lanes(col), .sum(out_beat[32*l +: 32]));
133    end
134  endgenerate
135
136  integer k;
137  // The scores and their exponentials in block RAM, each read and written at the position at hand.
138  localparam PA = (P_MAX > 1) ? $clog2(P_MAX) : 1;
139  logos_ram_w32 #(.D(P_MAX), .A(PA)) score_ram(.clk(clk), .ra(j[PA-1:0]), .rd(score_j), .we(!rst && state == FOLD),
140    .wa(j[PA-1:0]), .wd(score), .wm(4'hf));
141  logos_ram_w32 #(.D(P_MAX), .A(PA)) ex_ram(.clk(clk), .ra(j[PA-1:0]), .rd(ex_j), .we(!rst && state == EXP),
142    .wa(j[PA-1:0]), .wd(e_j), .wm(4'hf));
143  always @(posedge clk) begin
144    o_valid <= 1'b0;
145    o_last <= 1'b0;
146    if (rst) state <= IDLE;
147    else case (state)
148      IDLE: if (start) begin
149        qh <= q; sc <= scale; n <= positions; j <= 16'd0; b <= 16'd0; m <= 32'hff80_0000;
150        for (k = 0; k < 16*VD; k = k + 1) acc[k] <= 32'd0;
151        state <= (positions == 16'd0) ? OUT : KEY;
152      end
153      KEY: if (k_valid) begin
154        lanes <= k_next;
155        if (b + 16'd1 == KB) begin b <= 16'd0; state <= FOLD; end else b <= b + 16'd1;
156      end
157      FOLD: begin
158        if (!score_nan && m_lt) m <= score;
159        if (j + 16'd1 == n) begin j <= 16'd0; lanes <= 512'd0; state <= EXP; end
160        else begin j <= j + 16'd1; state <= KEY; end
161      end
162      EXP: begin
163        lanes[32*j[3:0] +: 32] <= e_next;
164        if (j + 16'd1 == n) begin j <= 16'd0; state <= INV; end else j <= j + 16'd1;
165      end
166      INV: begin inv <= total_inv; b <= 16'd0; state <= VALUE; end
167      VALUE: if (v_valid) begin
168        for (k = 0; k < 16; k = k + 1) acc[VD*j[3:0] + 16*b + k] <= v_next[32*k +: 32];
169        if (b + 16'd1 == VB) begin
170          b <= 16'd0;
171          if (j + 16'd1 == n) state <= OUT; else j <= j + 16'd1;
172        end else b <= b + 16'd1;
173      end
174      OUT: begin
175        o_valid <= 1'b1;
176        o_beat <= out_beat;
177        if (b + 16'd1 == VB) begin o_last <= 1'b1; b <= 16'd0; state <= IDLE; end else b <= b + 16'd1;
178      end
179    endcase
180  end
181endmodule
182"#;
183
184const ATTEND_INT8_VERILOG: &str = r#"
185// One query head of HD elements over up to P_MAX positions of an int8 cache on the array, the query one block (its
186// codes and its scale m·2^e): each key row's codes against the query's, exact, times the key's f16 significand and m,
187// rounded once and times the score scale; the softmax; then per block of 32 positions the probabilities times the
188// value rows' scales, their maximum's exponent, each an 8-bit code to nearest even, every element's exact integer sum
189// against the value codes, rounded once a block and added in order.
190module logos_attend_int8 #(parameter HD = 64, VD = 64, P_MAX = 64)
191  (input clk, input rst,
192   input start, input [15:0] positions, input [16*HD-1:0] q_codes, input [7:0] q_m, input signed [11:0] q_e,
193   input [31:0] scale, output busy,
194   input k_valid, input [8*HD-1:0] k_codes, input [15:0] k_scale, output k_ready,
195   input s_valid, input [15:0] v_scale, output s_ready,
196   input v_valid, input [8*VD-1:0] v_codes, output v_ready,
197   output reg o_valid, output reg [511:0] o_beat, output reg o_last);
198  localparam [15:0] OB = VD / 16, ELEMS = VD;
199  localparam IDLE = 4'd0, KEY = 4'd1, EXP = 4'd2, INV = 4'd3, WSCALE = 4'd4, WCODE = 4'd5, WVAL = 4'd6, WROUND = 4'd7, OUT = 4'd8;
200  reg [3:0] state;
201  reg [15:0] j, n, blk, b, el;
202  `LOGOS_RAMSTYLE_ASYNC(32) reg [31:0] w [0:31];
203  reg signed [31:0] sums [0:VD-1];
204  reg [31:0] out [0:VD-1];
205  reg [511:0] lanes;
206  reg [31:0] m, sc, inv, wmax;
207  reg [16*HD-1:0] qc;
208  reg [7:0] qm;
209  reg signed [11:0] qe;
210  reg signed [11:0] we;
211  assign busy = (state != IDLE);
212  assign k_ready = (state == KEY);
213  assign s_ready = (state == WSCALE);
214  assign v_ready = (state == WVAL);
215  // The key row's exact dot with the query's codes.
216  reg signed [47:0] kdot;
217  integer i;
218  always @* begin
219    kdot = 48'sd0;
220    for (i = 0; i < HD; i = i + 1) kdot = kdot + $signed(qc[16*i +: 16]) * $signed(k_codes[8*i +: 8]);
221  end
222  // The key's f16 scale as a signed significand and an exponent, a denormal's the least.
223  wire signed [12:0] k_sig = (k_scale[14:10] == 5'd0) ? (k_scale[15] ? -$signed({3'd0, k_scale[9:0]}) : $signed({3'd0, k_scale[9:0]}))
224                                                    : (k_scale[15] ? -$signed({2'd0, 1'b1, k_scale[9:0]}) : $signed({2'd0, 1'b1, k_scale[9:0]}));
225  wire signed [11:0] k_exp = (k_scale[14:10] == 5'd0) ? -12'sd24 : $signed({7'd0, k_scale[14:10]}) - 12'sd25;
226  wire signed [87:0] term = $signed(k_sig) * $signed({1'b0, qm}) * kdot;
227  wire [31:0] rounded, score;
228  logos_round88 rs(.v(term), .least(k_exp + qe), .f(rounded));
229  logos_fmul ms(.a(rounded), .b(sc), .y(score));
230  wire score_nan = (score[30:23] == 8'hff) && (score[22:0] != 23'd0);
231  wire m_lt;
232  logos_flt cmp(.a(m), .b(score), .y(m_lt));
233  wire [31:0] mx = (m[30:0] == 31'd0) ? 32'd0 : m;
234  wire [31:0] score_j, ex_j;
235  wire [31:0] diff, e_j, e_next;
236  logos_fsub sd(.a(score_j), .b(mx), .y(diff));
237  logos_exp xe(.x_in(diff), .y(e_j));
238  logos_fadd ae(.a(lanes[32*j[3:0] +: 32]), .b(e_j), .y(e_next));
239  wire [31:0] total, total_inv;
240  logos_tree16 te(.lanes(lanes), .sum(total));
241  logos_recip rr(.x_in(total), .y(total_inv));
242  // A position's weight: its probability times its value row's scale.
243  wire [31:0] prob, vs, wj;
244  logos_fmul mp(.a(ex_j), .b(inv), .y(prob));
245  logos_f16_to_f32 cv(.h(v_scale), .x(vs));
246  logos_fmul mw(.a(prob), .b(vs), .y(wj));
247  // The block's exponent: the least e with wmax <= 255·2^e.
248  wire [23:0] wm = {1'b1, wmax[22:0]};
249  wire signed [11:0] w_exp = $signed({4'd0, wmax[30:23]}) - 12'sd150 + 12'sd16 + ((wm > 24'd16711680) ? 12'sd1 : 12'sd0);
250  // A weight's 8-bit code: its significand shifted down to the block's exponent, nearest even.
251  wire [31:0] wv = w[j[4:0]];
252  wire [23:0] vm = {1'b1, wv[22:0]};
253  wire signed [11:0] shift = we - ($signed({4'd0, wv[30:23]}) - 12'sd150);
254  wire [23:0] kept = (shift > 12'sd24) ? 24'd0 : (vm >> shift);
255  wire [23:0] rem = (shift > 12'sd24) ? vm : (vm & ((24'd1 << shift) - 24'd1));
256  wire [24:0] halfw = (shift > 12'sd24) ? 25'h1000000 : (25'd1 << (shift - 1));
257  wire up = ({1'b0, rem} > halfw) || ({1'b0, rem} == halfw && kept[0]);
258  wire [8:0] code = (wv[30:0] == 31'd0) ? 9'd0 : (kept[8:0] + {8'd0, up});
259  // An element's block sum rounded once and added into its output.
260  // A block whose weights all vanish adds +0 to every element.
261  wire [31:0] block_value, added, zeroed;
262  logos_round88 rb(.v({{56{sums[el][31]}}, sums[el]}), .least(we), .f(block_value));
263  logos_fadd ao(.a(out[el]), .b(block_value), .y(added));
264  logos_fadd az(.a(out[el]), .b(32'd0), .y(zeroed));
265  // The output beat: sixteen elements.
266  wire [511:0] out_beat;
267  genvar l;
268  generate
269    for (l = 0; l < 16; l = l + 1) begin : ol
270      assign out_beat[32*l +: 32] = out[16*b + l];
271    end
272  endgenerate
273  integer k;
274  // The scores and their exponentials in block RAM, each read and written at the position at hand.
275  localparam PA = (P_MAX > 1) ? $clog2(P_MAX) : 1;
276  logos_ram_w32 #(.D(P_MAX), .A(PA)) score_ram(.clk(clk), .ra(j[PA-1:0]), .rd(score_j), .we(!rst && state == KEY && k_valid),
277    .wa(j[PA-1:0]), .wd(score), .wm(4'hf));
278  logos_ram_w32 #(.D(P_MAX), .A(PA)) ex_ram(.clk(clk), .ra(j[PA-1:0]), .rd(ex_j), .we(!rst && state == EXP),
279    .wa(j[PA-1:0]), .wd(e_j), .wm(4'hf));
280  always @(posedge clk) begin
281    o_valid <= 1'b0;
282    o_last <= 1'b0;
283    if (rst) state <= IDLE;
284    else case (state)
285      IDLE: if (start) begin
286        qc <= q_codes; qm <= q_m; qe <= q_e; sc <= scale; n <= positions; j <= 16'd0; m <= 32'hff80_0000;
287        for (k = 0; k < VD; k = k + 1) out[k] <= 32'd0;
288        state <= KEY;
289      end
290      KEY: if (k_valid) begin
291        if (!score_nan && m_lt) m <= score;
292        if (j + 16'd1 == n) begin j <= 16'd0; lanes <= 512'd0; state <= EXP; end else j <= j + 16'd1;
293      end
294      EXP: begin
295        lanes[32*j[3:0] +: 32] <= e_next;
296        if (j + 16'd1 == n) begin j <= 16'd0; state <= INV; end else j <= j + 16'd1;
297      end
298      INV: begin inv <= total_inv; blk <= 16'd0; wmax <= 32'd0; state <= WSCALE; end
299      WSCALE: if (s_valid) begin
300        w[j[4:0]] <= wj;
301        if (wj[30:0] > wmax[30:0]) wmax <= wj;
302        if (j[4:0] == 5'd31 || j + 16'd1 == n) begin
303          j <= {blk[10:0], 5'd0};
304          state <= WCODE;
305        end else j <= j + 16'd1;
306      end
307      WCODE: begin
308        we <= w_exp;
309        for (k = 0; k < VD; k = k + 1) sums[k] <= 32'sd0;
310        state <= WVAL;
311      end
312      WVAL: if (v_valid) begin
313        for (k = 0; k < VD; k = k + 1) sums[k] <= sums[k] + $signed({1'b0, code}) * $signed(v_codes[8*k +: 8]);
314        if (j[4:0] == 5'd31 || j + 16'd1 == n) begin el <= 16'd0; state <= WROUND; end else j <= j + 16'd1;
315      end
316      WROUND: begin
317        out[el] <= (wmax[30:0] == 31'd0) ? zeroed : added;
318        if (el + 16'd1 == ELEMS) begin
319          el <= 16'd0;
320          if ({blk[10:0], 5'd0} + 16'd32 >= n) begin b <= 16'd0; state <= OUT; end
321          else begin blk <= blk + 16'd1; j <= {blk[10:0], 5'd0} + 16'd32; wmax <= 32'd0; state <= WSCALE; end
322        end else el <= el + 16'd1;
323      end
324      OUT: begin
325        o_valid <= 1'b1;
326        o_beat <= out_beat;
327        if (b + 16'd1 == OB) begin o_last <= 1'b1; state <= IDLE; end else b <= b + 16'd1;
328      end
329    endcase
330  end
331endmodule
332"#;
333
334/// `logos_attend_int8_p`, `ATTEND_INT8_VERILOG`'s unit for the user clock, with the units it is built from.
335pub fn attend_int8_p_verilog() -> String {
336    use super::fp32::{FADD_LATENCY, FMUL_LATENCY, TREE16_LATENCY};
337    use super::vector::{EXP_LATENCY, RECIP_LATENCY};
338    let units = [
339        ("@FMUL@", FMUL_LATENCY),
340        ("@FADD@", FADD_LATENCY),
341        ("@TREE@", TREE16_LATENCY),
342        ("@EXP@", EXP_LATENCY),
343        ("@RECIP@", RECIP_LATENCY),
344        ("@ROUND@", super::lane::ROUND88_P_LATENCY),
345    ]
346    .iter()
347    .fold(ATTEND_INT8_P_VERILOG.to_string(), |t, (k, v)| t.replace(k, &v.to_string()));
348    format!("{}\n{}\n{}\n{units}", super::vector::vector_verilog(), super::lane::lane_verilog(), super::ram::bram_verilog(&[32]))
349}
350
351const ATTEND_INT8_P_VERILOG: &str = r#"
352// logos_attend_int8 for the user clock: one query head of HD elements over up to P_MAX positions of an int8 cache, a
353// line of 64 codes a clock. The query is one block (its codes, int8 or two-word, and its scale m·2^e). Each key row's
354// lines — the row's f16 scale with every line — make the row's exact dot with the query's codes, a line's 64 products
355// summed in a tree and the row's lines added; times the scale's significand and m, rounded once at their exponents,
356// times the score scale: the position's score, kept, and their maximum as the softmax takes it. Then each
357// exp(score − max), kept and summed in sixteen lanes by position, the lanes folded and the sum's reciprocal taken. Then
358// per block of 32 positions: each position's probability times its value row's f16 scale (a scale a clock), the block's
359// exponent from their maximum, each weight's 8-bit code to nearest even at it; the value rows' lines, each element's
360// exact integer sum against the codes; every element's sum rounded once at the exponent and added into its output, the
361// blocks in order (a block of no weight adds +0). The output VD/16 beats, busy until the last is out.
362module logos_attend_int8_p #(parameter HD = 256, VD = 256, P_MAX = 4096)
363  (input clk, input rst,
364   input start, input [15:0] positions, input [16*HD-1:0] q_codes, input [7:0] q_m, input signed [11:0] q_e,
365   input [31:0] scale, output busy,
366   input k_valid, input [511:0] k_line, input [15:0] k_scale, output k_ready,
367   input s_valid, input [15:0] v_scale, output s_ready,
368   input v_valid, input [511:0] v_line, output v_ready,
369   output reg o_valid, output reg [511:0] o_beat, output reg o_last);
370  localparam [15:0] KL = HD / 64, VL = VD / 64, OB = VD / 16, VN = VD;
371  localparam PA = (P_MAX > 1) ? $clog2(P_MAX) : 1;
372  localparam IDLE = 4'd0, KEY = 4'd1, KDRAIN = 4'd2, EXP = 4'd3, EDRAIN = 4'd4, INV = 4'd5, WSCALE = 4'd6, WCODE = 4'd7,
373             WVAL = 4'd8, WROUND = 4'd9, WDRAIN = 4'd10, OUT = 4'd11, WCODES = 4'd12, WVD = 4'd13;
374  // Clocks after an exponential's position is asked to its value (EX) and its lane's sum (EL); after a scale is taken
375  // to its weight (WL); after an element's sum is taken to its output's sum (RL).
376  localparam integer EX = 1 + @FADD@ + @EXP@, EL = EX + @FADD@, WL = 1 + 2 * @FMUL@, RL = @ROUND@ + @FADD@ + 1;
377  reg [3:0] state;
378  reg [15:0] n, j, li, blk, b, landed;
379  reg [16*HD-1:0] qc;
380  reg [7:0] qm;
381  reg signed [11:0] qe;
382  reg [31:0] sc;
383  assign busy = state != IDLE || o_valid;
384  assign k_ready = state == KEY;
385  assign s_ready = state == WSCALE && j != blk * 16'd32 + 16'd32 && j != n;
386  assign v_ready = state == WVAL;
387  // ── Keys: a line's products registered, the tree's two levels, the row's sum; the row's scale products; the
388  // rounding; the score scale.
389  wire k_take = k_valid && k_ready;
390  reg signed [23:0] kp [0:63];
391  reg kp_v, kp_first, kp_last;
392  reg [15:0] kp_scale;
393  integer i;
394  always @(posedge clk) begin
395    for (i = 0; i < 64; i = i + 1) kp[i] <= $signed(qc[16*(64*li + i) +: 16]) * $signed(k_line[8*i +: 8]);
396    kp_v <= !rst && k_take;
397    kp_first <= li == 16'd0;
398    kp_last <= li == KL - 1;
399    kp_scale <= k_scale;
400  end
401  // Eight products a partial and eight partials a line, each in a balanced tree of three adds.
402  function signed [29:0] tree8(input signed [23:0] a0, input signed [23:0] a1, input signed [23:0] a2,
403                               input signed [23:0] a3, input signed [23:0] a4, input signed [23:0] a5,
404                               input signed [23:0] a6, input signed [23:0] a7);
405    reg signed [24:0] b0, b1, b2, b3;
406    reg signed [25:0] c0, c1;
407    begin
408      b0 = a0 + a1; b1 = a2 + a3; b2 = a4 + a5; b3 = a6 + a7;
409      c0 = b0 + b1; c1 = b2 + b3;
410      tree8 = c0 + c1;
411    end
412  endfunction
413  reg signed [29:0] kt [0:7];
414  reg kt_v, kt_first, kt_last;
415  reg [15:0] kt_scale;
416  integer t;
417  always @(posedge clk) begin
418    for (t = 0; t < 8; t = t + 1)
419      kt[t] <= tree8(kp[8*t], kp[8*t+1], kp[8*t+2], kp[8*t+3], kp[8*t+4], kp[8*t+5], kp[8*t+6], kp[8*t+7]);
420    kt_v <= !rst && kp_v; kt_first <= kp_first; kt_last <= kp_last; kt_scale <= kp_scale;
421  end
422  reg signed [47:0] kdot;
423  reg kd_v;
424  reg [15:0] kd_scale;
425  wire signed [32:0] l0 = kt[0] + kt[1], l1 = kt[2] + kt[3], l2 = kt[4] + kt[5], l3 = kt[6] + kt[7];
426  wire signed [33:0] l4 = l0 + l1, l5 = l2 + l3;
427  wire signed [47:0] line_sum = l4 + l5;
428  always @(posedge clk) begin
429    if (kt_v) kdot <= (kt_first ? 48'sd0 : kdot) + line_sum;
430    kd_v <= !rst && kt_v && kt_last;
431    if (kt_v && kt_last) kd_scale <= kt_scale;
432  end
433  // The row's scale as a signed significand and an exponent, a denormal's the least; times m, times the dot.
434  wire signed [12:0] k_sig = (kd_scale[14:10] == 5'd0) ? (kd_scale[15] ? -$signed({3'd0, kd_scale[9:0]}) : $signed({3'd0, kd_scale[9:0]}))
435                                                      : (kd_scale[15] ? -$signed({2'd0, 1'b1, kd_scale[9:0]}) : $signed({2'd0, 1'b1, kd_scale[9:0]}));
436  wire signed [11:0] k_exp = (kd_scale[14:10] == 5'd0) ? -12'sd24 : $signed({7'd0, kd_scale[14:10]}) - 12'sd25;
437  reg signed [21:0] sm;
438  reg signed [47:0] kdot2;
439  reg signed [11:0] least2;
440  reg sm_v;
441  always @(posedge clk) begin
442    sm <= $signed(k_sig) * $signed({1'b0, qm});
443    kdot2 <= kdot;
444    least2 <= k_exp + qe;
445    sm_v <= !rst && kd_v;
446  end
447  reg signed [87:0] term;
448  reg signed [11:0] least3;
449  reg tm_v;
450  always @(posedge clk) begin
451    term <= $signed(sm) * $signed(kdot2);
452    least3 <= least2;
453    tm_v <= !rst && sm_v;
454  end
455  wire [31:0] rounded, score;
456  logos_round88_p rs(.clk(clk), .v(term), .least(least3), .f(rounded));
457  logos_fmul_p ms(.clk(clk), .a(rounded), .b(sc), .y(score));
458  reg [@ROUND@ + @FMUL@:1] sv_d;
459  always @(posedge clk) sv_d <= rst ? {(@ROUND@ + @FMUL@){1'b0}} : {sv_d[@ROUND@ + @FMUL@ - 1:1], tm_v};
460  wire score_in = sv_d[@ROUND@ + @FMUL@];
461  reg [15:0] sj;
462  // The scores' maximum: a NaN never larger, a largest zero +0.
463  reg [31:0] m;
464  wire score_nan = (score[30:23] == 8'hff) && (score[22:0] != 23'd0);
465  wire m_lt;
466  logos_flt cmp(.a(m), .b(score), .y(m_lt));
467  wire [31:0] mx = (m[30:0] == 31'd0) ? 32'd0 : m;
468  // ── Exponentials: a score read, less the maximum, through the exponential into its lane by position and kept.
469  wire [31:0] score_j, ex_j;
470  wire [31:0] diff, e_j;
471  reg [15:0] er;
472  logos_fadd_p ed(.clk(clk), .a(score_j), .b({~mx[31], mx[30:0]}), .y(diff));
473  logos_exp_p xe(.clk(clk), .x_in(diff), .y(e_j));
474  reg [EL:1] ev_d;
475  reg [15:0] ej_d [1:EL];
476  always @(posedge clk) begin
477    ev_d <= rst ? {EL{1'b0}} : {ev_d[EL-1:1], state == EXP && er != n};
478    ej_d[1] <= er;
479    for (i = 2; i <= EL; i = i + 1) ej_d[i] <= ej_d[i-1];
480  end
481  wire e_in = ev_d[EX];
482  wire [15:0] e_at = ej_d[EX];
483  reg [511:0] lanes;
484  wire [31:0] lane_sum;
485  logos_fadd_p el_add(.clk(clk), .a(lanes[32*e_at[3:0] +: 32]), .b(e_j), .y(lane_sum));
486  wire lane_in = ev_d[EL];
487  wire [15:0] lane_at = ej_d[EL];
488  wire [31:0] total, total_inv;
489  logos_tree16_p te(.clk(clk), .lanes(lanes), .sum(total));
490  logos_recip_p rr(.clk(clk), .x_in(total), .y(total_inv));
491  reg [15:0] wait_c;
492  reg [31:0] inv;
493  logos_bram_w32 #(.D(P_MAX), .A(PA)) score_ram(.clk(clk), .ra(er[PA-1:0]), .rd(score_j),
494    .we(score_in), .wa(sj[PA-1:0]), .wd(score), .wm(4'hf));
495  logos_bram_w32 #(.D(P_MAX), .A(PA)) ex_ram(.clk(clk), .ra(j[PA-1:0]), .rd(ex_j), .we(e_in), .wa(e_at[PA-1:0]), .wd(e_j),
496    .wm(4'hf));
497  // ── Weights: a position's probability times its value row's scale, into the block's weights.
498  wire [31:0] prob, vs_c, wj;
499  reg [31:0] vs;
500  reg [15:0] vs_q;
501  logos_fmul_p mp(.clk(clk), .a(ex_j), .b(inv), .y(prob));
502  reg [15:0] vs_d [1:@FMUL@];
503  always @(posedge clk) begin
504    vs_d[1] <= vs_q;
505    for (i = 2; i <= @FMUL@; i = i + 1) vs_d[i] <= vs_d[i-1];
506  end
507  // The row's scale widened from binary16 and registered, meeting its probability out of the multiply.
508  logos_f16_to_f32 cv(.h(vs_d[@FMUL@ - 1]), .x(vs_c));
509  always @(posedge clk) vs <= vs_c;
510  logos_fmul_p mw(.clk(clk), .a(prob), .b(vs), .y(wj));
511  reg [WL:1] wv_d;
512  reg [4:0] wj_d [1:WL];
513  always @(posedge clk) begin
514    wv_d <= rst ? {WL{1'b0}} : {wv_d[WL-1:1], s_valid && s_ready};
515    wj_d[1] <= j[4:0];
516    for (i = 2; i <= WL; i = i + 1) wj_d[i] <= wj_d[i-1];
517    vs_q <= v_scale;
518  end
519  wire w_in = wv_d[WL];
520  (* ramstyle = "logic" *) reg [31:0] w [0:31];
521  reg [31:0] wmax;
522  reg [5:0] w_got;
523  // The block's exponent: the least e with wmax <= 255·2^e.
524  wire [23:0] wm = {1'b1, wmax[22:0]};
525  wire signed [11:0] w_exp = $signed({4'd0, wmax[30:23]}) - 12'sd150 + 12'sd16 + ((wm > 24'd16711680) ? 12'sd1 : 12'sd0);
526  reg signed [11:0] we;
527  // ── Values: each position's code from its weight at the block's exponent, nearest even, for every position of the
528  // block once its exponent is set — four a clock, the weights registered, then each cut at the exponent (its kept
529  // bits, its remainder and the remainder's half), then rounded into its code (codes); then each value line's 64
530  // products with its position's code registered, and added into their sums the clock after, each sum by its own
531  // enable, so a block's last sum is in the clock after its last line.
532  function [58:0] code_cut(input [31:0] wv, input signed [11:0] e);
533    reg [23:0] vm, kept, rem;
534    reg [24:0] halfw;
535    reg signed [11:0] shift;
536    begin
537      vm = {1'b1, wv[22:0]};
538      shift = e - ($signed({4'd0, wv[30:23]}) - 12'sd150);
539      kept = (shift > 12'sd24) ? 24'd0 : (vm >> shift);
540      rem = (shift > 12'sd24) ? vm : (vm & ((24'd1 << shift) - 24'd1));
541      halfw = (shift > 12'sd24) ? 25'h1000000 : (25'd1 << (shift - 1));
542      code_cut = {wv[30:0] == 31'd0, kept[8:0], rem, halfw};
543    end
544  endfunction
545  function [8:0] code_round(input [58:0] c);
546    reg zero, up;
547    reg [8:0] kept;
548    reg [23:0] rem;
549    reg [24:0] halfw;
550    begin
551      {zero, kept, rem, halfw} = c;
552      up = ({1'b0, rem} > halfw) || ({1'b0, rem} == halfw && kept[0]);
553      code_round = zero ? 9'd0 : (kept + {8'd0, up});
554    end
555  endfunction
556  reg [8:0] codes [0:31];
557  reg [31:0] cw [0:3];
558  reg [58:0] cut [0:3];
559  reg [3:0] cj;
560  reg cw_v, cut_v;
561  reg [2:0] cw_at, cut_at;
562  reg signed [17:0] vp [0:63];
563  reg vp_v;
564  reg [15:0] vp_li;
565  wire v_take = v_valid && v_ready;
566  integer e;
567  always @(posedge clk) begin
568    for (e = 0; e < 4; e = e + 1) cut[e] <= code_cut(cw[e], we);
569    cut_v <= !rst && cw_v;
570    cut_at <= cw_at;
571    for (e = 0; e < 32; e = e + 1) if (cut_v && cut_at == e[4:2]) codes[e] <= code_round(cut[e % 4]);
572    vp_v <= !rst && v_take;
573    vp_li <= li;
574    for (e = 0; e < 64; e = e + 1) vp[e] <= $signed({1'b0, codes[j[4:0]]}) * $signed(v_line[8*e +: 8]);
575  end
576  reg signed [31:0] sums [0:VD-1];
577  // ── Rounding: an element's sum rounded once at the block's exponent, added into its output.
578  wire [31:0] block_value, added;
579  reg signed [31:0] rsum;
580  reg [15:0] rel;
581  reg r_v, r_zero;
582  logos_round88_p rb(.clk(clk), .v({{56{rsum[31]}}, rsum}), .least(we), .f(block_value));
583  reg [31:0] out [0:VD-1];
584  reg [@ROUND@:1] rv_d, rz_d;
585  reg [15:0] rel_d [1:RL];
586  always @(posedge clk) begin
587    rv_d <= rst ? {@ROUND@{1'b0}} : {rv_d[@ROUND@-1:1], r_v};
588    rz_d <= {rz_d[@ROUND@-1:1], r_zero};
589    rel_d[1] <= rel;
590    for (i = 2; i <= RL; i = i + 1) rel_d[i] <= rel_d[i-1];
591  end
592  reg [31:0] add_a;
593  logos_fadd_p ao(.clk(clk), .a(add_a), .b(rz_d[@ROUND@] ? 32'd0 : block_value), .y(added));
594  reg [@FADD@:1] av_d;
595  always @(posedge clk) begin
596    add_a <= out[rel_d[@ROUND@]];
597    av_d <= rst ? {@FADD@{1'b0}} : {av_d[@FADD@-1:1], rv_d[@ROUND@]};
598  end
599  wire out_in = av_d[@FADD@];
600  wire [15:0] out_at = rel_d[RL];
601  // The output beat: sixteen elements.
602  wire [511:0] out_beat;
603  genvar l;
604  generate
605    for (l = 0; l < 16; l = l + 1) begin : ol
606      assign out_beat[32*l +: 32] = out[16*b + l];
607    end
608  endgenerate
609  integer k;
610  always @(posedge clk) begin
611    o_valid <= 1'b0;
612    o_last <= 1'b0;
613    r_v <= 1'b0;
614    if (score_in) begin
615      sj <= sj + 16'd1;
616      if (!score_nan && m_lt) m <= score;
617    end
618    if (lane_in) lanes[32*lane_at[3:0] +: 32] <= lane_sum;
619    if (w_in) begin
620      w[wj_d[WL]] <= wj;
621      if (wj[30:0] > wmax[30:0]) wmax <= wj;
622      w_got <= w_got + 6'd1;
623    end
624    if (out_in) begin out[out_at] <= added; landed <= landed + 16'd1; end
625    if (rst) state <= IDLE;
626    else case (state)
627      IDLE: if (start) begin
628        qc <= q_codes; qm <= q_m; qe <= q_e; sc <= scale; n <= positions; j <= 16'd0; li <= 16'd0; sj <= 16'd0;
629        m <= 32'hff80_0000;
630        for (k = 0; k < VD; k = k + 1) out[k] <= 32'd0;
631        state <= KEY;
632      end
633      KEY: if (k_take) begin
634        if (li == KL - 1) begin
635          li <= 16'd0;
636          if (j + 16'd1 == n) begin j <= 16'd0; state <= KDRAIN; end else j <= j + 16'd1;
637        end else li <= li + 16'd1;
638      end
639      // Every score in; then each one's exponential, a position a clock.
640      KDRAIN: if (sj == n) begin er <= 16'd0; lanes <= 512'd0; state <= EXP; end
641      EXP: begin
642        if (er != n) er <= er + 16'd1;
643        else begin wait_c <= 16'd0; state <= EDRAIN; end
644      end
645      EDRAIN: begin
646        wait_c <= wait_c + 16'd1;
647        if (wait_c == EL + @TREE@ + @RECIP@ + 16'd2) state <= INV;
648      end
649      INV: begin inv <= total_inv; blk <= 16'd0; j <= 16'd0; wmax <= 32'd0; w_got <= 6'd0; state <= WSCALE; end
650      // A block's scales a clock, their weights in; then its exponent; then its value rows.
651      WSCALE: begin
652        if (s_valid && s_ready) j <= j + 16'd1;
653        if ((j == blk * 16'd32 + 16'd32 || j == n) && {10'd0, w_got} == j - blk * 16'd32) begin
654          j <= blk * 16'd32; state <= WCODE;
655        end
656      end
657      WCODE: begin
658        we <= w_exp;
659        for (k = 0; k < VD; k = k + 1) sums[k] <= 32'sd0;
660        li <= 16'd0;
661        cj <= 4'd0;
662        state <= WCODES;
663      end
664      WCODES: begin
665        if (cj != 4'd8) begin
666          for (k = 0; k < 4; k = k + 1) cw[k] <= w[4*cj[2:0] + k];
667          cw_v <= 1'b1; cw_at <= cj[2:0]; cj <= cj + 4'd1;
668        end else begin
669          cw_v <= 1'b0;
670          if (!cw_v && !cut_v) state <= WVAL;
671        end
672      end
673      WVAL: if (v_take) begin
674        if (li == VL - 1) begin
675          li <= 16'd0;
676          if (j[4:0] == 5'd31 || j + 16'd1 == n) begin rel <= 16'd0; landed <= 16'd0; state <= WVD; end
677          else j <= j + 16'd1;
678        end else li <= li + 16'd1;
679      end
680      WVD: state <= WROUND;
681      // Each element's sum through the rounding and the add, an element a clock.
682      WROUND: begin
683        r_v <= 1'b1; r_zero <= wmax[30:0] == 31'd0; rsum <= sums[rel];
684        rel <= rel + 16'd1;
685        if (rel + 16'd1 == VN) state <= WDRAIN;
686      end
687      WDRAIN: if (landed == VN) begin
688        if (blk * 16'd32 + 16'd32 >= n) begin b <= 16'd0; state <= OUT; end
689        else begin blk <= blk + 16'd1; j <= blk * 16'd32 + 16'd32; wmax <= 32'd0; w_got <= 6'd0; state <= WSCALE; end
690      end
691      OUT: begin
692        o_valid <= 1'b1;
693        o_beat <= out_beat;
694        if (b + 16'd1 == OB) begin o_last <= 1'b1; state <= IDLE; end else b <= b + 16'd1;
695      end
696      default: state <= IDLE;
697    endcase
698    for (k = 0; k < VD; k = k + 1) if (vp_v && vp_li == k[21:6]) sums[k] <= sums[k] + vp[k % 64];
699    if (rst) cw_v <= 1'b0;
700  end
701endmodule
702"#;
703
704/// `logos_attend_side`, a token-row's attention for a hybrid engine's fast side ([`super::engine::Streaming::attention`]),
705/// with the user clock's unit and quantizer it drives.
706pub fn attend_side_verilog() -> String {
707    format!("{}\n{}\n{}", attend_int8_p_verilog(), super::streaming::act_stream_p_verilog(), ATTEND_SIDE_VERILOG)
708}
709
710const ATTEND_SIDE_VERILOG: &str = r#"
711// A token-row's attention on a hybrid engine's fast side over an int8 cache: for each query head h, its cached head
712// g = h / (heads / kv_heads); the head's query from the store (HD/16 beats from q_beat + h·HD/16) quantized as one block
713// (logos_act_stream_p); then logos_attend_int8_p fed from the row's region through the port, the reads asked in the
714// order the unit takes them — per position its key row's scale line and its key lines, then per block of 32 positions
715// each value row's scale line, then each one's value lines — every line back in order, its tag naming what it is; the
716// head's output beats into the store from out_beat + h·VD/16. Positions are the row's own and every one before. A cache
717// row's first byte and its scale's are the region's (a line) plus the block's offsets plus (j·kv_heads + g) rows.
718module logos_attend_side #(parameter HD = 256, VD = 256, P_MAX = 4096, SA = 12, parameter [6:0] OUTSTANDING = 32, SRL = 1)
719  (input clk, input rst, input start, output busy,
720   input [31:0] region, input [31:0] keys, input [31:0] values, input [31:0] kscales, input [31:0] vscales,
721   input [31:0] kv_heads, input [15:0] heads, input [31:0] positions, input [31:0] scale,
722   input [SA-1:0] q_beat, input [SA-1:0] out_beat,
723   output [SA-1:0] st_ra, input [511:0] st_rd, output reg st_we, output reg [SA-1:0] st_wa, output reg [511:0] st_wd,
724   output rq_valid, input rq_ready, output [31:0] rq_line, input rs_valid, input [511:0] rs_data);
725  localparam [15:0] HB = HD / 16, OB = VD / 16, KL = HD / 64, VL = VD / 64;
726  localparam IDLE = 3'd0, GROUP = 3'd1, QLOAD = 3'd2, QUANT = 3'd3, RUN = 3'd4;
727  localparam [1:0] T_KS = 2'd0, T_K = 2'd1, T_VS = 2'd2, T_V = 2'd3;
728  reg [2:0] state;
729  // The job, the group size, the head and its cached head.
730  reg [31:0] rg, ko, vo, kso, vso, kv, n, sc;
731  reg [15:0] nh, grp, rem, h, hg;
732  reg [31:0] g;
733  reg [SA-1:0] qb, ob;
734  assign busy = state != IDLE || st_we;
735  // The query's beats into the quantizer, a beat the clock after its address.
736  reg [15:0] qj;
737  reg arriving, first;
738  // A query beat asked (arriving, first, the clock after its address) and out of the store SRL clocks after its address
739  // (landed, landed_first).
740  wire landed, landed_first, landing;
741  logos_landing #(.W(1), .D(SRL - 1)) qland(.clk(clk), .rst(rst), .in_v(arriving), .in_d(first), .out_v(landed),
742    .out_d(landed_first), .flight(landing));
743  wire aq_busy;
744  wire [7:0] aq_m;
745  wire signed [11:0] aq_e;
746  wire [16*HD-1:0] aq_codes;
747  reg aq_go;
748  wire [31:0] q_at = {{(32-SA){1'b0}}, qb} + h * HB + qj;
749  assign st_ra = q_at[SA-1:0];
750  logos_act_stream_p #(.LANES(HD)) aq(.clk(clk), .rst(rst), .clear(landed && landed_first), .in_v(landed), .in_beat(st_rd),
751    .go(aq_go), .two(1'b0), .busy(aq_busy), .m(aq_m), .e(aq_e), .codes(aq_codes));
752  // The unit.
753  reg u_start;
754  wire u_busy, k_ready, s_ready, v_ready, o_valid, o_last;
755  wire [511:0] o_beat;
756  reg k_valid, s_valid, v_valid;
757  reg [511:0] k_line, v_line;
758  reg [15:0] k_scale, v_scale;
759  logos_attend_int8_p #(.HD(HD), .VD(VD), .P_MAX(P_MAX)) u(.clk(clk), .rst(rst), .start(u_start), .positions(n[15:0]),
760    .q_codes(aq_codes), .q_m(aq_m), .q_e(aq_e), .scale(sc), .busy(u_busy), .k_valid(k_valid), .k_line(k_line),
761    .k_scale(k_scale), .k_ready(k_ready), .s_valid(s_valid), .v_scale(v_scale), .s_ready(s_ready), .v_valid(v_valid),
762    .v_line(v_line), .v_ready(v_ready), .o_valid(o_valid), .o_beat(o_beat), .o_last(o_last));
763  // The walk of the head's reads: its phase (keys, a block's scales, a block's values), position, line, and the cache row
764  // of its position (wj·kv_heads + g, kept as the walk steps) and of its block's first.
765  localparam [1:0] W_KEYS = 2'd0, W_SCALES = 2'd1, W_VALUES = 2'd2, W_DONE = 2'd3;
766  reg [1:0] wph;
767  reg [31:0] wj, wblk, w_row, b_row;
768  reg [15:0] wl;
769  reg walking;
770  wire [47:0] base = {16'd0, rg} * 48'd64;
771  wire [47:0] ks_byte = base + kso + {w_row, 1'b0};
772  wire [47:0] k_byte = base + ko + w_row * HD + {wl, 6'd0};
773  wire [47:0] vs_byte = base + vso + {w_row, 1'b0};
774  wire [47:0] v_byte = base + vo + w_row * VD + {wl, 6'd0};
775  // The keys' phase asks each position's scale line (wl == KL) then its lines; the values' each position's lines.
776  wire [1:0] w_kind = wph == W_KEYS ? (wl == KL ? T_KS : T_K) : wph == W_SCALES ? T_VS : T_V;
777  wire [47:0] w_byte = w_kind == T_KS ? ks_byte : w_kind == T_K ? k_byte : w_kind == T_VS ? vs_byte : v_byte;
778  reg [6:0] owed;
779  assign rq_valid = walking && owed < OUTSTANDING;
780  assign rq_line = w_byte[37:6];
781  wire rq_take = rq_valid && rq_ready;
782  wire [31:0] blk_end = ({wblk[26:0], 5'd0} + 32'd32 < n) ? {wblk[26:0], 5'd0} + 32'd32 : n;
783  // Every read's tag — its kind and its scale's byte — and its line, in order.
784  reg [7:0] tags [0:63];
785  reg [511:0] lines [0:63];
786  reg [5:0] t_wr, t_rd, l_wr;
787  always @(posedge clk) begin
788    if (rq_take) tags[t_wr] <= {w_kind, w_byte[5:0]};
789    if (rs_valid) lines[l_wr] <= rs_data;
790  end
791  wire have = l_wr != t_rd;
792  wire [1:0] t_kind = tags[t_rd][7:6];
793  wire [5:0] t_off = tags[t_rd][5:0];
794  wire [511:0] line = lines[t_rd];
795  wire [15:0] line_scale = line[8 * t_off +: 16];
796  // The line at the head handed to the unit as it asks: a key scale held for its row's lines.
797  reg [15:0] kscale_h;
798  wire to_k = have && t_kind == T_K, to_s = have && t_kind == T_VS, to_v = have && t_kind == T_V;
799  always @* begin
800    k_valid = to_k; k_line = line; k_scale = kscale_h;
801    s_valid = to_s; v_scale = line_scale;
802    v_valid = to_v; v_line = line;
803  end
804  wire pop = have && (t_kind == T_KS || (to_k && k_ready) || (to_s && s_ready) || (to_v && v_ready));
805  reg [15:0] ob_i;
806  wire [31:0] o_at = {{(32-SA){1'b0}}, ob} + h * OB + ob_i;
807  always @(posedge clk) begin
808    u_start <= 1'b0; aq_go <= 1'b0; arriving <= 1'b0; st_we <= 1'b0;
809    if (rq_take) t_wr <= t_wr + 6'd1;
810    if (rs_valid) l_wr <= l_wr + 6'd1;
811    if (pop) begin t_rd <= t_rd + 6'd1; if (t_kind == T_KS) kscale_h <= line_scale; end
812    owed <= owed + (rq_take ? 7'd1 : 7'd0) - (pop ? 7'd1 : 7'd0);
813    // The walk: keys, then each block's scales and values.
814    if (rq_take) begin
815      case (wph)
816        W_KEYS: if (wl == KL) wl <= 16'd0;
817                else if (wl + 16'd1 == KL) begin
818                  wl <= KL;
819                  if (wj + 32'd1 == n) begin wj <= 32'd0; w_row <= g; b_row <= g; wblk <= 32'd0; wph <= W_SCALES; end
820                  else begin wj <= wj + 32'd1; w_row <= w_row + kv; end
821                end else wl <= wl + 16'd1;
822        W_SCALES: if (wj + 32'd1 == blk_end) begin wj <= {wblk[26:0], 5'd0}; w_row <= b_row; wl <= 16'd0; wph <= W_VALUES; end
823                  else begin wj <= wj + 32'd1; w_row <= w_row + kv; end
824        W_VALUES: if (wl + 16'd1 == VL) begin
825                    wl <= 16'd0;
826                    if (wj + 32'd1 == blk_end) begin
827                      if (blk_end == n) begin walking <= 1'b0; wph <= W_DONE; end
828                      else begin wblk <= wblk + 32'd1; wj <= blk_end; w_row <= w_row + kv; b_row <= w_row + kv; wph <= W_SCALES; end
829                    end else begin wj <= wj + 32'd1; w_row <= w_row + kv; end
830                  end else wl <= wl + 16'd1;
831        default: ;
832      endcase
833    end
834    // The head's output beats into the store.
835    if (o_valid) begin st_we <= 1'b1; st_wa <= o_at[SA-1:0]; st_wd <= o_beat; ob_i <= ob_i + 16'd1; end
836    if (rst) begin
837      state <= IDLE; t_wr <= 6'd0; t_rd <= 6'd0; l_wr <= 6'd0; owed <= 7'd0; walking <= 1'b0;
838    end else case (state)
839      IDLE: if (start) begin
840        rg <= region; ko <= keys; vo <= values; kso <= kscales; vso <= vscales; kv <= kv_heads; nh <= heads;
841        n <= positions; sc <= scale; qb <= q_beat; ob <= out_beat; grp <= 16'd0; rem <= heads; state <= GROUP;
842      end
843      // heads / kv_heads by subtraction.
844      GROUP: if (rem >= kv[15:0]) begin rem <= rem - kv[15:0]; grp <= grp + 16'd1; end
845             else begin h <= 16'd0; hg <= 16'd0; g <= 32'd0; qj <= 16'd0; state <= QLOAD; end
846      QLOAD: begin
847        arriving <= 1'b1; first <= qj == 16'd0;
848        if (qj + 16'd1 == HB) begin qj <= 16'd0; state <= QUANT; end else qj <= qj + 16'd1;
849      end
850      QUANT: if (!arriving && !landing && !aq_go && !aq_busy) begin
851        if (qj == 16'd0) begin aq_go <= 1'b1; qj <= 16'd1; end
852        else begin
853          qj <= 16'd0; u_start <= 1'b1; ob_i <= 16'd0;
854          wph <= W_KEYS; wj <= 32'd0; w_row <= g; wl <= KL; walking <= 1'b1; state <= RUN;
855        end
856      end
857      RUN: if (o_valid && o_last) begin
858        if (h + 16'd1 == nh) state <= IDLE;
859        else begin
860          h <= h + 16'd1;
861          if (hg + 16'd1 == grp) begin hg <= 16'd0; g <= g + 32'd1; end else hg <= hg + 16'd1;
862          qj <= 16'd0; state <= QLOAD;
863        end
864      end
865      default: state <= IDLE;
866    endcase
867  end
868endmodule
869"#;