[executorch][muse-glimmer] Run DFlash sampling on CUDA

[Gasoonjia]

Stack from ghstack (oldest at bottom):

• -> [executorch][muse-glimmer] Run DFlash sampling on CUDA #22492
• [executorch][muse-glimmer] Add CUDA speculative sampling orchestration #22491
• [executorch][muse-glimmer] Add unified CUDA sample_token entry point #22490
• [executorch][muse-glimmer] Add CUDA residual correction sampling #22489
• [executorch][muse-glimmer] Add CUDA excluded-token correction #22488
• [executorch][muse-glimmer] Add CUDA speculative acceptance #22487
• [executorch][muse-glimmer] Add CUDA categorical sampling #22486
• [executorch][muse-glimmer] Add CUDA sampling probabilities #22485
• [executorch][muse-glimmer] Add CUDA argmax sampling primitive #22484

Keep the CUDA DFlash decode loop device-resident instead of round-tripping activations and logits through the host.

Export embed_text, target_forward_from_embeddings, and target_prefill_from_embeddings device-in/device-out, and keep draft_forward logits on device. Draft sampling, target sampling, speculative acceptance, and correction then run through the CUDA sampling primitives on the model's shared stream, so a verification step never touches host memory. The candidate, probability, position-table, and token-staging buffers are preallocated before CUDA graph capture so replay stays valid.

A CUDA build always samples on device; there is no runtime switch. The host sampler is reached only by non-CUDA (MLX) builds, so the host speculative cycle, run_target_only_cycle, and the run_target token wrapper are guarded with #ifndef EXECUTORCH_BUILD_CUDA rather than left as unreachable code behind an unconditional return.

The verifier reuses the existing embed_text + target_forward_from_embeddings chain rather than adding a token-input target method: candidate tokens already live on the device, so embed_text consumes them in place and hands its embeddings straight to the target with no copy. The remaining host callers (prefill and the vision splice) stage their own inputs into the preallocated device buffers. Keeping a single target implementation means the CUDA and MLX paths decode through the same method, and it removes both the duplicate complete-target weight blob and the export-time blob surgery that used to omit it -- surgery that depended on runtime method load order and on the AOTI delegate payload's line shape.

Sampling stays a runtime primitive: no sampling method is added to the .pte, and the DFlash timing mechanism is unchanged.

Differential Revision: D117826290

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

As of commit c1a4838 with merge base c65ad53 ():

NEW FAILURE - The following job has failed:

• Lint / lintrunner-mypy (gh)
  >>> Lint for backends/arm/tosa/partitioner.py:

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.