[harshaladhav-amd]

Purpose

Fixes #61964.

On AMX CPUs, CPUInt8ScaledMMLinearKernel routes symmetric W8A8 linears with aligned shapes to _apply_weights_sgl (the SGLang-derived int8_scaled_mm_with_quant kernel), which has been the default since #50801. That path passes activations straight to the kernel, and the kernel enforces CHECK_DIM(2, mat1). The Whisper encoder runs on B x T x D, so W8A8 Whisper fails on AMX machines with RuntimeError: mat1 must be a 2D tensor. Non-AMX CPUs fall back to _apply_weights_onednn, which already flattens to [M, K] and restores the leading dims, so the model works there.

This PR applies the same flatten/reshape in _apply_weights_sgl. It matches _apply_weights_onednn and CPUFp8BlockScaledMMKernel in the same file, and the W4A16 fix in #58268. 2-D inputs (all decoder-only models) are unaffected.

Not a duplicate: I searched open PRs and issues for _apply_weights_sgl, int8_scaled_mm_with_quant, and CPU/Whisper W8A8. The only related item is #56029, which fixes the ROCm Triton scaled_mm path (#56024) and doesn't touch the CPU kernel.

Test Plan

• New unit test test_cpu_sgl_int8_linear_accepts_3d_input in tests/kernels/quantization/test_scaled_mm_kernel_selection.py (already a cpu_test file). It builds a layer through CPUInt8ScaledMMLinearKernel.process_weights_for_sgl, runs apply_weights on a [2, 3, K] input, and checks the result is bit-identical to the flattened [6, K] call reshaped back. The test is skipped when int8_scaled_mm_with_quant is not built.
• pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py -v
• pre-commit run --files <changed files>

Test Result

Run on an AVX512-BF16 host (AMD EPYC 9454). It loads the same _C library, so the SGL int8 op runs there through its non-AMX fallback.

• Without the fix, the new test fails with RuntimeError: mat1 must be a 2D tensor.
• With the fix:

tests/kernels/quantization/test_scaled_mm_kernel_selection.py  11 passed

• pre-commit: all hooks passed (ruff check, ruff format, typos, mypy, SPDX, test tethering, ...).

I don't currently have AMX hardware available, so I haven't run an end-to-end W8A8 Whisper accuracy check on GNR/SPR. The change only affects inputs that previously raised, and the 2-D path is unchanged. I'd appreciate an AMX run from CPU CI or reviewers (cc @bigPYJ1151).

AI assistance was used to prepare this change. I reviewed every changed line and ran the tests above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[harshaladhav-amd]

@bigPYJ1151 could you please review this PR? It fixes the AMX int8 W8A8 path (_apply_weights_sgl) rejecting 3-D activations, which breaks W8A8 Whisper on AMX CPUs (#61964). Thanks!

[bigPYJ1151]

Thanks, LGTM : )

[bigPYJ1151]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95817 for commit bc093dee200c.

ℹ️ This PR is 5 commits behind upstream main. Merged upstream main into this branch as bc093dee200c before starting CI. Pull the branch before pushing more changes.