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

645 lines
27 KiB
C
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/**
* @file resident_attention_compute.c
* @brief Compute third of the resident flash-decode ATTENTION triad. For one query (packed Q,
* dim_tiles tiles) against a KV cache delivered in blocks of block_tiles position tiles,
* computes softmax(Q K^T * scale + mask) V with an online softmax, never materialising the
* whole score row. Per block:
*
* S = Qs . K_block^T (Qs = Q * scale, once per descriptor)
* S += mask only for position tiles not entirely live
* m_new = max(m_old, rowmax(S)) row-max reduce folded with the previous running max
* P = exp(S - m_new) in place in the score CB
* s = rowsum(P)
* first block: O = P.V, l = s
* later blocks: c = exp(m_old - m_new); l = c*l + s; O = c*O + P.V
*
* and finally out = O * (1/l).
*
* Non-obvious requirements:
* - The running max is seeded with the previous block's max (binary max in DST). A
* block-local max still gives correct-looking correlation, so nothing else would catch it.
* - The mask is additive and applied before the max, so stale cache data past the live end
* cannot win the max.
* - Every block processes exactly block_tiles tiles, dead ones included, so each CB's
* push/pop counts are fixed per descriptor; all trip counts come from the parameter block.
* - fp32 DST accumulation is required (the cb table provides it).
*
* Compiled three times (FT_TRISC_UNPACK / MATH / PACK); every hardware statement sits in one
* FT_TENSIX_<thread>(...) wrapper. Control flow, waits and the queue protocol run on all.
*
* Compile-time args: 0 cb_q, 1 cb_k, 2 cb_v, 3 cb_const (scaler, scale, tail mask, dead mask),
* 4 cb_s, 5 cb_scratch, 6 cb_m_a, 7 cb_m_b, 8 cb_l, 9 cb_lsum, 10 cb_c, 11 cb_o, 12 cb_pv,
* 13 cb_out, 14 cb_qs, 15 wait_mode (0 upstream, 1 stoppable, 2 bound-only).
*
* Build:
* CTARGS=0,1,2,3,4,5,6,7,8,9,10,11,12,16,14,1 \
* CB_TABLE=kernels/cb_tables/attention_float16b.cbtable \
* bash foundry/tt/build_tensix_kernel.sh \
* foundry/tt/kernels/resident_attention_compute.c out/ttk_attention
*
* The UNPACK image sits near the top of its 0x3000-byte text slab; the two noinline wait
* wrappers below are what keep it inside (measured).
*
* @status Clean-room rewrite (plan Phase 4.7).
* @issues None known.
* @todo None
*/
#include <stdbool.h>
#include <stdint.h>
#include "ft_tensix_cb_formats.h"
#include "ft_tensix_compute.h"
#include "ft_tensix_compute_broadcast.h"
#include "ft_tensix_compute_cb.h"
#include "ft_tensix_compute_matmul.h"
#include "ft_tensix_compute_reduce.h"
#include "ft_tensix_kernel_args.h"
#ifdef FT_TRISC_MATH
#include "ft_tensix_compute_math.h"
#include "ft_tensix_compute_sfpu.h"
#include "ft_tensix_math_eltwise_sfpu.h"
#endif
#ifdef FT_TRISC_UNPACK
#include "ft_tensix_compute_sfpu.h"
#include "ft_tensix_compute_unpack.h"
#endif
#ifdef FT_TRISC_PACK
#include "ft_tensix_compute_pack.h"
#endif
#include "resident_queue_protocol.h"
#include "resident_attention_protocol.h"
FTQ_STATIC_ASSERT(FT_CB_ANY_32BIT_UNPACK == 0,
"a CB with 32-bit unpack needs an unpack-to-DST copy path; this kernel only "
"has the unpack-then-datacopy path");
#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_s ft_tensix_compile_time_arg(4)
#define cb_scratch ft_tensix_compile_time_arg(5)
#define cb_m_a ft_tensix_compile_time_arg(6)
#define cb_m_b ft_tensix_compile_time_arg(7)
#define cb_l ft_tensix_compile_time_arg(8)
#define cb_lsum ft_tensix_compile_time_arg(9)
#define cb_c ft_tensix_compile_time_arg(10)
#define cb_o ft_tensix_compile_time_arg(11)
#define cb_pv ft_tensix_compile_time_arg(12)
#define cb_out ft_tensix_compile_time_arg(13)
#define cb_qs ft_tensix_compile_time_arg(14)
#define wait_mode ft_tensix_compile_time_arg(15)
enum { dst_register = 0, dst_previous = 1 };
#ifdef FT_TRISC_MATH
#define MATMUL_HIGH_FIDELITY (FT_CB_MATH_FIDELITY != 0)
#define MATMUL_INNER_LOOPS (MATMUL_HIGH_FIDELITY ? (uint32_t)FT_CB_MATH_FIDELITY : 1u)
#define MATMUL_FIDELITY_INCREMENT (MATMUL_HIGH_FIDELITY ? 1u : 0u)
#endif
/* Encoded FPU words handed to the add/sub init (not issued here). */
#define ATTENTION_ADD_WORD \
FT_BH_ENC_ELWADD(0, 0, FT_BH_ELWADD_SRCB_NO_BCAST, FT_BH_ADDR_MOD_0, 0)
#define ATTENTION_SUB_WORD \
FT_BH_ENC_ELWSUB(0, 0, FT_BH_ELWADD_SRCB_NO_BCAST, FT_BH_ADDR_MOD_0, 0)
/* ===== Out-of-line waits (noinline keeps the UNPACK image inside its slab) =============== */
static __attribute__((noinline)) bool attention_wait_front(ftq_context *ctx, uint32_t cb,
uint32_t pages)
{
return ftq_cb_wait_front_stoppable(ctx, (int32_t)cb, (int32_t)pages, wait_mode);
}
static __attribute__((noinline)) bool attention_reserve_back(ftq_context *ctx, uint32_t cb,
uint32_t pages)
{
return ftq_cb_reserve_back_stoppable(ctx, (int32_t)cb, (int32_t)pages, wait_mode);
}
/* ===== Compute verbs (thread order inside each is fixed) ================================ */
/** Matmul init; in0 feeds SrcB and in1 feeds SrcA. */
static FT_BH_INLINE void attention_matmul_init(uint32_t in0, uint32_t in1, uint32_t transpose)
{
(void)in0;
(void)in1;
(void)transpose;
FT_TENSIX_MATH(ft_tensix_math_matmul_init(transpose, MATMUL_FIDELITY_INCREMENT,
MATMUL_INNER_LOOPS, MATMUL_HIGH_FIDELITY));
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_matmul_init(
transpose,
(uint32_t)ft_cb_unpack_tile_num_faces[in1] * ft_cb_unpack_tile_face_rows[in1] *
FT_TENSIX_FACE_C_DIM -
1u,
(uint32_t)ft_cb_unpack_tile_num_faces[in0] * ft_cb_unpack_tile_face_rows[in0] *
FT_TENSIX_FACE_C_DIM -
1u,
1u /* kt_dim */));
}
/** Accumulate in0[index0] . in1[index1] into DST. */
static FT_BH_INLINE void attention_matmul_apply(uint32_t in0, uint32_t in1, uint32_t index0,
uint32_t index1)
{
(void)in0;
(void)in1;
(void)index0;
(void)index1;
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_matmul(in0, in1, index0, index1, 1 /* ct_dim */));
FT_TENSIX_MATH(ft_tensix_math_matmul(dst_register));
}
/** Multiply init: MATH first, then UNPACK. */
static FT_BH_INLINE void attention_mul_init(uint32_t a)
{
(void)a;
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_mul_init(FT_CB_MATH_FIDELITY,
ft_cb_unpack_tile_face_rows[a]));
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_init(ft_cb_unpack_tile_face_rows[a],
ft_cb_unpack_num_face_rows[a],
ft_cb_unpack_num_face_cols[a]));
}
/** a[index_a] * b[index_b] into DST. */
static FT_BH_INLINE void attention_mul_apply(uint32_t a, uint32_t b, uint32_t index_a,
uint32_t index_b)
{
(void)a;
(void)b;
(void)index_a;
(void)index_b;
FT_TENSIX_UNPACK(ft_tensix_unpack_ab(a, b, index_a, index_b));
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_mul(
dst_register, (uint32_t)ft_cb_unpack_tile_num_faces[a], ft_cb_unpack_tile_face_rows[a]));
}
/** Add/sub init (lowest fidelity): MATH first, then UNPACK. */
static FT_BH_INLINE void attention_addsub_init(uint32_t binary_op_word, uint32_t a)
{
(void)binary_op_word;
(void)a;
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_addsub_init(
binary_op_word, (uint32_t)ft_cb_unpack_tile_num_faces[a],
(uint32_t)ft_cb_unpack_tile_face_rows[a]));
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_init(ft_cb_unpack_tile_face_rows[a],
ft_cb_unpack_num_face_rows[a],
ft_cb_unpack_num_face_cols[a]));
}
/** a[index_a] +/- b[index_b] into DST (op chosen by the init). */
static FT_BH_INLINE void attention_addsub_apply(uint32_t a, uint32_t b, uint32_t index_a,
uint32_t index_b)
{
(void)a;
(void)b;
(void)index_a;
(void)index_b;
FT_TENSIX_UNPACK(ft_tensix_unpack_ab(a, b, index_a, index_b));
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_addsub(dst_register));
}
/** Subtract-broadcast-column init (lowest fidelity). */
static FT_BH_INLINE void attention_sub_bcast_col_init(uint32_t a)
{
(void)a;
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_sub_broadcast_col_init(
(uint32_t)ft_cb_unpack_num_face_cols[a], (uint32_t)ft_cb_unpack_tile_face_rows[a]));
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_broadcast_col_init(ft_cb_unpack_tile_face_rows[a],
ft_cb_unpack_num_face_rows[a],
ft_cb_unpack_num_face_cols[a]));
}
/** a - column(b) into DST. */
static FT_BH_INLINE void attention_sub_bcast_col_apply(uint32_t a, uint32_t b)
{
(void)a;
(void)b;
FT_TENSIX_UNPACK(ft_tensix_unpack_ab(a, b, 0, 0));
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_sub_broadcast_col(
dst_register, (uint32_t)ft_cb_unpack_num_face_rows[a]));
}
/** Multiply-broadcast-column init. */
static FT_BH_INLINE void attention_mul_bcast_col_init(uint32_t a)
{
(void)a;
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_mul_broadcast_col_init(
FT_CB_MATH_FIDELITY, (uint32_t)ft_cb_unpack_tile_face_rows[a]));
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_broadcast_col_init(ft_cb_unpack_tile_face_rows[a],
ft_cb_unpack_num_face_rows[a],
ft_cb_unpack_num_face_cols[a]));
}
/** a * column(b) into DST. */
static FT_BH_INLINE void attention_mul_bcast_col_apply(uint32_t a, uint32_t b)
{
(void)a;
(void)b;
FT_TENSIX_UNPACK(ft_tensix_unpack_ab(a, b, 0, 0));
FT_TENSIX_MATH(ft_tensix_math_eltwise_binary_mul_broadcast_col(
dst_register, FT_CB_MATH_FIDELITY, (uint32_t)ft_cb_unpack_num_face_rows[a],
(uint32_t)ft_cb_unpack_num_face_cols[a], (uint32_t)ft_cb_unpack_tile_face_rows[a]));
}
/** Copy init: UNPACK first, then MATH. */
static FT_BH_INLINE void attention_copy_init(uint32_t cb)
{
(void)cb;
FT_TENSIX_UNPACK(ft_tensix_unpack_a_init(
ft_cb_unpack_tile_face_rows[cb],
(uint32_t)ft_cb_unpack_num_face_rows[cb] * ft_cb_unpack_num_face_cols[cb]));
FT_TENSIX_MATH(ft_tensix_math_datacopy_init((uint32_t)ft_cb_unpack_num_face_rows[cb] *
ft_cb_unpack_num_face_cols[cb]));
}
/** Copy the front tile of `cb` into DST slot `dst_index`. */
static FT_BH_INLINE void attention_copy_apply(uint32_t cb, uint32_t dst_index)
{
(void)cb;
(void)dst_index;
FT_TENSIX_UNPACK(ft_tensix_unpack_a(cb, 0));
FT_TENSIX_MATH(ft_tensix_math_datacopy(dst_index));
}
/* Undo a reduce; the PACK clear is the only thing that removes the reduce packer edge mask.
* One-line head and column-0 braces so check_tensix_reduce_exit_state.py sees the wrapper. */
static FT_BH_INLINE void attention_reduce_uninit(void)
{
FT_TENSIX_MATH(ft_tensix_math_reduce_uninit());
FT_TENSIX_PACK(ft_tensix_pack_reduce_mask_clear());
}
/* ===== DST commits ====================================================================== */
/** Hand DST to PACK and pack it as one page of `cb`. */
static void commit_dst_to(ftq_context *ctx, uint32_t cb)
{
FT_TENSIX_MATH(ft_tensix_math_dest_section_done());
(void)attention_reserve_back(ctx, cb, 1);
FT_TENSIX_PACK(ft_tensix_packer_wait_for_math_done());
FT_TENSIX_PACK(ft_tensix_pack_tile(dst_register, cb));
FT_TENSIX_PACK(ft_tensix_pack_dest_section_done(FT_CB_DST_ACCUM_FP32));
FT_TENSIX_PACK(ft_tensix_compute_cb_push_back(cb, 1));
}
/** As commit_dst_to, over the front tile just consumed: the CB is full, so pop first. */
static void commit_dst_in_place(ftq_context *ctx, uint32_t cb)
{
FT_TENSIX_MATH(ft_tensix_math_dest_section_done());
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb, 1));
(void)attention_reserve_back(ctx, cb, 1);
FT_TENSIX_PACK(ft_tensix_packer_wait_for_math_done());
FT_TENSIX_PACK(ft_tensix_pack_tile(dst_register, cb));
FT_TENSIX_PACK(ft_tensix_pack_dest_section_done(FT_CB_DST_ACCUM_FP32));
FT_TENSIX_PACK(ft_tensix_compute_cb_push_back(cb, 1));
}
/* ===== Steps ============================================================================ */
/** cb_qs <- cb_q * scale, dim_tiles tiles. */
static void scale_query(ftq_context *ctx, const fta_shape *shape)
{
uint32_t dt;
for (dt = 0; dt < shape->dim_tiles; dt++) {
attention_mul_init(cb_q);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_mul_apply(cb_q, cb_const, dt, kAttnCbConstScale);
commit_dst_to(ctx, cb_qs);
}
(void)attention_wait_front(ctx, cb_qs, shape->dim_tiles);
}
/** cb_s <- Qs . K_block^T (+ mask where the position tile is not entirely live). */
static void score_block(ftq_context *ctx, const fta_shape *shape, uint32_t block)
{
uint32_t pt;
uint32_t dt;
for (pt = 0; pt < shape->block_tiles; pt++) {
const uint32_t mask = fta_mask_for_tile(shape, block * shape->block_tiles + pt);
attention_matmul_init(cb_qs, cb_k, 1u /* transpose B */);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
for (dt = 0; dt < shape->dim_tiles; dt++)
attention_matmul_apply(cb_qs, cb_k, dt, pt * shape->dim_tiles + dt);
if (mask >= (uint32_t)kAttnCbConstPages) {
commit_dst_to(ctx, cb_s);
continue;
}
commit_dst_to(ctx, cb_scratch);
(void)attention_wait_front(ctx, cb_scratch, 1);
attention_addsub_init(ATTENTION_ADD_WORD, cb_scratch);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_addsub_apply(cb_scratch, cb_const, 0, mask);
commit_dst_to(ctx, cb_s);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_scratch, 1));
}
}
/** destination <- max(rowmax(S), previous if seed). */
static void running_max(ftq_context *ctx, const fta_shape *shape, uint32_t destination,
uint32_t previous, bool seed)
{
uint32_t pt;
(void)previous;
(void)seed;
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_reduce_row_max_init(
(uint32_t)ft_cb_unpack_tile_num_faces[cb_s], (uint32_t)ft_cb_unpack_tile_face_rows[cb_s]));
FT_TENSIX_MATH(
ft_tensix_math_reduce_row_max_init(FT_CB_MATH_FIDELITY != 0 ? 1u : 0u, FT_CB_MATH_FIDELITY));
FT_TENSIX_PACK(
ft_tensix_pack_reduce_row_mask_config((uint32_t)ft_cb_pack_tile_face_rows[destination]));
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
for (pt = 0; pt < shape->block_tiles; pt++) {
FT_TENSIX_MATH(ft_tensix_math_reduce_row_max(dst_register,
(uint32_t)ft_cb_unpack_num_face_rows[cb_s],
(uint32_t)ft_cb_unpack_num_face_cols[cb_s]));
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_reduce_row_max(cb_s, cb_const, pt,
kAttnCbConstScaler));
}
/* The uninit must precede the copy below. */
attention_reduce_uninit();
if (seed) {
attention_copy_init(previous);
attention_copy_apply(previous, dst_previous);
FT_TENSIX_MATH(ft_tensix_math_binary_max_tile_init());
FT_TENSIX_MATH(ft_tensix_math_binary_max_tile(dst_register, dst_previous, dst_register));
}
commit_dst_to(ctx, destination);
(void)attention_wait_front(ctx, destination, 1);
}
/** cb_s <- exp(S - column(maximum)), in place. */
static void exponentiate_block(ftq_context *ctx, const fta_shape *shape, uint32_t maximum)
{
uint32_t pt;
for (pt = 0; pt < shape->block_tiles; pt++) {
attention_sub_bcast_col_init(cb_s);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_sub_bcast_col_apply(cb_s, maximum);
FT_TENSIX_MATH(ft_tensix_math_exp_tile_init());
FT_TENSIX_MATH(ft_tensix_math_exp_tile(dst_register));
commit_dst_in_place(ctx, cb_s);
}
(void)attention_wait_front(ctx, cb_s, shape->block_tiles);
}
/** cb_lsum <- rowsum(P). */
static void block_row_sum(ftq_context *ctx, const fta_shape *shape)
{
uint32_t pt;
FT_TENSIX_UNPACK(
ft_tensix_unpack_ab_reduce_row_init((uint32_t)ft_cb_unpack_tile_num_faces[cb_s]));
FT_TENSIX_MATH(
ft_tensix_math_reduce_row_init(FT_CB_MATH_FIDELITY != 0 ? 1u : 0u, FT_CB_MATH_FIDELITY));
FT_TENSIX_PACK(
ft_tensix_pack_reduce_row_mask_config((uint32_t)ft_cb_pack_tile_face_rows[cb_lsum]));
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
for (pt = 0; pt < shape->block_tiles; pt++) {
FT_TENSIX_MATH(ft_tensix_math_reduce_row(dst_register));
FT_TENSIX_UNPACK(ft_tensix_unpack_ab_reduce_row(cb_s, cb_const, pt, kAttnCbConstScaler));
}
attention_reduce_uninit();
commit_dst_to(ctx, cb_lsum);
(void)attention_wait_front(ctx, cb_lsum, 1);
}
/** cb_pv <- P . V_block, dim_tiles tiles. */
static void multiply_by_values(ftq_context *ctx, const fta_shape *shape)
{
uint32_t dt;
uint32_t pt;
for (dt = 0; dt < shape->dim_tiles; dt++) {
attention_matmul_init(cb_s, cb_v, 0u);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
for (pt = 0; pt < shape->block_tiles; pt++)
attention_matmul_apply(cb_s, cb_v, pt, pt * shape->dim_tiles + dt);
commit_dst_to(ctx, cb_pv);
}
(void)attention_wait_front(ctx, cb_pv, shape->dim_tiles);
}
/** First block: O <- P.V, l <- s. */
static void seed_accumulators(ftq_context *ctx, const fta_shape *shape)
{
uint32_t dt;
for (dt = 0; dt < shape->dim_tiles; dt++) {
attention_copy_init(cb_pv);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_copy_apply(cb_pv, dst_register);
commit_dst_to(ctx, cb_o);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_pv, 1));
}
attention_copy_init(cb_lsum);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_copy_apply(cb_lsum, dst_register);
commit_dst_to(ctx, cb_l);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_lsum, 1));
(void)attention_wait_front(ctx, cb_o, shape->dim_tiles);
(void)attention_wait_front(ctx, cb_l, 1);
}
/** Later blocks: c = exp(m_old - m_new); l = c*l + s; O = c*O + P.V. */
static void merge_block(ftq_context *ctx, const fta_shape *shape, uint32_t maximum,
uint32_t previous)
{
uint32_t dt;
attention_addsub_init(ATTENTION_SUB_WORD, previous);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_addsub_apply(previous, maximum, 0, 0);
FT_TENSIX_MATH(ft_tensix_math_exp_tile_init());
FT_TENSIX_MATH(ft_tensix_math_exp_tile(dst_register));
commit_dst_to(ctx, cb_c);
(void)attention_wait_front(ctx, cb_c, 1);
attention_mul_init(cb_l);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_mul_apply(cb_l, cb_c, 0, 0);
commit_dst_in_place(ctx, cb_l);
(void)attention_wait_front(ctx, cb_l, 1);
attention_addsub_init(ATTENTION_ADD_WORD, cb_l);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_addsub_apply(cb_l, cb_lsum, 0, 0);
commit_dst_in_place(ctx, cb_l);
(void)attention_wait_front(ctx, cb_l, 1);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_lsum, 1));
for (dt = 0; dt < shape->dim_tiles; dt++) {
attention_mul_bcast_col_init(cb_o);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_mul_bcast_col_apply(cb_o, cb_c);
commit_dst_in_place(ctx, cb_o);
}
(void)attention_wait_front(ctx, cb_o, shape->dim_tiles);
for (dt = 0; dt < shape->dim_tiles; dt++) {
attention_addsub_init(ATTENTION_ADD_WORD, cb_o);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_addsub_apply(cb_o, cb_pv, 0, 0);
/* Written out rather than commit_dst_in_place: this commit pops two CBs. */
FT_TENSIX_MATH(ft_tensix_math_dest_section_done());
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_o, 1));
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_pv, 1));
(void)attention_reserve_back(ctx, cb_o, 1);
FT_TENSIX_PACK(ft_tensix_packer_wait_for_math_done());
FT_TENSIX_PACK(ft_tensix_pack_tile(dst_register, cb_o));
FT_TENSIX_PACK(ft_tensix_pack_dest_section_done(FT_CB_DST_ACCUM_FP32));
FT_TENSIX_PACK(ft_tensix_compute_cb_push_back(cb_o, 1));
}
(void)attention_wait_front(ctx, cb_o, shape->dim_tiles);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_c, 1));
}
/** out <- O * (1/l), then release every per-descriptor CB. */
static void normalise_and_emit(ftq_context *ctx, const fta_shape *shape, uint32_t maximum)
{
uint32_t dt;
(void)maximum;
attention_copy_init(cb_l);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_copy_apply(cb_l, dst_register);
FT_TENSIX_MATH(ft_tensix_math_recip_tile_init());
FT_TENSIX_MATH(ft_tensix_math_recip_tile(dst_register));
commit_dst_to(ctx, cb_c);
(void)attention_wait_front(ctx, cb_c, 1);
for (dt = 0; dt < shape->dim_tiles; dt++) {
attention_mul_bcast_col_init(cb_o);
FT_TENSIX_MATH(ft_tensix_math_wait_for_dest_available());
attention_mul_bcast_col_apply(cb_o, cb_c);
FT_TENSIX_MATH(ft_tensix_math_dest_section_done());
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_o, 1));
(void)attention_reserve_back(ctx, cb_out, 1);
FT_TENSIX_PACK(ft_tensix_packer_wait_for_math_done());
FT_TENSIX_PACK(ft_tensix_pack_tile(dst_register, cb_out));
FT_TENSIX_PACK(ft_tensix_pack_dest_section_done(FT_CB_DST_ACCUM_FP32));
FT_TENSIX_PACK(ft_tensix_compute_cb_push_back(cb_out, 1));
}
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_c, 1));
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_l, 1));
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(maximum, 1));
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_qs, shape->dim_tiles));
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_q, shape->dim_tiles));
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_const, kAttnCbConstPages));
}
/** Run one ATTENTION descriptor; the two max CBs ping-pong between blocks. */
static bool compute_one_attention(ftq_context *ctx, volatile FTQ_L1_PTR uint32_t *parameters)
{
const fta_shape shape = fta_decode(parameters);
const uint32_t depth = shape.block_tiles * shape.dim_tiles;
uint32_t maximum = cb_m_a;
uint32_t previous = cb_m_b;
uint32_t block;
if (!fta_shape_valid(&shape)) {
ftq_abort(ctx, kErrBadParameters);
return false;
}
(void)attention_wait_front(ctx, cb_q, shape.dim_tiles);
(void)attention_wait_front(ctx, cb_const, kAttnCbConstPages);
scale_query(ctx, &shape);
for (block = 0; block < shape.blocks; block++) {
const uint32_t swap = maximum;
(void)attention_wait_front(ctx, cb_k, depth);
score_block(ctx, &shape, block);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_k, depth));
(void)attention_wait_front(ctx, cb_s, shape.block_tiles);
running_max(ctx, &shape, maximum, previous, block != 0u);
exponentiate_block(ctx, &shape, maximum);
block_row_sum(ctx, &shape);
(void)attention_wait_front(ctx, cb_v, depth);
multiply_by_values(ctx, &shape);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_v, depth));
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(cb_s, shape.block_tiles));
if (block == 0u) {
seed_accumulators(ctx, &shape);
} else {
merge_block(ctx, &shape, maximum, previous);
FT_TENSIX_UNPACK(ft_tensix_compute_cb_pop_front(previous, 1));
}
maximum = previous;
previous = swap;
}
/* After the loop the final max is in `previous`. */
normalise_and_emit(ctx, &shape, previous);
return !ctx->aborting;
}
void kernel_main(void)
{
ftq_context ctx;
bool running = true;
ftq_context_init(&ctx);
ftq_status_set(&ctx, kStateBoot);
if (!ftq_check_control(&ctx)) {
ftq_status_set(&ctx, kStateExit);
return;
}
/* Once, with the source order reversed for the dominant matmul (in0 -> SrcB, in1 ->
* SrcA): the unpacker pair is cb_k as A, cb_qs as B. */
FT_TENSIX_UNPACK(ft_tensix_unpack_hw_configure_ab(
ft_cb_unpack_src_format[cb_k], ft_cb_unpack_src_format[cb_qs],
ft_cb_unpack_dst_format[cb_k], ft_cb_unpack_dst_format[cb_qs],
ft_cb_unpack_tile_face_rows[cb_k], ft_cb_unpack_tile_face_rows[cb_qs],
ft_cb_unpack_tile_num_faces[cb_k], ft_cb_unpack_tile_num_faces[cb_qs],
ft_tensix_cb(cb_k)->fifo_page_size, ft_tensix_cb(cb_qs)->fifo_page_size));
FT_TENSIX_MATH(ft_tensix_math_pack_sync_init());
FT_TENSIX_MATH(ft_tensix_math_hw_configure(ft_cb_unpack_dst_format[cb_k],
ft_cb_unpack_dst_format[cb_qs],
FT_CB_DST_ACCUM_FP32));
FT_TENSIX_PACK(ft_tensix_pack_hw_configure(
ft_cb_pack_src_format[cb_s], ft_cb_pack_dst_format[cb_s],
ft_tensix_cb(cb_s)->fifo_page_size, ft_cb_pack_tile_num_faces[cb_s],
ft_cb_pack_partial_face[cb_s], 0 /* relu */, FT_CB_DST_ACCUM_FP32));
FT_TENSIX_PACK(ft_tensix_pack_init(ft_cb_pack_src_format[cb_s], ft_cb_pack_tile_face_rows[cb_s],
ft_cb_pack_tile_num_faces[cb_s], 1));
FT_TENSIX_PACK(ft_tensix_pack_dest_init());
attention_matmul_init(cb_qs, cb_k, 0u);
ctx.cursor = ctx.control[kFieldCompletedCount];
while (running) {
volatile FTQ_L1_PTR uint32_t *descriptor;
uint32_t opcode;
if (!ftq_wait_for_work(&ctx, kFieldHead))
break;
ctx.aborting = false;
ftq_status_set(&ctx, kStateExec);
descriptor = ftq_descriptor(&ctx);
opcode = descriptor[kDescOpcode];
ftq_status_opcode(&ctx, opcode);
if (opcode == kOpAttention) {
ftq_status_set(&ctx, kStateCbWaitFront);
(void)compute_one_attention(&ctx, ftq_parameters(descriptor));
ftq_status_set(&ctx, kStateExec);
if (ctx.aborting)
running = false;
} else if (opcode == kOpStop) {
running = false;
}
if (!ctx.aborting)
ctx.cursor++;
}
ftq_status_set(&ctx, kStateExit);
}
Reference in New Issue View Git Blame Copy Permalink