AstrAI/csrc/kernels/fp8/gemm.cuh

#pragma once
// FP8 GEMM umbrella: the kernel orchestrator and the host-side launch
// planning. Device layers live in gemm/ (policy / load / scheduler /
// mainloop / epilogue) — pure CUDA, no torch; launchers are plain functions
// shared by the torch binding and the C tests. Layout tags and the NN swap
// semantics are documented in common.h and the design notes
// (docs/developer/cuda_kernels.md).

#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <type_traits>

#include "../common/cp_async.cuh"
#include "common.h"
#include "gemm/epilogue.cuh"
#include "gemm/load.cuh"
#include "gemm/mainloop.cuh"
#include "gemm/policy.cuh"
#include "gemm/scheduler.cuh"

namespace astrai {
namespace fp8 {

template <typename Policy>
__global__ void __launch_bounds__(Policy::kCtaThreads, Policy::kMinCtas)
    fp8_gemm_kernel(FP8Params p) {
    using Traits = typename Policy::Traits;
    using Mainloop = Fp8CollectiveMainloop<Policy>;
    using Epilogue = Fp8CollectiveEpilogue<Policy>;
    // Stages live in dynamic shared memory so deep pipelines (> 48KB
    // static limit) opt in via cudaFuncSetAttribute in the launcher.
    extern __shared__ __align__(16) char fp8_gemm_smem[];

    // Batch slice (grid.z): broadcast operands carry a 0 stride, so the
    // same pointer serves every batch.
    using T8 = typename Mainloop::T8;
    const T8* a = reinterpret_cast<const T8*>(p.a_ptr) +
                  (int64_t)blockIdx.z * p.a_batch_stride;
    const T8* b = reinterpret_cast<const T8*>(p.b_ptr) +
                  (int64_t)blockIdx.z * p.b_batch_stride;
    auto* out_bf16 = reinterpret_cast<__nv_bfloat16*>(p.out_ptr) +
                     (int64_t)blockIdx.z * p.out_batch_stride;

    static_assert(Mainloop::kBlockM * Mainloop::kBlockN * 2 <=
                      Mainloop::kARing * Mainloop::kBlockM * Mainloop::kK +
                          Mainloop::kBRing * Mainloop::kBlockN * Mainloop::kK,
                  "output tile must fit the reclaimed operand smem");
    const int2 bn = Fp8GemmTileScheduler<Policy::kGroupRaster>::tile(blockIdx, gridDim);
    Mainloop mainloop(fp8_gemm_smem, a, b, p.m, p.n, p.k, p.a_ld, p.b_ld,
                      threadIdx.x, bn);
    float acc[Mainloop::kNt][Mainloop::kMt][4] = {};  // [nt][mt][acc]
    mainloop.prologue();
    mainloop.accumulate(acc);
    // Drain the pipeline before the epilogue reclaims the operand rings.
    astrai::cp_async_wait_all();
    Epilogue(fp8_gemm_smem, p, bn.x, bn.y, threadIdx.x).run(acc, out_bf16);
}

// ---------------------------------------------------------------------------
// Launchers — pure CUDA (no torch), usable from the binding and pure C tests.
// ---------------------------------------------------------------------------

// SM count of the current device (cached per device; benign init race —
// every writer stores the same value).
inline int device_sm_count() {
    static int cached[64] = {};
    int dev = 0;
    cudaGetDevice(&dev);
    const bool cacheable = dev >= 0 && dev < 64;
    int sms = cacheable ? cached[dev] : 0;
    if (!sms) {
        cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
        sms = sms > 0 ? sms : 1;
        if (cacheable) cached[dev] = sms;
    }
    return sms;
}

// Launch one kernel instantiation with its shared-memory budget: budgets
// beyond the 48KB static limit opt in once per instantiation via
// cudaFuncSetAttribute. Templated on the kernel *value* (auto NTTP) so
// every instantiation owns its own armed flag — same-signature kernels
// must not share it. A failed opt-in arms nothing, so the launch below
// fails loudly through the caller's error checks.
template <auto Kernel, typename... Args>
void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
                      cudaStream_t stream, Args... args) {
    if (smem_bytes > 48 * 1024) {
        static bool armed = false;  // per instantiation
        if (!armed) {
            const cudaError_t err = cudaFuncSetAttribute(
                Kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
                smem_bytes);
            armed = (err == cudaSuccess);
        }
    }
    Kernel<<<grid, block, smem_bytes, stream>>>(args...);
}

// Padding-driven small-CTA rule: m or n <= 64 wastes half a 128-row CTA's
// MMA work, and a non-128-divisible shape drags its edge tiles through the
// predicated generic path — when 64 divides both dims, the 64x64 CTA tiles
// exactly and wins that band.
inline bool small_cta_padding(int64_t m, int64_t n) {
    if (m <= 64 || n <= 64) return true;
    const bool big_div = (m % 128 == 0) && (n % 128 == 0);
    const bool small_div = (m % 64 == 0) && (n % 64 == 0);
    return !big_div && small_div;
}

// Launch configuration — a pure function of the problem (unit-testable
// without a GPU). Raster order is not a plan field: every canonical layout
// runs grouped raster; the plain-raster knob stays available through
// launch_plan's GroupRaster parameter for experiments.
struct Fp8GemmPlan {
    enum class Cta { kSmall64, kNarrow128x64, kBig128 };
    Cta cta;
    bool small_s3;  // kSmall64 only: cp.async pipeline depth (2 vs 3 stages)
};

// crosswise_ops counts the operands taking the direct crosswise load
// (A ColMajor / B RowMajor storage): 0 = dual-congruous NT, 1 = TN and the
// NN swap, 2 = TT. The layout shifts the crossovers (measured tables in
// the design notes): the small CTA hides the crosswise LDG+PRMT latency
// far better, while the big CTA's operand reuse buys back load bandwidth
// the crosswise path does not traffic in.
inline Fp8GemmPlan plan_gemm(const FP8Params& p, int crosswise_ops = 0) {
    const int64_t sm = device_sm_count();
    const int64_t tiles_128 =
        (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 128);
    const auto small = [&](bool s3) {
        return Fp8GemmPlan{Fp8GemmPlan::Cta::kSmall64, s3};
    };
    const auto big = [] {
        return Fp8GemmPlan{Fp8GemmPlan::Cta::kBig128, false};
    };
    const auto narrow = [] {
        return Fp8GemmPlan{Fp8GemmPlan::Cta::kNarrow128x64, false};
    };
    // Padding rules first: predication waste beats any wave-fill effect.
    if (small_cta_padding(p.m, p.n)) return small(crosswise_ops > 0);
    if (crosswise_ops > 0) {
        // Crosswise ladder (L20 measured): the small s3 CTA holds ~3/4 of
        // the big CTA's per-SM throughput but tiles 4x finer, so it owns
        // the whole sub-wave band and past it; the big CTA takes over once
        // its grid fills ~1.5 waves.
        if (tiles_128 >= sm * 3 / 2) return big();
        return small(true);
    }
    if (tiles_128 >= sm) {
        // Wave band: pick by the wave-quantization cost ceil(tiles/sm) *
        // T_tile. The narrow tile carries half the big tile's MMA work at
        // ~94% of its per-SM efficiency (T_narrow ~= 0.53 * T_big,
        // integer-scaled by 100 below) — reproduces every measured
        // crossover.
        const int64_t tiles_narrow =
            (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 63) / 64);
        const auto waves = [sm](int64_t tiles) { return (tiles + sm - 1) / sm; };
        if (waves(tiles_narrow) * 53 < waves(tiles_128) * 100) return narrow();
        return big();
    }
    // Sub-wave band: the narrow CTA fills the wave with N-tiles at full
    // warp depth once its grid passes ~3/8 of a wave; below that the plain
    // 64x64 CTA's extra parallelism wins, and past ~5/8 of a wave of
    // 128x128 tiles the big CTA's operand reuse wins instead.
    if (tiles_128 >= sm * 5 / 8) return big();
    const int64_t tiles_narrow =
        (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 63) / 64);
    if (tiles_narrow >= sm * 3 / 8) return narrow();
    // Full-ring small CTAs: the 24KB s2 variant keeps 4 CTAs/SM while the
    // whole grid stays resident; past that the 32KB s3 variant's deeper
    // pipeline wins on multi-wave grids.
    const int64_t tiles_64 =
        (int64_t)p.batch * ((p.m + 63) / 64) * ((p.n + 63) / 64);
    return small(tiles_64 > sm * 3);
}

