File size: 5,389 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | // 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);
}
|