Make the _apply_rotary_emb compatible with dynamo

[houseroad]

Summary:
There are two things to fix

1. some python files in vllm_flash_attn
2. move the import out of the dynamo scope.

Differential Revision: D73899797

[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.

🚀

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73899797

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D73899797

[jeejeelee]

Similar PR: #17370

[zou3519]

Yup this looks better than #17370, thank you

[jeejeelee]

Yup this looks better than #17370, thank you

Could you explain the reason? It seems that upgrading to torch 2.7.0 fixed it, and this PR doesn't seem to have any effect?

[zou3519]

Yup this looks better than #17370, thank you

Could you explain the reason? It seems that upgrading to torch 2.7.0 fixed it, and this PR doesn't seem to have any effect?

Ahh okay. Thanks for that piece of information. I agree with the assessment that this PR looks unrelated. I have a hypothesis for what the problem was

In PyTorch 2.7 we fixed inductor to respect the strides of Tensors inputs to triton kernels. Previously Inductor tried to do this but there were some edge cases where this didn't happen. If the strides of the Tensor inputs to triton kernels are incorrect then that can lead to CUDA device asserts.

The way to test this is to do something like torch.fx.experimental.proxy_tensor._should_save_eager_input_vals = lambda *args, **kwargs: False on PyTorch 2.7 (and if this causes the device assert to come back in PyTorch 2.7 then that validates the hypothesis). @jeejeelee are you able to do that experiment?

[jeejeelee]

@zou3519 PyTorch 2.7 doesn't have the function torch.fx.experimental.proxy_tensor._should_save_eager_input_val.
Coming back to the changes in this PR, I understand that lazy importing within a function would indeed increase the probability of errors when using torch.compile, is that correct?

[jeejeelee]

In PyTorch 2.7 we fixed inductor to respect the strides of Tensors inputs to triton kernels. Previously Inductor tried to do this but there were some edge cases where this didn't happen. If the strides of the Tensor inputs to triton kernels are incorrect then that can lead to CUDA device asserts.

@zou3519 It should be for this reason. When I was debugging earlier, I found that the issue was indeed caused by the input stride information being inconsistent with the actual tensor's stride. Then I thought about directly registering it as a custom op to avoid torch.compile tracing

[zou3519]

@zou3519 PyTorch 2.7 doesn't have the function torch.fx.experimental.proxy_tensor._should_save_eager_input_val.

Sorry I was looking at the wrong version of PyTorch. The experiment would be to comment out the following lines in your PyTorch install:
https://github.com/pytorch/pytorch/blob/924a247fbbe32063966533535b1cc0cf474ef1ed/torch/fx/experimental/proxy_tensor.py#L1115-L1124

Coming back to the changes in this PR, I understand that lazy importing within a function would indeed increase the probability of errors when using torch.compile, is that correct?

It looks like we have some bugs with importing a module inside of a function being compiled by torch.compile

[jeejeelee]

@zou3519 PyTorch 2.7 doesn't have the function torch.fx.experimental.proxy_tensor._should_save_eager_input_val.

Sorry I was looking at the wrong version of PyTorch. The experiment would be to comment out the following lines in your PyTorch install: https://github.com/pytorch/pytorch/blob/924a247fbbe32063966533535b1cc0cf474ef1ed/torch/fx/experimental/proxy_tensor.py#L1115-L1124

Coming back to the changes in this PR, I understand that lazy importing within a function would indeed increase the probability of errors when using torch.compile, is that correct?

It looks like we have some bugs with importing a module inside of a function being compiled by torch.compile

After commenting out this code, I can reproduce this error.

[zou3519]

After commenting out this code, I can reproduce this error.

Perfect, thanks for confirming. I'll write a summary on the original issue.