Hugging Face's logo

changh95
/
hamer-p150

Keypoint Detection
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-container
tenstorrent
ttnn
tt-metal
tt-nn
hand-pose
mano
3d-mesh
hamer
tt-model-catalog
Model card Files
xet
 Community
1
hamer-p150 / code /hamer /kernels /hattn /compute.cpp
changh95's picture
changh95
Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms
58ba6df verified
Raw History Blame Contribute Delete
5.39 kB
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn compute, round 6 (see reader.cpp): row 0 of every tile is the query token.
// s_j = sum_d q_d @ K_{j,d}^T (MT score tiles, fp32)
// m = rowmax(s); e_j = exp(s_j - m) (bf16, the PV operand)
// r = 1 / rowsum(e)
// o_d = (sum_j e_j @ V_{j,d}) * r -> 2 bf16 tiles
// Round 6 status: correct (PCC 0.999997 vs fp64) but NOT wired -- 10.1 us kernel / ~14 us per call vs ~10.7 us for
// head split + library SDPA + head concat: the accurate exp of 6 score tiles costs 3.6 us even in row mode.
// compile-time: MT
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/reduce.h"
#include "api/compute/bcast.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/exp.h"
#include "api/compute/eltwise_unary/recip.h"
#include "api/compute/compute_kernel_hw_startup.h"
#ifdef HA_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define PZ(n) DeviceZoneScopedN(n)
#else
#define PZ(n)
#endif
void kernel_main() {
constexpr uint32_t MT = get_compile_time_arg_val(0);
constexpr uint32_t c_q = 0, c_k = 1, c_v = 2, c_scal = 3, c_out = 16;
constexpr uint32_t c_s = 24, c_m = 25, c_e = 26, c_r = 27, c_o = 28;
compute_kernel_hw_startup(c_q, c_k, c_s);
{PZ("WAIT");
cb_wait_front(c_q, 2);
cb_wait_front(c_k, 2 * MT);
cb_wait_front(c_v, 2 * MT);
cb_wait_front(c_scal, 1);}
// scores
{PZ("S");
reconfig_data_format(c_k, c_q); // srcA = in1 = K, srcB = in0 = q
pack_reconfig_data_format(c_s);
matmul_init(c_q, c_k, 1);
cb_reserve_back(c_s, MT);
for (uint32_t j0 = 0; j0 < MT; j0 += 4) {
const uint32_t n = MT - j0 < 4 ? MT - j0 : 4;
tile_regs_acquire();
for (uint32_t j = 0; j < n; ++j) {
for (uint32_t d = 0; d < 2; ++d) {
matmul_tiles(c_q, c_k, d, (j0 + j) * 2 + d, j);
}
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t j = 0; j < n; ++j) {
pack_tile(j, c_s);
}
tile_regs_release();
}
cb_push_back(c_s, MT);
cb_wait_front(c_s, MT);}
// row max
{PZ("MAX");
reconfig_data_format(c_s, c_scal);
pack_reconfig_data_format(c_m);
reduce_init<PoolType::MAX, ReduceDim::REDUCE_ROW>(c_s, c_scal, c_m);
cb_reserve_back(c_m, 1);
tile_regs_acquire();
for (uint32_t j = 0; j < MT; ++j) {
reduce_tile<PoolType::MAX, ReduceDim::REDUCE_ROW>(c_s, c_scal, j, 0, 0);
}
reduce_uninit();
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_m);
tile_regs_release();
cb_push_back(c_m, 1);
cb_wait_front(c_m, 1);}
// e = exp(s - m)
{PZ("EXP");
reconfig_data_format(c_s, c_m);
pack_reconfig_data_format(c_e);
cb_reserve_back(c_e, MT);
for (uint32_t j0 = 0; j0 < MT; j0 += 4) {
const uint32_t n = MT - j0 < 4 ? MT - j0 : 4;
sub_bcast_cols_init(c_s, c_m);
tile_regs_acquire();
for (uint32_t j = 0; j < n; ++j) {
sub_tiles_bcast_cols(c_s, c_m, j0 + j, 0, j);
}
exp_tile_init<false>();
for (uint32_t j = 0; j < n; ++j) {
exp_tile<false>(j, VectorMode::R); // only row 0 (the query token) matters
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t j = 0; j < n; ++j) {
pack_tile(j, c_e);
}
tile_regs_release();
}
cb_push_back(c_e, MT);
cb_pop_front(c_m, 1);
cb_pop_front(c_s, MT);
cb_wait_front(c_e, MT);}
// r = 1 / rowsum(e)
{PZ("SUM"); // REDUCE_ROW SUM: scaler -> SrcA, data -> SrcB
reconfig_data_format(c_scal, c_e);
pack_reconfig_data_format(c_r);
reduce_init<PoolType::SUM, ReduceDim::REDUCE_ROW>(c_e, c_scal, c_r);
cb_reserve_back(c_r, 1);
tile_regs_acquire();
for (uint32_t j = 0; j < MT; ++j) {
reduce_tile<PoolType::SUM, ReduceDim::REDUCE_ROW>(c_e, c_scal, j, 0, 0);
}
reduce_uninit();
recip_tile_init();
recip_tile(0);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_r);
tile_regs_release();
cb_push_back(c_r, 1);}
// o = e @ V
{PZ("PV");
reconfig_data_format(c_v, c_e);
pack_reconfig_data_format(c_o);
matmul_init(c_e, c_v, 0);
cb_reserve_back(c_o, 2);
tile_regs_acquire();
for (uint32_t d = 0; d < 2; ++d) {
for (uint32_t j = 0; j < MT; ++j) {
matmul_tiles(c_e, c_v, j, j * 2 + d, d);
}
}
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_o);
pack_tile(1, c_o);
tile_regs_release();
cb_push_back(c_o, 2);}
// * r
{PZ("NORM");
cb_wait_front(c_o, 2);
cb_wait_front(c_r, 1);
reconfig_data_format(c_o, c_r);
pack_reconfig_data_format(c_out);
mul_bcast_cols_init(c_o, c_r);
cb_reserve_back(c_out, 2);
tile_regs_acquire();
mul_tiles_bcast_cols(c_o, c_r, 0, 0, 0);
mul_tiles_bcast_cols(c_o, c_r, 1, 0, 1);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_out);
pack_tile(1, c_out);
tile_regs_release();
cb_push_back(c_out, 2);}
cb_pop_front(c_o, 2);
cb_pop_front(c_r, 1);
cb_pop_front(c_e, MT);
cb_pop_front(c_q, 2);
cb_pop_front(c_k, 2 * MT);
cb_pop_front(c_v, 2 * MT);
cb_pop_front(c_scal, 1);
}