Add GPU-side Gumbel-max sampling for CUDA graph compatibility by Gasoonjia · Pull Request #18844 · pytorch/executorch

Gasoonjia · 2026-04-13T18:13:04Z

This PR replaces cpu sampler with CUDA sampler and fuse sampler with forward method to both eliminate unnecessary data transfer and improve sampling efficient. Decode performance increases from 113.8 token/s to 119.5 token/s

Once we land the device support pipeline, we should decompose the forward method with sampling.

pytorch-bot · 2026-04-13T18:13:09Z

🔗 Helpful Links

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

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

Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 6 New Failures, 1 Pending, 8 Unrelated Failures

As of commit eff4294 with merge base c48ea12 ():

NEW FAILURES - The following jobs have failed:

pull / unittest / macos / macos-job (gh)
backends/xnnpack/test/ops/test_conv2d.py::TestConv2d::test_fp16_conv2d
Test CUDA Builds / test-model-cuda-e2e (SocialLocalMobile, Qwen3.5-35B-A3B-HQQ-INT4, quantized-int4-tile-packed) / linux-job (gh)
RuntimeError: Command docker exec -t 802b0706f9db7e8675dbad4fc76b46a6ec211f3003fc7ddfa496fa3b1a9e81e1 /exec failed with exit code 127
Test CUDA Windows Export and E2E / export-model-cuda-windows-artifact (facebook, dinov2-small-imagenet1k-1-layer, non-quantized) / linux-job (gh)
RuntimeError: Command docker exec -t 72845024f28fbeec896ff02b30fedc8112672b255c504cb67a52c16a5dd1d737 /exec failed with exit code 1
Test CUDA Windows Export and E2E / export-model-cuda-windows-artifact (nvidia, parakeet-tdt, non-quantized) / linux-job (gh)
RuntimeError: Command docker exec -t 2570c9deee066bbc50ce1699e388e785d449dcd68f5a25ea3c35b1aa4864c425 /exec failed with exit code 1
Test CUDA Windows Export and E2E / export-model-cuda-windows-artifact (nvidia, parakeet-tdt, quantized-int4-weight-only) / linux-job (gh)
RuntimeError: Command docker exec -t af1db5990796c2666f85961a78f4e3712b62fc0213e1a2e9687faa3b1dfc2c5b /exec failed with exit code 1
Test Metal Backend / test-metal-qwen35-moe-tiny / macos-job (gh)
/Users/ec2-user/runner/_work/executorch/executorch/pytorch/executorch/examples/models/qwen3_5_moe/main.cpp:352:3: error: use of undeclared identifier 'cudaMemGetInfo'

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

pull / test-models-linux (emformer_join, portable, linux.4xlarge.memory) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
pull / test-models-linux (emformer_join, xnnpack-quantization-delegation, linux.4xlarge.memory) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
pull / test-models-linux (ic4, portable, linux.4xlarge.memory) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
pull / test-models-linux (ic4, xnnpack-quantization-delegation, linux.4xlarge.memory) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
pull / test-models-linux (llama3_2_vision_encoder, portable, linux.4xlarge.memory) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
pull / unittest / windows / windows-job (gh) (matched win rule in flaky-rules.json)
##[error]The operation was canceled.
Test Metal Backend / export-model-metal-artifact (mistralai, Voxtral-Mini-4B-Realtime-2602, quantized-int4-metal) / macos-job (gh) (matched macos rule in flaky-rules.json)
File doesn't exist

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

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

pull / unittest-editable / windows / windows-job (gh) (trunk failure)
##[error]The operation was canceled.

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

digantdesai · 2026-04-21T03:47:51Z

