/*
 * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Generated by tools/export-generated-programs (device kernels), merged into a
// single translation unit following the csrc/tinygemm2.cu convention (one TU =
// kernel family + TVM-FFI binding).
//
// Provenance: four generated Loom schedules, each an exact port of the
// TensorRT-LLM tinygemm2 kernel (csrc/tinygemm2.cu) with bit-identical
// outputs:
//   stage4     — schedule 'flashinfer_tinygemm2_stage4' (USE_PDL=0)
//   stage4_pdl — schedule 'flashinfer_tinygemm2_stage4' (USE_PDL=1)
//   stage8     — schedule 'flashinfer_tinygemm2'        (USE_PDL=0)
//   stage8_pdl — schedule 'flashinfer_tinygemm2'        (USE_PDL=1)
//
// The four generated TUs are concatenated here by a mechanical transform that
// does not touch kernel code:
//   1. the extern "C" symbol kernel_flashinfer_tinygemm2 is renamed
//      kernel_tinygemm2_sm100_<variant> (one rename per section);
//   2. each section keeps its generated #define block verbatim and #undef's
//      every section-local macro at the section end; SMEM_TOTAL/USE_PDL are
//      captured into per-variant constexprs for the launcher first;
//   3. the fixed-width typedefs and the opaque CUtensorMap typedef of the
//      generated prelude are dropped (the host headers below supply them),
//      and the LoomTensorMap struct + __device__ helpers — byte-identical
//      across the four generated TUs — are deduplicated to one copy.
// Kernel bodies are byte-identical to the generated output modulo the symbol
// rename.

#include <cuda.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <math_constants.h>

#include <cstdint>
#include <cstring>
#include <limits>

#include "tvm_ffi_utils.h"

// clang-format off

// ---- Generated device prelude (shared by all four variant sections) ----
struct __align__(128) LoomTensorMap { uint64_t opaque[16]; };


__device__ __forceinline__ int make_warp_uniform(int x) {
    int result;
    asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;"
                 : "=r"(result) : "r"(x));
    return result;
}

__device__ __forceinline__ uint32_t elect_sync() {
    uint32_t pred = 0;
    asm volatile(
        "{\n\t"
        ".reg .pred %%px;\n\t"
        "elect.sync _|%%px, %1;\n\t"
        "@%%px mov.s32 %0, 1;\n\t"
        "}\n"
        : "+r"(pred)
        : "r"(0xFFFFFFFF));
    return pred;
}


__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;"
        :: "r"(mbar_addr), "r"(count));
}


__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) {
    uint32_t token;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64"
        " P1, [%1], %2;\n\t"
        "selp.u32 %0, 1, 0, P1;\n\t"
        "}\n"
        : "=r"(token)
        : "r"(mbar_addr), "r"(phase) : "memory");
    return token;
}

__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) {
    uint32_t token;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64"
        " P1, [%1], %2;\n\t"
        "selp.u32 %0, 1, 0, P1;\n\t"
        "}\n"
        : "=r"(token)
        : "r"(mbar_addr), "r"(phase) : "memory");
    return token;
}

__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) {
    uint32_t ticks = 0x989680;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "LAB_WAIT:\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64"
        " P1, [%0], %1, %2;\n\t"
        "@P1 bra.uni DONE;\n\t"
        "bra.uni LAB_WAIT;\n\t"
        "DONE:\n\t"
        "}\n"
        :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory");
}

__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) {
    uint32_t ticks = 0x989680;
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "LAB_WAIT_CLUSTER:\n\t"
        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64"
        " P1, [%0], %1, %2;\n\t"
        "@P1 bra.uni DONE_CLUSTER;\n\t"
        "bra.uni LAB_WAIT_CLUSTER;\n\t"
        "DONE_CLUSTER:\n\t"
        "}\n"
        :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory");
}

__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) {
    if (token == 0) {
        mbarrier_wait(mbar_addr, phase);
    }
}

__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) {
    if (token == 0) {
        mbarrier_wait_cluster(mbar_addr, phase);
    }
}


__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) {
    asm volatile(
        "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];"
        :: "r"(mbar_addr) : "memory");
}


__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) {
    asm volatile(
        "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
        :: "r"(mbar_addr), "r"(bytes) : "memory");
}


__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) {
    asm volatile(
        "{\n\t"
        ".reg .pred p;\n\t"
        "setp.ne.b32 p, %2, 0;\n\t"
        "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t"
        "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred));
}


__device__ __forceinline__ void fence_async_shared() {
    asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
}


__device__ __forceinline__ uint64_t desc_encode(uint64_t x) {
    return (x & 0x3FFFFULL) >> 4ULL;
}


__device__ __forceinline__ uint64_t make_smem_desc(int addr) {
    const int SBO = 1024;
    return desc_encode(addr)
         | (desc_encode(SBO) << 32ULL)
         | (1ULL << 46ULL)
         | (2ULL << 61ULL);
}


__device__ __forceinline__ void tma_2d_gmem2smem(
    int dst, const void *tmap_ptr, int x, int y, int mbar_addr) {
    asm volatile(
        "cp.async.bulk.tensor.2d.shared::cta.global"
        ".mbarrier::complete_tx::bytes"
        " [%0], [%1, {%2, %3}], [%4];"
        :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y),
           "r"(mbar_addr) : "memory");
}


__device__ __forceinline__ unsigned int __as_u32(float v) {
    unsigned int u;
    asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v));
    return u;
}
__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; }
__device__ __forceinline__ unsigned int __as_u32(int v) {
    unsigned int u;
    asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v));
    return u;
}

// ============================================================================
// Section stage4 — generated Loom schedule 'flashinfer_tinygemm2_stage4' (USE_PDL=0).
// ============================================================================

#define LOOM_INF CUDART_INF_F
#define NUM_MAIN_STAGES 1
#define SMEM_SMEM_WT_OFF 1024
#define SMEM_SMEM_WT_STAGE_BYTES 2048
#define SMEM_SMEM_WT_STRIDE 2048
#define SMEM_SMEM_ACT_OFF 33792
#define SMEM_SMEM_ACT_STAGE_BYTES 1024
#define SMEM_SMEM_ACT_STRIDE 1024
#define SMEM_SMEM_RED_OFF 50176
#define SMEM_SMEM_RED_STAGE_BYTES 2048
#define SMEM_SMEM_RED_STRIDE 2048
#define SMEM_SMEM_BIAS_OFF 52224
#define SMEM_SMEM_BIAS_STAGE_BYTES 32
#define SMEM_SMEM_BIAS_STRIDE 32
#define SMEM_TOTAL 52352
#define THREADS 384
#define USE_PDL 0

extern "C" {

__global__ __launch_bounds__(384, 1) void
kernel_tinygemm2_sm100_stage4(__grid_constant__ LoomTensorMap const tmap_wt, __grid_constant__ LoomTensorMap const tmap_act, __nv_bfloat16* __restrict__ output, __nv_bfloat16* __restrict__ bias, int M, int N, int K)
{
    const int tid = threadIdx.x;
    const int warp = make_warp_uniform(tid / 32);
    const int lane = tid % 32;

    extern __shared__ __align__(1024) char smem_raw[];
    int smem;
    smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw);

    const int bid = blockIdx.x;
    const int num_bids = gridDim.x;

    // Kernel setup ops
    __nv_bfloat16* smem_wt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024);
    const int smem_wt_addr = smem + 1024;
    __nv_bfloat16* smem_act = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792);
    const int smem_act_addr = smem + 33792;
    float* smem_red = reinterpret_cast<float*>(smem_raw + 50176);
    const int smem_red_addr = smem + 50176;
    __nv_bfloat16* smem_bias = reinterpret_cast<__nv_bfloat16*>(smem_raw + 52224);
    const int smem_bias_addr = smem + 52224;
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_wt))) : "memory"); }
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_act))) : "memory"); }

    // Mbarrier init (3 groups, 12 barriers)
    // Mbarriers at smem_raw[0..96)

    if (warp == 0) {
        uint32_t leader = elect_sync();
        // wt_ready: 4 barriers, init_count=1
        mbarrier_init_pred(smem + 0, 1, leader);
        mbarrier_init_pred(smem + 8, 1, leader);
        mbarrier_init_pred(smem + 16, 1, leader);
        mbarrier_init_pred(smem + 24, 1, leader);
        // act_ready: 4 barriers, init_count=1
        mbarrier_init_pred(smem + 32, 1, leader);
        mbarrier_init_pred(smem + 40, 1, leader);
        mbarrier_init_pred(smem + 48, 1, leader);
        mbarrier_init_pred(smem + 56, 1, leader);
        // data_consumed: 4 barriers, init_count=32
        mbarrier_init_pred(smem + 64, 32, leader);
        mbarrier_init_pred(smem + 72, 32, leader);
        mbarrier_init_pred(smem + 80, 32, leader);
        mbarrier_init_pred(smem + 88, 32, leader);
        asm volatile("fence.mbarrier_init.release.cluster;");
    }

    __syncthreads();

    __syncthreads();

    const int mbar_base = smem;
    #define wt_ready_addr (mbar_base + 0)
    #define act_ready_addr (mbar_base + 32)
    #define data_consumed_addr (mbar_base + 64)

    // ---- Role: compute ----
    if (warp <= 3) {
        { // compute_main
            int k_loops_c = (K + 1024 - 1) / 1024;
            int mib_c = blockIdx.x * 16;
            int ni_c = blockIdx.y * 8;
            if (tid < 16) {
                smem_bias[tid] = bias[mib_c + tid];
            }
            float accum[4];
            #pragma unroll
            for (int z = 0; z < 4; z++) {
                accum[z] = 0.0f;
            }
            unsigned int lane_div8 = lane / 8;
            unsigned int lane_mod8 = lane % 8;
            unsigned int row_wt = lane_mod8 + lane_div8 % 2 * 8;
            unsigned int col_off_wt = lane_div8 / 2;
            unsigned int row_act = lane_mod8;
            #pragma unroll 2
            for (unsigned int ki = 0; ki < k_loops_c; ki++) {
                unsigned int stage_c = (unsigned int)warp;
                unsigned int phase_c = ki & 1;
                mbarrier_wait(wt_ready_addr + (stage_c) * 8, phase_c);
                mbarrier_wait(act_ready_addr + (stage_c) * 8, phase_c);
                #pragma unroll
                for (int su = 0; su < 4; su++) {
                    unsigned int base_wt = smem_wt_addr + (stage_c * 4 + (unsigned int)su) * 2048;
                    unsigned int base_act = smem_act_addr + (stage_c * 4 + (unsigned int)su) * 1024;
                    #pragma unroll
                    for (int kii = 0; kii < 4; kii++) {
                        unsigned int a_frag[4];
                        unsigned int b_frag[2];
                        unsigned int col_w = (unsigned int)(2 * kii) + col_off_wt;
                        unsigned int col_sw_w = row_wt % 8 ^ col_w;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
                            : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3])
                            : "r"(base_wt + row_wt * 128 + col_sw_w * 16)
                            : "memory");
                        unsigned int col_a = (unsigned int)(2 * kii) + lane_div8;
                        unsigned int col_sw_a = row_act % 8 ^ col_a;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
                            : "=r"(b_frag[0]), "=r"(b_frag[1])
                            : "r"(base_act + row_act * 128 + col_sw_a * 16)
                            : "memory");
                        asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
                            : "+f"(accum[0]), "+f"(accum[1]), "+f"(accum[2]), "+f"(accum[3])
                            : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1]));
                    }
                }
                asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
                mbarrier_arrive(data_consumed_addr + (stage_c) * 8);
            }
            asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_red_addr + (unsigned int)(tid * 16)), "r"(__as_u32(accum[0])), "r"(__as_u32(accum[1])), "r"(__as_u32(accum[2])), "r"(__as_u32(accum[3])) : "memory");
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
            if (warp == 0) {
                float part[12];
                #pragma unroll
                for (int w = 0; w < 3; w++) {
                    asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];"
                        : "=r"(*reinterpret_cast<uint32_t*>(&part[w * 4])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 1])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 2])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 3]))
                        : "r"(smem_red_addr + (unsigned int)((32 + w * 32 + tid) * 16)));
                }
                #pragma unroll
                for (int z_1 = 0; z_1 < 4; z_1++) {
                    accum[z_1] = accum[z_1] + part[z_1] + part[4 + z_1] + part[8 + z_1];
                }
                int tm = mib_c + lane / 4;
                int tn = ni_c + 2 * (lane % 4);
                float bias_lo = smem_bias[lane / 4];
                float bias_hi = smem_bias[lane / 4 + 8];
                float o00 = accum[0] + bias_lo;
                float o01 = accum[1] + bias_lo;
                float o10 = accum[2] + bias_hi;
                float o11 = accum[3] + bias_hi;
                if (tn < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm))) = __float2bfloat16_rn(o00);
                    }
                }
                if (tn + 1 < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm))) = __float2bfloat16_rn(o01);
                    }
                }
                if (tn < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm + 8))) = __float2bfloat16_rn(o10);
                    }
                }
                if (tn + 1 < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm + 8))) = __float2bfloat16_rn(o11);
                    }
                }
            }
        }
    // ---- Role: load_wt ----
    } else if (warp >= 4 && warp <= 7) {
        { // load_wt_main
            int k_loops = (K + 1024 - 1) / 1024;
            int mib = blockIdx.x * 16;
            unsigned int wslot = warp % 4;
            if (elect_sync()) {
                #pragma unroll 1
                for (unsigned int ki_1 = 0; ki_1 < k_loops; ki_1++) {
                    unsigned int stage = wslot;
                    unsigned int phase = ki_1 & 1;
                    int k_base = (ki_1 * 4 + wslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage) * 8, phase ^ 1);
                    mbarrier_arrive_expect_tx(wt_ready_addr + (stage) * 8, 8192);
                    #pragma unroll
                    for (int i = 0; i < 4; i++) {
                        int dst_wt = smem_wt_addr + (stage * 4 + (unsigned int)i) * 2048;
                        tma_2d_gmem2smem(dst_wt, (&tmap_wt), k_base + i * 64, mib, wt_ready_addr + (stage) * 8);
                    }
                }
                #pragma unroll
                for (int di = 0; di < 1; di++) {
                    if (di + 1 < 1) {
                        unsigned int dki = k_loops + di;
                        unsigned int dstage = wslot;
                        unsigned int dphase = dki & 1;
                        mbarrier_wait(data_consumed_addr + (dstage) * 8, dphase ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    // ---- Role: load_act ----
    } else if (warp >= 8 && warp <= 11) {
        { // load_act_main
            int k_loops_a = (K + 1024 - 1) / 1024;
            int ni = blockIdx.y * 8;
            unsigned int aslot = warp % 4;
            if (elect_sync()) {
                #pragma unroll 1
                for (unsigned int ki_2 = 0; ki_2 < k_loops_a; ki_2++) {
                    unsigned int stage_a = aslot;
                    unsigned int phase_a = ki_2 & 1;
                    int k_base_a = (ki_2 * 4 + aslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage_a) * 8, phase_a ^ 1);
                    mbarrier_arrive_expect_tx(act_ready_addr + (stage_a) * 8, 4096);
                    #pragma unroll
                    for (int i_1 = 0; i_1 < 4; i_1++) {
                        int dst_act = smem_act_addr + (stage_a * 4 + (unsigned int)i_1) * 1024;
                        tma_2d_gmem2smem(dst_act, (&tmap_act), k_base_a + i_1 * 64, ni, act_ready_addr + (stage_a) * 8);
                    }
                }
                #pragma unroll
                for (int di_1 = 0; di_1 < 1; di_1++) {
                    if (di_1 + 1 < 1) {
                        unsigned int dki_a = k_loops_a + di_1;
                        unsigned int dstage_a = aslot;
                        unsigned int dphase_a = dki_a & 1;
                        mbarrier_wait(data_consumed_addr + (dstage_a) * 8, dphase_a ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    }

    // Cleanup
}

} // extern "C"

// End of generated section stage4: capture the launch metadata, then unscope
// the section-local macros so the next section can redefine them.
namespace flashinfer {
namespace tinygemm2_sm100 {
static_assert(THREADS == 384, "generated stage4 variant disagrees on block size");
static_assert(USE_PDL == 0, "generated stage4 variant must be the non-PDL flavor");
constexpr int kSmemBytesStage4 = SMEM_TOTAL;
static_assert(kSmemBytesStage4 == 52352,
              "generated stage4 variant disagrees on dynamic smem footprint");
}  // namespace tinygemm2_sm100
}  // namespace flashinfer
#undef LOOM_INF
#undef NUM_MAIN_STAGES
#undef SMEM_SMEM_WT_OFF
#undef SMEM_SMEM_WT_STAGE_BYTES
#undef SMEM_SMEM_WT_STRIDE
#undef SMEM_SMEM_ACT_OFF
#undef SMEM_SMEM_ACT_STAGE_BYTES
#undef SMEM_SMEM_ACT_STRIDE
#undef SMEM_SMEM_RED_OFF
#undef SMEM_SMEM_RED_STAGE_BYTES
#undef SMEM_SMEM_RED_STRIDE
#undef SMEM_SMEM_BIAS_OFF
#undef SMEM_SMEM_BIAS_STAGE_BYTES
#undef SMEM_SMEM_BIAS_STRIDE
#undef SMEM_TOTAL
#undef THREADS
#undef USE_PDL
#undef wt_ready_addr
#undef act_ready_addr
#undef data_consumed_addr

// ============================================================================
// Section stage4_pdl — generated Loom schedule 'flashinfer_tinygemm2_stage4' (USE_PDL=1).
// ============================================================================

#define LOOM_INF CUDART_INF_F
#define NUM_MAIN_STAGES 1
#define SMEM_SMEM_WT_OFF 1024
#define SMEM_SMEM_WT_STAGE_BYTES 2048
#define SMEM_SMEM_WT_STRIDE 2048
#define SMEM_SMEM_ACT_OFF 33792
#define SMEM_SMEM_ACT_STAGE_BYTES 1024
#define SMEM_SMEM_ACT_STRIDE 1024
#define SMEM_SMEM_RED_OFF 50176
#define SMEM_SMEM_RED_STAGE_BYTES 2048
#define SMEM_SMEM_RED_STRIDE 2048
#define SMEM_SMEM_BIAS_OFF 52224
#define SMEM_SMEM_BIAS_STAGE_BYTES 32
#define SMEM_SMEM_BIAS_STRIDE 32
#define SMEM_TOTAL 52352
#define THREADS 384
#define USE_PDL 1

extern "C" {

__global__ __launch_bounds__(384, 1) void
kernel_tinygemm2_sm100_stage4_pdl(__grid_constant__ LoomTensorMap const tmap_wt, __grid_constant__ LoomTensorMap const tmap_act, __nv_bfloat16* __restrict__ output, __nv_bfloat16* __restrict__ bias, int M, int N, int K)
{
    const int tid = threadIdx.x;
    const int warp = make_warp_uniform(tid / 32);
    const int lane = tid % 32;

    extern __shared__ __align__(1024) char smem_raw[];
    int smem;
    smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw);

    const int bid = blockIdx.x;
    const int num_bids = gridDim.x;

    // Kernel setup ops
    __nv_bfloat16* smem_wt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024);
    const int smem_wt_addr = smem + 1024;
    __nv_bfloat16* smem_act = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792);
    const int smem_act_addr = smem + 33792;
    float* smem_red = reinterpret_cast<float*>(smem_raw + 50176);
    const int smem_red_addr = smem + 50176;
    __nv_bfloat16* smem_bias = reinterpret_cast<__nv_bfloat16*>(smem_raw + 52224);
    const int smem_bias_addr = smem + 52224;
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_wt))) : "memory"); }
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_act))) : "memory"); }

    // Mbarrier init (3 groups, 12 barriers)
    // Mbarriers at smem_raw[0..96)

    if (warp == 0) {
        uint32_t leader = elect_sync();
        // wt_ready: 4 barriers, init_count=1
        mbarrier_init_pred(smem + 0, 1, leader);
        mbarrier_init_pred(smem + 8, 1, leader);
        mbarrier_init_pred(smem + 16, 1, leader);
        mbarrier_init_pred(smem + 24, 1, leader);
        // act_ready: 4 barriers, init_count=1
        mbarrier_init_pred(smem + 32, 1, leader);
        mbarrier_init_pred(smem + 40, 1, leader);
        mbarrier_init_pred(smem + 48, 1, leader);
        mbarrier_init_pred(smem + 56, 1, leader);
        // data_consumed: 4 barriers, init_count=32
        mbarrier_init_pred(smem + 64, 32, leader);
        mbarrier_init_pred(smem + 72, 32, leader);
        mbarrier_init_pred(smem + 80, 32, leader);
        mbarrier_init_pred(smem + 88, 32, leader);
        asm volatile("fence.mbarrier_init.release.cluster;");
    }

    __syncthreads();

    __syncthreads();

    const int mbar_base = smem;
    #define wt_ready_addr (mbar_base + 0)
    #define act_ready_addr (mbar_base + 32)
    #define data_consumed_addr (mbar_base + 64)

    // ---- Role: compute ----
    if (warp <= 3) {
        { // compute_main
            int k_loops_c = (K + 1024 - 1) / 1024;
            int mib_c = blockIdx.x * 16;
            int ni_c = blockIdx.y * 8;
            if (tid < 16) {
                smem_bias[tid] = bias[mib_c + tid];
            }
            float accum[4];
            #pragma unroll
            for (int z = 0; z < 4; z++) {
                accum[z] = 0.0f;
            }
            unsigned int lane_div8 = lane / 8;
            unsigned int lane_mod8 = lane % 8;
            unsigned int row_wt = lane_mod8 + lane_div8 % 2 * 8;
            unsigned int col_off_wt = lane_div8 / 2;
            unsigned int row_act = lane_mod8;
            #pragma unroll 2
            for (unsigned int ki = 0; ki < k_loops_c; ki++) {
                unsigned int stage_c = (unsigned int)warp;
                unsigned int phase_c = ki & 1;
                mbarrier_wait(wt_ready_addr + (stage_c) * 8, phase_c);
                mbarrier_wait(act_ready_addr + (stage_c) * 8, phase_c);
                #pragma unroll
                for (int su = 0; su < 4; su++) {
                    unsigned int base_wt = smem_wt_addr + (stage_c * 4 + (unsigned int)su) * 2048;
                    unsigned int base_act = smem_act_addr + (stage_c * 4 + (unsigned int)su) * 1024;
                    #pragma unroll
                    for (int kii = 0; kii < 4; kii++) {
                        unsigned int a_frag[4];
                        unsigned int b_frag[2];
                        unsigned int col_w = (unsigned int)(2 * kii) + col_off_wt;
                        unsigned int col_sw_w = row_wt % 8 ^ col_w;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
                            : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3])
                            : "r"(base_wt + row_wt * 128 + col_sw_w * 16)
                            : "memory");
                        unsigned int col_a = (unsigned int)(2 * kii) + lane_div8;
                        unsigned int col_sw_a = row_act % 8 ^ col_a;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
                            : "=r"(b_frag[0]), "=r"(b_frag[1])
                            : "r"(base_act + row_act * 128 + col_sw_a * 16)
                            : "memory");
                        asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
                            : "+f"(accum[0]), "+f"(accum[1]), "+f"(accum[2]), "+f"(accum[3])
                            : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1]));
                    }
                }
                asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
                mbarrier_arrive(data_consumed_addr + (stage_c) * 8);
            }
            asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_red_addr + (unsigned int)(tid * 16)), "r"(__as_u32(accum[0])), "r"(__as_u32(accum[1])), "r"(__as_u32(accum[2])), "r"(__as_u32(accum[3])) : "memory");
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
            if (warp == 0) {
                float part[12];
                #pragma unroll
                for (int w = 0; w < 3; w++) {
                    asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];"
                        : "=r"(*reinterpret_cast<uint32_t*>(&part[w * 4])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 1])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 2])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 3]))
                        : "r"(smem_red_addr + (unsigned int)((32 + w * 32 + tid) * 16)));
                }
                #pragma unroll
                for (int z_1 = 0; z_1 < 4; z_1++) {
                    accum[z_1] = accum[z_1] + part[z_1] + part[4 + z_1] + part[8 + z_1];
                }
                int tm = mib_c + lane / 4;
                int tn = ni_c + 2 * (lane % 4);
                float bias_lo = smem_bias[lane / 4];
                float bias_hi = smem_bias[lane / 4 + 8];
                float o00 = accum[0] + bias_lo;
                float o01 = accum[1] + bias_lo;
                float o10 = accum[2] + bias_hi;
                float o11 = accum[3] + bias_hi;
                if (tn < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm))) = __float2bfloat16_rn(o00);
                    }
                }
                if (tn + 1 < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm))) = __float2bfloat16_rn(o01);
                    }
                }
                if (tn < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm + 8))) = __float2bfloat16_rn(o10);
                    }
                }
                if (tn + 1 < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm + 8))) = __float2bfloat16_rn(o11);
                    }
                }
            }
        }
    // ---- Role: load_wt ----
    } else if (warp >= 4 && warp <= 7) {
        { // load_wt_main
            int k_loops = (K + 1024 - 1) / 1024;
            int mib = blockIdx.x * 16;
            unsigned int wslot = warp % 4;
            if (elect_sync()) {
                #pragma unroll 1
                for (unsigned int ki_1 = 0; ki_1 < k_loops; ki_1++) {
                    unsigned int stage = wslot;
                    unsigned int phase = ki_1 & 1;
                    int k_base = (ki_1 * 4 + wslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage) * 8, phase ^ 1);
                    mbarrier_arrive_expect_tx(wt_ready_addr + (stage) * 8, 8192);
                    #pragma unroll
                    for (int i = 0; i < 4; i++) {
                        int dst_wt = smem_wt_addr + (stage * 4 + (unsigned int)i) * 2048;
                        tma_2d_gmem2smem(dst_wt, (&tmap_wt), k_base + i * 64, mib, wt_ready_addr + (stage) * 8);
                    }
                }
                #pragma unroll
                for (int di = 0; di < 1; di++) {
                    if (di + 1 < 1) {
                        unsigned int dki = k_loops + di;
                        unsigned int dstage = wslot;
                        unsigned int dphase = dki & 1;
                        mbarrier_wait(data_consumed_addr + (dstage) * 8, dphase ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    // ---- Role: load_act ----
    } else if (warp >= 8 && warp <= 11) {
        { // load_act_main
            int k_loops_a = (K + 1024 - 1) / 1024;
            int ni = blockIdx.y * 8;
            unsigned int aslot = warp % 4;
            if (elect_sync()) {
                {
                    asm volatile("griddepcontrol.wait;" ::: "memory");
                    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
                }
                #pragma unroll 1
                for (unsigned int ki_2 = 0; ki_2 < k_loops_a; ki_2++) {
                    unsigned int stage_a = aslot;
                    unsigned int phase_a = ki_2 & 1;
                    int k_base_a = (ki_2 * 4 + aslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage_a) * 8, phase_a ^ 1);
                    mbarrier_arrive_expect_tx(act_ready_addr + (stage_a) * 8, 4096);
                    #pragma unroll
                    for (int i_1 = 0; i_1 < 4; i_1++) {
                        int dst_act = smem_act_addr + (stage_a * 4 + (unsigned int)i_1) * 1024;
                        tma_2d_gmem2smem(dst_act, (&tmap_act), k_base_a + i_1 * 64, ni, act_ready_addr + (stage_a) * 8);
                    }
                }
                #pragma unroll
                for (int di_1 = 0; di_1 < 1; di_1++) {
                    if (di_1 + 1 < 1) {
                        unsigned int dki_a = k_loops_a + di_1;
                        unsigned int dstage_a = aslot;
                        unsigned int dphase_a = dki_a & 1;
                        mbarrier_wait(data_consumed_addr + (dstage_a) * 8, dphase_a ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    }

    // Cleanup
}

} // extern "C"

// End of generated section stage4_pdl: capture the launch metadata, then unscope
// the section-local macros so the next section can redefine them.
namespace flashinfer {
namespace tinygemm2_sm100 {
static_assert(THREADS == 384, "generated stage4_pdl variant disagrees on block size");
static_assert(USE_PDL == 1, "generated stage4_pdl variant must be the PDL flavor");
constexpr int kSmemBytesStage4Pdl = SMEM_TOTAL;
static_assert(kSmemBytesStage4Pdl == 52352,
              "generated stage4_pdl variant disagrees on dynamic smem footprint");
}  // namespace tinygemm2_sm100
}  // namespace flashinfer
#undef LOOM_INF
#undef NUM_MAIN_STAGES
#undef SMEM_SMEM_WT_OFF
#undef SMEM_SMEM_WT_STAGE_BYTES
#undef SMEM_SMEM_WT_STRIDE
#undef SMEM_SMEM_ACT_OFF
#undef SMEM_SMEM_ACT_STAGE_BYTES
#undef SMEM_SMEM_ACT_STRIDE
#undef SMEM_SMEM_RED_OFF
#undef SMEM_SMEM_RED_STAGE_BYTES
#undef SMEM_SMEM_RED_STRIDE
#undef SMEM_SMEM_BIAS_OFF
#undef SMEM_SMEM_BIAS_STAGE_BYTES
#undef SMEM_SMEM_BIAS_STRIDE
#undef SMEM_TOTAL
#undef THREADS
#undef USE_PDL
#undef wt_ready_addr
#undef act_ready_addr
#undef data_consumed_addr

// ============================================================================
// Section stage8 — generated Loom schedule 'flashinfer_tinygemm2' (USE_PDL=0).
// ============================================================================

#define LOOM_INF CUDART_INF_F
#define NUM_MAIN_STAGES 1
#define SMEM_SMEM_WT_OFF 1024
#define SMEM_SMEM_WT_STAGE_BYTES 2048
#define SMEM_SMEM_WT_STRIDE 2048
#define SMEM_SMEM_ACT_OFF 66560
#define SMEM_SMEM_ACT_STAGE_BYTES 1024
#define SMEM_SMEM_ACT_STRIDE 1024
#define SMEM_SMEM_RED_OFF 99328
#define SMEM_SMEM_RED_STAGE_BYTES 2048
#define SMEM_SMEM_RED_STRIDE 2048
#define SMEM_SMEM_BIAS_OFF 101376
#define SMEM_SMEM_BIAS_STAGE_BYTES 32
#define SMEM_SMEM_BIAS_STRIDE 32
#define SMEM_TOTAL 101504
#define THREADS 384
#define USE_PDL 0

extern "C" {

__global__ __launch_bounds__(384, 1) void
kernel_tinygemm2_sm100_stage8(__grid_constant__ LoomTensorMap const tmap_wt, __grid_constant__ LoomTensorMap const tmap_act, __nv_bfloat16* __restrict__ output, __nv_bfloat16* __restrict__ bias, int M, int N, int K)
{
    const int tid = threadIdx.x;
    const int warp = make_warp_uniform(tid / 32);
    const int lane = tid % 32;

    extern __shared__ __align__(1024) char smem_raw[];
    int smem;
    smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw);

    const int bid = blockIdx.x;
    const int num_bids = gridDim.x;

    // Kernel setup ops
    __nv_bfloat16* smem_wt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024);
    const int smem_wt_addr = smem + 1024;
    __nv_bfloat16* smem_act = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560);
    const int smem_act_addr = smem + 66560;
    float* smem_red = reinterpret_cast<float*>(smem_raw + 99328);
    const int smem_red_addr = smem + 99328;
    __nv_bfloat16* smem_bias = reinterpret_cast<__nv_bfloat16*>(smem_raw + 101376);
    const int smem_bias_addr = smem + 101376;
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_wt))) : "memory"); }
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_act))) : "memory"); }

    // Mbarrier init (3 groups, 24 barriers)
    // Mbarriers at smem_raw[0..192)

    if (warp == 0) {
        uint32_t leader = elect_sync();
        // wt_ready: 8 barriers, init_count=1
        mbarrier_init_pred(smem + 0, 1, leader);
        mbarrier_init_pred(smem + 8, 1, leader);
        mbarrier_init_pred(smem + 16, 1, leader);
        mbarrier_init_pred(smem + 24, 1, leader);
        mbarrier_init_pred(smem + 32, 1, leader);
        mbarrier_init_pred(smem + 40, 1, leader);
        mbarrier_init_pred(smem + 48, 1, leader);
        mbarrier_init_pred(smem + 56, 1, leader);
        // act_ready: 8 barriers, init_count=1
        mbarrier_init_pred(smem + 64, 1, leader);
        mbarrier_init_pred(smem + 72, 1, leader);
        mbarrier_init_pred(smem + 80, 1, leader);
        mbarrier_init_pred(smem + 88, 1, leader);
        mbarrier_init_pred(smem + 96, 1, leader);
        mbarrier_init_pred(smem + 104, 1, leader);
        mbarrier_init_pred(smem + 112, 1, leader);
        mbarrier_init_pred(smem + 120, 1, leader);
        // data_consumed: 8 barriers, init_count=32
        mbarrier_init_pred(smem + 128, 32, leader);
        mbarrier_init_pred(smem + 136, 32, leader);
        mbarrier_init_pred(smem + 144, 32, leader);
        mbarrier_init_pred(smem + 152, 32, leader);
        mbarrier_init_pred(smem + 160, 32, leader);
        mbarrier_init_pred(smem + 168, 32, leader);
        mbarrier_init_pred(smem + 176, 32, leader);
        mbarrier_init_pred(smem + 184, 32, leader);
        asm volatile("fence.mbarrier_init.release.cluster;");
    }

    __syncthreads();

    __syncthreads();

    const int mbar_base = smem;
    #define wt_ready_addr (mbar_base + 0)
    #define act_ready_addr (mbar_base + 64)
    #define data_consumed_addr (mbar_base + 128)

    // ---- Role: compute ----
    if (warp <= 3) {
        { // compute_main
            int k_loops_c = (K + 1024 - 1) / 1024;
            int mib_c = blockIdx.x * 16;
            int ni_c = blockIdx.y * 8;
            if (tid < 16) {
                smem_bias[tid] = bias[mib_c + tid];
            }
            float accum[4];
            #pragma unroll
            for (int z = 0; z < 4; z++) {
                accum[z] = 0.0f;
            }
            unsigned int lane_div8 = lane / 8;
            unsigned int lane_mod8 = lane % 8;
            unsigned int row_wt = lane_mod8 + lane_div8 % 2 * 8;
            unsigned int col_off_wt = lane_div8 / 2;
            unsigned int row_act = lane_mod8;
            #pragma unroll 2
            for (unsigned int ki = 0; ki < k_loops_c; ki++) {
                unsigned int stage_c = (unsigned int)warp + 4 * (ki % 2);
                unsigned int phase_c = ki / 2 & 1;
                mbarrier_wait(wt_ready_addr + (stage_c) * 8, phase_c);
                mbarrier_wait(act_ready_addr + (stage_c) * 8, phase_c);
                #pragma unroll
                for (int su = 0; su < 4; su++) {
                    unsigned int base_wt = smem_wt_addr + (stage_c * 4 + (unsigned int)su) * 2048;
                    unsigned int base_act = smem_act_addr + (stage_c * 4 + (unsigned int)su) * 1024;
                    #pragma unroll
                    for (int kii = 0; kii < 4; kii++) {
                        unsigned int a_frag[4];
                        unsigned int b_frag[2];
                        unsigned int col_w = (unsigned int)(2 * kii) + col_off_wt;
                        unsigned int col_sw_w = row_wt % 8 ^ col_w;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
                            : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3])
                            : "r"(base_wt + row_wt * 128 + col_sw_w * 16)
                            : "memory");
                        unsigned int col_a = (unsigned int)(2 * kii) + lane_div8;
                        unsigned int col_sw_a = row_act % 8 ^ col_a;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
                            : "=r"(b_frag[0]), "=r"(b_frag[1])
                            : "r"(base_act + row_act * 128 + col_sw_a * 16)
                            : "memory");
                        asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
                            : "+f"(accum[0]), "+f"(accum[1]), "+f"(accum[2]), "+f"(accum[3])
                            : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1]));
                    }
                }
                asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
                mbarrier_arrive(data_consumed_addr + (stage_c) * 8);
            }
            asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_red_addr + (unsigned int)(tid * 16)), "r"(__as_u32(accum[0])), "r"(__as_u32(accum[1])), "r"(__as_u32(accum[2])), "r"(__as_u32(accum[3])) : "memory");
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
            if (warp == 0) {
                float part[12];
                #pragma unroll
                for (int w = 0; w < 3; w++) {
                    asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];"
                        : "=r"(*reinterpret_cast<uint32_t*>(&part[w * 4])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 1])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 2])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 3]))
                        : "r"(smem_red_addr + (unsigned int)((32 + w * 32 + tid) * 16)));
                }
                #pragma unroll
                for (int z_1 = 0; z_1 < 4; z_1++) {
                    accum[z_1] = accum[z_1] + part[z_1] + part[4 + z_1] + part[8 + z_1];
                }
                int tm = mib_c + lane / 4;
                int tn = ni_c + 2 * (lane % 4);
                float bias_lo = smem_bias[lane / 4];
                float bias_hi = smem_bias[lane / 4 + 8];
                float o00 = accum[0] + bias_lo;
                float o01 = accum[1] + bias_lo;
                float o10 = accum[2] + bias_hi;
                float o11 = accum[3] + bias_hi;
                if (tn < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm))) = __float2bfloat16_rn(o00);
                    }
                }
                if (tn + 1 < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm))) = __float2bfloat16_rn(o01);
                    }
                }
                if (tn < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm + 8))) = __float2bfloat16_rn(o10);
                    }
                }
                if (tn + 1 < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm + 8))) = __float2bfloat16_rn(o11);
                    }
                }
            }
        }
    // ---- Role: load_wt ----
    } else if (warp >= 4 && warp <= 7) {
        { // load_wt_main
            int k_loops = (K + 1024 - 1) / 1024;
            int mib = blockIdx.x * 16;
            unsigned int wslot = warp % 4;
            if (elect_sync()) {
                #pragma unroll 1
                for (unsigned int ki_1 = 0; ki_1 < k_loops; ki_1++) {
                    unsigned int stage = wslot + 4 * (ki_1 % 2);
                    unsigned int phase = ki_1 / 2 & 1;
                    int k_base = (ki_1 * 4 + wslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage) * 8, phase ^ 1);
                    mbarrier_arrive_expect_tx(wt_ready_addr + (stage) * 8, 8192);
                    #pragma unroll
                    for (int i = 0; i < 4; i++) {
                        int dst_wt = smem_wt_addr + (stage * 4 + (unsigned int)i) * 2048;
                        tma_2d_gmem2smem(dst_wt, (&tmap_wt), k_base + i * 64, mib, wt_ready_addr + (stage) * 8);
                    }
                }
                #pragma unroll
                for (int di = 0; di < 2; di++) {
                    if (di + 1 < 2) {
                        unsigned int dki = k_loops + di;
                        unsigned int dstage = wslot + 4 * (dki % 2);
                        unsigned int dphase = dki / 2 & 1;
                        mbarrier_wait(data_consumed_addr + (dstage) * 8, dphase ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    // ---- Role: load_act ----
    } else if (warp >= 8 && warp <= 11) {
        { // load_act_main
            int k_loops_a = (K + 1024 - 1) / 1024;
            int ni = blockIdx.y * 8;
            unsigned int aslot = warp % 4;
            if (elect_sync()) {
                #pragma unroll 1
                for (unsigned int ki_2 = 0; ki_2 < k_loops_a; ki_2++) {
                    unsigned int stage_a = aslot + 4 * (ki_2 % 2);
                    unsigned int phase_a = ki_2 / 2 & 1;
                    int k_base_a = (ki_2 * 4 + aslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage_a) * 8, phase_a ^ 1);
                    mbarrier_arrive_expect_tx(act_ready_addr + (stage_a) * 8, 4096);
                    #pragma unroll
                    for (int i_1 = 0; i_1 < 4; i_1++) {
                        int dst_act = smem_act_addr + (stage_a * 4 + (unsigned int)i_1) * 1024;
                        tma_2d_gmem2smem(dst_act, (&tmap_act), k_base_a + i_1 * 64, ni, act_ready_addr + (stage_a) * 8);
                    }
                }
                #pragma unroll
                for (int di_1 = 0; di_1 < 2; di_1++) {
                    if (di_1 + 1 < 2) {
                        unsigned int dki_a = k_loops_a + di_1;
                        unsigned int dstage_a = aslot + 4 * (dki_a % 2);
                        unsigned int dphase_a = dki_a / 2 & 1;
                        mbarrier_wait(data_consumed_addr + (dstage_a) * 8, dphase_a ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    }

    // Cleanup
}

} // extern "C"

// End of generated section stage8: capture the launch metadata, then unscope
// the section-local macros so the next section can redefine them.
namespace flashinfer {
namespace tinygemm2_sm100 {
static_assert(THREADS == 384, "generated stage8 variant disagrees on block size");
static_assert(USE_PDL == 0, "generated stage8 variant must be the non-PDL flavor");
constexpr int kSmemBytesStage8 = SMEM_TOTAL;
static_assert(kSmemBytesStage8 == 101504,
              "generated stage8 variant disagrees on dynamic smem footprint");
}  // namespace tinygemm2_sm100
}  // namespace flashinfer
#undef LOOM_INF
#undef NUM_MAIN_STAGES
#undef SMEM_SMEM_WT_OFF
#undef SMEM_SMEM_WT_STAGE_BYTES
#undef SMEM_SMEM_WT_STRIDE
#undef SMEM_SMEM_ACT_OFF
#undef SMEM_SMEM_ACT_STAGE_BYTES
#undef SMEM_SMEM_ACT_STRIDE
#undef SMEM_SMEM_RED_OFF
#undef SMEM_SMEM_RED_STAGE_BYTES
#undef SMEM_SMEM_RED_STRIDE
#undef SMEM_SMEM_BIAS_OFF
#undef SMEM_SMEM_BIAS_STAGE_BYTES
#undef SMEM_SMEM_BIAS_STRIDE
#undef SMEM_TOTAL
#undef THREADS
#undef USE_PDL
#undef wt_ready_addr
#undef act_ready_addr
#undef data_consumed_addr

// ============================================================================
// Section stage8_pdl — generated Loom schedule 'flashinfer_tinygemm2' (USE_PDL=1).
// ============================================================================

#define LOOM_INF CUDART_INF_F
#define NUM_MAIN_STAGES 1
#define SMEM_SMEM_WT_OFF 1024
#define SMEM_SMEM_WT_STAGE_BYTES 2048
#define SMEM_SMEM_WT_STRIDE 2048
#define SMEM_SMEM_ACT_OFF 66560
#define SMEM_SMEM_ACT_STAGE_BYTES 1024
#define SMEM_SMEM_ACT_STRIDE 1024
#define SMEM_SMEM_RED_OFF 99328
#define SMEM_SMEM_RED_STAGE_BYTES 2048
#define SMEM_SMEM_RED_STRIDE 2048
#define SMEM_SMEM_BIAS_OFF 101376
#define SMEM_SMEM_BIAS_STAGE_BYTES 32
#define SMEM_SMEM_BIAS_STRIDE 32
#define SMEM_TOTAL 101504
#define THREADS 384
#define USE_PDL 1

extern "C" {

__global__ __launch_bounds__(384, 1) void
kernel_tinygemm2_sm100_stage8_pdl(__grid_constant__ LoomTensorMap const tmap_wt, __grid_constant__ LoomTensorMap const tmap_act, __nv_bfloat16* __restrict__ output, __nv_bfloat16* __restrict__ bias, int M, int N, int K)
{
    const int tid = threadIdx.x;
    const int warp = make_warp_uniform(tid / 32);
    const int lane = tid % 32;

    extern __shared__ __align__(1024) char smem_raw[];
    int smem;
    smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw);

    const int bid = blockIdx.x;
    const int num_bids = gridDim.x;

    // Kernel setup ops
    __nv_bfloat16* smem_wt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024);
    const int smem_wt_addr = smem + 1024;
    __nv_bfloat16* smem_act = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560);
    const int smem_act_addr = smem + 66560;
    float* smem_red = reinterpret_cast<float*>(smem_raw + 99328);
    const int smem_red_addr = smem + 99328;
    __nv_bfloat16* smem_bias = reinterpret_cast<__nv_bfloat16*>(smem_raw + 101376);
    const int smem_bias_addr = smem + 101376;
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_wt))) : "memory"); }
    if (tid == 0) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&tmap_act))) : "memory"); }

    // Mbarrier init (3 groups, 24 barriers)
    // Mbarriers at smem_raw[0..192)

    if (warp == 0) {
        uint32_t leader = elect_sync();
        // wt_ready: 8 barriers, init_count=1
        mbarrier_init_pred(smem + 0, 1, leader);
        mbarrier_init_pred(smem + 8, 1, leader);
        mbarrier_init_pred(smem + 16, 1, leader);
        mbarrier_init_pred(smem + 24, 1, leader);
        mbarrier_init_pred(smem + 32, 1, leader);
        mbarrier_init_pred(smem + 40, 1, leader);
        mbarrier_init_pred(smem + 48, 1, leader);
        mbarrier_init_pred(smem + 56, 1, leader);
        // act_ready: 8 barriers, init_count=1
        mbarrier_init_pred(smem + 64, 1, leader);
        mbarrier_init_pred(smem + 72, 1, leader);
        mbarrier_init_pred(smem + 80, 1, leader);
        mbarrier_init_pred(smem + 88, 1, leader);
        mbarrier_init_pred(smem + 96, 1, leader);
        mbarrier_init_pred(smem + 104, 1, leader);
        mbarrier_init_pred(smem + 112, 1, leader);
        mbarrier_init_pred(smem + 120, 1, leader);
        // data_consumed: 8 barriers, init_count=32
        mbarrier_init_pred(smem + 128, 32, leader);
        mbarrier_init_pred(smem + 136, 32, leader);
        mbarrier_init_pred(smem + 144, 32, leader);
        mbarrier_init_pred(smem + 152, 32, leader);
        mbarrier_init_pred(smem + 160, 32, leader);
        mbarrier_init_pred(smem + 168, 32, leader);
        mbarrier_init_pred(smem + 176, 32, leader);
        mbarrier_init_pred(smem + 184, 32, leader);
        asm volatile("fence.mbarrier_init.release.cluster;");
    }

    __syncthreads();

    __syncthreads();

    const int mbar_base = smem;
    #define wt_ready_addr (mbar_base + 0)
    #define act_ready_addr (mbar_base + 64)
    #define data_consumed_addr (mbar_base + 128)

    // ---- Role: compute ----
    if (warp <= 3) {
        { // compute_main
            int k_loops_c = (K + 1024 - 1) / 1024;
            int mib_c = blockIdx.x * 16;
            int ni_c = blockIdx.y * 8;
            if (tid < 16) {
                smem_bias[tid] = bias[mib_c + tid];
            }
            float accum[4];
            #pragma unroll
            for (int z = 0; z < 4; z++) {
                accum[z] = 0.0f;
            }
            unsigned int lane_div8 = lane / 8;
            unsigned int lane_mod8 = lane % 8;
            unsigned int row_wt = lane_mod8 + lane_div8 % 2 * 8;
            unsigned int col_off_wt = lane_div8 / 2;
            unsigned int row_act = lane_mod8;
            #pragma unroll 2
            for (unsigned int ki = 0; ki < k_loops_c; ki++) {
                unsigned int stage_c = (unsigned int)warp + 4 * (ki % 2);
                unsigned int phase_c = ki / 2 & 1;
                mbarrier_wait(wt_ready_addr + (stage_c) * 8, phase_c);
                mbarrier_wait(act_ready_addr + (stage_c) * 8, phase_c);
                #pragma unroll
                for (int su = 0; su < 4; su++) {
                    unsigned int base_wt = smem_wt_addr + (stage_c * 4 + (unsigned int)su) * 2048;
                    unsigned int base_act = smem_act_addr + (stage_c * 4 + (unsigned int)su) * 1024;
                    #pragma unroll
                    for (int kii = 0; kii < 4; kii++) {
                        unsigned int a_frag[4];
                        unsigned int b_frag[2];
                        unsigned int col_w = (unsigned int)(2 * kii) + col_off_wt;
                        unsigned int col_sw_w = row_wt % 8 ^ col_w;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
                            : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3])
                            : "r"(base_wt + row_wt * 128 + col_sw_w * 16)
                            : "memory");
                        unsigned int col_a = (unsigned int)(2 * kii) + lane_div8;
                        unsigned int col_sw_a = row_act % 8 ^ col_a;
                        asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
                            : "=r"(b_frag[0]), "=r"(b_frag[1])
                            : "r"(base_act + row_act * 128 + col_sw_a * 16)
                            : "memory");
                        asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
                            : "+f"(accum[0]), "+f"(accum[1]), "+f"(accum[2]), "+f"(accum[3])
                            : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1]));
                    }
                }
                asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
                mbarrier_arrive(data_consumed_addr + (stage_c) * 8);
            }
            asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"(smem_red_addr + (unsigned int)(tid * 16)), "r"(__as_u32(accum[0])), "r"(__as_u32(accum[1])), "r"(__as_u32(accum[2])), "r"(__as_u32(accum[3])) : "memory");
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
            if (warp == 0) {
                float part[12];
                #pragma unroll
                for (int w = 0; w < 3; w++) {
                    asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];"
                        : "=r"(*reinterpret_cast<uint32_t*>(&part[w * 4])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 1])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 2])), "=r"(*reinterpret_cast<uint32_t*>(&part[(w * 4) + 3]))
                        : "r"(smem_red_addr + (unsigned int)((32 + w * 32 + tid) * 16)));
                }
                #pragma unroll
                for (int z_1 = 0; z_1 < 4; z_1++) {
                    accum[z_1] = accum[z_1] + part[z_1] + part[4 + z_1] + part[8 + z_1];
                }
                int tm = mib_c + lane / 4;
                int tn = ni_c + 2 * (lane % 4);
                float bias_lo = smem_bias[lane / 4];
                float bias_hi = smem_bias[lane / 4 + 8];
                float o00 = accum[0] + bias_lo;
                float o01 = accum[1] + bias_lo;
                float o10 = accum[2] + bias_hi;
                float o11 = accum[3] + bias_hi;
                if (tn < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm))) = __float2bfloat16_rn(o00);
                    }
                }
                if (tn + 1 < N) {
                    if (tm < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm))) = __float2bfloat16_rn(o01);
                    }
                }
                if (tn < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + (tn * M + tm + 8))) = __float2bfloat16_rn(o10);
                    }
                }
                if (tn + 1 < N) {
                    if (tm + 8 < M) {
                        *((__nv_bfloat16*)(output + ((tn + 1) * M + tm + 8))) = __float2bfloat16_rn(o11);
                    }
                }
            }
        }
    // ---- Role: load_wt ----
    } else if (warp >= 4 && warp <= 7) {
        { // load_wt_main
            int k_loops = (K + 1024 - 1) / 1024;
            int mib = blockIdx.x * 16;
            unsigned int wslot = warp % 4;
            if (elect_sync()) {
                #pragma unroll 1
                for (unsigned int ki_1 = 0; ki_1 < k_loops; ki_1++) {
                    unsigned int stage = wslot + 4 * (ki_1 % 2);
                    unsigned int phase = ki_1 / 2 & 1;
                    int k_base = (ki_1 * 4 + wslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage) * 8, phase ^ 1);
                    mbarrier_arrive_expect_tx(wt_ready_addr + (stage) * 8, 8192);
                    #pragma unroll
                    for (int i = 0; i < 4; i++) {
                        int dst_wt = smem_wt_addr + (stage * 4 + (unsigned int)i) * 2048;
                        tma_2d_gmem2smem(dst_wt, (&tmap_wt), k_base + i * 64, mib, wt_ready_addr + (stage) * 8);
                    }
                }
                #pragma unroll
                for (int di = 0; di < 2; di++) {
                    if (di + 1 < 2) {
                        unsigned int dki = k_loops + di;
                        unsigned int dstage = wslot + 4 * (dki % 2);
                        unsigned int dphase = dki / 2 & 1;
                        mbarrier_wait(data_consumed_addr + (dstage) * 8, dphase ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    // ---- Role: load_act ----
    } else if (warp >= 8 && warp <= 11) {
        { // load_act_main
            int k_loops_a = (K + 1024 - 1) / 1024;
            int ni = blockIdx.y * 8;
            unsigned int aslot = warp % 4;
            if (elect_sync()) {
                {
                    asm volatile("griddepcontrol.wait;" ::: "memory");
                    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
                }
                #pragma unroll 1
                for (unsigned int ki_2 = 0; ki_2 < k_loops_a; ki_2++) {
                    unsigned int stage_a = aslot + 4 * (ki_2 % 2);
                    unsigned int phase_a = ki_2 / 2 & 1;
                    int k_base_a = (ki_2 * 4 + aslot) * 256;
                    mbarrier_wait(data_consumed_addr + (stage_a) * 8, phase_a ^ 1);
                    mbarrier_arrive_expect_tx(act_ready_addr + (stage_a) * 8, 4096);
                    #pragma unroll
                    for (int i_1 = 0; i_1 < 4; i_1++) {
                        int dst_act = smem_act_addr + (stage_a * 4 + (unsigned int)i_1) * 1024;
                        tma_2d_gmem2smem(dst_act, (&tmap_act), k_base_a + i_1 * 64, ni, act_ready_addr + (stage_a) * 8);
                    }
                }
                #pragma unroll
                for (int di_1 = 0; di_1 < 2; di_1++) {
                    if (di_1 + 1 < 2) {
                        unsigned int dki_a = k_loops_a + di_1;
                        unsigned int dstage_a = aslot + 4 * (dki_a % 2);
                        unsigned int dphase_a = dki_a / 2 & 1;
                        mbarrier_wait(data_consumed_addr + (dstage_a) * 8, dphase_a ^ 1);
                    }
                }
            }
            asm volatile("barrier.sync 2, %0;" :: "r"(384));
        }
    }

    // Cleanup
}

} // extern "C"

// End of generated section stage8_pdl: capture the launch metadata, then unscope
// the section-local macros so the next section can redefine them.
namespace flashinfer {
namespace tinygemm2_sm100 {
static_assert(THREADS == 384, "generated stage8_pdl variant disagrees on block size");
static_assert(USE_PDL == 1, "generated stage8_pdl variant must be the PDL flavor");
constexpr int kSmemBytesStage8Pdl = SMEM_TOTAL;
static_assert(kSmemBytesStage8Pdl == 101504,
              "generated stage8_pdl variant disagrees on dynamic smem footprint");
}  // namespace tinygemm2_sm100
}  // namespace flashinfer
#undef LOOM_INF
#undef NUM_MAIN_STAGES
#undef SMEM_SMEM_WT_OFF
#undef SMEM_SMEM_WT_STAGE_BYTES
#undef SMEM_SMEM_WT_STRIDE
#undef SMEM_SMEM_ACT_OFF
#undef SMEM_SMEM_ACT_STAGE_BYTES
#undef SMEM_SMEM_ACT_STRIDE
#undef SMEM_SMEM_RED_OFF
#undef SMEM_SMEM_RED_STAGE_BYTES
#undef SMEM_SMEM_RED_STRIDE
#undef SMEM_SMEM_BIAS_OFF
#undef SMEM_SMEM_BIAS_STAGE_BYTES
#undef SMEM_SMEM_BIAS_STRIDE
#undef SMEM_TOTAL
#undef THREADS
#undef USE_PDL
#undef wt_ready_addr
#undef act_ready_addr
#undef data_consumed_addr

// clang-format on

namespace flashinfer {
namespace tinygemm2_sm100 {

using tvm::ffi::TensorView;

// Fixed tile geometry shared by every generated variant. These mirror the
// TensorRT-LLM tinygemm2 template constants (WARP_TILE_M=16, TILE_N=8,
// TILE_K=64) that the generated schedules were ported from.
constexpr int kTileM = 16;  // output-features tile
constexpr int kTileN = 8;   // batch tile
constexpr int kTileK = 64;  // reduction tile (one TMA box)
constexpr int kThreads = 384;

struct ProblemDims {
  int batch;
  int in_features;
  int out_features;
};

inline void CheckCuda(cudaError_t status, const char* operation) {
  TVM_FFI_ICHECK(status == cudaSuccess) << operation << " failed: " << cudaGetErrorString(status);
}

inline void CheckSm100Family(int device_id) {
  int major = 0;
  int minor = 0;
  CheckCuda(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id),
            "cudaDeviceGetAttribute(compute capability major)");
  CheckCuda(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id),
            "cudaDeviceGetAttribute(compute capability minor)");
  TVM_FFI_ICHECK(major == 10 && (minor == 0 || minor == 3))
      << "tinygemm2_sm100 requires an SM100/SM103 (B200/B300 class) device, got sm_" << major
      << minor;
}

inline void CheckBf16(const TensorView& t, const char* name) {
  const DLDataType d = t.dtype();
  TVM_FFI_ICHECK(d.code == kDLBfloat && d.bits == 16 && d.lanes == 1)
      << name << " must be bfloat16, got (code=" << int(d.code) << ", bits=" << int(d.bits)
      << ", lanes=" << int(d.lanes) << ")";
}

inline void CheckCudaBf16Contiguous(const TensorView& t, int ndim, const char* name) {
  TVM_FFI_ICHECK(t.device().device_type == kDLCUDA) << name << " must be a CUDA tensor";
  TVM_FFI_ICHECK(t.ndim() == ndim) << name << " must be " << ndim << "D, got ndim=" << t.ndim();
  TVM_FFI_ICHECK(t.IsContiguous()) << name << " must be contiguous";
  CheckBf16(t, name);
}

// Validate the public `out = input @ weight.T + bias` contract plus the
// coverage guards of the generated kernels (mirroring the Loom host shim):
// in_features must fit one TMA box; out_features must be a positive multiple
// of the kTileM output tile. The batch axis has NO lower guard — the
// activation descriptor deliberately allows an out-of-bounds box on that axis
// and TMA zero-fills rows past the end, so batch 1..7 inputs are valid.
inline ProblemDims CheckInputs(const TensorView& input, const TensorView& weight,
                               const TensorView& bias, const TensorView& out) {
  CheckCudaBf16Contiguous(input, 2, "input");
  CheckCudaBf16Contiguous(weight, 2, "weight");
  CheckCudaBf16Contiguous(bias, 1, "bias");
  CheckCudaBf16Contiguous(out, 2, "out");
  const int device_id = input.device().device_id;
  TVM_FFI_ICHECK(weight.device().device_id == device_id && bias.device().device_id == device_id &&
                 out.device().device_id == device_id)
      << "input/weight/bias/out must live on the same CUDA device";
  CheckSm100Family(device_id);

  const int64_t batch = input.size(0);
  const int64_t in_features = input.size(1);
  const int64_t out_features = weight.size(0);
  TVM_FFI_ICHECK(weight.size(1) == in_features)
      << "weight.shape[1] (" << weight.size(1) << ") must equal input.shape[1] (" << in_features
      << ")";
  TVM_FFI_ICHECK(bias.size(0) == out_features)
      << "bias.shape[0] (" << bias.size(0) << ") must equal weight.shape[0] (" << out_features
      << ")";
  TVM_FFI_ICHECK(out.size(0) == batch && out.size(1) == out_features)
      << "out must have shape (" << batch << ", " << out_features << "), got (" << out.size(0)
      << ", " << out.size(1) << ")";

  TVM_FFI_ICHECK(batch > 0) << "batch must be positive, got " << batch;
  TVM_FFI_ICHECK(in_features >= kTileK)
      << "in_features (" << in_features << ") must be at least " << kTileK << " (one TMA box)";
  TVM_FFI_ICHECK(out_features >= kTileM && out_features % kTileM == 0)
      << "out_features (" << out_features << ") must be a positive multiple of " << kTileM;
  TVM_FFI_ICHECK(batch <= std::numeric_limits<int>::max() &&
                 in_features <= std::numeric_limits<int>::max() &&
                 out_features <= std::numeric_limits<int>::max())
      << "problem dimensions exceed the kernel's i32 scalar range";

  return ProblemDims{static_cast<int>(batch), static_cast<int>(in_features),
                     static_cast<int>(out_features)};
}

// 2D TMA descriptor for the weight matrix — field-for-field the descriptor
// the Loom host shim encodes for 'tmap_wt': box (kTileK, kTileM), 128B
// swizzle, no L2 promotion, no OOB fill. Both boxed axes stay in bounds
// (CheckInputs guarantees in_features >= kTileK and out_features >= kTileM).
inline CUtensorMap EncodeWeightTma(const TensorView& weight) {
  const uint64_t global_dim[2] = {static_cast<uint64_t>(weight.size(1)),
                                  static_cast<uint64_t>(weight.size(0))};
  const uint64_t global_strides[1] = {static_cast<uint64_t>(weight.stride(0)) *
                                      sizeof(__nv_bfloat16)};
  const uint32_t box_dim[2] = {static_cast<uint32_t>(kTileK), static_cast<uint32_t>(kTileM)};
  const uint32_t elem_strides[2] = {1u, 1u};
  CUtensorMap tm;
  const CUresult r = cuTensorMapEncodeTiled(
      &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, weight.data_ptr(), global_dim, global_strides,
      box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
      CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
  TVM_FFI_ICHECK(r == CUDA_SUCCESS)
      << "cuTensorMapEncodeTiled failed for the weight descriptor: CUresult=" << int(r);
  return tm;
}

// 2D TMA descriptor for the activation matrix — the Loom host shim's
// 'tmap_act' descriptor. The batch axis opts into an out-of-bounds box
// (box kTileN may exceed batch for batch 1..7); TMA zero-fills those rows.
inline CUtensorMap EncodeActivationTma(const TensorView& input) {
  const uint64_t global_dim[2] = {static_cast<uint64_t>(input.size(1)),
                                  static_cast<uint64_t>(input.size(0))};
  const uint64_t global_strides[1] = {static_cast<uint64_t>(input.stride(0)) *
                                      sizeof(__nv_bfloat16)};
  const uint32_t box_dim[2] = {static_cast<uint32_t>(kTileK), static_cast<uint32_t>(kTileN)};
  const uint32_t elem_strides[2] = {1u, 1u};
  CUtensorMap tm;
  const CUresult r = cuTensorMapEncodeTiled(
      &tm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, input.data_ptr(), global_dim, global_strides,
      box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
      CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
  TVM_FFI_ICHECK(r == CUDA_SUCCESS)
      << "cuTensorMapEncodeTiled failed for the activation descriptor: CUresult=" << int(r);
  return tm;
}

// Launch one generated variant. LoomTensorMap is the generated kernels'
// by-value __grid_constant__ tensor-map parameter type; it is
// layout-compatible with CUtensorMap. PDL variants launch through
// cudaLaunchKernelEx with programmatic stream serialization, matching the
// in-kernel griddepcontrol pair compiled into those sections.
using GeneratedKernel = void (*)(LoomTensorMap, LoomTensorMap, __nv_bfloat16*, __nv_bfloat16*, int,
                                 int, int);

inline void LaunchVariant(GeneratedKernel kernel, int smem_bytes, bool pdl,
                          const CUtensorMap& weight_map, const CUtensorMap& activation_map,
                          __nv_bfloat16* out, __nv_bfloat16* bias, const ProblemDims& dims,
                          cudaStream_t stream) {
  static_assert(sizeof(LoomTensorMap) == sizeof(CUtensorMap),
                "generated tensor-map parameter must be layout-compatible with CUtensorMap");
  LoomTensorMap wt_param, act_param;
  std::memcpy(&wt_param, &weight_map, sizeof(LoomTensorMap));
  std::memcpy(&act_param, &activation_map, sizeof(LoomTensorMap));

  CheckCuda(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes),
            "cudaFuncSetAttribute(tinygemm2_sm100 dynamic smem)");

  const dim3 grid((dims.out_features + kTileM - 1) / kTileM, (dims.batch + kTileN - 1) / kTileN);
  const dim3 block(kThreads);

  if (pdl) {
    cudaLaunchConfig_t config;
    cudaLaunchAttribute attrs[1];
    config.gridDim = grid;
    config.blockDim = block;
    config.dynamicSmemBytes = smem_bytes;
    config.stream = stream;
    attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
    attrs[0].val.programmaticStreamSerializationAllowed = 1;
    config.attrs = attrs;
    config.numAttrs = 1;
    CheckCuda(cudaLaunchKernelEx(&config, kernel, wt_param, act_param, out, bias, dims.out_features,
                                 dims.batch, dims.in_features),
              "cudaLaunchKernelEx(tinygemm2_sm100)");
  } else {
    kernel<<<grid, block, smem_bytes, stream>>>(wt_param, act_param, out, bias, dims.out_features,
                                                dims.batch, dims.in_features);
    CheckCuda(cudaGetLastError(), "tinygemm2_sm100 kernel launch");
  }
}

// out = input @ weight.T + bias (bf16, fp32 accumulation), column-major
// epilogue identical to csrc/tinygemm2.cu.
void RunStage4(TensorView input, TensorView weight, TensorView bias, TensorView out) {
  const ProblemDims dims = CheckInputs(input, weight, bias, out);
  const CUtensorMap weight_map = EncodeWeightTma(weight);
  const CUtensorMap activation_map = EncodeActivationTma(input);
  const cudaStream_t stream = get_stream(input.device());
  LaunchVariant(&kernel_tinygemm2_sm100_stage4, kSmemBytesStage4, /*pdl=*/false, weight_map,
                activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
                reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream);
}

