New PR number #25693

[charlifu]

This PR adds a few fusion pass for Aiter to fusion layernorm + fp8 block quant and silu + fp8 block quant.

[gemini-code-assist]

Code Review

This pull request adds support for fusing layernorm and silu operations with FP8 block quantization for AITemplate on ROCm. The changes introduce new fusion patterns for torch.compile and the corresponding custom operator implementations. The changes also expand AITemplate support to non-MI300 ROCm devices by providing a Triton-based fallback for GEMM operations.

My review found a critical issue in the type hints for the newly added custom operator implementations in vllm/model_executor/layers/layernorm.py. The residual parameter is incorrectly typed as torch.Tensor instead of Optional[torch.Tensor], which will lead to a TypeError during compilation and tracing. I've provided suggestions to fix this. The rest of the changes look good and are consistent with the goal of the pull request.

[gemini-code-assist]

The residual parameter is typed as torch.Tensor, but it can be None when called from the AiterRMSGroupQuantFP8Pattern fusion pass. This will cause a TypeError during compilation. Please change the type hint to Optional[torch.Tensor] to reflect that it can be None.

Suggested change

	x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
	variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
	x: torch.Tensor, residual: Optional[torch.Tensor], weight: torch.Tensor,
	variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:

[gemini-code-assist]

The residual parameter in this fake implementation is typed as torch.Tensor, but it can be None. This will cause a TypeError during fake-tensor tracing for the fusion pass. Please change the type hint to Optional[torch.Tensor].

Suggested change

	x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
	variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
	x: torch.Tensor, residual: Optional[torch.Tensor], weight: torch.Tensor,
	variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:

[mergify]

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

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

[ProExpertProg]

Thanks for submitting this, this is really exciting!

I'm currently overhauling custom op matching in #24604. We also recently added a torch implementation of group quant, could you compare its performance with AITER? Also could you compare the perf of the fused AITER kernel to the fused torch.compile kernel for rmsnorm+quant. Happy to help out with instructions, but overall:

• you'll need [Performance] Move apply_w8a8_block_fp8_linear to an op class #24666 reapplied (it was recently reverted)
• you'll need to disable quant_fp8 using-O.custom_ops+=-quant_fp8
• you'll have to replace the AITER block quant with QuantFP8
  • we should refactor this after [Performance] Move apply_w8a8_block_fp8_linear to an op class #24666 is re-merged so that the aiter op is under QuantFP8 as well

[charlifu]

New PR #25693