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 /hattn2 /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.67 kB
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn2 compute, round 15 (D5): 1-query attention of one head with a DENSE score layout (round 6's hattn exp'd six
// row-form score tiles; here the 192 scores are the first 192 elements = face 0, rows 0..11 of ONE fp32 tile):
// s_j = q @ K_j^T six [1, 32] tiny fp32 tiles (q: two [1, 32] tiny tiles, K tiles transposed in the matmul),
// packed back to back into c_st = the start of the dense tile c_sd (the reader pre-filled the
// rest of c_sd with a large negative pad, so the whole-tile max sees only the scores)
// m = max(c_sd) (REDUCE_SCALAR)
// e = exp((s - m) * scale) on face 0 only (VectorMode::None), packed fp32 -> c_e; its first 192 elements are six
// [1, 32] tiny fp32 tiles (alias c_et); the other faces are never read
// o_d = sum_j e_j @ V_{j,d}, z = sum_j e_j @ ONES (one DST acquire; z = sum(e) in every column)
// o_d *= 1 / z (SFPU in DST: recip of the tiny z tile, then two tiny multiplies) -> two [1, 32] bf16 -> c_out
// compile-time: scale bits (fp32; exact as bf16 for 1/8)
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/reduce.h"
#include "api/compute/bcast.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_binary_sfpu.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"
#include "api/compute/reconfig_data_format.h"
#ifdef HA2_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define PZ(n) DeviceZoneScopedN(n)
#else
#define PZ(n)
#endif
void kernel_main() {
constexpr uint32_t SCALE_BITS = get_compile_time_arg_val(0);
constexpr uint16_t SCALE_BF16 = (uint16_t)(SCALE_BITS >> 16);
static_assert((SCALE_BITS & 0xFFFFu) == 0, "the exp scale must be exact in bf16");
constexpr uint32_t MT = 6;
constexpr uint32_t c_qt = 0, c_k = 1, c_v = 2, c_scal = 3, c_st = 4, c_sd = 5, c_one = 6, c_et = 7, c_out = 16;
constexpr uint32_t c_m = 24, c_e = 25;
compute_kernel_hw_startup(c_qt, c_k, c_st);
{PZ("WAIT");
cb_wait_front(c_qt, 2);
cb_wait_front(c_k, 2 * MT);} // (V arrives separately, waited for before the PV phase)
// scores: six tiny tiles
{PZ("S");
reconfig_full_operand(c_k, c_qt); // srcA = in1 = K, srcB = in0 = q
pack_reconfig_data_format<true>(c_st);
matmul_init(c_qt, c_k, 1);
cb_reserve_back(c_st, MT);
for (uint32_t j = 0; j < MT; ++j) {
tile_regs_acquire();
matmul_tiles(c_qt, c_k, 0, j * 2, 0);
matmul_tiles(c_qt, c_k, 1, j * 2 + 1, 0);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_st);
tile_regs_release();
}
cb_push_back(c_st, MT);
// the same memory as one dense fp32 tile
cb_reserve_back(c_sd, 1);
cb_push_back(c_sd, 1);
}
cb_wait_front(c_scal, 1);
// m = max over the dense tile
{PZ("MAX");
reconfig_full_operand(c_sd, c_scal); // REDUCE_SCALAR: data -> SrcA, scaler -> SrcB
pack_reconfig_data_format<true>(c_m);
reduce_init<PoolType::MAX, ReduceDim::REDUCE_SCALAR>(c_sd, c_scal, c_m);
cb_wait_front(c_sd, 1);
cb_reserve_back(c_m, 1);
tile_regs_acquire();
reduce_tile<PoolType::MAX, ReduceDim::REDUCE_SCALAR>(c_sd, c_scal, 0, 0, 0);
reduce_uninit();
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_m);
tile_regs_release();
cb_push_back(c_m, 1);}
// e = exp((s - m) * scale), face 0 only
{PZ("EXP");
cb_wait_front(c_m, 1);
reconfig_data_format(c_sd, c_m);
pack_reconfig_data_format(c_e);
sub_bcast_scalar_init(c_sd, c_m);
cb_reserve_back(c_e, 1);
tile_regs_acquire();
sub_tiles_bcast_scalar(c_sd, c_m, 0, 0, 0);
exp_tile_init<false>();
exp_tile<false, true>(0, VectorMode::None, SCALE_BF16);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, c_e);
tile_regs_release();
cb_push_back(c_e, 1);
cb_pop_front(c_m, 1);
cb_reserve_back(c_et, MT);
cb_push_back(c_et, MT);}
// o = (e @ V) / (e @ ONES)
{PZ("PV");
cb_wait_front(c_v, 2 * MT);
cb_wait_front(c_one, 1);
cb_wait_front(c_et, MT);
reconfig_full_operand(c_v, c_et); // srcA = in1 = V / ONES, srcB = in0 = e
pack_reconfig_data_format<true>(c_out);
matmul_init(c_et, c_v, 0);
cb_reserve_back(c_out, 2);
tile_regs_acquire();
for (uint32_t j = 0; j < MT; ++j) {
matmul_tiles(c_et, c_v, j, j * 2, 0);
matmul_tiles(c_et, c_v, j, j * 2 + 1, 1);
matmul_tiles(c_et, c_one, j, 0, 2);
}
// the tiny [1, 32] tiles span DST faces 0 and 1 only: VectorMode::R
recip_tile_init();
recip_tile(2, VectorMode::R);
mul_binary_tile_init();
MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_sfpu_binary_mul,
(APPROX, ckernel::BinaryOp::MUL, 8, DST_ACCUM_MODE), 0, 2, 0, VectorMode::R)));
MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, calculate_sfpu_binary_mul,
(APPROX, ckernel::BinaryOp::MUL, 8, DST_ACCUM_MODE), 1, 2, 1, VectorMode::R)));
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_et, MT);
cb_pop_front(c_e, 1);
cb_pop_front(c_sd, 1);
cb_pop_front(c_one, 1);
cb_pop_front(c_qt, 2);
cb_pop_front(c_k, 2 * MT);
cb_pop_front(c_v, 2 * MT);
cb_pop_front(c_scal, 1);
}