File size: 1,792 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 | // SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn2 writer (BRISC), round 15: the two [1, 32] tiny context tiles of head h -> row 0 of tiles 2h, 2h+1 of the
// head-split output [1, H, 32, 64] (bf16; read by the out-projection GEMV as the row-form [1, 1, 32, H*64] vector).
// It first reads head h's 12 V tiles (needed last) on its own NoC, in parallel with the reader's q / K reads.
// compile-time: TensorAccessorArgs(out), TensorAccessorArgs(v). runtime: out_addr, h, v_addr, v_nc, v_base (see reader)
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
void kernel_main() {
constexpr auto o_args = TensorAccessorArgs<0>();
constexpr auto v_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
constexpr uint32_t c_v = 2, c_out = 16, page = 2048, FB = 512, MT = 6;
const uint32_t o_addr = get_arg_val<uint32_t>(0), h = get_arg_val<uint32_t>(1);
const auto out = TensorAccessor(o_args, o_addr, page);
{
const auto v = TensorAccessor(v_args, get_arg_val<uint32_t>(2), page);
cb_reserve_back(c_v, 2 * MT);
uint32_t vd = get_write_ptr(c_v);
for (uint32_t t = 0; t < 2 * MT; ++t) {
noc_async_read(v.get_noc_addr((t / 2) * get_arg_val<uint32_t>(3) + get_arg_val<uint32_t>(4) + (t % 2)), vd, page);
vd += page;
}
noc_async_read_barrier();
cb_push_back(c_v, 2 * MT);
}
cb_wait_front(c_out, 2);
const uint32_t src = get_read_ptr(c_out);
for (uint32_t d = 0; d < 2; ++d) {
const uint64_t dst = out.get_noc_addr(2 * h + d);
noc_async_write(src + d * 64, dst, 32);
noc_async_write(src + d * 64 + 32, dst + FB, 32);
}
noc_async_write_barrier();
cb_pop_front(c_out, 2);
}
|