ULT7RA-FLUXNATION/flux_opt.cu

#include "SINGLE_STREAM_BLOCK.H"
#include <algorithm>
#include <cfloat>
#include <cmath>
#include <cstdlib>

#include "cutlass/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/gemm/gemm.h"
#include "cutlass/gemm/device/gemm_universal.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
#include "cutlass/layout/matrix.h"

namespace transformer {

using CutlassElementA = cutlass::float_e4m3_t;
using CutlassElementB = cutlass::float_e4m3_t;
using CutlassElementAcc = float;
using CutlassLayoutA = cutlass::layout::RowMajor;
using CutlassLayoutB = cutlass::layout::ColumnMajor;
using CutlassLayoutC = cutlass::layout::RowMajor;

using CutlassThreadblockShape = cutlass::gemm::GemmShape<64, 128, 64>;
using CutlassWarpShape = cutlass::gemm::GemmShape<32, 64, 64>;
using CutlassInstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
static int const kCutlassStages = 4;
static int const kCutlassAlignmentA = 16;
static int const kCutlassAlignmentB = 16;
static int const kKernelThreads = 256;
static_assert((kKernelThreads % 32) == 0, "kKernelThreads must be a multiple of 32.");

using CutlassElementC = cutlass::bfloat16_t;
using CutlassEpilogueOp = cutlass::epilogue::thread::LinearCombination<
    CutlassElementC,
    128 / cutlass::sizeof_bits<CutlassElementC>::value,
    CutlassElementAcc, CutlassElementAcc>;
using CutlassGemm = cutlass::gemm::device::GemmUniversal<
    CutlassElementA, CutlassLayoutA,
    CutlassElementB, CutlassLayoutB,
    CutlassElementC, CutlassLayoutC,
    CutlassElementAcc,
    cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
    CutlassThreadblockShape, CutlassWarpShape, CutlassInstructionShape,
    CutlassEpilogueOp,
    cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
    kCutlassStages, kCutlassAlignmentA, kCutlassAlignmentB>;

// Vectorized load/store helpers — 128-bit (8 half elements) for max bandwidth
union VecHalf8 {
    float4 f4;
    half   h[8];
};
union VecHalf4 {
    float2 f2;
    half   h[4];
};
union VecHalf2 {
    unsigned int u;
    half h[2];
};
union VecBF16_4 {
    float2 f2;
    __nv_bfloat16 b[4];
};
union VecBF16_8 {
    float4 f4;
    __nv_bfloat16 b[8];
};

// Forward declarations
__global__ void bias_add_kernel(half* __restrict__ out,
                                const half* __restrict__ bias,
                                int rows,
                                int cols);
__global__ void bf16_to_fp16_bias_add_kernel(void* __restrict__ data,
                                             const half* __restrict__ bias,
                                             int rows,
                                             int cols);

static inline void linear_gemm(cublasHandle_t cub,
                               int m,
                               int n,
                               int k,
                               const half* x,
                               const half* w_torch,
                               half* y)
{
    const float alpha = 1.0f;
    const float beta = 0.0f;
    BK(cublasGemmEx(
        cub,
        CUBLAS_OP_T,
        CUBLAS_OP_N,
        n, m, k,
        &alpha,
        w_torch, CUDA_R_16F, k,
        x, CUDA_R_16F, k,
        &beta,
        y, CUDA_R_16F, n,
        CUBLAS_COMPUTE_32F_FAST_16F,
        CUBLAS_GEMM_DEFAULT_TENSOR_OP));
}

__forceinline__ __device__ float warp_reduce_sum(float v)
{
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

__forceinline__ __device__ float warp_reduce_max(float v)
{
    for (int offset = 16; offset > 0; offset >>= 1) {
        v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset));
    }
    return v;
}

__forceinline__ __device__ float block_reduce_sum(float v)
{
    constexpr int kWarps = kKernelThreads / 32;
    __shared__ float warp_vals[kWarps];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;

    v = warp_reduce_sum(v);
    if (lane == 0) {
        warp_vals[wid] = v;
    }
    __syncthreads();

    if (wid == 0) {
        float x = (lane < kWarps) ? warp_vals[lane] : 0.0f;
        x = warp_reduce_sum(x);
        if (lane == 0) {
            warp_vals[0] = x;
        }
    }
    __syncthreads();
    return warp_vals[0];
}

__forceinline__ __device__ float block_reduce_max(float v)
{
    constexpr int kWarps = kKernelThreads / 32;
    __shared__ float warp_vals[kWarps];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;

    v = warp_reduce_max(v);
    if (lane == 0) {
        warp_vals[wid] = v;
    }
    __syncthreads();

    if (wid == 0) {
        float x = (lane < kWarps) ? warp_vals[lane] : -FLT_MAX;
        x = warp_reduce_max(x);
        if (lane == 0) {
            warp_vals[0] = x;
        }
    }
    __syncthreads();
    return warp_vals[0];
}

__global__ __launch_bounds__(256, 4)
void amax_abs_kernel(const half* __restrict__ in,
                     unsigned int* __restrict__ amax_uint,
                     int n)
{
    float local_max = 0.0f;
    int stride = blockDim.x * gridDim.x;
    int base = blockIdx.x * blockDim.x + threadIdx.x;
    // Vectorized: process 8 elements per iteration via float4
    int n8 = n & ~7;
    for (int i = base * 8; i < n8; i += stride * 8) {
        VecHalf8 v;
        v.f4 = *reinterpret_cast<const float4*>(in + i);
        #pragma unroll
        for (int j = 0; j < 8; j++)
            local_max = fmaxf(local_max, fabsf(__half2float(v.h[j])));
    }
    // Scalar tail
    for (int i = n8 + base; i < n; i += stride) {
        local_max = fmaxf(local_max, fabsf(__half2float(in[i])));
    }
    float block_max = block_reduce_max(local_max);
    if (threadIdx.x == 0) {
        atomicMax(amax_uint, __float_as_uint(block_max));
    }
}

__global__ __launch_bounds__(256, 4)
void quantize_act_fp8_kernel(const half* __restrict__ in,
                             __nv_fp8_e4m3* __restrict__ out,
                             const unsigned int* __restrict__ amax_uint,
                             float* __restrict__ scale_out,
                             const float* __restrict__ b_scale,
                             float* __restrict__ combined_scale_out,
                             int n)
{
    float amax = __uint_as_float(*amax_uint);
    float act_s = (amax < 1.0e-12f) ? 1.0f : (amax / 448.0f);
    float inv_scale = (amax < 1.0e-12f) ? 1.0f : (448.0f / amax);
    if (blockIdx.x == 0 && threadIdx.x == 0) {
        *scale_out = act_s;
        *combined_scale_out = act_s * (*b_scale);
    }
    int base = blockIdx.x * blockDim.x + threadIdx.x;
    // Vectorized: process 8 elements via float4 load
    int idx8 = base * 8;
    if (idx8 + 7 < n) {
        VecHalf8 v;
        v.f4 = *reinterpret_cast<const float4*>(in + idx8);
        #pragma unroll
        for (int j = 0; j < 8; j++) {
            float val = __half2float(v.h[j]) * inv_scale;
            val = fminf(fmaxf(val, -448.0f), 448.0f);
            out[idx8 + j] = __nv_fp8_e4m3(val);
        }
    } else {
        for (int i = idx8; i < min(idx8 + 8, n); i++) {
            float val = __half2float(in[i]) * inv_scale;
            val = fminf(fmaxf(val, -448.0f), 448.0f);
            out[i] = __nv_fp8_e4m3(val);
        }
    }
}

