[Bugfix] Fix GPT-OSS AR+NORM fusion

[elvischenv]

Purpose

The slice op between all_reduce and rms_norm breaks the AR+NORM fusion with GPT-OSS+B200.

# moe
auto_functionalized_1 = torch.ops.higher_order.auto_functionalized(torch.ops.vllm.moe_forward.default, hidden_states = constant_pad_nd, router_logits = addmm_1, layer_name = 'model.layers.0.mlp.experts');  constant_pad_nd = addmm_1 = None
getitem_3: "bf16[s72, 3072]" = auto_functionalized_1[0];  auto_functionalized_1 = None

# all_reduce
all_reduce_1: "bf16[s72, 3072]" = torch.ops.vllm.all_reduce.default(getitem_3, 'tp:0');  getitem_3 = None

# slice
slice_1: "bf16[s72, 2880]" = torch.ops.aten.slice.Tensor(all_reduce_1, 1, 0, 2880)

# rms_norm
auto_functionalized_2 = torch.ops.higher_order.auto_functionalized(torch.ops._C.fused_add_rms_norm.default, input = slice_1, residual = getitem_18, weight = arg8_1, epsilon = 1e-05);  slice_1 = getitem_18 = arg8_1 = None

This PR move the slice op before the all_reduce to enable the possible fusion:

# moe
auto_functionalized = torch.ops.higher_order.auto_functionalized(torch.ops.vllm.moe_forward.default, hidden_states = constant_pad_nd, router_logits = addmm_1, layer_name = 'model.layers.1.mlp.experts');  constant_pad_nd = addmm_1 = None
getitem: "bf16[s72, 3072]" = auto_functionalized[0];  auto_functionalized = None

# slice
slice_1: "bf16[s72, 2880]" = torch.ops.aten.slice.Tensor(getitem, 1, 0, 2880);  getitem = None

# all_reduce + rms_norm
auto_functionalized_1 = torch.ops.higher_order.auto_functionalized(torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default, allreduce_in = slice_1, residual = getitem_14, norm_out = None, quant_out = None, scale_out = None, rms_gamma = arg8_1, rms_eps = 1e-05, pattern_code = 1, world_rank = 0, world_size = 4, launch_with_pdl = True, trigger_completion_at_end = True, fp32_acc = True, max_token_num = 5825);  slice_1 = getitem_14 = arg8_1 = None

Test Plan && Result

Accuracy:

PR:

[{'eval_name': 'aime25', 'model_name': 'gpt-oss-20b-medium_temp1.0_20251117_060221', 'metric': 0.7166666666666667}]

main:

[{'eval_name': 'aime25', 'model_name': 'gpt-oss-20b-medium_temp1.0_20251117_071029', 'metric': 0.7375}]

Perf (TP8 con8):

PR: 3.8% improvement

============ Serving Benchmark Result ============
Successful requests:                     40
Benchmark duration (s):                  10.84
Total input tokens:                      40960
Total generated tokens:                  40960
Request throughput (req/s):              3.69
Output token throughput (tok/s):         3779.37
Total Token throughput (tok/s):          7558.74
---------------Time to First Token----------------
Mean TTFT (ms):                          54.30
Median TTFT (ms):                        56.31
P99 TTFT (ms):                           71.81
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          2.06
Median TPOT (ms):                        2.07
P99 TPOT (ms):                           2.11
---------------Inter-token Latency----------------
Mean ITL (ms):                           2.07
Median ITL (ms):                         2.06
P99 ITL (ms):                            2.38
----------------End-to-end Latency----------------
Mean E2EL (ms):                          2165.61
Median E2EL (ms):                        2168.17
P99 E2EL (ms):                           2211.53
==================================================

main:

============ Serving Benchmark Result ============
Successful requests:                     40
Benchmark duration (s):                  11.24
Total input tokens:                      40960
Total generated tokens:                  40960
Request throughput (req/s):              3.56
Output token throughput (tok/s):         3644.69
Total Token throughput (tok/s):          7289.38
---------------Time to First Token----------------
Mean TTFT (ms):                          54.18
Median TTFT (ms):                        63.99
P99 TTFT (ms):                           72.94
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          2.14
Median TPOT (ms):                        2.15
P99 TPOT (ms):                           2.17
---------------Inter-token Latency----------------
Mean ITL (ms):                           2.14
Median ITL (ms):                         2.14
P99 ITL (ms):                            2.48
----------------End-to-end Latency----------------
Mean E2EL (ms):                          2245.97
Median E2EL (ms):                        2240.58
P99 E2EL (ms):                           2277.84
==================================================

