Refactor Qwen3 decode forward path

[lyfne123]

Summary

• Rename the Qwen3 fused decode entry to decode_fwd.py
• Keep token embedding and greedy sampling on the inline SPMD path
• Feed final RMSNorm from the FP32 carry before the LM head
• Remove obsolete standalone greedy/token/L3 generate entries

Related Issues

[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: 10b40217-f98f-4095-95be-105b60282411

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

This PR propagates FP32 hidden states through the LM-head computation across Qwen3-14B decode and prefill kernels, removing BF16 casts before rms_lm_head. It also refactors decode_fwd's token embedding and greedy sampling to use pl.spmd, removes token_embed.py and qwen3_14b_l3_generate.py, and updates docs/tests/comments from decode_layer.py to decode_fwd.py.

Changes

FP32 carry into rms_lm_head

Layer / File(s)	Summary
rms_lm_head FP32 signature and RMSNorm loops models/qwen3/14b/rms_lm_head.py	hidden_states parameter changes from BF16 to FP32; explicit FP32 casts removed from both RMSNorm loops since input is already FP32.
decode_fwd LM-head call uses FP32 cur directly models/qwen3/14b/decode_fwd.py	Removes tiled BF16 cast (cur_bf16) and calls rms_lm_head(cur, ...) directly with FP32; comments clarify FP32 carry handling at the last layer.
decode_fwd_tq, prefill_fwd, prefill_fwd_tq final_hidden cast to FP32 models/qwen3/14b/qwen3_14b_decode_tq.py, models/qwen3/14b/prefill_fwd.py, models/qwen3/14b/qwen3_14b_prefill_tq.py	final_hidden buffers allocated as FP32 and populated via FP32-cast slices before calling rms_lm_head.

decode_fwd SPMD refactor and file cleanup

Layer / File(s)	Summary
Token embedding SPMD loop models/qwen3/14b/decode_fwd.py	_token_embed_inline switches from pl.parallel/pl.at(CORE_GROUP) to pl.spmd(BATCH, name_hint="token_embed").
Greedy sampling SPMD loop models/qwen3/14b/decode_fwd.py	_greedy_sample_inline restructured to pl.spmd(BATCH, name_hint="greedy_sample"), covering chunked top-k selection, cross-chunk reduction, and winner-token write logic.
Removed standalone files and decode_layer → decode_fwd references models/qwen3/14b/token_embed.py, models/qwen3/14b/qwen3_14b_l3_generate.py, docs/compile-runtime-workflow.md, models/qwen3/14b/config.py, tests/golden/test_qwen3_greedy_source.py	token_embed.py and qwen3_14b_l3_generate.py are deleted; docs, comments, and tests are updated to reference decode_fwd.py instead of decode_layer.py.

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

Possibly related issues

• [Tracking] Perf comparison update: Qwen3-14B single decode layer (decode_layer.py) vs ASC NPU small operators (seq=4096) #465: Concerns the Qwen3-14B decode pipeline's FP32 carry/final RMSNorm handling in decode_fwd/rms_lm_head, directly addressed by this PR's changes.

Possibly related PRs

• hw-native-sys/pypto-lib#387: Both modify models/qwen3/14b/rms_lm_head.py around the same final-norm/LM-head kernel.
• hw-native-sys/pypto-lib#449: Both heavily modify models/qwen3/14b/decode_fwd.py's decode pipeline structure.
• hw-native-sys/pypto-lib#639: Directly related—removes token_embed.py (which the retrieved PR added) and updates decode_layer→decode_fwd references.

Poem

Hop, hop, through FP32 fields I bound,
No more BF16 casts weighing me down,
SPMD loops now spin so neat,
Sampling tokens, swift and sweet,
Old files gone, new paths I've found! 🐇✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main refactor: updating the Qwen3 decode forward path and related decode-side kernels.
Description check	✅ Passed	The description is directly related to the changeset and summarizes the key refactor and removals accurately.
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 Qwen3-14B generation pipeline by transitioning the final RMSNorm and LM-head inputs to FP32, which removes redundant BF16 casting and round-trips. Additionally, it cleans up the codebase by deleting several standalone files (greedy_sample.py, qwen3_14b_generate_l3.py, and token_embed.py) in favor of inline implementations, updates documentation and tests, and adopts pl.spmd in inline functions. The review feedback suggests a valuable refactoring to eliminate duplicated loop logic when scanning winning logits in decode_fwd.py, along with a corresponding update to the test assertions.

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.