__global__ void write_combined_scale_kernel(const float* __restrict__ a_scale,
                                            const float* __restrict__ b_scale,
                                            float* __restrict__ ab_scale)
{
    if (threadIdx.x == 0 && blockIdx.x == 0)
        *ab_scale = (*a_scale) * (*b_scale);
}

// FP8 GEMM: device-pointer alpha, ZERO host syncs, BF16 output
static inline void linear_gemm_fp8(int m, int n, int k,
                                    const half* a_half,
                                    const __nv_fp8_e4m3* b_fp8,
                                    const float* b_scale,
                                    __nv_bfloat16* y_bf16,
                                    __nv_fp8_e4m3* act_fp8,
                                    unsigned int* act_amax,
                                    float* act_scale,
                                    float* d_alpha,
                                    void* workspace,
                                    size_t workspace_bytes,
                                    cudaStream_t stream)
{
    const int act_n = m * k;
    const int threads = kKernelThreads;
    const int amax_blocks = min((act_n + threads - 1) / threads, 4096);
    const int q_blocks = (act_n + threads * 8 - 1) / (threads * 8);

    CK(cudaMemsetAsync(act_amax, 0, sizeof(unsigned int), stream));
    amax_abs_kernel<<<amax_blocks, threads, 0, stream>>>(a_half, act_amax, act_n);
    quantize_act_fp8_kernel<<<q_blocks, threads, 0, stream>>>(a_half, act_fp8, act_amax, act_scale, b_scale, d_alpha, act_n);

    CutlassGemm gemm_op;
    typename CutlassGemm::Arguments args{
        cutlass::gemm::GemmUniversalMode::kGemm,
        {m, n, k}, 1,
        {d_alpha, (const float*)nullptr},
        reinterpret_cast<CutlassElementA const*>(act_fp8),
        reinterpret_cast<CutlassElementB const*>(b_fp8),
        reinterpret_cast<CutlassElementC const*>(y_bf16),
        reinterpret_cast<CutlassElementC*>(y_bf16),
        static_cast<int64_t>(m) * k, static_cast<int64_t>(k) * n,
        static_cast<int64_t>(m) * n, static_cast<int64_t>(m) * n,
        k, k, n, n
    };

    cutlass::Status st = gemm_op.can_implement(args);
    if (st != cutlass::Status::kSuccess)
        throw std::runtime_error("CUTLASS FP8 can_implement failed");
    size_t ws_needed = gemm_op.get_workspace_size(args);
    if (ws_needed > workspace_bytes)
        throw std::runtime_error("CUTLASS workspace too small");
    st = gemm_op.initialize(args, workspace, stream);
    if (st != cutlass::Status::kSuccess)
        throw std::runtime_error("CUTLASS FP8 initialize failed");
    st = gemm_op(stream);
    if (st != cutlass::Status::kSuccess)
        throw std::runtime_error("CUTLASS FP8 run failed");
    CK(cudaGetLastError());
}

// In-place bf16→fp16 conversion with bias add. Both types are 2 bytes so
// this overwrites each element in-place. Vectorized: 4 elements per thread.
__global__ __launch_bounds__(256, 4)
void bf16_to_fp16_bias_add_kernel(void* __restrict__ data,
                                  const half* __restrict__ bias,
                                  int rows,
                                  int cols)
{
    int total = rows * cols;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if (base + 3 < total) {
        int col = base % cols;
        // Check we don't cross a row boundary
        if (col + 3 < cols) {
            VecBF16_4 bv;
            bv.f2 = *reinterpret_cast<const float2*>(reinterpret_cast<__nv_bfloat16*>(data) + base);
            VecHalf4 bias_v;
            bias_v.f2 = *reinterpret_cast<const float2*>(bias + col);
            VecHalf4 out_v;
            #pragma unroll
            for (int j = 0; j < 4; j++)
                out_v.h[j] = __float2half_rn(__bfloat162float(bv.b[j]) + __half2float(bias_v.h[j]));
            *reinterpret_cast<float2*>(reinterpret_cast<half*>(data) + base) = out_v.f2;
            return;
        }
    }
    // Scalar fallback for tail/row-boundary
    for (int j = 0; j < 4 && base + j < total; j++) {
        int idx = base + j;
        int col = idx % cols;
        float v = __bfloat162float(reinterpret_cast<__nv_bfloat16*>(data)[idx])
                + __half2float(bias[col]);
        reinterpret_cast<half*>(data)[idx] = __float2half_rn(v);
    }
}

static inline void qk_gemm(cublasHandle_t cub,
                           int batch,
                           int seq,
                           int d,
                           const half* k,
                           const half* q,
                           half* scores)
{
    const float alpha = ATT_SC;
    const float beta = 0.0f;
    const long long stride_qk = static_cast<long long>(seq) * d;
    const long long stride_sc = static_cast<long long>(seq) * seq;

    BK(cublasGemmStridedBatchedEx(
        cub,
        CUBLAS_OP_T,
        CUBLAS_OP_N,
        seq, seq, d,
        &alpha,
        k, CUDA_R_16F, d, stride_qk,
        q, CUDA_R_16F, d, stride_qk,
        &beta,
        scores, CUDA_R_16F, seq, stride_sc,
        batch,
        CUBLAS_COMPUTE_32F_FAST_16F,
        CUBLAS_GEMM_DEFAULT_TENSOR_OP));
}

static inline void sv_gemm(cublasHandle_t cub,
                           int batch,
                           int seq,
                           int d,
                           const half* v,
                           const half* scores,
                           half* out_heads)
{
    const float alpha = 1.0f;
    const float beta = 0.0f;
    const long long stride_v = static_cast<long long>(seq) * d;
    const long long stride_sc = static_cast<long long>(seq) * seq;

    BK(cublasGemmStridedBatchedEx(
        cub,
        CUBLAS_OP_N,
        CUBLAS_OP_N,
        d, seq, seq,
        &alpha,
        v, CUDA_R_16F, d, stride_v,
        scores, CUDA_R_16F, seq, stride_sc,
        &beta,
        out_heads, CUDA_R_16F, d, stride_v,
        batch,
        CUBLAS_COMPUTE_32F_FAST_16F,
        CUBLAS_GEMM_DEFAULT_TENSOR_OP));
}


