code/hamer/kernels/hattn/reader.cpp · changh95/hamer-p150 at main

hamer-p150 / code /hamer /kernels /hattn /reader.cpp

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

58ba6df verified about 1 month ago

2.7 kB

// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn reader (NCRISC), round 6: 1-token cross-attention of the MANO head, one core per head h.
//   q: row 0 of tiles 2h, 2h+1 of the row-form query [1, 1, 32, H*64] (bf16; already scaled by 1/sqrt(64) on the host)
//   K / V: straight from the batched K/V projection output kv [1, 1, MT*32, L*1024] (TILE, interleaved): layer `l`,
//          head h: K tile (j, d) = page j * NC + l*32 + 2h + d, V tile = page j * NC + l*32 + 16 + 2h + d  (j < MT)
// compile-time: MT, NC (kv tile columns), TensorAccessorArgs(q), TensorAccessorArgs(kv)
// runtime: q_addr, kv_addr, h, l
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    constexpr uint32_t MT = get_compile_time_arg_val(0);
    constexpr uint32_t NC = get_compile_time_arg_val(1);
    constexpr auto q_args = TensorAccessorArgs<2>();
    constexpr auto kv_args = TensorAccessorArgs<q_args.next_compile_time_args_offset()>();
    constexpr uint32_t c_q = 0, c_k = 1, c_v = 2, c_scal = 3;
    const uint32_t q_addr = get_arg_val<uint32_t>(0), kv_addr = get_arg_val<uint32_t>(1);
    const uint32_t h = get_arg_val<uint32_t>(2), l = get_arg_val<uint32_t>(3);
    constexpr uint32_t page = 2048, FB = 512;  // bf16 tile / face bytes
    const auto q = TensorAccessor(q_args, q_addr, page);
    const auto kv = TensorAccessor(kv_args, kv_addr, page);
    cb_reserve_back(c_q, 2);
    const uint32_t qd = get_write_ptr(c_q);
    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 * page, 32);
        noc_async_read(s + FB, qd + d * page + FB, 32);
    }
    cb_reserve_back(c_k, 2 * MT);
    cb_reserve_back(c_v, 2 * MT);
    uint32_t kd = get_write_ptr(c_k), vd = get_write_ptr(c_v);
    for (uint32_t j = 0; j < MT; ++j) {
        for (uint32_t d = 0; d < 2; ++d) {
            noc_async_read(kv.get_noc_addr(j * NC + l * 32 + 2 * h + d), kd, page);
            noc_async_read(kv.get_noc_addr(j * NC + l * 32 + 16 + 2 * h + d), vd, page);
            kd += page;
            vd += page;
        }
    }
    // reduce scaler: bf16 1.0 in row 0 of every face
    cb_reserve_back(c_scal, 1);
    volatile tt_l1_ptr uint32_t* sc = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(c_scal));
    for (uint32_t i = 0; i < 512; ++i) {
        sc[i] = 0;
    }
    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_scal, 1);
    noc_async_read_barrier();
    cb_push_back(c_q, 2);
    cb_push_back(c_k, 2 * MT);
    cb_push_back(c_v, 2 * MT);
}

	// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
	// SPDX-License-Identifier: Apache-2.0
	// hattn reader (NCRISC), round 6: 1-token cross-attention of the MANO head, one core per head h.
	// q: row 0 of tiles 2h, 2h+1 of the row-form query [1, 1, 32, H*64] (bf16; already scaled by 1/sqrt(64) on the host)
	// K / V: straight from the batched K/V projection output kv [1, 1, MT32, L1024] (TILE, interleaved): layer `l`,
	// head h: K tile (j, d) = page j * NC + l32 + 2h + d, V tile = page j NC + l*32 + 16 + 2h + d (j < MT)
	// compile-time: MT, NC (kv tile columns), TensorAccessorArgs(q), TensorAccessorArgs(kv)
	// runtime: q_addr, kv_addr, h, l
	#include <stdint.h>
	#include "api/dataflow/dataflow_api.h"

	void kernel_main() {
	constexpr uint32_t MT = get_compile_time_arg_val(0);
	constexpr uint32_t NC = get_compile_time_arg_val(1);
	constexpr auto q_args = TensorAccessorArgs<2>();
	constexpr auto kv_args = TensorAccessorArgs<q_args.next_compile_time_args_offset()>();
	constexpr uint32_t c_q = 0, c_k = 1, c_v = 2, c_scal = 3;
	const uint32_t q_addr = get_arg_val<uint32_t>(0), kv_addr = get_arg_val<uint32_t>(1);
	const uint32_t h = get_arg_val<uint32_t>(2), l = get_arg_val<uint32_t>(3);
	constexpr uint32_t page = 2048, FB = 512; // bf16 tile / face bytes
	const auto q = TensorAccessor(q_args, q_addr, page);
	const auto kv = TensorAccessor(kv_args, kv_addr, page);
	cb_reserve_back(c_q, 2);
	const uint32_t qd = get_write_ptr(c_q);
	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 * page, 32);
	noc_async_read(s + FB, qd + d * page + FB, 32);
	}
	cb_reserve_back(c_k, 2 * MT);
	cb_reserve_back(c_v, 2 * MT);
	uint32_t kd = get_write_ptr(c_k), vd = get_write_ptr(c_v);
	for (uint32_t j = 0; j < MT; ++j) {
	for (uint32_t d = 0; d < 2; ++d) {
	noc_async_read(kv.get_noc_addr(j * NC + l * 32 + 2 * h + d), kd, page);
	noc_async_read(kv.get_noc_addr(j * NC + l * 32 + 16 + 2 * h + d), vd, page);
	kd += page;
	vd += page;
	}
	}
	// reduce scaler: bf16 1.0 in row 0 of every face
	cb_reserve_back(c_scal, 1);
	volatile tt_l1_ptr uint32_t* sc = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(c_scal));
	for (uint32_t i = 0; i < 512; ++i) {
	sc[i] = 0;
	}
	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_scal, 1);
	noc_async_read_barrier();
	cb_push_back(c_q, 2);
	cb_push_back(c_k, 2 * MT);
	cb_push_back(c_v, 2 * MT);
	}