File size: 5,672 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 | // 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);
}
|