Repository navigation
LLM runner: stop() during prefill is ignored #25100
Description
🐛 Describe the bug
Summary: stop() on an LLM runner is ignored if it's called while the prompt is still being prefilled. Generation then runs all the way to max_new_tokens.
What happens
TextLLMRunner::generate()prefills the prompt, then callsTextTokenGenerator::generate().TextTokenGenerator::generate()setsshould_stop_ = falseright before its decode loop (text_token_generator.h).- So a
stop()that arrived during prefill is wiped, and every decode step still runs. MultimodalRunner::generate()has the same pattern.
Why it matters
- Prefill of a long prompt can take seconds (much longer on CPU or WebAssembly). That's exactly when users press Stop.
- Anything waiting for the generation to end (shutting the runner down, the next request) waits for the full reply.
Repro (as a unit test in extension/llm/runner/test/test_text_llm_runner.cpp)
- A mock prefiller that calls
runner.stop()insideprefill(). generate()withmax_new_tokens = 10.- Actual: the token callback runs 10 times and the decoder runs 9 steps.
- Expected: 1 callback (the token prefill produced), no decoder steps.
Fix and test: PR coming next.
Versions
- ExecuTorch
mainat d3369c4 (also in v1.6.0). C++ runtime, any platform. - Found from Flutter (Dart FFI) on macOS arm64; reproduced with the gtest above.
collect_env output
Collecting environment information...
PyTorch version: 2.14.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A
OS: macOS 27.2.1 (arm64)
GCC version: Could not collect
Clang version: 21.0.0 (clang-2100.4.1.1)
CMake version: version 4.5.0
Libc version: N/A
Python version: 3.10.11 (v3.10.11:7d4cc5aa85, Apr 4 2023, 19:05:19) [Clang 13.0.0 (clang-1300.0.29.30)] (64-bit runtime)
Python platform: macOS-27.2.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False
Caching allocator config: N/A
CPU:
Apple M2 Pro
Versions of relevant libraries:
[pip3] executorch==1.5.0
[pip3] flake8==6.1.0
[pip3] flake8-breakpoint==1.1.0
[pip3] flake8-bugbear==24.4.26
[pip3] flake8-comprehensions==3.14.0
[pip3] flake8-plugin-utils==1.3.3
[pip3] flake8-pyi==23.5.0
[pip3] mypy==1.14.1
[pip3] mypy_extensions==1.1.0
[pip3] numpy==2.2.6
[pip3] pytorch_tokenizers==1.5.0
[pip3] torch==2.14.0
[pip3] torchao==0.18.0+git5f2baf9d5
[pip3] torchvision==0.29.0
[conda] Could not collect
Activity