perf(qwen3): compile prefill/decode/greedy straight from the signature

[lyfne123]

What

The three Qwen3-14B HOST wrappers (qwen3_prefill_host, qwen3_decode_host, qwen3_greedy_sample_host) were compiled by hand-building large throwaway torch.empty(...) tensors that .compile() only inspects for shape/dtype — the data is never read, yet the memory is really allocated. Several are large:

• prefill/decode k_cache/v_cache dummies are sized like a full KV cache (num_layers·batch·cache_blocks·kv_heads·page_size × head_dim),
• plus lm_head/embed (VOCAB × HIDDEN) and the stacked projection weights.

So every model compile wasted that much host RAM (and time), and re-declared the shape contract already implied by the kernels.

With pypto #2014 (compile() straight from the signature), a fully-annotated wrapper is compiled with no sample tensors — the shape/dtype contract is read from the parameter annotations.

Changes

• qwen3_l3_dispatch.py — full pl.Tensor[[...], dtype] annotations on all three wrappers. Static extents come from validated model constants; runtime-varying dims reuse the kernels' shared pl.dynamic vars so host↔kernel dims unify. Dep-graph inference still re-marks kernel-bound dims dynamic, so the compiled artifact is identical to the old compile(sample_tensors) call. Decode's first return is bound to the out param name so signature-mode specialization can fill the return-tuple element type (a renamed local stays a bare pl.Tensor, which the parser rejects inside tuple[...]).
• npu_executor.py — inject the annotation globals (model constants + the kernels' pl.dynamic instances) before compiling; delete the three dummy-arg factory methods (_compile_{prefill,decode,greedy_sample}_*_callable); compile() with no tensors (−208 / +156 lines net in the two runner files).
• test_batching.py — refresh two tests that referenced the removed decode hidden buffer / the pre-fusion decode kernel arg order.

Verification

• Artifact equivalence (offline, codegen_only): for each wrapper, signature-mode compile() returns the same cached DistributedCompiledProgram as the old compile(*dummy_tensors) — proving the annotations reproduce the exact static/dynamic shape contract.
• Unit tests: tests/test_batching.py + tests/test_parallel.py → 28 passed.
• Lint: ruff / headers / english-only clean.
• On-device e2e: npu_generate.py --max-seq-len 512 --max-new-tokens 5 (the runtime-max_seq < kernel MAX_SEQ case) → correct coherent output.

Dependency

Requires pypto with #2014 (signature-mode compile()).

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: b1718f69-acdc-4b33-a17a-ff61fa1fb8a2

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Qwen3 host wrappers now use explicit tensor shape contracts and signature-based JIT compilation. Dispatch dimensions are injected before compilation, specialized dummy arguments are removed, and batching tests reflect fused decode inputs and flexible copy-kernel signatures.

Changes

Qwen3 signature compilation

Layer / File(s)	Summary
Dispatch shape contracts examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py	Prefill, decode, and greedy-sampling wrappers now declare explicit tensor shapes, dtypes, dynamic dimensions, and typed outputs.
Signature compilation wiring examples/model/qwen3_14b/runner/npu_executor.py, tests/test_batching.py	The executor injects dispatch shapes and compiles host wrappers without dummy tensors; batching tests reflect fused decode inputs and variable dispatch signatures.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

• hw-native-sys/pypto-serving#68: Updates the same Qwen3 executor and decode dispatch wiring around the fused decode path.

Sequence Diagram(s)

sequenceDiagram
  participant Qwen314BPyptoExecutor
  participant qwen3_l3_dispatch
  participant prefill_kernel
  Qwen314BPyptoExecutor->>qwen3_l3_dispatch: Inject static shapes and dynamic dimensions
  qwen3_l3_dispatch->>prefill_kernel: Read loaded kernel dynamic-dimension objects
  Qwen314BPyptoExecutor->>qwen3_l3_dispatch: Compile typed host wrappers by signature

Loading

Poem

A rabbit hops through shapes so bright,
While kernels compile by type tonight.
No dummy tensors block the way,
Decode hops on fused lanes of gray.
Qwen3 bounds the path just right—
“Squeak!” says the bunny, “ship the flight!”

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: compiling Qwen3 wrappers directly from their signatures.
Description check	✅ Passed	The description matches the changeset, explaining the signature-based compilation, executor update, and test adjustments.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request refactors the compilation of the Qwen3-14B host wrappers to use signature-mode compilation instead of compiling with dummy tensors. It introduces _inject_dispatch_shapes to dynamically bind shape and dynamic dimension variables to the dispatch module's globals prior to compilation, and updates the host wrappers with explicit tensor annotations. The feedback suggests wrapping the dynamic attribute lookup in a try-except block to provide a clearer error message if any expected dynamic dimension variables are missing from the prefill module.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When dynamically retrieving dynamic dimension variables from the loaded prefill module, if any of the expected attributes are missing (e.g., due to a version mismatch between the runner and the kernel module), a generic AttributeError will be raised. Wrapping this lookup in a try-except block to raise a clearer, more descriptive error message will greatly improve debuggability.

        for dyn_name in (
            "USER_BATCH_DYN",
            "PREFILL_TOKENS_DYN",
            "KV_CACHE_ROWS_DYN",
            "BLOCK_TABLE_FLAT_DYN",
            "LAYER_DYN",
            "LAYER_HIDDEN_ROWS_DYN",
            "LAYER_INTER_ROWS_DYN",
        ):
            try:
                dyn_var = getattr(prefill_module, dyn_name)
            except AttributeError as e:
                raise AttributeError(
                    f"Prefill kernel module is missing required dynamic dimension variable '{dyn_name}'. "
                    "Please ensure your kernel module is up to date."
                ) from e
            setattr(qwen3_l3_dispatch, dyn_name, dyn_var)