clean up some runtime potential bugs (#19447)

[billmguo]

Summary:

1. Zero cache on allocation (line 65-66): std::fill on cache_data_ and update_data_ after allocator_.allocate() —eliminates uninitialized memory garbage that varies across devices.
2. Zero cache on reset (line 191): std::fill on cache_data_ in reset() — ensures stale KV cache from a previous prompt is fully cleared, not just the position counters.
3. Zero padding in last prefill chunk (line 618-621): When batch_len < input_len, fill the tail of the input buffer with zeros — prevents stale tokens from a previous chunk leaking through the embedding layer. sa_runner.cpp
4. Call runner.reset() before each prompt in the multi-prompt loop, stdin prompt loop, and stdin tokens loop —ensures the KV cache, masks, and input_pos_ are fully reset between prompts

Reviewed By: Andriyluck, derekdixu

Differential Revision: D104615993

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 1cd3d16 with merge base 7127953 ():

NEW FAILURE - The following job has failed:

• Docathon Labels Sync / check-labels (gh)
  Process completed with exit code 1.

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

[meta-codesync]

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

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

[lucylq]

@claude review

[claude]

Claude finished @lucylq's task in 1m 23s —— View job

---

Review of PR #19447: clean up some runtime potential bugs

• Gather context and view diff
• Review code changes
• Post review feedback

---

Overall: The changes are defensive improvements to zero-initialize memory in the static KV cache and input buffer. They address real potential issues with uninitialized/stale data. A few observations:

---

1. Zero cache on allocation (lines 67-68) ✅

std::fill(cache_data_, cache_data_ + cache_data_size_, T(0));
std::fill(update_data_, update_data_ + update_data_size_, T(0));

Good. The AllocatorT::allocate() is a template parameter with no guarantee of zero-initialization (e.g., std::allocator::allocate does not zero memory). This eliminates non-deterministic behavior from uninitialized reads if any code path accesses cache regions before they're written.

---

2. Zero cache on reset (line 191) ✅

std::fill(cache_data_, cache_data_ + cache_data_size_, T(0));

This makes reset() truly clear the cache data rather than just resetting position counters. Important for multi-turn/multi-prompt scenarios where stale KV values could influence attention if the mask logic has any gaps.

Nit: Should update_data_ also be zeroed in reset()? Currently only cache_data_ is cleared. If the reasoning is that update_data_ is always overwritten before being read (since it's used as output buffer), that's fine — but it's worth confirming that invariant holds in all code paths.

---

3. Zero padding in last prefill chunk (lines 619-622) ✅

if (batch_len < input_len) {
  std::fill(
      input_buffer.begin() + batch_len, input_buffer.end(), TokenT(0));
}

Correct fix. When the last chunk of tokens is smaller than input_len, the tail of input_buffer would otherwise contain stale tokens from a previous iteration of the loop (the prior std::copy only writes batch_len elements). Zero-padding prevents these from leaking through the embedding layer.

---

4. PR description vs. actual diff mismatch

The PR description mentions "Call runner.reset() before each prompt in the multi-prompt loop, stdin prompt loop, and stdin tokens loop" but this change is not present in the diff. Only the header file is modified. Was this change intended for a separate file (e.g., sa_runner.cpp) and accidentally omitted, or is it being landed separately?

---

Summary: The three changes that are present look correct and low-risk. They add O(n) zeroing overhead on allocation, reset, and the last prefill chunk — negligible compared to model inference time. The main question is whether the runner.reset() calls mentioned in the description are missing from this PR.