__global__ __launch_bounds__(256, 4)
void silu_kernel(const half* __restrict__ in,
                 half* __restrict__ out,
                 int n)
{
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base + 7 < n) {
        VecHalf8 v;
        v.f4 = *reinterpret_cast<const float4*>(in + base);
        #pragma unroll
        for (int j = 0; j < 8; j++) {
            float x = __half2float(v.h[j]);
            v.h[j] = __float2half_rn(x / (1.0f + expf(-x)));
        }
        *reinterpret_cast<float4*>(out + base) = v.f4;
    } else {
        for (int i = base; i < min(base + 8, n); i++) {
            float x = __half2float(in[i]);
            out[i] = __float2half_rn(x / (1.0f + expf(-x)));
        }
    }
}

__global__ __launch_bounds__(256, 4)
void bias_add_kernel(half* __restrict__ out,
                     const half* __restrict__ bias,
                     int rows,
                     int cols)
{
    int total = rows * cols;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if (base + 3 < total) {
        int col = base % cols;
        if (col + 3 < cols) {
            VecHalf4 ov, bv;
            ov.f2 = *reinterpret_cast<const float2*>(out + base);
            bv.f2 = *reinterpret_cast<const float2*>(bias + col);
            #pragma unroll
            for (int j = 0; j < 4; j++)
                ov.h[j] = __float2half_rn(__half2float(ov.h[j]) + __half2float(bv.h[j]));
            *reinterpret_cast<float2*>(out + base) = ov.f2;
            return;
        }
    }
    for (int j = 0; j < 4 && base + j < total; j++) {
        int idx = base + j;
        int col = idx % cols;
        float v = __half2float(out[idx]) + __half2float(bias[col]);
        out[idx] = __float2half_rn(v);
    }
}

// Input x is bf16 (model runs in bf16, values exceed fp16 max of 65504).
// Output stays fp16 — layernorm normalizes to small range. Vectorized 4-wide.
__global__ __launch_bounds__(256, 2)
void layernorm_mod_kernel(const __nv_bfloat16* __restrict__ x,
                          const half* __restrict__ mod,
                          half* __restrict__ out,
                          int seq)
{
    int row = blockIdx.x;
    int b = row / seq;

    const __nv_bfloat16* x_row = x + static_cast<long long>(row) * HIDDEN;
    half* out_row = out + static_cast<long long>(row) * HIDDEN;
    const half* shift = mod + static_cast<long long>(b) * MODDIM;
    const half* scale = shift + HIDDEN;

    float sum = 0.0f;
    float sq = 0.0f;
    // Vectorized reduction: 4 bf16 elements per iteration
    for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
        VecBF16_4 bv;
        bv.f2 = *reinterpret_cast<const float2*>(x_row + i);
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float v = __bfloat162float(bv.b[j]);
            sum += v;
            sq += v * v;
        }
    }
    float mean = block_reduce_sum(sum) / HIDDEN;
    float var = block_reduce_sum(sq) / HIDDEN - mean * mean;
    float inv = rsqrtf(var + LN_EPS);

    // Vectorized writeback
    for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
        VecBF16_4 xv;
        xv.f2 = *reinterpret_cast<const float2*>(x_row + i);
        VecHalf4 sc_v, sh_v, ov;
        sc_v.f2 = *reinterpret_cast<const float2*>(scale + i);
        sh_v.f2 = *reinterpret_cast<const float2*>(shift + i);
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float v = __bfloat162float(xv.b[j]);
            float nrm = (v - mean) * inv;
            ov.h[j] = __float2half_rn(nrm * (1.0f + __half2float(sc_v.h[j])) + __half2float(sh_v.h[j]));
        }
        *reinterpret_cast<float2*>(out_row + i) = ov.f2;
    }
}

// Fused: reads mod as BF16 (raw GEMM output), adds bias inline, applies layernorm+mod.
// Eliminates the separate bf16_to_fp16_bias_add pass for Mod.
__global__ __launch_bounds__(256, 2)
void layernorm_mod_bf16bias_kernel(const __nv_bfloat16* __restrict__ x,
                                   const void* __restrict__ mod_raw,
                                   const half* __restrict__ mod_bias,
                                   half* __restrict__ out,
                                   int seq)
{
    int row = blockIdx.x;
    int b = row / seq;

    const __nv_bfloat16* x_row = x + static_cast<long long>(row) * HIDDEN;
    half* out_row = out + static_cast<long long>(row) * HIDDEN;
    const __nv_bfloat16* mod_bf = reinterpret_cast<const __nv_bfloat16*>(mod_raw) + static_cast<long long>(b) * MODDIM;
    // shift = mod[0..HIDDEN), scale = mod[HIDDEN..2*HIDDEN)

    float sum = 0.0f;
    float sq = 0.0f;
    for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
        VecBF16_4 bv;
        bv.f2 = *reinterpret_cast<const float2*>(x_row + i);
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float v = __bfloat162float(bv.b[j]);
            sum += v;
            sq += v * v;
        }
    }
    float mean = block_reduce_sum(sum) / HIDDEN;
    float var = block_reduce_sum(sq) / HIDDEN - mean * mean;
    float inv = rsqrtf(var + LN_EPS);

    for (int i = threadIdx.x * 4; i < HIDDEN; i += blockDim.x * 4) {
        VecBF16_4 xv;
        xv.f2 = *reinterpret_cast<const float2*>(x_row + i);
        // Read scale and shift as BF16, add bias inline
        VecBF16_4 sc_bf, sh_bf;
        sc_bf.f2 = *reinterpret_cast<const float2*>(mod_bf + HIDDEN + i);  // scale portion
        sh_bf.f2 = *reinterpret_cast<const float2*>(mod_bf + i);           // shift portion
        VecHalf4 sc_bias, sh_bias;
        sc_bias.f2 = *reinterpret_cast<const float2*>(mod_bias + HIDDEN + i);
        sh_bias.f2 = *reinterpret_cast<const float2*>(mod_bias + i);
        VecHalf4 ov;
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float v = __bfloat162float(xv.b[j]);
            float nrm = (v - mean) * inv;
            float sc = __bfloat162float(sc_bf.b[j]) + __half2float(sc_bias.h[j]);
            float sh = __bfloat162float(sh_bf.b[j]) + __half2float(sh_bias.h[j]);
            ov.h[j] = __float2half_rn(nrm * (1.0f + sc) + sh);
        }
        *reinterpret_cast<float2*>(out_row + i) = ov.f2;
    }
}

