File size: 3,192 Bytes
b9db84c | 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 | // SPDX-License-Identifier: Apache-2.0
// sam3 backbone: fused head split + RoPE reader. Work unit u = (b, r, h): b = window / batch, r = row tile, h = head.
// Per unit: the 6 qkv tiles of head h in row tile (b, r) -> cb_in (q0 q1 k0 k1 v0 v1); when r changes: the 4 cos/sin
// tiles of row r -> cb_cs (cos0 cos1 sin0 sin1). First: the scalar (-1) tile -> cb_scalar.
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t cb_in = get_compile_time_arg_val(0);
constexpr uint32_t cb_cs = get_compile_time_arg_val(1);
constexpr uint32_t cb_scalar = get_compile_time_arg_val(2);
constexpr uint32_t Lt = get_compile_time_arg_val(3); // row tiles per batch entry
constexpr uint32_t NH = get_compile_time_arg_val(4); // heads
constexpr uint32_t ROW_TILES = get_compile_time_arg_val(5); // qkv tiles per row (3 * NH * 2)
constexpr uint32_t TILE_BYTES = get_compile_time_arg_val(6);
constexpr auto qkv_args = TensorAccessorArgs<7>();
constexpr auto cos_args = TensorAccessorArgs<qkv_args.next_compile_time_args_offset()>();
constexpr auto sin_args = TensorAccessorArgs<cos_args.next_compile_time_args_offset()>();
void kernel_main() {
const uint32_t qkv_addr = get_arg_val<uint32_t>(0);
const uint32_t cos_addr = get_arg_val<uint32_t>(1);
const uint32_t sin_addr = get_arg_val<uint32_t>(2);
const uint32_t start = get_arg_val<uint32_t>(3);
const uint32_t n = get_arg_val<uint32_t>(4);
const auto qkv = TensorAccessor(qkv_args, qkv_addr, TILE_BYTES);
const auto cs_c = TensorAccessor(cos_args, cos_addr, TILE_BYTES);
const auto cs_s = TensorAccessor(sin_args, sin_addr, TILE_BYTES);
cb_reserve_back(cb_scalar, 1);
volatile tt_l1_ptr uint16_t* sc = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(get_write_ptr(cb_scalar));
sc[0] = 0xBF80; // bf16 -1.0 (the value ttnn's rotary_embedding broadcasts)
cb_push_back(cb_scalar, 1);
uint32_t cur_r = 0xFFFFFFFF;
for (uint32_t u = start; u < start + n; ++u) {
const uint32_t h = u % NH;
const uint32_t br = u / NH; // b * Lt + r
const uint32_t r = br % Lt;
if (r != cur_r) {
cb_reserve_back(cb_cs, 4);
uint32_t l1 = get_write_ptr(cb_cs);
noc_async_read(cs_c.get_noc_addr(r * 2), l1, TILE_BYTES);
noc_async_read(cs_c.get_noc_addr(r * 2 + 1), l1 + TILE_BYTES, TILE_BYTES);
noc_async_read(cs_s.get_noc_addr(r * 2), l1 + 2 * TILE_BYTES, TILE_BYTES);
noc_async_read(cs_s.get_noc_addr(r * 2 + 1), l1 + 3 * TILE_BYTES, TILE_BYTES);
noc_async_read_barrier();
cb_push_back(cb_cs, 4);
cur_r = r;
}
cb_reserve_back(cb_in, 6);
uint32_t l1 = get_write_ptr(cb_in);
const uint32_t base = br * ROW_TILES + h * 2;
for (uint32_t p = 0; p < 3; ++p) { // q, k, v blocks of the row
const uint32_t t = base + p * NH * 2;
noc_async_read(qkv.get_noc_addr(t), l1 + (2 * p) * TILE_BYTES, TILE_BYTES);
noc_async_read(qkv.get_noc_addr(t + 1), l1 + (2 * p + 1) * TILE_BYTES, TILE_BYTES);
}
noc_async_read_barrier();
cb_push_back(cb_in, 6);
}
}
|