[Gasoonjia]

Stack from ghstack (oldest at bottom):

• [executorch][cuda] Grow sequence-layout ring layers on demand #24076
• [executorch][muse-glimmer] Bound the CUDA prefill logits selector by the sampler rows #24107
• [executorch][muse] Serve Muse Glimmer CUDA batching on a batched-sequence cache #24075
• [executorch][cuda] Lower off-graph KV to the batched-sequence layout #24074
• [executorch][cuda] Add a batched-sequence CUDA KV cache #24073
• -> [executorch][cuda] Add a length-aware Triton SDPA over per-sequence KV columns #24072

Adds triton::sdpa_varlen, the attention kernel the CUDA batched-sequence KV cache needs. Each resident sequence owns one column of a [rows, cols, H, D] store, and a device segment table seq_info Int32[S, 4] = (q_start, q_len, col, kv_end) says which packed query rows belong to which column and how long that sequence is. Lengths are read on device, so one captured CUDA graph serves any mix of sequence lengths, as the MLX per-sequence SDPA does per group.

Covers causal, non-causal (DFlash draft blocks), sliding window and ring storage, GQA, rows no segment covers (zeroed), and split-K for short steps. A symbolic query length takes the generic path without host-side length branches, so exporting it with a dynamic step width does not pin that width.

It is a new op in its own file; triton::sdpa and its dispatch are untouched.

Differential Revision: D125987296

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 10 Unrelated Failures

As of commit 15a539b with merge base 4eb094b ():

NEW FAILURE - The following job has failed:

• MLX / test-mlx-llm-offgraph (google/gemma-4-E2B-it, gemma4-e2b, gemma4) / test-mlx-llm-offgraph-gemma4-e2b (gh)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• MLX / test-mlx-llm (unsloth/Qwen3-0.6B, qwen3-0.6b, true, nvfp4) / test-mlx-llm-qwen3-0.6b-custom-nvfp4 (gh) (detected as infra flaky with no log or failing log classifier)

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

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

• MLX / test-mlx / test-mlx (gh) (trunk failure)
  [ FAILED ] Implementations/MLXSequenceControlTest.DFlashRejectsCellOrPersistsOnlyContext/Sequence, where GetParam() = 4-byte object <01-00 00-00>
• MLX / test-mlx-llm-offgraph (unsloth/gemma-3-1b-it, gemma3-1b, gemma) / test-mlx-llm-offgraph-gemma3-1b (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• MLX / test-mlx-llm-offgraph (unsloth/Llama-3.2-1B-Instruct, llama-1b, llama3) / test-mlx-llm-offgraph-llama-1b (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• MLX / test-mlx-parakeet / test-mlx-parakeet (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• MLX / test-mlx-stories110m / test-mlx-stories110m (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• MLX / test-mlx-supertonic / test-mlx-supertonic (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 8
• MLX / test-mlx-voxtral / test-mlx-voxtral (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
• MLX / test-mlx-voxtral-realtime / test-mlx-voxtral-realtime (gh) (trunk failure)
  RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 134
• Test CUDA Builds / test-muse-glimmer-cuda-e2e-kquant-17gb-dflash-image-batching / linux-job (gh) (trunk failure)
  ##[error]received SIGINT

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.