// Fused: reads gate portion of mod as BF16, adds bias, writes FP16
__global__ __launch_bounds__(256, 6)
void extract_gate_bf16bias_kernel(const void* __restrict__ mod_raw,
                                  const half* __restrict__ mod_bias,
                                  half* __restrict__ gate,
                                  int batch)
{
    int total = batch * HIDDEN;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    const __nv_bfloat16* mod_bf = reinterpret_cast<const __nv_bfloat16*>(mod_raw);
    if (base + 7 < total) {
        int b = base / HIDDEN;
        int h = base % HIDDEN;
        if (h + 7 < HIDDEN) {
            VecBF16_8 mv;
            mv.f4 = *reinterpret_cast<const float4*>(mod_bf + static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h);
            VecHalf8 bv;
            bv.f4 = *reinterpret_cast<const float4*>(mod_bias + 2 * HIDDEN + h);
            VecHalf8 ov;
            #pragma unroll
            for (int j = 0; j < 8; j++)
                ov.h[j] = __float2half_rn(__bfloat162float(mv.b[j]) + __half2float(bv.h[j]));
            *reinterpret_cast<float4*>(gate + base) = ov.f4;
            return;
        }
    }
    for (int j = 0; j < 8 && base + j < total; j++) {
        int idx = base + j;
        int b = idx / HIDDEN;
        int h = idx % HIDDEN;
        float val = __bfloat162float(mod_bf[static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h])
                  + __half2float(mod_bias[2 * HIDDEN + h]);
        gate[idx] = __float2half_rn(val);
    }
}

// Fused: reads L1 output as BF16, adds bias, splits into Q/K/V
__global__ void split_qkv_bf16bias_kernel(const void* __restrict__ l1o_raw,
                                          const half* __restrict__ l1_bias,
                                          half* __restrict__ q,
                                          half* __restrict__ k,
                                          half* __restrict__ v,
                                          int seq)
{
    const __nv_bfloat16* l1o_bf = reinterpret_cast<const __nv_bfloat16*>(l1o_raw);
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = seq * HIDDEN;
    int batch_total = gridDim.y;
    int linear = blockIdx.y * total + idx;
    int max_linear = batch_total * total;
    if (linear >= max_linear) return;

    int token = linear / HIDDEN;
    int h = linear % HIDDEN;
    int b = token / seq;
    int s = token % seq;
    int head = h / HDIM;
    int d = h % HDIM;

    long long base_l1 = static_cast<long long>(token) * LIN1DIM;
    long long dst = (((static_cast<long long>(b) * HEADS + head) * seq) + s) * HDIM + d;

    q[dst] = __float2half_rn(__bfloat162float(l1o_bf[base_l1 + h]) + __half2float(l1_bias[h]));
    k[dst] = __float2half_rn(__bfloat162float(l1o_bf[base_l1 + HIDDEN + h]) + __half2float(l1_bias[HIDDEN + h]));
    v[dst] = __float2half_rn(__bfloat162float(l1o_bf[base_l1 + 2 * HIDDEN + h]) + __half2float(l1_bias[2 * HIDDEN + h]));
}

// Original: reads L1 output as FP16 (already converted), splits into Q/K/V
__global__ void split_qkv_kernel(const half* __restrict__ l1o,
                                 half* __restrict__ q,
                                 half* __restrict__ k,
                                 half* __restrict__ v,
                                 int seq)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = seq * HIDDEN;
    int batch_total = gridDim.y;
    int linear = blockIdx.y * total + idx;
    int max_linear = batch_total * total;
    if (linear >= max_linear) return;

    int token = linear / HIDDEN;
    int h = linear % HIDDEN;
    int b = token / seq;
    int s = token % seq;
    int head = h / HDIM;
    int d = h % HDIM;

    long long base_l1 = static_cast<long long>(token) * LIN1DIM;
    long long dst = (((static_cast<long long>(b) * HEADS + head) * seq) + s) * HDIM + d;

    q[dst] = l1o[base_l1 + h];
    k[dst] = l1o[base_l1 + HIDDEN + h];
    v[dst] = l1o[base_l1 + 2 * HIDDEN + h];
}

__global__ __launch_bounds__(256, 6)
void copy_mlp_kernel(const half* __restrict__ l1o,
                     half* __restrict__ mlp,
                     int rows)
{
    int total = rows * FFDIM;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base + 7 < total) {
        int row = base / FFDIM;
        int col = base % FFDIM;
        if (col + 7 < FFDIM) {
            float4 v = *reinterpret_cast<const float4*>(l1o + static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col);
            *reinterpret_cast<float4*>(mlp + base) = v;
            return;
        }
    }
    for (int j = 0; j < 8 && base + j < total; j++) {
        int idx = base + j;
        int row = idx / FFDIM;
        int col = idx % FFDIM;
        mlp[idx] = l1o[static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col];
    }
}

// Fused: reads L1 MLP portion as BF16, adds bias, writes FP16
__global__ __launch_bounds__(256, 6)
void copy_mlp_bf16bias_kernel(const void* __restrict__ l1o_raw,
                              const half* __restrict__ l1_bias,
                              half* __restrict__ mlp,
                              int rows)
{
    const __nv_bfloat16* l1o_bf = reinterpret_cast<const __nv_bfloat16*>(l1o_raw);
    int total = rows * FFDIM;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base + 7 < total) {
        int row = base / FFDIM;
        int col = base % FFDIM;
        if (col + 7 < FFDIM) {
            VecBF16_8 sv;
            sv.f4 = *reinterpret_cast<const float4*>(l1o_bf + static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col);
            VecHalf8 bv;
            bv.f4 = *reinterpret_cast<const float4*>(l1_bias + 3 * HIDDEN + col);
            VecHalf8 ov;
            #pragma unroll
            for (int j = 0; j < 8; j++)
                ov.h[j] = __float2half_rn(__bfloat162float(sv.b[j]) + __half2float(bv.h[j]));
            *reinterpret_cast<float4*>(mlp + base) = ov.f4;
            return;
        }
    }
    for (int j = 0; j < 8 && base + j < total; j++) {
        int idx = base + j;
        int row = idx / FFDIM;
        int col = idx % FFDIM;
        float val = __bfloat162float(l1o_bf[static_cast<long long>(row) * LIN1DIM + 3 * HIDDEN + col])
                  + __half2float(l1_bias[3 * HIDDEN + col]);
        mlp[idx] = __float2half_rn(val);
    }
}

// Simple bulk copy with raw float4 — minimal registers
__global__ __launch_bounds__(256, 6)
void extract_gate_kernel(const half* __restrict__ mod,
                         half* __restrict__ gate,
                         int batch)
{
    int total = batch * HIDDEN;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base + 7 < total) {
        int b = base / HIDDEN;
        int h = base % HIDDEN;
        if (h + 7 < HIDDEN) {
            float4 v = *reinterpret_cast<const float4*>(mod + static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h);
            *reinterpret_cast<float4*>(gate + base) = v;
            return;
        }
    }
    for (int j = 0; j < 8 && base + j < total; j++) {
        int idx = base + j;
        int b = idx / HIDDEN;
        int h = idx % HIDDEN;
        gate[idx] = mod[static_cast<long long>(b) * MODDIM + 2 * HIDDEN + h];
    }
}

