Harden existing audio prefill APIs (#20136) by kirklandsign · Pull Request #20136 · pytorch/executorch

kirklandsign · 2026-06-08T22:44:27Z

Summary:

Improve the existing audio prefill methods (prefillAudio, prefillRawAudio) with input validation and code cleanup. No API signature changes.

Kotlin:

Add require checks for positive dimensions and array size >= expected element count
Use Math.multiplyExact to detect overflow in dimension multiplication (consistent with ByteBuffer variants)
Improve docstrings: clarify data types (uint8/float32), document throws IllegalArgumentException

JNI (jni_layer_llama.cpp):

Add dimension validation (batch_size > 0, n_bins > 0, etc.)
Add array size consistency check against batchSize * nBins * nFrames
Replace double-allocation + per-element copy with single allocation + reinterpret_cast / direct getRegion
Fix emplace_back lint (modernize-use-emplace)
Remove silent success on empty data (now returns InvalidArgument)
Use size_t casts to prevent integer overflow in size calculations

Differential Revision: D107929913

pytorch-bot · 2026-06-08T22:44:30Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 4 New Failures

As of commit 69dba93 with merge base a9d5674 ():

NEW FAILURES - The following jobs have failed:

pull / test-lora-multimethod-linux / linux-job (gh)
RuntimeError: Command docker exec -t 32a0bf59cdf4ef4d4f59cf41aaceb86e50310a29a6ae99b091c8b2a22b4eaf3d /exec failed with exit code 1
pull / test-qnn-passes-linux / linux-job (gh)
RuntimeError: Command docker exec -t 50d5c1b2c6ebdf94c55406ad882edffe2e820ab91bae284dca449d37a34a3ae1 /exec failed with exit code 1
pull / test-sqnr-static-llm-qnn-linux (smollm2_135m) / linux-job (gh)
RuntimeError: Command docker exec -t 56ae3b08a7dff95ace067fe34b9a0cdb67dd22bb46825113915a32e14fc4366c /exec failed with exit code 1
pull / test-static-llama-qnn-linux (stories_260k_bc) / linux-job (gh)
RuntimeError: Command docker exec -t 8c508bbec0ca7e5f96eba155f4ea0a2df85d3c05b76187518505d3f521857d3c /exec failed with exit code 1

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

meta-codesync · 2026-06-08T22:44:38Z

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

github-actions · 2026-06-08T22:45:21Z

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.

Copilot

Pull request overview

This PR aims to harden the Android audio prefill APIs by adding input validation and reducing unnecessary JNI copies while keeping API signatures unchanged.

Changes:

Add dimension/size validation for prefillAudio / prefillRawAudio in Kotlin and JNI.
Replace JNI double-allocation + per-element conversions with direct getRegion into native vectors.
Update Kotlin KDoc to clarify data types and thrown exceptions.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 7 comments.

File	Description
extension/android/jni/jni_layer_llama.cpp	Adds validation and refactors audio/raw-audio JNI array handling to reduce copies (but currently has a correctness issue when arrays are larger than expected).
extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.kt	Adds `require(...)` validation for audio prefills and improves documentation.

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

    auto data_size = data->size();
-    if (data_size == 0) {
-      return 0;
+    size_t expected =
+        static_cast<size_t>(batch_size) * static_cast<size_t>(n_bins) *
+        static_cast<size_t>(n_frames);
+    if (static_cast<size_t>(data_size) < expected) {
+      return static_cast<jint>(Error::InvalidArgument);
    }
-    std::vector<jbyte> data_jbyte(data_size);
    std::vector<uint8_t> data_u8(data_size);
-    data->getRegion(0, data_size, data_jbyte.data());
-    for (int i = 0; i < data_size; i++) {
-      data_u8[i] = data_jbyte[i];
-    }
+    data->getRegion(
+        0, data_size, reinterpret_cast<jbyte*>(data_u8.data()));
    llm::Audio audio{std::move(data_u8), batch_size, n_bins, n_frames};


    auto data_size = data->size();
-    if (data_size == 0) {
-      return 0;
+    size_t expected =
+        static_cast<size_t>(batch_size) * static_cast<size_t>(n_bins) *
+        static_cast<size_t>(n_frames);
+    if (static_cast<size_t>(data_size) < expected) {
+      return static_cast<jint>(Error::InvalidArgument);
    }
-    std::vector<jfloat> data_jfloat(data_size);
    std::vector<float> data_f(data_size);
-    data->getRegion(0, data_size, data_jfloat.data());
-    for (int i = 0; i < data_size; i++) {
-      data_f[i] = data_jfloat[i];
-    }
+    data->getRegion(0, data_size, data_f.data());
    llm::Audio audio{std::move(data_f), batch_size, n_bins, n_frames};


    auto data_size = data->size();
-    if (data_size == 0) {
-      return 0;
+    size_t expected =
+        static_cast<size_t>(batch_size) * static_cast<size_t>(n_channels) *
+        static_cast<size_t>(n_samples);
+    if (static_cast<size_t>(data_size) < expected) {
+      return static_cast<jint>(Error::InvalidArgument);
    }
-    std::vector<jbyte> data_jbyte(data_size);
    std::vector<uint8_t> data_u8(data_size);
-    data->getRegion(0, data_size, data_jbyte.data());
-    for (int i = 0; i < data_size; i++) {
-      data_u8[i] = data_jbyte[i];
-    }
+    data->getRegion(
+        0, data_size, reinterpret_cast<jbyte*>(data_u8.data()));
    llm::RawAudio audio{std::move(data_u8), batch_size, n_channels, n_samples};


-   * @param batchSize Input batch size
-   * @param nBins Input number of bins
-   * @param nFrames Input number of frames
+   * @param audio Input audio as a byte array (uint8 values)


-   * @param batchSize Input batch size
-   * @param nBins Input number of bins
-   * @param nFrames Input number of frames
+   * @param audio Input audio as a float array (float32 values)


-   * Prefill the KV cache with the given raw audio input.
+   * Prefill the KV cache with raw audio data.
   *
   * @param audio Input raw audio as a byte array


+      require(batchSize > 0 && nBins > 0 && nFrames > 0) {
+        "batchSize, nBins, and nFrames must all be positive"
+      }


kirklandsign · 2026-06-08T22:57:40Z

@claude review

claude · 2026-06-08T22:58:25Z

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

Code Review: Harden existing audio prefill APIs

Read the PR diff and changed files
Analyze JNI changes for correctness and safety
Analyze Kotlin changes for correctness and safety
Post review feedback

Overall this is a solid hardening PR — the validation additions, emplace_back fix, and copy elimination are all welcome improvements. I have one correctness issue to flag and a few minor notes.

Bug: Vector size mismatch with `llm::Audio` constructor

Severity: High — this will cause an abort at runtime for valid (but oversized) inputs.

The Kotlin layer validates audio.size >= expected (allows arrays larger than needed), and the JNI code allocates a vector of size data_size (the full JNI array length) and copies all elements. However, llm::Audio's constructor (audio.h:57-63) enforces strict equality:

ET_CHECK_MSG(
    ... std::get<std::vector<uint8_t>>(data_).size() ==
        static_cast<size_t>(batch_size * n_bins * n_frames), ...);

So if a caller passes an array with size > expected, the Kotlin check passes, the JNI check passes, but the Audio constructor's ET_CHECK_MSG will abort the process.

Affected lines: jni_layer_llama.cpp:487-490, jni_layer_llama.cpp:520-522

Suggested fix (pick one):

Strict equality everywhere — change Kotlin >= to == and JNI < to !=. This is the simplest and most defensive option.
Copy only expected elements — allocate expected and getRegion(0, expected, ...). This allows callers to pass padded buffers.

Note: RawAudio is a plain struct with no size assertion, so the third function (prefill_raw_audio_input) does not have this problem — but using >= there is still inconsistent with option 1.

Fix this →

Minor Notes

1. Overflow in audio.h (pre-existing, not introduced here)

The Audio constructor computes batch_size * n_bins * n_frames using int32_t arithmetic (audio.h:60), which can overflow for large dimensions. The new Kotlin/JNI validation catches this before reaching the constructor, so it's safe in practice for this call path, but worth noting for future hardening of audio.h itself.

2. reinterpret_cast<jbyte*>(data_u8.data()) is correct

jbyte is signed char; uint8_t is unsigned char. They share the same size and alignment on all platforms, and reinterpret_cast preserves the bit pattern. The raw bytes are correctly reinterpreted as uint8_t after the copy. Good optimization eliminating the intermediate vector + per-element loop.

3. Float path getRegion (line 521)

jfloat is float (IEEE 754, same type), so data->getRegion(0, data_size, data_f.data()) works directly without a cast. Correct.

4. emplace_back fix

Changing inputs.emplace_back(llm::MultimodalInput{std::move(audio)}) to inputs.emplace_back(std::move(audio)) — good, this avoids the intermediate temporary and satisfies modernize-use-emplace.

Summary

The PR successfully hardens input validation and eliminates unnecessary copies. The one blocking issue is the vector-size mismatch between the >= check at the JNI/Kotlin layer and the == assertion inside llm::Audio's constructor. Once that's aligned, this looks good to land.

Summary: Improve the existing audio prefill methods (`prefillAudio`, `prefillRawAudio`) with input validation and code cleanup. No API signature changes. Kotlin: - Add `require` checks for positive dimensions and array size >= expected element count - Use `Math.multiplyExact` to detect overflow in dimension multiplication (consistent with ByteBuffer variants) - Improve docstrings: clarify data types (uint8/float32), document `throws IllegalArgumentException` JNI (`jni_layer_llama.cpp`): - Add dimension validation (`batch_size > 0`, `n_bins > 0`, etc.) - Add array size consistency check against `batchSize * nBins * nFrames` - Replace double-allocation + per-element copy with single allocation + `reinterpret_cast` / direct `getRegion` - Fix `emplace_back` lint (`modernize-use-emplace`) - Remove silent success on empty data (now returns `InvalidArgument`) - Use `size_t` casts to prevent integer overflow in size calculations Differential Revision: D107929913

Copilot AI review requested due to automatic review settings June 8, 2026 22:44

kirklandsign requested a review from psiddh as a code owner June 8, 2026 22:44

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 8, 2026

Copilot started reviewing on behalf of kirklandsign June 8, 2026 22:44 View session

meta-codesync Bot added the meta-exported label Jun 8, 2026

Copilot AI reviewed Jun 8, 2026

View reviewed changes

meta-codesync Bot changed the title ~~Harden existing audio prefill APIs~~ Harden existing audio prefill APIs (#20136) Jun 8, 2026

meta-codesync Bot force-pushed the export-D107929913 branch from 12ab659 to d165042 Compare June 8, 2026 23:09

meta-codesync Bot force-pushed the export-D107929913 branch from d165042 to 1692652 Compare June 8, 2026 23:13

Copilot AI review requested due to automatic review settings June 8, 2026 23:18

meta-codesync Bot force-pushed the export-D107929913 branch from 1692652 to 69dba93 Compare June 8, 2026 23:18

kirklandsign review requested due to automatic review settings June 8, 2026 23:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Harden existing audio prefill APIs (#20136)#20136

Harden existing audio prefill APIs (#20136)#20136
kirklandsign wants to merge 1 commit into
mainfrom
export-D107929913

kirklandsign commented Jun 8, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

pytorch-bot Bot commented Jun 8, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented Jun 8, 2026

Uh oh!

github-actions Bot commented Jun 8, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

kirklandsign commented Jun 8, 2026

Uh oh!

claude Bot commented Jun 8, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

kirklandsign commented Jun 8, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 4 New Failures

Uh oh!

meta-codesync Bot commented Jun 8, 2026

Uh oh!

github-actions Bot commented Jun 8, 2026

This PR needs a release notes: label

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

kirklandsign commented Jun 8, 2026

Uh oh!

claude Bot commented Jun 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review: Harden existing audio prefill APIs

Bug: Vector size mismatch with llm::Audio constructor

Minor Notes

Summary

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kirklandsign commented Jun 8, 2026 •

edited by meta-codesync Bot

Loading

pytorch-bot Bot commented Jun 8, 2026 •

edited

Loading

This PR needs a `release notes:` label

claude Bot commented Jun 8, 2026 •

edited

Loading

Bug: Vector size mismatch with `llm::Audio` constructor