tt-model push bevformer-p150 (container)

55474c5 verified 17 days ago

23 kB

// SPDX-License-Identifier: Apache-2.0
// msda_mk (tt/msda_mk.py, BEVFORMER_MSDA_MK; OPT_PLAN section 4 MK4a / MK4b, D19): the MSDA block module megakernel,
// data movement. One program per encoder MSDA call does what six (TSA) / four (SCA) programs did:
//   the query concat [A | B] (TSA: [v0 | query + pos]) -> the compact sampling + attention projections in ONE
//   K-blocked matmul (fp32 weights resident in L1) -> + bias -> loc = int32(cols + base [+ shift row]) and the fp32
//   logits -> (the msda_k steps) in-kernel softmax, pixel-pair gathers, fp32 corner weights, packer-L1 sum -> the
//   merged-head output.
// The same source runs on both data-movement RISC-Vs: PART 0 (NCRISC) owns the points [0, P/2), the per-block inputs
// (the K tiles of the query rows from A / B, the LOC_TW base tiles) and the reduce scalers / masks / E tiles; PART 1
// (BRISC) owns the points [P/2, P), loads the weights / bias / shift once and writes the output tiles.
//
// Work: the items (b, j) (32-query block b, head block j) in BLOCK-MAJOR order w = b * H + j, contiguous ranges per
// core (core index from the logical coordinates; common RT args only, probe P3). A core computes the front (the
// projections and the glue) of every block its range touches (at most two partial blocks per core are recomputed by
// a neighbour: identical values) and the MSDA of its own items.
// Front -> gathers hand-off: the compute kernel packs each block's LOC_TW int32 loc tiles into cb_loc0 and cb_loc1
// (one copy per RISC) and its AT_TW fp32 logit tiles into cb_lgc (PART 0); the readers wait for them at the block's
// first item and pop them after its last. The compute kernel runs the front one block ahead (CB depth 2 blocks).
#include <stdint.h>

#include "api/dataflow/dataflow_api.h"
#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_dataflow.hpp"

constexpr uint32_t PART = get_compile_time_arg_val(0);
constexpr int32_t HV = get_compile_time_arg_val(1);
constexpr int32_t WV = get_compile_time_arg_val(2);
constexpr uint32_t NQ = get_compile_time_arg_val(3);
constexpr uint32_t NB = get_compile_time_arg_val(4);
constexpr uint32_t H = get_compile_time_arg_val(5);
constexpr uint32_t SUB = get_compile_time_arg_val(6);
constexpr uint32_t P = get_compile_time_arg_val(7);
constexpr uint32_t NC = get_compile_time_arg_val(8);
constexpr uint32_t GX = get_compile_time_arg_val(9);
constexpr uint32_t KA = get_compile_time_arg_val(10);      // K tiles of the query rows from A (per block)
constexpr uint32_t KBT = get_compile_time_arg_val(11);     // K tiles from B (0: A only)
constexpr uint32_t LOC_TW = get_compile_time_arg_val(12);  // loc tiles per block (H * 2P / 32)
constexpr uint32_t AT_TW = get_compile_time_arg_val(13);   // logit tiles per block (H * P / 32)
constexpr uint32_t PAD_BITS = get_compile_time_arg_val(14);
constexpr uint32_t HAS_SHIFT = get_compile_time_arg_val(15);
// ACT (BEVFORMER_SCA_ACTIVE, SCA only): the per-calibration runtime tensor `nblk` (uint32 page, entries 0..SUB-1 =
// the active 32-row blocks of each camera's rebatched rows, valid rows first) restricts the work to the active
// blocks: item w = a * H + j over the NA = sum(nblk) active blocks a, block(a) = camera c * BPS + (a - prefix c).
// The padding blocks behind each camera's rows are never computed (their output rows are never gathered).
constexpr uint32_t ACT = get_compile_time_arg_val(16);
// RAW (BEVFORMER_SCA_BASE=kernel): `base` is the raw grid base [Q, 32] fp32 (one tile per block, the 2P point columns
// in 0..2P-1); PART 0 reads each block's tile with faces 0 / 2 also written into faces 1 / 3 (both column halves =
// the point columns: the compact layout of two heads per tile), PART 1 loads the `aff` tiles (S, O rows replicated)
constexpr uint32_t RAW = get_compile_time_arg_val(17);
constexpr auto v_args = TensorAccessorArgs<18>();
constexpr auto a_args = TensorAccessorArgs<v_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<a_args.next_compile_time_args_offset()>();
constexpr auto w_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr auto bi_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr auto ba_args = TensorAccessorArgs<bi_args.next_compile_time_args_offset()>();
constexpr auto s_args = TensorAccessorArgs<ba_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();
constexpr auto n_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
constexpr auto f_args = TensorAccessorArgs<n_args.next_compile_time_args_offset()>();

constexpr uint32_t KT = KA + KBT;
constexpr uint32_t NOUT = LOC_TW + AT_TW;
constexpr uint32_t LOC_HS = 2 * P;
constexpr uint32_t cb_in = PART ? 1 : 0, cb_geo = PART ? 3 : 2, cb_e = 4, cb_out = 16;
constexpr uint32_t cb_lg = 9, cb_scmax = 10, cb_scsum = 11, cb_msk = 12;
constexpr uint32_t cb_x = 21, cb_wsa = 22, cb_bias = 23, cb_base = 25, cb_shift = 26;
constexpr uint32_t cb_loc = PART ? 29 : 28, cb_lgc = 30;
constexpr uint32_t cb_cnt = PART ? 8 : 7;
constexpr uint32_t cb_aff = 31;
constexpr uint32_t BASE_TILES = RAW ? 1 : LOC_TW;   // ACT: the nblk page (PART 0's copy is pushed to the compute kernel)

constexpr uint32_t PH = P / 2;
constexpr uint32_t P0 = PART * PH;
constexpr uint32_t NS = PH * 4;
constexpr uint32_t TB16 = 2048, TB32 = 4096;
constexpr uint32_t NWORK = NB * H;
constexpr uint32_t BPS = NB / SUB;
constexpr int32_t LOC_LIM = 1 << 28;
static_assert(P % 2 == 0 && P <= 8, "msda_mk: P even, <= 8");
static_assert(NB % SUB == 0, "msda_mk: blocks split evenly over the value maps of a head");
static_assert(!RAW || LOC_HS == 16, "msda_mk: the raw base needs 16 point columns (two heads per tile)");
static_assert(32 % LOC_HS == 0 && (H * LOC_HS) % 32 == 0 && (H * P) % 32 == 0, "msda_mk: compact layouts");

