@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f1f0d2fab8

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Revert "[Kernel] Add cuda kernel for gpt_oss activation" (#22948)

Simon Mo committed Aug 14, 2025 at 17:38 UTC f1f0d2fab8a1b1ced68ccf5b7197393cf01e1a02

8 files changed +24 -150

csrc/activation_kernels.cu

-59

@@ -128,45 +128,6 @@ __global__ void act_and_mul_kernel_with_param(

128 }

129 }

130

131 -template <typename T>

132 -__device__ __forceinline__ T swigluoai_and_mul(const T& gate, const T& up,

133 - float alpha, float limit) {

134 - // clamp gate: min=None, max=limit

135 - const float gate_f = (float)gate;

136 - const float clamped_gate = gate_f > limit ? limit : gate_f;

137 -

138 - // clamp up: min=-limit, max=limit

139 - const float up_f = (float)up;

140 - const float clamped_up =

141 - up_f > limit ? limit : (up_f < -limit ? -limit : up_f);

142 -

143 - // glu = gate * sigmoid(gate * alpha)

144 - const float sigmoid_val = 1.0f / (1.0f + expf(-clamped_gate * alpha));

145 - const float glu = clamped_gate * sigmoid_val;

146 -

147 - // (up + 1) * glu

148 - return (T)((clamped_up + 1.0f) * glu);

149 -}

150 -

151 -template <typename scalar_t,

152 - scalar_t (*ACT_FN)(const scalar_t&, const scalar_t&, const float,

153 - const float)>

154 -__global__ void swigluoai_and_mul_kernel(

155 - scalar_t* __restrict__ out, // [..., d]

156 - const scalar_t* __restrict__ input, // [..., 2, d]

157 - const int d, const float alpha, const float limit) {

158 - const int64_t token_idx = blockIdx.x;

159 - // TODO: Vectorize loads and stores.

160 - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {

161 - // gate = x[..., ::2] (even indices)

162 - const scalar_t gate = VLLM_LDG(&input[token_idx * 2 * d + 2 * idx]);

163 - // up = x[..., 1::2] (odd indices)

164 - const scalar_t up = VLLM_LDG(&input[token_idx * 2 * d + 2 * idx + 1]);

165 -

166 - out[token_idx * d + idx] = ACT_FN(gate, up, alpha, limit);

167 - }

168 -}

169 -

131 } // namespace vllm

132

133 #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM) \

@@ -184,31 +145,11 @@ __global__ void swigluoai_and_mul_kernel(

145 PARAM); \

146 });

147

187 -#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \

188 - int d = input.size(-1) / 2; \

189 - int64_t num_tokens = input.numel() / input.size(-1); \

190 - dim3 grid(num_tokens); \

191 - dim3 block(std::min(d, 1024)); \

192 - const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \

193 - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \

194 - VLLM_DISPATCH_FLOATING_TYPES( \

195 - input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] { \

196 - vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>> \

197 - <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \

198 - input.data_ptr<scalar_t>(), d, ALPHA, \

199 - LIMIT); \

200 - });

201 -

148 void fatrelu_and_mul(torch::Tensor& out, // [..., d],

149 torch::Tensor& input, // [..., 2 * d]

150 double threshold) {

151 LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm::fatrelu_kernel, threshold);

152 }

207 -void swigluoai_and_mul(torch::Tensor& out, // [..., d]

208 - torch::Tensor& input, // [..., 2 * d]

209 - double alpha, double limit) {

210 - LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);

211 -}

153 namespace vllm {

154

155 // Element-wise activation kernel template.

csrc/ops.h

-2

@@ -138,8 +138,6 @@ void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);

138

139 void fatrelu_and_mul(torch::Tensor& out, torch::Tensor& input,

140 double threshold);

141 -void swigluoai_and_mul(torch::Tensor& out, torch::Tensor& input,

142 - double alpha = 1.702, double limit = 7.0);

141

142 void gelu_new(torch::Tensor& out, torch::Tensor& input);

143

csrc/torch_bindings.cpp

-5

@@ -130,11 +130,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

130 ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");

131 ops.impl("fatrelu_and_mul", torch::kCUDA, &fatrelu_and_mul);

132

133 - ops.def(

134 - "swigluoai_and_mul(Tensor! out, Tensor input, float alpha, float limit) "

135 - "-> ()");

136 - ops.impl("swigluoai_and_mul", torch::kCUDA, &swigluoai_and_mul);

