Fix N-D input handling in default backend int8 ops

[2sumtech]

Summary

The default backend implementations of int8_mm_dequant, int8_vectorwise_quant and int8_vectorwise_dequant assume a 2-D input. This makes them disagree with their own registered fake kernels on any input with more than two dimensions.

• int8_mm_dequant flattens the leading dimensions and returns the flattened result, so a 3-D input silently comes back 2-D
• int8_vectorwise_quant reduces over dim=1 instead of the last dimension and raises a RuntimeError
• int8_vectorwise_dequant broadcasts the row scales against the wrong dimension and raises a RuntimeError

Root cause

All three fake kernels treat the leading dimensions as rows and preserve the shape of A. The CUDA kernels for the first two do the same, allocating with torch.empty_like(A) / torch.empty(A.shape) and passing A.numel() // A.shape[-1] as the row count. Only the default implementations flatten, so the disagreement shows up on CPU, MPS and XPU but not on CUDA. torch.library.opcheck fails its FakeTensor comparison for a 3-D input.

int8_scaled_mm inherits the problem through int8_mm_dequant, returning a 2-D result where the contract promises 3-D.

Fix

Flatten the leading dimensions for the computation and restore the original shape on the way out. int8_mm_dequant also switches view to reshape so a non-contiguous input does not throw.

Impact

Output for 2-D inputs is bit for bit unchanged, verified by hashing all three ops over several shapes on CPU and MPS against main. 1-D, 3-D and 4-D inputs now match the fake kernel, opcheck passes, and torch.compile with fullgraph=True works over a 3-D int8_mm_dequant.

No new tests added. The existing test_int8_mm_dequant already asserts shape and runs opcheck, and the current suites cover the change. Happy to parametrize a 3-D shape into it if you want the case pinned.

Tests

• pytest tests/test_ops.py -k int8 (325 passed, 21 skipped)
• pytest tests/test_functional.py -k "int8 or igemmlt or coo" (49 passed, 5 skipped)
• pytest tests/test_linear8bitlt.py tests/test_autograd.py (1640 passed)
• pre-commit ruff and ruff-format clean on both changed files

Tested on Apple Silicon, so CPU and MPS both exercise the default kernels. I do not have CUDA hardware, and the CUDA paths are untouched.

---

Drafted with AI assistance. All analysis, the fix and the test runs were verified locally by me.