__global__ __launch_bounds__(256, 2)
void rmsnorm_scale_kernel(half* __restrict__ data,
                          const half* __restrict__ scale)
{
    int row = blockIdx.x;
    half* row_ptr = data + static_cast<long long>(row) * HDIM;

    float sq = 0.0f;
    // Vectorized reduction: 4 elements per iteration
    for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
        VecHalf4 v;
        v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float val = __half2float(v.h[j]);
            sq += val * val;
        }
    }
    float inv = rsqrtf(block_reduce_sum(sq) / HDIM + RMS_EPS);

    // Vectorized writeback
    for (int i = threadIdx.x * 4; i < HDIM; i += blockDim.x * 4) {
        VecHalf4 v, sc;
        v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
        sc.f2 = *reinterpret_cast<const float2*>(scale + i);
        #pragma unroll
        for (int j = 0; j < 4; j++)
            v.h[j] = __float2half_rn(__half2float(v.h[j]) * inv * __half2float(sc.h[j]));
        *reinterpret_cast<float2*>(row_ptr + i) = v.f2;
    }
}

__global__ __launch_bounds__(256, 4)
void rope_kernel(half* __restrict__ q,
                 half* __restrict__ k,
                 const half* __restrict__ pe,
                 int batch,
                 int seq)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch * HEADS * seq * (HDIM / 2);
    if (idx >= total) {
        return;
    }

    int pair = idx % (HDIM / 2);
    int t0 = idx / (HDIM / 2);
    int s = t0 % seq;
    int t1 = t0 / seq;
    int h = t1 % HEADS;
    int b = t1 / HEADS;

    long long base = (((static_cast<long long>(b) * HEADS + h) * seq) + s) * HDIM + 2 * pair;
    long long pe_off = (static_cast<long long>(s) * (HDIM / 2) + pair) * 2;

    // Load cos/sin pair
    VecHalf2 pe_v;
    pe_v.u = *reinterpret_cast<const unsigned int*>(pe + pe_off);
    float c = __half2float(pe_v.h[0]);
    float sn = __half2float(pe_v.h[1]);

    // Load and rotate q pair
    VecHalf2 qv;
    qv.u = *reinterpret_cast<const unsigned int*>(q + base);
    float q0 = __half2float(qv.h[0]);
    float q1 = __half2float(qv.h[1]);
    qv.h[0] = __float2half_rn(q0 * c - q1 * sn);
    qv.h[1] = __float2half_rn(q0 * sn + q1 * c);
    *reinterpret_cast<unsigned int*>(q + base) = qv.u;

    // Load and rotate k pair
    VecHalf2 kv;
    kv.u = *reinterpret_cast<const unsigned int*>(k + base);
    float k0 = __half2float(kv.h[0]);
    float k1 = __half2float(kv.h[1]);
    kv.h[0] = __float2half_rn(k0 * c - k1 * sn);
    kv.h[1] = __float2half_rn(k0 * sn + k1 * c);
    *reinterpret_cast<unsigned int*>(k + base) = kv.u;
}

// 3-pass vectorized softmax: proven faster than online variant
__global__ __launch_bounds__(256, 2)
void softmax_kernel(half* __restrict__ scores,
                    int seq)
{
    int row = blockIdx.x;
    half* row_ptr = scores + static_cast<long long>(row) * seq;

    // Pass 1: find max
    float local_max = -1.0e30f;
    int seq4 = seq & ~3;
    for (int i = threadIdx.x * 4; i < seq4; i += blockDim.x * 4) {
        VecHalf4 v;
        v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
        #pragma unroll
        for (int j = 0; j < 4; j++)
            local_max = fmaxf(local_max, __half2float(v.h[j]));
    }
    for (int i = seq4 + threadIdx.x; i < seq; i += blockDim.x)
        local_max = fmaxf(local_max, __half2float(row_ptr[i]));
    float maxv = block_reduce_max(local_max);

    // Pass 2: exp + sum
    float local_sum = 0.0f;
    for (int i = threadIdx.x * 4; i < seq4; i += blockDim.x * 4) {
        VecHalf4 v;
        v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            float e = expf(__half2float(v.h[j]) - maxv);
            v.h[j] = __float2half_rn(e);
            local_sum += e;
        }
        *reinterpret_cast<float2*>(row_ptr + i) = v.f2;
    }
    for (int i = seq4 + threadIdx.x; i < seq; i += blockDim.x) {
        float e = expf(__half2float(row_ptr[i]) - maxv);
        row_ptr[i] = __float2half_rn(e);
        local_sum += e;
    }
    float inv = 1.0f / block_reduce_sum(local_sum);

    // Pass 3: normalize
    for (int i = threadIdx.x * 4; i < seq4; i += blockDim.x * 4) {
        VecHalf4 v;
        v.f2 = *reinterpret_cast<const float2*>(row_ptr + i);
        #pragma unroll
        for (int j = 0; j < 4; j++)
            v.h[j] = __float2half_rn(__half2float(v.h[j]) * inv);
        *reinterpret_cast<float2*>(row_ptr + i) = v.f2;
    }
    for (int i = seq4 + threadIdx.x; i < seq; i += blockDim.x)
        row_ptr[i] = __float2half_rn(__half2float(row_ptr[i]) * inv);
}

__global__ __launch_bounds__(256, 6)
void concat_heads_kernel(const half* __restrict__ heads,
                         half* __restrict__ out,
                         int batch,
                         int seq)
{
    int total = batch * seq * HIDDEN;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if (base + 3 >= total) {
        for (int j = 0; j < 4 && base + j < total; j++) {
            int idx = base + j;
            int h = idx % HIDDEN;
            int token = idx / HIDDEN;
            int s = token % seq;
            int b = token / seq;
            long long src = (((static_cast<long long>(b) * HEADS + h / HDIM) * seq) + s) * HDIM + h % HDIM;
            out[idx] = heads[src];
        }
        return;
    }

    int h = base % HIDDEN;
    int token = base / HIDDEN;
    int s = token % seq;
    int b = token / seq;
    int d = h % HDIM;

    // HDIM=128 is a multiple of 4, so boundary crossings are at d=125,126,127 only (rare)
    if (d + 3 < HDIM) {
        int head = h / HDIM;
        long long src = (((static_cast<long long>(b) * HEADS + head) * seq) + s) * HDIM + d;
        float2 v = *reinterpret_cast<const float2*>(heads + src);
        *reinterpret_cast<float2*>(out + base) = v;
    } else {
        #pragma unroll
        for (int j = 0; j < 4; j++) {
            int hj = h + j;
            long long src = (((static_cast<long long>(b) * HEADS + hj / HDIM) * seq) + s) * HDIM + hj % HDIM;
            out[base + j] = heads[src];
        }
    }
}

