File size: 4,399 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | // SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn2 reader (NCRISC), round 15: MANO-head 1-query cross-attention, one core per head h (D5).
// q: row 0 of tiles 2h, 2h+1 of the head-split query [1, H, 32, 64] (bf16, only row 0 meaningful) -> two [1, 32]
// tiny tiles (c_qt, contiguous 64 B each: face-0 row 0 then face-1 row 0)
// K: the 12 tiles (key tile j, dim tile d) of head h -> c_k, page j * NC + base + d (runtime NC / base: the head-split
// [1, H, 192, 64] K is NC = 2, base = 12 h; round 15: straight from the K/V projection output [1, 1, 192, L*1024],
// NC = its tile columns, base = 32 l + 2 h -- no head split op). V: the writer, the same way.
// the dense score tile (c_sd, fp32 32x32, which the compute fills with the 192 scores as six [1, 32] tiny tiles at its
// start) is first filled with a large negative constant (a 4 KB fp32 tile in L1, read once per launch), so the reductions
// over the whole tile see only the scores
// reduce scaler: bf16 1.0 in row 0 of every face (c_scal); a bf16 tile of ones (c_one, the sum-of-e matmul operand)
// compile-time: TensorAccessorArgs(q), (k), (v), (consts: tile 0 = fp32 pad, 4 KB; tile 1 = bf16 ones, 2 KB)
// runtime: q_addr, k_addr, v_addr, consts_addr, h, k_nc, k_base
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
void kernel_main() {
constexpr auto q_args = TensorAccessorArgs<0>();
constexpr auto k_args = TensorAccessorArgs<q_args.next_compile_time_args_offset()>();
constexpr auto v_args = TensorAccessorArgs<k_args.next_compile_time_args_offset()>();
constexpr auto n_args = TensorAccessorArgs<v_args.next_compile_time_args_offset()>();
constexpr uint32_t c_qt = 0, c_k = 1, c_v = 2, c_scal = 3, c_st = 4, c_one = 6;
constexpr uint32_t MT = 6; // key tiles (192 keys)
constexpr uint32_t page = 2048, FB = 512; // bf16 tile / face bytes
const uint32_t q_addr = get_arg_val<uint32_t>(0), k_addr = get_arg_val<uint32_t>(1);
const uint32_t v_addr = get_arg_val<uint32_t>(2), n_addr = get_arg_val<uint32_t>(3), h = get_arg_val<uint32_t>(4);
const uint32_t k_nc = get_arg_val<uint32_t>(5), k_base = get_arg_val<uint32_t>(6);
const auto q = TensorAccessor(q_args, q_addr, page);
const auto k = TensorAccessor(k_args, k_addr, page);
// consts: one 6 KB row-major uint32 L1 buffer (page = 6 KB): [0, 4096) the fp32 pad tile, [4096, 6144) the bf16 ones
const auto cst = TensorAccessor(n_args, n_addr, 6144);
const uint64_t cs = cst.get_noc_addr(0);
// pad fill of the dense score tile (its memory = c_st's), before anything is pushed to the compute
noc_async_read(cs, get_write_ptr(c_st), 4096);
cb_reserve_back(c_one, 1);
noc_async_read(cs + 4096, get_write_ptr(c_one), 2048);
cb_reserve_back(c_k, 2 * MT);
uint32_t kd = get_write_ptr(c_k);
for (uint32_t t = 0; t < 2 * MT; ++t) {
noc_async_read(k.get_noc_addr((t / 2) * k_nc + k_base + (t % 2)), kd, page);
kd += page;
}
cb_reserve_back(c_scal, 1);
{
const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
for (uint32_t off = 0; off < 2048; off += MEM_ZEROS_SIZE) {
noc_async_read(zeros, get_write_ptr(c_scal) + off, MEM_ZEROS_SIZE);
}
}
#ifdef HA2_WAIT
// round 15 (fused into the q GEMV program): the two q GEMV cores of head h write its query columns, then each
// increments this core's semaphore 0 -- wait for both before reading q (K / pad / ones / scaler reads are in flight)
{
volatile tt_l1_ptr uint32_t* sem = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(0));
noc_semaphore_wait(sem, 2);
}
#endif
cb_reserve_back(c_qt, 2);
const uint32_t qd = get_write_ptr(c_qt);
for (uint32_t d = 0; d < 2; ++d) {
const uint64_t s = q.get_noc_addr(2 * h + d);
noc_async_read(s, qd + d * 64, 32);
noc_async_read(s + FB, qd + d * 64 + 32, 32);
}
noc_async_read_barrier();
volatile tt_l1_ptr uint32_t* sc = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(c_scal));
for (uint32_t f = 0; f < 4; ++f) {
for (uint32_t j = 0; j < 8; ++j) {
sc[f * 128 + j] = 0x3F803F80u;
}
}
cb_push_back(c_qt, 2);
cb_push_back(c_k, 2 * MT);
cb_push_back(c_scal, 1);
cb_push_back(c_one, 1);
}
|