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.

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

⚠️ 16 Awaiting Approval

As of commit 776ee17 with merge base ed65b12 ():

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

• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows Wheels (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML 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.