[executorch][cuda] Run INT6 decode linears on autotuned QuantizedGemmFamily Triton kernels and delete the int6_plain_mm C shim

[Gasoonjia]

Stack from ghstack (oldest at bottom):

• [executorch][cuda] Run INT5 decode linears on autotuned QuantizedGemmFamily Triton kernels and delete the last C shim (int5_plain_mm) #23516
• [executorch][cuda] Run INT8 decode linears on autotuned QuantizedGemmFamily Triton kernels and delete the int8_plain_mm C shim #23515
• -> [executorch][cuda] Run INT6 decode linears on autotuned QuantizedGemmFamily Triton kernels and delete the int6_plain_mm C shim #23514
• [executorch][cuda] Run INT4 decode linears on autotuned QuantizedGemmFamily Triton kernels and delete the int4_plain_mm C shim #23513
• [executorch][cuda] Add QuantizedGemmFamily: per-M quantized GEMM ops with format-owned legality and shared dispatch #23512
• [executorch][cuda] Autotune data-dependent kernels on representative inputs declared by the kernel author #23511
• [executorch][cuda] Time AOTI autotune candidates with CUDA graphs so a busy CPU does not decide the pick #23510

Decode-sized (M <= 4) INT6 linears (CudaDp4aPlanarInt6Tensor, GGUF Q6_K) move onto QuantizedGemmFamily as triton::int6_quantized_gemm_m{1,2,3,4}, the same way as INT4 in the parent diff. This removes the int6_plain_mm C shim everywhere.

Kernels, triton/kernels/int6_quantized_gemm.py. W6A8 DP4A.

• The activation is quantized once per K32 block with the shared quantize_activations_q8.
• Each output row is one warp. The planar 6-bit weights are rebuilt from ql/qh, the constant -32 offset is folded into the INT32 dot product through a DP4A activation sum, and both scale levels are applied in FP32.
• Explicit 1/2/3/4-row kernels share the rebuilt weights across activation rows. Under a dynamic M they branch uniformly to the kernel with exactly the runtime M rows.
• Autotune space: the shared generic space over two implementations (K16 or K32 per lane per trip) x rows per CTA {1, 2, 4, 8} x pipeline stages {1, 2, 3}, pruned by trip count. Split-K is the shared SM-count rule. There are no architecture or shape tables.
• Every shared piece is imported from quantized_gemm_utils.py / quantized_gemm_family.py: activation quantization, DP4A/warp-sum helpers, split-K rule and reduce, the autotune space and pruning, the legality checks, and the launch skeleton. Only the INT6 decode, the main kernels and the INT6 rules live here.
• The kernels were generated with KernelAgent from the earlier Triton INT6 bucket kernels.

Dispatch, quantize_op_dispatch/int6_dispatch.py. It uses the shared quantized_linear and chunked_dequant_linear (the INT6 dequant is now chunked along N like the others). Unsupported inputs fall back to dequant + F.linear and never raise. The int6_plain_mm schema and its Meta/CUDA impls are removed.

C shim removal (fbcode + xplat):

• runtime/shims/int6_plain_mm.{h,cu,cuh}, its gtest and its benchmark;
• the entries in runtime/targets.bzl, CMakeLists.txt and the shim tests' CMakeLists.txt;
• the INT6 cases of gen_plain_mm_test_vectors.py;
• the fallback-kernel and custom_ops_to_c_shims entries in cuda_backend.py, and the test_sort_shim expectations;
• the OSS wheel symbol list;
• the gemma4_31b Q6_K export tests now assert triton.int6_quantized_gemm_m1 in the graph.

Op level (A100; full sequence = activation quantization + GEMM (+ reduce); the shim built from the pre-diff sources; each case runs the best config of the pruned space, then shim and Triton alternate for 7 rounds and medians are compared). All five Q6_K linear shapes of the gemma4_31b, muse-glimmer and dflash GGUFs, M = 1..4, plus the 4-row bucket at a dynamic M = 2, 3:

shape (N x K)	M=1	M=2	M=3	M=4	M=2 on bucket 4	M=3 on bucket 4
gemma attn_v 4096 x 5376	1.062x	1.006x	0.993x	1.016x	1.000x	1.073x
gemma ffn_down 5376 x 21504	1.043x	1.127x	1.445x	1.173x	1.104x	1.471x
muse attn_v 256 x 6656	1.173x	1.221x	1.298x	1.380x	1.184x	1.269x
muse ffn_down 6656 x 19968	1.065x	1.117x	1.303x	1.052x	1.087x	1.342x
dflash attn_v 1024 x 6656	1.019x	1.087x	1.157x	1.067x	1.084x	1.206x

Geomean 1.134x for static M and 1.174x for dynamic M. The one case below 1.0x is gemma attn_v at M = 3: 0.993x here, and a 0.996x tie in KernelAgent's own measurement. Mean relative difference vs the shim is < 0.0005 everywhere.

SEE_AB_TABLE

Differential Revision: D123644066

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23514

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 2 Cancelled Jobs, 1 Unrelated Failure

As of commit 6907db4 with merge base 8789aa5 ():

NEW FAILURE - The following job has failed:

• Test ROCm AOTI / unittest-rocm-gfx950-rocm7.2 / linux-job (gh)
  backends/cuda/tests/test_sdpa_splitk_replacement.py::TestUnsupportedSDPAStaysWithRegularLowering::test_float32_inputs_are_not_replaced

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• pull / test-qnn-python-imports-linux / linux-job (gh)
• Test CUDA Builds / unittest-cuda / linux-job (gh)
  [runner-container-hooks] FATAL: received SIGINT

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

• trunk / test-samsung-models-linux / linux-job (gh) (#23507)
  test_yolo26_fp16

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.