[executorch][muse-glimmer] Bound the CUDA prefill logits selector by the sampler rows

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #24076
* __->__ #24107
* #24075
* #24074
* #24073
* #24072

`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

Pull Request resolved: #24107

Author: Gasoonjia

examples/models/muse-glimmer/export/export_solo.py

@@ -527,13 +527,16 @@ def export_cuda_batching(  # noqa: C901
         materialize_runtime_buffers(model, dtype=torch.bfloat16)
     geometry = offgraph_kv_cache_geometry(model)
 
-    # A sequence prefill samples a few rows; tracing all rows creates a
-    # max_step_tokens * vocab FP32 output during AOTI's JIT validation.
-    example_rows = _CUDA_MIN_PREFILL_CHUNK if sequence_layout else max_step_tokens
+    # The selected logits feed the device samplers, which take at most
+    # MAX_SAMPLE_ROWS rows, so a step never selects more. Bounding the selector
+    # there instead of at the step width keeps the memory-planned FP32 logits
+    # output at MAX_SAMPLE_ROWS * vocab rather than max_step_tokens * vocab.
+    max_selected_rows = min(max_step_tokens, MAX_SAMPLE_ROWS)
+    example_rows = _CUDA_MIN_PREFILL_CHUNK if sequence_layout else max_selected_rows
     width = Dim(
         "width", min=max(widest_static, _CUDA_MIN_PREFILL_CHUNK), max=max_step_tokens
     )
-    rows = Dim("rows", min=_CUDA_MIN_PREFILL_CHUNK, max=max_step_tokens)
+    rows = Dim("rows", min=_CUDA_MIN_PREFILL_CHUNK, max=max_selected_rows)
     programs = {}
     with (
         common.BoundMethodForward(model, types.MethodType(step_forward, model)),