137 -

133 // GELU implementation used in GPT-2.

134 ops.def("gelu_new(Tensor! out, Tensor input) -> ()");

135 ops.impl("gelu_new", torch::kCUDA, &gelu_new);

tests/kernels/core/test_activation.py

+6 -39

@@ -11,7 +11,7 @@ from tests.kernels.utils import opcheck

11 from vllm.model_executor.layers.activation import (FastGELU, FatreluAndMul,

12 GeluAndMul, MulAndSilu,

13 NewGELU, QuickGELU,

14 - SiluAndMul, SwigluOAIAndMul)

14 + SiluAndMul)

15 from vllm.platforms import current_platform

17 DTYPES = [torch.half, torch.bfloat16, torch.float]

@@ -25,15 +25,7 @@ CUDA_DEVICES = [

26 @pytest.mark.parametrize(

27 "activation",

28 - [

29 - "silu_and_mul",

30 - "mul_and_silu",

31 - "gelu",

32 - "gelu_tanh",

33 - "fatrelu",

34 - "swigluoai_and_mul",

35 - ],

36 -)

28 + ["silu_and_mul", "mul_and_silu", "gelu", "gelu_tanh", "fatrelu"])

29 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)

30 @pytest.mark.parametrize("d", D)

31 @pytest.mark.parametrize("dtype", DTYPES)

@@ -67,43 +59,18 @@ def test_act_and_mul(

59 threshold = random.uniform(0, 1)

60 layer = FatreluAndMul(threshold)

61 fn = torch.ops._C.fatrelu_and_mul

70 - elif activation == "swigluoai_and_mul":

71 - layer = SwigluOAIAndMul()

72 - fn = torch.ops._C.swigluoai_and_mul

62 out = layer(x)

63 ref_out = layer.forward_native(x)

75 - if activation == "swigluoai_and_mul":

76 -

77 - rtol = {

78 - #For fp16, change the relative tolerance from 1e-3 to 2e-3

79 - torch.float16:

80 - 2e-3,

81 - torch.bfloat16:

82 - 2e-2,

83 - torch.float:

84 - 1.3e-6

85 - }

86 -

87 - def _get_rtol(output) -> float:

88 - return rtol[output.dtype]

89 -

90 - torch.testing.assert_close(out,

91 - ref_out,

92 - atol=get_default_atol(out),

93 - rtol=_get_rtol(out))

94 - else:

95 - # The SiluAndMul, MulAndSilu, GELU and FatReLU implementations are

96 - # equivalent to the native PyTorch implementations, so we can do exact

97 - # comparison.

98 - torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0)

64 + # The SiluAndMul, MulAndSilu, GELU and FatReLU implementations are

65 + # equivalent to the native PyTorch implementations, so we can do exact

66 + # comparison.

67 + torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0)

69 d = x.shape[-1] // 2

70 output_shape = (x.shape[:-1] + (d, ))

71 out = torch.empty(output_shape, dtype=x.dtype, device=x.device)

72 if activation == "fatrelu":

73 opcheck(fn, (out, x, threshold))

105 - elif activation == "swigluoai_and_mul":

106 - opcheck(fn, (out, x, layer.alpha, layer.limit))

74 else:

75 opcheck(fn, (out, x))

vllm/model_executor/layers/activation.py

+3 -38

@@ -239,35 +239,6 @@ class GeluAndMul(CustomOp):

239 return f'approximate={repr(self.approximate)}'

240

241

242 -@CustomOp.register("swigluoai_and_mul")

243 -class SwigluOAIAndMul(CustomOp):

244 - # https://github.com/huggingface/transformers/blob/v4.55.0/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L106-L110

245 - def __init__(self, alpha: float = 1.702, limit: float = 7.0):

246 - super().__init__()

247 - self.alpha = alpha

248 - self.limit = limit

249 -

250 - def forward_native(self, x: torch.Tensor) -> torch.Tensor:

251 - """PyTorch-native implementation equivalent to forward()."""

252 -

253 - gate, up = x[..., ::2], x[..., 1::2]

254 - gate = gate.clamp(min=None, max=self.limit)

255 - up = up.clamp(min=-self.limit, max=self.limit)

256 - glu = gate * torch.sigmoid(gate * self.alpha)

257 - gated_output = (up + 1) * glu

258 - return gated_output

259 -

260 - def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:

261 - d = x.shape[-1] // 2

