Fix QNN runner KV cache bitwidth in Android JNI

[infil00p]

Summary

The QNN runner was hardcoded to use Runner<uint16_t>, but all current Llama quantization recipes use annotate_kv_8bit for 8-bit KV cache. This mismatch caused the KV cache data to be misinterpreted, resulting in degenerate output (repetitive real words like "Nigeria Nigeria...") while the model otherwise ran correctly on the HTP NPU.

Authored with Claude

Test plan

I built the AAR and used the AAR in the LlamaDemo in executorch-examples. This does require QAIRT 2.43.0 to test, which is later than the last QAIRT used to build the last Executorch release. This was tested on a OnePlus 15 running a Snapdragon 8 Elite Gen 5 (SM8850).

cc @kirklandsign @cbilgin @cccclai

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 7 New Failures, 1 Cancelled Job, 5 Unrelated Failures

As of commit af1f892 with merge base 298311e ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner / linux-job (gh)
  >>> Lint for extension/android/jni/jni_layer_llama.cpp:
• pull / test-mediatek-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 13e6317b64440c11ac079be38a04f5d459d3076cd7b417b6597ee414d9760798 /exec failed with exit code 2
• pull / test-openvino-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 5c28183eeb078144ba4c98072e2623263021a7987d59d0743e2fb6ab998596ab /exec failed with exit code 1
• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t d5cd45c752afdf7decc6f8528a37493a737396f313eeef7581da958747c424a0 /exec failed with exit code 1
• pull / test-samsung-quantmodels-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 5762e91d59e9852c32abf28059375dedd1968840014defe77110246e872eb4b6 /exec failed with exit code 1
• pull / unittest-buck / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 830970ce33a33e546f6cab430f4585fdbaee0ccd3959377772a275bf1388b485 /exec failed with exit code 3
• pull / unittest-buck / macos / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 3

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

• periodic (gh)

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / test-models-linux (resnet50, portable, linux.2xlarge) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• pull / test-moshi-linux / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• Test CUDA Builds / test-models-cuda (mv2) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• Test CUDA Builds / test-models-cuda (resnet18) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• Test CUDA Builds / unittest-cuda / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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

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

[cccclai]

I think we need a better way to handle this..I remember llama uses 16 bit kv cache and qwen uses 8 bit kv cache cc: @haowhsu-quic

[haowhsu-quic]

Yes, probably need a branch here to dispatch runner correctly.

[abhinaykukkadapu]

@haowhsu-quic Ccan we detect it by dynamically querying get_kv_io_bit_width from the model if the method exists and do something like (default to 8 bit):

example::KvBitWidth kv_bitwidth = example::KvBitWidth::kWidth8;
if (module->method_names()->count("get_kv_io_bit_width") > 0) {
  kv_bitwidth = static_cast<example::KvBitWidth>(
      module->get("get_kv_io_bit_width")
          .get()
          .toScalar()
          .to<int64_t>());
}

if (kv_bitwidth == example::KvBitWidth::kWidth16) {
  runner_ = std::make_unique<example::Runner<uint16_t>>(...)
} else {
  runner_ = std::make_unique<example::Runner<uint8_t>>(...)
}

[nil-is-all]

@haowhsu-quic there an update on dynamically querying the bit-width? Is there anything to recommend for the PR's author to do?

[haowhsu-quic]

Hi @abhinaykukkadapu , we actually do this in our runner (https://github.com/pytorch/executorch/blob/main/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp).

example::KvBitWidth kv_bitwidth = example::KvBitWidth::kWidth8;
if (module->method_names()->count("get_kv_io_bit_width") > 0) {
  kv_bitwidth = static_cast<example::KvBitWidth>(
    module->get("get_kv_io_bit_width").get().toScalar().to<int64_t>());
}

Maybe PR's author can try to follow this apporach.

[nil-is-all]

@infil00p could you try this approach and request review once done?

[nil-is-all]

@infil00p bringing this up again