Fix Int4OpaqueTensor A16W4 crash on mismatched activation dtype

[tritsystem]

Summary

Int4OpaqueTensor's A16W4 (weight-only) docstring documents support for
"float16/bfloat16/float32 activation + int4 weight", implying any of these
activation dtypes works regardless of the dtype the weight happened to be
in at quantization time. In practice, the aten.linear.default forward
implementation (torchao/prototype/quantization/int4/int4_opaque_tensor.py)
passed the activation straight into
torch.ops.aten._weight_int4pack_mm_for_cpu without casting it to match
scale_and_zero.dtype. That raw aten op requires an exact dtype match with
no implicit promotion at this boundary (unlike plain elementwise ops on
current PyTorch, which silently type-promote), so any mismatched
(weight_dtype, activation_dtype) pair raises:

RuntimeError: expected scalar type X but found Y

Mechanism

act_mat = input_tensor                      # whatever dtype the caller passes
...
act_mat = act_mat.reshape(-1, act_mat.shape[-1])
y = torch.ops.aten._weight_int4pack_mm_for_cpu(
    act_mat.contiguous(), packed_weight, groupsize, scale_and_zero
)

scale_and_zero is fixed at whatever dtype weight was quantized from
(Int4OpaqueTensor.from_hp), e.g. float32 if the nn.Linear wasn't cast
before quantize_() was called. If the caller later runs inference with a
different float dtype (a very normal thing to do — e.g. quantize a model
as-loaded, then call it with a bfloat16 activation), the op call above
crashes with a cryptic dtype error that gives no hint the real cause is a
weight/activation dtype mismatch, contradicting the class's own documented
contract.

The sibling CUDA tensor class, Int4TilePackedTo4dTensor, already handles
this correctly by casting its activation to its one supported dtype
(bfloat16) before calling its own tinygemm op. This PR applies the same
cast-before-op pattern to Int4OpaqueTensor, casting to scale_and_zero.dtype
(since this class supports three activation dtypes, not just one).

Verification

Directly exercised Int4OpaqueTensor.from_hp() + F.linear() for the full
3x3 (weight_dtype, act_dtype) grid over {float32, bfloat16, float16}:

• Before the fix: only the 3 matching-dtype pairs succeed; all 6
  mismatched pairs raise RuntimeError: expected scalar type X but found Y.
• After the fix: all 9 pairs succeed and return output in the caller's
  activation dtype.

Added test_linear_mismatched_activation_dtype to
test/prototype/test_int4_opaque_tensor.py, parametrized over all 9
(weight_dtype, act_dtype) pairs x both use_hqq settings (18 cases).

Red/green confirmed by reverting only the source fix: exactly the 12
mismatched-dtype cases fail (6 pairs x 2 algorithms); the 6 matching-dtype
cases pass regardless. With the fix, all 18 pass.

Ran the full existing test_int4_opaque_tensor.py suite before and after
this change: 36 pre-existing failures are present identically on unmodified
main (appear to be torch.compile-related in this local environment —
Python 3.14 / torch 2.13 CPU build — unrelated to this change) with no new
failures introduced by this fix.

Ran ruff check and ruff format on both changed files — clean.

Verification limitation

Could not test the from_hp_da8w4 (DA8W4, int8-activation) path, or any
CUDA/XPU/NPU int4 tensor class, since this change only touches the
CPU-only A16W4 path (_weight_int4pack_mm_for_cpu). This machine has no
CUDA-enabled PyTorch build installed, so GPU-only code paths in this repo
could not be exercised at all.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.