ULT7RA-FLUXNATION/csrc/FLUXNATION.cu

/*
 * flux_ext_kernels.cu — PyTorch C++ extension for FLUXNATION
 *
 * Wraps SSBEngine (from flux_opt.cu) as torch.ops custom operators.
 * Adds fused gate+residual+bias kernel to eliminate Python-side overhead.
 */

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

// SSBEngine class + constants from existing code
#include "SINGLE_STREAM_BLOCK.H"

#include <unordered_map>
#include <memory>
#include <mutex>

using namespace transformer;

// =========================================================================
// Vectorized load/store types (same as flux_opt.cu but local to this TU)
// =========================================================================
union VecBF16_8_ext {
    float4 f4;
    __nv_bfloat16 b[8];
};
union VecHalf8_ext {
    float4 f4;
    half h[8];
};

// =========================================================================
// NEW: Fused gate + residual + L2-bias kernel
//
// result[b,s,h] = x[b,s,h] + gate[b,h] * (l2_out[b,s,h] + l2b[h])
//
// All reads/writes vectorized 8-wide. Eliminates 5 intermediate float32
// tensors that the Python path previously created per block.
// =========================================================================
__global__ void gate_residual_bias_bf16_kernel(
    const __nv_bfloat16* __restrict__ d_l2out,   // [rows, HIDDEN] L2 GEMM output (bf16)
    const __nv_bfloat16* __restrict__ d_x,       // [rows, HIDDEN] original input / residual (bf16)
    const half*          __restrict__ d_gate,     // [B, HIDDEN] gate values (fp16)
    const half*          __restrict__ l2b,        // [HIDDEN] L2 bias (fp16)
    __nv_bfloat16*       __restrict__ d_result,   // [rows, HIDDEN] output (bf16)
    int S,
    int total_elems)
{
    int base = (blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base >= total_elems) return;

    int row = base / HIDDEN;
    int col = base % HIDDEN;
    int batch = row / S;

    // Vectorized loads
    VecBF16_8_ext out_v = *reinterpret_cast<const VecBF16_8_ext*>(&d_l2out[base]);
    VecBF16_8_ext x_v   = *reinterpret_cast<const VecBF16_8_ext*>(&d_x[base]);
    VecHalf8_ext  gate_v = *reinterpret_cast<const VecHalf8_ext*>(&d_gate[batch * HIDDEN + col]);
    VecHalf8_ext  bias_v = *reinterpret_cast<const VecHalf8_ext*>(&l2b[col]);

    VecBF16_8_ext result_v;
    #pragma unroll
    for (int i = 0; i < 8; i++) {
        float l2_f = __bfloat162float(out_v.b[i]);
        float x_f  = __bfloat162float(x_v.b[i]);
        float g_f  = __half2float(gate_v.h[i]);
        float b_f  = __half2float(bias_v.h[i]);
        result_v.b[i] = __float2bfloat16(x_f + g_f * (l2_f + b_f));
    }
    *reinterpret_cast<VecBF16_8_ext*>(&d_result[base]) = result_v;
}

// =========================================================================
// Global engine cache — keyed by (device, batch, seq_len)
// =========================================================================
struct EngineEntry {
    std::unique_ptr<SSBEngine> engine;
    at::Tensor workspace;  // prevents deallocation
};

static std::unordered_map<int64_t, EngineEntry> g_engine_cache;
static std::mutex g_engine_mutex;

static int64_t make_engine_key(int device, int batch, int seq) {
    return (int64_t)device * 1000000000LL + (int64_t)batch * 1000000LL + seq;
}