// out = input @ weight.T + bias (bf16, fp32 accumulation), column-major
// epilogue identical to csrc/tinygemm2.cu.
void RunStage4Pdl(TensorView input, TensorView weight, TensorView bias, TensorView out) {
  const ProblemDims dims = CheckInputs(input, weight, bias, out);
  const CUtensorMap weight_map = EncodeWeightTma(weight);
  const CUtensorMap activation_map = EncodeActivationTma(input);
  const cudaStream_t stream = get_stream(input.device());
  LaunchVariant(&kernel_tinygemm2_sm100_stage4_pdl, kSmemBytesStage4Pdl, /*pdl=*/true, weight_map,
                activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
                reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream);
}

// out = input @ weight.T + bias (bf16, fp32 accumulation), column-major
// epilogue identical to csrc/tinygemm2.cu.
void RunStage8(TensorView input, TensorView weight, TensorView bias, TensorView out) {
  const ProblemDims dims = CheckInputs(input, weight, bias, out);
  const CUtensorMap weight_map = EncodeWeightTma(weight);
  const CUtensorMap activation_map = EncodeActivationTma(input);
  const cudaStream_t stream = get_stream(input.device());
  LaunchVariant(&kernel_tinygemm2_sm100_stage8, kSmemBytesStage8, /*pdl=*/false, weight_map,
                activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
                reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream);
}

