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

372 lines
16 KiB
C
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/**
* @file resident_attention_reader.c
* @brief Reader third of the RESIDENT FLASH-DECODE ATTENTION triad: the group's Q, packed into
* `dim_tiles` tiles whose rows are Q heads; the four constant tiles this descriptor
* selects; and then the KV cache, ONE BLOCK of K and one of V at a time, each block a
* SINGLE contiguous `noc_async_read`.
*
* THE KV BLOCK IS ONE TRANSACTION AND THAT IS THE WHOLE POINT OF THE LAYOUT.
* ft_tt_op_memory.h PART 1 §1.2: positions [32a, 32b) of one head are one contiguous byte
* range, walked linearly in exactly the order the matmul reader already wants. So a block
* of `block_tiles` position tiles is `block_tiles * dim_tiles` consecutive tiles, i.e.
* ONE read and ONE barrier -- the shape that was worth 5.9x on the matmul
* (docs/tt_kernel_model.md §12) rather than a strided walk.
*
* THE Q PACK IS THE ONE PIECE OF REAL WORK THE READER DOES.
* Q arrives in the model's activation layout -- head h's dim tile d at tile column
* `h*dim_tiles + d`, the decode token at row `position % 32` -- and the compute thread
* wants ONE tile per dim tile whose ROW h is head `q_first + h`. The group's tiles are
* consecutive, so the whole set is one read into a staging buffer, and the pack is
* `q_rows * dim_tiles * 2` copies of a 32-byte face row. It reads 64 KiB to use 2 KiB of
* it, and that is the right trade: it is ~0.4% of one layer's KV traffic at 8k context,
* and the alternative -- a 32-byte strided DRAM read per face row -- is not legal at all
* (64 B read alignment, ft_tt_noc_alignment()).
*
* THIS PROCESSOR OWNS `tail`, AND NOTHING ELSE IN THE CONTROL BLOCK.
*
*
* C99 on Foundry's own `ft_tensix_dataflow.h` — no tt-metal include (plan Phase 6.6).
* The wait mode, which used to be a `template <uint32_t kMode>` argument, is now an
* ordinary parameter carrying the same compile-time value; `-Os` folds it identically.
* Compile-time args:
* 0: CB_Q index 1: CB_K index 2: CB_V index 3: CB_CONST index
* 4: CB_QSTAGE index (reader-private staging; never pushed, never popped)
* 5: tile size in bytes
* 6: wait mode -- 0 upstream, 1 stoppable (ships), 2 bound-only
* Runtime args: the ten of resident_queue_protocol.h, identical on all five processors
*
* Build (BRISC, NoC 0):
* CTARGS=0,1,2,3,13,2048,1 NOC=0 KERNEL_TEXT_BASE=0x15000 bash \
* foundry/tt/build_tensix_dataflow_kernel.sh \
* foundry/tt/kernels/resident_attention_reader.c brisc /tmp/ttk_attn
*
* @status GREEN ON SILICON on Tensix 70/71/72/73, 0 checks failed.
* @issues None
* @todo None
*/
#include <stdbool.h>
#include <stdint.h>
#include "ft_tensix_dataflow.h"
#include "resident_queue_protocol.h"
#include "resident_attention_protocol.h"
#define CB_Q ft_tensix_compile_time_arg(0)
#define CB_K ft_tensix_compile_time_arg(1)
#define CB_V ft_tensix_compile_time_arg(2)
#define CB_CONST ft_tensix_compile_time_arg(3)
#define CB_QSTAGE ft_tensix_compile_time_arg(4)
#define TILE_BYTES ft_tensix_compile_time_arg(5)
#define WAIT_MODE ft_tensix_compile_time_arg(6)
/**
* Refuse a descriptor the circular buffers this residency was STAGED with cannot serve.
*
* The page counts are fixed when the resident program is staged; the shape is a
* per-descriptor word. So the host CAN ask for a block bigger than the buffer -- it did, on
* another op's first hardware run, and the result was ten seconds of blocked polling
* reported as an idle timeout (resident_topk_reader.c). And the two block-prefetch
* conditions are the §5 invariant: `pages % depth == 0` so every wrap lands exactly on the
* limit, and `tile_count % pages == 0` so the descriptor leaves the write pointer where it
* found it. Without the second, a later descriptor's block steps OVER the limit and walks
* up L1 -- measured, on this family, landing on the queue control block with nothing
* faulting.
*/
static bool descriptor_is_executable(ftq_context *ctx, const fta_shape *shape)
{
const uint32_t depth = shape->block_tiles * shape->dim_tiles;
const uint32_t k_pages = ft_tensix_cb(CB_K)->fifo_num_pages;
const uint32_t v_pages = ft_tensix_cb(CB_V)->fifo_num_pages;
const uint32_t q_pages = ft_tensix_cb(CB_Q)->fifo_num_pages;
const uint32_t stage_pages = ft_tensix_cb(CB_QSTAGE)->fifo_num_pages;
const uint32_t const_pages = ft_tensix_cb(CB_CONST)->fifo_num_pages;
const uint32_t tile_count = shape->padded_tiles * shape->dim_tiles;
/* CHUNK mode reads the Q tiles straight into CB_Q, so it never touches the staging
* buffer and must not be measured against it: at C = 32 and dim_tiles = 8 the decode
* residency's 32 staging pages would refuse every chunk descriptor for a buffer the
* descriptor does not use. */
const uint32_t stage_needed = fta_is_chunk(shape) ? 0u : shape->q_rows * shape->dim_tiles;
if (!fta_shape_valid(shape) || TILE_BYTES == 0u || depth == 0u)
ctx->error = kErrBadParameters;
else if (k_pages == 0u || v_pages == 0u || q_pages < shape->dim_tiles ||
const_pages < kAttnCbConstPages || stage_pages < stage_needed)
ctx->error = kErrBadParameters;
else if (depth > k_pages || (k_pages % depth) != 0u || (tile_count % k_pages) != 0u)
ctx->error = kErrBadParameters;
else if (depth > v_pages || (v_pages % depth) != 0u || (tile_count % v_pages) != 0u)
ctx->error = kErrBadParameters;
else
return true;
return false;
}
/**
* Read `tiles` consecutive tiles at `base + first*TILE_BYTES` into `cb` as ONE block.
*
* @return false if a reservation was abandoned, or if the length came out zero -- which is
* a refusal and never a skipped transfer (§6)
*/
static bool read_block(ftq_context *ctx, uint32_t cb, uint32_t x, uint32_t y, uint32_t base,
uint32_t first, uint32_t tiles)
{
const uint32_t bytes = tiles * TILE_BYTES;
if (!fta_length_is_issuable(bytes)) {
ftq_abort(ctx, kErrBadParameters);
return false;
}
if (!ftq_cb_reserve_back_stoppable(ctx, (int32_t)cb, (int32_t)tiles, WAIT_MODE))
return false;
// Rule 3: while draining, issue NOTHING. The push still happens so the page counters
// stay balanced, but a transaction issued during a drain would still be in flight when
// the kernel returns.
if (!ctx->aborting) {
ft_tensix_noc_read(ft_tensix_noc_address(x, y, base + first * TILE_BYTES),
ft_tensix_cb_write_pointer(cb), bytes);
// ONE barrier for the whole block: it is per-core, not per-transaction.
ft_tensix_noc_read_barrier();
}
ft_tensix_cb_push_back(cb, tiles);
return true;
}
/**
* The four constants this descriptor uses, in the order the compute thread indexes them:
* the reduce scaler, the 1/sqrt(head_dim) scale, the tail mask and the all-dead mask.
*
* Two of the four are adjacent in the operand so they are one transaction; the two masks
* are selected by `kv_positions % 32` and by nothing else, which is why the mask INDEX is
* derived here rather than sent (a second field that must agree with the first is a second
* field that can disagree).
*/
static bool read_constants(ftq_context *ctx, const fta_shape *shape, uint32_t x, uint32_t y,
uint32_t base, uint32_t rows_mask)
{
/* The TAIL slot carries the causal tile in chunk mode and the tail mask otherwise --
* one function, resident_attention_protocol.h, so this and the off-hardware suite
* cannot spell the selection differently. */
const uint32_t tail_mask = fta_tail_slot_source_tile(shape);
const uint32_t bytes = kAttnCbConstPages * TILE_BYTES;
if (!fta_length_is_issuable(bytes)) {
ftq_abort(ctx, kErrBadParameters);
return false;
}
if (!ftq_cb_reserve_back_stoppable(ctx, (int32_t)CB_CONST,
(int32_t)kAttnCbConstPages, WAIT_MODE))
return false;
if (!ctx->aborting) {
const uint32_t write = ft_tensix_cb_write_pointer(CB_CONST);
/* Tiles 0 and 1 are adjacent in the operand: one transaction. */
ft_tensix_noc_read(ft_tensix_noc_address(x, y, base + kAttnConstScalerTile * TILE_BYTES),
write,
2u * TILE_BYTES);
ft_tensix_noc_read(ft_tensix_noc_address(x, y,
(shape->flags & (uint32_t)kAttnFlagRows) != 0u
? rows_mask
: base + tail_mask * TILE_BYTES),
write + kAttnCbConstTailMask * TILE_BYTES, TILE_BYTES);
ft_tensix_noc_read(ft_tensix_noc_address(x, y, base + kAttnConstMaskTile * TILE_BYTES),
write + kAttnCbConstDeadMask * TILE_BYTES, TILE_BYTES);
ft_tensix_noc_read_barrier();
}
ft_tensix_cb_push_back(CB_CONST, kAttnCbConstPages);
return true;
}
/**
* Build the packed Q: `dim_tiles` tiles whose row h is Q head `q_first + h`.
*
* The group's `q_rows * dim_tiles` Q tiles are consecutive in the activation buffer, so
* they arrive in ONE transaction; the pack is then a face-row scatter inside L1. Source
* tile (h, dt) is at staging index `h * dim_tiles + dt` and contributes its row
* `query_row` to row h of packed tile dt.
*
* IN CHUNK MODE THERE IS NO PACK AT ALL — see the branch below.
*/
static bool read_and_pack_q(ftq_context *ctx, const fta_shape *shape, uint32_t x, uint32_t y,
uint32_t base)
{
/* Chunk mode reads ONE head's `dim_tiles` tiles; decode reads the whole group's
* `q_rows * dim_tiles` into staging and packs one row out of each. */
const uint32_t source_tiles =
fta_is_chunk(shape) ? shape->dim_tiles : shape->q_rows * shape->dim_tiles;
const uint32_t bytes = source_tiles * TILE_BYTES;
const uint32_t chunk = fta_chunk_bytes(TILE_BYTES);
uint32_t stage;
uint32_t packed;
if (!fta_length_is_issuable(bytes) || !fta_length_is_issuable(chunk)) {
ftq_abort(ctx, kErrBadParameters);
return false;
}
if (!ftq_cb_reserve_back_stoppable(ctx, (int32_t)CB_Q, (int32_t)shape->dim_tiles, WAIT_MODE))
return false;
/* CHUNK MODE: THERE IS NOTHING TO PACK. The activation tile already has position t in
* row t, which IS the packed Q -- `dim_tiles` tiles whose rows are what the score
* matmul's M axis should be. One contiguous read, no staging buffer, no face-row
* copies: 16 transactions a layer instead of 64 KiB read to use 2 KiB of it. */
if (fta_is_chunk(shape)) {
if (!ctx->aborting) {
ft_tensix_noc_read(ft_tensix_noc_address(x, y, base),
ft_tensix_cb_write_pointer(CB_Q), bytes);
ft_tensix_noc_read_barrier();
}
ft_tensix_cb_push_back(CB_Q, shape->dim_tiles);
return true;
}
stage = ft_tensix_cb_write_pointer(CB_QSTAGE);
packed = ft_tensix_cb_write_pointer(CB_Q);
if (!ctx->aborting) {
ft_tensix_noc_read(ft_tensix_noc_address(x, y, base), stage, bytes);
ft_tensix_noc_read_barrier();
for (uint32_t dt = 0; dt < shape->dim_tiles; dt++) {
for (uint32_t p = 0; p < shape->positions; p++) {
for (uint32_t head = 0; head < shape->q_rows; head++) {
const uint32_t source_tile = (head * shape->dim_tiles + dt) * TILE_BYTES;
for (uint32_t half = 0; half < 2u; half++) {
const uint32_t from = stage + source_tile +
fta_chunk_offset(shape->query_row + p, half,
TILE_BYTES);
const uint32_t to = packed + dt * TILE_BYTES +
fta_chunk_offset(p * shape->q_rows + head, half,
TILE_BYTES);
fta_copy_words((uint32_t *)(uintptr_t)(to),
(const uint32_t *)(uintptr_t)(from), chunk / 4u);
}
}
}
}
}
ft_tensix_cb_push_back(CB_Q, shape->dim_tiles);
return true;
}
/**
* The reader half of one ATTENTION descriptor.
*
* Every trip count comes from the parameter block and none from a wait, which is what makes
* the drain safe (resident_queue_protocol.h rule 3).
*
* @return true if the whole descriptor was read; false if a wait was abandoned, in which
* case the loop still ran to completion and the caller must not retire it
*/
static bool read_one_attention(ftq_context *ctx, volatile FTQ_L1_PTR uint32_t *parameters)
{
const fta_shape shape = fta_decode(parameters);
const uint32_t q_x = fta_coord_x(parameters[kAttnQCoord]);
const uint32_t q_y = fta_coord_y(parameters[kAttnQCoord]);
const uint32_t k_x = fta_coord_x(parameters[kAttnKCoord]);
const uint32_t k_y = fta_coord_y(parameters[kAttnKCoord]);
const uint32_t v_x = fta_coord_x(parameters[kAttnVCoord]);
const uint32_t v_y = fta_coord_y(parameters[kAttnVCoord]);
const uint32_t c_x = fta_coord_x(parameters[kAttnConstCoord]);
const uint32_t c_y = fta_coord_y(parameters[kAttnConstCoord]);
const uint32_t depth = shape.block_tiles * shape.dim_tiles;
if (!descriptor_is_executable(ctx, &shape)) {
ftq_abort(ctx, ctx->error);
return false;
}
if (!read_and_pack_q(ctx, &shape, q_x, q_y, parameters[kAttnQBase]))
return false;
if (!read_constants(ctx, &shape, c_x, c_y, parameters[kAttnConstBase],
parameters[kAttnReserved]))
return false;
for (uint32_t block = 0; block < shape.blocks; block++) {
const uint32_t first = block * depth;
if (!read_block(ctx, CB_K, k_x, k_y, parameters[kAttnKBase], first, depth))
return false;
if (!read_block(ctx, CB_V, v_x, v_y, parameters[kAttnVBase], first, depth))
return false;
}
return !ctx->aborting;
}
void kernel_main(void)
{
ftq_context ctx;
bool running = true;
ftq_context_init(&ctx);
ftq_status_set(&ctx, kStateBoot);
ctx.control[kFieldState] = kStateBoot;
if (!ftq_check_control(&ctx)) {
ctx.control[kFieldError] = ctx.error;
ctx.control[kFieldState] = kStateExit;
ftq_status_set(&ctx, kStateExit);
return;
}
// RESUME rather than restart; resident_matmul_reader.c carries the argument and the
// refusal that enforces it.
if (ctx.control[kFieldTail] != ctx.control[kFieldCompletedCount]) {
ctx.error = kErrBadParameters;
ctx.control[kFieldError] = ctx.error;
ctx.control[kFieldState] = kStateExit;
ftq_status_set(&ctx, kStateExit);
return;
}
ctx.cursor = ctx.control[kFieldCompletedCount];
while (running) {
if (!ftq_wait_for_work(&ctx, kFieldHead))
break;
ctx.aborting = false;
ftq_status_set(&ctx, kStateExec);
volatile FTQ_L1_PTR uint32_t *descriptor = ftq_descriptor(&ctx);
const uint32_t opcode = descriptor[kDescOpcode];
bool done;
ctx.control[kFieldLastOpcode] = opcode;
ftq_status_opcode(&ctx, opcode);
if (opcode == kOpAttention) {
ftq_status_set(&ctx, kStateCbReserveBack);
done = read_one_attention(&ctx, ftq_parameters(descriptor));
ftq_status_set(&ctx, kStateExec);
if (!done)
running = false;
} else if (opcode == kOpNop) {
done = true;
} else if (opcode == kOpStop) {
done = true;
running = false;
} else {
// A binary is bound to one slab, so an opcode another program implements is
// genuinely not available here. Refusing names that.
ctx.error = kErrBadOpcode;
done = false;
running = false;
}
if (done) {
ctx.cursor++;
ctx.control[kFieldTail] = ctx.cursor;
}
}
ft_tensix_noc_read_barrier();
ft_tensix_noc_write_barrier();
ctx.control[kFieldError] = ctx.error;
ctx.control[kFieldState] = kStateExit;
ftq_status_set(&ctx, kStateExit);
}
Reference in New Issue View Git Blame Copy Permalink