[FEAT] [ROCm]: Support AITER Linear

[tjtanaa]

Description

This PR is to introduce AITER Linear kernel so that any up-coming optimization in AITER kernel could be directly use and evaluated within vLLM framework.

(Updates) This PR has been updated to also support V1.

Given that the AITER ops is being used by multiple files, and requires to be registered as custom op through direct_register_custom_op for it to be used with V1, we proposed to keep all AITER related flags (helper functions) and ops within vllm/_aiter_ops.py. We would like to propose that moving forward, AITER ops will be first defined in this file and registered as custom ops using direct_register_custom_op. AITER related flags (helper functions) will be defined here as well.

We have added tests/v1/rocm/test_aiter_ops.py . It tests if the aiter ops are

1. correctly registered as custom ops
2. correctly defined the relationship between implementation and fake function
3. can be used with torch.compile
   This file will be skipped if AITER is not installed and the platform is not ROCm.

NOTE:
This unit tests is by no means to check the correctness of the AITER ops. It only checks if the AITER ops are correctly registered and if torch.compile can be used with the AITER ops. The correctness of the AITER ops is tested in the https://github.com/ROCm/aiter

To meet those two criteria, the following checks are done.

Criteria 1

To check for fake tensor implementation is correct or not, it can be done using torch.library.check as follows:

    # Verify the op's fake implementation with FP8 inputs
    # Disable test_schema as fp8 datatype is not supported by
    # torch.library.opcheck
    # Related error:
    #      OpCheckError: opcheck(op, ...): test_schema failed with
    #      "mul_cuda" not implemented for 'Float8_e4m3fnuz'
    torch.library.opcheck(torch.ops.vllm.rocm_aiter_tuned_gemm,
                          (input_fp8, weight_fp8),
                          kwargs={
                              "out_dtype": torch.float16,
                              "scale_a": scale_a,
                              "scale_b": scale_b
                          },
                          test_utils=("test_faketensor"))

Criteria 2

Ensure the torch.compile mode and the eager mode generates the same results

    # Run both compiled (V1 graph mode) and uncompiled versions (V1 eager mode)
    result_original = gemm_fp8_fn(input_fp8, weight_fp8, scale_a, scale_b)
    result_compiled = compiled_fp8_fn(input_fp8, weight_fp8, scale_a, scale_b)

What is AITER Linear?

AITER Linear is a kernel from the AI Tensor Engine for ROCm (AITER) that has been integrated into vLLM for the unquantized linear method and for per-tensor-weight and per-tensor-activation quantization FP8 Scaled GEMM.

How to Enable AITER Linear

• It is enabled by default when the environment variable VLLM_ROCM_USE_AITER=1 is set
• It can be specifically enabled or disabled using its dedicated environment variable VLLM_ROCM_USE_AITER_LINEAR

Out-dated-results

Performance Comparison with No-AITER

All tests are V0 Engine

Llama-3.1-8B-Instruct (with FP8 per-tensor dynamic quantization)

• With Linear only: -2% to 2% performance improvement

Llama-3.1-8B-Instruct-BF16

• With Linear only: -5% to 4% performance improvement

Llama-3.1-70B-Instruct (with FP8 per-tensor dynamic quantization)

• With Linear only: -1% to 0.5% performance improvement

Llama-3.1-70B-Instruct-BF16

• With Linear only: -0.3% to 0.2% performance improvement

Throughput Performance Comparison with No-AITER

Before PR: [Performance][ROCm] Add skinny gemms for unquantized linear on ROCm #15830

Settings:

• Model: Llama-3.1-8B-Instruct
• Cases (Input Token Length: Output Token Length):
  • 128: 128
  • 128: 2048
  • 2048: 128
  • 2048: 2048

Llama-3.1-8B-Instruct (with FP8 per-tensor dynamic quantization)

• With Linear only (Performance gain):
  • 128: 128: 1.93x
  • 128: 2048: 1x
  • 2048: 128: 1x
  • 2048: 2048: 1x

Llama-3.1-8B-Instruct-BF16

• With Linear only (Performance gain):
  • 128: 128: 1.75x
  • 128: 2048: 1x
  • 2048: 128: 1x
  • 2048: 2048: 1x

V1 Accuracy Test

Settings:

• Model: Llama-3.1-8B-Instruct

Unquantized No AITER

vllm (pretrained=meta-llama/Llama-3.1-8B-Instruct,tensor_parallel_size=1,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.7824	±	0.0114
		strict-match	5	exact_match	↑	0.7566	±	0.0118

AITER Linear unquantized

vllm (pretrained=meta-llama/Llama-3.1-8B-Instruct,tensor_parallel_size=1,max_model_len=10000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.7824	±	0.0114
		strict-match	5	exact_match	↑	0.7566	±	0.0118

AITER Linear per-tensor FP8

vllm (pretrained=meta-llama/Llama-3.1-8B-Instruct,tensor_parallel_size=1,max_model_len=10000,quantization=fp8,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.6414	±	0.0132
		strict-match	5	exact_match	↑	0.6020	±	0.0135

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

🚀

[ProExpertProg]

Could you please split up the linear between unquantized and quantized? That way we can land unquantized now and the quantized soon after, once I'm done with the FP8 scaledmm refactor.

[ProExpertProg]

I think it'd be good to specify the exact signature here

[tjtanaa]

@ProExpertProg We have fix the typo and specified the exact signature

def dispatch_unquantized_linear_func(
) -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]:
    from vllm._aiter_ops import is_rocm_aiter_linear_enabled
    if is_rocm_aiter_linear_enabled():
        return aiter_ops.rocm_aiter_tuned_gemm
    return F.linear

[mergify]

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

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

[SageMoore]

This patch generally looks fine to me. Lets rebase it and clean it up a bit. Same as the other AITER PRs, can you post lm_eval results from models that this kernel should support?

[tjtanaa]

This patch generally looks fine to me. Lets rebase it and clean it up a bit. Same as the other AITER PRs, can you post lm_eval results from models that this kernel should support?

This kernel is for generic use case. So we will pick just one model and compute its lmeval to show the correctness of the implementation which is llama3.1.

[tjtanaa]

@SageMoore @ProExpertProg We have updated the PR description with the latest information.
In this PR it will also enable the AITER kernels to be used with V1 Engine.

Should we keep the is_rocm_aiter_xxx_enabled flag in the platforms/rocm.py in the same place where use_rocm_custom_paged_attention and use_custom_allreduce are stored?

[sijiac]

cutlass_w8a8_scaled_mm(*, qinput: torch.Tensor
->
cutlass_w8a8_scaled_mm(qinput: torch.Tensor

why did we remove the starting *? is it on purpose?

[tjtanaa]

Yes this is on purpose. The input argument name of the AITER tuned gemm is not call qinput thus will break the dispatcher function logic. So removing the * and avoid specifying the first argument as keyword argument.

[mergify]

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

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

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!