Skip to content

fix(kernels): rename misnamed gate_silu_mul* to gate_geglu* #181

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

Open
hebangwen wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(kernels): rename misnamed gate_silu_mul* to gate_geglu* #181

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
4 changes: 2 additions & 2 deletions cpp/models/pi05/src/targets/sm110/operation_driver.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -358,7 +358,7 @@ modalities::Status Sm110OperationDriver::gate_gelu_fp16(
if (!merged || !output || rows <= 0 || hidden <= 0) {
return invalid("SM110 FP16 gated GELU arguments are invalid");
}
::gate_silu_mul_merged_fp16(
::gate_geglu_merged_fp16(
static_cast<const __half*>(merged), static_cast<__half*>(output),
rows, hidden, reinterpret_cast<cudaStream_t>(stream));
return launch_status();
Expand All @@ -370,7 +370,7 @@ modalities::Status Sm110OperationDriver::gate_gelu_fp8(
if (!merged || !output || !scale || rows <= 0 || hidden <= 0) {
return invalid("SM110 FP8 gated GELU arguments are invalid");
}
::gate_silu_mul_merged_fp8_fp16(
::gate_geglu_merged_fp8_fp16(
static_cast<const __half*>(merged),
static_cast<__nv_fp8_e4m3*>(output), rows, hidden, scale,
reinterpret_cast<cudaStream_t>(stream));
Expand Down
6 changes: 3 additions & 3 deletions cpp/models/pi05/src/targets/sm120/target.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -787,7 +787,7 @@ modalities::Status frontend_gated_activation(
if (!scale) return invalid("SM120 gated-activation scale is invalid");
*linear_input = binding->fp8_linear->scratch_data();
*prequantized = true;
::gate_silu_mul_merged_fp8(
::gate_geglu_merged_fp8(
static_cast<const __nv_bfloat16*>(gate),
static_cast<__nv_fp8_e4m3*>(
binding->fp8_linear->scratch_data()),
Expand All @@ -797,12 +797,12 @@ modalities::Status frontend_gated_activation(
*linear_input = output;
*prequantized = false;
if (merged) {
::gate_silu_mul_merged(
::gate_geglu_merged(
static_cast<const __nv_bfloat16*>(gate),
static_cast<__nv_bfloat16*>(output), rows, hidden_width,
cuda_stream);
} else {
::gate_silu_mul(
::gate_geglu(
static_cast<const __nv_bfloat16*>(gate),
static_cast<const __nv_bfloat16*>(up),
static_cast<__nv_bfloat16*>(output), rows * hidden_width,
Expand Down
12 changes: 6 additions & 6 deletions csrc/bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -960,12 +960,12 @@ PYBIND11_MODULE(flash_rt_kernels, m) {

// Activation — GEGLU (tanh-approx GELU(gate) * up), not SiLU.
m.def("gate_geglu", [](uintptr_t gate, uintptr_t up, uintptr_t out, int n, uintptr_t stream) {
gate_silu_mul(typed_ptr<__nv_bfloat16>(gate), typed_ptr<__nv_bfloat16>(up),
gate_geglu(typed_ptr<__nv_bfloat16>(gate), typed_ptr<__nv_bfloat16>(up),
typed_ptr<__nv_bfloat16>(out), n, to_stream(stream));
}, py::arg("gate"), py::arg("up"), py::arg("out"), py::arg("n"), py::arg("stream") = 0);

m.def("gate_geglu_fp16", [](uintptr_t gate, uintptr_t up, uintptr_t out, int n, uintptr_t stream) {
gate_silu_mul_fp16(typed_ptr<__half>(gate), typed_ptr<__half>(up),
gate_geglu_fp16(typed_ptr<__half>(gate), typed_ptr<__half>(up),
typed_ptr<__half>(out), n, to_stream(stream));
}, py::arg("gate"), py::arg("up"), py::arg("out"), py::arg("n"), py::arg("stream") = 0);

Expand Down Expand Up @@ -1030,14 +1030,14 @@ PYBIND11_MODULE(flash_rt_kernels, m) {

m.def("gate_geglu_merged", [](uintptr_t merged, uintptr_t out,
int seq, int half_dim, uintptr_t stream) {
gate_silu_mul_merged(typed_ptr<__nv_bfloat16>(merged),
gate_geglu_merged(typed_ptr<__nv_bfloat16>(merged),
typed_ptr<__nv_bfloat16>(out), seq, half_dim, to_stream(stream));
}, py::arg("merged"), py::arg("out"), py::arg("seq"), py::arg("half_dim"), py::arg("stream") = 0);

m.def("gate_geglu_merged_fp8", [](uintptr_t merged, uintptr_t out,
int seq, int half_dim,
uintptr_t d_scale, uintptr_t stream) {
gate_silu_mul_merged_fp8(typed_ptr<__nv_bfloat16>(merged),
gate_geglu_merged_fp8(typed_ptr<__nv_bfloat16>(merged),
typed_ptr<__nv_fp8_e4m3>(out), seq, half_dim,
reinterpret_cast<const float*>(d_scale), to_stream(stream));
}, py::arg("merged"), py::arg("out"), py::arg("seq"), py::arg("half_dim"),
Expand Down Expand Up @@ -1838,7 +1838,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) {

m.def("gate_geglu_merged_fp16", [](uintptr_t merged, uintptr_t out,
int seq, int half_dim, uintptr_t stream) {
gate_silu_mul_merged_fp16(reinterpret_cast<const __half*>(merged),
gate_geglu_merged_fp16(reinterpret_cast<const __half*>(merged),
reinterpret_cast<__half*>(out), seq, half_dim, to_stream(stream));
}, py::arg("merged"), py::arg("out"), py::arg("seq"), py::arg("half_dim"), py::arg("stream") = 0);

Expand All @@ -1853,7 +1853,7 @@ PYBIND11_MODULE(flash_rt_kernels, m) {
m.def("gate_geglu_merged_fp8_fp16", [](uintptr_t merged, uintptr_t out,
int seq, int half_dim,
uintptr_t d_scale, uintptr_t stream) {
gate_silu_mul_merged_fp8_fp16(reinterpret_cast<const __half*>(merged),
gate_geglu_merged_fp8_fp16(reinterpret_cast<const __half*>(merged),
typed_ptr<__nv_fp8_e4m3>(out), seq, half_dim,
reinterpret_cast<const float*>(d_scale), to_stream(stream));
}, py::arg("merged"), py::arg("out"), py::arg("seq"), py::arg("half_dim"),
Expand Down
48 changes: 25 additions & 23 deletions csrc/kernels/activation.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@

// ── Gate GELU Multiply ──
// GELU(x) approx: x * sigmoid(1.5957691216 * x * (1 + 0.044715 * x^2))
// NOTE: previously misnamed `gate_silu_mul_kernel`; this is the tanh/sigmoid
// approx GELU, not SiLU. Renamed to reflect the actual activation.
template<typename T>
__global__ void gate_silu_mul_kernel(const T* __restrict__ gate,
__global__ void gate_geglu_kernel(const T* __restrict__ gate,
const T* __restrict__ up,
T* __restrict__ out, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
Expand All @@ -22,16 +24,16 @@ __global__ void gate_silu_mul_kernel(const T* __restrict__ gate,
}
}

template __global__ void gate_silu_mul_kernel<__half>(const __half*, const __half*, __half*, int);
template __global__ void gate_silu_mul_kernel<__nv_bfloat16>(const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, int);
template __global__ void gate_geglu_kernel<__half>(const __half*, const __half*, __half*, int);
template __global__ void gate_geglu_kernel<__nv_bfloat16>(const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, int);

void gate_silu_mul(const __nv_bfloat16* gate, const __nv_bfloat16* up,
void gate_geglu(const __nv_bfloat16* gate, const __nv_bfloat16* up,
__nv_bfloat16* out, int n, cudaStream_t stream) {
gate_silu_mul_kernel<__nv_bfloat16><<<(n + 255) / 256, 256, 0, stream>>>(gate, up, out, n);
gate_geglu_kernel<__nv_bfloat16><<<(n + 255) / 256, 256, 0, stream>>>(gate, up, out, n);
}
void gate_silu_mul_fp16(const __half* gate, const __half* up,
void gate_geglu_fp16(const __half* gate, const __half* up,
__half* out, int n, cudaStream_t stream) {
gate_silu_mul_kernel<__half><<<(n + 255) / 256, 256, 0, stream>>>(gate, up, out, n);
gate_geglu_kernel<__half><<<(n + 255) / 256, 256, 0, stream>>>(gate, up, out, n);
}

// ── GELU in-place ──
Expand Down Expand Up @@ -139,7 +141,7 @@ void bias_gelu_inplace_bf16_strict(__nv_bfloat16* x,
// ── Gate GELU Mul Merged ──
// Input: (seq, 2*half_dim), gate = [:, :half_dim], up = [:, half_dim:]
template<typename T>
__global__ void gate_silu_mul_merged_kernel(const T* __restrict__ merged,
__global__ void gate_geglu_merged_kernel(const T* __restrict__ merged,
T* __restrict__ out,
int seq, int half_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
Expand All @@ -155,24 +157,24 @@ __global__ void gate_silu_mul_merged_kernel(const T* __restrict__ merged,
}
}

template __global__ void gate_silu_mul_merged_kernel<__half>(const __half*, __half*, int, int);
template __global__ void gate_silu_mul_merged_kernel<__nv_bfloat16>(const __nv_bfloat16*, __nv_bfloat16*, int, int);
template __global__ void gate_geglu_merged_kernel<__half>(const __half*, __half*, int, int);
template __global__ void gate_geglu_merged_kernel<__nv_bfloat16>(const __nv_bfloat16*, __nv_bfloat16*, int, int);

void gate_silu_mul_merged(const __nv_bfloat16* merged, __nv_bfloat16* out,
void gate_geglu_merged(const __nv_bfloat16* merged, __nv_bfloat16* out,
int seq, int half_dim, cudaStream_t stream) {
int total = seq * half_dim;
int blocks = (total + 255) / 256;
gate_silu_mul_merged_kernel<__nv_bfloat16><<<blocks, 256, 0, stream>>>(merged, out, seq, half_dim);
gate_geglu_merged_kernel<__nv_bfloat16><<<blocks, 256, 0, stream>>>(merged, out, seq, half_dim);
}
void gate_silu_mul_merged_fp16(const __half* merged, __half* out,
void gate_geglu_merged_fp16(const __half* merged, __half* out,
int seq, int half_dim, cudaStream_t stream) {
int total = seq * half_dim;
int blocks = (total + 255) / 256;
gate_silu_mul_merged_kernel<__half><<<blocks, 256, 0, stream>>>(merged, out, seq, half_dim);
gate_geglu_merged_kernel<__half><<<blocks, 256, 0, stream>>>(merged, out, seq, half_dim);
}

// Vectorized 8-half / thread element-wise multiply. BW-bound; pairs
// with two split-G7 GEMMs in R3.1 to replace gate_silu_mul_merged_fp16.
// with two split-G7 GEMMs in R3.1 to replace gate_geglu_merged_fp16.
__global__ void mul_fp16_kernel(const __half* __restrict__ a,
const __half* __restrict__ b,
__half* __restrict__ out, int n) {
Expand Down Expand Up @@ -216,7 +218,7 @@ void mul_fp16(const __half* a, const __half* b, __half* out, int n, cudaStream_t
// ── Gate GELU Mul Merged -> FP8 ──
// 4 elem/thread vectorized, matching production silu_mul_split_fp8_k throughput.
// Merged layout: merged[s, 0..H-1] = gate, merged[s, H..2H-1] = up
__global__ void gate_silu_mul_merged_fp8_kernel_fp16(const __half* merged, __nv_fp8_e4m3* out, int S, int H,
__global__ void gate_geglu_merged_fp8_kernel_fp16(const __half* merged, __nv_fp8_e4m3* out, int S, int H,
const float* descale_ptr) {
int i = (blockIdx.x * blockDim.x + threadIdx.x) * 4; // 4 elements per thread
if (i >= S * H) return;
Expand Down Expand Up @@ -246,7 +248,7 @@ __global__ void gate_silu_mul_merged_fp8_kernel_fp16(const __half* merged, __nv_

// BF16 generic version (non-encoder paths)
template<typename T>
__global__ void gate_silu_mul_merged_fp8_kernel(const T* __restrict__ merged,
__global__ void gate_geglu_merged_fp8_kernel(const T* __restrict__ merged,
__nv_fp8_e4m3* __restrict__ out,
int seq, int half_dim,
const float* __restrict__ d_scale) {
Expand All @@ -266,24 +268,24 @@ __global__ void gate_silu_mul_merged_fp8_kernel(const T* __restrict__ merged,
}
}

template __global__ void gate_silu_mul_merged_fp8_kernel<__half>(const __half*, __nv_fp8_e4m3*, int, int, const float*);
template __global__ void gate_silu_mul_merged_fp8_kernel<__nv_bfloat16>(const __nv_bfloat16*, __nv_fp8_e4m3*, int, int, const float*);
template __global__ void gate_geglu_merged_fp8_kernel<__half>(const __half*, __nv_fp8_e4m3*, int, int, const float*);
template __global__ void gate_geglu_merged_fp8_kernel<__nv_bfloat16>(const __nv_bfloat16*, __nv_fp8_e4m3*, int, int, const float*);

void gate_silu_mul_merged_fp8(const __nv_bfloat16* merged, __nv_fp8_e4m3* out,
void gate_geglu_merged_fp8(const __nv_bfloat16* merged, __nv_fp8_e4m3* out,
int seq, int half_dim,
const float* d_scale, cudaStream_t stream) {
int total = seq * half_dim;
int blocks = (total + 255) / 256;
gate_silu_mul_merged_fp8_kernel<__nv_bfloat16><<<blocks, 256, 0, stream>>>(
gate_geglu_merged_fp8_kernel<__nv_bfloat16><<<blocks, 256, 0, stream>>>(
merged, out, seq, half_dim, d_scale);
}
void gate_silu_mul_merged_fp8_fp16(const __half* merged, __nv_fp8_e4m3* out,
void gate_geglu_merged_fp8_fp16(const __half* merged, __nv_fp8_e4m3* out,
int seq, int half_dim,
const float* d_scale, cudaStream_t stream) {
// 4 elem/thread, matching production throughput
int total = seq * half_dim;
int blocks = (total / 4 + 255) / 256;
gate_silu_mul_merged_fp8_kernel_fp16<<<blocks, 256, 0, stream>>>(
gate_geglu_merged_fp8_kernel_fp16<<<blocks, 256, 0, stream>>>(
merged, out, seq, half_dim, d_scale);
}

Expand Down
14 changes: 8 additions & 6 deletions csrc/kernels/activation.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@

// ── BF16 (original signatures, backward compatible) ──

void gate_silu_mul(const __nv_bfloat16* gate, const __nv_bfloat16* up,
// NOTE: these were previously misnamed `gate_silu_mul*`; they compute the
// tanh/sigmoid-approx GELU (GeGLU), not SiLU. Renamed to `gate_geglu*`.
void gate_geglu(const __nv_bfloat16* gate, const __nv_bfloat16* up,
__nv_bfloat16* out, int n, cudaStream_t stream = 0);

void gelu_inplace(__nv_bfloat16* x, int n, cudaStream_t stream = 0);
Expand All @@ -30,21 +32,21 @@ void bias_gelu_inplace_bf16_strict(__nv_bfloat16* x,
const __nv_bfloat16* bias,
int M, int N, cudaStream_t stream = 0);

void gate_silu_mul_merged(const __nv_bfloat16* merged, __nv_bfloat16* out,
void gate_geglu_merged(const __nv_bfloat16* merged, __nv_bfloat16* out,
int seq, int half_dim, cudaStream_t stream = 0);

void gate_silu_mul_merged_fp8(const __nv_bfloat16* merged, __nv_fp8_e4m3* out,
void gate_geglu_merged_fp8(const __nv_bfloat16* merged, __nv_fp8_e4m3* out,
int seq, int half_dim,
const float* d_scale, cudaStream_t stream = 0);

// ── FP16 variants ──

void gate_silu_mul_fp16(const __half* gate, const __half* up,
void gate_geglu_fp16(const __half* gate, const __half* up,
__half* out, int n, cudaStream_t stream = 0);

void gelu_inplace_fp16(__half* x, int n, cudaStream_t stream = 0);

void gate_silu_mul_merged_fp16(const __half* merged, __half* out,
void gate_geglu_merged_fp16(const __half* merged, __half* out,
int seq, int half_dim, cudaStream_t stream = 0);

// Element-wise multiply: out[i] = a[i] * b[i] for i in [0, n).
Expand All @@ -53,7 +55,7 @@ void gate_silu_mul_merged_fp16(const __half* merged, __half* out,
void mul_fp16(const __half* a, const __half* b, __half* out,
int n, cudaStream_t stream = 0);

void gate_silu_mul_merged_fp8_fp16(const __half* merged, __nv_fp8_e4m3* out,
void gate_geglu_merged_fp8_fp16(const __half* merged, __nv_fp8_e4m3* out,
int seq, int half_dim,
const float* d_scale, cudaStream_t stream = 0);

Expand Down