uint32_t nblk[SUB > 0 ? SUB : 1];

// active block index -> absolute 32-row block (identity without ACT)
static inline uint32_t blk(uint32_t a) {
    if constexpr (ACT) {
        uint32_t c = 0;
        while (c + 1 < SUB && a >= nblk[c]) {
            a -= nblk[c];
            ++c;
        }
        return c * BPS + a;
    } else {
        return a;
    }
}

static inline uint32_t tile_elem(uint32_t r, uint32_t c) {
    return (((r >> 4) << 1) + (c >> 4)) * 256 + ((r & 15) << 4) + (c & 15);
}

constexpr uint32_t FRAC16_ONE = 0x3F800000;   // fp32 1.0 (MSDA_MK_FRAC16: 1 + q / 2^16 has q << 7 as mantissa)

static inline uint32_t q16_to_f32(uint32_t q) {
    if (q == 0) {
        return 0;
    }
    const uint32_t n = 31 - __builtin_clz(q);
    return ((127 - 16 + n) << 23) | ((q << (23 - n)) & 0x7FFFFF);
}

static inline void zero_l1(uint32_t addr, uint32_t bytes) {
    const uint64_t z = get_noc_addr(MEM_ZEROS_BASE);
    for (uint32_t o = 0; o < bytes; o += MEM_ZEROS_SIZE) {
        noc_async_read(z, addr + o, MEM_ZEROS_SIZE);
    }
}

// PART 0: the KT query tiles of block b (A then B) and its LOC_TW base tiles
template <typename TA, typename TB_, typename TC>
static inline void push_block_inputs(uint32_t b, const TA& A, const TB_& B, const TC& BA) {
    cb_reserve_back(cb_x, KT);
    cb_reserve_back(cb_base, BASE_TILES);
    const uint32_t xa = get_write_ptr(cb_x);
    const uint32_t ba = get_write_ptr(cb_base);
    for (uint32_t k = 0; k < KA; ++k) {
        noc_async_read_tile(b * KA + k, A, xa + k * TB16);
    }
    if constexpr (KBT > 0) {
        for (uint32_t k = 0; k < KBT; ++k) {
            noc_async_read_tile(b * KBT + k, B, xa + (KA + k) * TB16);
        }
    }
    if constexpr (RAW) {
        // fp32 faces are 1 KB: face f of the source tile -> faces f and f + 1 (f = 0, 2)
        for (uint32_t f = 0; f < 4; f += 2) {
            const uint64_t src = BA.get_noc_addr(b, f * 1024);
            noc_async_read(src, ba + f * 1024, 1024);
            noc_async_read(src, ba + (f + 1) * 1024, 1024);
        }
    } else {
        for (uint32_t t = 0; t < LOC_TW; ++t) {
            noc_async_read_tile(b * LOC_TW + t, BA, ba + t * TB32);
        }
    }
    noc_async_read_barrier();
    cb_push_back(cb_x, KT);
    cb_push_back(cb_base, BASE_TILES);
}

#ifdef MSDA_MK_MCAST
// BEVFORMER_MSDA_MK_MCAST (D20 round 2): PART 1's resident weights are read from DRAM once per grid row by the row's
// first core (logical x 0) and multicast to the GX - 1 other cores of the row (mmres_k's / bevfusion K6's handshake:
// each receiver zeroes its 'valid' semaphore, increments the sender's 'ready' semaphore and waits for valid = 1; the
// sender waits for GX - 1, multicasts, flushes, multicasts valid = 1; the rectangle never contains the sender).
// Every core of the grid takes part, also one without work (NC = the whole grid).
constexpr uint32_t MC_VYS[] = {MC_VY};
#endif

template <typename WA>
inline void load_weights(const WA& W, uint32_t wa) {
#ifdef MSDA_MK_MCAST
    const bool sender = get_absolute_logical_x() == 0;
    const uint32_t y = MC_VYS[get_absolute_logical_y()];
    const uint32_t ready_addr = get_semaphore(0);
    const uint32_t valid_addr = get_semaphore(1);
    volatile tt_l1_ptr uint32_t* ready_ptr = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ready_addr);
    volatile tt_l1_ptr uint32_t* valid_ptr = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(valid_addr);
    constexpr uint32_t NRECV = GX - 1;
    if (sender) {
        for (uint32_t t = 0; t < KT * NOUT; ++t) {
            noc_async_read_tile(t, W, wa + t * TB32);
        }
        noc_async_read_barrier();
        const bool n1 = noc_index == 1;
        const uint32_t x0 = n1 ? MC_VXE : MC_VX1, x1 = n1 ? MC_VX1 : MC_VXE;
        noc_semaphore_wait(ready_ptr, NRECV);
        noc_semaphore_set(ready_ptr, 0);
        noc_async_write_multicast(wa, get_noc_multicast_addr(x0, y, x1, y, wa), KT * NOUT * TB32, NRECV, true);
        noc_async_writes_flushed();
        *valid_ptr = 1;
        noc_semaphore_set_multicast(valid_addr, get_noc_multicast_addr(x0, y, x1, y, valid_addr), NRECV);
        noc_async_write_barrier();
    } else {
        noc_semaphore_set(valid_ptr, 0);
        noc_semaphore_inc(get_noc_addr(MC_VX0, y, ready_addr), 1);
        noc_semaphore_wait(valid_ptr, 1);
    }
#else
#ifndef MSDA_MK_NOWLOAD
    for (uint32_t t = 0; t < KT * NOUT; ++t) {
        noc_async_read_tile(t, W, wa + t * TB32);
    }