static SSBEngine* get_or_create_engine(int batch, int seq, at::Device device) {
    int dev_idx = device.index();
    int64_t key = make_engine_key(dev_idx, batch, seq);

    std::lock_guard<std::mutex> lock(g_engine_mutex);
    auto it = g_engine_cache.find(key);
    if (it != g_engine_cache.end()) {
        return it->second.engine.get();
    }

    // Create new engine
    at::cuda::CUDAGuard guard(device);
    size_t ws_bytes = ssb_workspace_bytes(batch, seq);
    auto ws_tensor = at::empty({(int64_t)ws_bytes}, at::TensorOptions().dtype(at::kByte).device(device));

    auto engine = std::make_unique<SSBEngine>(batch, seq, ws_tensor.data_ptr(), ws_bytes);
    auto* ptr = engine.get();

    EngineEntry entry;
    entry.engine = std::move(engine);
    entry.workspace = ws_tensor;
    g_engine_cache[key] = std::move(entry);

    fprintf(stderr, "[flux_ext] Created engine B=%d S=%d ws=%.1fMB dev=%d\n",
            batch, seq, ws_bytes / (1024.0 * 1024.0), dev_idx);
    fflush(stderr);

    return ptr;
}

// =========================================================================
// torch.ops implementation: ssb_prepare
//
// Runs: silu → mod GEMM → layernorm_mod → linear1 GEMM → split_qkv →
//       copy_mlp → extract_gate → rmsnorm → rope
//
// Returns: (q, k, v, mlp, gate)
// =========================================================================
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
ssb_prepare_impl(
    at::Tensor x,            // [B, S, 3072] bf16
    at::Tensor vec,          // [B, 3072] or [B, 1, 3072] fp16
    at::Tensor pe,           // [S, 64, 2] fp16 (pre-converted)
    at::Tensor mod_w,        // FP8 [MODDIM, HIDDEN]
    at::Tensor mod_w_scale,  // float32 [1]
    at::Tensor mod_b,        // fp16 [MODDIM]
    at::Tensor l1w,          // FP8 [LIN1DIM, HIDDEN]
    at::Tensor l1w_scale,    // float32 [1]
    at::Tensor l1b,          // fp16 [LIN1DIM]
    at::Tensor qns,          // fp16 [128]
    at::Tensor kns)          // fp16 [128]
{
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");

    int B = x.size(0);
    int S = x.size(1);
    auto device = x.device();
    at::cuda::CUDAGuard guard(device);
    auto stream = at::cuda::getCurrentCUDAStream(device.index());

    // Ensure correct dtypes and contiguity
    auto x_bf16 = x.to(at::kBFloat16).contiguous();
    auto vec_fp16 = vec.dim() == 3 ? vec.squeeze(1).to(at::kHalf).contiguous()
                                   : vec.to(at::kHalf).contiguous();
    auto pe_fp16 = pe.to(at::kHalf).contiguous();

    // Allocate output tensors
    auto opts_fp16 = at::TensorOptions().dtype(at::kHalf).device(device);
    auto q    = at::empty({B, HEADS, S, HDIM}, opts_fp16);
    auto k    = at::empty({B, HEADS, S, HDIM}, opts_fp16);
    auto v    = at::empty({B, HEADS, S, HDIM}, opts_fp16);
    auto mlp  = at::empty({B, S, FFDIM}, opts_fp16);
    auto gate = at::empty({B, HIDDEN}, opts_fp16);

    // Get or create engine (manages workspace internally)
    SSBEngine* engine = get_or_create_engine(B, S, device);

    // Call prepare — all pointer casts match existing kernel expectations
    engine->prepare(
        reinterpret_cast<const half*>(x_bf16.data_ptr()),
        reinterpret_cast<const half*>(vec_fp16.data_ptr()),
        reinterpret_cast<const half*>(pe_fp16.data_ptr()),
        mod_w.data_ptr(),                                               // FP8 weights (void*)
        mod_w_scale.defined() ? reinterpret_cast<const float*>(mod_w_scale.data_ptr()) : nullptr,
        reinterpret_cast<const half*>(mod_b.data_ptr()),
        l1w.data_ptr(),                                                 // FP8 weights (void*)
        l1w_scale.defined() ? reinterpret_cast<const float*>(l1w_scale.data_ptr()) : nullptr,
        reinterpret_cast<const half*>(l1b.data_ptr()),
        reinterpret_cast<const half*>(qns.data_ptr()),
        reinterpret_cast<const half*>(kns.data_ptr()),
        7,  // fp8_mask = all FP8
        reinterpret_cast<half*>(q.data_ptr()),
        reinterpret_cast<half*>(k.data_ptr()),
        reinterpret_cast<half*>(v.data_ptr()),
        reinterpret_cast<half*>(mlp.data_ptr()),
        reinterpret_cast<half*>(gate.data_ptr()),
        stream);

    return std::make_tuple(q, k, v, mlp, gate);
}

