[TPU] Add support for online w8a8 quantization

[kyuyeunk]

Expand existing online w8a16 quantization to support w8a8

• Using {'activation_scheme': 'dynamic'} utilizes dynamic activation quantization
• Using {'activation_scheme': 'none'} utilizes weight only quantization

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

🚀

[gemini-code-assist]

Code Review

This pull request adds support for w8a8 quantization on TPUs by introducing a 'dynamic' activation scheme. The changes correctly plumb this new option down to a new torch_xla operator. My review focuses on improving the error handling for cases where an incompatible version of torch_xla is installed, to provide a better user experience.

[gemini-code-assist]

The current try...except block only catches an ImportError if torch_xla or the custom_kernel module is missing. However, it doesn't handle the case where torch_xla is installed but is an older version that doesn't have the quantized_matmul_int8 operator. In that scenario, an AttributeError will be raised later when the operator is called, which can be confusing for the user.

To provide a better user experience and a more informative error message, it's best to check for the operator's existence within the try block and catch both ImportError and AttributeError. This ensures that users with an incompatible torch_xla version get a clear message on how to resolve the issue.

        try:
            import torch_xla.experimental.custom_kernel  # noqa: F401
            # Eagerly check for the op to provide a better error message.
            _ = torch.ops.xla.quantized_matmul_int8
        except (ImportError, AttributeError) as err:
            raise ImportError(
                "torch_xla is not installed or is too old to support w8a8 "
                "quantization. Please install/update torch_xla by following "
                "the instructions at "
                "https://docs.vllm.ai/en/latest/getting_started/tpu-installation.html "  # noqa: E501
                "to run vLLM on TPU.") from err

[yaochengji]

Thanks for the contribution, @kyuyeunk .

Could you also add a test for this feature?

[kyuyeunk]

Thanks for the contribution, @kyuyeunk .

Could you also add a test for this feature?

done. @yaochengji , can you take a look again?

[yaochengji]

Thanks for adding the test!

Could move it to tests/v1/tpu?
And then add the test in CI .buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh?

Also could we also check the response, at least to make sure it doesn't generate garbage?

[kyuyeunk]

Thanks for adding the test!

Could move it to tests/v1/tpu? And then add the test in CI .buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh?

Also could we also check the response, at least to make sure it doesn't generate garbage?

Done. And verified that it produces correct response.

• prompt: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\n'
• response: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\nThe core of the engine is a distributed memory system'

[yaochengji]

prompt: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\n'
response: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\nThe core of the engine is a distributed memory system'

@kyuyeunk Could you add the response check in the test code?

[kyuyeunk]

prompt: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\n'
response: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\nThe core of the engine is a distributed memory system'

@kyuyeunk Could you add the response check in the test code?

Do you have a suggested to achieve this? Not sure what is the best automated approach to check if an output is garbage or not.

[kyuyeunk]

prompt: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\n'
response: 'vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.\nThe core of the engine is a distributed memory system'

@kyuyeunk Could you add the response check in the test code?

Do you have a suggested to achieve this? Not sure what is the best automated approach to check if an output is garbage or not.

Done!

[yaochengji]

LGTM, thanks!