---

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 refactors the Mixture-of-Experts layer to enable an important fusion optimization (all_reduce + norm). The change involves moving a tensor slicing operation to occur before the all_reduce call, which allows the compiler to fuse all_reduce with the subsequent normalization. The implementation correctly adjusts all call sites to account for this reordering. The logic is sound, and the provided performance metrics show a clear benefit without any significant accuracy degradation. The changes are well-contained and address the issue described.

[ZJY0516]

cc @ProExpertProg

[ProExpertProg]

Nice, this should also reduce the size of the reduction in general! cc @ilmarkov @varun-sundar-rabindranath @bnellnm can you take a look please

[ProExpertProg]

@elvischenv @mgoin should we add GPT-OSS to nightly E2E fusion tests?

[mgoin]

Yes definitely, the 20b is good for CI

[elvischenv]

@ProExpertProg @mgoin Added the 20b e2e fusion test.
Also tested on main and got expected failure:

>   assert int(log_matches[0]) == matches.allreduce_fusion
      ^^^^^^^^^^^^^^^^^
E   AssertionError: assert 25 == 49
E    +  where 25 = int('25')
E    +  and   49 = Matches(attention_fusion=0, allreduce_fusion=49, sequence_parallel=49, async_tp=48).allreduce_fusion

[bnellnm]

Is og_hidden_states going to be the same across all ranks? Also, is the call to contiguous necessary?

[elvischenv]

Please see the above few lines. og_hidden_states is just used by padding before the MoE kernel. trtllm ar kernel just raise error if the memory is not contiguous.

[bnellnm]

Can we move the contiguous call so that it only applies for the trtllm kernel?

[elvischenv]

Sure. Does that mean other all reduce kernels support non-continuous memory?

[bnellnm]

Not sure but there didn't seem to be a call to contiguous before?

[elvischenv]

Sorry, I just thought the issue may not a trtllm specific issue, but a general issue with symm_mem:

vllm/vllm/distributed/device_communicators/symm_mem.py

Line 134 in f6aa122

self.buffer[: inp.numel()].copy_(inp.view(-1))

  File "/workspace/vllm/vllm/compilation/collective_fusion.py", line 560, in call_trtllm_fused_allreduce_norm
    allreduce_out = tensor_model_parallel_all_reduce(allreduce_in)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/vllm/vllm/distributed/communication_op.py", line 14, in tensor_model_parallel_all_reduce
    return get_tp_group().all_reduce(input_)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/vllm/vllm/distributed/parallel_state.py", line 499, in all_reduce
    return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/.venv/lib/python3.12/site-packages/torch/_ops.py", line 1255, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/vllm/vllm/distributed/parallel_state.py", line 122, in all_reduce
    return group._all_reduce_out_place(tensor)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/vllm/vllm/distributed/parallel_state.py", line 506, in _all_reduce_out_place
    return self.device_communicator.all_reduce(input_)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/vllm/vllm/distributed/device_communicators/cuda_communicator.py", line 157, in all_reduce
    out = symm_mem_comm.all_reduce(input_)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspace/vllm/vllm/distributed/device_communicators/symm_mem.py", line 134, in all_reduce
    self.buffer[: inp.numel()].copy_(inp.view(-1))
                                     ^^^^^^^^^^^^
RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

Just pushed a fix with .reshape(). Or maybe there are better solutions?

[ProExpertProg]

It seems like the "Blackwell Compile and Fusion tests" did not get triggered in CI. Could you add fused_moe/layer.py to the list of dependencies?

[elvischenv]

Resolved merge conflict.

[nvpohanh]

@ProExpertProg Could you review again and see if your comment was addressed?

[nvpohanh]

@mgoin could you help us to trigger the failing pipelines? I can't see the pipeline logs. thanks!

[DarkLight1337]

This PR has been reverted by #29483 as it broke LoRA TP tests. Please open a new version of the PR that passes the test.