// =========================================================================
// torch.ops implementation: ssb_finish
//
// Runs: gelu → concat → L2 GEMM → FUSED gate+residual+bias
//
// Returns: result [B, S, 3072] bf16
// =========================================================================
at::Tensor ssb_finish_impl(
    at::Tensor x,            // [B, S, 3072] bf16 — original input for residual
    at::Tensor attn,         // [B, S, 3072] fp16 — attention output
    at::Tensor mlp,          // [B, S, 12288] fp16 — MLP intermediate (pre-GELU)
    at::Tensor gate,         // [B, 3072] fp16 — gate values
    at::Tensor l2w,          // FP8 [HIDDEN, LIN2DIM]
    at::Tensor l2w_scale,    // float32 [1]
    at::Tensor l2b)          // fp16 [3072] — L2 bias
{
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.dim() == 3 && x.size(2) == HIDDEN, "x must be [B, S, 3072]");

    int B = x.size(0);
    int S = x.size(1);
    int rows = B * S;
    auto device = x.device();
    at::cuda::CUDAGuard guard(device);
    auto stream = at::cuda::getCurrentCUDAStream(device.index());

    auto x_bf16 = x.to(at::kBFloat16).contiguous();
    auto attn_fp16 = attn.to(at::kHalf).contiguous();
    // mlp is already fp16 from prepare — ensure contiguous
    auto mlp_c = mlp.contiguous();
    auto gate_c = gate.contiguous();
    auto l2b_c = l2b.to(at::kHalf).contiguous();

    // Temp buffer for L2 GEMM output (bf16) — ATen caching allocator makes this free
    auto l2_out = at::empty({B, S, HIDDEN}, at::TensorOptions().dtype(at::kBFloat16).device(device));

    SSBEngine* engine = get_or_create_engine(B, S, device);

    // Call existing finish() — does gelu, concat, L2 GEMM
    // NOTE: finish() FP8 path writes bf16 to d_out, does NOT add bias/gate/residual
    engine->finish(
        reinterpret_cast<const half*>(x_bf16.data_ptr()),    // d_x (bf16 in half*)
        reinterpret_cast<const half*>(attn_fp16.data_ptr()), // d_attn
        reinterpret_cast<half*>(mlp_c.data_ptr()),           // d_mlp (in/out)
        reinterpret_cast<const half*>(gate_c.data_ptr()),    // d_gate
        reinterpret_cast<half*>(l2_out.data_ptr()),          // d_out (bf16)
        l2w.data_ptr(),                                       // FP8 weights
        l2w_scale.defined() ? reinterpret_cast<const float*>(l2w_scale.data_ptr()) : nullptr,
        reinterpret_cast<const half*>(l2b_c.data_ptr()),     // l2b (unused by finish FP8 path)
        7,  // fp8_mask = all FP8
        stream);

    // FUSED: gate + residual + L2 bias in one kernel
    // result = x + gate * (l2_out + l2b)
    auto result = at::empty_like(x_bf16);
    int total_elems = rows * HIDDEN;
    int threads = 256;
    int blocks = (total_elems + threads * 8 - 1) / (threads * 8);

    gate_residual_bias_bf16_kernel<<<blocks, threads, 0, stream>>>(
        reinterpret_cast<const __nv_bfloat16*>(l2_out.data_ptr()),
        reinterpret_cast<const __nv_bfloat16*>(x_bf16.data_ptr()),
        reinterpret_cast<const half*>(gate_c.data_ptr()),
        reinterpret_cast<const half*>(l2b_c.data_ptr()),
        reinterpret_cast<__nv_bfloat16*>(result.data_ptr()),
        S,
        total_elems);

    return result;
}

// =========================================================================
// Cleanup function — destroys all cached engines
// =========================================================================
void ssb_cleanup_impl() {
    std::lock_guard<std::mutex> lock(g_engine_mutex);
    g_engine_cache.clear();
    fprintf(stderr, "[flux_ext] All engines destroyed\n");
    fflush(stderr);
}