[Kernel] Add cuda kernel for gpt_oss activation

[jeejeelee]

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Test Plan

lm eval

VLLM_ATTENTION_BACKEND=TRITON_ATTN_VLLM_V1 lm_eval --model vllm --model_args "pretrained=unsloth/gpt-oss-20b-BF16,max_model_len=32768" --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto

benchmrk

# serve VLLM_ATTENTION_BACKEND=TRITON_ATTN_VLLM_V1 vllm serve unsloth/gpt-oss-20b-BF16 # bench script vllm bench serve \ --backend vllm \ --model unsloth/gpt-oss-20b-BF16 \ --endpoint /v1/completions \ --dataset-name random \ --random-input 2048 \ --random-output 1024 \ --max-concurrency 10 \ --num-prompt 100 \ 

Test Result(A800)

This PR

vllm (pretrained=unsloth/gpt-oss-20b-BF16,max_model_len=32768,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.3654|± |0.0133| | | |strict-match | 5|exact_match|↑ |0.2593|± |0.0121| ============ Serving Benchmark Result ============ Successful requests: 100 Maximum request concurrency: 10 Benchmark duration (s): 146.79 Total input tokens: 204483 Total generated tokens: 65236 Request throughput (req/s): 0.68 Output token throughput (tok/s): 444.42 Total Token throughput (tok/s): 1837.46 ---------------Time to First Token---------------- Mean TTFT (ms): 389.82 Median TTFT (ms): 178.96 P99 TTFT (ms): 2009.92 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 21.35 Median TPOT (ms): 21.27 P99 TPOT (ms): 25.40 ---------------Inter-token Latency---------------- Mean ITL (ms): 21.16 Median ITL (ms): 19.71 P99 ITL (ms): 137.77 ================================================== 

The main branch

vllm (pretrained=unsloth/gpt-oss-20b-BF16,max_model_len=32768,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.3457|± |0.0131| | | |strict-match | 5|exact_match|↑ |0.2328|± |0.0116| ============ Serving Benchmark Result ============ Successful requests: 100 Maximum request concurrency: 10 Benchmark duration (s): 149.69 Total input tokens: 204483 Total generated tokens: 64990 Request throughput (req/s): 0.67 Output token throughput (tok/s): 434.16 Total Token throughput (tok/s): 1800.19 ---------------Time to First Token---------------- Mean TTFT (ms): 371.65 Median TTFT (ms): 188.81 P99 TTFT (ms): 2388.54 -----Time per Output Token (excl. 1st token)------ Mean TPOT (ms): 24.39 Median TPOT (ms): 21.64 P99 TPOT (ms): 41.43 ---------------Inter-token Latency---------------- Mean ITL (ms): 21.58 Median ITL (ms): 19.94 P99 ITL (ms): 147.82 ================================================== 

@mgoin The value of GSM8K is too low. We also tested using the transformers backend and found that the scores were similar.

 hf result |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.3829|± |0.0134| | | |strict-match | 5|exact_match|↑ |0.2570|± |0.0120| 

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request introduces a new CUDA kernel for the swiglu_oai activation function, which is used in gpt_oss models. The changes span the CUDA kernel implementation, PyTorch bindings, a Python wrapper class, and corresponding tests. My review has identified a critical issue in the Python implementation where the data layout assumption is inconsistent with the CUDA kernel, and a high-severity bug in the test logic that inadvertently breaks existing tests for another activation function.

[gemini-code-assist]

The forward_native implementation for SwiGLUOAI assumes an interleaved layout for gate and up tensors (x[..., ::2], x[..., 1::2]). However, the corresponding CUDA kernel clamp_swiglu_kernel_with_params expects a concatenated layout (first half is gate, second half is up), similar to other activation functions in this file like SiluAndMul. This inconsistency will lead to incorrect results when using the native PyTorch path and will cause the tests to fail. The implementation should be updated to match the CUDA kernel's expectation.

Suggested change

	gate, up = x[..., ::2], x[..., 1::2]
	gate = gate.clamp(min=None, max=self.limit)
	up = up.clamp(min=-self.limit, max=self.limit)
	glu = gate * torch.sigmoid(gate * self.alpha)
	gated_output = (up + 1) * glu
	return gated_output
	d = x.shape[-1] // 2
	gate, up = x[..., :d], x[..., d:]
	gate = gate.clamp(min=None, max=self.limit)
	up = up.clamp(min=-self.limit, max=self.limit)
	glu = gate * torch.sigmoid(gate * self.alpha)
	return (up + 1) * glu

[jeejeelee]

Yes, you are right, but we should fix the kernel rather than this part

[gemini-code-assist]

The conditional logic for opcheck is incorrect. By adding a new if statement instead of an elif, you've broken the logic for the fatrelu activation. When activation == "fatrelu", opcheck is now called twice: once correctly, and a second time inside the final else block with the wrong number of arguments, which will cause the test to fail. This should be a single if/elif/else chain.

Suggested change

	if activation == "fatrelu":
	opcheck(fn, (out, x, threshold))
	if activation == "swiglu_oai":
	opcheck(fn, (out, x, layer.alpha, layer.limit))
	else:
	opcheck(fn, (out, x))
	if activation == "fatrelu":
	opcheck(fn, (out, x, threshold))
	elif activation == "swiglu_oai":
	opcheck(fn, (out, x, layer.alpha, layer.limit))
	else:
	opcheck(fn, (out, x))

[jeejeelee]

Good catchs

[jeejeelee]

@WoosukKwon @zyongye It looks like if we can use triton kernel, is it still necessary to add cuda kernel?

[zyongye]

That's right. The triton kernel should support this, but I guess it can use that only for SM80+?

[mgoin]

Seems like this should be a SwigluOAIAndMul. Don't forget to put it in the registry at the bottom

[mgoin]

LGTM, just a few nits

[mgoin]

LGTM, thanks!

[zyongye]

could you change the name here as well?

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

Line 71 in c9232d4

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

[jeejeelee]

could you change the name here as well?

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

Line 71 in c9232d4

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

Done in 8471138

[simon-mo]

@jeejeelee please re-open, I ran into issue after merge, see #22948 (comment)

[jeejeelee]

the merge order case, I'll try to resolve it.