262 - output_shape = (x.shape[:-1] + (d, ))

263 - out = torch.empty(output_shape, dtype=x.dtype, device=x.device)

264 - torch.ops._C.swigluoai_and_mul(out, x, self.alpha, self.limit)

265 - return out

266 -

267 - def extra_repr(self) -> str:

268 - return f"alpha={repr(self.alpha)}, limit={repr(self.limit)}"

269 -

270 -

242 @CustomOp.register("gelu_new")

243 class NewGELU(CustomOp):

244

@@ -359,7 +330,6 @@ class ReLUSquaredActivation(CustomOp):

330 return torch.square(F.relu(x))

331

332 def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:

362 - #TODO : implement cuda kenrels

333 return self.forward_native(x)

334

335

@@ -436,14 +406,9 @@ def get_act_fn(act_fn_name: str) -> nn.Module:

406

407

408 _ACTIVATION_AND_MUL_REGISTRY = LazyDict({

439 - "gelu":

440 - lambda: GeluAndMul(),

441 - "silu":

442 - lambda: SiluAndMul(),

443 - "geglu":

444 - lambda: GeluAndMul(),

445 - "swigluoai_and_mul":

446 - lambda *args, **kwargs: SwigluOAIAndMul(*args, **kwargs),

409 + "gelu": lambda: GeluAndMul(),

410 + "silu": lambda: SiluAndMul(),

411 + "geglu": lambda: GeluAndMul(),

412 })

413

414

vllm/model_executor/layers/fused_moe/fused_moe.py

+13 -5

@@ -1633,6 +1633,17 @@ def fused_experts_impl(

1633 block_shape=block_shape,

1634 B_bias=w1_bias)

1635

1636 + # TODO fused kernel

1637 + def swiglu_oai(gate_up):

1638 + alpha = 1.702

1639 + limit = 7.0

1640 + gate, up = gate_up[..., ::2], gate_up[..., 1::2]

1641 + gate = gate.clamp(min=None, max=limit)

1642 + up = up.clamp(min=-limit, max=limit)

1643 + glu = gate * torch.sigmoid(gate * alpha)

1644 + gated_output = (up + 1) * glu

1645 + return gated_output

1646 +

1647 # Activation function with multiplication

1648 if activation == "silu" and is_act_and_mul:

1649 torch.ops._C.silu_and_mul(intermediate_cache2,

@@ -1640,16 +1651,13 @@ def fused_experts_impl(

1651 elif activation == "gelu" and is_act_and_mul:

1652 torch.ops._C.gelu_and_mul(intermediate_cache2,

1653 intermediate_cache1.view(-1, N))

1643 - elif activation == "swigluoai" and is_act_and_mul:

1644 - # alpha = 1.702, limit = 7.0

1645 - torch.ops._C.swigluoai_and_mul(intermediate_cache2,

1646 - intermediate_cache1.view(-1, N))

1654 # Activation function without multiplication

1655 elif activation == "silu":

1656 intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))

1657 elif activation == "gelu":

1658 intermediate_cache2 = F.gelu(intermediate_cache1.view(-1, N))

1652 -

1659 + elif activation == "swiglu_oai":

1660 + intermediate_cache2 = swiglu_oai(intermediate_cache1.view(-1, N))

1661 else:

1662 raise ValueError(f"Unsupported FusedMoe activation: {activation}, "

1663 f"with is_act_and_mul={is_act_and_mul}.")

vllm/model_executor/layers/quantization/utils/mxfp4_utils.py

+1 -1

@@ -68,7 +68,7 @@ def _can_support_mxfp4(use_grouped_topk: bool = False,

68 return not (use_grouped_topk or topk_group or num_expert_group

69 or expert_map or custom_routing_function

70 or e_score_correction_bias or apply_router_weight_on_input

71 - or scoring_func != "softmax" or activation != "swigluoai"

71 + or scoring_func != "softmax" or activation != "swiglu_oai"

72 or expert_load_view or logical_to_physical_map

73 or logical_replica_count)

vllm/model_executor/models/gpt_oss.py

+1 -1

@@ -159,7 +159,7 @@ class MLPBlock(torch.nn.Module):

159 prefix=f"{prefix}.experts",

160 apply_router_weight_on_input=False,

161 has_bias=True,

162 - activation="swigluoai")

162 + activation="swiglu_oai")

163

164 def forward(self, x: torch.Tensor) -> torch.Tensor:

165 t = self.norm(x)