[adenzhou1350]

Summary

VectorQuantizerFunction moves all three accumulation buffers to the default CUDA device, even when the input and codebook are on CPU. This makes the exported CodeBookQuantizer and a CPU QuantizedLinear using it fail during forward.

Allocate the FP32 buffers on the existing tensor device instead. Keep the centroid update, empty-cluster handling and straight-through backward unchanged. Replace the two placeholder codebook tests and add a CPU quantized-linear regression. Device subtests also exercise available CUDA devices when run on a GPU host.

Validation

• On current main 3972ed015091f659418dedf12edb980a8ca56b53, the three new CPU regressions fail at .cuda() without the production fix.

• With the fix, native package tests pass: 16 passed, 1 skipped, 2 subtests passed. The skipped test requires Faiss. Python 3.12 / PyTorch 2.12.1+cu130, CUDA hidden:

  CUDA_VISIBLE_DEVICES='' python -m pytest -q -p no:cacheprovider \
    test/prototype/module_swap_quantization/test_quantizers.py \
    test/prototype/module_swap_quantization/test_module_swap.py \
    test/prototype/module_swap_quantization/test_kmeans_codebook.py

• Ruff 0.11.6 lint/format, copyright-header check and git diff --check pass.

• The entire directory could not be collected in this source-only environment: Transformers rejects the uninstalled TorchAO version unknown in a separate test file. No CUDA runtime or Faiss initialization result is claimed.

This is a device-correctness fix for the prototype API, not a performance change.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4962

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[adenzhou1350]

@andrewor14 could you take a look? This keeps the prototype codebook buffers on the input device. The three new CPU regressions fail on main and pass with the fix, including the QuantizedLinear caller.