#endif
#endif
}

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    uint32_t nwork = NWORK;
    if constexpr (ACT) {
        const auto n_acc = TensorAccessor(n_args, get_common_arg_val<uint32_t>(8), 128);
        cb_reserve_back(cb_cnt, 1);
        const uint32_t na = get_write_ptr(cb_cnt);
        noc_async_read(n_acc.get_noc_addr(0), na, 128);
        noc_async_read_barrier();
        volatile tt_l1_ptr uint32_t* nv = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(na);
        uint32_t tot = 0;
        for (uint32_t c = 0; c < SUB; ++c) {
            const uint32_t v = nv[c] > BPS ? BPS : nv[c];
            nv[c] = v;                       // the compute kernel reads the clamped counts
            nblk[c] = v;
            tot += v;
        }
        if constexpr (PART == 0) {
            cb_push_back(cb_cnt, 1);
        }
        nwork = tot * H;
    }
    const uint32_t w0 = ci * nwork / NC;
    const uint32_t w1 = (ci + 1) * nwork / NC;
    if (w0 >= w1) {
#ifdef MSDA_MK_MCAST
        if constexpr (PART == 1) {          // no work here, but the row's weight multicast needs every core
            const auto W = TensorAccessor(w_args, get_common_arg_val<uint32_t>(3), TB32);
            cb_reserve_back(cb_wsa, KT * NOUT);
            load_weights(W, get_write_ptr(cb_wsa));
        }
#endif
        return;
    }
    const uint32_t bf = w0 / H;           // active block indices (absolute blocks without ACT)
    const uint32_t bl = (w1 - 1) / H;
    const auto v_acc = TensorAccessor(v_args, get_common_arg_val<uint32_t>(0), 128);
    const auto A = TensorAccessor(a_args, get_common_arg_val<uint32_t>(1), TB16);
    const auto B = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), TB16);
    const auto W = TensorAccessor(w_args, get_common_arg_val<uint32_t>(3), TB32);
    const auto BI = TensorAccessor(bi_args, get_common_arg_val<uint32_t>(4), TB32);
    const auto BA = TensorAccessor(ba_args, get_common_arg_val<uint32_t>(5), TB32);
    const auto S = TensorAccessor(s_args, get_common_arg_val<uint32_t>(6), TB32);
    const auto o_acc = TensorAccessor(o_args, get_common_arg_val<uint32_t>(7), TB16);
    const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);

    if constexpr (PART == 1) {
        // the projection weights [KT, NOUT] tiles (row-major tile order k * NOUT + n), the bias row and the shift
        // row (row 0 replicated over the 32 rows: the value binary_ng's row broadcast adds to every row)
        cb_reserve_back(cb_wsa, KT * NOUT);
        const uint32_t wa = get_write_ptr(cb_wsa);
        load_weights(W, wa);
        cb_reserve_back(cb_bias, NOUT);
        const uint32_t bia = get_write_ptr(cb_bias);
        for (uint32_t t = 0; t < NOUT; ++t) {
            noc_async_read_tile(t, BI, bia + t * TB32);
        }
        if constexpr (RAW) {
            const auto F = TensorAccessor(f_args, get_common_arg_val<uint32_t>(9), TB32);
            cb_reserve_back(cb_aff, 2);
            const uint32_t fa = get_write_ptr(cb_aff);
            noc_async_read_tile(0, F, fa);
            noc_async_read_tile(1, F, fa + TB32);
        }
        noc_async_read_barrier();
        cb_push_back(cb_wsa, KT * NOUT);
        cb_push_back(cb_bias, NOUT);
        if constexpr (RAW) {
            cb_push_back(cb_aff, 2);
        }
        if constexpr (HAS_SHIFT) {
            cb_reserve_back(cb_shift, LOC_TW);
            const uint32_t sa = get_write_ptr(cb_shift);
            for (uint32_t j = 0; j < LOC_TW; ++j) {
                noc_async_read_tile(j, S, sa + j * TB32);
            }
            noc_async_read_barrier();
            for (uint32_t j = 0; j < LOC_TW; ++j) {
                volatile tt_l1_ptr uint32_t* t = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sa + j * TB32);
                for (uint32_t f = 0; f < 4; ++f) {
                    const uint32_t src = (f & 1) * 256;
                    for (uint32_t r = (f < 2 ? 1 : 0); r < 16; ++r) {
                        for (uint32_t c = 0; c < 16; ++c) {
                            t[f * 256 + r * 16 + c] = t[src + c];
                        }
                    }
                }
            }
            cb_push_back(cb_shift, LOC_TW);
        }
    } else {
        // E_p (bf16): row p all ones; the softmax's reduce scalers; the two half masks; the logit pages pre-filled
        // with PAD_LOGIT (the msda_k PART-0 setup)
        cb_reserve_back(cb_e, P);
        const uint32_t e_addr = get_write_ptr(cb_e);
        zero_l1(e_addr, P * TB16);
        noc_async_read_barrier();
        volatile tt_l1_ptr uint16_t* e = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(e_addr);
        for (uint32_t p = 0; p < P; ++p) {
            for (uint32_t c = 0; c < 32; ++c) {
                e[p * (TB16 / 2) + tile_elem(p, c)] = 0x3F80;
            }
        }
        cb_push_back(cb_e, P);
        dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scmax, ckernel::PoolType::MAX,
                                                                 ckernel::ReduceDim::REDUCE_ROW>();
        dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scsum, ckernel::PoolType::SUM,
                                                                 ckernel::ReduceDim::REDUCE_ROW>();
        cb_reserve_back(cb_msk, 2);
        volatile tt_l1_ptr uint32_t* m = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_msk));
        for (uint32_t h = 0; h < 2; ++h) {
            for (uint32_t r = 0; r < 32; ++r) {
                for (uint32_t c = 0; c < 32; ++c) {
                    m[h * 1024 + tile_elem(r, c)] = (c >= h * PH && c < (h + 1) * PH) ? 0x3F800000u : 0u;
                }
            }
        }
        cb_push_back(cb_msk, 2);
        cb_reserve_back(cb_lg, 2);
        volatile tt_l1_ptr uint32_t* lg = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
        for (uint32_t i = 0; i < 2048; ++i) {
            lg[i] = PAD_BITS;
        }
        // the first two blocks' inputs (the compute kernel runs the front one block ahead)
        push_block_inputs(blk(bf), A, B, BA);
        if (bf + 1 <= bl) {
            push_block_inputs(blk(bf + 1), A, B, BA);
        }
    }
    cb_reserve_back(cb_geo, 6);
    zero_l1(get_write_ptr(cb_geo), 6 * TB32);
    noc_async_read_barrier();

    uint32_t prev = 0xFFFFFFFF;
    for (uint32_t a = bf; a <= bl; ++a) {
        const uint32_t b = blk(a);
        const uint32_t j0 = a == bf ? w0 - bf * H : 0;
        const uint32_t j1 = a == bl ? w1 - bl * H : H;
        cb_wait_front(cb_loc, LOC_TW);
        const uint32_t lbuf = get_read_ptr(cb_loc);
        uint32_t lgbuf = 0;
        if constexpr (PART == 0) {
            cb_wait_front(cb_lgc, AT_TW);
            lgbuf = get_read_ptr(cb_lgc);
        }
        for (uint32_t j = j0; j < j1; ++j) {
            const uint32_t vbase = (j * SUB + b / BPS) * HV * (WV + 1);
            cb_reserve_back(cb_geo, 3);
            cb_reserve_back(cb_in, NS);
            const uint32_t g_addr = get_write_ptr(cb_geo);
            const uint32_t in_base = get_write_ptr(cb_in);
            volatile tt_l1_ptr int32_t* loc =
                reinterpret_cast<volatile tt_l1_ptr int32_t*>(lbuf + ((j * LOC_HS) / 32) * TB32);
            const uint32_t lcol = (j * LOC_HS) % 32;
            if constexpr (PART == 0) {
                // the head's logit tile: columns 0..P-1 <- compact columns (j P) % 32 + p; columns >= P keep PAD
                cb_reserve_back(cb_lg, 1);
                volatile tt_l1_ptr uint32_t* lg =
                    reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
                volatile tt_l1_ptr uint32_t* src =
                    reinterpret_cast<volatile tt_l1_ptr uint32_t*>(lgbuf + ((j * P) / 32) * TB32);
                const uint32_t acol = (j * P) % 32;
                for (uint32_t r = 0; r < 32; ++r) {
                    const bool ok = b * 32 + r < NQ;
                    for (uint32_t p = 0; p < P; ++p) {
                        lg[tile_elem(r, p)] = ok ? src[tile_elem(r, acol + p)] : 0u;
                    }
                }
                cb_push_back(cb_lg, 1);
            }
            volatile tt_l1_ptr uint32_t* fx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr);
            volatile tt_l1_ptr uint32_t* fy = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + TB32);
