[gemma4_31b][cuda] Export Gemma4-31B @128k on 5090

[Gasoonjia]

Current gemma4-31b can not be successfully exported on consumer gpu like 5090 with three reasons:

1. During int4_dispatch we need to dequant whole matmul weight to bf16 for prefill in one step for lm_head, leading to weight duplcation;
2. When lowering to AOTI-CUDA, we moved the whole model, including kv cache, onto gpu. With context length increased, the gpu memory consumption will also be increased dramatically.
3. No autotune config for kernels like sdpa work for consumer gpu like 5090.

Three CUDA-export memory optimizations, all gated behind the existing low_memory_mode compile spec (no impact on other models or on runtime):

• int4_dispatch: chunk the inline _dequant_matmul along N for vocab-sized weights, gated behind a low-memory flag with an N>65536 threshold so only the lm_head crosses it. Avoids transiently materializing the full ~10 GiB bf16 lm_head during AOTI autotune / cpp_wrapper. The prefill MLP path is untouched -> zero runtime impact.

• cuda_backend / aoti_backend: skip occupying the GPU with the KV-cache buffers during AOTI compile. A new move_program_to_device hook places KV constants on the target device but immediately frees their storage (resize_(0)), so the fake-tensor device check passes while no real KV bytes sit on the GPU during autotune. The emptied buffers are re-synthesized as zeros at the unlift_graph clone and at serialization, and excluded from constant dedup (resize(0) gives every KV data_ptr 0, which would otherwise collapse same-shape caches across layers). All gated behind low_memory_mode.

• tq4_sdpa: add BLOCK_N=16 (and a BLOCK_M=32) autotune config. The superset is kept for big-shared-memory GPUs (A100/H100); the Triton autotuner auto-prunes configs that exceed a GPU's shared memory (OutOfResources -> inf), so the same config list also works on the 5090 (Blackwell, ~101 KB SMEM) where the previous smallest config did not fit.

Full Gemma4-31B on 128k TQ export: peak 28.0 GiB, runtime output correct ("...Paris.").

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

❌ 3 New Failures, 3 Unrelated Failures

As of commit 993cff5 with merge base 1b726b2 ():

NEW FAILURES - The following jobs have failed:

• pull / test-arm-backend-no-driver (test_pytest_ops_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t 240605ac100851bf3f29b75b801e35308a3562f7979541483ce6d8133c30e0d5 /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 485e4fcfd099959114159bdfd57ae874b14a1aa5dbb354818ebb69da8cb8ed06 /exec failed with exit code 1
• pull / unittest-editable / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 66f95c521d611bd0dd325b49e963b2632bd6467386ee969ca66a6240c99ca7ad /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh) (trunk failure)
• pull / unittest / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

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.

[digantdesai]

Aren't these kind of device specific?

[digantdesai]

I wish there is a better way to do this i.e. why does this logic needs to be aware of export issues?

[metascroy]

Is this just a crude way of doing tile level dequant?