ULT7RA-FLUXNATION/FLUX_STREAM.H

#ifndef FLUX_STREAM_H
#define FLUX_STREAM_H

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <cublas_v2.h>
#include <cstdio>
#include <stdexcept>
#include <string>

namespace transformer {

constexpr int HIDDEN  = 3072;
constexpr int HEADS   = 24;
constexpr int HDIM    = 128;
constexpr int FFDIM   = 12288;
constexpr int MODDIM  = 3 * HIDDEN;
constexpr int LIN1DIM = 3 * HIDDEN + FFDIM;
constexpr int LIN2DIM = HIDDEN + FFDIM;
constexpr float LN_EPS  = 1e-6f;
constexpr float RMS_EPS = 1e-6f;
constexpr float ATT_SC  = 0.0883883476f;
constexpr int FP8_MOD_W = 1;
constexpr int FP8_L1_W  = 2;
constexpr int FP8_L2_W  = 4;
#define CK(call) do { \
    cudaError_t err__ = (call); \
    if (err__ != cudaSuccess) { \
        throw std::runtime_error(std::string("CUDA ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + cudaGetErrorString(err__)); \
    } \
} while (0)

#define BK(call) do { \
    cublasStatus_t st__ = (call); \
    if (st__ != CUBLAS_STATUS_SUCCESS) { \
        throw std::runtime_error(std::string("cuBLAS ") + __FILE__ + ":" + std::to_string(__LINE__) + ": " + std::to_string((int)st__)); \
    } \
} while (0)

// Returns minimum workspace bytes needed for ssb_create
static inline size_t ssb_workspace_bytes(int batch, int seq) {
    size_t rows = static_cast<size_t>(batch) * seq;
    size_t b = static_cast<size_t>(batch);
    size_t sz = 0;
    // Each sub-buffer 256-byte aligned
    #define ADDBUF(count, elem_size) \
        sz = (sz + 255) & ~size_t(255); sz += (count) * (elem_size);
    ADDBUF(b * HIDDEN, 2);              // d_silu_
    ADDBUF(b * MODDIM, 2);              // d_mod_
    ADDBUF(rows * HIDDEN, 2);           // d_xmod_
    ADDBUF(rows * LIN1DIM, 2);          // d_l1o_
    ADDBUF(rows * LIN2DIM, 1);          // d_act_fp8_
    ADDBUF(1, 4);                       // d_act_amax_
    ADDBUF(1, 4);                       // d_act_scale_
    ADDBUF(1, 4);                       // d_out_scale_
    sz = (sz + 255) & ~size_t(255);
    sz += 32ull << 20;                  // 32MB CUTLASS/cuBLAS workspace
    #undef ADDBUF
    return sz;
}

class SSBEngine {
public:
    // workspace_buf: Python-allocated GPU memory (torch tensor .data_ptr())
    // workspace_buf_size: must be >= ssb_workspace_bytes(batch, seq)
    SSBEngine(int batch, int seq, void* workspace_buf, size_t workspace_buf_size);
    ~SSBEngine();

    void 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 = 0);

    void 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 = 0);

    void 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 = 0);

    int batch() const { return B_; }
    int seq() const { return S_; }

private:
    int B_;
    int S_;
    cublasHandle_t cub_;
    void* workspace_;
    size_t workspace_bytes_;

    half* d_silu_;
    half* d_mod_;
    half* d_xmod_;
    half* d_l1o_;
    __nv_fp8_e4m3* d_act_fp8_;
    unsigned int* d_act_amax_;
    float* d_act_scale_;
    float* d_out_scale_;
};

} // namespace transformer

#endif // FLUX_STREAM_H