From a2026f464f422c2b03e41c66b7fe17f0fac4d949 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Thu, 7 May 2026 11:07:52 +0000 Subject: [PATCH 1/3] feat: add more activation func --- .../pytorch/kernels/activation/geglu_impl.py | 138 +++++- .../pytorch/kernels/activation/gelu_impl.py | 106 +++++ .../kernels/activation/quick_geglu_impl.py | 132 ++++++ .../pytorch/kernels/activation/swiglu_impl.py | 140 +++++- primus_turbo/pytorch/ops/activation.py | 421 ++++++++++++++++-- .../triton/activation/bias_geglu_kernel.py | 216 +++++++++ .../triton/activation/bias_gelu_kernel.py | 160 +++++++ .../triton/activation/bias_swiglu_kernel.py | 222 +++++++++ .../triton/activation/geglu_kernel.py | 10 +- primus_turbo/triton/activation/gelu_kernel.py | 186 ++++++++ .../triton/activation/quick_geglu_kernel.py | 273 ++++++++++++ .../triton/activation/swiglu_kernel.py | 4 +- primus_turbo/triton/utils/gelu.py | 15 + tests/pytorch/ops/test_activation.py | 303 +++++++++---- tests/pytorch/ops/test_attention.py | 2 +- 15 files changed, 2138 insertions(+), 190 deletions(-) create mode 100644 primus_turbo/pytorch/kernels/activation/gelu_impl.py create mode 100644 primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py create mode 100644 primus_turbo/triton/activation/bias_geglu_kernel.py create mode 100644 primus_turbo/triton/activation/bias_gelu_kernel.py create mode 100644 primus_turbo/triton/activation/bias_swiglu_kernel.py create mode 100644 primus_turbo/triton/activation/gelu_kernel.py create mode 100644 primus_turbo/triton/activation/quick_geglu_kernel.py diff --git a/primus_turbo/pytorch/kernels/activation/geglu_impl.py b/primus_turbo/pytorch/kernels/activation/geglu_impl.py index 9cb976556..8140e9cc2 100644 --- a/primus_turbo/pytorch/kernels/activation/geglu_impl.py +++ b/primus_turbo/pytorch/kernels/activation/geglu_impl.py @@ -9,6 +9,12 @@ import torch import triton +from primus_turbo.triton.activation.bias_geglu_kernel import ( + bias_geglu_bwd_kernel, + bias_geglu_fwd_kernel, + bias_geglu_with_mask_bwd_kernel, + bias_geglu_with_mask_fwd_kernel, +) from primus_turbo.triton.activation.geglu_kernel import ( geglu_bwd_kernel, geglu_fwd_kernel, @@ -16,13 +22,18 @@ geglu_with_mask_fwd_kernel, ) +MASK_BLOCK_SIZE = 8192 -def geglu_fwd_with_probs(x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None): - num_tokens, double_hidden_size = x.size() - probs = probs.unsqueeze(-1) +# ── Weighted GeGLU (with probs/weights) ────────────────────────────────────── + - out = torch.empty(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) +def geglu_fwd(x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None): + num_tokens, double_hidden_size = x.size() + probs = probs.unsqueeze(-1) + alloc_fn = torch.zeros if row_mask is not None else torch.empty + out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + load_width = triton.next_power_of_2(double_hidden_size // 2) if row_mask is None: grid = (num_tokens,) @@ -34,13 +45,11 @@ def geglu_fwd_with_probs(x: torch.Tensor, probs: torch.Tensor, row_mask: Optiona stride_x_token=x.stride(0), stride_probs_token=probs.stride(0), stride_out_token=out.stride(0), - LOAD_WIDTH=triton.next_power_of_2(double_hidden_size // 2), + LOAD_WIDTH=load_width, ) else: assert row_mask.is_cuda, "row_mask must be a CUDA tensor" - - BLOCK_SIZE = 8192 - grid = (BLOCK_SIZE,) + grid = (MASK_BLOCK_SIZE,) geglu_with_mask_fwd_kernel[grid]( x, probs, @@ -50,23 +59,24 @@ def geglu_fwd_with_probs(x: torch.Tensor, probs: torch.Tensor, row_mask: Optiona stride_x_token=x.stride(0), stride_probs_token=probs.stride(0), stride_out_token=out.stride(0), - LOAD_WIDTH=triton.next_power_of_2(double_hidden_size // 2), - BLOCK_SIZE=BLOCK_SIZE, + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, ) return out -def geglu_bwd_with_probs( +def geglu_bwd( grad_out: torch.Tensor, x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None, ): num_tokens, hidden_size = grad_out.size() - - grad_x = torch.empty_like(x) - grad_probs = torch.empty_like(probs) + alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like + grad_x = alloc_fn(x) + grad_probs = alloc_fn(probs) + load_width = triton.next_power_of_2(hidden_size) if row_mask is None: grid = (num_tokens,) @@ -82,13 +92,11 @@ def geglu_bwd_with_probs( stride_probs_token=probs.stride(0), stride_grad_x_token=grad_x.stride(0), stride_grad_probs_token=grad_probs.stride(0), - LOAD_WIDTH=triton.next_power_of_2(hidden_size), + LOAD_WIDTH=load_width, ) else: assert row_mask.is_cuda, "row_mask must be a CUDA tensor" - - BLOCK_SIZE = 8192 - grid = (BLOCK_SIZE,) + grid = (MASK_BLOCK_SIZE,) geglu_with_mask_bwd_kernel[grid]( grad_out, x, @@ -102,8 +110,98 @@ def geglu_bwd_with_probs( stride_probs_token=probs.stride(0), stride_grad_x_token=grad_x.stride(0), stride_grad_probs_token=grad_probs.stride(0), - LOAD_WIDTH=triton.next_power_of_2(hidden_size), - BLOCK_SIZE=BLOCK_SIZE, + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, ) return grad_x, grad_probs + + +# ── Bias GeGLU ─────────────────────────────────────────────────────────────── + + +def bias_geglu_fwd( + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens, double_hidden_size = x.size() + alloc_fn = torch.zeros if row_mask is not None else torch.empty + out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + has_bias = bias is not None + load_width = triton.next_power_of_2(double_hidden_size // 2) + + if row_mask is None: + grid = (num_tokens,) + bias_geglu_fwd_kernel[grid]( + x, + bias, + out, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + ) + else: + grid = (MASK_BLOCK_SIZE,) + bias_geglu_with_mask_fwd_kernel[grid]( + x, + bias, + row_mask, + out, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + ) + + return out + + +def bias_geglu_bwd( + grad_out: torch.Tensor, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens, hidden_size = grad_out.size() + alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like + grad_x = alloc_fn(x) + has_bias = bias is not None + load_width = triton.next_power_of_2(hidden_size) + + if row_mask is None: + grid = (num_tokens,) + bias_geglu_bwd_kernel[grid]( + grad_out, + x, + bias, + grad_x, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_grad_x_token=grad_x.stride(0), + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + ) + else: + grid = (MASK_BLOCK_SIZE,) + bias_geglu_with_mask_bwd_kernel[grid]( + grad_out, + x, + bias, + row_mask, + grad_x, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_grad_x_token=grad_x.stride(0), + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + ) + + return grad_x diff --git a/primus_turbo/pytorch/kernels/activation/gelu_impl.py b/primus_turbo/pytorch/kernels/activation/gelu_impl.py new file mode 100644 index 000000000..adb955222 --- /dev/null +++ b/primus_turbo/pytorch/kernels/activation/gelu_impl.py @@ -0,0 +1,106 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from typing import Optional + +import torch +import triton + +from primus_turbo.triton.activation.bias_gelu_kernel import ( + bias_gelu_bwd_kernel, + bias_gelu_fwd_kernel, + bias_gelu_with_mask_bwd_kernel, + bias_gelu_with_mask_fwd_kernel, +) + +MASK_BLOCK_SIZE = 8192 + + +def bias_gelu_fwd( + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens, hidden_size = x.size() + alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like + out = alloc_fn(x) + has_bias = bias is not None + load_width = triton.next_power_of_2(hidden_size) + + if row_mask is None: + grid = (num_tokens,) + bias_gelu_fwd_kernel[grid]( + x, + bias, + out, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + ) + else: + grid = (MASK_BLOCK_SIZE,) + bias_gelu_with_mask_fwd_kernel[grid]( + x, + bias, + row_mask, + out, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + ) + + return out + + +def bias_gelu_bwd( + grad_out: torch.Tensor, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens, hidden_size = grad_out.size() + alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like + grad_x = alloc_fn(x) + has_bias = bias is not None + load_width = triton.next_power_of_2(hidden_size) + + if row_mask is None: + grid = (num_tokens,) + bias_gelu_bwd_kernel[grid]( + grad_out, + x, + bias, + grad_x, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_grad_x_token=grad_x.stride(0), + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + ) + else: + grid = (MASK_BLOCK_SIZE,) + bias_gelu_with_mask_bwd_kernel[grid]( + grad_out, + x, + bias, + row_mask, + grad_x, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_grad_x_token=grad_x.stride(0), + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + ) + + return grad_x diff --git a/primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py b/primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py new file mode 100644 index 000000000..4dc4c034f --- /dev/null +++ b/primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py @@ -0,0 +1,132 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from typing import Optional + +import torch +import triton + +from primus_turbo.triton.activation.quick_geglu_kernel import ( + quick_geglu_bwd_kernel, + quick_geglu_fwd_kernel, + quick_geglu_with_mask_bwd_kernel, + quick_geglu_with_mask_fwd_kernel, +) + +MASK_BLOCK_SIZE = 8192 + + +def quick_geglu_fwd( + x: torch.Tensor, + linear_offset: float, + bias: Optional[torch.Tensor] = None, + weights: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens, double_hidden_size = x.size() + alloc_fn = torch.zeros if row_mask is not None else torch.empty + out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + has_bias = bias is not None + has_weights = weights is not None + load_width = triton.next_power_of_2(double_hidden_size // 2) + + if row_mask is None: + grid = (num_tokens,) + quick_geglu_fwd_kernel[grid]( + x, + weights, + bias, + out, + linear_offset, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_weights_token=weights.stride(0) if has_weights else 0, + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + HAS_WEIGHTS=has_weights, + ) + else: + grid = (MASK_BLOCK_SIZE,) + quick_geglu_with_mask_fwd_kernel[grid]( + x, + weights, + bias, + row_mask, + out, + linear_offset, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_weights_token=weights.stride(0) if has_weights else 0, + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + HAS_WEIGHTS=has_weights, + ) + + return out + + +def quick_geglu_bwd( + grad_out: torch.Tensor, + x: torch.Tensor, + linear_offset: float, + bias: Optional[torch.Tensor] = None, + weights: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +): + num_tokens, hidden_size = grad_out.size() + grad_x = torch.empty_like(x) + has_bias = bias is not None + has_weights = weights is not None + grad_weights = torch.empty_like(weights) if has_weights else None + load_width = triton.next_power_of_2(hidden_size) + + if row_mask is None: + grid = (num_tokens,) + quick_geglu_bwd_kernel[grid]( + grad_out, + x, + weights, + bias, + grad_x, + grad_weights, + linear_offset, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_weights_token=weights.stride(0) if has_weights else 0, + stride_grad_x_token=grad_x.stride(0), + stride_grad_weights_token=grad_weights.stride(0) if has_weights else 0, + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + HAS_WEIGHTS=has_weights, + ) + else: + grid = (MASK_BLOCK_SIZE,) + quick_geglu_with_mask_bwd_kernel[grid]( + grad_out, + x, + weights, + bias, + row_mask, + grad_x, + grad_weights, + linear_offset, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_weights_token=weights.stride(0) if has_weights else 0, + stride_grad_x_token=grad_x.stride(0), + stride_grad_weights_token=grad_weights.stride(0) if has_weights else 0, + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + HAS_WEIGHTS=has_weights, + ) + + return grad_x, grad_weights diff --git a/primus_turbo/pytorch/kernels/activation/swiglu_impl.py b/primus_turbo/pytorch/kernels/activation/swiglu_impl.py index 17c09b587..a9e8ac396 100644 --- a/primus_turbo/pytorch/kernels/activation/swiglu_impl.py +++ b/primus_turbo/pytorch/kernels/activation/swiglu_impl.py @@ -9,6 +9,12 @@ import torch import triton +from primus_turbo.triton.activation.bias_swiglu_kernel import ( + bias_swiglu_bwd_kernel, + bias_swiglu_fwd_kernel, + bias_swiglu_with_mask_bwd_kernel, + bias_swiglu_with_mask_fwd_kernel, +) from primus_turbo.triton.activation.swiglu_kernel import ( swiglu_bwd_kernel, swiglu_fwd_kernel, @@ -16,13 +22,18 @@ swiglu_with_mask_fwd_kernel, ) +MASK_BLOCK_SIZE = 8192 -def swiglu_fwd_with_probs(x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None): - num_tokens, double_hidden_size = x.size() - probs = probs.unsqueeze(-1) +# ── Weighted SwiGLU (with probs/weights) ───────────────────────────────────── + - out = torch.empty(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) +def swiglu_fwd(x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None): + num_tokens, double_hidden_size = x.size() + probs = probs.unsqueeze(-1) + alloc_fn = torch.zeros if row_mask is not None else torch.empty + out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + load_width = triton.next_power_of_2(double_hidden_size // 2) if row_mask is None: grid = (num_tokens,) @@ -34,13 +45,11 @@ def swiglu_fwd_with_probs(x: torch.Tensor, probs: torch.Tensor, row_mask: Option stride_x_token=x.stride(0), stride_probs_token=probs.stride(0), stride_out_token=out.stride(0), - LOAD_WIDTH=triton.next_power_of_2(double_hidden_size // 2), + LOAD_WIDTH=load_width, ) else: assert row_mask.is_cuda, "row_mask must be a CUDA tensor" - - BLOCK_SIZE = 8192 - grid = (BLOCK_SIZE,) + grid = (MASK_BLOCK_SIZE,) swiglu_with_mask_fwd_kernel[grid]( x, probs, @@ -50,23 +59,24 @@ def swiglu_fwd_with_probs(x: torch.Tensor, probs: torch.Tensor, row_mask: Option stride_x_token=x.stride(0), stride_probs_token=probs.stride(0), stride_out_token=out.stride(0), - LOAD_WIDTH=triton.next_power_of_2(double_hidden_size // 2), - BLOCK_SIZE=BLOCK_SIZE, + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, ) return out -def swiglu_bwd_with_probs( +def swiglu_bwd( grad_out: torch.Tensor, x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None, ): num_tokens, hidden_size = grad_out.size() - - grad_x = torch.empty_like(x) - grad_probs = torch.empty_like(probs) + alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like + grad_x = alloc_fn(x) + grad_probs = alloc_fn(probs) + load_width = triton.next_power_of_2(hidden_size) if row_mask is None: grid = (num_tokens,) @@ -82,13 +92,11 @@ def swiglu_bwd_with_probs( stride_probs_token=probs.stride(0), stride_grad_x_token=grad_x.stride(0), stride_grad_probs_token=grad_probs.stride(0), - LOAD_WIDTH=triton.next_power_of_2(hidden_size), + LOAD_WIDTH=load_width, ) else: - assert row_mask.is_cuda, "tokens_per_expert must be a CUDA tensor" - - BLOCK_SIZE = 8192 - grid = (BLOCK_SIZE,) + assert row_mask.is_cuda, "row_mask must be a CUDA tensor" + grid = (MASK_BLOCK_SIZE,) swiglu_with_mask_bwd_kernel[grid]( grad_out, x, @@ -102,8 +110,98 @@ def swiglu_bwd_with_probs( stride_probs_token=probs.stride(0), stride_grad_x_token=grad_x.stride(0), stride_grad_probs_token=grad_probs.stride(0), - LOAD_WIDTH=triton.next_power_of_2(hidden_size), - BLOCK_SIZE=BLOCK_SIZE, + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, ) return grad_x, grad_probs + + +# ── Bias SwiGLU ────────────────────────────────────────────────────────────── + + +def bias_swiglu_fwd( + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens, double_hidden_size = x.size() + alloc_fn = torch.zeros if row_mask is not None else torch.empty + out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + has_bias = bias is not None + load_width = triton.next_power_of_2(double_hidden_size // 2) + + if row_mask is None: + grid = (num_tokens,) + bias_swiglu_fwd_kernel[grid]( + x, + bias, + out, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + ) + else: + grid = (MASK_BLOCK_SIZE,) + bias_swiglu_with_mask_fwd_kernel[grid]( + x, + bias, + row_mask, + out, + num_tokens=num_tokens, + stride_x_token=x.stride(0), + stride_out_token=out.stride(0), + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + ) + + return out + + +def bias_swiglu_bwd( + grad_out: torch.Tensor, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + num_tokens, hidden_size = grad_out.size() + alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like + grad_x = alloc_fn(x) + has_bias = bias is not None + load_width = triton.next_power_of_2(hidden_size) + + if row_mask is None: + grid = (num_tokens,) + bias_swiglu_bwd_kernel[grid]( + grad_out, + x, + bias, + grad_x, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_grad_x_token=grad_x.stride(0), + LOAD_WIDTH=load_width, + HAS_BIAS=has_bias, + ) + else: + grid = (MASK_BLOCK_SIZE,) + bias_swiglu_with_mask_bwd_kernel[grid]( + grad_out, + x, + bias, + row_mask, + grad_x, + num_tokens=num_tokens, + stride_grad_out_token=grad_out.stride(0), + stride_x_token=x.stride(0), + stride_grad_x_token=grad_x.stride(0), + LOAD_WIDTH=load_width, + BLOCK_SIZE=MASK_BLOCK_SIZE, + HAS_BIAS=has_bias, + ) + + return grad_x diff --git a/primus_turbo/pytorch/ops/activation.py b/primus_turbo/pytorch/ops/activation.py index 91ed35ce9..b6a1eef7d 100644 --- a/primus_turbo/pytorch/ops/activation.py +++ b/primus_turbo/pytorch/ops/activation.py @@ -4,80 +4,405 @@ # See LICENSE for license information. ############################################################################### -from typing import Union +from typing import Optional import torch from primus_turbo.pytorch.kernels.activation.geglu_impl import ( - geglu_bwd_with_probs, - geglu_fwd_with_probs, + bias_geglu_bwd, + bias_geglu_fwd, + geglu_bwd, + geglu_fwd, +) +from primus_turbo.pytorch.kernels.activation.gelu_impl import ( + bias_gelu_bwd, + bias_gelu_fwd, +) +from primus_turbo.pytorch.kernels.activation.quick_geglu_impl import ( + quick_geglu_bwd, + quick_geglu_fwd, ) from primus_turbo.pytorch.kernels.activation.swiglu_impl import ( - swiglu_bwd_with_probs, - swiglu_fwd_with_probs, + bias_swiglu_bwd, + bias_swiglu_fwd, + swiglu_bwd, + swiglu_fwd, ) -__all__ = ["swiglu_with_probs", "geglu_with_probs"] +__all__ = [ + "bias_gelu_impl", + "bias_geglu_impl", + "bias_swiglu_impl", + "weighted_bias_swiglu_impl", + "weighted_bias_geglu_impl", + "weighted_bias_quick_geglu_impl", +] + + +def _validate_row_mask(row_mask: torch.Tensor, num_rows: int): + assert row_mask.is_cuda, "row_mask must be a CUDA tensor" + assert row_mask.ndim == 1, "row_mask must be 1D tensor" + assert row_mask.size(0) == num_rows, "first dimension of input and row_mask must be the same" + assert row_mask.dtype == torch.int64, "row_mask must be torch.int64" + + +class BiasSwiGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias, row_mask): + ctx.save_for_backward(input, bias, row_mask) + return bias_swiglu_fwd(input, bias, row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, bias, row_mask = ctx.saved_tensors + grad_x = bias_swiglu_bwd(grad_output, input, bias, row_mask) + return grad_x, (grad_x if bias is not None else None), None + + +def _validate_weight_and_row_mask( + input: torch.Tensor, + weights: Optional[torch.Tensor], + row_mask: Optional[torch.Tensor], +): + if weights is not None: + assert input.size(0) == weights.size(0), "first dimension of input and weights must be the same" + assert weights.dtype == torch.float32, "weights must be float32" + if row_mask is not None: + _validate_row_mask(row_mask, input.size(0)) + + +def _clamp_glu_input(input: torch.Tensor, clamp_value: Optional[float]) -> torch.Tensor: + if clamp_value is None: + return input + x_glu, x_linear = input.chunk(2, -1) + return torch.cat( + ( + x_glu.clamp(min=None, max=clamp_value), + x_linear.clamp(min=-clamp_value, max=clamp_value), + ), + -1, + ) + + +class BiasGeLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias, row_mask): + ctx.save_for_backward(input, bias, row_mask) + return bias_gelu_fwd(input, bias, row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, bias, row_mask = ctx.saved_tensors + grad_x = bias_gelu_bwd(grad_output, input, bias, row_mask) + return grad_x, (grad_x if bias is not None else None), None + + +class BiasGeGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias, row_mask): + ctx.save_for_backward(input, bias, row_mask) + return bias_geglu_fwd(input, bias, row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, bias, row_mask = ctx.saved_tensors + grad_x = bias_geglu_bwd(grad_output, input, bias, row_mask) + return grad_x, (grad_x if bias is not None else None), None + + +class BiasQuickGeGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias, linear_offset, row_mask): + ctx.save_for_backward(input, bias, row_mask) + ctx.linear_offset = linear_offset + return quick_geglu_fwd(input, linear_offset, bias=bias, row_mask=row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, bias, row_mask = ctx.saved_tensors + grad_x, _ = quick_geglu_bwd(grad_output, input, ctx.linear_offset, bias=bias, row_mask=row_mask) + return grad_x, (grad_x if bias is not None else None), None, None + + +def bias_gelu_impl( + input: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """ + Bias GELU fusion. + input: [N, D] or [B, S, D] + bias: Optional [D] + row_mask: Optional [N] int64 + output: [N, D] or [B, S, D] + """ + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + + input = input.view(-1, ori_shape[-1]) + _validate_weight_and_row_mask(input, None, row_mask) + + output = BiasGeLUFunction.apply(input, bias, row_mask) + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) + + +def bias_geglu_impl( + input: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, + clamp_value: Optional[float] = None, +) -> torch.Tensor: + """ + Bias GEGLU fusion. + input: [N, 2H] or [B, S, 2H] + bias: Optional [2H] + row_mask: Optional [N] int64 + output: [N, H] or [B, S, H] + """ + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + + input = _clamp_glu_input(input, clamp_value) + input = input.view(-1, ori_shape[-1]) + _validate_weight_and_row_mask(input, None, row_mask) + + output = BiasGeGLUFunction.apply(input, bias, row_mask) + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) + +def bias_swiglu_impl( + input: torch.Tensor, + bias: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, + clamp_value: Optional[float] = None, +) -> torch.Tensor: + """ + Bias SwiGLU fusion. + input: [N, 2H] or [B, S, 2H] + bias: Optional [2H] + row_mask: Optional [N] int64 + output: [N, H] or [B, S, H] + """ + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + + input = _clamp_glu_input(input, clamp_value) + input = input.view(-1, ori_shape[-1]) + _validate_weight_and_row_mask(input, None, row_mask) + + output = BiasSwiGLUFunction.apply(input, bias, row_mask) + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) + + +class WeightedSwiGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weights, row_mask): + ctx.save_for_backward(input, weights, row_mask) + + x = input.view(-1, input.size(-1)) + w = weights.view(-1) + return swiglu_fwd(x, w, row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, weights, row_mask = ctx.saved_tensors + + x = input.view(-1, input.size(-1)) + w = weights.view(-1) + grad_x, grad_w = swiglu_bwd(grad_output, x, w, row_mask) + return grad_x.view(input.shape), grad_w.view(weights.shape), None -class GLUWithProbs(torch.autograd.Function): + +class WeightedGeGLUFunction(torch.autograd.Function): @staticmethod - def forward( - ctx, x: torch.Tensor, probs: torch.Tensor, row_mask: Union[torch.Tensor, None], act_type: str - ): - x_origin_shape = x.size() - probs_origin_shape = probs.size() + def forward(ctx, input, weights, row_mask): + ctx.save_for_backward(input, weights, row_mask) + + x = input.view(-1, input.size(-1)) + w = weights.view(-1) + return geglu_fwd(x, w, row_mask) - x = x.view(-1, x.size(-1)) - probs = probs.view(-1) + @staticmethod + def backward(ctx, grad_output): + input, weights, row_mask = ctx.saved_tensors - assert x.size(0) == probs.size(0), "first dimension of x and probs must be the same" - assert probs.dtype == torch.float32, "probs must be float32" + x = input.view(-1, input.size(-1)) + w = weights.view(-1) + grad_x, grad_w = geglu_bwd(grad_output, x, w, row_mask) + return grad_x.view(input.shape), grad_w.view(weights.shape), None - SUPPORTED_ACT_TYPES = ["silu", "gelu"] - assert ( - act_type in SUPPORTED_ACT_TYPES - ), f"Unsupported act_type: {act_type}. Supported types: {SUPPORTED_ACT_TYPES}" - if row_mask is not None: - assert row_mask.is_cuda, "row_mask must be a CUDA tensor" - assert x.size(0) == row_mask.size(0), "first dimension of x and row_mask must be the same" - assert row_mask.ndim == 1, "row_mask must be 1D tensor" - assert row_mask.dtype == torch.int64, "The dtype of row_mask must be torch.int64." +class WeightedQuickGeGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weights, linear_offset, row_mask): + ctx.save_for_backward(input, weights, row_mask) + ctx.linear_offset = linear_offset + return quick_geglu_fwd(input, linear_offset, weights=weights, row_mask=row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, weights, row_mask = ctx.saved_tensors + grad_x, grad_w = quick_geglu_bwd( + grad_output, input, ctx.linear_offset, weights=weights, row_mask=row_mask + ) + return grad_x, grad_w, None, None - if act_type == "silu": - out = swiglu_fwd_with_probs(x, probs, row_mask) - elif act_type == "gelu": - out = geglu_fwd_with_probs(x, probs, row_mask) - ctx.save_for_backward(x, probs, row_mask) - ctx.act_type = act_type - ctx.x_origin_shape = x_origin_shape - ctx.probs_origin_shape = probs_origin_shape +class WeightedBiasSwiGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias, weights, row_mask): + ctx.save_for_backward(input, bias, weights, row_mask) - return out + y = (input + bias).view(-1, input.size(-1)) + w = weights.view(-1) + return swiglu_fwd(y, w, row_mask) @staticmethod - def backward(ctx, grad_output: torch.Tensor): - assert grad_output.ndim == 2 + def backward(ctx, grad_output): + input, bias, weights, row_mask = ctx.saved_tensors - x, probs, row_mask = ctx.saved_tensors + y = (input + bias).view(-1, input.size(-1)) + w = weights.view(-1) + grad_y, grad_w = swiglu_bwd(grad_output, y, w, row_mask) + grad_y = grad_y.view(input.shape) + return grad_y, grad_y, grad_w.view(weights.shape), None - if ctx.act_type == "silu": - grad_x, grad_probs = swiglu_bwd_with_probs(grad_output, x, probs, row_mask) - elif ctx.act_type == "gelu": - grad_x, grad_probs = geglu_bwd_with_probs(grad_output, x, probs, row_mask) - return grad_x.view(ctx.x_origin_shape), grad_probs.view(ctx.probs_origin_shape), None, None +class WeightedBiasGeGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias, weights, row_mask): + ctx.save_for_backward(input, bias, weights, row_mask) + y = (input + bias).view(-1, input.size(-1)) + w = weights.view(-1) + return geglu_fwd(y, w, row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, bias, weights, row_mask = ctx.saved_tensors -def swiglu_with_probs( - x: torch.Tensor, probs: torch.Tensor, row_mask: Union[torch.Tensor, None] + y = (input + bias).view(-1, input.size(-1)) + w = weights.view(-1) + grad_y, grad_w = geglu_bwd(grad_output, y, w, row_mask) + grad_y = grad_y.view(input.shape) + return grad_y, grad_y, grad_w.view(weights.shape), None + + +class WeightedBiasQuickGeGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias, weights, linear_offset, row_mask): + ctx.save_for_backward(input, bias, weights, row_mask) + ctx.linear_offset = linear_offset + return quick_geglu_fwd(input, linear_offset, bias=bias, weights=weights, row_mask=row_mask) + + @staticmethod + def backward(ctx, grad_output): + input, bias, weights, row_mask = ctx.saved_tensors + grad_x, grad_w = quick_geglu_bwd( + grad_output, + input, + ctx.linear_offset, + bias=bias, + weights=weights, + row_mask=row_mask, + ) + return grad_x, grad_x, grad_w, None, None + + +def weighted_bias_swiglu_impl( + input: torch.Tensor, + bias: Optional[torch.Tensor] = None, + weights: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, + clamp_value: Optional[float] = None, ) -> torch.Tensor: - return GLUWithProbs.apply(x, probs, row_mask, "silu") + """ + Token-wise-weighted bias SwiGLU fusion. + input: [N, 2H] or [B, S, 2H] + bias: Optional [2H] + weights: Optional [N, 1] + row_mask: Optional [N] int64 + output: [N, H] or [B, S, H] + """ + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + + input = _clamp_glu_input(input, clamp_value) + input = input.view(-1, ori_shape[-1]) + _validate_weight_and_row_mask(input, weights, row_mask) + + if weights is not None and bias is not None: + output = WeightedBiasSwiGLUFunction.apply(input, bias, weights, row_mask) + elif weights is not None: + output = WeightedSwiGLUFunction.apply(input, weights, row_mask) + else: + output = BiasSwiGLUFunction.apply(input, bias, row_mask) + + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) -def geglu_with_probs( - x: torch.Tensor, probs: torch.Tensor, row_mask: Union[torch.Tensor, None] +def weighted_bias_geglu_impl( + input: torch.Tensor, + bias: Optional[torch.Tensor] = None, + weights: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, + clamp_value: Optional[float] = None, ) -> torch.Tensor: - return GLUWithProbs.apply(x, probs, row_mask, "gelu") + """ + Token-wise-weighted bias GEGLU fusion. + input: [N, 2H] or [B, S, 2H] + bias: Optional [2H] + weights: Optional [N, 1] + row_mask: Optional [N] int64 + output: [N, H] or [B, S, H] + """ + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + + input = _clamp_glu_input(input, clamp_value) + input = input.view(-1, ori_shape[-1]) + _validate_weight_and_row_mask(input, weights, row_mask) + + if weights is not None and bias is not None: + output = WeightedBiasGeGLUFunction.apply(input, bias, weights, row_mask) + elif weights is not None: + output = WeightedGeGLUFunction.apply(input, weights, row_mask) + else: + output = BiasGeGLUFunction.apply(input, bias, row_mask) + + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) + + +def weighted_bias_quick_geglu_impl( + input: torch.Tensor, + bias: Optional[torch.Tensor] = None, + weights: Optional[torch.Tensor] = None, + row_mask: Optional[torch.Tensor] = None, + linear_offset: float = 0.0, + clamp_value: Optional[float] = None, +) -> torch.Tensor: + """ + Token-wise-weighted bias Quick-GEGLU (sigmoid approximation) fusion. + input: [N, 2H] or [B, S, 2H] + bias: Optional [2H] + weights: Optional [N, 1] + row_mask: Optional [N] int64 + linear_offset: float + output: [N, H] or [B, S, H] + """ + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + + input = _clamp_glu_input(input, clamp_value) + input = input.view(-1, ori_shape[-1]) + linear_offset = float(linear_offset) + _validate_weight_and_row_mask(input, weights, row_mask) + + if weights is not None and bias is not None: + output = WeightedBiasQuickGeGLUFunction.apply(input, bias, weights, linear_offset, row_mask) + elif weights is not None: + output = WeightedQuickGeGLUFunction.apply(input, weights, linear_offset, row_mask) + else: + output = BiasQuickGeGLUFunction.apply(input, bias, linear_offset, row_mask) + + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) diff --git a/primus_turbo/triton/activation/bias_geglu_kernel.py b/primus_turbo/triton/activation/bias_geglu_kernel.py new file mode 100644 index 000000000..c4747fdb8 --- /dev/null +++ b/primus_turbo/triton/activation/bias_geglu_kernel.py @@ -0,0 +1,216 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import triton +import triton.language as tl + +from primus_turbo.triton.utils.gelu import gelu_bwd_none, gelu_none + + +@triton.jit +def bias_geglu_fwd_kernel( + x_ptr, + bias_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token: tl.constexpr, + stride_out_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + + half_stride = stride_x_token // 2 + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + mask = row_mask & col_mask + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + up = up + bias_up + down = down + bias_down + + out = gelu_none(up) * down + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def bias_geglu_with_mask_fwd_kernel( + x_ptr, + bias_ptr, + row_mask_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token, + stride_out_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + idx_type = tl.int64 + + half_stride = stride_x_token // 2 + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + up = up + bias_up + down = down + bias_down + + out = gelu_none(up) * down + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def bias_geglu_bwd_kernel( + grad_out_ptr, + x_ptr, + bias_ptr, + grad_x_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token: tl.constexpr, + stride_x_token: tl.constexpr, + stride_grad_x_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_data_type = grad_x_ptr.dtype.element_ty + + half_stride = stride_x_token // 2 + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + mask = row_mask & col_mask + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + up = up + bias_up + down = down + bias_down + + gelu_up = gelu_none(up) + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + + grad_down = g * gelu_up + grad_up = g * down * gelu_bwd_none(up, tl.full(up.shape, 1.0, dtype=compute_type)) + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_up.to(grad_data_type), + mask=mask, + ) + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + half_stride + col_off, + grad_down.to(grad_data_type), + mask=mask, + ) + + +@triton.jit +def bias_geglu_with_mask_bwd_kernel( + grad_out_ptr, + x_ptr, + bias_ptr, + row_mask_ptr, + grad_x_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token, + stride_x_token, + stride_grad_x_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_data_type = grad_x_ptr.dtype.element_ty + idx_type = tl.int64 + + half_stride = stride_x_token // 2 + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + up = up + bias_up + down = down + bias_down + + gelu_up = gelu_none(up) + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + + grad_down = g * gelu_up + grad_up = g * down * gelu_bwd_none(up, tl.full(up.shape, 1.0, dtype=compute_type)) + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_up.to(grad_data_type), + mask=mask, + ) + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + half_stride + col_off, + grad_down.to(grad_data_type), + mask=mask, + ) diff --git a/primus_turbo/triton/activation/bias_gelu_kernel.py b/primus_turbo/triton/activation/bias_gelu_kernel.py new file mode 100644 index 000000000..8aa7e5e0b --- /dev/null +++ b/primus_turbo/triton/activation/bias_gelu_kernel.py @@ -0,0 +1,160 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import triton +import triton.language as tl + +from primus_turbo.triton.utils.gelu import gelu_bwd_none, gelu_none + + +@triton.jit +def bias_gelu_fwd_kernel( + x_ptr, + bias_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token: tl.constexpr, + stride_out_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + mask = row_mask & col_mask + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + if HAS_BIAS: + b = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + x = x + b + + out = gelu_none(x) + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def bias_gelu_with_mask_fwd_kernel( + x_ptr, + bias_ptr, + row_mask_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token, + stride_out_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + idx_type = tl.int64 + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + if HAS_BIAS: + b = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + x = x + b + + out = gelu_none(x) + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def bias_gelu_bwd_kernel( + grad_out_ptr, + x_ptr, + bias_ptr, + grad_x_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token: tl.constexpr, + stride_x_token: tl.constexpr, + stride_grad_x_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_data_type = grad_x_ptr.dtype.element_ty + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + mask = row_mask & col_mask + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + if HAS_BIAS: + b = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + x = x + b + + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + grad_x = gelu_bwd_none(x, g) + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_x.to(grad_data_type), + mask=mask, + ) + + +@triton.jit +def bias_gelu_with_mask_bwd_kernel( + grad_out_ptr, + x_ptr, + bias_ptr, + row_mask_ptr, + grad_x_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token, + stride_x_token, + stride_grad_x_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_data_type = grad_x_ptr.dtype.element_ty + idx_type = tl.int64 + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + if HAS_BIAS: + b = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + x = x + b + + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + grad_x = gelu_bwd_none(x, g) + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_x.to(grad_data_type), + mask=mask, + ) diff --git a/primus_turbo/triton/activation/bias_swiglu_kernel.py b/primus_turbo/triton/activation/bias_swiglu_kernel.py new file mode 100644 index 000000000..6bd7b3580 --- /dev/null +++ b/primus_turbo/triton/activation/bias_swiglu_kernel.py @@ -0,0 +1,222 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import triton +import triton.language as tl + + +@triton.jit +def bias_swiglu_fwd_kernel( + x_ptr, + bias_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token: tl.constexpr, + stride_out_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + + half_stride = stride_x_token // 2 + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + mask = row_mask & col_mask + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + up = up + bias_up + down = down + bias_down + + silu = tl.fdiv(up, (1.0 + tl.exp(-up))) + out = silu * down + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def bias_swiglu_with_mask_fwd_kernel( + x_ptr, + bias_ptr, + row_mask_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token, + stride_out_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + idx_type = tl.int64 + + half_stride = stride_x_token // 2 + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + up = up + bias_up + down = down + bias_down + + silu = tl.fdiv(up, (1.0 + tl.exp(-up))) + out = silu * down + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def bias_swiglu_bwd_kernel( + grad_out_ptr, + x_ptr, + bias_ptr, + grad_x_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token: tl.constexpr, + stride_x_token: tl.constexpr, + stride_grad_x_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_data_type = grad_x_ptr.dtype.element_ty + + half_stride = stride_x_token // 2 + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + mask = row_mask & col_mask + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + up = up + bias_up + down = down + bias_down + + sigmoid = tl.sigmoid(up) + silu = sigmoid * up + + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + + grad_down = g * silu + grad_silu = sigmoid * (1.0 + up * (1.0 - sigmoid)) + grad_up = g * down * grad_silu + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_up.to(grad_data_type), + mask=mask, + ) + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + half_stride + col_off, + grad_down.to(grad_data_type), + mask=mask, + ) + + +@triton.jit +def bias_swiglu_with_mask_bwd_kernel( + grad_out_ptr, + x_ptr, + bias_ptr, + row_mask_ptr, + grad_x_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token, + stride_x_token, + stride_grad_x_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_data_type = grad_x_ptr.dtype.element_ty + idx_type = tl.int64 + + half_stride = stride_x_token // 2 + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + up = up + bias_up + down = down + bias_down + + sigmoid = tl.sigmoid(up) + silu = sigmoid * up + + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + + grad_down = g * silu + grad_silu = sigmoid * (1.0 + up * (1.0 - sigmoid)) + grad_up = g * down * grad_silu + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_up.to(grad_data_type), + mask=mask, + ) + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + half_stride + col_off, + grad_down.to(grad_data_type), + mask=mask, + ) diff --git a/primus_turbo/triton/activation/geglu_kernel.py b/primus_turbo/triton/activation/geglu_kernel.py index 97b87bc5a..25be93144 100644 --- a/primus_turbo/triton/activation/geglu_kernel.py +++ b/primus_turbo/triton/activation/geglu_kernel.py @@ -53,7 +53,7 @@ def geglu_with_mask_fwd_kernel( up = gelu_none(up) out = up * down - probs = tl.load(probs_ptr + row_idx * stride_probs_token) + probs = tl.load(probs_ptr + row_idx * stride_probs_token, mask=row_mask) out = out * probs tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) @@ -119,12 +119,11 @@ def geglu_with_mask_bwd_kernel( mask=row_mask, ) - probs = tl.load(probs_ptr + row_idx * stride_probs_token).to(compute_type) + probs = tl.load(probs_ptr + row_idx * stride_probs_token, mask=row_mask).to(compute_type) grad_out_with_probs = grad_out * probs grad_down = grad_out_with_probs * gelu - grad_gelu = gelu_bwd_none(up, grad_out) - grad_up = grad_out_with_probs * down * grad_gelu + grad_up = gelu_bwd_none(up, grad_out_with_probs * down) tl.store( grad_x_ptr + row_idx * stride_grad_x_token + col_off, grad_up.to(grad_x_data_type), mask=mask @@ -237,8 +236,7 @@ def geglu_bwd_kernel( grad_out_with_probs = grad_out * probs grad_down = grad_out_with_probs * gelu - grad_gelu = gelu_bwd_none(up, grad_out) - grad_up = grad_out_with_probs * down * grad_gelu + grad_up = gelu_bwd_none(up, grad_out_with_probs * down) tl.store(grad_x_ptr + row_idx * stride_grad_x_token + col_off, grad_up.to(grad_x_data_type), mask=mask) tl.store( diff --git a/primus_turbo/triton/activation/gelu_kernel.py b/primus_turbo/triton/activation/gelu_kernel.py new file mode 100644 index 000000000..fb5d66aad --- /dev/null +++ b/primus_turbo/triton/activation/gelu_kernel.py @@ -0,0 +1,186 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import triton +import triton.language as tl + +from primus_turbo.triton.utils.gelu import gelu_bwd_none, gelu_none + + +@triton.jit +def gelu_fwd_kernel( + x_ptr, + probs_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token: tl.constexpr, + stride_probs_token: tl.constexpr, + stride_out_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + mask = row_mask & col_mask + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + out = gelu_none(x) + + probs = tl.load(probs_ptr + row_idx * stride_probs_token) + out = out * probs + + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def gelu_with_mask_fwd_kernel( + x_ptr, + probs_ptr, + row_mask_ptr, + out_ptr, + num_tokens: tl.constexpr, + stride_x_token, + stride_probs_token, + stride_out_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + idx_type = tl.int64 + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + out = gelu_none(x) + + probs = tl.load(probs_ptr + row_idx * stride_probs_token, mask=row_mask) + out = out * probs + + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def gelu_bwd_kernel( + grad_out_ptr, + x_ptr, + probs_ptr, + grad_x_ptr, + grad_probs_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token: tl.constexpr, + stride_x_token: tl.constexpr, + stride_probs_token: tl.constexpr, + stride_grad_x_token: tl.constexpr, + stride_grad_probs_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_x_data_type = grad_x_ptr.dtype.element_ty + grad_probs_data_type = grad_probs_ptr.dtype.element_ty + + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + mask = row_mask & col_mask + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + gelu_out = gelu_none(x) + + grad_out = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + + # grad_probs = sum(grad_out * gelu(x)) + grad_probs = grad_out * gelu_out + grad_probs_sum = tl.sum(grad_probs, dtype=compute_type) + tl.store( + grad_probs_ptr + row_idx * stride_grad_probs_token, + grad_probs_sum.to(grad_probs_data_type), + mask=row_mask, + ) + + # grad_x = grad_out * probs * gelu'(x) + probs = tl.load(probs_ptr + row_idx * stride_probs_token).to(compute_type) + grad_x = gelu_bwd_none(x, grad_out * probs) + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_x.to(grad_x_data_type), + mask=mask, + ) + + +@triton.jit +def gelu_with_mask_bwd_kernel( + grad_out_ptr, + x_ptr, + probs_ptr, + row_mask_ptr, + grad_x_ptr, + grad_probs_ptr, + num_tokens: tl.constexpr, + stride_grad_out_token, + stride_x_token, + stride_probs_token, + stride_grad_x_token, + stride_grad_probs_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_x_data_type = grad_x_ptr.dtype.element_ty + grad_probs_data_type = grad_probs_ptr.dtype.element_ty + idx_type = tl.int64 + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < stride_x_token + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + x = tl.load(x_ptr + row_idx * stride_x_token + col_off, mask=mask).to(compute_type) + gelu_out = gelu_none(x) + + grad_out = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to( + compute_type + ) + + grad_probs = grad_out * gelu_out + grad_probs_sum = tl.sum(grad_probs, dtype=compute_type) + tl.store( + grad_probs_ptr + row_idx * stride_grad_probs_token, + grad_probs_sum.to(grad_probs_data_type), + mask=row_mask, + ) + + probs = tl.load(probs_ptr + row_idx * stride_probs_token, mask=row_mask).to(compute_type) + grad_x = gelu_bwd_none(x, grad_out * probs) + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + grad_x.to(grad_x_data_type), + mask=mask, + ) diff --git a/primus_turbo/triton/activation/quick_geglu_kernel.py b/primus_turbo/triton/activation/quick_geglu_kernel.py new file mode 100644 index 000000000..0c50e8bec --- /dev/null +++ b/primus_turbo/triton/activation/quick_geglu_kernel.py @@ -0,0 +1,273 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import triton +import triton.language as tl + +from primus_turbo.triton.utils.gelu import quick_gelu + + +@triton.jit +def quick_geglu_fwd_kernel( + x_ptr, + weights_ptr, + bias_ptr, + out_ptr, + linear_offset, + num_tokens: tl.constexpr, + stride_x_token: tl.constexpr, + stride_weights_token: tl.constexpr, + stride_out_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_WEIGHTS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + + half_stride = stride_x_token // 2 + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + mask = row_mask & col_mask + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + up = up + bias_up + down = down + bias_down + + out = quick_gelu(up) * (down + linear_offset) + + if HAS_WEIGHTS: + w = tl.load(weights_ptr + row_idx * stride_weights_token) + out = out * w + + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def quick_geglu_with_mask_fwd_kernel( + x_ptr, + weights_ptr, + bias_ptr, + row_mask_ptr, + out_ptr, + linear_offset, + num_tokens: tl.constexpr, + stride_x_token, + stride_weights_token, + stride_out_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_WEIGHTS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + data_type = x_ptr.dtype.element_ty + idx_type = tl.int64 + + half_stride = stride_x_token // 2 + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + up = up + bias_up + down = down + bias_down + + out = quick_gelu(up) * (down + linear_offset) + + if HAS_WEIGHTS: + w = tl.load(weights_ptr + row_idx * stride_weights_token, mask=row_mask) + out = out * w + + tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) + + +@triton.jit +def quick_geglu_bwd_kernel( + grad_out_ptr, + x_ptr, + weights_ptr, + bias_ptr, + grad_x_ptr, + grad_weights_ptr, + linear_offset, + num_tokens: tl.constexpr, + stride_grad_out_token: tl.constexpr, + stride_x_token: tl.constexpr, + stride_weights_token: tl.constexpr, + stride_grad_x_token: tl.constexpr, + stride_grad_weights_token: tl.constexpr, + LOAD_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_WEIGHTS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_x_dtype = grad_x_ptr.dtype.element_ty + tl.int64 + + half_stride = stride_x_token // 2 + row_idx = pid + row_mask = row_idx < num_tokens + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + mask = row_mask & col_mask + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + up = up + bias_up + down = down + bias_down + + sigmoid_out = tl.sigmoid(1.702 * up) + quick_gelu_up = up * sigmoid_out + + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + + if HAS_WEIGHTS: + activation_val = quick_gelu_up * (down + linear_offset) + grad_w = tl.sum(g * activation_val, dtype=compute_type) + grad_w_dtype = grad_weights_ptr.dtype.element_ty + tl.store( + grad_weights_ptr + row_idx * stride_grad_weights_token, + grad_w.to(grad_w_dtype), + mask=row_mask, + ) + w = tl.load(weights_ptr + row_idx * stride_weights_token).to(compute_type) + g = g * w + + dy_up = g * sigmoid_out * (1.0 + 1.702 * up * (1.0 - sigmoid_out)) * (down + linear_offset) + dy_down = g * quick_gelu_up + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + dy_up.to(grad_x_dtype), + mask=mask, + ) + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + half_stride + col_off, + dy_down.to(grad_x_dtype), + mask=mask, + ) + + +@triton.jit +def quick_geglu_with_mask_bwd_kernel( + grad_out_ptr, + x_ptr, + weights_ptr, + bias_ptr, + row_mask_ptr, + grad_x_ptr, + grad_weights_ptr, + linear_offset, + num_tokens: tl.constexpr, + stride_grad_out_token, + stride_x_token, + stride_weights_token, + stride_grad_x_token, + stride_grad_weights_token, + LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_WEIGHTS: tl.constexpr, +): + pid = tl.program_id(0) + compute_type = tl.float32 + grad_x_dtype = grad_x_ptr.dtype.element_ty + idx_type = tl.int64 + + half_stride = stride_x_token // 2 + col_off = tl.arange(0, LOAD_WIDTH) + col_mask = col_off < half_stride + + if HAS_BIAS: + bias_up = tl.load(bias_ptr + col_off, mask=col_mask).to(compute_type) + bias_down = tl.load(bias_ptr + half_stride + col_off, mask=col_mask).to(compute_type) + + loop = (num_tokens + BLOCK_SIZE - 1) // BLOCK_SIZE + for i in range(0, loop): + row_idx = (i * BLOCK_SIZE + pid).to(idx_type) + row_mask = row_idx < num_tokens + + extra_mask = tl.load(row_mask_ptr + row_idx, mask=row_mask) + mask = (row_mask & col_mask) & (extra_mask != 0) + + up_ptr = x_ptr + row_idx * stride_x_token + down_ptr = up_ptr + half_stride + + up = tl.load(up_ptr + col_off, mask=mask).to(compute_type) + down = tl.load(down_ptr + col_off, mask=mask).to(compute_type) + + if HAS_BIAS: + up = up + bias_up + down = down + bias_down + + sigmoid_out = tl.sigmoid(1.702 * up) + quick_gelu_up = up * sigmoid_out + + g = tl.load(grad_out_ptr + row_idx * stride_grad_out_token + col_off, mask=mask).to(compute_type) + + if HAS_WEIGHTS: + activation_val = quick_gelu_up * (down + linear_offset) + grad_w = tl.sum(g * activation_val, dtype=compute_type) + grad_w_dtype = grad_weights_ptr.dtype.element_ty + tl.store( + grad_weights_ptr + row_idx * stride_grad_weights_token, + grad_w.to(grad_w_dtype), + mask=row_mask, + ) + w = tl.load(weights_ptr + row_idx * stride_weights_token, mask=row_mask).to(compute_type) + g = g * w + + dy_up = g * sigmoid_out * (1.0 + 1.702 * up * (1.0 - sigmoid_out)) * (down + linear_offset) + dy_down = g * quick_gelu_up + + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + col_off, + dy_up.to(grad_x_dtype), + mask=mask, + ) + tl.store( + grad_x_ptr + row_idx * stride_grad_x_token + half_stride + col_off, + dy_down.to(grad_x_dtype), + mask=mask, + ) diff --git a/primus_turbo/triton/activation/swiglu_kernel.py b/primus_turbo/triton/activation/swiglu_kernel.py index bf28d0070..83a1e4cdf 100644 --- a/primus_turbo/triton/activation/swiglu_kernel.py +++ b/primus_turbo/triton/activation/swiglu_kernel.py @@ -51,7 +51,7 @@ def swiglu_with_mask_fwd_kernel( up = tl.fdiv(up, (1.0 + tl.exp(-up))) out = up * down - probs = tl.load(probs_ptr + row_idx * stride_probs_token) + probs = tl.load(probs_ptr + row_idx * stride_probs_token, mask=row_mask) out = out * probs tl.store(out_ptr + row_idx * stride_out_token + col_off, out.to(data_type), mask=mask) @@ -118,7 +118,7 @@ def swiglu_with_mask_bwd_kernel( mask=row_mask, ) - probs = tl.load(probs_ptr + row_idx * stride_probs_token).to(compute_type) + probs = tl.load(probs_ptr + row_idx * stride_probs_token, mask=row_mask).to(compute_type) grad_out_with_probs = grad_out * probs grad_down = grad_out_with_probs * silu diff --git a/primus_turbo/triton/utils/gelu.py b/primus_turbo/triton/utils/gelu.py index 1755ab0ea..bc48db98c 100644 --- a/primus_turbo/triton/utils/gelu.py +++ b/primus_turbo/triton/utils/gelu.py @@ -47,3 +47,18 @@ def gelu_bwd_none(x, dy): dydx = scale2 * x_fp32 * exp(-pow(scale1 * x_fp32, 2)) + 0.5 * tl.erf(scale1 * x_fp32) + 0.5 dx = dydx * dy return dx + + +@triton.jit +def quick_gelu(x): + """quick_gelu(x) = x * sigmoid(1.702 * x)""" + return x * tl.sigmoid(1.702 * x.to(tl.float32)) + + +@triton.jit +def quick_gelu_bwd(x, dy): + """Gradient of quick_gelu: dy * sigmoid(1.702*x) * (1 + 1.702*x*(1 - sigmoid(1.702*x)))""" + x_fp32 = x.to(tl.float32) + sigmoid_out = tl.sigmoid(1.702 * x_fp32) + dydx = sigmoid_out * (1.0 + 1.702 * x_fp32 * (1.0 - sigmoid_out)) + return dydx * dy diff --git a/tests/pytorch/ops/test_activation.py b/tests/pytorch/ops/test_activation.py index fb18323a3..4d29e3583 100644 --- a/tests/pytorch/ops/test_activation.py +++ b/tests/pytorch/ops/test_activation.py @@ -10,119 +10,238 @@ import torch import torch.nn.functional as F -from primus_turbo.pytorch.ops.activation import geglu_with_probs, swiglu_with_probs +from primus_turbo.pytorch.ops.activation import ( + bias_geglu_impl, + bias_gelu_impl, + bias_swiglu_impl, + weighted_bias_geglu_impl, + weighted_bias_quick_geglu_impl, + weighted_bias_swiglu_impl, +) from tests.pytorch.test_utils import get_tolerances -torch.manual_seed(42) -random.seed(42) + +def _reset_seed(): + torch.manual_seed(42) + random.seed(42) + + +# ── Reference implementations (plain PyTorch, fp32 intermediate) ───────────── -# NOTE: Align precision with torch.compile -@torch.compile -def swiglu_with_probs_ref(x: torch.Tensor, probs: torch.Tensor): +def swiglu_ref(x: torch.Tensor): dtype = x.dtype - x = torch.chunk(x, 2, dim=-1) - res = F.silu(x[0]) * x[1] - return (res * probs).to(dtype) + x = x.float() + x_1, x_2 = torch.chunk(x, 2, dim=-1) + return (F.silu(x_1) * x_2).to(dtype) -def generate_tokens_per_expert_list(num_experts: int, num_tokens: int): +def geglu_ref(x: torch.Tensor): + dtype = x.dtype + x = x.float() + x_1, x_2 = torch.chunk(x, 2, dim=-1) + return (F.gelu(x_1) * x_2).to(dtype) + - if num_experts == 1: - return [num_tokens] +def weighted_swiglu_ref(x: torch.Tensor, weights: torch.Tensor): + dtype = x.dtype + x = x.float() + x_1, x_2 = torch.chunk(x, 2, dim=-1) + return (F.silu(x_1) * x_2 * weights.float()).to(dtype) - parts = [] - remaining = num_tokens - for _ in range(num_experts - 1): - val = random.randint(0, remaining) - parts.append(val) - remaining -= val - parts.append(remaining) - return parts +def weighted_geglu_ref(x: torch.Tensor, weights: torch.Tensor): + dtype = x.dtype + x = x.float() + x_1, x_2 = torch.chunk(x, 2, dim=-1) + return (F.gelu(x_1) * x_2 * weights.float()).to(dtype) -# NOTE: Align precision with torch.compile -@torch.compile -def geglu_with_probs_ref(x: torch.Tensor, probs: torch.Tensor): +def weighted_quick_geglu_ref(x: torch.Tensor, weights: torch.Tensor, linear_offset: float = 0.0): dtype = x.dtype - x = torch.chunk(x, 2, dim=-1) - res = F.gelu(x[0]) * x[1] - return (res * probs).to(dtype) + x = x.float() + x_1, x_2 = torch.chunk(x, 2, dim=-1) + return ((x_1 * torch.sigmoid(1.702 * x_1)) * (x_2 + linear_offset) * weights.float()).to(dtype) -@pytest.mark.parametrize( - "batch_size", - [1, 8], -) -@pytest.mark.parametrize( - "sequence_length", - [ - 1, - 128, - 2048, - 2025, - 8192, - ], -) -@pytest.mark.parametrize( - "hidden_size", - [ - 128, - 256, - 2048, - ], -) +# ── Helpers ────────────────────────────────────────────────────────────────── + + +def get_fwd_tolerances(dtype): + if dtype == torch.bfloat16: + return dict(rtol=2e-2, atol=2e-2) + return get_tolerances(dtype) + + +def get_bwd_tolerances(dtype): + if dtype in (torch.bfloat16, torch.float16): + return dict(rtol=3.5e-1, atol=3.5e-1) + return get_tolerances(dtype) + + +def make_row_mask(num_tokens: int, device: str): + """Create a row_mask with ~75% active rows.""" + mask = torch.ones(num_tokens, device=device, dtype=torch.int64) + num_masked = max(1, num_tokens // 4) + indices = torch.randperm(num_tokens, device=device)[:num_masked] + mask[indices] = 0 + return mask + + +# ── Tests: bias_gelu_impl ──────────────────────────────────────────────────── + + +@pytest.mark.parametrize("num_tokens", [1, 128, 2048]) +@pytest.mark.parametrize("hidden_size", [128, 256, 2048]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_bias", [False, True]) +@pytest.mark.parametrize("has_row_mask", [False, True]) +def test_bias_gelu_impl(num_tokens, hidden_size, dtype, has_bias, has_row_mask): + _reset_seed() + device = "cuda" + + x = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype, requires_grad=True) + bias = torch.randn(hidden_size, device=device, dtype=dtype) if has_bias else None + row_mask = make_row_mask(num_tokens, device) if has_row_mask else None + + x_ref = x.clone().detach().float().requires_grad_(True) + + out = bias_gelu_impl(x, bias, row_mask) + + y_ref = x_ref + bias.float() if has_bias else x_ref + out_ref = F.gelu(y_ref).to(dtype) + if has_row_mask: + out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) + + torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) + + grad_out = torch.randn_like(out) + out.backward(grad_out) + out_ref.backward(grad_out) + torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) + + +# ── Tests: bias_geglu_impl / bias_swiglu_impl ──────────────────────────────── + + +@pytest.mark.parametrize("num_tokens", [1, 128, 2048]) +@pytest.mark.parametrize("hidden_size", [128, 256]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("with_tokens_per_expert", [False, True]) -@pytest.mark.parametrize("act_type", ["swiglu", "geglu"]) -def test_glu_with_probs(batch_size, sequence_length, hidden_size, dtype, with_tokens_per_expert, act_type): - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - if act_type == "swiglu": - func = swiglu_with_probs - ref_func = swiglu_with_probs_ref - elif act_type == "geglu": - func = geglu_with_probs - ref_func = geglu_with_probs_ref +@pytest.mark.parametrize("has_bias", [False, True]) +@pytest.mark.parametrize("has_row_mask", [False, True]) +@pytest.mark.parametrize("act_type", ["geglu", "swiglu"]) +def test_bias_glu_impl(num_tokens, hidden_size, dtype, has_bias, has_row_mask, act_type): + _reset_seed() + device = "cuda" + + x = torch.randn(num_tokens, hidden_size * 2, device=device, dtype=dtype, requires_grad=True) + bias = torch.randn(hidden_size * 2, device=device, dtype=dtype) if has_bias else None + row_mask = make_row_mask(num_tokens, device) if has_row_mask else None + + x_ref = x.clone().detach().float().requires_grad_(True) + + if act_type == "geglu": + out = bias_geglu_impl(x, bias, row_mask) + else: + out = bias_swiglu_impl(x, bias, row_mask) + + y_ref = x_ref + bias.float() if has_bias else x_ref + y1, y2 = torch.chunk(y_ref, 2, dim=-1) + if act_type == "geglu": + out_ref = (F.gelu(y1) * y2).to(dtype) + else: + out_ref = (F.silu(y1) * y2).to(dtype) + if has_row_mask: + out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) + + torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) + + grad_out = torch.randn_like(out) + out.backward(grad_out) + out_ref.backward(grad_out) + torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) + + +# ── Tests: weighted_bias_swiglu_impl / weighted_bias_geglu_impl ────────────── + +@pytest.mark.parametrize("num_tokens", [1, 64, 128]) +@pytest.mark.parametrize("hidden_size", [128, 256]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_bias", [False, True]) +@pytest.mark.parametrize("has_row_mask", [False, True]) +@pytest.mark.parametrize("act_type", ["geglu", "swiglu"]) +def test_weighted_bias_glu_impl(num_tokens, hidden_size, dtype, has_bias, has_row_mask, act_type): + _reset_seed() device = "cuda" - probs_dtype = torch.float32 - x = torch.randn( - batch_size, sequence_length, hidden_size * 2, device=device, dtype=dtype, requires_grad=True - ) - probs = torch.rand(batch_size, sequence_length, device=device, dtype=probs_dtype, requires_grad=True) - - num_tokens = batch_size * sequence_length - - x_ref = x.view(-1, x.size(-1)).clone().detach() - x_ref.requires_grad_() - probs_ref = probs.view(-1).unsqueeze(-1).clone().detach() - probs_ref.requires_grad_() - - if with_tokens_per_expert: - num_experts = 64 - tokens_per_expert = torch.tensor( - generate_tokens_per_expert_list(num_experts, num_tokens), device=device, requires_grad=False - ) - row_mask = torch.zeros(num_tokens, device=device, dtype=torch.int64, requires_grad=False) - row_mask[: torch.sum(tokens_per_expert)] = 1 + x = torch.randn(num_tokens, hidden_size * 2, device=device, dtype=dtype, requires_grad=True) + bias = torch.randn(hidden_size * 2, device=device, dtype=dtype) if has_bias else None + weights = torch.rand(num_tokens, 1, device=device, dtype=torch.float32, requires_grad=True) + row_mask = make_row_mask(num_tokens, device) if has_row_mask else None + + x_ref = x.clone().detach().float().requires_grad_(True) + w_ref = weights.clone().detach().requires_grad_(True) + + if act_type == "geglu": + out = weighted_bias_geglu_impl(x, bias, weights, row_mask) + else: + out = weighted_bias_swiglu_impl(x, bias, weights, row_mask) + + y_ref = x_ref + bias.float() if has_bias else x_ref + y1, y2 = torch.chunk(y_ref, 2, dim=-1) + if act_type == "geglu": + act_ref = F.gelu(y1) * y2 else: - row_mask = None + act_ref = F.silu(y1) * y2 + out_ref = (act_ref * w_ref.float()).to(dtype) + if has_row_mask: + out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) + + torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) + + grad_out = torch.randn_like(out) + out.backward(grad_out) + out_ref.backward(grad_out) + torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) + torch.testing.assert_close(weights.grad, w_ref.grad, **get_bwd_tolerances(dtype)) + + +# ── Tests: weighted_bias_quick_geglu_impl ──────────────────────────────────── + + +@pytest.mark.parametrize("num_tokens", [16, 64, 128]) +@pytest.mark.parametrize("hidden_size", [128, 256]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_bias", [False, True]) +@pytest.mark.parametrize("has_row_mask", [False, True]) +@pytest.mark.parametrize("linear_offset", [0.0, 1.0]) +def test_weighted_bias_quick_geglu_impl( + num_tokens, hidden_size, dtype, has_bias, has_row_mask, linear_offset +): + _reset_seed() + device = "cuda" + + x = torch.randn(num_tokens, hidden_size * 2, device=device, dtype=dtype, requires_grad=True) + bias = torch.randn(hidden_size * 2, device=device, dtype=dtype) if has_bias else None + weights = torch.rand(num_tokens, 1, device=device, dtype=torch.float32, requires_grad=True) + row_mask = make_row_mask(num_tokens, device) if has_row_mask else None + + x_ref = x.clone().detach().float().requires_grad_(True) + w_ref = weights.clone().detach().requires_grad_(True) - out = func(x, probs, row_mask) - out_ref = ref_func(x_ref, probs_ref) - torch.testing.assert_close(out, out_ref, **get_tolerances(dtype)) + out = weighted_bias_quick_geglu_impl(x, bias, weights, row_mask, linear_offset) - out.backward(torch.ones_like(out)) - grad_x = x.grad.clone() - grad_probs = probs.grad.clone() + y_ref = x_ref + bias.float() if has_bias else x_ref + y1, y2 = torch.chunk(y_ref, 2, dim=-1) + out_ref = ((y1 * torch.sigmoid(1.702 * y1)) * (y2 + linear_offset) * w_ref.float()).to(dtype) + if has_row_mask: + out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) - out_ref.backward(torch.ones_like(out_ref)) - grad_x_ref = x_ref.grad.clone() - grad_probs_ref = probs_ref.grad.clone() + torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) - torch.testing.assert_close(grad_x, grad_x_ref.view_as(grad_x), **get_tolerances(dtype)) - torch.testing.assert_close(grad_probs, grad_probs_ref.view_as(grad_probs), **get_tolerances(probs_dtype)) + grad_out = torch.randn_like(out) + out.backward(grad_out) + out_ref.backward(grad_out) + torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) + torch.testing.assert_close(weights.grad, w_ref.grad, **get_bwd_tolerances(dtype)) diff --git a/tests/pytorch/ops/test_attention.py b/tests/pytorch/ops/test_attention.py index a031d58bc..bfa642d38 100644 --- a/tests/pytorch/ops/test_attention.py +++ b/tests/pytorch/ops/test_attention.py @@ -48,7 +48,7 @@ @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("config", test_cases) @pytest.mark.parametrize("causal", [True, False]) -@pytest.mark.parametrize("enable_sink", [False, True]) +@pytest.mark.parametrize("enable_sink", [True]) @pytest.mark.parametrize("window_size_left", [-1, 32, 64, 128]) @pytest.mark.parametrize("qkv_format", ["bshd", "sbhd", "bhsd"]) @pytest.mark.parametrize("is_v3_atomic_fp32", [False, True]) From 1cfb424a5c7fd9f3abe0b01de0347daab3ba9350 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Fri, 8 May 2026 02:40:42 +0000 Subject: [PATCH 2/3] fix --- .../pytorch/kernels/activation/geglu_impl.py | 14 +++---- .../pytorch/kernels/activation/gelu_impl.py | 6 +-- .../kernels/activation/quick_geglu_impl.py | 3 +- .../pytorch/kernels/activation/swiglu_impl.py | 14 +++---- tests/pytorch/ops/test_activation.py | 37 ++++++++++++++----- 5 files changed, 40 insertions(+), 34 deletions(-) diff --git a/primus_turbo/pytorch/kernels/activation/geglu_impl.py b/primus_turbo/pytorch/kernels/activation/geglu_impl.py index 8140e9cc2..17c46c663 100644 --- a/primus_turbo/pytorch/kernels/activation/geglu_impl.py +++ b/primus_turbo/pytorch/kernels/activation/geglu_impl.py @@ -31,8 +31,7 @@ def geglu_fwd(x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None): num_tokens, double_hidden_size = x.size() probs = probs.unsqueeze(-1) - alloc_fn = torch.zeros if row_mask is not None else torch.empty - out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + out = torch.empty(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) load_width = triton.next_power_of_2(double_hidden_size // 2) if row_mask is None: @@ -73,9 +72,8 @@ def geglu_bwd( row_mask: Optional[torch.Tensor] = None, ): num_tokens, hidden_size = grad_out.size() - alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like - grad_x = alloc_fn(x) - grad_probs = alloc_fn(probs) + grad_x = torch.empty_like(x) + grad_probs = torch.empty_like(probs) load_width = triton.next_power_of_2(hidden_size) if row_mask is None: @@ -126,8 +124,7 @@ def bias_geglu_fwd( row_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens, double_hidden_size = x.size() - alloc_fn = torch.zeros if row_mask is not None else torch.empty - out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + out = torch.empty(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) has_bias = bias is not None load_width = triton.next_power_of_2(double_hidden_size // 2) @@ -168,8 +165,7 @@ def bias_geglu_bwd( row_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens, hidden_size = grad_out.size() - alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like - grad_x = alloc_fn(x) + grad_x = torch.empty_like(x) has_bias = bias is not None load_width = triton.next_power_of_2(hidden_size) diff --git a/primus_turbo/pytorch/kernels/activation/gelu_impl.py b/primus_turbo/pytorch/kernels/activation/gelu_impl.py index adb955222..48152d7fc 100644 --- a/primus_turbo/pytorch/kernels/activation/gelu_impl.py +++ b/primus_turbo/pytorch/kernels/activation/gelu_impl.py @@ -25,8 +25,7 @@ def bias_gelu_fwd( row_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens, hidden_size = x.size() - alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like - out = alloc_fn(x) + out = torch.empty_like(x) has_bias = bias is not None load_width = triton.next_power_of_2(hidden_size) @@ -67,8 +66,7 @@ def bias_gelu_bwd( row_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens, hidden_size = grad_out.size() - alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like - grad_x = alloc_fn(x) + grad_x = torch.empty_like(x) has_bias = bias is not None load_width = triton.next_power_of_2(hidden_size) diff --git a/primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py b/primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py index 4dc4c034f..89dd1d896 100644 --- a/primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py +++ b/primus_turbo/pytorch/kernels/activation/quick_geglu_impl.py @@ -27,8 +27,7 @@ def quick_geglu_fwd( row_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens, double_hidden_size = x.size() - alloc_fn = torch.zeros if row_mask is not None else torch.empty - out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + out = torch.empty(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) has_bias = bias is not None has_weights = weights is not None load_width = triton.next_power_of_2(double_hidden_size // 2) diff --git a/primus_turbo/pytorch/kernels/activation/swiglu_impl.py b/primus_turbo/pytorch/kernels/activation/swiglu_impl.py index a9e8ac396..9d48a8ba2 100644 --- a/primus_turbo/pytorch/kernels/activation/swiglu_impl.py +++ b/primus_turbo/pytorch/kernels/activation/swiglu_impl.py @@ -31,8 +31,7 @@ def swiglu_fwd(x: torch.Tensor, probs: torch.Tensor, row_mask: Optional[torch.Tensor] = None): num_tokens, double_hidden_size = x.size() probs = probs.unsqueeze(-1) - alloc_fn = torch.zeros if row_mask is not None else torch.empty - out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + out = torch.empty(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) load_width = triton.next_power_of_2(double_hidden_size // 2) if row_mask is None: @@ -73,9 +72,8 @@ def swiglu_bwd( row_mask: Optional[torch.Tensor] = None, ): num_tokens, hidden_size = grad_out.size() - alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like - grad_x = alloc_fn(x) - grad_probs = alloc_fn(probs) + grad_x = torch.empty_like(x) + grad_probs = torch.empty_like(probs) load_width = triton.next_power_of_2(hidden_size) if row_mask is None: @@ -126,8 +124,7 @@ def bias_swiglu_fwd( row_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens, double_hidden_size = x.size() - alloc_fn = torch.zeros if row_mask is not None else torch.empty - out = alloc_fn(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) + out = torch.empty(num_tokens, double_hidden_size // 2, dtype=x.dtype, device=x.device) has_bias = bias is not None load_width = triton.next_power_of_2(double_hidden_size // 2) @@ -168,8 +165,7 @@ def bias_swiglu_bwd( row_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens, hidden_size = grad_out.size() - alloc_fn = torch.zeros_like if row_mask is not None else torch.empty_like - grad_x = alloc_fn(x) + grad_x = torch.empty_like(x) has_bias = bias is not None load_width = triton.next_power_of_2(hidden_size) diff --git a/tests/pytorch/ops/test_activation.py b/tests/pytorch/ops/test_activation.py index 4d29e3583..92f6db652 100644 --- a/tests/pytorch/ops/test_activation.py +++ b/tests/pytorch/ops/test_activation.py @@ -88,6 +88,15 @@ def make_row_mask(num_tokens: int, device: str): return mask +def masked_assert_close(actual, expected, row_mask, **kwargs): + """assert_close only on rows where row_mask != 0, ignoring masked-out rows.""" + if row_mask is None: + torch.testing.assert_close(actual, expected, **kwargs) + return + bool_mask = row_mask.bool() + torch.testing.assert_close(actual[bool_mask], expected[bool_mask], **kwargs) + + # ── Tests: bias_gelu_impl ──────────────────────────────────────────────────── @@ -113,12 +122,14 @@ def test_bias_gelu_impl(num_tokens, hidden_size, dtype, has_bias, has_row_mask): if has_row_mask: out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) - torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) + masked_assert_close(out, out_ref, row_mask, **get_fwd_tolerances(dtype)) grad_out = torch.randn_like(out) + if has_row_mask: + grad_out[~row_mask.bool()] = 0 out.backward(grad_out) out_ref.backward(grad_out) - torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) + masked_assert_close(x.grad, x_ref.grad.to(dtype), row_mask, **get_bwd_tolerances(dtype)) # ── Tests: bias_geglu_impl / bias_swiglu_impl ──────────────────────────────── @@ -154,12 +165,14 @@ def test_bias_glu_impl(num_tokens, hidden_size, dtype, has_bias, has_row_mask, a if has_row_mask: out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) - torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) + masked_assert_close(out, out_ref, row_mask, **get_fwd_tolerances(dtype)) grad_out = torch.randn_like(out) + if has_row_mask: + grad_out[~row_mask.bool()] = 0 out.backward(grad_out) out_ref.backward(grad_out) - torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) + masked_assert_close(x.grad, x_ref.grad.to(dtype), row_mask, **get_bwd_tolerances(dtype)) # ── Tests: weighted_bias_swiglu_impl / weighted_bias_geglu_impl ────────────── @@ -198,13 +211,15 @@ def test_weighted_bias_glu_impl(num_tokens, hidden_size, dtype, has_bias, has_ro if has_row_mask: out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) - torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) + masked_assert_close(out, out_ref, row_mask, **get_fwd_tolerances(dtype)) grad_out = torch.randn_like(out) + if has_row_mask: + grad_out[~row_mask.bool()] = 0 out.backward(grad_out) out_ref.backward(grad_out) - torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) - torch.testing.assert_close(weights.grad, w_ref.grad, **get_bwd_tolerances(dtype)) + masked_assert_close(x.grad, x_ref.grad.to(dtype), row_mask, **get_bwd_tolerances(dtype)) + masked_assert_close(weights.grad, w_ref.grad, row_mask, **get_bwd_tolerances(dtype)) # ── Tests: weighted_bias_quick_geglu_impl ──────────────────────────────────── @@ -238,10 +253,12 @@ def test_weighted_bias_quick_geglu_impl( if has_row_mask: out_ref = out_ref * row_mask.unsqueeze(-1).to(out_ref.dtype) - torch.testing.assert_close(out, out_ref, **get_fwd_tolerances(dtype)) + masked_assert_close(out, out_ref, row_mask, **get_fwd_tolerances(dtype)) grad_out = torch.randn_like(out) + if has_row_mask: + grad_out[~row_mask.bool()] = 0 out.backward(grad_out) out_ref.backward(grad_out) - torch.testing.assert_close(x.grad, x_ref.grad.to(dtype), **get_bwd_tolerances(dtype)) - torch.testing.assert_close(weights.grad, w_ref.grad, **get_bwd_tolerances(dtype)) + masked_assert_close(x.grad, x_ref.grad.to(dtype), row_mask, **get_bwd_tolerances(dtype)) + masked_assert_close(weights.grad, w_ref.grad, row_mask, **get_bwd_tolerances(dtype)) From e700abe7e33e9ed8f028351e2c0e9f155e01919b Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Fri, 8 May 2026 02:49:23 +0000 Subject: [PATCH 3/3] fix --- tests/pytorch/ops/test_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/ops/test_attention.py b/tests/pytorch/ops/test_attention.py index bfa642d38..a031d58bc 100644 --- a/tests/pytorch/ops/test_attention.py +++ b/tests/pytorch/ops/test_attention.py @@ -48,7 +48,7 @@ @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("config", test_cases) @pytest.mark.parametrize("causal", [True, False]) -@pytest.mark.parametrize("enable_sink", [True]) +@pytest.mark.parametrize("enable_sink", [False, True]) @pytest.mark.parametrize("window_size_left", [-1, 32, 64, 128]) @pytest.mark.parametrize("qkv_format", ["bshd", "sbhd", "bhsd"]) @pytest.mark.parametrize("is_v3_atomic_fp32", [False, True])