Fix Qwen3.5 static-shape prefill (#18832)

[yenhao-huang]

Summary

Fixes #18832

The Python native Llama runner always sent the complete prompt in a single
prefill invocation. This fails for static-shape KV-cache PTEs whose token input
is fixed to shape [1, 1].

The C++ TextPrefiller handles this configuration by reading
enable_dynamic_shape from the PTE metadata
(extension/llm/runner/llm_runner_helper.cpp:94-129) and passing it as
enable_parallel_prefill when constructing TextPrefiller
(extension/llm/runner/llm_runner_helper.cpp:320-324).

TextPrefiller::prefill_chunk() uses this value to select a prefill strategy
that is consistent with the exported graph's input contract and KV-cache state
(extension/llm/runner/text_prefiller.cpp:111-164). The Python runner now
reads the contract during initialization
(examples/models/llama/runner/native.py:38-52) and applies it at prefill time
(examples/models/llama/runner/generation.py:94-136).

The Python runner changes are:

• Load the PTE before initializing the base runner and read enable_dynamic_shape,
  defaulting to false when the metadata method is not present.
• Add enable_dynamic_shape to LlamaRunner and route prompt processing
  through a dedicated _prefill_chunk() method.
• Preserve the existing full-prompt prefill when dynamic shapes are enabled or
  the KV cache is disabled.
• For static-shape KV-cache models, invoke the model once per prompt token with
  shape [1, 1] and increment input_pos for each invocation.

Regression tests cover static-shape KV-cache prefill, dynamic-shape parallel
prefill, and the non-KV-cache path.

Test plan

• Run lintrunner on the affected files:

  lintrunner \
    examples/models/llama/runner/generation.py \
    examples/models/llama/runner/native.py \
    examples/models/llama/tests/BUCK \
    examples/models/llama/tests/test_generation.py

  Result: No lint issues

• Run the focused generation regression tests:

  PYTHONPATH="$PWD/src" python -m pytest -q \
    examples/models/llama/tests/test_generation.py

  Result: 3 passed

• Run the Qwen3.5 static-shape PTE end to end with the KV cache enabled:

  env -u PYTHONPATH OMP_NUM_THREADS=8 timeout 300 \
    /data/executorch_venv_18832_fixed/bin/python -m \
    executorch.examples.models.llama.runner.native \
    --model qwen3_5_0_8b \
    --pte /data/executorch-issues/18832/outputs/qwen3_5_0_8b_fp32.pte \
    --tokenizer /models/Qwen-Qwen3.5-0.8B/tokenizer.json \
    --tokenizer_config /models/Qwen-Qwen3.5-0.8B/tokenizer_config.json \
    --prompt $'<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n' \
    --params examples/models/qwen3_5/config/0_8b_config.json \
    --max_len 128 \
    -kv \
    --temperature 0.3

  Result: Generated a response successfully without the input-shape error.

cc @mergennachin @iseeyuan @lucylq @helunwencser @tarun292 @kimishpatel @jackzhxng @larryliu0820 @cccclai @digantdesai

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 21 Awaiting Approval

As of commit d2bc53d with merge base 20fa24e ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux CUDA Wheels (gh)
• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux CUDA Wheels (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows Wheels (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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.

[nil-is-all]

Thanks for the PR, @yenhao-huang. Re-running CI tests
@digantdesai could you review this?

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 1m 5s —— View job

---

Summary: The fix is correct. It matches the C++ TextPrefiller behavior and its false default for enable_dynamic_shape, and I found one small simplification and nothing that blocks merging.

Detailed review

PR Review: #21873

Summary

The Python native Llama runner now reads enable_dynamic_shape from the PTE. When that is false and the KV cache is on, it prefills one [1, 1] token at a time, matching TextPrefiller::prefill_chunk() in extension/llm/runner/text_prefiller.cpp. Verdict: correct. The only finding is a minor code-quality point.

Code Quality

• examples/models/llama/runner/generation.py:113-136: the sequential branch copies the C++ structure: it runs prompt_tokens[0] on its own, then loops over prompt_tokens[1:] with enumerate(..., start=1). The C++ version splits out the first token because of how it reuses the from_blob tensor and its "first token is BOS" comment. Neither reason applies in Python, so a single loop gives the same calls with about half the code:

  for offset, token in enumerate(prompt_tokens):
      logits = self.forward(
          tokens=torch.tensor([[token]], dtype=torch.long, device=self.device),
          input_pos=torch.tensor([start_pos + offset], dtype=torch.long, device=self.device),
      )
  return logits

  The else: after an early return can go too. The existing tests ([([10], 7), ([11], 8), ([12], 9), ...]) should pass unchanged. Fix this →

Recommendation

Approve. The simplification above is optional.

[digantdesai]

can we fix c issues, LGTM

[yenhao-huang]

Hi @digantdesai, all review feedback addressed and ready to merge. Thanks!

[nil-is-all]

Thanks for the ping @yenhao-huang. Could you update the branch and give a ping?

[yenhao-huang]

@nil-is-all Branch updated. Thanks.