[Jiajia Qin (qjia7)]

Description

Extend the existing K-parallel DP4A dispatch to narrow INT2/4/8 projections (N <= 32, K >= 2048), with FP32 accumulation of their partial sums. Allow projections to share activation quantization and subgroup-matrix input prepacking within one invocation.

The fixed 64×64 DP4AMatMulNBitsProgram, its dispatch, and its WGSL shader are unchanged by this PR.

Motivation and Context

Extract reusable kernel changes from #33370 so its QLoRA operator and fused shaders can be reviewed separately. #33370 will be rebased after this merges.

Performance

RTX 5080, Windows, block size 32, graph capture enabled. GPU timestamp medians include preparation and multiplication (two rounds, 100 replays each, warm-up samples discarded):

| M × K × N | Weights / activations | D3D12, main → PR | Vulkan, main → PR |

|---|---|---:|---:|

| 128 × 2048 × 32 | INT8 / FP16 | 47 → 10 µs | 52 → 14 µs |

| 512 × 4096 × 32 | INT8 / FP16 | 96 → 36 µs | 97.5 → 32 µs |

| 128 × 2048 × 4096 | INT4 / FP16 | 84.5 → 84 µs | 84 → 85 µs |

| 128 × 2048 × 4096 | INT4 / FP32 | 74 → 74 µs | 73.5 → 73 µs |

These measurements cover isolated GPU operators; timestamp resolution is 1 µs.

Validation

• Windows Release build with built-in WebGPU, D3D12 and Vulkan; lint and git diff --check passed.

• Provider regressions: 15 passed, five skipped (four Arm64/KleidiAI cases and subgroup-matrix reuse on D3D12). All 28 QKV/MLP checks passed.

• All three shared-input tests passed with Vulkan, including subgroup-matrix prepacking.

• Retained INT2/4/8 tests cover K-parallel dispatch thresholds, K tails, 16/32-column outputs, batching, zero points, bias, and FP32 inputs. Shared-input tests also cover FP16 when supported by the device. Tests specific to the removed tile variants are excluded.

• Dedicated FP16 INT2/4/8 regressions cover cancellation after partial sums exceed the FP16 range, with the general FP32-accumulation option disabled. They check device capabilities and disable CPU fallback. All fail when the narrow-path FP32 override is removed and pass with it restored.

[Copilot]

🟡 Changes recommended

The new INT4 and INT8 regressions silently execute no cases in combined CUDA/WebGPU builds, leaving their intended coverage absent.

2 open issues

• Ensure combined CUDA/WebGPU tests validate WebGPU output · New
• Ensure 8-bit tests exercise data types on WebGPU · New

What changed in this PR

This PR extracts reusable WebGPU MatMulNBits optimizations from #33370, separating shared input preparation and narrow-projection tuning from the QLoRA implementation.

Changes:

• Adds invocation-local reuse of activation quantization and subgroup-matrix prepacking.
• Tunes narrow DP4A projections with rectangular tiles, wider INT8 loads, and K-parallel reduction while preserving full-width loading.
• Adds shared-input and INT2/4/8 regression tests covering tails, long K, zero points, and bias.

File	Description
onnxruntime/​test/​providers/​webgpu/​matmul_nbits_input_preparation_test.cc	Tests shared versus fresh input preparation.
onnxruntime/​test/​contrib_ops/​matmul_8bits_test.cc	Adds narrow INT8 projection regressions.
onnxruntime/​test/​contrib_ops/​matmul_4bits_test.cc	Adds narrow INT4 regressions and updates a shader-name comment.
onnxruntime/​test/​contrib_ops/​matmul_2bits_test.cc	Adds narrow INT2 regressions and FP16 WebGPU execution.
onnxruntime/​contrib_ops/​webgpu/​quantization/​subgroup_matrix_matmul_nbits.h	Declares reusable prepacking interfaces.
onnxruntime/​contrib_ops/​webgpu/​quantization/​subgroup_matrix_matmul_nbits.cc	Extracts tiling selection and accepts prepacked activations.
onnxruntime/​contrib_ops/​webgpu/​quantization/​matmul_nbits.h	Adds an optional quantized-input workspace parameter.
onnxruntime/​contrib_ops/​webgpu/​quantization/​matmul_nbits.cc	Forwards the workspace to DP4A execution.
onnxruntime/​contrib_ops/​webgpu/​quantization/​dp4a_matmul.wgsl.template	Supports rectangular tiles and wider K loads.
onnxruntime/​contrib_ops/​webgpu/​quantization/​dp4a_matmul_nbits.h	Defines shared workspace and configurable DP4A programs.
onnxruntime/​contrib_ops/​webgpu/​quantization/​dp4a_matmul_nbits.cc	Reuses quantization and selects narrow-projection strategies.
onnxruntime/​contrib_ops/​webgpu/​quantization/​dp4a_matmul_k_parallel.wgsl.template	Renames the reduction shader and bounds scale loading.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The new regression tests need FP16 capability guards to avoid failing on supported FP32-only WebGPU adapters.

1 open issue

• Gate FP16 tests on ShaderF16 support · New

2 resolved since last review

• Ensure 8-bit tests exercise data types on WebGPU
• Ensure combined CUDA/WebGPU tests validate WebGPU output

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

The forced FP32 accumulation for narrow FP16 outputs needs a numerical regression that distinguishes it from FP16 accumulation.

1 open issue

• Add FP16 large-K FP32 accumulation regression test · New

1 resolved since last review

• Gate FP16 tests on ShaderF16 support

_{Review effort: Balanced}