code/hamer/kernels/cmm/compute_k2w.cpp · changh95/hamer-p150 at main

hamer-p150 / code /hamer /kernels /cmm /compute_k2w.cpp

Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms

58ba6df verified about 1 month ago

5.56 kB

// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// Round 11: cmm compute for the k2 K-split layout with a WIDE output block (NT > 4 tile columns per core, e.g. qkv:
// 2 token tile rows x 12 tile columns).  The 2 x NT fp32 accumulators do not fit DST, so the matmul runs in subblocks of
// SB = 4 columns (2 x 4 = 8 fp32 DST tiles, full-sync DST) and accumulates over the K blocks with packer L1
// accumulation into c_part (fp32, [row i][col n] = tile i * NT + n; block 0 overwrites).  Then:
//   - the partner row's partial (c_part tiles NT .. 2 NT - 1) is written by the BRISC into the partner's c_int;
//   - this core's own partial (tiles 0 .. NT - 1) is copied into DST exactly (c_part is unpacked to DST in fp32,
//     UnpackToDestFp32) and packed with L1 accumulation onto the partner's partial in c_int (exact fp32 sum) once the
//     NCRISC released the c_gate token (partner partial landed);
//   - + bias (row broadcast) -> c_out (the output dtype, e.g. bfp8), in chunks of 4 tiles.
// c_part is never popped (the BRISC reads its partner half after the compute is done with it; CB state resets per launch).
// compile-time: NB (K blocks of this core), KB, NT
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/bcast.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/cb_api.h"
#include "api/compute/pack.h"
#include "api/compute/reg_api.h"
#ifdef CMM_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define KZONE(n) DeviceZoneScopedN(n)
#else
#define KZONE(n)
#endif

