diff --git a/cpp/models/pi05/src/targets/sm110/operation_driver.cu b/cpp/models/pi05/src/targets/sm110/operation_driver.cu index 83af00148..fd14112a4 100644 --- a/cpp/models/pi05/src/targets/sm110/operation_driver.cu +++ b/cpp/models/pi05/src/targets/sm110/operation_driver.cu @@ -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(merged), static_cast<__half*>(output), rows, hidden, reinterpret_cast(stream)); return launch_status(); @@ -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(merged), static_cast<__nv_fp8_e4m3*>(output), rows, hidden, scale, reinterpret_cast(stream)); diff --git a/cpp/models/pi05/src/targets/sm120/target.cpp b/cpp/models/pi05/src/targets/sm120/target.cpp index 641f32b00..221963f0c 100644 --- a/cpp/models/pi05/src/targets/sm120/target.cpp +++ b/cpp/models/pi05/src/targets/sm120/target.cpp @@ -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(gate), static_cast<__nv_fp8_e4m3*>( binding->fp8_linear->scratch_data()), @@ -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(gate), static_cast<__nv_bfloat16*>(output), rows, hidden_width, cuda_stream); } else { - ::gate_silu_mul( + ::gate_geglu( static_cast(gate), static_cast(up), static_cast<__nv_bfloat16*>(output), rows * hidden_width, diff --git a/csrc/bindings.cpp b/csrc/bindings.cpp index 3a1e0bcac..781a0bc2c 100644 --- a/csrc/bindings.cpp +++ b/csrc/bindings.cpp @@ -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); @@ -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(d_scale), to_stream(stream)); }, py::arg("merged"), py::arg("out"), py::arg("seq"), py::arg("half_dim"), @@ -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(merged), + gate_geglu_merged_fp16(reinterpret_cast(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); @@ -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(merged), + gate_geglu_merged_fp8_fp16(reinterpret_cast(merged), typed_ptr<__nv_fp8_e4m3>(out), seq, half_dim, reinterpret_cast(d_scale), to_stream(stream)); }, py::arg("merged"), py::arg("out"), py::arg("seq"), py::arg("half_dim"), diff --git a/csrc/kernels/activation.cu b/csrc/kernels/activation.cu index 8f5895b18..40717169d 100644 --- a/csrc/kernels/activation.cu +++ b/csrc/kernels/activation.cu @@ -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 -__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; @@ -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 ── @@ -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 -__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; @@ -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><<>>(merged, out, seq, half_dim); + gate_geglu_merged_kernel<__nv_bfloat16><<>>(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><<>>(merged, out, seq, half_dim); + gate_geglu_merged_kernel<__half><<>>(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) { @@ -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; @@ -246,7 +248,7 @@ __global__ void gate_silu_mul_merged_fp8_kernel_fp16(const __half* merged, __nv_ // BF16 generic version (non-encoder paths) template -__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) { @@ -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><<>>( + gate_geglu_merged_fp8_kernel<__nv_bfloat16><<>>( 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<<>>( + gate_geglu_merged_fp8_kernel_fp16<<>>( merged, out, seq, half_dim, d_scale); } diff --git a/csrc/kernels/activation.cuh b/csrc/kernels/activation.cuh index 5c320db39..9335f527c 100644 --- a/csrc/kernels/activation.cuh +++ b/csrc/kernels/activation.cuh @@ -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); @@ -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). @@ -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);