+
+extern "C" {
+
+AOTITorchError aoti_torch_cuda_rand(


Is this from PyTorch/Aten or we are rolling our own?

rolling our own. I didn;t see there's an aten version.

digantdesai · 2026-04-21T03:48:48Z

-    example_prefill_len = config.max_seq_len - 1
-    prefill_tokens = torch.zeros((1, example_prefill_len), dtype=torch.long)
-    prefill_pos = torch.arange(example_prefill_len, dtype=torch.long)
+    prefill_tokens = torch.tensor([[0, 1]], dtype=torch.long)


why do we need to revert these changes or just a cherry-pick artifact?

just need another round of main branch update. Have bring it back.

digantdesai · 2026-04-21T03:50:58Z

-def _sample(logits, temperature):
-    """Sample from logits with temperature."""
-    if temperature <= 0:
-        return logits.argmax(dim=-1)
-    probs = torch.softmax(logits / temperature, dim=-1)
-    return torch.multinomial(probs, num_samples=1).squeeze(-1)


Do we need to verify if we are close to this using some unit tests?

digantdesai · 2026-04-21T03:51:25Z

 using SizesType = executorch::aten::SizesType;

+// Read a sampled token from the model output tensor [B, 1].
+// The model performs Gumbel-max sampling on-device and returns a single


can you add a comment to the paper or something for Gumbel-max?

digantdesai · 2026-04-21T03:52:33Z

    prev_token = cur_token;
-    stats.on_sampling_begin();
-    cur_token = llm::logits_to_token(*step_logits_ptr, FLAGS_temperature);
-    stats.on_sampling_end();


I guess we can't report sampling time separately anymore.

We can't until we land the device support pipeline. With device support we can make the model output in cuda and pull out sampler

digantdesai · 2026-04-21T03:53:45Z

+        # GPU-side Gumbel-max sampling: argmax(logits/T + gumbel_noise)
+        # Equivalent to sampling from softmax(logits/T) but fully on-device.
+        logits = logits / temperature.clamp(min=1e-6)
+        noise = torch.rand_like(logits)


Does this one need rand.cu?

Yes torch.rand_like need it

digantdesai · 2026-04-21T03:54:11Z

+ * advances a GPU-resident counter for unique offsets, making this fully
+ * compatible with CUDA graph capture and replay.
+ */
+AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_randint_low_out(


No unit-tests?

digantdesai · 2026-04-21T03:54:41Z

@claude review this PR.

claude · 2026-04-21T03:55:22Z

Claude finished @digantdesai's task in 2m 46s —— View job

PR Review: Add GPU-side Gumbel-max sampling for CUDA graph compatibility

Gather context and read changed files
Review rand.cu and rand.h (new CUDA sampling implementation)
Review changes to examples/models/qwen3_5_moe/ files
Review build system changes
Post review feedback

Summary

This PR fuses sampling into the model's forward() method using Gumbel-max trick on GPU, replacing the previous CPU-side logits_to_token / _sample approach. This eliminates a GPU→CPU→GPU roundtrip for the full [B, V] logits tensor during decode, and makes the sampling operation compatible with CUDA graph capture (since the RNG state lives on device). The new rand.cu/rand.h files implement custom CUDA RNG shims that use GPU-resident Philox state with atomic counter advancement.

Issues

1. Prefill exports with minimal example shape instead of max sequence length

export.py:655-656 — The prefill example was changed from torch.zeros((1, max_seq_len-1)) to torch.tensor([[0, 1]]). The original comment (still present at line 649-651) explains why the large example was needed:

Example T must equal max_seq_len-1 so AOTI compiles kernels (especially chunk_gated_delta_rule with CHUNK_SIZE=64) for the full range of sequence lengths. Smaller examples cause AOTI to bake in intermediate buffer sizes that reject longer prompts at runtime.

If this comment is still accurate, the change will break prefill for prompts longer than 2 tokens. If the constraint has been relaxed, the comment should be updated.

Fix this →

2. CUDA::curand not linked on MSVC path

CMakeLists.txt:154-157 — The MSVC branch of target_link_libraries(aoti_cuda_shims ...) doesn't include CUDA::curand, but the non-MSVC path (line 166) does. Since rand.cu uses curand_kernel.h, this will cause link failures on Windows MSVC builds.

Fix this →

3. Global mutable state without thread safety or cleanup

rand.cu:48-49 — d_rng and g_rng_init_done are bare globals with no synchronization. If ensure_rng_init is called concurrently from multiple threads (e.g. multiple models or streams), there's a data race on both the pointer and the flag. Consider using std::once_flag / std::call_once, or at minimum a std::atomic<bool>. Also, d_rng is never freed — not critical for a long-running process, but worth noting.

4. atomicAdd per-element is expensive for large tensors

rand.cu:88,104,118 — Each thread calls atomicAdd(&rng->counter, 1ULL). The comment at line 83 acknowledges this is fine for numel=1, but the kernels support arbitrary numel. For large tensors, the serial contention on a single counter will be a bottleneck. A common pattern is to have one thread (or one per block) grab a range via atomicAdd(&rng->counter, (unsigned long long)numel) and then distribute offsets within the block.

5. Gumbel noise epsilon values are ad-hoc

model.py:640 — The 1e-20 epsilon in torch.log(-torch.log(noise + 1e-20) + 1e-20) works for float32 but could be brittle. torch.rand_like returns values in [0, 1) (i.e., it can return exactly 0.0), so the inner log needs protection. The outer log also needs protection since -log(0 + 1e-20) is a large positive number, not problematic. The values seem fine for float32 but it would be good to add a brief note that this assumes float32 logits (which is guaranteed by the .float() cast on line 635).

6. No error checking on CUDA API calls in read_token

main.cpp:54 — cudaMemcpy return value is unchecked. If the copy fails for any reason, val will contain uninitialized data, silently producing garbage tokens.

7. Missing temperature in MLX export path

export.py:537-549 — The MLX export still uses the old 2-argument signature (example_tokens, example_input_pos). Since the model's forward() now requires temperature as a third argument, the MLX path will fail at export time. The forward signature change should either be gated on the backend, or the MLX export should be updated to pass temperature.

Fix this →

Suggestions

8. Seeding with time(nullptr) gives poor entropy

rand.cu:57 — time(nullptr) has second-level resolution and is predictable. For a sampling RNG this doesn't affect correctness, but for reproducibility it would be better to allow an explicit seed parameter, and for better entropy consider mixing in the device clock or using curandGenerateSeeds.

9. Consider making the forward method composable

As the PR description notes, once the device support pipeline lands, forward() should be decomposed. Currently, the model's forward() couples the transformer computation with sampling, which means:

You can't get raw logits for tasks like perplexity evaluation or beam search
The output is float token IDs rather than Long, requiring the read_token float→uint64 cast
Temperature is baked into the model graph

This is reasonable for the current CUDA-graph-focused use case but should be tracked for future cleanup.

10. read_token could use the output dtype directly

main.cpp:45-58 — The model returns argmax(...).float() which means the output is a float representation of an integer. The cast chain float → uint64_t works but is unusual. If the model returned Long instead, you could avoid the float intermediary and read an int64_t directly.

CI Failures

The strides_vec compilation errors in CI are from the base cuda-graph branch (cuda_backend.cpp:697 uses strides_vec before it's defined at line 715), not from this PR's changes. The 2 Windows failures are also from trunk.

Addressing Reviewer Questions

Re: @digantdesai's questions:

rand.cu origin (line 134): This is a custom implementation, not from PyTorch/ATen. It implements the AOTI shim interface (aoti_torch_cuda_rand, aoti_torch_cuda_randint_low_out) using cuRAND Philox primitives with GPU-resident state for CUDA graph compatibility.
Unit tests (rand.h:56): Agree — the new RNG shims need unit tests, particularly verifying (1) output is in expected range, (2) different replay invocations produce different sequences, and (3) basic statistical properties (uniformity).
export.py reverts (line 655): See Issue Add support for quantized LeakyReLU #1 above — the prefill example shape change may break long prompts.
rand.cu needed by model.py? (line 639): Yes — torch.rand_like(logits) in the forward method compiles to an AOTI rand call, which dispatches to aoti_torch_cuda_rand at runtime.
Sampling time reporting (main.cpp:216): Correct, sampling is now part of the model execution time. You could add a cudaEvent pair around the read_token call to measure the D2H transfer time, but the sampling kernel time would be inseparable from the model forward pass.

digantdesai

Thanks Songhao. Can we make a util such that this can be used by other models easily?

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 13, 2026

Gasoonjia force-pushed the cuda-graph-sampling branch 2 times, most recently from b4f9eca to c7450dd Compare April 13, 2026 20:34

Gasoonjia added ciflow/cuda labels Apr 13, 2026

Gasoonjia force-pushed the cuda-graph-sampling branch 2 times, most recently from 1bf973d to 028894e Compare April 13, 2026 21:26

Gasoonjia marked this pull request as ready for review April 13, 2026 23:57

Gasoonjia requested review from kirklandsign, larryliu0820 and lucylq as code owners April 13, 2026 23:57

Gasoonjia mentioned this pull request Apr 14, 2026

CUDA sampler support in Qwen 3.5 MoE #18885

Open

digantdesai reviewed Apr 21, 2026

View reviewed changes

digantdesai approved these changes Apr 21, 2026

View reviewed changes

Gasoonjia force-pushed the cuda-graph-sampling branch from 93bee20 to d3bca0d Compare April 23, 2026 03:58

gasoonjia and others added 3 commits April 22, 2026 22:16

Add GPU-side Gumbel-max sampling for CUDA graph compatibility

0492e8d

lintrunner

8c0bbf3

remove git info

5245f64

Gasoonjia force-pushed the cuda-graph-sampling branch from d3bca0d to 5245f64 Compare April 23, 2026 05:20

Gasoonjia added 2 commits April 23, 2026 15:44

reintro llm headers

880391d

lint

6f411af

add top-p and top-k arg

eff4294

Conversation

Gasoonjia commented Apr 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Apr 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

❌ 6 New Failures, 1 Pending, 8 Unrelated Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

digantdesai Apr 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

digantdesai commented Apr 21, 2026

Uh oh!

claude Bot commented Apr 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Review: Add GPU-side Gumbel-max sampling for CUDA graph compatibility

Summary

Issues

Suggestions

CI Failures

Addressing Reviewer Questions

Uh oh!

digantdesai left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Gasoonjia commented Apr 13, 2026 •

edited

Loading

pytorch-bot Bot commented Apr 13, 2026 •

edited

Loading

digantdesai Apr 21, 2026 •

edited

Loading

claude Bot commented Apr 21, 2026 •

edited

Loading

digantdesai left a comment •

edited

Loading