void kernel_main() {
    constexpr uint32_t NB = get_compile_time_arg_val(0);
    constexpr uint32_t KB = get_compile_time_arg_val(1);
    constexpr uint32_t NT = get_compile_time_arg_val(2);
    constexpr uint32_t MPC = CMM_MPC;
#ifdef CMM_K2W_SB
    constexpr uint32_t SB = CMM_K2W_SB;  // 2: half-sync DST (4 fp32 tiles, math / pack overlap); 4: full-sync DST
#else
    constexpr uint32_t SB = 4;
#endif
    constexpr uint32_t DSTC = MPC * SB;  // fp32 DST tiles per acquire
    constexpr uint32_t CH = 4;  // tiles per DST acquire in the sum / bias passes (fits half- and full-sync fp32 DST)
    static_assert(NT % SB == 0 && NT % CH == 0 && DSTC <= 8, "subblocks");
    constexpr uint32_t NSB = NT / SB;
    constexpr uint32_t c_a = 0, c_w = 1, c_a1 = 5, c_w1 = 6, c_bias = 2, c_out = 16, c_int = 24, c_part = 26, c_gate = 22;

    compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_part);
    cb_reserve_back(c_gate, 1);  // arm the gate (1-page CB, full until the NCRISC pops it)
    cb_push_back(c_gate, 1);
    matmul_block_init(c_a, c_w, 0, SB, MPC, KB);
    cb_reserve_back(c_part, MPC * NT);
    for (uint32_t b = 0; b < NB; ++b) {
        KZONE("CBLK");
        const uint32_t ca = (b & 1) ? c_a1 : c_a, cw = (b & 1) ? c_w1 : c_w;
        {
            KZONE("CWA");
            cb_wait_front(ca, KB * MPC);
        }
        {
            KZONE("CWW");
            cb_wait_front(cw, KB * NT);
        }
        for (uint32_t sb = 0; sb < NSB; ++sb) {
            tile_regs_acquire();
#ifndef CMM_NO_MM
            for (uint32_t k = 0; k < KB; ++k) {
                matmul_block(ca, cw, k, k * NT + sb * SB, 0, 0, SB, MPC, KB);
            }
#endif
            tile_regs_commit();
            tile_regs_wait();
            if (b == 1 && sb == 0) {
                PACK((llk_pack_reconfig_l1_acc(1)));  // block 0 overwrote, every later block accumulates
            }
            for (uint32_t i = 0; i < MPC; ++i) {
                for (uint32_t c = 0; c < SB; ++c) {
                    pack_tile<true>(i * SB + c, c_part, i * NT + sb * SB + c);
                }
            }
            tile_regs_release();
        }
        cb_pop_front(ca, KB * MPC);
        cb_pop_front(cw, KB * NT);
    }
    PACK((llk_pack_reconfig_l1_acc(0)));
    cb_push_back(c_part, MPC * NT);

    // own partial + the partner's partial (exact fp32: unpack-to-dest copy, packer L1 accumulation)
    cb_wait_front(c_part, MPC * NT);
    {
    KZONE("CSUM");
    reconfig_data_format_srca(c_part);
    copy_tile_to_dst_init_short(c_part);
    pack_reconfig_data_format(c_int);
    {
        KZONE("CXWAIT");
        cb_reserve_back(c_gate, 1);  // the partner's partial is in c_int
    }
    cb_reserve_back(c_int, NT);
    PACK((llk_pack_reconfig_l1_acc(1)));
    for (uint32_t q = 0; q < NT; q += CH) {
        constexpr uint32_t n = CH;
        tile_regs_acquire();
        for (uint32_t t = 0; t < n; ++t) {
            copy_tile(c_part, q + t, t);
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t t = 0; t < n; ++t) {
            pack_tile<true>(t, c_int, q + t);
        }
        tile_regs_release();
    }
    PACK((llk_pack_reconfig_l1_acc(0)));
    cb_push_back(c_int, NT);
    }

    // + bias -> output
    KZONE("CBIAS");
    cb_wait_front(c_int, NT);
    cb_wait_front(c_bias, NT);
    reconfig_data_format(c_int, c_bias);
    pack_reconfig_data_format(c_out);
    add_bcast_rows_init(c_int, c_bias);
    cb_reserve_back(c_out, NT);
    for (uint32_t q = 0; q < NT; q += CH) {
        tile_regs_acquire();
        for (uint32_t i = 0; i < CH; ++i) {
            add_tiles_bcast_rows(c_int, c_bias, q + i, q + i, i);
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t i = 0; i < CH; ++i) {
            pack_tile(i, c_out);
        }
        tile_regs_release();
    }
    cb_push_back(c_out, NT);
    cb_pop_front(c_int, NT);
    cb_pop_front(c_bias, NT);
}

	// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
	// SPDX-License-Identifier: Apache-2.0
	// Round 11: cmm compute for the k2 K-split layout with a WIDE output block (NT > 4 tile columns per core, e.g. qkv:
	// 2 token tile rows x 12 tile columns). The 2 x NT fp32 accumulators do not fit DST, so the matmul runs in subblocks of
	// SB = 4 columns (2 x 4 = 8 fp32 DST tiles, full-sync DST) and accumulates over the K blocks with packer L1
	// accumulation into c_part (fp32, [row i][col n] = tile i * NT + n; block 0 overwrites). Then:
	// - the partner row's partial (c_part tiles NT .. 2 NT - 1) is written by the BRISC into the partner's c_int;
	// - this core's own partial (tiles 0 .. NT - 1) is copied into DST exactly (c_part is unpacked to DST in fp32,
	// UnpackToDestFp32) and packed with L1 accumulation onto the partner's partial in c_int (exact fp32 sum) once the
	// NCRISC released the c_gate token (partner partial landed);
	// - + bias (row broadcast) -> c_out (the output dtype, e.g. bfp8), in chunks of 4 tiles.
	// c_part is never popped (the BRISC reads its partner half after the compute is done with it; CB state resets per launch).
	// compile-time: NB (K blocks of this core), KB, NT
	#include <cstdint>
	#include "api/compute/matmul.h"
	#include "api/compute/bcast.h"
	#include "api/compute/tile_move_copy.h"
	#include "api/compute/compute_kernel_hw_startup.h"
	#include "api/compute/cb_api.h"
	#include "api/compute/pack.h"
	#include "api/compute/reg_api.h"
	#ifdef CMM_PROF
	#include "tools/profiler/kernel_profiler.hpp"
	#define KZONE(n) DeviceZoneScopedN(n)
	#else
	#define KZONE(n)
	#endif

	void kernel_main() {
	constexpr uint32_t NB = get_compile_time_arg_val(0);
	constexpr uint32_t KB = get_compile_time_arg_val(1);
	constexpr uint32_t NT = get_compile_time_arg_val(2);
	constexpr uint32_t MPC = CMM_MPC;
	#ifdef CMM_K2W_SB
	constexpr uint32_t SB = CMM_K2W_SB; // 2: half-sync DST (4 fp32 tiles, math / pack overlap); 4: full-sync DST
	#else
	constexpr uint32_t SB = 4;
	#endif
	constexpr uint32_t DSTC = MPC * SB; // fp32 DST tiles per acquire
	constexpr uint32_t CH = 4; // tiles per DST acquire in the sum / bias passes (fits half- and full-sync fp32 DST)
	static_assert(NT % SB == 0 && NT % CH == 0 && DSTC <= 8, "subblocks");
	constexpr uint32_t NSB = NT / SB;
	constexpr uint32_t c_a = 0, c_w = 1, c_a1 = 5, c_w1 = 6, c_bias = 2, c_out = 16, c_int = 24, c_part = 26, c_gate = 22;

	compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_part);
	cb_reserve_back(c_gate, 1); // arm the gate (1-page CB, full until the NCRISC pops it)
	cb_push_back(c_gate, 1);
	matmul_block_init(c_a, c_w, 0, SB, MPC, KB);
	cb_reserve_back(c_part, MPC * NT);
	for (uint32_t b = 0; b < NB; ++b) {
	KZONE("CBLK");
	const uint32_t ca = (b & 1) ? c_a1 : c_a, cw = (b & 1) ? c_w1 : c_w;
	{
	KZONE("CWA");
	cb_wait_front(ca, KB * MPC);
	}
	{
	KZONE("CWW");
	cb_wait_front(cw, KB * NT);
	}
	for (uint32_t sb = 0; sb < NSB; ++sb) {
	tile_regs_acquire();
	#ifndef CMM_NO_MM
	for (uint32_t k = 0; k < KB; ++k) {
	matmul_block(ca, cw, k, k * NT + sb * SB, 0, 0, SB, MPC, KB);
	}
	#endif
	tile_regs_commit();
	tile_regs_wait();
	if (b == 1 && sb == 0) {
	PACK((llk_pack_reconfig_l1_acc(1))); // block 0 overwrote, every later block accumulates
	}
	for (uint32_t i = 0; i < MPC; ++i) {
	for (uint32_t c = 0; c < SB; ++c) {
	pack_tile<true>(i * SB + c, c_part, i * NT + sb * SB + c);
	}
	}
	tile_regs_release();
	}
	cb_pop_front(ca, KB * MPC);
	cb_pop_front(cw, KB * NT);
	}
	PACK((llk_pack_reconfig_l1_acc(0)));
	cb_push_back(c_part, MPC * NT);

	// own partial + the partner's partial (exact fp32: unpack-to-dest copy, packer L1 accumulation)
	cb_wait_front(c_part, MPC * NT);
	{
	KZONE("CSUM");
	reconfig_data_format_srca(c_part);
	copy_tile_to_dst_init_short(c_part);
	pack_reconfig_data_format(c_int);
	{
	KZONE("CXWAIT");
	cb_reserve_back(c_gate, 1); // the partner's partial is in c_int
	}
	cb_reserve_back(c_int, NT);
	PACK((llk_pack_reconfig_l1_acc(1)));
	for (uint32_t q = 0; q < NT; q += CH) {
	constexpr uint32_t n = CH;
	tile_regs_acquire();
	for (uint32_t t = 0; t < n; ++t) {
	copy_tile(c_part, q + t, t);
	}
	tile_regs_commit();
	tile_regs_wait();
	for (uint32_t t = 0; t < n; ++t) {
	pack_tile<true>(t, c_int, q + t);
	}
	tile_regs_release();
	}
	PACK((llk_pack_reconfig_l1_acc(0)));
	cb_push_back(c_int, NT);
	}

	// + bias -> output
	KZONE("CBIAS");
	cb_wait_front(c_int, NT);
	cb_wait_front(c_bias, NT);
	reconfig_data_format(c_int, c_bias);
	pack_reconfig_data_format(c_out);
	add_bcast_rows_init(c_int, c_bias);
	cb_reserve_back(c_out, NT);
	for (uint32_t q = 0; q < NT; q += CH) {
	tile_regs_acquire();
	for (uint32_t i = 0; i < CH; ++i) {
	add_tiles_bcast_rows(c_int, c_bias, q + i, q + i, i);
	}
	tile_regs_commit();
	tile_regs_wait();
	for (uint32_t i = 0; i < CH; ++i) {
	pack_tile(i, c_out);
	}
	tile_regs_release();
	}
	cb_push_back(c_out, NT);
	cb_pop_front(c_int, NT);
	cb_pop_front(c_bias, NT);
	}