Skip to content

fix(opus): guard gfx942 bf16ws splitk reduce #3684

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
valarLip merged 1 commit into from
Jun 12, 2026
Merged

fix(opus): guard gfx942 bf16ws splitk reduce #3684

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 47 additions & 13 deletions csrc/opus_gemm/codegen/gen_instances_gfx942.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,13 @@ def _uses_bf16_workspace(k):

def splitk_reduce_extra_forward_decls():
return (
"template<int VEC_, int BLOCK_, typename D_OUT,\n"
" bool HAS_BIAS_, typename D_BIAS_, bool HAS_OOB_>\n"
"__global__ void splitk_reduce_kernel_bf16ws_fallback(\n"
" const opus_splitk_ws_handle* ws_handle, D_OUT* c_out,\n"
" int split_k, int M, int N, int batch,\n"
" int padded_M, int padded_N,\n"
" const D_BIAS_* bias, int stride_bias_batch);\n"
"template<int SPLIT_K, int N_VEC, int ROWS_PER_BLOCK, int VEC_,\n"
" typename D_WS, typename D_OUT, bool HAS_BIAS_, typename D_BIAS_>\n"
"__global__ void splitk_reduce_kernel_exact_n_rowblock(\n"
Expand All @@ -123,6 +130,15 @@ def splitk_reduce_extra_forward_decls():

def splitk_reduce_extra_device_instantiations():
contents = "// Exact-N row-block reduce instantiations (BLOCK=N_VEC*ROWS)\n"
for out_type in ("__bf16", "float"):
contents += (
f"template __global__ void splitk_reduce_kernel_bf16ws_fallback<16, 64, {out_type}, true, {out_type}, true>(\n"
f" const opus_splitk_ws_handle*, {out_type}*, int, int, int, int, int, int,\n"
f" const {out_type}*, int);\n"
f"template __global__ void splitk_reduce_kernel_bf16ws_fallback<16, 64, {out_type}, false, {out_type}, true>(\n"
f" const opus_splitk_ws_handle*, {out_type}*, int, int, int, int, int, int,\n"
f" const {out_type}*, int);\n"
)
for vec, nvec, rows in EXACT_N_ROWBLOCK_REDUCE_CONFIGS:
for sk in SPLITK_REDUCE_SUPPORTED_SPLITKS:
for ws_type in ("float", "__bf16"):
Expand Down Expand Up @@ -240,8 +256,13 @@ def _baseline_call(dtype, hasbias, indent):
if hasbias
else f"\n{indent} nullptr, 0);"
)
reduce_kernel = (
"splitk_reduce_kernel_bf16ws_fallback"
if bf16ws
else "splitk_reduce_kernel_fallback"
)
return (
f"{indent}splitk_reduce_kernel_fallback<REDUCE_VEC, REDUCE_BS, {dtype}, {hb}, {dtype}, true>\n"
f"{indent}{reduce_kernel}<REDUCE_VEC, REDUCE_BS, {dtype}, {hb}, {dtype}, true>\n"
f"{indent} <<<grid_reduce, block_reduce, 0, stream>>>(\n"
f"{indent} ws_handle_,\n"
f"{indent} reinterpret_cast<{dtype}*>(Y.data_ptr()),\n"
Expand All @@ -254,25 +275,37 @@ def _baseline_call(dtype, hasbias, indent):
fp32_t = _baseline_call("float", True, " ")
fp32_f = _baseline_call("float", False, " ")
bf16_y_check = ""
bf16ws_reduce_check = ""
bf16ws_fallback_decl = ""
bf16ws_host_redirect = ""
if bf16ws:
n_conditions = " || ".join(
fp32ws_name = k.name.replace("_bf16ws", "")
exact_reduce_shape_conditions = " ||\n ".join(
f"(N == {vec * nvec} && (M % {rows} == 0))"
for vec, nvec, rows in EXACT_N_ROWBLOCK_REDUCE_CONFIGS
)
sk_conditions = " || ".join(
f"split_k == {sk}" for sk in SPLITK_REDUCE_SUPPORTED_SPLITKS
)
bf16ws_fallback_decl = f"""
#if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__)
template <typename D_C>
void {fp32ws_name}(
aiter_tensor_t &XQ,
aiter_tensor_t &WQ,
aiter_tensor_t &Y,
std::optional<aiter_tensor_t> bias,
int splitK);
#endif
"""
bf16ws_host_redirect = f"""
const bool bf16ws_exact_reduce_shape =
{exact_reduce_shape_conditions};
if (!bf16ws_exact_reduce_shape) {{
{fp32ws_name}<D_C>(XQ, WQ, Y, bias, splitK);
return;
}}
"""
bf16_y_check = (
" AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16,\n"
f' "{err_label} bf16 workspace currently supports only bf16 Y");\n'
)
bf16ws_reduce_check = f"""
AITER_CHECK(padded_N == N && ({n_conditions}),
"{err_label} bf16 workspace requires an exact-N row-block reduce shape");
AITER_CHECK({sk_conditions},
"{err_label} bf16 workspace requires an instantiated split_k");
"""
reduce_launch = f"""
constexpr int REDUCE_VEC = 16;
constexpr int REDUCE_BS = 64;
Expand Down Expand Up @@ -314,6 +347,7 @@ def _baseline_call(dtype, hasbias, indent):
#else
#include "{pipeline_header}"
#endif
{bf16ws_fallback_decl}
{traits_aliases}
#if !defined(__HIP_DEVICE_COMPILE__) && !defined(__HIPCC_RTC__)
template <typename D_C>
Expand All @@ -333,6 +367,7 @@ def _baseline_call(dtype, hasbias, indent):
int N = WQ.size(1);
int K = XQ.size(2);

{bf16ws_host_redirect}
{bf16_y_check}
AITER_CHECK(Y.dtype() == AITER_DTYPE_bf16
|| Y.dtype() == AITER_DTYPE_fp32,
Expand Down Expand Up @@ -418,7 +453,6 @@ def _baseline_call(dtype, hasbias, indent):
int num_tiles_n = (N + {k.B_N} - 1) / {k.B_N};
int padded_M = num_tiles_m * {k.B_M};
int padded_N = num_tiles_n * {k.B_N};
{bf16ws_reduce_check}

auto stream = aiter::getCurrentHIPStream();
hipStreamCaptureStatus capture_status = hipStreamCaptureStatusNone;
Expand Down
67 changes: 57 additions & 10 deletions csrc/opus_gemm/include/gfx942/splitk_reduce_gfx942.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,10 @@
#include "opus_gemm_traits_a16w16.cuh" // opus_splitk_ws_handle
#include <cstdint>

template<int VEC_ = 16, int BLOCK_ = 64, typename D_OUT = __bf16,
bool HAS_BIAS_ = false, typename D_BIAS_ = D_OUT,
bool HAS_OOB_ = true>
__global__ void splitk_reduce_kernel_fallback(
template<int VEC_ = 16, int BLOCK_ = 64, typename D_WS = float,
typename D_OUT = __bf16, bool HAS_BIAS_ = false,
typename D_BIAS_ = D_OUT, bool HAS_OOB_ = true>
__device__ __forceinline__ void splitk_reduce_kernel_fallback_body(
const opus_splitk_ws_handle* __restrict__ ws_handle,
D_OUT* __restrict__ c_out,
int split_k, int M, int N, int batch,
Expand All @@ -26,6 +26,8 @@ __global__ void splitk_reduce_kernel_fallback(
constexpr bool HAS_BIAS = HAS_BIAS_;
constexpr bool HAS_OOB = HAS_OOB_;
using D_BIAS = D_BIAS_;
static_assert(sizeof(D_WS) == 2 || sizeof(D_WS) == 4,
"splitk_reduce fallback supports fp32/bf16 workspace");

constexpr int STEP = 16 / sizeof(D_OUT);
static_assert(STEP * sizeof(D_OUT) == 16,
Expand All @@ -45,7 +47,7 @@ __global__ void splitk_reduce_kernel_fallback(

const int b = bm_id / M;
const int m = bm_id - b * M;
const float* workspace = reinterpret_cast<const float*>(ws_handle->ptr);
const D_WS* workspace = reinterpret_cast<const D_WS*>(ws_handle->ptr);

opus::vector_t<float, VEC> bias_fp32;
if constexpr (HAS_BIAS) {
Expand All @@ -68,19 +70,34 @@ __global__ void splitk_reduce_kernel_fallback(
const long split_stride = (long)batch * padded_M * padded_N;

auto g_ws = opus::make_gmem(workspace,
(unsigned int)(split_stride * split_k * sizeof(float)));
(unsigned int)(split_stride * split_k * sizeof(D_WS)));

opus::vector_t<float, VEC> acc;
#pragma unroll
for (int t = 0; t < VEC; ++t) acc[t] = 0.0f;

for (int s = 0; s < split_k; ++s) {
int ws_idx = ws_row_base + (int)(s * split_stride);
#pragma unroll
for (int g = 0; g < VEC / 4; ++g) {
auto v4 = g_ws.template load<4>(ws_idx + g * 4);
if constexpr (sizeof(D_WS) == 2 && VEC % 8 == 0) {
#pragma unroll
for (int g = 0; g < VEC / 8; ++g) {
auto v8 = g_ws.template load<8>(ws_idx + g * 8);
#pragma unroll
for (int j = 0; j < 8; ++j) {
acc[g * 8 + j] += static_cast<float>(v8[j]);
}
}
} else {
static_assert(VEC % 4 == 0,
"splitk_reduce fallback fp32 workspace needs VEC % 4 == 0");
#pragma unroll
for (int j = 0; j < 4; ++j) acc[g * 4 + j] += v4[j];
for (int g = 0; g < VEC / 4; ++g) {
auto v4 = g_ws.template load<4>(ws_idx + g * 4);
#pragma unroll
for (int j = 0; j < 4; ++j) {
acc[g * 4 + j] += static_cast<float>(v4[j]);
}
}
}
}

Expand Down Expand Up @@ -173,6 +190,36 @@ __global__ void splitk_reduce_kernel_fallback(
#endif // __HIP_DEVICE_COMPILE__
}

template<int VEC_ = 16, int BLOCK_ = 64, typename D_OUT = __bf16,
bool HAS_BIAS_ = false, typename D_BIAS_ = D_OUT,
bool HAS_OOB_ = true>
__global__ void splitk_reduce_kernel_fallback(
const opus_splitk_ws_handle* __restrict__ ws_handle,
D_OUT* __restrict__ c_out,
int split_k, int M, int N, int batch,
int padded_M, int padded_N,
const D_BIAS_* __restrict__ bias,
int bias_stride_batch)
{
splitk_reduce_kernel_fallback_body<VEC_, BLOCK_, float, D_OUT, HAS_BIAS_, D_BIAS_, HAS_OOB_>(
ws_handle, c_out, split_k, M, N, batch, padded_M, padded_N, bias, bias_stride_batch);
}

template<int VEC_ = 16, int BLOCK_ = 64, typename D_OUT = __bf16,
bool HAS_BIAS_ = false, typename D_BIAS_ = D_OUT,
bool HAS_OOB_ = true>
__global__ void splitk_reduce_kernel_bf16ws_fallback(
const opus_splitk_ws_handle* __restrict__ ws_handle,
D_OUT* __restrict__ c_out,
int split_k, int M, int N, int batch,
int padded_M, int padded_N,
const D_BIAS_* __restrict__ bias,
int bias_stride_batch)
{
splitk_reduce_kernel_fallback_body<VEC_, BLOCK_, __bf16, D_OUT, HAS_BIAS_, D_BIAS_, HAS_OOB_>(
ws_handle, c_out, split_k, M, N, batch, padded_M, padded_N, bias, bias_stride_batch);
}

// Exact-N row-block fast path: static split_k unroll, no N tail/OOB.
template<int SPLIT_K, int N_VEC, int ROWS_PER_BLOCK, int VEC_ = 8,
typename D_WS = float, typename D_OUT = __bf16,
Expand Down
Loading