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 12, 14:45 UTC
[Bugfix] Fix GPT-OSS AR+NORM fusion (#28841)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Nov 25, 2025 at 15:59 UTC
6330f9477db214477004df6546f86e3f14f8eab9
4 files changed
+24
-7
.buildkite/test-pipeline.yaml
+1
@@ -971,6 +971,7 @@ steps:
971
- vllm/model_executor/layers/layernorm.py
972
- vllm/model_executor/layers/activation.py
973
- vllm/model_executor/layers/quantization/input_quant_fp8.py
974
+ - vllm/model_executor/layers/fused_moe/layer.py
975
- tests/compile/test_fusion_attn.py
976
- tests/compile/test_silu_mul_quant_fusion.py
977
- tests/compile/distributed/test_fusion_all_reduce.py
tests/compile/distributed/test_fusions_e2e.py
+11
@@ -111,6 +111,17 @@ if current_platform.is_cuda():
111
async_tp=96, # MLP is MoE, half the fusions of dense
112
),
113
),
114
+ ModelBackendTestCase(
115
+ model_name="openai/gpt-oss-20b",
116
+ model_kwargs=dict(max_model_len=1024, kv_cache_dtype="fp8"),
117
+ backend=AttentionBackendEnum.FLASHINFER,
118
+ matches=Matches(
119
+ attention_fusion=0,
120
+ allreduce_fusion=49,
121
+ sequence_parallel=49,
122
+ async_tp=48,
123
+ ),
124
+ ),
125
]
126
127
elif current_platform.is_rocm():
vllm/distributed/device_communicators/symm_mem.py
+1
-1
@@ -131,7 +131,7 @@ class SymmMemCommunicator:
131
return None
132
if out is None:
133
out = torch.empty_like(inp)
134
- self.buffer[: inp.numel()].copy_(inp.view(-1))
134
+ self.buffer[: inp.numel()].copy_(inp.reshape(-1))
135
136
# Determine which algorithm to use
137
use_multimem = False
vllm/model_executor/layers/fused_moe/layer.py
+11
-6
@@ -1690,6 +1690,10 @@ class FusedMoE(CustomOp):
1690
)
1691
1692
def reduce_output(states: torch.Tensor) -> torch.Tensor:
1693
+ # Slice before all_reduce to enable possible fusion
1694
+ if self.hidden_size != og_hidden_states:
1695
+ states = states[..., :og_hidden_states]
1696
+
1697
if (
1698
not self.is_sequence_parallel
1699
and not self.use_dp_chunking
@@ -1712,11 +1716,12 @@ class FusedMoE(CustomOp):
1716
if self.zero_expert_num is not None and self.zero_expert_num > 0:
1717
assert isinstance(fused_output, tuple)
1718
fused_output, zero_expert_result = fused_output
1715
- return (reduce_output(fused_output) + zero_expert_result)[
1716
- ..., :og_hidden_states
1717
- ]
1719
+ return (
1720
+ reduce_output(fused_output)
1721
+ + zero_expert_result[..., :og_hidden_states]
1722
+ )
1723
else:
1719
- return reduce_output(fused_output)[..., :og_hidden_states]
1724
+ return reduce_output(fused_output)
1725
else:
1726
if current_platform.is_tpu():
1727
# TODO: Once the OOM issue for the TPU backend is resolved, we
@@ -1729,8 +1734,8 @@ class FusedMoE(CustomOp):
1734
hidden_states, router_logits, self.layer_name
1735
)
1736
return (
1732
- reduce_output(shared_output)[..., :og_hidden_states],
1733
- reduce_output(fused_output)[..., :og_hidden_states],
1737
+ reduce_output(shared_output),
1738
+ reduce_output(fused_output),
1739
)
1740
1741
def forward_cuda(