perf: tune bf16 gemv and add opt-in fused swiglu
- deepen common-shape BF16 GEMV tuning with warp-row tiling for LLaMA/Qwen2/GPT-NeoX/OPT decode projections - add fused BF16 up/gate SwiGLU CUDA primitive with ASTRAI_SWIGLU=0/1/auto dispatch - keep the unfused linear backend as the default path; auto enables no shape until per-architecture checkpoint gates pass - fall back to the linear/torch chain when kernels are absent, on CPU, in training, or outside supported M/K/dtype shapes - add gemv/swiglu benchmark scripts, dispatch and parity tests, and kernel documentation Benchmark: NVIDIA L20 (sm_89), CUDA 12.8, PyTorch 2.11.0+cu128, idle GPU. AstrAI 1B config (24 layers, hidden 1536, vocab 100000), BF16, prompt 128, 32 greedy decode tokens, CUDA graphs enabled, A/B in separate interleaved processes (3 rounds, 8 trials each, medians). Default vs ASTRAI_SWIGLU=1 per generate call: batch 1 134.8->129.1 ms (+4.44%), batch 2 136.2->130.9 ms (+4.06%), batch 4 145.5->140.3 ms (+3.66%). Greedy output identical at batch 1, differs at batch 2/4, so auto stays unfused by default; kernelless fallback verified bit-identical greedy.
This commit is contained in:
0z5a
@@ -12,7 +12,9 @@
|
||||
namespace {
|
||||
|
||||
constexpr int kThreads = 256;
|
||||
constexpr int kHalfCtaThreads = 128;
|
||||
constexpr int kWarpSize = 32;
|
||||
constexpr int kWarpTiledThreads = 128;
|
||||
|
||||
__device__ __forceinline__ float warp_sum(float value) {
|
||||
#pragma unroll
|
||||
@@ -22,7 +24,7 @@ __device__ __forceinline__ float warp_sum(float value) {
|
||||
return value;
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
template <int Rows, int Threads>
|
||||
__global__ void bf16_gemv_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ weight,
|
||||
@@ -36,7 +38,7 @@ __global__ void bf16_gemv_kernel(
|
||||
const int warp = threadIdx.x / kWarpSize;
|
||||
|
||||
float sums[Rows] = {};
|
||||
__shared__ float warp_sums[Rows][kThreads / kWarpSize];
|
||||
__shared__ float warp_sums[Rows][Threads / kWarpSize];
|
||||
// Weight row: scalar head/tail around a 16-byte-aligned uint4 middle so
|
||||
// any K is accepted while keeping 128-bit weight loads, which dominate
|
||||
// bandwidth on decode shapes. x pairs with scalar loads: it is a tiny
|
||||
@@ -129,7 +131,6 @@ __global__ void bf16_gemv_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
sums[row] = warp_sum(sums[row]);
|
||||
@@ -146,7 +147,7 @@ __global__ void bf16_gemv_kernel(
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
float sum =
|
||||
lane < (kThreads / kWarpSize) ? warp_sums[row][lane] : 0.0f;
|
||||
lane < (Threads / kWarpSize) ? warp_sums[row][lane] : 0.0f;
|
||||
sum = warp_sum(sum);
|
||||
if (lane == 0) {
|
||||
if (bias != nullptr) {
|
||||
@@ -158,6 +159,120 @@ __global__ void bf16_gemv_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
__global__ void bf16_gemv_aligned_warp_tiled_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ weight,
|
||||
const __nv_bfloat16* __restrict__ bias,
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
int n,
|
||||
int k
|
||||
) {
|
||||
constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize;
|
||||
const int lane = threadIdx.x & (kWarpSize - 1);
|
||||
const int warp = threadIdx.x / kWarpSize;
|
||||
const int output_index = blockIdx.x * kWarpsPerBlock + warp;
|
||||
if (output_index >= n) {
|
||||
return;
|
||||
}
|
||||
|
||||
// The launcher selects this path only when each row is 16-byte aligned.
|
||||
// Four independent output rows per CTA remove the block-wide reduction
|
||||
// barrier and improve occupancy for the medium LLaMA projection bands.
|
||||
const int vectors = k / 8;
|
||||
const auto* x4 = reinterpret_cast<const uint4*>(x);
|
||||
const auto* w4 = reinterpret_cast<const uint4*>(weight) +
|
||||
static_cast<int64_t>(output_index) * vectors;
|
||||
float sums[Rows] = {};
|
||||
for (int vector = lane; vector < vectors; vector += kWarpSize) {
|
||||
const uint4 wv_raw = w4[vector];
|
||||
const auto* wv = reinterpret_cast<const __nv_bfloat162*>(&wv_raw);
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
const uint4 xv_raw =
|
||||
x4[static_cast<int64_t>(row) * vectors + vector];
|
||||
const auto* xv = reinterpret_cast<const __nv_bfloat162*>(&xv_raw);
|
||||
#pragma unroll
|
||||
for (int pair = 0; pair < 4; ++pair) {
|
||||
sums[row] = fmaf(
|
||||
__bfloat162float(__low2bfloat16(xv[pair])),
|
||||
__bfloat162float(__low2bfloat16(wv[pair])),
|
||||
sums[row]
|
||||
);
|
||||
sums[row] = fmaf(
|
||||
__bfloat162float(__high2bfloat16(xv[pair])),
|
||||
__bfloat162float(__high2bfloat16(wv[pair])),
|
||||
sums[row]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
sums[row] = warp_sum(sums[row]);
|
||||
if (lane == 0) {
|
||||
if (bias != nullptr) {
|
||||
sums[row] += __bfloat162float(bias[output_index]);
|
||||
}
|
||||
output[row * n + output_index] = __float2bfloat16_rn(sums[row]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
constexpr bool use_warp_tiled_kernel(int n, int k) {
|
||||
// These bands are intentionally narrow and are validated by the common
|
||||
// transformer benchmark. The 256-thread cooperative kernel remains the
|
||||
// fallback for arbitrary K, larger projections, and M=2 (where the
|
||||
// single-warp reduction regresses the current vectorized kernel).
|
||||
if constexpr (Rows == 4) {
|
||||
return (n == 1024 && k == 4096) ||
|
||||
(n == 4096 && k == 4096) ||
|
||||
(n == 11008 && k == 4096) ||
|
||||
(n == 4096 && k == 11008);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
constexpr bool use_half_cta_kernel(int n, int k) {
|
||||
// A 128-thread CTA reduces synchronization and scheduling overhead for
|
||||
// selected medium decode projections. Keep the selector exact: long-K
|
||||
// and bandwidth-saturated shapes regress, and the winning bands differ
|
||||
// materially with the number of reused input rows.
|
||||
if constexpr (Rows == 1) {
|
||||
return n == 8192 && k == 2048;
|
||||
}
|
||||
if constexpr (Rows == 2) {
|
||||
return (n == 4096 && k == 4096) ||
|
||||
(n == 11008 && k == 4096) ||
|
||||
(n == 3584 && k == 3584) ||
|
||||
(n == 2048 && k == 2048) ||
|
||||
(n == 8192 && k == 2048);
|
||||
}
|
||||
if constexpr (Rows == 4) {
|
||||
return (n == 5120 && k == 5120) ||
|
||||
(n == 3584 && k == 3584) ||
|
||||
(n == 2048 && k == 2048) ||
|
||||
(n == 8192 && k == 2048);
|
||||
}
|
||||
if constexpr (Rows == 8) {
|
||||
return (n == 4096 && k == 4096) ||
|
||||
(n == 11008 && k == 4096) ||
|
||||
(n == 4096 && k == 11008) ||
|
||||
(n == 1024 && k == 4096) ||
|
||||
(n == 5120 && k == 5120) ||
|
||||
(n == 512 && k == 3584) ||
|
||||
(n == 3584 && k == 3584) ||
|
||||
(n == 1024 && k == 8192) ||
|
||||
(n == 2048 && k == 2048) ||
|
||||
(n == 8192 && k == 2048) ||
|
||||
(n == 2048 && k == 8192);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
void launch_bf16_gemv(
|
||||
const __nv_bfloat16* x,
|
||||
@@ -168,7 +283,28 @@ void launch_bf16_gemv(
|
||||
int k,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
bf16_gemv_kernel<Rows><<<n, kThreads, 0, stream>>>(
|
||||
const bool aligned_rows = k % 8 == 0 &&
|
||||
(reinterpret_cast<uintptr_t>(x) & 15u) == 0u &&
|
||||
(reinterpret_cast<uintptr_t>(weight) & 15u) == 0u;
|
||||
if constexpr (Rows == 4) {
|
||||
if (aligned_rows && use_warp_tiled_kernel<Rows>(n, k)) {
|
||||
constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize;
|
||||
const int blocks = (n + kWarpsPerBlock - 1) / kWarpsPerBlock;
|
||||
bf16_gemv_aligned_warp_tiled_kernel<Rows>
|
||||
<<<blocks, kWarpTiledThreads, 0, stream>>>(
|
||||
x, weight, bias, output, n, k
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
if (aligned_rows && use_half_cta_kernel<Rows>(n, k)) {
|
||||
bf16_gemv_kernel<Rows, kHalfCtaThreads>
|
||||
<<<n, kHalfCtaThreads, 0, stream>>>(
|
||||
x, weight, bias, output, n, k
|
||||
);
|
||||
return;
|
||||
}
|
||||
bf16_gemv_kernel<Rows, kThreads><<<n, kThreads, 0, stream>>>(
|
||||
x, weight, bias, output, n, k
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
// Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers.
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kThreads = 256;
|
||||
constexpr int kWarpSize = 32;
|
||||
constexpr int kWarps = kThreads / kWarpSize;
|
||||
|
||||
__device__ __forceinline__ float warp_sum(float value) {
|
||||
#pragma unroll
|
||||
for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
|
||||
value += __shfl_down_sync(0xffffffff, value, offset);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float round_bf16(float value) {
|
||||
return __bfloat162float(__float2bfloat16_rn(value));
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
__global__ void bf16_swiglu_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ up_weight,
|
||||
const __nv_bfloat16* __restrict__ gate_weight,
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
int n,
|
||||
int k
|
||||
) {
|
||||
const int output_index = blockIdx.x;
|
||||
const int lane = threadIdx.x & (kWarpSize - 1);
|
||||
const int warp = threadIdx.x / kWarpSize;
|
||||
const int vector_count = k / 8;
|
||||
|
||||
float up_sums[Rows] = {};
|
||||
float gate_sums[Rows] = {};
|
||||
__shared__ float up_warp_sums[Rows][kWarps];
|
||||
__shared__ float gate_warp_sums[Rows][kWarps];
|
||||
|
||||
const auto* x4 = reinterpret_cast<const uint4*>(x);
|
||||
const auto* up4 = reinterpret_cast<const uint4*>(
|
||||
up_weight + static_cast<int64_t>(output_index) * k
|
||||
);
|
||||
const auto* gate4 = reinterpret_cast<const uint4*>(
|
||||
gate_weight + static_cast<int64_t>(output_index) * k
|
||||
);
|
||||
|
||||
// Read each pair of up/gate weight chunks once per CTA, then reuse it for
|
||||
// every active decode row. The fused epilogue removes two [M, N]
|
||||
// intermediates and the standalone SiLU and multiply launches.
|
||||
for (int vector_index = threadIdx.x;
|
||||
vector_index < vector_count;
|
||||
vector_index += blockDim.x) {
|
||||
const uint4 up_raw = up4[vector_index];
|
||||
const uint4 gate_raw = gate4[vector_index];
|
||||
const auto* up_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&up_raw);
|
||||
const auto* gate_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&gate_raw);
|
||||
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
const uint4 x_raw =
|
||||
x4[static_cast<int64_t>(row) * vector_count + vector_index];
|
||||
const auto* x_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&x_raw);
|
||||
#pragma unroll
|
||||
for (int pair = 0; pair < 4; ++pair) {
|
||||
const float2 xv = __bfloat1622float2(x_values[pair]);
|
||||
const float2 uv = __bfloat1622float2(up_values[pair]);
|
||||
const float2 gv = __bfloat1622float2(gate_values[pair]);
|
||||
up_sums[row] = fmaf(xv.x, uv.x, up_sums[row]);
|
||||
up_sums[row] = fmaf(xv.y, uv.y, up_sums[row]);
|
||||
gate_sums[row] = fmaf(xv.x, gv.x, gate_sums[row]);
|
||||
gate_sums[row] = fmaf(xv.y, gv.y, gate_sums[row]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
up_sums[row] = warp_sum(up_sums[row]);
|
||||
gate_sums[row] = warp_sum(gate_sums[row]);
|
||||
}
|
||||
if (lane == 0) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
up_warp_sums[row][warp] = up_sums[row];
|
||||
gate_warp_sums[row][warp] = gate_sums[row];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (warp == 0) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
float up = lane < kWarps ? up_warp_sums[row][lane] : 0.0f;
|
||||
float gate = lane < kWarps ? gate_warp_sums[row][lane] : 0.0f;
|
||||
up = warp_sum(up);
|
||||
gate = warp_sum(gate);
|
||||
if (lane == 0) {
|
||||
// Match the public composition's BF16 rounding boundaries:
|
||||
// BF16 linear outputs, BF16 SiLU output, then BF16 multiply.
|
||||
up = round_bf16(up);
|
||||
gate = round_bf16(gate);
|
||||
const float silu = round_bf16(gate / (1.0f + expf(-gate)));
|
||||
output[static_cast<int64_t>(row) * n + output_index] =
|
||||
__float2bfloat16_rn(up * silu);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
void launch_bf16_swiglu(
|
||||
const __nv_bfloat16* x,
|
||||
const __nv_bfloat16* up_weight,
|
||||
const __nv_bfloat16* gate_weight,
|
||||
__nv_bfloat16* output,
|
||||
int n,
|
||||
int k,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
bf16_swiglu_kernel<Rows><<<n, kThreads, 0, stream>>>(
|
||||
x, up_weight, gate_weight, output, n, k
|
||||
);
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
__global__ void bf16_swiglu_warp_rows_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ up_weight,
|
||||
const __nv_bfloat16* __restrict__ gate_weight,
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
int n,
|
||||
int k
|
||||
) {
|
||||
const int output_index = blockIdx.x;
|
||||
const int row = threadIdx.x / kWarpSize;
|
||||
const int lane = threadIdx.x & (kWarpSize - 1);
|
||||
const int vector_count = k / 8;
|
||||
|
||||
float up_sum = 0.0f;
|
||||
float gate_sum = 0.0f;
|
||||
const auto* x4 = reinterpret_cast<const uint4*>(
|
||||
x + static_cast<int64_t>(row) * k
|
||||
);
|
||||
const auto* up4 = reinterpret_cast<const uint4*>(
|
||||
up_weight + static_cast<int64_t>(output_index) * k
|
||||
);
|
||||
const auto* gate4 = reinterpret_cast<const uint4*>(
|
||||
gate_weight + static_cast<int64_t>(output_index) * k
|
||||
);
|
||||
|
||||
// A warp owns one decode row. Same-address weight reads from sibling
|
||||
// warps are served through the read-only/L1 path, while each row avoids
|
||||
// CTA-wide shared-memory reductions and synchronization.
|
||||
for (int vector_index = lane;
|
||||
vector_index < vector_count;
|
||||
vector_index += kWarpSize) {
|
||||
const uint4 x_raw = x4[vector_index];
|
||||
const uint4 up_raw = up4[vector_index];
|
||||
const uint4 gate_raw = gate4[vector_index];
|
||||
const auto* x_values = reinterpret_cast<const __nv_bfloat162*>(&x_raw);
|
||||
const auto* up_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&up_raw);
|
||||
const auto* gate_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&gate_raw);
|
||||
#pragma unroll
|
||||
for (int pair = 0; pair < 4; ++pair) {
|
||||
const float2 xv = __bfloat1622float2(x_values[pair]);
|
||||
const float2 uv = __bfloat1622float2(up_values[pair]);
|
||||
const float2 gv = __bfloat1622float2(gate_values[pair]);
|
||||
up_sum = fmaf(xv.x, uv.x, up_sum);
|
||||
up_sum = fmaf(xv.y, uv.y, up_sum);
|
||||
gate_sum = fmaf(xv.x, gv.x, gate_sum);
|
||||
gate_sum = fmaf(xv.y, gv.y, gate_sum);
|
||||
}
|
||||
}
|
||||
up_sum = warp_sum(up_sum);
|
||||
gate_sum = warp_sum(gate_sum);
|
||||
if (lane == 0) {
|
||||
up_sum = round_bf16(up_sum);
|
||||
gate_sum = round_bf16(gate_sum);
|
||||
const float silu =
|
||||
round_bf16(gate_sum / (1.0f + expf(-gate_sum)));
|
||||
output[static_cast<int64_t>(row) * n + output_index] =
|
||||
__float2bfloat16_rn(up_sum * silu);
|
||||
}
|
||||
}
|
||||
|
||||
template <int Rows>
|
||||
void launch_bf16_swiglu_warp_rows(
|
||||
const __nv_bfloat16* x,
|
||||
const __nv_bfloat16* up_weight,
|
||||
const __nv_bfloat16* gate_weight,
|
||||
__nv_bfloat16* output,
|
||||
int n,
|
||||
int k,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
bf16_swiglu_warp_rows_kernel<Rows><<<n, Rows * kWarpSize, 0, stream>>>(
|
||||
x, up_weight, gate_weight, output, n, k
|
||||
);
|
||||
}
|
||||
|
||||
torch::Tensor bf16_swiglu(
|
||||
torch::Tensor x,
|
||||
torch::Tensor up_weight,
|
||||
torch::Tensor gate_weight
|
||||
) {
|
||||
TORCH_CHECK(
|
||||
x.is_cuda() && up_weight.is_cuda() && gate_weight.is_cuda(),
|
||||
"x, up_weight, and gate_weight must be CUDA tensors"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.device() == up_weight.device() && x.device() == gate_weight.device(),
|
||||
"x and weights must share a device"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.scalar_type() == torch::kBFloat16 &&
|
||||
up_weight.scalar_type() == torch::kBFloat16 &&
|
||||
gate_weight.scalar_type() == torch::kBFloat16,
|
||||
"x and weights must be bf16"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.dim() == 1 || x.dim() == 2,
|
||||
"x must have shape [K] or [M, K]"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
up_weight.dim() == 2 && gate_weight.dim() == 2,
|
||||
"weights must have shape [N, K]"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.is_contiguous() && up_weight.is_contiguous() &&
|
||||
gate_weight.is_contiguous(),
|
||||
"x and weights must be contiguous"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
!x.requires_grad() && !up_weight.requires_grad() &&
|
||||
!gate_weight.requires_grad(),
|
||||
"bf16_swiglu is inference-only and does not support autograd"
|
||||
);
|
||||
|
||||
const int64_t m = x.dim() == 1 ? 1 : x.size(0);
|
||||
const int64_t k = x.size(-1);
|
||||
const int64_t n = up_weight.size(0);
|
||||
TORCH_CHECK(m >= 1 && m <= 8, "M must be in [1, 8]");
|
||||
TORCH_CHECK(
|
||||
gate_weight.sizes() == up_weight.sizes(),
|
||||
"up_weight and gate_weight must have identical shapes"
|
||||
);
|
||||
TORCH_CHECK(up_weight.size(1) == k, "weight K must match x K");
|
||||
TORCH_CHECK(k > 0 && n > 0, "N and K must be positive");
|
||||
TORCH_CHECK(k % 8 == 0, "K must be divisible by 8");
|
||||
TORCH_CHECK(
|
||||
k <= std::numeric_limits<int>::max() &&
|
||||
n <= std::numeric_limits<int>::max(),
|
||||
"N or K exceeds the CUDA launcher limit"
|
||||
);
|
||||
|
||||
const at::cuda::OptionalCUDAGuard guard(x.device());
|
||||
const auto* properties = at::cuda::getDeviceProperties(x.device().index());
|
||||
TORCH_CHECK(
|
||||
properties->major >= 8,
|
||||
"bf16_swiglu requires compute capability 8.0+"
|
||||
);
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
auto output = x.dim() == 1 ? torch::empty({n}, x.options())
|
||||
: torch::empty({m, n}, x.options());
|
||||
|
||||
const auto* x_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
|
||||
const auto* up_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(up_weight.data_ptr());
|
||||
const auto* gate_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(gate_weight.data_ptr());
|
||||
auto* output_ptr =
|
||||
reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
|
||||
const int n_int = static_cast<int>(n);
|
||||
const int k_int = static_cast<int>(k);
|
||||
const bool use_warp_rows = n_int == 6912 && k_int == 1536;
|
||||
|
||||
switch (m) {
|
||||
case 1:
|
||||
launch_bf16_swiglu<1>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 2:
|
||||
if (use_warp_rows) {
|
||||
launch_bf16_swiglu_warp_rows<2>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
} else {
|
||||
launch_bf16_swiglu<2>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
}
|
||||
break;
|
||||
case 3:
|
||||
launch_bf16_swiglu<3>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 4:
|
||||
if (use_warp_rows) {
|
||||
launch_bf16_swiglu_warp_rows<4>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
} else {
|
||||
launch_bf16_swiglu<4>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
}
|
||||
break;
|
||||
case 5:
|
||||
launch_bf16_swiglu<5>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 6:
|
||||
launch_bf16_swiglu<6>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 7:
|
||||
launch_bf16_swiglu<7>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 8:
|
||||
if (use_warp_rows) {
|
||||
launch_bf16_swiglu_warp_rows<8>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
} else {
|
||||
launch_bf16_swiglu<8>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
}
|
||||
break;
|
||||
}
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
|
||||
module.def(
|
||||
"bf16_swiglu",
|
||||
&bf16_swiglu,
|
||||
py::arg("x"),
|
||||
py::arg("up_weight"),
|
||||
py::arg("gate_weight"),
|
||||
"M in [1, 8] fused BF16 up/gate projection and SwiGLU"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user