/* SPDX-License-Identifier: AGPL-3.0-or-later */ /* Copyright (C) 2026 David (Corey) Herrera */ /** * @file ft_tt_op_attention.h * @brief Host driver for FLASH-DECODE ATTENTION: one GQA group's query against the KV * cache, with an ONLINE softmax so the score row is never materialised. Places the * triad's circular buffers, encodes its parameter block, and builds the constants — the * reduce scaler, the 1/sqrt(head_dim) scale and the 32 tail masks. * * The device half is `kernels/resident_attention_protocol.h`; the two must agree word for * word and `tests/test_attention_ops.c` pins the pair off-hardware. * * --------------------------------------------------------------------------- * THE UNIT OF WORK IS A KV GROUP, NOT A Q HEAD, AND THAT IS THE DESIGN DECISION * --------------------------------------------------------------------------- * MiniCPM5-1B is GQA — **16 Q heads, 2 KV heads, head_dim 128, 24 layers, context * 131072**, read from the GGUF header (docs/tt_phase0_decisions.md §3) and not assumed. So * each KV head serves EIGHT Q heads, and the choice of what one core owns is the choice of * how many times the KV cache is read: * * - **a core owns a Q HEAD** — the cache is read 8x per layer per token, because each of * the 8 Q heads sharing a KV head walks it separately; * - **a core owns the KV GROUP** — the cache is read ONCE and serves all 8. * * Decode is bandwidth-shaped: at 8k context one layer's K+V for one KV head is 8 MiB, so * the first design pays 64 MiB per layer per token against the second's 8. The group is * the unit. The 8 Q heads become the 8 ROWS of one tile, so the score matmul is issued once * rather than eight times — which also keeps the op data-movement bound rather than MATH * bound (8 HiFi4 tile matmuls per K tile pair would be 640 ns against a reader delivering a * pair every ~104 ns, docs/tt_kernel_model.md §12, §16). * * --------------------------------------------------------------------------- * THE ONLINE SOFTMAX IS A CAPACITY REQUIREMENT, NOT A PREFERENCE * --------------------------------------------------------------------------- * A standalone softmax holds its row twice over and tops out at Wt = 325, i.e. W ≤ 10,400 * (`ft_tt_op_reduce_max_tile_cols()`, docs/tt_bringup_log.md §22.3). A 32k-context * attention row is **Wt = 1024 and does not fit at any tile size**. So this op accumulates * a running max, row sum and output across KV blocks and never holds more than one block: * `ft_tt_op_attention_l1_bytes()` is flat in the context length and depends only on the * residency. That is the whole reason attention is its own op rather than "matmul, then * softmax, then matmul". * * --------------------------------------------------------------------------- * WHAT IT CONSUMES, AND WHAT IT IS ALLOWED TO ASSUME * --------------------------------------------------------------------------- * K and V are `ft_tt_op_memory.h` PART 1's contract, unchanged: F1 (K and V share one * orientation, position on the tile's H axis), F2 (axis order HEAD → POSITION-TILE → * DIM-TILE, so a position range of one head is ONE contiguous run), F3 (face-ordered * tiles), F4 (preallocated, growth never relayouts). `ft_tt_kv_cache_run()` gives the * (offset, bytes) pair this driver folds into the parameter block. * * §1.3's ONE inferred assumption — that `matmul_init(in0, in1, transpose)` transposes B on * Blackhole — is **MEASURED and true** as of this op: `tests/test_attention_transpose_probe.c` * runs both arms of one binary against a non-symmetric operand pair and the transposed arm * reproduces A·Bᵀ (PCC 0.999998668) while the plain arm reproduces A·B. The fallback §1.3 * describes (transpose Q instead) is therefore not needed and not implemented. * * Q and the output are the model's ordinary activation layout: a tile grid * `q_head_count * dim_tiles` wide, head h's dim tile d at tile column `h*dim_tiles + d`, * the decode token at row `ft_tt_op_memory_position_row(position)` — §1.5b's row-phase * invariant, inherited rather than re-decided. * * @status GREEN ON SILICON — Tensix 70, 71, 72 and 73, 0 checks failed on every one, PCC * 0.99998-1.0 and max-abs <= 0.002 against a CPU fp32 reference at kv = 1, 33, 256, 1000 * and 2048. Core 71 reproduced core 70's PCCs to NINE DIGITS on five shapes. Four kernel * mutations were built into their own program directories and driven through the SAME * unmodified test binary: all four go red (6, 6, 10 and 2 failures). Timing at * block_tiles = 16: **0.133 / 0.236 / 0.842 ms per descriptor at kv = 1024 / 2048 / 8192**, * i.e. 40.4 ms and 24.7 tok/s per token at 8k on ONE core (24 layers x 2 KV groups), with * the two groups independent so two cores halve it. * @issues The op is COMPUTE bound, not reader bound: at kv = 8192 it moves 4 MiB of KV in * 0.842 ms = 5.0 GB/s, against the matmul reader's measured ~28 GB/s at the same tile * size. INFERRED, not measured: the SFPU exponential and the two row reduces, which the * matmul does not have. That makes LoFi and an approximate exp real levers here for the * first time in this backend — the opposite of docs/tt_bringup_log.md §16's finding on * the matmul, and it should be measured before it is believed. * @todo None */ #ifndef FT_TT_OP_ATTENTION_H #define FT_TT_OP_ATTENTION_H #include #include #include #include "ft_tt_launch.h" #include "ft_tt_op_dispatch.h" #include "ft_tt_queue.h" /* ================================================================ * The parameter block — mirrors kernels/resident_attention_protocol.h * ================================================================ */ /** The sixteen words a FT_TT_QUEUE_OP_ATTENTION descriptor points at. * * Every operand's NoC coordinate is ONE word, `(y << 16) | x`. The other families spend * two words on a coordinate; this op has five operands and sixteen words, so it packs — * symmetrically on both sides of the wire, and pinned by the off-hardware suite. */ typedef enum ft_tt_op_attention_field { /** byte 0 q_rows, byte 1 query_row, byte 2 flags. */ FT_TT_OP_ATTENTION_HEADS = 0, FT_TT_OP_ATTENTION_DIM_TILES = 1, /** LIVE positions attended to. The live tile count and the tail width are DERIVED from * it on both sides — two fields that must agree are two fields that can disagree. */ FT_TT_OP_ATTENTION_KV_POSITIONS = 2, /** Position tiles per K/V block. Fixed by the RESIDENCY: it fixes the page counts, and * the kernel refuses a descriptor that disagrees with the buffers it was staged with. */ FT_TT_OP_ATTENTION_BLOCK_TILES = 3, /** Position tiles the cache region holds, so the kernel can refuse a padded read that * would run past it. */ FT_TT_OP_ATTENTION_CAPACITY_TILES = 4, FT_TT_OP_ATTENTION_RESERVED = 5, FT_TT_OP_ATTENTION_Q_COORD = 6, FT_TT_OP_ATTENTION_Q_BASE = 7, FT_TT_OP_ATTENTION_K_COORD = 8, FT_TT_OP_ATTENTION_K_BASE = 9, FT_TT_OP_ATTENTION_V_COORD = 10, FT_TT_OP_ATTENTION_V_BASE = 11, FT_TT_OP_ATTENTION_CONST_COORD = 12, FT_TT_OP_ATTENTION_CONST_BASE = 13, FT_TT_OP_ATTENTION_OUT_COORD = 14, FT_TT_OP_ATTENTION_OUT_BASE = 15, FT_TT_OP_ATTENTION_WORDS = 16 } ft_tt_op_attention_field; /** Tile layout of the constants operand, built once per model. * * THE CAUSAL BLOCK IS APPENDED. Tiles 0..33 are byte for byte what the decode arm already * uploads, so adding chunk mode does not move a single constant decode reads. */ #define FT_TT_OP_ATTENTION_CONST_SCALER_TILE 0u #define FT_TT_OP_ATTENTION_CONST_SCALE_TILE 1u #define FT_TT_OP_ATTENTION_CONST_MASK_TILE 2u /** Causal mask `+ (C - 1)`: element (row r, column c) live iff `c <= r && c < C`. */ #define FT_TT_OP_ATTENTION_CONST_CAUSAL_TILE (FT_TT_OP_ATTENTION_CONST_MASK_TILE + 32u) #define FT_TT_OP_ATTENTION_CONST_TILES (FT_TT_OP_ATTENTION_CONST_CAUSAL_TILE + 32u) /** `flags` byte 2 of word 0. CHUNK: one descriptor is ONE Q head over `q_rows` consecutive * positions instead of one KV group at one position — kernels/resident_attention_protocol.h * carries the design and the exhaustive list of what it changes. */ #define FT_TT_OP_ATTENTION_FLAG_CHUNK 1u /** S2f ROWS mode — mirrors kAttnFlagRows (resident_attention_protocol.h). */ #define FT_TT_OP_ATTENTION_FLAG_ROWS 2u /** How the reader stacks the four constants it fetches per descriptor. */ #define FT_TT_OP_ATTENTION_CONST_PAGES 4u /** The value a masked-out score is driven to before the row max sees it. Large enough that * `exp(dead - live_max)` is zero in bf16 at any plausible score, small enough that the * subtraction cannot overflow. */ #define FT_TT_OP_ATTENTION_MASK_VALUE (-1.0e30f) /* CB indices, matching the compile-time args the triad is built with. */ #define FT_TT_OP_ATTENTION_CB_Q 0u #define FT_TT_OP_ATTENTION_CB_K 1u #define FT_TT_OP_ATTENTION_CB_V 2u #define FT_TT_OP_ATTENTION_CB_CONST 3u #define FT_TT_OP_ATTENTION_CB_S 4u #define FT_TT_OP_ATTENTION_CB_SCRATCH 5u #define FT_TT_OP_ATTENTION_CB_M_A 6u #define FT_TT_OP_ATTENTION_CB_M_B 7u #define FT_TT_OP_ATTENTION_CB_L 8u #define FT_TT_OP_ATTENTION_CB_LSUM 9u #define FT_TT_OP_ATTENTION_CB_C 10u #define FT_TT_OP_ATTENTION_CB_O 11u #define FT_TT_OP_ATTENTION_CB_PV 12u #define FT_TT_OP_ATTENTION_CB_QSTAGE 13u #define FT_TT_OP_ATTENTION_CB_QS 14u #define FT_TT_OP_ATTENTION_CB_OUT 16u #define FT_TT_OP_ATTENTION_CB_COUNT 16u /* ================================================================ * The residency — what ONE staged program can serve * ================================================================ */ /** * The circular buffers are placed once, at launch, so the RESIDENCY and not the descriptor * fixes the group width, the head dimension and the KV block size. A descriptor that * disagrees is REFUSED by the kernel rather than run against buffers sized for something * else — the page counts are what a block-prefetch reader's correctness rests on * (kernels/resident_attention_protocol.h §5). */ typedef struct ft_tt_op_attention_residency { uint32_t q_rows; /**< Q heads per KV group; 8 for MiniCPM5-1B. 1..32 */ uint32_t dim_tiles; /**< head_dim / 32; 4 for MiniCPM5-1B */ uint32_t block_tiles; /**< position tiles per K/V block — the prefetch depth */ uint32_t tile_bytes; /**< 2048 Float16_b; the KV cache is bf16 (§1.6) */ } ft_tt_op_attention_residency; /** L1 the residency's circular buffers occupy. FLAT IN THE CONTEXT LENGTH — that is the * online softmax's whole payoff, and the number that makes 32k possible at all. */ size_t ft_tt_op_attention_l1_bytes(const ft_tt_op_attention_residency *residency); /** * The largest `block_tiles` that fits the CB arena for this group shape. * * @return 0 if even one position tile per block does not fit * * Thread-safe: yes Reentrant: yes Signal-safe: yes */ uint32_t ft_tt_op_attention_max_block_tiles(uint32_t q_rows, uint32_t dim_tiles, uint32_t tile_bytes); /** * Place the triad's sixteen circular buffers in the arena. * * K and V get TWO blocks of `block_tiles * dim_tiles` pages, so the reader fills one while * the compute thread drains the other — and two, rather than any other multiple, is what * makes `pages % depth == 0` hold (§5). * * @return FT_TT_OK; FT_TT_EINVAL; FT_TT_ENOMEM if it does not fit the arena */ int ft_tt_op_attention_place_cbs(const ft_tt_op_attention_residency *residency, ft_tt_launch_cb_entry out_entries[FT_TT_OP_ATTENTION_CB_COUNT], void *buffer, size_t buffer_bytes, uint64_t *out_mask); /* ================================================================ * One call * ================================================================ */ /** One attention call: the group's query at `position`, against `kv_positions` cached * positions of ONE KV head. Call it once per KV group per layer per token. */ typedef struct ft_tt_op_attention_request { /** DECODE: Q heads in this group. CHUNK: live positions C in this chunk, 1..32. */ uint32_t q_rows; uint32_t dim_tiles; /**< head_dim / 32 */ uint32_t block_tiles; /**< must equal the residency's */ uint32_t kv_positions; /**< LIVE positions; the query attends to all of them */ /** Position tiles the K (and V) region holds — `ft_tt_kv_cache_layout.capacity_tiles`. * The padded read window may exceed the live length by up to `2*block_tiles - 1` * tiles, and this is what stops that becoming a read past the region. */ uint32_t capacity_tiles; /** Where this token's query row sits: `ft_tt_op_memory_position_row(position)`, and it * is not optional — RoPE and the embedding put it there (§1.5b). */ uint32_t query_row; /** FT_TT_OP_ATTENTION_FLAG_CHUNK, or 0 for an ordinary decode descriptor. */ uint32_t flags; /** ROWS mode: strip rows served (query_row .. +positions), and the byte address of the * staircase mask tile on the constants' coordinate. */ uint32_t positions; uint32_t rows_mask_address; /** Q, already rotated: the FIRST tile of this group, i.e. head `group * q_rows`, dim * tile 0, of the row block the token sits in. `address` carries the head offset. * CHUNK: the first tile of ONE Q head, i.e. head `h`, dim tile 0. */ ft_tt_op_operand query; /** This KV head's K region at position 0 — base + `ft_tt_kv_cache_run()`'s offset. */ ft_tt_op_operand keys; /** This KV head's V region, same geometry (contract F1). */ ft_tt_op_operand values; /** ft_tt_op_attention_build_constants()' buffer, built once per model. */ ft_tt_op_operand constants; /** The attention output, in the same layout as `query`. */ ft_tt_op_operand output; } ft_tt_op_attention_request; /** Position tiles the reader actually fetches: the live count rounded up to a whole number * of DOUBLE blocks. The rounding is what keeps every descriptor's circular-buffer push and * pop counts identical, which is the block-prefetch invariant's second condition. */ uint32_t ft_tt_op_attention_padded_tiles(uint32_t kv_positions, uint32_t block_tiles); /** Position tiles that hold at least one live position: ceil(kv_positions / 32). */ uint32_t ft_tt_op_attention_live_tiles(uint32_t kv_positions); /** * Encode `request` into the sixteen words a FT_TT_QUEUE_OP_ATTENTION descriptor points at. * * Pure: no device access, no allocation. That is what makes the wire format testable * off-hardware, which matters more here than usual — the card is the expensive place to * find a field in the wrong word. * * @param out_words FT_TT_OP_ATTENTION_WORDS values, fully written (unused words zeroed) * @return FT_TT_OK; FT_TT_EINVAL on a zero or oversized group, a zero head * dimension, a zero KV length, a query row past a tile, a tile size that * is not the cache's, or a padded read window past `capacity_tiles` * * Thread-safe: yes Reentrant: yes Signal-safe: yes */ int ft_tt_op_attention_parameters(const ft_tt_op_attention_request *request, uint32_t out_words[FT_TT_OP_ATTENTION_WORDS]); /** Bytes ft_tt_op_attention_build_constants() writes. */ size_t ft_tt_op_attention_constants_bytes(uint32_t tile_bytes); /** * Build the constants operand: the reduce scaler, the query scale, and all 32 tail masks. * * THE MASKS ARE BUILT ONCE PER MODEL AND SELECTED BY ARITHMETIC. Mask `t` is 0 in columns * below `t` and FT_TT_OP_ATTENTION_MASK_VALUE at and above it, so mask `kv_positions % 32` * kills the dead columns of the last live position tile and mask 0 — every column dead — * kills a tile entirely past the end. The kernel derives which one it needs from * `kv_positions`; nothing about the mask is on the wire. Uploading a per-call mask instead * would be 2 KiB per layer per token over a 287 MB/s link, i.e. ~7 µs of pure host cost per * layer that buys nothing. * * THE SCALER IS 1.0, FOR BOTH REDUCES. Upstream's reduce scaler is 1.0 for SUM and for MAX * alike (`reduce_helpers_dataflow.hpp`), and both of this op's reduces want a plain sum and * a plain max — so one tile serves both, exactly as in the REDUCE family. * * @param head_dim the real head dimension in ELEMENTS: the query scale is 1/sqrt of it, * so a dim_tiles-derived value would silently be wrong for a head_dim * that is not a multiple of 32 * @param out_bytes at least ft_tt_op_attention_constants_bytes() * @return FT_TT_OK, FT_TT_EINVAL * * Thread-safe: yes Reentrant: yes Signal-safe: no */ int ft_tt_op_attention_build_constants(uint32_t head_dim, uint32_t tile_bytes, void *out, size_t out_bytes); #endif /* FT_TT_OP_ATTENTION_H */