[DP/EP][GPTOSS] Use triton matmul-ogs kernels for GPTOSS DP/EP

[varun-sundar-rabindranath]

Purpose

Enable DP/EP for gpt-oss model via DeepEPHightThroughput and triton_kernel's matmul_ogs

Test Plan

Run gpt-oss eval from link

Test Result

<style type="text/css"></style>

GPT-OSS-120B
Reasoning Effort	GPQA
Low	0.66
Medium	0.72
High	0.8

Benchmark

Please find benchmark results here
TLDR - Good TTFT numbers. Bad TPOT numbers as there is no cuda graphs for decode.

---

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces support for Triton matmul-ogs kernels for GPT-OSS models with Data Parallelism and Expert Parallelism, specifically for mxfp4 quantization. This is done by adding a new OAITritonExperts class. The PR also includes substantial changes to dependency management, removing pinned versions of PyTorch and related packages across various requirement files. My review identified some leftover debugging print statements that should be removed before merging.

[varun-sundar-rabindranath]

Status:

command : CUDA_LAUNCH_BLOCKING=1 VLLM_ALL2ALL_BACKEND=deepep_high_throughput canhazgpu run -g2 -- vllm serve openai/gpt-oss-120b --tensor_parallel_size 1 --data-parallel-size 2 --enable-expert-parallel --no-enable-prefix-caching

getting an IMA during model startup

...
(EngineCore_0 pid=1309549)     return self._do_fused_experts(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm/model_executor/layers/fused_moe/modular_kernel.py", line 575, in _do_fused_experts
(EngineCore_0 pid=1309549)     self.fused_experts.apply(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py", line 336, in apply
(EngineCore_0 pid=1309549)     experts_output = triton_kernel_fused_experts(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py", line 171, in triton_kernel_fused_experts
(EngineCore_0 pid=1309549)     intermediate_cache1 = matmul_ogs(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm-test/lib/python3.10/site-packages/triton_kernels/matmul_ogs.py", line 550, in matmul_ogs
(EngineCore_0 pid=1309549)     (kernels._p_matmul_ogs if opt_flags.is_persistent else kernels._matmul_ogs)[(grid,)](
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/.deps/triton/python/triton/runtime/jit.py", line 420, in <lambda>
(EngineCore_0 pid=1309549)     return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/.deps/triton/python/triton/runtime/jit.py", line 742, in run
(EngineCore_0 pid=1309549)     kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata,
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/.deps/triton/python/triton/backends/nvidia/driver.py", line 717, in __call__
(EngineCore_0 pid=1309549)     self.launch(gridX, gridY, gridZ, stream, function, self.launch_cooperative_grid, self.launch_pdl,
(EngineCore_0 pid=1309549) RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered
terminate called after throwing an instance of 'c10::AcceleratorError'
  what():  CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

[varun-sundar-rabindranath]

Status:

command : CUDA_LAUNCH_BLOCKING=1 VLLM_ALL2ALL_BACKEND=deepep_high_throughput canhazgpu run -g2 -- vllm serve openai/gpt-oss-120b --tensor_parallel_size 1 --data-parallel-size 2 --enable-expert-parallel --no-enable-prefix-caching

getting an IMA during model startup

...
(EngineCore_0 pid=1309549)     return self._do_fused_experts(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm/model_executor/layers/fused_moe/modular_kernel.py", line 575, in _do_fused_experts
(EngineCore_0 pid=1309549)     self.fused_experts.apply(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py", line 336, in apply
(EngineCore_0 pid=1309549)     experts_output = triton_kernel_fused_experts(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py", line 171, in triton_kernel_fused_experts
(EngineCore_0 pid=1309549)     intermediate_cache1 = matmul_ogs(
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/vllm-test/lib/python3.10/site-packages/triton_kernels/matmul_ogs.py", line 550, in matmul_ogs
(EngineCore_0 pid=1309549)     (kernels._p_matmul_ogs if opt_flags.is_persistent else kernels._matmul_ogs)[(grid,)](
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/.deps/triton/python/triton/runtime/jit.py", line 420, in <lambda>
(EngineCore_0 pid=1309549)     return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/.deps/triton/python/triton/runtime/jit.py", line 742, in run
(EngineCore_0 pid=1309549)     kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata,
(EngineCore_0 pid=1309549)   File "/home/varun/code/vllm/.deps/triton/python/triton/backends/nvidia/driver.py", line 717, in __call__
(EngineCore_0 pid=1309549)     self.launch(gridX, gridY, gridZ, stream, function, self.launch_cooperative_grid, self.launch_pdl,
(EngineCore_0 pid=1309549) RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered
terminate called after throwing an instance of 'c10::AcceleratorError'
  what():  CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

commit e1e55e09d68eaab088a0433a2328f5eb8675ff4d fixes this.

[varun-sundar-rabindranath]

Update. It appears DeepEP kernels dont support the gpt-oss hidden-size (2880). I have a PR here in the DeepEP repo to fix this deepseek-ai/DeepEP#408 . I see correctness with this.

[zyongye]

You can add padding here so we don't need to change DeepEP side.

vllm/vllm/model_executor/layers/fused_moe/layer.py

Lines 837 to 848 in 7ae9887

	if quant_config and quant_config.get_name() == "mxfp4":
	from vllm.model_executor.layers.quantization.mxfp4 import (
	Mxfp4Backend, get_mxfp4_backend)
	current_mxfp4_backend = get_mxfp4_backend()
	if (current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
	or current_mxfp4_backend
	== Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS):
	hidden_size = round_up(hidden_size, 128)
	elif (current_platform.is_rocm() or current_mxfp4_backend
	== Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM or
	current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
	hidden_size = round_up(hidden_size, 256)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @varun-sundar-rabindranath.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[varun-sundar-rabindranath]

You can add padding here so we don't need to change DeepEP side.

vllm/vllm/model_executor/layers/fused_moe/layer.py

Lines 837 to 848 in 7ae9887

	if quant_config and quant_config.get_name() == "mxfp4":
	from vllm.model_executor.layers.quantization.mxfp4 import (
	Mxfp4Backend, get_mxfp4_backend)
	current_mxfp4_backend = get_mxfp4_backend()
	if (current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
	or current_mxfp4_backend
	== Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS):
	hidden_size = round_up(hidden_size, 128)
	elif (current_platform.is_rocm() or current_mxfp4_backend
	== Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM or
	current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
	hidden_size = round_up(hidden_size, 256)

Nice Idea 🙌 I have updated the code to add this padding. PTAL! Thanks 👍

[bnellnm]

Can this be replaced with cdiv(hidden_size_bytes, xfer_atom_size)?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @varun-sundar-rabindranath.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[varun-sundar-rabindranath]

We need a better way to infer the activation datatype entering the MoE layers. At the moment, all of the MoE code interprets model_config.dtype as the activation datatype as can be seen in the construction of FusedMoEConfig object. I think is very brittle and prone to errors.

cc @mgoin @bnellnm is there a better way that you know of ? Thanks.

[bnellnm]

I think I wrote this bit of code. I don't know of a better way to get the type but I agree that it would be good to make this more robust. The else clause is literally only for one test. I wouldn't worry too much about whether or not that is brittle.

[varun-sundar-rabindranath]

The issue I think is it really depends on the output dtype of the layer before the MoE layer (usually attention). If that layer chooses to quantize the activations, then this would be wrong 😟

[mgoin]

I have little clue beyond this

[mgoin]

I have little clue beyond this