[TPU] Add Load-time W8A16 quantization for TPU Backend

[lsy323]

Add Load-time W8A16 quantization for TPU Backend. The workflow is similar to the existing load-time fp8 quantization. Open the PR to help discussion process.

• Added a new quantization type tpu_int8 for load-time int8 weight only quantization for tpu Backend. (e.g. LLM(model="google/gemma-2b", quantization="tpu_int8")
• Added TPUInt8LinearMethod which quantizes bfloat16 weights to int8 weights for linear layers, and calls TPU quantized ops in forward.

[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 consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[mgoin]

Interesting, this looks reasonable to me. The important note is to lazily import the torch xla function when needed, rather than at the top of the quant file.

[robertgshaw2-neuralmagic]

Super cool!

As a follow up, we can work on hooking this up to some of the existing checkpoints we have in addition to inplace quantization

[robertgshaw2-neuralmagic]

By chance, what schemes does the following support:

• torch.ops.xla.quantized_matmul(x, weight, scale)

Channelwise?
Activations?

[lsy323]

Hi @mgoin, @robertgshaw2-neuralmagic,

Thank your for reviewing my PR! Excited to work with you to enable quantization for TPU backend through compressed-tensors!

By chance, what schemes does the following support:

• torch.ops.xla.quantized_matmul(x, weight, scale)

Channelwise? Activations?

We have the quantized ops (Equivalent to the quantized cuda kernels in vLLM, but for TPU) in PyTorch/XLA here. The quantized matmul kernel is registered as a torch op and is compatible with torch.compile, it can be configured to support per-channel/blockwise quantization, both int8 and int4 are supported (int4 is not optimized now) The quantized matmul supporting matrix is here.

[robertgshaw2-neuralmagic]

Hi @mgoin, @robertgshaw2-neuralmagic,

Thank your for reviewing my PR! Excited to work with you to enable quantization for TPU backend through compressed-tensors!

By chance, what schemes does the following support:

• torch.ops.xla.quantized_matmul(x, weight, scale)

Channelwise? Activations?

We have the quantized ops (Equivalent to the quantized cuda kernels in vLLM, but for TPU) in PyTorch/XLA here. The quantized matmul kernel is registered as a torch op and is compatible with torch.compile, it can be configured to support per-channel/blockwise quantization, both int8 and int4 are supported (int4 is not optimial) The quantized matmul supporting matrix is here.

This is so awesome!!!! Running the same compressed models on various hardware backends is going to be an awesome feature

[WoosukKwon]

@lsy323 Thanks for the PR! I'm really looking forward to using this feature!

I think we have two things to figure out on this PR:

1. Where to put this quantization config and the linear method? Do we want to put this as a new quantization config (like in the current PR) or in compressed-tensors?
2. IIRC, this currently does not support the case when the BF16 weights exceed the HBM size of the TPU while INT8 weights do (e.g., Llama 8B on TPUv5e which has 16 GB HBM). Could you please remind us of why this isn't supported?

[miladm]

TorchXLA:TPU FP8 support is WIP (partially supported). @lsy323 do we have an outlined plan somewhere that extends this effort to FP8?

[lsy323]

1. Where to put this quantization config and the linear method? Do we want to put this as a new quantization config (like in the current PR) or in compressed-tensors?

I think for compressed-tensors config, we assume all the checkpoints are in compressed-tensors format. This load-time quantization doesn't seem to belong that flow, hence I think keeping it in a separate file looks cleaner.

2. IIRC, this currently does not support the case when the BF16 weights exceed the HBM size of the TPU while INT8 weights do (e.g., Llama 8B on TPUv5e which has 16 GB HBM). Could you please remind us of why this isn't supported?

Sure, the current flow is:

1. move bfloat16 weights from host to TPU
2. Quantize bfloat16 weights to int8

When the BF16 weights exceed the HBM size of the TPU, step 1 would hit OOM. To avoid this problem, we can delay weight transferring if load-time quantization is enabled.

[lsy323]

TorchXLA:TPU FP8 support is WIP (partially supported). @lsy323 do we have an outlined plan somewhere that extends this effort to FP8?

I don't have a crystal plan for this, alternatives are as follows:

1. Reuse fp8.py which works for CUDA workflow. (Support both load-time fp8 quantization and offline quantized fp8 ckpt using TensorRT-LLM ref)
2. Have another quantizaiotn config for fp8 on TPU (e.g. fp8_tpu)
3. Extend compressed_tensors to support TPU backend (This will support fp8 ckpt in compressed_tensors format)

[robertgshaw2-neuralmagic]

TorchXLA:TPU FP8 support is WIP (partially supported). @lsy323 do we have an outlined plan somewhere that extends this effort to FP8?

I don't have a crystal plan for this, alternatives are as follows:

1. Reuse fp8.py which works for CUDA workflow. (Support both load-time fp8 quantization and offline quantized fp8 ckpt using TensorRT-LLM ref)
2. Have another quantizaiotn config for fp8 on TPU (e.g. fp8_tpu)
3. Extend compressed_tensors to support TPU backend (This will support fp8 ckpt in compressed_tensors format)

Hey guys - there are a couple considerations here. For vLLM, we want to support both cases:

• In-place quantization
• Pre-quantized checkpoints

We will be making all go forward checkpoints inside the compressed-tensors integration for mixed precision, integer activation quantization, and floating point activation quantization. And so, I think we should focus on this pathway.

Both fp8.py and compressed-tensors share the same backend code (you can see they both use apply_fp8_linear. We factored this utility out, such that the kernel calls are shared by the various integrations. If you add the TPU calls to this function, then you should get all these integrations "for free"

[lsy323]

2. IIRC, this currently does not support the case when the BF16 weights exceed the HBM size of the TPU while INT8 weights do (e.g., Llama 8B on TPUv5e which has 16 GB HBM). Could you please remind us of why this isn't supported?

hi @WoosukKwon, I looked into this in detail, it doesn't look like to be a straightforward change, I think we can consider support that in a separate PR.

In the current flow, weights are moved to device as the model is initialized (ref), then load time quant will be done on device ref. We need to introduce a new flow to support this case.

[WoosukKwon]

@lsy323 Seems like my previous comment was not addressed for some reason. Can you please check it again?

[lsy323]

@lsy323 Seems like my previous comment was not addressed for some reason. Can you please check it again?

@WoosukKwon Somehow I force pushed without the suggested change commits. Now should be fixed. Thank you for reminding!

[WoosukKwon]

LGTM! Thanks for the PR! It works well on my machine 🎉 🎉

Looking forward to the next step! (adding INT8 activation quantization in tpu-int8).