// Grid + launch for one concrete Policy — the only place a GEMM kernel goes
// to the wire.
template <typename Policy>
void launch_policy(const FP8Params& p, cudaStream_t stream) {
    using Traits = typename Policy::Traits;
    dim3 grid((p.n + Traits::kBlockN - 1) / Traits::kBlockN,
              (p.m + Traits::kBlockM - 1) / Traits::kBlockM, p.batch);
    launch_with_smem<fp8_gemm_kernel<Policy>>(
        Policy::kSmemBytes, grid, dim3(Traits::kCtaThreads), stream, p);
}

// Plan -> Policy: the production-tuned configs. Big CTA: 128x128 of 8 warps
// x 64x32, kK=64, 2-stage full ring, fast loop only for dual-congruous
// layouts. Narrow: 128x64. Small CTA: 64x64 of 4 warps x 32x32, kK=64,
// kFastLoop always on.
template <FP8Format Fmt, typename LayoutA, typename LayoutB, int GroupRaster>
void launch_plan(const FP8Params& p, const Fp8GemmPlan& plan,
                 cudaStream_t stream) {
    constexpr bool kBigFast = !std::is_same_v<LayoutA, ColMajor> &&
                              !std::is_same_v<LayoutB, RowMajor>;
    switch (plan.cta) {
    case Fp8GemmPlan::Cta::kBig128: {
        using Policy =
            Fp8GemmPolicy<Fmt, 128, 128, LayoutA, LayoutB, 64, 32, 64, 2,
                          GroupRaster, false, kBigFast>;
        launch_policy<Policy>(p, stream);
        break;
    }
    case Fp8GemmPlan::Cta::kNarrow128x64: {
        using Policy =
            Fp8GemmPolicy<Fmt, 128, 64, LayoutA, LayoutB, 32, 32, 64, 2,
                          GroupRaster, false, true>;
        launch_policy<Policy>(p, stream);
        break;
    }
    case Fp8GemmPlan::Cta::kSmall64: {
        if (plan.small_s3) {
            using Policy = Fp8GemmPolicy<Fmt, 64, 64, LayoutA, LayoutB, 32, 32,
                                         64, 3, GroupRaster, false, true>;
            launch_policy<Policy>(p, stream);
        } else {
            using Policy = Fp8GemmPolicy<Fmt, 64, 64, LayoutA, LayoutB, 32, 32,
                                         64, 2, GroupRaster, false, true>;
            launch_policy<Policy>(p, stream);
        }
        break;
    }
    }
}

