Add w8a8 quantized matmul kernel #9278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Merged

vanbasten23 merged 10 commits into master from xiowei/add_w8a8_quantized_matmul_kernel_v5e

Jun 4, 2025

Collaborator

vanbasten23 commented Jun 2, 2025 •

edited

Loading

This PR adds:

a Pallas kernel for doing w8a8 quantized matmul
tuned table for v6e

Test plan

pip install pytest
pytest pytorch/xla/test/test_quantized_matmul_pallas_kernel.py -s

vanbasten23 added 4 commits

June 2, 2025 16:57


          added the tests and the quantized matmul kernel. All test passed.

b6f5d4b


          added tuned table and unit test

f1c96aa


          fix the test

3609fa7


          fix a test

b0d3aa4

vanbasten23 requested review from yaochengji and lsy323

June 2, 2025 18:47

vanbasten23 added 2 commits

June 2, 2025 19:04


          convert scalar type to f32

6adfdd9


          fix linter

ad184f4

yaochengji reviewed

View reviewed changes

Collaborator

yaochengji left a comment

Thanks Xiongfei for your contribution! Left a few comments.

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Outdated Show resolved Hide resolved

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Outdated Show resolved Hide resolved

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Show resolved Hide resolved

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Show resolved Hide resolved

test/test_quantized_matmul_pallas_kernel.py Outdated Show resolved Hide resolved

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Show resolved Hide resolved

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Show resolved Hide resolved

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Show resolved Hide resolved

vanbasten23 added 3 commits

June 3, 2025 22:33


          add vmem_limit_bytes to static_argnames

f699c0c


          removed int4_weight

317c728


          fix comments

6836bae

lsy323 reviewed

View reviewed changes

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Outdated Show resolved Hide resolved

lsy323 reviewed

View reviewed changes

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Show resolved Hide resolved

lsy323 reviewed

View reviewed changes

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Show resolved Hide resolved

lsy323 reviewed

View reviewed changes

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Outdated Show resolved Hide resolved

lsy323 reviewed

View reviewed changes

torch_xla/experimental/pallas_kernels/quantized_matmul_kernel.py Outdated Show resolved Hide resolved

vanbasten23 requested review from yaochengji and lsy323

June 4, 2025 16:30


          fix comments

a48e36d

yaochengji approved these changes

View reviewed changes

Collaborator

yaochengji left a comment

LGTM, thanks!

Collaborator

lsy323 commented Jun 4, 2025

Quantization logic looks good to me!

lsy323 approved these changes

View reviewed changes

Collaborator Author

vanbasten23 commented Jun 4, 2025

Thanks for the review!

vanbasten23 merged commit 2504888 into master

22 checks passed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet