code/tt_sam3/tt/kernels/qkv_rope_reader.cpp

// 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);
    }
}