From e1ed3afaefa0d84425b9e04f7ec16a2c5af3d5a7 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 9 Nov 2026 21:05:27 -0800 Subject: [PATCH] [executorch][muse-glimmer] Bound the CUDA prefill logits selector by the sampler rows Pull Request resolved: https://github.com/pytorch/executorch/pull/24107 `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. ghstack-source-id: 469359754 @exported-using-ghexport Differential Revision: [D127114350](https://our.internmc.facebook.com/intern/diff/D127114350/) --- examples/models/muse-glimmer/export/export_solo.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index 6dc5fbb6f8d..fb1c94dcd92 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/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)),