// Pure problem rewrite: the dual-N-contiguous problem (trans_a/trans_b both
// false) has no dedicated instantiation — it runs as its transpose
// E[N][M] = B^T @ A^T (CUTLASS-sm90's is_swapAB) over swapped operands,
// with p.out_transposed making the epilogue scatter into the caller's
// [M][N] row-major buffer. The rewritten trans flags become the layout tags
// the launcher instantiates; the NN path pays a scalar-store scatter, which
// its rare usage makes the right trade.
inline void canonicalize_gemm(FP8Params& p, bool& trans_a, bool& trans_b) {
    if (!trans_a && !trans_b) {
        FP8Params s = p;  // E = B^T * A^T: swap roles, M <-> N
        s.m = p.n;
        s.n = p.m;
        s.a_ptr = p.b_ptr;
        s.b_ptr = p.a_ptr;
        s.a_ld = p.b_ld;
        s.b_ld = p.a_ld;
        s.a_batch_stride = p.b_batch_stride;
        s.b_batch_stride = p.a_batch_stride;
        s.out_transposed = 1;
        p = s;
        trans_a = trans_b = true;
    }
}

// Entry point: canonicalize the problem, plan the launch, wire the layout
// tags through.
template <FP8Format Fmt>
void gemm(FP8Params p, cudaStream_t stream, bool trans_a, bool trans_b) {
    canonicalize_gemm(p, trans_a, trans_b);
    // Crosswise operand count for the plan: transposed-A storage (ColMajor)
    // and plain-B storage (RowMajor) both take the direct crosswise load.
    const int crosswise = (trans_a ? 1 : 0) + (trans_b ? 0 : 1);
    const Fp8GemmPlan plan = plan_gemm(p, crosswise);
    if (trans_a && trans_b)
        launch_plan<Fmt, ColMajor, ColMajor, 8>(p, plan, stream);
    else if (trans_b)
        launch_plan<Fmt, RowMajor, ColMajor, 8>(p, plan, stream);
    else
        launch_plan<Fmt, ColMajor, RowMajor, 8>(p, plan, stream);
}

}  // namespace fp8
}  // namespace astrai