[executorch][cuda] Run INT4 decode linears on autotuned QuantizedGemmFamily Triton kernels and delete the int4_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) INT4 linears (CudaCoalescedInt4Tensor) move onto QuantizedGemmFamily (parent diff) as triton::int4_quantized_gemm_m{1,2,3,4}. This removes the int4_plain_mm C shim everywhere.

Shared pieces, triton/kernels/quantized_gemm_utils.py. These are written once for INT4/5/6/8 and moved here from the INT4 kernels:

• the inline-PTX helpers (DP4A, warp sum, round-to-nearest-even);
• the BF16 -> INT8 K32 activation quantization kernel and its launcher;
• the deterministic split-K reduce;
• split_k_for, the device's-SM-count rule, now rounded to a power of two under the workspace bound;
• the generic autotune space autotune_configs(implementations): implementations x rows per CTA {1, 2, 4, 8} (= num_warps) x pipeline stages s in {1, 2, 3} (PIPELINE_STAGES = num_stages = s);
• prune_by_main_loop_trips for prune_configs_by (drops stage counts above the main loop's trip count);
• the legality checks the formats compose.

quantized_gemm_family.launch_split_k_gemm is the shared launch skeleton: output, bucket-sized split-K workspace, grid, and reduce.

INT4, triton/kernels/int4_quantized_gemm.py. It keeps only the INT4 kernels and rules.

• _unsupported_reason covers every rule the old _validate_inputs raised on, plus static K, K % 256, gs == 32, the M range (static M == bucket, or a dynamic M provably in [1, bucket]), and CUDA-or-fake. The family derives supports and validate from it; every op validates before launching.
• Each bucket's triton.autotune covers the generic space over three implementations: the generic row-blocked kernel, the explicit per-row-accumulator kernel with the bucket's rows, and the next-larger explicit kernel. The extra row is masked, but its code is faster at M = 3 with large K. Configs are pruned by trip count (AOTInductor honors prune_configs_by on user kernels).
• Under a dynamic M, an explicit kernel branches uniformly on the runtime M to the explicit kernel with exactly M rows.

Dispatch, quantize_op_dispatch/int4_dispatch.py. It uses the shared quantized_linear and chunked_dequant_linear; INT4 provides only its single-chunk dequantization. Unsupported inputs fall back to dequant + F.linear and never raise. The int4_plain_mm schema and its Meta/CUDA impls are removed.

C shim removal (fbcode + xplat):

• runtime/shims/int4_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 INT4 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.

_weight_int4pack_mm is a different op and is untouched. aoti_cuda_shims.lib (Windows) is left as is; its extra symbol is unreferenced.

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 16 INT4 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:

	geomean vs shim	worst	cases slower than the shim
static M (64 cases)	1.111x	0.961x	7, by 0.5-4%: M=1 K>=16384 (gemma attn_out 16384, ffn_down; 1-2%), M=1 dflash attn_k (1.7%), M=2 gemma qkv (0.5%), M=3 gemma attn_out 8192 (1.4%), M=3 dflash fc K=33280 (3.9%)
dynamic M (32 cases)	1.058x	0.911x	7: mostly M=3 on bucket 4 at K>=8192 (3-9%), M=2 muse q/gate (7%)

No device-agnostic split-K rule closes the remaining cases: a sweep of split rules shows only a per-shape split would. They are left as is; the model A/B below is the check.

SEE_AB_TABLE

Differential Revision: D123569750

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 130 Pending

As of commit 7e8cbe6 with merge base fc9a1c6 ():

NEW FAILURE - The following job has failed:

• pull / test-arm-backend-no-driver (test_pytest_ops_no_target) / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

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

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.