[Bug] [ROCm] Fix Llama 4 Enablement Bug on ROCm: V0 ROCmFlashAttentionImpl and Triton Fused MoE bugs

[tjtanaa]

Description

This PR fixes two bugs:

1. TypeError: ROCmFlashAttentionImpl.__init__() got an unexpected keyword argument 'use_irope'
2. Fix the topk_weights in invoke_fused_moe_kernel being not contiguous under V1 + ROCm + torch.compile + Dynamo + hipgraph mode.

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

🚀

[hongxiayang]

cc @houseroad

[hongxiayang]

@simon-mo @houseroad @SageMoore Can you help to merge this? This will unblock us from aiter integration for performance improvement. Thanks!

[ProExpertProg]

Is it possible that this is related to Inductor always putting matrices in row-major order? And we should add a modifier to the custom op?

See comment in torch_bindings.cpp:

  // The default behavior in PyTorch 2.6 is "requires_contiguous", so we need
  // to override this for many GEMMs with the following tag. Otherwise,
  // torch.compile will force all input tensors to be contiguous(), which
  // will break many custom ops that require column-major weight matrices.
  // TODO: remove this for PyTorch 2.8, when the default is planned to switch
  // to match exact eager-mode strides.
  at::Tag stride_tag = at::Tag::needs_fixed_stride_order;

[tjtanaa]

I think the issue might not be related to inductor as it does not happen on CUDA.
the topk_weights.stride() on CUDA returns (1,1) but on ROCm returns (1,1024)

[houseroad]

Can we create an issue to track this hack if @ProExpertProg's suggestion doesn't work

[hongxiayang]

there is an issue created: ROCm/pytorch#2020

[tjtanaa]

@ProExpertProg @houseroad
The topk_weights is generated using Llama4MoE.custom_routing_function which is just a series of native PyTorch operator. So, there is no custom ops involved in generating topk_weights.

vllm/vllm/model_executor/models/llama4.py

Line 44 in 027b204

class Llama4MoE(nn.Module):

class Llama4MoE(nn.Module):

    @staticmethod
    def custom_routing_function(
        hidden_states: torch.Tensor,
        gating_output: torch.Tensor,
        topk: int,
        renormalize: bool,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        router_scores, router_indices = torch.topk(gating_output, topk, dim=-1)
        router_scores = torch.sigmoid(router_scores.float()).to(
            hidden_states.dtype)
        return (router_scores, router_indices.to(torch.int32))

The router_scores is the topk_weights of the fused_moe input.

[ProExpertProg]

It's not about the custom op generating but about consuming a tensor, so the wna16 op consumes this tensor and it might get transposed, unless that IP is marked with the tag

[tjtanaa]

@ProExpertProg
Thank you for the leads. It seems there is a way to add the tag through the PyTorch Python API as well. We have expose the tags interface through direct_register_custom_op from vllm/utils.py which is a functioned proposed by Kaichao to register custom ops that are not traceable by torch.compile. Adding the tags tags=(torch.Tag.needs_fixed_stride_order,), does resolve the issue.

direct_register_custom_op(
    op_name="inplace_fused_experts",
    op_func=inplace_fused_experts,
    mutates_args=["hidden_states"],
    fake_impl=inplace_fused_experts_fake,
+    tags=(torch.Tag.needs_fixed_stride_order,),
)

[ProExpertProg]

Yep this looks right, great work!

[SageMoore]

This looks reasonable. Let's try @ProExpertProg's suggestion for fixing the topk_weights issue.

[houseroad]

Looks fine for unblocking now. We need to create 2 follow ups.

[houseroad]

Can we create an issue to track this hack if @ProExpertProg's suggestion doesn't work

[houseroad]

Create an issue to trace this progress?

[xw285cornell]

i think the output will be incorrect with global attention

[houseroad]

I remember it seems reasonable, but we should definitely have the right approach here.

[hongxiayang]

Agreed about tracking this issue if we want to fully support V0. We will create one internally. Does that sound good to you?

[ProExpertProg]

Thanks for adding support for tags! LGTM assuming this tag fixed the original issue!

[DarkLight1337]

Stamp