[torch.compile] Fix RMSNorm + quant fusion in the non-cutlass-fp8 case, rename RedundantReshapesPass to NoopEliminationPass

[ProExpertProg]

This PR fixes the fp8 case, when cutlass_mm is not available. It contains the following fixes:

• Removes the padding for fp8 torch._scaled_mm in the torch.compile case, as branch specialization might not work correctly, and it makes fusion difficult.
• Implements redundant slice and slice_scatter elimination, which is implemented in PyTorch but does not cover all cases. It renames the RedundantReshapesPass to NoopEliminationPass.
• Minor custom pass improvements.

This PR is a pre-requisite PR to #10836, which enables torch.compile on AMD and uses the non-cutlass-fp8 path.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
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 do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[mergify]

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

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

[bnellnm]

Are these always the right ops to use? e.g. is there a torch.ops.aten.slice.default or a torch.ops.aten.slice_scatter.Tensor?

[ProExpertProg]

I haven't seen them so I am not sure - I just went off what I saw. The other overloads could be added easily if we ever see them in the graph

[mgoin]

Is this cached? It could be expensive each forward call

[ProExpertProg]

Yes, in eager mode this will get called on every forward pass, but it will only happen once when compiled. In eager mode there isn't really a better way that's still correct - the only way is to check the config context. I don't think this getter is significant but I haven't measured it.

[tlrmchlsmth]

We could pass in a allow_input_padding flag? and pass it in? I do think this is annoying though. I think it's woth it to do a quick check for performance regressions on a small model eager mode benchmark with cutlass_scaled_mm disabled?

[ProExpertProg]

I think we'd have to pass that flag through the whole call stack though so I don't think it's worth it. I'll run a small model.

[tlrmchlsmth]

We could pass in a allow_input_padding flag? and pass it in? I do think this is annoying though. I think it's woth it to do a quick check for performance regressions on a small model eager mode benchmark with cutlass_scaled_mm disabled?

[tlrmchlsmth]

Looks good overall but I had a few minor comments

[tlrmchlsmth]

Thanks for the great work! LGTM assuming we don't see any performance regression

[ProExpertProg]

Yep will post perf numbers once I have them, thanks!

[BoyuanFeng]

nit: noop elimination for slice errors when end = -1.

repro:

import torch

def dims_equivalent(dim, i_dim) -> bool:
    # Case 1 and 2
    if dim == i_dim or dim == -1:
        return True
    # Case 3
    return isinstance(dim, torch.fx.Node) and dim.meta["val"] == i_dim

input = torch.randn((2,3,2))
dim_index, start, end = 0, 0, -1

input_shape = input.shape

i_dim = input_shape[dim_index]

# following https://github.com/vllm-project/vllm/blob/main/vllm/compilation/noop_elimination.py#L79
if start == 0 and dims_equivalent(end, i_dim):
    is_noop = True
else:
    is_noop = False

if is_noop:
    # input.shape: (2,3,2)
    # torch.ops.aten.slice.Tensor(input, dim_index, start, end).shape: (1,3,2)
    assert input.shape == torch.ops.aten.slice.Tensor(input, dim_index, start, end).shape

similarly for slice_scatter.

[ProExpertProg]

nit: noop elimination for slice errors when end = -1.

repro:

import torch

def dims_equivalent(dim, i_dim) -> bool:
    # Case 1 and 2
    if dim == i_dim or dim == -1:
        return True
    # Case 3
    return isinstance(dim, torch.fx.Node) and dim.meta["val"] == i_dim

input = torch.randn((2,3,2))
dim_index, start, end = 0, 0, -1

input_shape = input.shape

i_dim = input_shape[dim_index]

# following https://github.com/vllm-project/vllm/blob/main/vllm/compilation/noop_elimination.py#L79
if start == 0 and dims_equivalent(end, i_dim):
    is_noop = True
else:
    is_noop = False

if is_noop:
    # input.shape: (2,3,2)
    # torch.ops.aten.slice.Tensor(input, dim_index, start, end).shape: (1,3,2)
    assert input.shape == torch.ops.aten.slice.Tensor(input, dim_index, start, end).shape

similarly for slice_scatter.

Great find - I didn't realize slice handles end differently. Could you file an issue for this please?

[BoyuanFeng]

issue: #17078

Btw, pytorch supports noop elimination for view, slice, and slice_scatter now, which should be equivalent with noop_elimination.py.

[ProExpertProg]

Okay, sounds good, we can probably deprecate the pass, although it's nice to have an easy place to add noop transforms to unblock pattern matching in the short term before they fixes are upstreamed to torch.