[TPU][Quantization] TPU W8A8

[robertgshaw2-neuralmagic]

SUMMARY:

• support TPU for compressed-tensors w8a8 models.
• To run, just load a W8A8 model:

from vllm import LLM
model = LLM("neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8", max_model_len=2048)
model.generate("Hello my name is")

TESTING:

• verified accuracy on TPU for Llama-8B on TP=1 (exact score as GPU)
• verified accuracy on TPU for Llama-8B on TP=4 (exact score as GPU)
• verified accuracy on TPU for Llama-70B on TP=1 (exact score as GPU)
• verified accuracy on TPU for Qwen on TP=1 (exact score as GPU) --- note: bias in model
• confirmed all schemes still work on GPU, including:
  • nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Dynamic-Asym
  • nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Sym
  • nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym
• add Llama TP=1 tests to CI/CD
  • FOLLOW UP: Add more than one model once we enable lm-eval framework on TPU
  • FOLLOW UP: add TP>1 once we enable this machine type in the CI
• figure out workaround for user warning re: cond

FOLLOW UP

• [TPU] Mixed precision
• [TPU] Estimated memory usage is elevated due to peak_bytes capturing some intermediate tensors, fix it.
• [Software Quality] Add TritonScaledMMLinear abstraction
• [Software Quality] Convert Fp8 methods to use Kernel abstraction

[robertgshaw2-neuralmagic]

@mgoin this is ready to go.

[robertgshaw2-neuralmagic]

NOTE for reviewer - this file is not changed, it is just moved

[mgoin]

LGTM, excellent work