Allow TextDecoderRunner to reuse a persistent Sampler

[kirklandsign]

Summary:
TextDecoderRunner now optionally takes a std::unique_ptr in its
constructor, owning the Sampler and reusing it across token generation steps.
This avoids creating a temporary Sampler on every logits_to_token() call,
preserving RNG state across tokens and laying groundwork for future grammar-
constrained decoding.

Changes:

• sampler.h: add set_temperature() so the owned Sampler's temperature can be
  updated per-call without reconstructing it
• util.h: extract sample_from_logits(Tensor, Sampler&) for dtype dispatch with
  an existing Sampler; logits_to_token delegates to it
• text_decoder_runner.h/cpp: ctor takes optional unique_ptr; when set,
  logits_to_token updates its temperature and reuses it
• multimodal_decoder_runner.h: forward Sampler to base class ctor
• llm_runner_helper.cpp: factory creates Sampler and passes to TextDecoderRunner
• text_llm_runner.cpp: temperature resolved to local variable per-call
• test_text_decoder_runner.cpp: add tests for ctor-injected Sampler

Differential Revision: D103748189

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 1 Cancelled Job, 11 Pending, 11 Unrelated Failures

As of commit 0a1b6b7 with merge base 851cffb ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner (gh)
  >>> Lint for extension/llm/sampler/sampler.h:
• pull / android / run-emulator (gh)
  The process '/usr/bin/sh' failed with exit code 1
• pull / test-qnn-delegate-linux / linux-job (gh)
  backends/qualcomm/tests/test_qnn_delegate.py::TestQNNFloatingPointOperator::test_qnn_backend_argmax

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-lora-multimethod-linux / linux-job (gh)

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-lora-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-models-linux (ic4, portable, linux.4xlarge.memory) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-models-linux (ic4, xnnpack-quantization-delegation, linux.4xlarge.memory) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-models-linux (mobilebert, portable, linux.2xlarge) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-models-linux (mobilebert, xnnpack-quantization-delegation, linux.2xlarge) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-models-linux (phi_4_mini, portable, linux.4xlarge.memory) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-moshi-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-phi-3-mini-runner-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-vulkan-models-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest / linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest-editable / linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[meta-codesync]

@kirklandsign has exported this pull request. If you are a Meta employee, you can view the originating Diff in D103748189.

[kirklandsign]

@claude review

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

[claude]

Claude finished @kirklandsign's task in 1m 37s —— View job

---

Code Review: Allow TextDecoderRunner to reuse a persistent Sampler

• Read changed files and understand the design
• Review sampler.h changes
• Review util.h changes
• Review text_decoder_runner.h/cpp changes
• Review multimodal_decoder_runner.h changes
• Review llm_runner_helper.cpp changes
• Review text_llm_runner.cpp changes
• Review test file
• Post review feedback

---

Overall: Clean, well-structured change. The design is sound — constructor-injected Sampler ownership with a fallback to per-call temporary creation preserves backward compatibility while enabling RNG state persistence. A few minor observations below.

---

sampler.h — set_temperature()

The new set_temperature() at line 47 uses static_cast<bool>(temperature) to match the constructor pattern in sampler.cpp:179. This is fine and consistent, though a reader might initially expect temperature != 0.0f. The static_cast<bool> cast is idiomatic for this pattern in the codebase (matches both constructors), so no issue.

One nit: the 0.0f literal in set_temperature vs 0 (no f suffix) in the constructor initializer lists (sampler.cpp:179,185). Functionally identical but worth noting for consistency.

util.h — sample_from_logits() extraction

Good refactor. The extracted function cleanly separates dtype dispatch from Sampler construction. The retained logits_to_token() wrapper is minimal and correctly delegates.

text_decoder_runner.h — logits_to_token() dispatch

if (sampler_) {
  sampler_->set_temperature(temperature);
  return ::executorch::extension::llm::sample_from_logits(logits_tensor, *sampler_);
}
return ::executorch::extension::llm::logits_to_token(logits_tensor, temperature);

This is clean. The only subtlety: when sampler_ is set, the vocab_size used for sampling is fixed at construction time (in llm_runner_helper.cpp). The fallback path (logits_to_token) reads vocab_size from the tensor shape on every call. If the model's output vocab_size ever differed from metadata.at(kVocabSize), the persistent Sampler could misbehave — but this is an extremely unlikely edge case and not a real concern for production models.

llm_runner_helper.cpp — Sampler factory

int32_t vocab_size = static_cast<int32_t>(metadata.at(kVocabSize));
float init_temp = temperature == -1.0f ? 0.0f : temperature;
auto sampler = std::make_unique<Sampler>(vocab_size, init_temp);

The -1.0f sentinel handling here mirrors what text_llm_runner.cpp does for the generate loop. The initial temperature doesn't matter much since logits_to_token() calls set_temperature() before every sample, but setting a reasonable initial value is still good practice.

For the multimodal path (line 344), the sampler is created with hardcoded 0.0f temperature. This is fine since it'll be overwritten per-call.

text_llm_runner.cpp

Extracting resolved_temp to a local variable is a minor readability improvement. No concerns.

multimodal_decoder_runner.h

The constructor signature expansion correctly forwards all parameters to the base class. The default sampler = nullptr preserves backward compatibility.

Tests

Good coverage:

• Basic injected sampler with argmax (greedy)
• Temperature switching on a persistent sampler
• 3D tensor handling with injected sampler

One observation: the tests pass nullptr for both module and io_manager when constructing TextDecoderRunner with an injected sampler — this works because logits_to_token doesn't touch those members. This is fine for unit testing the sampling path but worth noting that it implicitly tests that the sampling logic is decoupled from the model execution logic.

---

Summary

This is a well-designed incremental change that:

1. Preserves full backward compatibility (nullptr default for sampler)
2. Enables RNG state persistence across tokens (key for reproducible generation)
3. Lays groundwork for grammar-constrained decoding (persistent sampler state)
4. Has good test coverage for the new code paths

No blocking issues. LGTM.

---

[Copilot]

Pull request overview

This PR updates the LLM decoding path to optionally reuse a persistent Sampler inside TextDecoderRunner, avoiding per-token sampler reconstruction, preserving RNG state across tokens, and enabling future extensions like grammar-constrained decoding.

Changes:

• Added Sampler::set_temperature() and refactored sampling utilities to support reusing an existing Sampler (sample_from_logits()).
• Extended TextDecoderRunner (and MultimodalDecoderRunner) constructors to optionally accept and own a std::unique_ptr<Sampler> and reuse it during token selection.
• Updated runner factories to construct and inject a Sampler, and added unit tests covering injected sampler behavior.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
extension/llm/sampler/util.h	Extracts sample_from_logits() and routes logits_to_token() through it.
extension/llm/sampler/sampler.h	Adds set_temperature() to update an existing sampler per call.
extension/llm/runner/text_llm_runner.cpp	Resolves temperature into a local variable before generation.
extension/llm/runner/text_decoder_runner.h	Adds optional owned Sampler and reuses it in logits_to_token().
extension/llm/runner/text_decoder_runner.cpp	Wires new constructor parameter into member initialization.
extension/llm/runner/test/test_text_decoder_runner.cpp	Adds tests for constructor-injected sampler reuse and temperature switching.
extension/llm/runner/multimodal_decoder_runner.h	Forwards optional sampler/method name to TextDecoderRunner base ctor.
extension/llm/runner/llm_runner_helper.cpp	Creates a Sampler in factories and injects it into decoder runners.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.