[Feature][Hardware][AMD] Enable level 3 compilation on rocm

[charlifu]

This PR fixs the fusion pass not enabled on rocm by:

• add fp8 dtype selection to the fusion pass, since rocm use torch.float8funz
• use tensor slice operation to replace the torch.narrow op which creates extra ops in the IR generated by torch.compile and makes the rms+fp8_quant fusion [torch.compile] Fuse RMSNorm with quant #9138 not working.

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

🚀

[mgoin]

Could you consider adding AMD to this fusion test case?

vllm/tests/compile/test_fusion.py

Lines 45 to 47 in 9b14d97

	@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda",
	reason="Only test on CUDA")
	def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps):

[charlifu]

@ProExpertProg @mgoin Some updates:

I was trying to enable the unit test of the fusion pass on rocm. I found that with num_token < 17, we are still seeing the extra ops even with the slice operations which fail the test. I think this is because we are padding the input when num_token < 17 and this adds the extra slice_scatter and slice ops into the generated IR.

May I ask it is ok to disable the padding by default? Link

[ProExpertProg]

I am looking into this and found the same issue.

There's also another problem, unrelated to fusion. When we compile with dynamic shape, the max expression will not be a part of the trace, and the graph will contain only the taken branch inside max. A dimension that's marked dynamic still has an underlying value, and that value is erroneously used in the max to pick the larger one, even though that shouldn't be known.

That means if we compile with a dynamic num_tokens that's larger than 17, the graph will always use s0 (the dynamic dimension in place of num_tokens) for the size of the tensor, even if the actual value of s0 is less than 17. The other case is worse: if we originally compile with s0 < 17, the size will always be 17, even if s0 > 17 during execution.

@mgoin I'd advocate for removing the padding in the short term - how often do we deal with num_tokens < 17? In the long term, we can fix the tracing and the fusion for the padded case

[ProExpertProg]

Minor correction: the max does get traced properly into a torch.sym_max(s0, 17) - but for some reason it isn't used so it gets optimized out during the autograd phase.

[mgoin]

Could we simply remove the padding in the non-CUDA case? The reason why it is there is because of bad scaled_mm performance on CUDA

[ProExpertProg]

Ok if we want to keep the padding I think I have a solution for the fusion (but not for dynamic shape compilation). I can implement it tomorrow

[ProExpertProg]

For what it's worth, torch.narrow gets lowered to a slice operation anyway, so at least in the torch.compile regime there should be no difference with slicing.