// out = input @ weight.T + bias (bf16, fp32 accumulation), column-major
// epilogue identical to csrc/tinygemm2.cu.
void RunStage8Pdl(TensorView input, TensorView weight, TensorView bias, TensorView out) {
  const ProblemDims dims = CheckInputs(input, weight, bias, out);
  const CUtensorMap weight_map = EncodeWeightTma(weight);
  const CUtensorMap activation_map = EncodeActivationTma(input);
  const cudaStream_t stream = get_stream(input.device());
  LaunchVariant(&kernel_tinygemm2_sm100_stage8_pdl, kSmemBytesStage8Pdl, /*pdl=*/true, weight_map,
                activation_map, reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
                reinterpret_cast<__nv_bfloat16*>(bias.data_ptr()), dims, stream);
}

}  // namespace tinygemm2_sm100
}  // namespace flashinfer

TVM_FFI_DLL_EXPORT_TYPED_FUNC(stage4_op, flashinfer::tinygemm2_sm100::RunStage4);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(stage4_pdl_op, flashinfer::tinygemm2_sm100::RunStage4Pdl);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(stage8_op, flashinfer::tinygemm2_sm100::RunStage8);
TVM_FFI_DLL_EXPORT_TYPED_FUNC(stage8_pdl_op, flashinfer::tinygemm2_sm100::RunStage8Pdl);
