Fix PyTorch 2.8 API compatibility for symm_mem.fused_scaled_matmul_reduce_scatter

[louiswang524]

Summary

Fixes PyTorch 2.8 API compatibility issue in collective fusion patterns by adding the missing orig_scatter_dim
parameter to torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter calls.

Problem

PyTorch 2.8 introduced a breaking change where fused_scaled_matmul_reduce_scatter now requires an
orig_scatter_dim parameter. This was causing compilation test failures in tests/compile/test_async_tp.py.

Solution

Added orig_scatter_dim=0 parameter to both instances of the function call in:

• ScaledMMReduceScatterPattern.replacement() (line 168)
• CutlassScaledMMReduceScatterPattern.replacement() (line 281)

Testing

The fix maintains backward compatibility and addresses the specific error mentioned in issue #24376.

Fixes #24376

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 ask your reviewers to trigger select CI tests on top of fastcheck CI.

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 either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request aims to fix an API compatibility issue with PyTorch 2.8 by adding the orig_scatter_dim parameter to fused_scaled_matmul_reduce_scatter calls. While the changes are correct for PyTorch 2.8, they break backward compatibility with older versions of PyTorch, which will cause runtime errors. My review includes critical feedback on how to address this by adding the new parameter conditionally based on the PyTorch version, ensuring the fix works across different environments.

[gemini-code-assist]

While this change fixes compatibility with PyTorch 2.8, it introduces a backward compatibility issue with older PyTorch versions (e.g., 2.7) that do not have the orig_scatter_dim argument. This will raise a TypeError on older versions.

To ensure backward compatibility, this argument should be added conditionally. I recommend refactoring the call to use a kwargs dictionary and add orig_scatter_dim only if the PyTorch version is 2.8 or newer.

Example:

from vllm.utils import is_torch_equal_or_newer

kwargs = {
    "input": input,
    "mat2": mat2,
    "scale_a": scale_a,
    "scale_b": scale_b,
    "reduce_op": "avg",
    "scatter_dim": 0,
    "out_dtype": self.dtype,
    "group_name": self.tp.device_group.group_name,
}
if is_torch_equal_or_newer("2.8"):
    kwargs["orig_scatter_dim"] = 0

gemm_rs = torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(**kwargs)

[gemini-code-assist]

This change, while fixing compatibility for PyTorch 2.8, breaks backward compatibility for older versions like 2.7. The orig_scatter_dim argument does not exist in fused_scaled_matmul_reduce_scatter in older PyTorch versions, which will lead to a TypeError. This argument should be added conditionally based on the PyTorch version to maintain compatibility.

Example:

from vllm.utils import is_torch_equal_or_newer

kwargs = {
    # ... other args
}
if is_torch_equal_or_newer("2.8"):
    kwargs["orig_scatter_dim"] = 0

gemm_rs = torch.ops.symm_mem.fused_scaled_matmul_reduce_scatter(**kwargs)

[louiswang524]

Closing due to incorrect analysis of the issue

[andoorve]

Hey @louiswang524 why was the analysis incorrect?

cc: @jasonlizhengjian