Use aiter triton fused_add_rmsnorm_pad for gpt-oss (#30976)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
2026-01-28 14:47:47 -06:00
parent 3e440786af
commit 59bcc5b6f2
9 changed files with 327 additions and 11 deletions
--- a/vllm/model_executor/models/gpt_oss.py
+++ b/vllm/model_executor/models/gpt_oss.py
@@ -187,7 +187,7 @@ class MLPBlock(torch.nn.Module):
            )
        else:
            g = self.router(x)
-        x = self.experts(hidden_states=x, router_logits=g)
+        x = self.experts(hidden_states=x, router_logits=g)[:, : self.hidden_size]

        if self.is_sequence_parallel:
            x = tensor_model_parallel_all_gather(x.contiguous(), 0)