__global__ __launch_bounds__(256, 4)
void gelu_tanh_kernel(half* __restrict__ data,
                      int n)
{
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base + 7 < n) {
        VecHalf8 v;
        v.f4 = *reinterpret_cast<const float4*>(data + base);
        #pragma unroll
        for (int j = 0; j < 8; j++) {
            float x = __half2float(v.h[j]);
            float x3 = x * x * x;
            v.h[j] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
        }
        *reinterpret_cast<float4*>(data + base) = v.f4;
    } else {
        for (int i = base; i < min(base + 8, n); i++) {
            float x = __half2float(data[i]);
            float x3 = x * x * x;
            data[i] = __float2half_rn(0.5f * x * (1.0f + tanhf(0.7978845608f * (x + 0.044715f * x3))));
        }
    }
}

__global__ __launch_bounds__(256, 6)
void concat_final_kernel(const half* __restrict__ attn,
                         const half* __restrict__ mlp,
                         half* __restrict__ out,
                         int rows)
{
    int total = rows * LIN2DIM;
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base + 7 < total) {
        int row = base / LIN2DIM;
        int col = base % LIN2DIM;
        if (col + 7 < LIN2DIM) {
            float4 v;
            if (col + 7 < HIDDEN) {
                v = *reinterpret_cast<const float4*>(attn + static_cast<long long>(row) * HIDDEN + col);
            } else if (col >= HIDDEN) {
                v = *reinterpret_cast<const float4*>(mlp + static_cast<long long>(row) * FFDIM + (col - HIDDEN));
            } else {
                // Crosses attn/mlp boundary — scalar
                VecHalf8 vh;
                #pragma unroll
                for (int j = 0; j < 8; j++) {
                    int c = col + j;
                    if (c < HIDDEN)
                        vh.h[j] = attn[static_cast<long long>(row) * HIDDEN + c];
                    else
                        vh.h[j] = mlp[static_cast<long long>(row) * FFDIM + (c - HIDDEN)];
                }
                v = vh.f4;
            }
            *reinterpret_cast<float4*>(out + base) = v;
            return;
        }
    }
    for (int j = 0; j < 8 && base + j < total; j++) {
        int idx = base + j;
        int row = idx / LIN2DIM;
        int col = idx % LIN2DIM;
        if (col < HIDDEN)
            out[idx] = attn[static_cast<long long>(row) * HIDDEN + col];
        else
            out[idx] = mlp[static_cast<long long>(row) * FFDIM + (col - HIDDEN)];
    }
}

__global__ void gate_residual_kernel(const half* __restrict__ residual,
                                     const half* __restrict__ gate,
                                     half* __restrict__ out,
                                     int seq)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = gridDim.y * seq * HIDDEN;
    int linear = blockIdx.y * (seq * HIDDEN) + idx;
    if (linear >= total) {
        return;
    }

    int h = linear % HIDDEN;
    int token = linear / HIDDEN;
    int b = token / seq;
    float gate_v = __half2float(gate[static_cast<long long>(b) * HIDDEN + h]);
    float res = __half2float(residual[linear]);
    float v = __half2float(out[linear]);
    out[linear] = __float2half_rn(res + gate_v * v);
}

// (gate_residual_rescale_bias_kernel removed — Python handles gate+residual in bf16)

SSBEngine::SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size)
    : B_(batch),
      S_(seq),
      cub_(nullptr),
      workspace_(nullptr),
      workspace_bytes_(0),
      d_silu_(nullptr),
      d_mod_(nullptr),
      d_xmod_(nullptr),
      d_l1o_(nullptr),
      d_act_fp8_(nullptr),
      d_act_amax_(nullptr),
      d_act_scale_(nullptr),
      d_out_scale_(nullptr)
{
    BK(cublasCreate(&cub_));
    BK(cublasSetMathMode(cub_, CUBLAS_TENSOR_OP_MATH));

    size_t min_needed = ssb_workspace_bytes(B_, S_);
    if (workspace_buf_size < min_needed) {
        throw std::runtime_error("Workspace buffer too small: need " +
            std::to_string(min_needed) + " got " + std::to_string(workspace_buf_size));
    }

    // Carve the Python-allocated workspace into sub-buffers (zero cudaMalloc)
    char* p = static_cast<char*>(workspace_buf);
    size_t off = 0;
    size_t rows = static_cast<size_t>(B_) * S_;

    #define CARVE(ptr, type, count) \
        off = (off + 255) & ~size_t(255); \
        ptr = reinterpret_cast<type*>(p + off); \
        off += static_cast<size_t>(count) * sizeof(type);

    CARVE(d_silu_, half, B_ * HIDDEN)
    CARVE(d_mod_, half, static_cast<size_t>(B_) * MODDIM)
    CARVE(d_xmod_, half, rows * HIDDEN)
    CARVE(d_l1o_, half, rows * LIN1DIM)
    CARVE(d_act_fp8_, __nv_fp8_e4m3, rows * LIN2DIM)
    CARVE(d_act_amax_, unsigned int, 1)
    CARVE(d_act_scale_, float, 1)
    CARVE(d_out_scale_, float, 1)  // unused by fused GEMM, kept for header compat

    #undef CARVE

    off = (off + 255) & ~size_t(255);
    workspace_ = p + off;
    workspace_bytes_ = workspace_buf_size - off;
}

SSBEngine::~SSBEngine()
{
    // No cudaFree — all memory owned by Python tensors
    if (cub_) cublasDestroy(cub_);
}

