[abdelaziz-mahdy]

Summary

Fixes #25100.

Bug: a stop() that arrives while the prompt is being prefilled is lost. TextTokenGenerator::generate() clears the stop flag when its decode loop starts, after prefill, so generation runs to max_new_tokens.

Fix

• TextTokenGenerator:
  • new reset_stop() clears the flag;
  • generate() takes reset_stop (default true, so existing callers behave as before);
  • the loop also checks the flag before each step.
• TextLLMRunner::generate() and MultimodalRunner::generate():
  • call reset_stop() when they start, before prefill;
  • then call the generator with reset_stop=false.

Result

• Stop during prefill: only the token prefill produced is emitted, and no decode steps run.
• A stop from before generate() still doesn't carry into it.

Test plan

New tests in extension/llm/runner/test/test_text_llm_runner.cpp:

• StopDuringPrefillEndsGeneration:
  • fails without the fix: the decoder runs 9 steps and the callback fires 10 times;
  • passes with it: 1 callback, 0 steps.
• StopBeforeGenerateDoesNotCarryOver: passes.

Run on macOS arm64 (Release, -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON):

cmake --build cmake-out --target test_runner
./test_runner --gtest_filter='RunnerTest.*:*TextPrefiller*:*TextDecoderRunner*:*GenerationConfig*:*TextStream*'
[  PASSED  ] 57 tests.

The changed files are clang-format clean.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 25 Awaiting Approval

As of commit bc35054 with merge base d3369c4 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux CUDA Wheels (gh)
• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux CUDA Wheels (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows CUDA Wheels (gh)
• Build Windows Wheels (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• MLX (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test CPU Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test Vulkan Backend (specialized runners) (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[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.