code/hamer/kernels/hattn2/writer.cpp · changh95/hamer-p150 at main

hamer-p150 / code /hamer /kernels /hattn2 /writer.cpp

Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms

58ba6df verified about 1 month ago

1.79 kB

// 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);
}

	// 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);
	}