void SSBEngine::prepare(const half* d_x,
                        const half* d_vec,
                        const half* d_pe,
                        const void* mod_w,
                        const float* mod_w_scale,
                        const half* mod_b,
                        const void* l1w,
                        const float* l1w_scale,
                        const half* l1b,
                        const half* qns,
                        const half* kns,
                        int fp8_mask,
                        half* d_q,
                        half* d_k,
                        half* d_v,
                        half* d_mlp,
                        half* d_gate,
                        cudaStream_t stream)
{
    BK(cublasSetStream(cub_, stream));

    int rows = B_ * S_;
    int threads = kKernelThreads;
    bool mod_fp8 = (fp8_mask & FP8_MOD_W) != 0;
    bool l1_fp8 = (fp8_mask & FP8_L1_W) != 0;

    // Vectorized silu: each thread handles 8 elements
    silu_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_vec, d_silu_, B_ * HIDDEN);

    // Modulation GEMM: FP8→BF16 output (bias fused into downstream kernels)
    if (mod_fp8) {
        linear_gemm_fp8(B_, MODDIM, HIDDEN,
            d_silu_, static_cast<const __nv_fp8_e4m3*>(mod_w), mod_w_scale,
            reinterpret_cast<__nv_bfloat16*>(d_mod_),
            d_act_fp8_, d_act_amax_, d_act_scale_, d_out_scale_, workspace_, workspace_bytes_, stream);
        // No bf16_to_fp16_bias_add — bias fused into layernorm_mod and extract_gate below
    } else {
        linear_gemm(cub_, B_, MODDIM, HIDDEN, d_silu_, static_cast<const half*>(mod_w), d_mod_);
        int out_n = B_ * MODDIM;
        bias_add_kernel<<<(out_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_mod_, mod_b, B_, MODDIM);
    }

    // d_x contains bf16 bits (model runs in bf16). Kernel reads as bf16, outputs fp16.
    if (mod_fp8) {
        // Fused: reads mod as BF16, adds bias inline during layernorm
        layernorm_mod_bf16bias_kernel<<<rows, threads, 0, stream>>>(reinterpret_cast<const __nv_bfloat16*>(d_x), d_mod_, mod_b, d_xmod_, S_);
    } else {
        layernorm_mod_kernel<<<rows, threads, 0, stream>>>(reinterpret_cast<const __nv_bfloat16*>(d_x), d_mod_, d_xmod_, S_);
    }

    // Linear1 GEMM: FP8→BF16 output (bias fused into split_qkv and copy_mlp)
    if (l1_fp8) {
        linear_gemm_fp8(rows, LIN1DIM, HIDDEN,
            d_xmod_, static_cast<const __nv_fp8_e4m3*>(l1w), l1w_scale,
            reinterpret_cast<__nv_bfloat16*>(d_l1o_),
            d_act_fp8_, d_act_amax_, d_act_scale_, d_out_scale_, workspace_, workspace_bytes_, stream);
        // No bf16_to_fp16_bias_add — bias fused into split_qkv and copy_mlp below
    } else {
        linear_gemm(cub_, rows, LIN1DIM, HIDDEN, d_xmod_, static_cast<const half*>(l1w), d_l1o_);
        int out_n = rows * LIN1DIM;
        bias_add_kernel<<<(out_n + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(d_l1o_, l1b, rows, LIN1DIM);
    }

    dim3 split_grid((S_ * HIDDEN + threads - 1) / threads, B_);
    if (l1_fp8) {
        // Fused: reads L1 output as BF16, adds bias inline during split
        split_qkv_bf16bias_kernel<<<split_grid, threads, 0, stream>>>(d_l1o_, l1b, d_q, d_k, d_v, S_);
        copy_mlp_bf16bias_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_l1o_, l1b, d_mlp, rows);
    } else {
        split_qkv_kernel<<<split_grid, threads, 0, stream>>>(d_l1o_, d_q, d_k, d_v, S_);
        copy_mlp_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_l1o_, d_mlp, rows);
    }
    if (mod_fp8) {
        // Fused: reads gate from mod BF16, adds bias inline
        extract_gate_bf16bias_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mod_, mod_b, d_gate, B_);
    } else {
        extract_gate_kernel<<<(B_ * HIDDEN + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mod_, d_gate, B_);
    }

    int qk_rows = B_ * HEADS * S_;
    rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_q, qns);
    rmsnorm_scale_kernel<<<qk_rows, threads, 0, stream>>>(d_k, kns);

    int rope_total = B_ * HEADS * S_ * (HDIM / 2);
    rope_kernel<<<(rope_total + threads - 1) / threads, threads, 0, stream>>>(d_q, d_k, d_pe, B_, S_);

    CK(cudaGetLastError());
}

void SSBEngine::finish(const half* d_x,
                       const half* d_attn,
                       half* d_mlp,
                       const half* d_gate,
                       half* d_out,
                       const void* l2w,
                       const float* l2w_scale,
                       const half* l2b,
                       int fp8_mask,
                       cudaStream_t stream)
{
    BK(cublasSetStream(cub_, stream));

    int rows = B_ * S_;
    int threads = kKernelThreads;
    bool l2_fp8 = (fp8_mask & FP8_L2_W) != 0;

    // Vectorized: each thread handles 8 elements
    gelu_tanh_kernel<<<(rows * FFDIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_mlp, rows * FFDIM);
    concat_final_kernel<<<(rows * LIN2DIM + threads * 8 - 1) / (threads * 8), threads, 0, stream>>>(d_attn, d_mlp, d_l1o_, rows);

    if (l2_fp8) {
        // Linear2: FP8→BF16 output, no bias (Python handles bias+gate+residual in fp32)
        linear_gemm_fp8(rows, HIDDEN, LIN2DIM,
            d_l1o_, static_cast<const __nv_fp8_e4m3*>(l2w), l2w_scale,
            reinterpret_cast<__nv_bfloat16*>(d_out),
            d_act_fp8_, d_act_amax_, d_act_scale_, d_out_scale_, workspace_, workspace_bytes_, stream);
        // d_out contains BF16 linear2 output. Python reads as torch.bfloat16.
    } else {
        linear_gemm(cub_, rows, HIDDEN, LIN2DIM, d_l1o_, static_cast<const half*>(l2w), d_out);
        int out_blocks = (rows * HIDDEN + threads * 4 - 1) / (threads * 4);
        bias_add_kernel<<<out_blocks, threads, 0, stream>>>(d_out, l2b, rows, HIDDEN);

        dim3 gate_grid((S_ * HIDDEN + threads - 1) / threads, B_);
        gate_residual_kernel<<<gate_grid, threads, 0, stream>>>(d_x, d_gate, d_out, S_);
    }

    CK(cudaGetLastError());
}

void SSBEngine::forward(const half* d_x,
                        const half* d_vec,
                        half* d_out,
                        const half* d_pe,
                        const void* mod_w,
                        const float* mod_w_scale,
                        const half* mod_b,
                        const void* l1w,
                        const float* l1w_scale,
                        const half* l1b,
                        const void* l2w,
                        const float* l2w_scale,
                        const half* l2b,
                        const half* qns,
                        const half* kns,
                        int fp8_mask,
                        cudaStream_t stream)
{
    size_t rows = static_cast<size_t>(B_) * S_;
    int threads = kKernelThreads;
    int qk_rows = B_ * HEADS * S_;
    size_t qkv_elems = static_cast<size_t>(B_) * HEADS * S_ * HDIM;

    // Temporary attention buffers (test driver only — not called from ComfyUI)
    half *t_q, *t_k, *t_v, *t_mlp, *t_gate, *t_scores, *t_attn;
    CK(cudaMalloc(&t_q, qkv_elems * sizeof(half)));
    CK(cudaMalloc(&t_k, qkv_elems * sizeof(half)));
    CK(cudaMalloc(&t_v, qkv_elems * sizeof(half)));
    CK(cudaMalloc(&t_mlp, rows * FFDIM * sizeof(half)));
    CK(cudaMalloc(&t_gate, static_cast<size_t>(B_) * HIDDEN * sizeof(half)));
    CK(cudaMalloc(&t_scores, static_cast<size_t>(B_) * HEADS * S_ * S_ * sizeof(half)));
    CK(cudaMalloc(&t_attn, rows * HIDDEN * sizeof(half)));

    prepare(d_x, d_vec, d_pe, mod_w, mod_w_scale, mod_b, l1w, l1w_scale, l1b, qns, kns, fp8_mask, t_q, t_k, t_v, t_mlp, t_gate, stream);
    qk_gemm(cub_, B_ * HEADS, S_, HDIM, t_k, t_q, t_scores);
    softmax_kernel<<<qk_rows, threads, 0, stream>>>(t_scores, S_);
    sv_gemm(cub_, B_ * HEADS, S_, HDIM, t_v, t_scores, t_q);
    concat_heads_kernel<<<(rows * HIDDEN + threads * 4 - 1) / (threads * 4), threads, 0, stream>>>(t_q, t_attn, B_, S_);
    finish(d_x, t_attn, t_mlp, t_gate, d_out, l2w, l2w_scale, l2b, fp8_mask, stream);

    cudaFree(t_attn);
    cudaFree(t_scores);
    cudaFree(t_gate);
    cudaFree(t_mlp);
    cudaFree(t_v);
    cudaFree(t_k);
    cudaFree(t_q);
    CK(cudaGetLastError());
}

} // namespace transformer

