[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
• [executorch][muse] Accept images in CUDA DFlash batching #24066
• [executorch][muse] Serve images on CUDA batching #24065
• [executorch][muse] Export vision with the CUDA batching contract #24064
• [executorch][cuda] Run CudaExecutor steps from prepared embeddings #24063
• [executorch][muse] Share the Muse Glimmer embeddings engine across backends #24062

forward_others and forward_from_embeddings take up to max_step_tokens
(4096) tokens and select which rows' logits to return. The selector's dynamic
bound was also max_step_tokens, so memory planning sized each method's FP32
logits output at 4096 x 202048 x 4 B = 3157 MiB, allocated at load whether or
not a request ever uses it.

A step never selects more rows than the device samplers take
(MAX_SAMPLE_ROWS = 64: one row per prefilling session, or max_owners x block for DFlash verify), and the executor already rejects a step with more.
Bound the selector at min(max_step_tokens, MAX_SAMPLE_ROWS). Each planned
logits output drops to 49 MiB; token width and the DFlash taps keep their
max_step_tokens bound.

Differential Revision: D127114350

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Unclassified Failure

As of commit dedcd5f with merge base 2c3e231 ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / macos / macos-job (gh)
  examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity
• pull / unittest-editable / macos / macos-job (gh)
  examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• Test CUDA Builds / test-muse-glimmer-cuda-e2e-kquant-17gb-solo-text-batching / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /__w/_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.