David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/tt/ft_tt_op_attention.c

361 lines
16 KiB
C
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file ft_tt_op_attention.c
* @brief Implementation of flash-decode attention's host driver. See ft_tt_op_attention.h,
* which carries the decision that one core owns a KV GROUP rather than a Q head, and the
* capacity argument that makes the online softmax a requirement.
*
* Everything here is PURE — no device access, no allocation. The enqueue itself is
* ft_tt_queue_write_op_parameters() plus ft_tt_queue_enqueue(), which the caller does,
* because a resident op has no launch of its own to own.
*
* @status See docs/tt_bringup_log.md and the Phase 6 attention report.
* @issues None
* @todo None
*/
#include "ft_tt_op_attention.h"
#include <math.h>
#include <string.h>
#include "convert/ft_float_convert.h"
#include "ft_tt_device.h"
#include "ft_tt_op_l1_layout.h"
#include "ft_tt_tilize.h"
/** K and V hold TWO blocks: the reader fills one while the compute thread drains the
* other. Two is also what makes `pages % depth == 0` hold with a `depth`-sized push. */
#define ATTENTION_KV_BLOCKS 2u
/** Pages in the buffers that carry one tile at a time. */
#define ATTENTION_CB_PAGES_SMALL 1u
/** Elements in one 32x32 tile. */
#define ATTENTION_TILE_ELEMENTS (FT_TT_OP_TILE_DIM * FT_TT_OP_TILE_DIM)
/** `(y << 16) | x`, the packed coordinate the parameter block carries. */
static uint32_t pack_coordinate(const ft_tt_op_operand *operand)
{
return ((uint32_t)operand->y << 16) | (uint32_t)operand->x;
}
/** Is `operand` filled in, and does it carry the tile size the whole op runs at? */
static bool operand_is_valid(const ft_tt_op_operand *operand, uint32_t tile_bytes)
{
return operand->tile_bytes == tile_bytes && tile_bytes != 0u;
}
uint32_t ft_tt_op_attention_live_tiles(uint32_t kv_positions)
{
return (kv_positions + FT_TT_OP_TILE_DIM - 1u) / FT_TT_OP_TILE_DIM;
}
uint32_t ft_tt_op_attention_padded_tiles(uint32_t kv_positions, uint32_t block_tiles)
{
uint32_t live;
uint32_t window;
if (block_tiles == 0u)
return 0u;
live = ft_tt_op_attention_live_tiles(kv_positions);
window = ATTENTION_KV_BLOCKS * block_tiles;
return ((live + window - 1u) / window) * window;
}
/** Pages each of the sixteen buffers wants, in ft_tt_op_attention_place_cbs()' order. This
* is the ONE place the page counts live; the L1 arithmetic and the placement both read it,
* so a residency's cost and its layout cannot disagree. */
static void page_counts(const ft_tt_op_attention_residency *residency,
uint32_t out_pages[FT_TT_OP_ATTENTION_CB_COUNT],
unsigned out_index[FT_TT_OP_ATTENTION_CB_COUNT])
{
const uint32_t depth = residency->block_tiles * residency->dim_tiles;
out_index[0] = FT_TT_OP_ATTENTION_CB_Q;
out_pages[0] = residency->dim_tiles;
out_index[1] = FT_TT_OP_ATTENTION_CB_QS;
out_pages[1] = residency->dim_tiles;
out_index[2] = FT_TT_OP_ATTENTION_CB_K;
out_pages[2] = ATTENTION_KV_BLOCKS * depth;
out_index[3] = FT_TT_OP_ATTENTION_CB_V;
out_pages[3] = ATTENTION_KV_BLOCKS * depth;
out_index[4] = FT_TT_OP_ATTENTION_CB_CONST;
out_pages[4] = FT_TT_OP_ATTENTION_CONST_PAGES;
out_index[5] = FT_TT_OP_ATTENTION_CB_S;
out_pages[5] = residency->block_tiles;
out_index[6] = FT_TT_OP_ATTENTION_CB_SCRATCH;
out_pages[6] = ATTENTION_CB_PAGES_SMALL;
out_index[7] = FT_TT_OP_ATTENTION_CB_M_A;
out_pages[7] = ATTENTION_CB_PAGES_SMALL;
out_index[8] = FT_TT_OP_ATTENTION_CB_M_B;
out_pages[8] = ATTENTION_CB_PAGES_SMALL;
out_index[9] = FT_TT_OP_ATTENTION_CB_L;
out_pages[9] = ATTENTION_CB_PAGES_SMALL;
out_index[10] = FT_TT_OP_ATTENTION_CB_LSUM;
out_pages[10] = ATTENTION_CB_PAGES_SMALL;
out_index[11] = FT_TT_OP_ATTENTION_CB_C;
out_pages[11] = ATTENTION_CB_PAGES_SMALL;
out_index[12] = FT_TT_OP_ATTENTION_CB_O;
out_pages[12] = residency->dim_tiles;
out_index[13] = FT_TT_OP_ATTENTION_CB_PV;
out_pages[13] = residency->dim_tiles;
/* The reader's private staging for the group's raw Q tiles. Declared as a circular
* buffer so its L1 comes out of the arena allocator like everything else's and shows up
* in the same accounting — nothing flows through it, and it is never pushed or popped. */
out_index[14] = FT_TT_OP_ATTENTION_CB_QSTAGE;
out_pages[14] = residency->q_rows * residency->dim_tiles;
out_index[15] = FT_TT_OP_ATTENTION_CB_OUT;
out_pages[15] = 2u * residency->dim_tiles;
}
/** Is this a residency that can be placed at all, before the arena is consulted? */
static bool residency_is_valid(const ft_tt_op_attention_residency *residency)
{
if (!residency)
return false;
if (residency->q_rows == 0u || residency->q_rows > FT_TT_OP_TILE_DIM)
return false;
if (residency->dim_tiles == 0u || residency->block_tiles == 0u)
return false;
if (residency->tile_bytes == 0u)
return false;
/* The KV cache is bf16 and block float is refused by name, not ignored
* (ft_tt_op_memory.h §1.6) — a shared exponent along a face row cannot survive a
* single-position insert. */
if (residency->tile_bytes != FT_TT_OP_TILE_BYTES_BF16 &&
residency->tile_bytes != FT_TT_OP_TILE_BYTES_FP32)
return false;
return true;
}
size_t ft_tt_op_attention_l1_bytes(const ft_tt_op_attention_residency *residency)
{
uint32_t pages[FT_TT_OP_ATTENTION_CB_COUNT];
unsigned indices[FT_TT_OP_ATTENTION_CB_COUNT];
size_t total = 0;
unsigned index;
if (!residency_is_valid(residency))
return 0;
page_counts(residency, pages, indices);
for (index = 0; index < FT_TT_OP_ATTENTION_CB_COUNT; index++) {
/* The allocator rounds each buffer up to 64 bytes, so the estimate must too or it
* under-reports the residency that will actually be placed. */
const size_t bytes = (size_t)pages[index] * residency->tile_bytes;
total += (bytes + 63u) & ~(size_t)63u;
}
return total;
}
uint32_t ft_tt_op_attention_max_block_tiles(uint32_t q_rows, uint32_t dim_tiles,
uint32_t tile_bytes)
{
ft_tt_op_attention_residency residency;
uint32_t best = 0u;
uint32_t candidate;
residency.q_rows = q_rows;
residency.dim_tiles = dim_tiles;
residency.tile_bytes = tile_bytes;
/* Searched rather than solved: the arena cost is piecewise because every buffer is
* rounded up to 64 bytes, and a closed form would have to reproduce that rounding —
* which is exactly the kind of second copy of an arithmetic that drifts. 256 is far
* past anything that fits. */
for (candidate = 1u; candidate <= 256u; candidate++) {
residency.block_tiles = candidate;
if (!residency_is_valid(&residency))
return 0u;
if (ft_tt_op_attention_l1_bytes(&residency) >
(size_t)(FT_TT_OP_CB_ARENA_END - FT_TT_OP_CB_ARENA_BASE))
break;
best = candidate;
}
return best;
}
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)
{
ft_tt_op_cb_request requests[FT_TT_OP_ATTENTION_CB_COUNT];
uint32_t pages[FT_TT_OP_ATTENTION_CB_COUNT];
unsigned indices[FT_TT_OP_ATTENTION_CB_COUNT];
unsigned index;
if (!residency_is_valid(residency) || !out_entries || !buffer || !out_mask)
return FT_TT_EINVAL;
page_counts(residency, pages, indices);
for (index = 0; index < FT_TT_OP_ATTENTION_CB_COUNT; index++) {
requests[index].index = indices[index];
requests[index].page_bytes = residency->tile_bytes;
requests[index].page_count = pages[index];
}
return ft_tt_op_cb_place(requests, FT_TT_OP_ATTENTION_CB_COUNT, out_entries, buffer,
buffer_bytes, out_mask);
}
int ft_tt_op_attention_parameters(const ft_tt_op_attention_request *request,
uint32_t out_words[FT_TT_OP_ATTENTION_WORDS])
{
uint32_t tile_bytes;
uint32_t padded;
if (!request || !out_words)
return FT_TT_EINVAL;
tile_bytes = request->query.tile_bytes;
if (request->q_rows == 0u || request->q_rows > FT_TT_OP_TILE_DIM)
return FT_TT_EINVAL;
if (request->dim_tiles == 0u || request->block_tiles == 0u || request->kv_positions == 0u)
return FT_TT_EINVAL;
if (request->query_row >= FT_TT_OP_TILE_DIM)
return FT_TT_EINVAL;
if ((request->flags & ~(uint32_t)(FT_TT_OP_ATTENTION_FLAG_CHUNK |
FT_TT_OP_ATTENTION_FLAG_ROWS)) != 0u)
return FT_TT_EINVAL;
if ((request->flags & (uint32_t)FT_TT_OP_ATTENTION_FLAG_ROWS) != 0u &&
(request->positions == 0u || request->q_rows * request->positions > FT_TT_OP_TILE_DIM ||
request->query_row + request->positions > FT_TT_OP_TILE_DIM ||
request->kv_positions < request->positions ||
(request->kv_positions - request->positions) / FT_TT_OP_TILE_DIM !=
(request->kv_positions - 1u) / FT_TT_OP_TILE_DIM ||
request->rows_mask_address == 0u))
return FT_TT_EINVAL;
/* THE CHUNK GEOMETRY, REFUSED HERE AS WELL AS ON THE DEVICE. The chunk's C positions
* END the live range and must START on a position-tile boundary, so the chunk occupies
* exactly the LAST LIVE TILE and row t is position `first + t`. A chunk straddling two
* position tiles would need two masks and there is one mask slot — and it would be a
* wrong mask on a live tile, i.e. a plausible wrong answer with nothing out of range. */
if ((request->flags & (uint32_t)FT_TT_OP_ATTENTION_FLAG_CHUNK) != 0u) {
if (request->query_row != 0u)
return FT_TT_EINVAL;
if (request->kv_positions < request->q_rows)
return FT_TT_EINVAL;
if (((request->kv_positions - request->q_rows) % FT_TT_OP_TILE_DIM) != 0u)
return FT_TT_EINVAL;
}
if (!operand_is_valid(&request->query, tile_bytes) ||
!operand_is_valid(&request->keys, tile_bytes) ||
!operand_is_valid(&request->values, tile_bytes) ||
!operand_is_valid(&request->constants, tile_bytes) ||
!operand_is_valid(&request->output, tile_bytes))
return FT_TT_EINVAL;
/* The PADDED window is what the reader fetches, so it — not the live length — is what
* has to fit the region. Getting this wrong is a read outside the KV cache. */
padded = ft_tt_op_attention_padded_tiles(request->kv_positions, request->block_tiles);
if (padded == 0u || padded > request->capacity_tiles)
return FT_TT_EINVAL;
memset(out_words, 0, FT_TT_OP_ATTENTION_WORDS * sizeof(out_words[0]));
out_words[FT_TT_OP_ATTENTION_HEADS] =
request->q_rows | (request->query_row << 8) | (request->flags << 16) |
((request->flags & (uint32_t)FT_TT_OP_ATTENTION_FLAG_ROWS) != 0u
? (request->positions & 0xFFu) << 24 : 0u);
if ((request->flags & (uint32_t)FT_TT_OP_ATTENTION_FLAG_ROWS) != 0u)
out_words[FT_TT_OP_ATTENTION_RESERVED] = request->rows_mask_address;
out_words[FT_TT_OP_ATTENTION_DIM_TILES] = request->dim_tiles;
out_words[FT_TT_OP_ATTENTION_KV_POSITIONS] = request->kv_positions;
out_words[FT_TT_OP_ATTENTION_BLOCK_TILES] = request->block_tiles;
out_words[FT_TT_OP_ATTENTION_CAPACITY_TILES] = request->capacity_tiles;
out_words[FT_TT_OP_ATTENTION_Q_COORD] = pack_coordinate(&request->query);
out_words[FT_TT_OP_ATTENTION_Q_BASE] = request->query.address;
out_words[FT_TT_OP_ATTENTION_K_COORD] = pack_coordinate(&request->keys);
out_words[FT_TT_OP_ATTENTION_K_BASE] = request->keys.address;
out_words[FT_TT_OP_ATTENTION_V_COORD] = pack_coordinate(&request->values);
out_words[FT_TT_OP_ATTENTION_V_BASE] = request->values.address;
out_words[FT_TT_OP_ATTENTION_CONST_COORD] = pack_coordinate(&request->constants);
out_words[FT_TT_OP_ATTENTION_CONST_BASE] = request->constants.address;
out_words[FT_TT_OP_ATTENTION_OUT_COORD] = pack_coordinate(&request->output);
out_words[FT_TT_OP_ATTENTION_OUT_BASE] = request->output.address;
return FT_TT_OK;
}
size_t ft_tt_op_attention_constants_bytes(uint32_t tile_bytes)
{
return (size_t)FT_TT_OP_ATTENTION_CONST_TILES * tile_bytes;
}
/** Fill one whole face-ordered tile with `value`. Only column 0 of the scaler is read by
* the reduce LLK and only the mask's columns matter for a mask, but filling the whole tile
* is a superset and costs nothing at build time. */
static void fill_tile(uint16_t *tile, float value)
{
const uint16_t encoded = ft_bf16_from_float(value);
uint32_t index;
for (index = 0; index < ATTENTION_TILE_ELEMENTS; index++)
tile[index] = encoded;
}
/** Mask tile `tail`: 0 in columns below `tail`, the mask value at and above it. `tail` 0 is
* every column dead, which is what a position tile entirely past the end needs. */
static void fill_mask_tile(uint16_t *tile, uint32_t tail)
{
const uint16_t dead = ft_bf16_from_float(FT_TT_OP_ATTENTION_MASK_VALUE);
const uint16_t live = ft_bf16_from_float(0.0f);
uint32_t row;
for (row = 0; row < FT_TT_OP_TILE_DIM; row++) {
uint32_t column;
for (column = 0; column < FT_TT_OP_TILE_DIM; column++)
tile[ft_tt_tile_face_index(row, column)] = (column < tail) ? live : dead;
}
}
/** Causal mask tile for a chunk of `live` positions: element (row r, column c) is 0 when
* `c <= r && c < live` and the mask value otherwise.
*
* ROW r IS POSITION `first + r` AND COLUMN c IS POSITION `first + c`, which is the whole
* of why one constant tile is the entire causality mechanism: `c <= r` is "position r may
* attend to position c", and `c < live` kills the columns of a ragged chunk that hold no
* position at all. Rows at and above `live` are the strip's padding rows; they see the
* whole chunk, which keeps their softmax denominator non-zero and their output finite —
* garbage in a row every consumer discards, rather than a NaN that could propagate. */
static void fill_causal_mask_tile(uint16_t *tile, uint32_t live)
{
const uint16_t dead = ft_bf16_from_float(FT_TT_OP_ATTENTION_MASK_VALUE);
const uint16_t alive = ft_bf16_from_float(0.0f);
uint32_t row;
for (row = 0; row < FT_TT_OP_TILE_DIM; row++) {
uint32_t column;
for (column = 0; column < FT_TT_OP_TILE_DIM; column++)
tile[ft_tt_tile_face_index(row, column)] =
(column <= row && column < live) ? alive : dead;
}
}
int ft_tt_op_attention_build_constants(uint32_t head_dim, uint32_t tile_bytes, void *out,
size_t out_bytes)
{
uint16_t *tiles = (uint16_t *)out;
uint32_t tail;
if (!out || head_dim == 0u || tile_bytes != FT_TT_OP_TILE_BYTES_BF16)
return FT_TT_EINVAL;
if (out_bytes < ft_tt_op_attention_constants_bytes(tile_bytes))
return FT_TT_EINVAL;
memset(out, 0, ft_tt_op_attention_constants_bytes(tile_bytes));
fill_tile(tiles + FT_TT_OP_ATTENTION_CONST_SCALER_TILE * ATTENTION_TILE_ELEMENTS, 1.0f);
fill_tile(tiles + FT_TT_OP_ATTENTION_CONST_SCALE_TILE * ATTENTION_TILE_ELEMENTS,
(float)(1.0 / sqrt((double)head_dim)));
for (tail = 0; tail < FT_TT_OP_TILE_DIM; tail++)
fill_mask_tile(tiles + (FT_TT_OP_ATTENTION_CONST_MASK_TILE + tail) *
ATTENTION_TILE_ELEMENTS,
tail);
/* Indexed by C - 1, so C = 1..32 rather than 0..31: a chunk of no positions is not a
* descriptor, and the all-dead tile already exists for a tile past the end. */
for (tail = 0; tail < FT_TT_OP_TILE_DIM; tail++)
fill_causal_mask_tile(tiles + (FT_TT_OP_ATTENTION_CONST_CAUSAL_TILE + tail) *
ATTENTION_TILE_ELEMENTS,
tail + 1u);
return FT_TT_OK;
}
Reference in New Issue View Git Blame Copy Permalink