#ifndef FLUX_TORCH_EXT  // Skip DLL exports when building as PyTorch extension

extern "C" {

using namespace transformer;

namespace {
thread_local std::string g_last_error;

void clear_last_error()
{
    g_last_error.clear();
}

void set_last_error(const std::string& message)
{
    g_last_error = message;
}
}

__declspec(dllexport) size_t ssb_workspace_size(int batch, int seq_len)
{
    return ssb_workspace_bytes(batch, seq_len);
}

__declspec(dllexport) void* ssb_create(int batch, int seq_len, void* workspace_buf, size_t workspace_buf_size)
{
    try {
        clear_last_error();
        return new SSBEngine(batch, seq_len, workspace_buf, workspace_buf_size);
    } catch (const std::exception& e) {
        set_last_error(e.what());
        fprintf(stderr, "[flux_cuda_kernel] ssb_create failed: %s\n", e.what());
        fflush(stderr);
        return nullptr;
    }
}

__declspec(dllexport) int ssb_destroy(void* engine)
{
    try {
        clear_last_error();
        delete static_cast<SSBEngine*>(engine);
        return 0;
    } catch (const std::exception& e) {
        set_last_error(e.what());
        fprintf(stderr, "[flux_cuda_kernel] ssb_destroy failed: %s\n", e.what());
        fflush(stderr);
        return -1;
    }
}

__declspec(dllexport) int ssb_prepare(void* engine,
                                      const half* d_x,
                                      const half* d_vec,
                                      const half* d_pe,
                                      const void* mod_w,
                                      const float* mod_w_scale,
                                      const half* mod_b,
                                      const void* l1w,
                                      const float* l1w_scale,
                                      const half* l1b,
                                      const half* qns,
                                      const half* kns,
                                      int fp8_mask,
                                      half* d_q,
                                      half* d_k,
                                      half* d_v,
                                      half* d_mlp,
                                      half* d_gate,
                                      void* stream_handle)
{
    try {
        clear_last_error();
        if (engine == nullptr) {
            set_last_error("SSBEngine handle is null");
            fprintf(stderr, "[flux_cuda_kernel] ssb_prepare failed: %s\n", g_last_error.c_str());
            fflush(stderr);
            return -1;
        }
        cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_handle);
        static_cast<SSBEngine*>(engine)->prepare(
            d_x, d_vec, d_pe,
            mod_w, mod_w_scale, mod_b,
            l1w, l1w_scale, l1b,
            qns, kns,
            fp8_mask,
            d_q, d_k, d_v, d_mlp, d_gate,
            stream);
        return 0;
    } catch (const std::exception& e) {
        set_last_error(e.what());
        fprintf(stderr, "[flux_cuda_kernel] ssb_prepare failed: %s\n", e.what());
        fflush(stderr);
        return -1;
    }
}

__declspec(dllexport) int ssb_finish(void* engine,
                                     const half* d_x,
                                     const half* d_attn,
                                     half* d_mlp,
                                     const half* d_gate,
                                     half* d_out,
                                     const void* l2w,
                                     const float* l2w_scale,
                                     const half* l2b,
                                     int fp8_mask,
                                     void* stream_handle)
{
    try {
        clear_last_error();
        if (engine == nullptr) {
            set_last_error("SSBEngine handle is null");
            fprintf(stderr, "[flux_cuda_kernel] ssb_finish failed: %s\n", g_last_error.c_str());
            fflush(stderr);
            return -1;
        }
        cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_handle);
        static_cast<SSBEngine*>(engine)->finish(
            d_x, d_attn, d_mlp, d_gate, d_out,
            l2w, l2w_scale, l2b,
            fp8_mask,
            stream);
        return 0;
    } catch (const std::exception& e) {
        set_last_error(e.what());
        fprintf(stderr, "[flux_cuda_kernel] ssb_finish failed: %s\n", e.what());
        fflush(stderr);
        return -1;
    }
}

__declspec(dllexport) int ssb_forward(void* engine,
                                      const half* d_x,
                                      const half* d_vec,
                                      half* d_out,
                                      const half* d_pe,
                                      const void* mod_w,
                                      const float* mod_w_scale,
                                      const half* mod_b,
                                      const void* l1w,
                                      const float* l1w_scale,
                                      const half* l1b,
                                      const void* l2w,
                                      const float* l2w_scale,
                                      const half* l2b,
                                      const half* qns,
                                      const half* kns,
                                      int fp8_mask,
                                      void* stream_handle)
{
    try {
        clear_last_error();
        if (engine == nullptr) {
            set_last_error("SSBEngine handle is null");
            fprintf(stderr, "[flux_cuda_kernel] ssb_forward failed: %s\n", g_last_error.c_str());
            fflush(stderr);
            return -1;
        }
        cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_handle);
        static_cast<SSBEngine*>(engine)->forward(
            d_x, d_vec, d_out,
            d_pe,
            mod_w, mod_w_scale, mod_b,
            l1w, l1w_scale, l1b,
            l2w, l2w_scale, l2b,
            qns, kns,
            fp8_mask,
            stream);
        return 0;
    } catch (const std::exception& e) {
        set_last_error(e.what());
        fprintf(stderr, "[flux_cuda_kernel] ssb_forward failed: %s\n", e.what());
        fflush(stderr);
        return -1;
    }
}

__declspec(dllexport) const char* ssb_last_error()
{
    return g_last_error.empty() ? nullptr : g_last_error.c_str();
}

} // extern "C"

#endif // FLUX_TORCH_EXT