#ifdef MSDA_MK_NOPTS
            for (uint32_t r = 0; r < 0; ++r) {
#else
            for (uint32_t r = 0; r < 32; ++r) {
#endif
                const bool row_ok = b * 32 + r < NQ;
                for (uint32_t p = P0; p < P0 + PH; ++p) {
                    const uint32_t ew = tile_elem(r, p);
                    const uint32_t blk = in_base + (p - P0) * 4 * TB16 + r * 128;
                    if (!row_ok) {
#ifdef MSDA_MK_FRAC16
                        fx[ew] = FRAC16_ONE;      // decodes to 0.0
                        fy[ew] = FRAC16_ONE;
#else
                        fx[ew] = 0;
                        fy[ew] = 0;
#endif
                        noc_async_read(zeros, blk, 128);
                        noc_async_read(zeros, blk + 2 * TB16, 128);
                        continue;
                    }
                    int32_t vx = loc[tile_elem(r, lcol + 2 * p)];
                    int32_t vy = loc[tile_elem(r, lcol + 2 * p + 1)];
                    vx = vx > LOC_LIM ? LOC_LIM : (vx < -LOC_LIM ? -LOC_LIM : vx);
                    vy = vy > LOC_LIM ? LOC_LIM : (vy < -LOC_LIM ? -LOC_LIM : vy);
                    const int32_t x0 = vx >> 16;
                    const int32_t y0 = vy >> 16;
#ifdef MSDA_MK_FRAC16
                    // BEVFORMER_MSDA_FRAC16 (D20 round 3): the fraction q (16 bits) as the fp32 bits of 1 + q / 2^16
                    // (exponent of 1.0, q in the top mantissa bits): one shift and or; the compute kernel's SFPU
                    // weight step takes 1 + q / 2^16 - 1 = q / 2^16 exactly (= q16_to_f32(q), bit for bit)
                    fx[ew] = FRAC16_ONE | ((static_cast<uint32_t>(vx) & 0xFFFF) << 7);
                    fy[ew] = FRAC16_ONE | ((static_cast<uint32_t>(vy) & 0xFFFF) << 7);
#elif defined(MSDA_MK_RAWFRAC)
                    fx[ew] = static_cast<uint32_t>(vx) & 0xFFFF;   // (anatomy probe: no fp32 conversion)
                    fy[ew] = static_cast<uint32_t>(vy) & 0xFFFF;
#elif defined(MSDA_MK_NOGEO)
#else
                    fx[ew] = q16_to_f32(static_cast<uint32_t>(vx) & 0xFFFF);
                    fy[ew] = q16_to_f32(static_cast<uint32_t>(vy) & 0xFFFF);
#endif
                    const bool yv0 = y0 >= 0 && y0 < HV, yv1 = y0 + 1 >= 0 && y0 + 1 < HV;
                    const bool xp = x0 >= -1 && x0 < WV;  // page i = x0 + 1 in [0, WV]
                    const uint32_t pg = vbase + static_cast<uint32_t>(y0 * (WV + 1) + x0 + 1);
#ifdef MSDA_MK_NOREQ
                    if (pg == 0xFFFFFFFF) {
                        noc_async_read(zeros, blk, 128);
                    }
                    continue;
#endif
#ifdef MSDA_MK_NOREAD
                    if (false) {
#else
                    if (xp && yv0) {
#endif
                        noc_async_read(v_acc.get_noc_addr(pg), blk, 128);
                    } else {
                        noc_async_read(zeros, blk, 128);
                    }
#ifdef MSDA_MK_NOREAD
                    if (false) {
#else
                    if (xp && yv1) {
#endif
                        noc_async_read(v_acc.get_noc_addr(pg + WV + 1), blk + 2 * TB16, 128);
                    } else {
                        noc_async_read(zeros, blk + 2 * TB16, 128);
                    }
                }
            }
            noc_async_read_barrier();
            cb_push_back(cb_geo, 3);
            cb_push_back(cb_in, NS);
            if constexpr (PART == 1) {
                if (prev != 0xFFFFFFFF) {
                    cb_wait_front(cb_out, 1);
                    noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(prev), TB16);
                    noc_async_writes_flushed();
                    cb_pop_front(cb_out, 1);
                }
                prev = b * H + j;   // output tile (b, j) of [Q, H * 32]
            }
        }
        cb_pop_front(cb_loc, LOC_TW);
        if constexpr (PART == 0) {
            cb_pop_front(cb_lgc, AT_TW);
            if (a + 2 <= bl) {
                push_block_inputs(blk(a + 2), A, B, BA);
            }
        }
    }
    if constexpr (PART == 1) {
        cb_wait_front(cb_out, 1);
        noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(prev), TB16);
        noc_async_write_barrier();
        cb_pop_front(cb_out, 1);
    }
}

	// SPDX-License-Identifier: Apache-2.0
	// msda_mk (tt/msda_mk.py, BEVFORMER_MSDA_MK; OPT_PLAN section 4 MK4a / MK4b, D19): the MSDA block module megakernel,
	// data movement. One program per encoder MSDA call does what six (TSA) / four (SCA) programs did:
	// the query concat [A \| B] (TSA: [v0 \| query + pos]) -> the compact sampling + attention projections in ONE
	// K-blocked matmul (fp32 weights resident in L1) -> + bias -> loc = int32(cols + base [+ shift row]) and the fp32
	// logits -> (the msda_k steps) in-kernel softmax, pixel-pair gathers, fp32 corner weights, packer-L1 sum -> the
	// merged-head output.
	// The same source runs on both data-movement RISC-Vs: PART 0 (NCRISC) owns the points [0, P/2), the per-block inputs
	// (the K tiles of the query rows from A / B, the LOC_TW base tiles) and the reduce scalers / masks / E tiles; PART 1
	// (BRISC) owns the points [P/2, P), loads the weights / bias / shift once and writes the output tiles.
	//
	// Work: the items (b, j) (32-query block b, head block j) in BLOCK-MAJOR order w = b * H + j, contiguous ranges per
	// core (core index from the logical coordinates; common RT args only, probe P3). A core computes the front (the
	// projections and the glue) of every block its range touches (at most two partial blocks per core are recomputed by
	// a neighbour: identical values) and the MSDA of its own items.
	// Front -> gathers hand-off: the compute kernel packs each block's LOC_TW int32 loc tiles into cb_loc0 and cb_loc1
	// (one copy per RISC) and its AT_TW fp32 logit tiles into cb_lgc (PART 0); the readers wait for them at the block's
	// first item and pop them after its last. The compute kernel runs the front one block ahead (CB depth 2 blocks).
	#include <stdint.h>

	#include "api/dataflow/dataflow_api.h"
	#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_dataflow.hpp"

	constexpr uint32_t PART = get_compile_time_arg_val(0);
	constexpr int32_t HV = get_compile_time_arg_val(1);
	constexpr int32_t WV = get_compile_time_arg_val(2);
	constexpr uint32_t NQ = get_compile_time_arg_val(3);
	constexpr uint32_t NB = get_compile_time_arg_val(4);
	constexpr uint32_t H = get_compile_time_arg_val(5);
	constexpr uint32_t SUB = get_compile_time_arg_val(6);
	constexpr uint32_t P = get_compile_time_arg_val(7);
	constexpr uint32_t NC = get_compile_time_arg_val(8);
	constexpr uint32_t GX = get_compile_time_arg_val(9);
	constexpr uint32_t KA = get_compile_time_arg_val(10); // K tiles of the query rows from A (per block)
	constexpr uint32_t KBT = get_compile_time_arg_val(11); // K tiles from B (0: A only)
	constexpr uint32_t LOC_TW = get_compile_time_arg_val(12); // loc tiles per block (H * 2P / 32)
	constexpr uint32_t AT_TW = get_compile_time_arg_val(13); // logit tiles per block (H * P / 32)
	constexpr uint32_t PAD_BITS = get_compile_time_arg_val(14);
	constexpr uint32_t HAS_SHIFT = get_compile_time_arg_val(15);
	// ACT (BEVFORMER_SCA_ACTIVE, SCA only): the per-calibration runtime tensor `nblk` (uint32 page, entries 0..SUB-1 =
	// the active 32-row blocks of each camera's rebatched rows, valid rows first) restricts the work to the active
	// blocks: item w = a * H + j over the NA = sum(nblk) active blocks a, block(a) = camera c * BPS + (a - prefix c).
	// The padding blocks behind each camera's rows are never computed (their output rows are never gathered).
	constexpr uint32_t ACT = get_compile_time_arg_val(16);
	// RAW (BEVFORMER_SCA_BASE=kernel): `base` is the raw grid base [Q, 32] fp32 (one tile per block, the 2P point columns
	// in 0..2P-1); PART 0 reads each block's tile with faces 0 / 2 also written into faces 1 / 3 (both column halves =
	// the point columns: the compact layout of two heads per tile), PART 1 loads the `aff` tiles (S, O rows replicated)
	constexpr uint32_t RAW = get_compile_time_arg_val(17);
	constexpr auto v_args = TensorAccessorArgs<18>();
	constexpr auto a_args = TensorAccessorArgs<v_args.next_compile_time_args_offset()>();
	constexpr auto b_args = TensorAccessorArgs<a_args.next_compile_time_args_offset()>();
	constexpr auto w_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
	constexpr auto bi_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
	constexpr auto ba_args = TensorAccessorArgs<bi_args.next_compile_time_args_offset()>();
	constexpr auto s_args = TensorAccessorArgs<ba_args.next_compile_time_args_offset()>();
	constexpr auto o_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();
	constexpr auto n_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
	constexpr auto f_args = TensorAccessorArgs<n_args.next_compile_time_args_offset()>();

	constexpr uint32_t KT = KA + KBT;
	constexpr uint32_t NOUT = LOC_TW + AT_TW;
	constexpr uint32_t LOC_HS = 2 * P;
	constexpr uint32_t cb_in = PART ? 1 : 0, cb_geo = PART ? 3 : 2, cb_e = 4, cb_out = 16;
	constexpr uint32_t cb_lg = 9, cb_scmax = 10, cb_scsum = 11, cb_msk = 12;
	constexpr uint32_t cb_x = 21, cb_wsa = 22, cb_bias = 23, cb_base = 25, cb_shift = 26;
	constexpr uint32_t cb_loc = PART ? 29 : 28, cb_lgc = 30;
	constexpr uint32_t cb_cnt = PART ? 8 : 7;
	constexpr uint32_t cb_aff = 31;
	constexpr uint32_t BASE_TILES = RAW ? 1 : LOC_TW; // ACT: the nblk page (PART 0's copy is pushed to the compute kernel)

	constexpr uint32_t PH = P / 2;
	constexpr uint32_t P0 = PART * PH;
	constexpr uint32_t NS = PH * 4;
	constexpr uint32_t TB16 = 2048, TB32 = 4096;
	constexpr uint32_t NWORK = NB * H;
	constexpr uint32_t BPS = NB / SUB;
	constexpr int32_t LOC_LIM = 1 << 28;
	static_assert(P % 2 == 0 && P <= 8, "msda_mk: P even, <= 8");
	static_assert(NB % SUB == 0, "msda_mk: blocks split evenly over the value maps of a head");
	static_assert(!RAW \|\| LOC_HS == 16, "msda_mk: the raw base needs 16 point columns (two heads per tile)");
	static_assert(32 % LOC_HS == 0 && (H * LOC_HS) % 32 == 0 && (H * P) % 32 == 0, "msda_mk: compact layouts");

	uint32_t nblk[SUB > 0 ? SUB : 1];

	// active block index -> absolute 32-row block (identity without ACT)
	static inline uint32_t blk(uint32_t a) {
	if constexpr (ACT) {
	uint32_t c = 0;
	while (c + 1 < SUB && a >= nblk[c]) {
	a -= nblk[c];
	++c;
	}
	return c * BPS + a;
	} else {
	return a;
	}
	}

	static inline uint32_t tile_elem(uint32_t r, uint32_t c) {
	return (((r >> 4) << 1) + (c >> 4)) * 256 + ((r & 15) << 4) + (c & 15);
	}

	constexpr uint32_t FRAC16_ONE = 0x3F800000; // fp32 1.0 (MSDA_MK_FRAC16: 1 + q / 2^16 has q << 7 as mantissa)

	static inline uint32_t q16_to_f32(uint32_t q) {
	if (q == 0) {
	return 0;
	}
	const uint32_t n = 31 - __builtin_clz(q);
	return ((127 - 16 + n) << 23) \| ((q << (23 - n)) & 0x7FFFFF);
	}

	static inline void zero_l1(uint32_t addr, uint32_t bytes) {
	const uint64_t z = get_noc_addr(MEM_ZEROS_BASE);
	for (uint32_t o = 0; o < bytes; o += MEM_ZEROS_SIZE) {
	noc_async_read(z, addr + o, MEM_ZEROS_SIZE);
	}
	}

	// PART 0: the KT query tiles of block b (A then B) and its LOC_TW base tiles
	template <typename TA, typename TB_, typename TC>
	static inline void push_block_inputs(uint32_t b, const TA& A, const TB_& B, const TC& BA) {
	cb_reserve_back(cb_x, KT);
	cb_reserve_back(cb_base, BASE_TILES);
	const uint32_t xa = get_write_ptr(cb_x);
	const uint32_t ba = get_write_ptr(cb_base);
	for (uint32_t k = 0; k < KA; ++k) {
	noc_async_read_tile(b * KA + k, A, xa + k * TB16);
	}
	if constexpr (KBT > 0) {
	for (uint32_t k = 0; k < KBT; ++k) {
	noc_async_read_tile(b * KBT + k, B, xa + (KA + k) * TB16);
	}
	}
	if constexpr (RAW) {
	// fp32 faces are 1 KB: face f of the source tile -> faces f and f + 1 (f = 0, 2)
	for (uint32_t f = 0; f < 4; f += 2) {
	const uint64_t src = BA.get_noc_addr(b, f * 1024);
	noc_async_read(src, ba + f * 1024, 1024);
	noc_async_read(src, ba + (f + 1) * 1024, 1024);
	}
	} else {
	for (uint32_t t = 0; t < LOC_TW; ++t) {
	noc_async_read_tile(b * LOC_TW + t, BA, ba + t * TB32);
	}
	}
	noc_async_read_barrier();
	cb_push_back(cb_x, KT);
	cb_push_back(cb_base, BASE_TILES);
	}

	#ifdef MSDA_MK_MCAST
	// BEVFORMER_MSDA_MK_MCAST (D20 round 2): PART 1's resident weights are read from DRAM once per grid row by the row's
	// first core (logical x 0) and multicast to the GX - 1 other cores of the row (mmres_k's / bevfusion K6's handshake:
	// each receiver zeroes its 'valid' semaphore, increments the sender's 'ready' semaphore and waits for valid = 1; the
	// sender waits for GX - 1, multicasts, flushes, multicasts valid = 1; the rectangle never contains the sender).
	// Every core of the grid takes part, also one without work (NC = the whole grid).
	constexpr uint32_t MC_VYS[] = {MC_VY};
	#endif

	template <typename WA>
	inline void load_weights(const WA& W, uint32_t wa) {
	#ifdef MSDA_MK_MCAST
	const bool sender = get_absolute_logical_x() == 0;
	const uint32_t y = MC_VYS[get_absolute_logical_y()];
	const uint32_t ready_addr = get_semaphore(0);
	const uint32_t valid_addr = get_semaphore(1);
	volatile tt_l1_ptr uint32_t* ready_ptr = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ready_addr);
	volatile tt_l1_ptr uint32_t* valid_ptr = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(valid_addr);
	constexpr uint32_t NRECV = GX - 1;
	if (sender) {
	for (uint32_t t = 0; t < KT * NOUT; ++t) {
	noc_async_read_tile(t, W, wa + t * TB32);
	}
	noc_async_read_barrier();
	const bool n1 = noc_index == 1;
	const uint32_t x0 = n1 ? MC_VXE : MC_VX1, x1 = n1 ? MC_VX1 : MC_VXE;
	noc_semaphore_wait(ready_ptr, NRECV);
	noc_semaphore_set(ready_ptr, 0);
	noc_async_write_multicast(wa, get_noc_multicast_addr(x0, y, x1, y, wa), KT * NOUT * TB32, NRECV, true);
	noc_async_writes_flushed();
	*valid_ptr = 1;
	noc_semaphore_set_multicast(valid_addr, get_noc_multicast_addr(x0, y, x1, y, valid_addr), NRECV);
	noc_async_write_barrier();
	} else {
	noc_semaphore_set(valid_ptr, 0);
	noc_semaphore_inc(get_noc_addr(MC_VX0, y, ready_addr), 1);
	noc_semaphore_wait(valid_ptr, 1);
	}
	#else
	#ifndef MSDA_MK_NOWLOAD
	for (uint32_t t = 0; t < KT * NOUT; ++t) {
	noc_async_read_tile(t, W, wa + t * TB32);
	}
	#endif
	#endif
	}

	void kernel_main() {
	const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
	if (ci >= NC) {
	return;
	}
	uint32_t nwork = NWORK;
	if constexpr (ACT) {
	const auto n_acc = TensorAccessor(n_args, get_common_arg_val<uint32_t>(8), 128);
	cb_reserve_back(cb_cnt, 1);
	const uint32_t na = get_write_ptr(cb_cnt);
	noc_async_read(n_acc.get_noc_addr(0), na, 128);
	noc_async_read_barrier();
	volatile tt_l1_ptr uint32_t* nv = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(na);
	uint32_t tot = 0;
	for (uint32_t c = 0; c < SUB; ++c) {
	const uint32_t v = nv[c] > BPS ? BPS : nv[c];
	nv[c] = v; // the compute kernel reads the clamped counts
	nblk[c] = v;
	tot += v;
	}
	if constexpr (PART == 0) {
	cb_push_back(cb_cnt, 1);
	}
	nwork = tot * H;
	}
	const uint32_t w0 = ci * nwork / NC;
	const uint32_t w1 = (ci + 1) * nwork / NC;
	if (w0 >= w1) {
	#ifdef MSDA_MK_MCAST
	if constexpr (PART == 1) { // no work here, but the row's weight multicast needs every core
	const auto W = TensorAccessor(w_args, get_common_arg_val<uint32_t>(3), TB32);
	cb_reserve_back(cb_wsa, KT * NOUT);
	load_weights(W, get_write_ptr(cb_wsa));
	}
	#endif
	return;
	}
	const uint32_t bf = w0 / H; // active block indices (absolute blocks without ACT)
	const uint32_t bl = (w1 - 1) / H;
	const auto v_acc = TensorAccessor(v_args, get_common_arg_val<uint32_t>(0), 128);
	const auto A = TensorAccessor(a_args, get_common_arg_val<uint32_t>(1), TB16);
	const auto B = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), TB16);
	const auto W = TensorAccessor(w_args, get_common_arg_val<uint32_t>(3), TB32);
	const auto BI = TensorAccessor(bi_args, get_common_arg_val<uint32_t>(4), TB32);
	const auto BA = TensorAccessor(ba_args, get_common_arg_val<uint32_t>(5), TB32);
	const auto S = TensorAccessor(s_args, get_common_arg_val<uint32_t>(6), TB32);
	const auto o_acc = TensorAccessor(o_args, get_common_arg_val<uint32_t>(7), TB16);
	const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);

	if constexpr (PART == 1) {
	// the projection weights [KT, NOUT] tiles (row-major tile order k * NOUT + n), the bias row and the shift
	// row (row 0 replicated over the 32 rows: the value binary_ng's row broadcast adds to every row)
	cb_reserve_back(cb_wsa, KT * NOUT);
	const uint32_t wa = get_write_ptr(cb_wsa);
	load_weights(W, wa);
	cb_reserve_back(cb_bias, NOUT);
	const uint32_t bia = get_write_ptr(cb_bias);
	for (uint32_t t = 0; t < NOUT; ++t) {
	noc_async_read_tile(t, BI, bia + t * TB32);
	}
	if constexpr (RAW) {
	const auto F = TensorAccessor(f_args, get_common_arg_val<uint32_t>(9), TB32);
	cb_reserve_back(cb_aff, 2);
	const uint32_t fa = get_write_ptr(cb_aff);
	noc_async_read_tile(0, F, fa);
	noc_async_read_tile(1, F, fa + TB32);
	}
	noc_async_read_barrier();
	cb_push_back(cb_wsa, KT * NOUT);
	cb_push_back(cb_bias, NOUT);
	if constexpr (RAW) {
	cb_push_back(cb_aff, 2);
	}
	if constexpr (HAS_SHIFT) {
	cb_reserve_back(cb_shift, LOC_TW);
	const uint32_t sa = get_write_ptr(cb_shift);
	for (uint32_t j = 0; j < LOC_TW; ++j) {
	noc_async_read_tile(j, S, sa + j * TB32);
	}
	noc_async_read_barrier();
	for (uint32_t j = 0; j < LOC_TW; ++j) {
	volatile tt_l1_ptr uint32_t* t = reinterpret_cast<volatile tt_l1_ptr uint32_t>(sa + j TB32);
	for (uint32_t f = 0; f < 4; ++f) {
	const uint32_t src = (f & 1) * 256;
	for (uint32_t r = (f < 2 ? 1 : 0); r < 16; ++r) {
	for (uint32_t c = 0; c < 16; ++c) {
	t[f * 256 + r * 16 + c] = t[src + c];
	}
	}
	}
	}
	cb_push_back(cb_shift, LOC_TW);
	}
	} else {
	// E_p (bf16): row p all ones; the softmax's reduce scalers; the two half masks; the logit pages pre-filled
	// with PAD_LOGIT (the msda_k PART-0 setup)
	cb_reserve_back(cb_e, P);
	const uint32_t e_addr = get_write_ptr(cb_e);
	zero_l1(e_addr, P * TB16);
	noc_async_read_barrier();
	volatile tt_l1_ptr uint16_t* e = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(e_addr);
	for (uint32_t p = 0; p < P; ++p) {
	for (uint32_t c = 0; c < 32; ++c) {
	e[p * (TB16 / 2) + tile_elem(p, c)] = 0x3F80;
	}
	}
	cb_push_back(cb_e, P);
	dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scmax, ckernel::PoolType::MAX,
	ckernel::ReduceDim::REDUCE_ROW>();
	dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scsum, ckernel::PoolType::SUM,
	ckernel::ReduceDim::REDUCE_ROW>();
	cb_reserve_back(cb_msk, 2);
	volatile tt_l1_ptr uint32_t* m = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_msk));
	for (uint32_t h = 0; h < 2; ++h) {
	for (uint32_t r = 0; r < 32; ++r) {
	for (uint32_t c = 0; c < 32; ++c) {
	m[h * 1024 + tile_elem(r, c)] = (c >= h * PH && c < (h + 1) * PH) ? 0x3F800000u : 0u;
	}
	}
	}
	cb_push_back(cb_msk, 2);
	cb_reserve_back(cb_lg, 2);
	volatile tt_l1_ptr uint32_t* lg = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
	for (uint32_t i = 0; i < 2048; ++i) {
	lg[i] = PAD_BITS;
	}
	// the first two blocks' inputs (the compute kernel runs the front one block ahead)
	push_block_inputs(blk(bf), A, B, BA);
	if (bf + 1 <= bl) {
	push_block_inputs(blk(bf + 1), A, B, BA);
	}
	}
	cb_reserve_back(cb_geo, 6);
	zero_l1(get_write_ptr(cb_geo), 6 * TB32);
	noc_async_read_barrier();

	uint32_t prev = 0xFFFFFFFF;
	for (uint32_t a = bf; a <= bl; ++a) {
	const uint32_t b = blk(a);
	const uint32_t j0 = a == bf ? w0 - bf * H : 0;
	const uint32_t j1 = a == bl ? w1 - bl * H : H;
	cb_wait_front(cb_loc, LOC_TW);
	const uint32_t lbuf = get_read_ptr(cb_loc);
	uint32_t lgbuf = 0;
	if constexpr (PART == 0) {
	cb_wait_front(cb_lgc, AT_TW);
	lgbuf = get_read_ptr(cb_lgc);
	}
	for (uint32_t j = j0; j < j1; ++j) {
	const uint32_t vbase = (j * SUB + b / BPS) * HV * (WV + 1);
	cb_reserve_back(cb_geo, 3);
	cb_reserve_back(cb_in, NS);
	const uint32_t g_addr = get_write_ptr(cb_geo);
	const uint32_t in_base = get_write_ptr(cb_in);
	volatile tt_l1_ptr int32_t* loc =
	reinterpret_cast<volatile tt_l1_ptr int32_t>(lbuf + ((j LOC_HS) / 32) * TB32);
	const uint32_t lcol = (j * LOC_HS) % 32;
	if constexpr (PART == 0) {
	// the head's logit tile: columns 0..P-1 <- compact columns (j P) % 32 + p; columns >= P keep PAD
	cb_reserve_back(cb_lg, 1);
	volatile tt_l1_ptr uint32_t* lg =
	reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
	volatile tt_l1_ptr uint32_t* src =
	reinterpret_cast<volatile tt_l1_ptr uint32_t>(lgbuf + ((j P) / 32) * TB32);
	const uint32_t acol = (j * P) % 32;
	for (uint32_t r = 0; r < 32; ++r) {
	const bool ok = b * 32 + r < NQ;
	for (uint32_t p = 0; p < P; ++p) {
	lg[tile_elem(r, p)] = ok ? src[tile_elem(r, acol + p)] : 0u;
	}
	}
	cb_push_back(cb_lg, 1);
	}
	volatile tt_l1_ptr uint32_t* fx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr);
	volatile tt_l1_ptr uint32_t* fy = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + TB32);
	#ifdef MSDA_MK_NOPTS
	for (uint32_t r = 0; r < 0; ++r) {
	#else
	for (uint32_t r = 0; r < 32; ++r) {
	#endif
	const bool row_ok = b * 32 + r < NQ;
	for (uint32_t p = P0; p < P0 + PH; ++p) {
	const uint32_t ew = tile_elem(r, p);
	const uint32_t blk = in_base + (p - P0) * 4 * TB16 + r * 128;
	if (!row_ok) {
	#ifdef MSDA_MK_FRAC16
	fx[ew] = FRAC16_ONE; // decodes to 0.0
	fy[ew] = FRAC16_ONE;
	#else
	fx[ew] = 0;
	fy[ew] = 0;
	#endif
	noc_async_read(zeros, blk, 128);
	noc_async_read(zeros, blk + 2 * TB16, 128);
	continue;
	}
	int32_t vx = loc[tile_elem(r, lcol + 2 * p)];
	int32_t vy = loc[tile_elem(r, lcol + 2 * p + 1)];
	vx = vx > LOC_LIM ? LOC_LIM : (vx < -LOC_LIM ? -LOC_LIM : vx);
	vy = vy > LOC_LIM ? LOC_LIM : (vy < -LOC_LIM ? -LOC_LIM : vy);
	const int32_t x0 = vx >> 16;
	const int32_t y0 = vy >> 16;
	#ifdef MSDA_MK_FRAC16
	// BEVFORMER_MSDA_FRAC16 (D20 round 3): the fraction q (16 bits) as the fp32 bits of 1 + q / 2^16
	// (exponent of 1.0, q in the top mantissa bits): one shift and or; the compute kernel's SFPU
	// weight step takes 1 + q / 2^16 - 1 = q / 2^16 exactly (= q16_to_f32(q), bit for bit)
	fx[ew] = FRAC16_ONE \| ((static_cast<uint32_t>(vx) & 0xFFFF) << 7);
	fy[ew] = FRAC16_ONE \| ((static_cast<uint32_t>(vy) & 0xFFFF) << 7);
	#elif defined(MSDA_MK_RAWFRAC)
	fx[ew] = static_cast<uint32_t>(vx) & 0xFFFF; // (anatomy probe: no fp32 conversion)
	fy[ew] = static_cast<uint32_t>(vy) & 0xFFFF;
	#elif defined(MSDA_MK_NOGEO)
	#else
	fx[ew] = q16_to_f32(static_cast<uint32_t>(vx) & 0xFFFF);
	fy[ew] = q16_to_f32(static_cast<uint32_t>(vy) & 0xFFFF);
	#endif
	const bool yv0 = y0 >= 0 && y0 < HV, yv1 = y0 + 1 >= 0 && y0 + 1 < HV;
	const bool xp = x0 >= -1 && x0 < WV; // page i = x0 + 1 in [0, WV]
	const uint32_t pg = vbase + static_cast<uint32_t>(y0 * (WV + 1) + x0 + 1);
	#ifdef MSDA_MK_NOREQ
	if (pg == 0xFFFFFFFF) {
	noc_async_read(zeros, blk, 128);
	}
	continue;
	#endif
	#ifdef MSDA_MK_NOREAD
	if (false) {
	#else
	if (xp && yv0) {
	#endif
	noc_async_read(v_acc.get_noc_addr(pg), blk, 128);
	} else {
	noc_async_read(zeros, blk, 128);
	}
	#ifdef MSDA_MK_NOREAD
	if (false) {
	#else
	if (xp && yv1) {
	#endif
	noc_async_read(v_acc.get_noc_addr(pg + WV + 1), blk + 2 * TB16, 128);
	} else {
	noc_async_read(zeros, blk + 2 * TB16, 128);
	}
	}
	}
	noc_async_read_barrier();
	cb_push_back(cb_geo, 3);
	cb_push_back(cb_in, NS);
	if constexpr (PART == 1) {
	if (prev != 0xFFFFFFFF) {
	cb_wait_front(cb_out, 1);
	noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(prev), TB16);
	noc_async_writes_flushed();
	cb_pop_front(cb_out, 1);
	}
	prev = b * H + j; // output tile (b, j) of [Q, H * 32]
	}
	}
	cb_pop_front(cb_loc, LOC_TW);
	if constexpr (PART == 0) {
	cb_pop_front(cb_lgc, AT_TW);
	if (a + 2 <= bl) {
	push_block_inputs(blk(a + 2), A, B, BA);
	}
	}
	}
	if constexpr (PART == 1) {
	cb_wait_front(cb_out, 1);
	noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(prev), TB16);
	noc_async_write_barrier();
	cb_pop_front(cb_out, 1);
	}
	}