perf(qwen3-14b decode): early-dispatch infra + critical-path restructuring

[ChaoWao]

Summary

Five mechanisms for minimum critical-path latency in the qwen3-14B decode layer.

1. SPMD q/k/v_proj

Originally per-tile discrete tasks so downstream rope could consume each tile's output independently and start early. Now that rope + QK-norm are fused into fa_fused (one big kernel), the tiles can no longer provide early-start benefit — all outputs are consumed atomically by the fused kernel. Consolidating to SPMD eliminates the per-tile dependency-resolution overhead (fewer edges in the task graph).

2. Merged seeds (q_seed / kv_seed / mlp_out_seed)

These zero-fill tasks have long dependency chains and excessive successor counts. Their successors all depend on fa_fused anyway, so funneling through a single merged seed + fa dep suffices. This reduces orchestrator generation, scheduler dependency resolution, and AICore launch overhead.

3. seed_dummy barrier

seed_dummy (unflagged task_dummy) gates rms_recip / kv_seed / mlp_out_seed so they are NOT early-dispatched. At the moment prev_out completes, only q_proj should get the dispatch window (it is fa's longest predecessor). Without the barrier, kv/mlp seeds would compete for the window, delaying q_proj.

4. allow_early_resolve on critical path

Flags the critical path (q_proj → fa_fused → out_proj → next-layer dcr_xgamma) so each stage propagates dispatch_fanin to the next, enabling pre-staging before producers truly complete. Fa's other producers (rms_recip, k/v_proj, fa_work_build, mlp_out_seed, q_seed, down_proj) are also flagged so fa's dispatch_fanin can reach fanin_actual_count.

5. out_proj 24/26 critical-path split

Of the 50 out_proj tasks, ~24 feed the critical downstream chain (residual_rms_cast → gate/up_proj). These 24 connect directly to fa (deps=[attn\_done\_tid]) for first-wave dispatch. The remaining 26 go through out\_proj\_dummy (task_dummy, unflagged → no early dispatch) to defer their scheduling priority, ensuring the critical 24 execute first.

Test

• python decode_fwd.py -p a2a3 -d <dev> --validate-fwd --fwd-layers 2 --enable-l2-swimlane — golden argmax passes.

🤖 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: bc7eb971-32a1-4c15-9979-d80a065e757e

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 restructures the _decode_layer task-graph in models/qwen3/14b/decode_fwd.py to shorten the critical path. It adds scheduling barriers, converts several tasks to pl.spmd kernels, enables early resolution on multiple stages, and splits out_proj dependency wiring for prioritized dispatch.

Changes

Decode layer scheduling rework

Layer / File(s)	Summary
Critical-path design documentation models/qwen3/14b/decode_fwd.py	Adds a comment block documenting five scheduling mechanisms and the out_proj critical-path split strategy.
Seed barrier and q/k/v projection restructuring models/qwen3/14b/decode_fwd.py	Introduces seed_dummy barrier gating kv/mlp seeds, updates rms_recip deps and early-resolve; converts q_seed/kv_seed and q_proj/k_proj/v_proj to pl.spmd kernels with explicit deps; updates rope dependency task ID propagation and mlp_out_seed early-resolve.
Fused attention early resolve models/qwen3/14b/decode_fwd.py	Enables allow_early_resolve=True on the fa_fused dispatch.
Out_proj critical-path split models/qwen3/14b/decode_fwd.py	Replaces uniform out_proj dependency wiring with a split using out_proj_dummy for deferred tasks and direct attn_done_tid deps for a prioritized subset.
Down_proj and dcr_xgamma early resolve models/qwen3/14b/decode_fwd.py	Sets allow_early_resolve=True on down_proj and the consolidated dcr_xgamma dispatch.

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

Sequence Diagram(s)

sequenceDiagram
  participant SeedDummy as seed_dummy
  participant QSeed as q_seed
  participant QProj as q_proj (spmd)
  participant KVSeed as kv_seed
  participant KProj as k_proj (spmd)
  participant VProj as v_proj (spmd)
  participant FaFused as fa_fused
  participant OutProjDummy as out_proj_dummy
  participant OutProj as out_proj tasks
  participant DownProj as down_proj

  SeedDummy->>KVSeed: gates dispatch
  QSeed->>QProj: allow_early_resolve
  KVSeed->>KProj: deps + early dispatch prevention
  KVSeed->>VProj: deps + early dispatch prevention
  QProj-->>FaFused: rope dependency task ids
  KProj-->>FaFused: rope dependency task ids
  VProj-->>FaFused: rope dependency task ids
  FaFused->>OutProjDummy: attn_done_tid
  FaFused->>OutProj: direct deps (priority subset)
  OutProjDummy->>OutProj: deferred subset
  OutProj->>DownProj: allow_early_resolve

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#489: Both PRs modify the attention-completion-to-out_proj dependency wiring in the same decode task graph.
• hw-native-sys/pypto-lib#602: Both PRs touch the Q/K/RoPE decode pipeline with early-dispatch/allow_early_resolve changes.
• hw-native-sys/pypto-lib#656: Both PRs target the same fa_fused dependency wiring and early-resolve behavior.

Suggested labels: enhancement

Poem

A rabbit hops through task and seed,
Untangling deps at critical speed,
q_proj darts first, kv trails behind,
out_proj splits — the fast and the timed.
🐇⚡ Early resolve, hop, hop, done!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change: early-dispatch infrastructure and critical-path restructuring for qwen3-14b decode.
Description check	✅ Passed	The description accurately summarizes the same perf-focused decode-layer changes and test result.
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 restructures the decode layer of the Qwen3-14B model to minimize critical-path latency. Key changes include consolidating q/k/v_proj to SPMD, merging zero-fill seeds, introducing a seed_dummy barrier to prevent early-dispatch competition, flagging critical path stages with allow_early_resolve, and splitting out_proj tasks into critical and deferred waves. Feedback focuses on improving code robustness and maintainability by replacing hardcoded task dependency indices with dynamic list comprehensions and refactoring duplicated loop logic in the out_proj split into an inline helper function.

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]

Hardcoding the indices 0 to 4 of prev_out_tids is fragile and will break if DOWN_ON changes. Consider using a single list comprehension to dynamically construct the dependency list up to DOWN_ON and append seed_dummy without using list concatenation.

Suggested change

	deps=[prev_out_tids[0], prev_out_tids[1], prev_out_tids[2], prev_out_tids[3], prev_out_tids[4], seed_dummy],
	deps=[prev_out_tids[i] if i < DOWN_ON else seed_dummy for i in range(DOWN_ON + 1)],

[gemini-code-assist]

Hardcoding the indices 0 to 4 of prev_normed_tids is fragile and will break if DOWN_ON changes. Consider using a single list comprehension to dynamically construct the dependency list up to DOWN_ON and append q_seed_tid without using list concatenation.

Suggested change

	deps=[prev_normed_tids[0], prev_normed_tids[1], prev_normed_tids[2], prev_normed_tids[3], prev_normed_tids[4], q_seed_tid],
	deps=[prev_normed_tids[i] if i < DOWN_ON else q_seed_tid for i in range(DOWN_ON + 1)],

[gemini-code-assist]

Hardcoding the indices 0 to 4 of prev_out_tids is fragile and will break if DOWN_ON changes. Consider using a single list comprehension to dynamically construct the dependency list up to DOWN_ON and append seed_dummy without using list concatenation.

Suggested change

	with pl.at(level=pl.Level.CORE_GROUP, name_hint="kv_seed", deps=[prev_out_tids[0], prev_out_tids[1], prev_out_tids[2], prev_out_tids[3], prev_out_tids[4], seed_dummy]) as kv_seed_tid:
	with pl.at(level=pl.Level.CORE_GROUP, name_hint="kv_seed", deps=[prev_out_tids[i] if i < DOWN_ON else seed_dummy for i in range(DOWN_ON + 1)]) as kv_seed_tid:

[gemini-code-assist]

Hardcoding the indices 0 to 4 of prev_out_tids is fragile and will break if DOWN_ON changes. Consider using a single list comprehension to dynamically construct the dependency list up to DOWN_ON and append seed_dummy without using list concatenation.

Suggested change

	with pl.at(level=pl.Level.CORE_GROUP, name_hint="mlp_out_seed", allow_early_resolve=True, deps=[prev_out_tids[0], prev_out_tids[1], prev_out_tids[2], prev_out_tids[3], prev_out_tids[4], seed_dummy]) as mlp_out_seed_tid:
	with pl.at(level=pl.Level.CORE_GROUP, name_hint="mlp_out_seed", allow_early_resolve=True, deps=[prev_out_tids[i] if i < DOWN_ON else seed_dummy for i in range(DOWN_ON + 1)]) as mlp_out_seed_tid:

[gemini-code-assist]

Hardcoding the indices 0 to 4 of prev_normed_tids is fragile and will break if DOWN_ON changes. Consider using a single list comprehension to dynamically construct the dependency list up to DOWN_ON and append kv_seed_tid without using list concatenation.

Suggested change

	deps=[prev_normed_tids[0], prev_normed_tids[1], prev_normed_tids[2], prev_normed_tids[3], prev_normed_tids[4], kv_seed_tid],
	deps=[prev_normed_tids[i] if i < DOWN_ON else kv_seed_tid for i in range(DOWN_ON + 1)],

[gemini-code-assist]

Hardcoding the indices 0 to 4 of prev_normed_tids is fragile and will break if DOWN_ON changes. Consider using a single list comprehension to dynamically construct the dependency list up to DOWN_ON and append kv_seed_tid without using list concatenation.

Suggested change

	deps=[prev_normed_tids[0], prev_normed_tids[1], prev_normed_tids[2], prev_normed_tids[3], prev_normed_tids[4], kv_seed_tid],
	deps=[prev_normed_tids[i] if i < DOWN_ON else kv_seed_tid for i in range(DOWN_ON + 1)],

[Little-oil]

fixed

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

models/qwen3/14b/decode_fwd.py (1)

932-933: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Derive the split from the tile count rather than hardcoding 24.

N_OUT_DIRECT = N_SPLITS_OUT * K_SPLITS_OUT - 24 bakes in the assumption that there are exactly 50 out_proj tasks (the "First 26" comment). If N_SPLITS_OUT or K_SPLITS_OUT are retuned, the 24/26 balance drifts silently, and a total below 24 makes N_OUT_DIRECT negative. Consider a named constant or a fraction of the total so the split scales with the tiling.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@models/qwen3/14b/decode_fwd.py` around lines 932 - 933, The out-proj task
split in decode_fwd.py is hardcoded with a fixed 24-task remainder, which will
drift or go negative if the tiling changes. Update the logic around N_OUT_DIRECT
in the decode path to derive the direct-vs-first split from the total tile count
(N_SPLITS_OUT * K_SPLITS_OUT) using a named constant or computed fraction
instead of embedding 24, and ensure the calculation stays valid for smaller or
retuned split configurations.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@models/qwen3/14b/decode_fwd.py`:
- Around line 526-530: The q_proj stage is missing the allow_early_resolve flag,
which breaks propagation through the primary attention path. Update the pl.spmd
call in q_proj within decode_fwd.py to include allow_early_resolve=True,
matching the neighboring attention stages so dispatch_fanin can continue through
q_proj → fa_fused → out_proj → dcr_xgamma.

---

Nitpick comments:
In `@models/qwen3/14b/decode_fwd.py`:
- Around line 932-933: The out-proj task split in decode_fwd.py is hardcoded
with a fixed 24-task remainder, which will drift or go negative if the tiling
changes. Update the logic around N_OUT_DIRECT in the decode path to derive the
direct-vs-first split from the total tile count (N_SPLITS_OUT * K_SPLITS_OUT)
using a named constant or computed fraction instead of embedding 24, and ensure
the calculation stays valid for smaller or retuned split configurations.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: c2020571-38a8-454f-b4af-86a637fb7b6e

📥 Commits

Reviewing files that changed from the base of the PR and between bbd7d93 and eab322e.

📒 Files selected for processing (1)

• models/qwen3/14b/decode_fwd.py

[ChaoWao]

Note: the seed_dummy barrier must carry a real dependency — pypto elides an empty-deps task_dummy, so it never gates

Verified this PR's early-dispatch barrier empirically (2-layer decode_fwd --validate-fwd --fwd-layers 2 --enable-l2-swimlane 4 on a2a3, against direct-only early-dispatch simpler). It works as intended — kv_seed / mlp_out_seed / rms_recip are correctly kept off early-dispatch. But the mechanism is subtle and worth a note so it isn't accidentally broken.

How the gate works. Under direct-only early-dispatch, a consumer becomes an early-dispatch candidate only when every producer in its fanin is flagged (allow_early_resolve); a single unflagged producer disqualifies it. seed_dummy gates its consumers by being an unflagged producer edge in their fanin — but only if that edge actually exists.

Why deps=[prev_out_tids…] is load-bearing (a codegen behavior, not a device-runtime one). pypto generates a phase-fence barrier dummy behind a guard: submit the dummy only if it collected ≥1 dependency, and add it to a consumer only if its id is valid. From the generated decode_fwd.cpp:

uint32_t ..._deps_count = 0;                 // deps=[] -> stays 0
PTO2TaskId seed_dummy = PTO2TaskId::invalid();
if (..._deps_count > 0) {                     // 0 > 0 -> false
    rt_submit_dummy_task(...);                 // never runs
    seed_dummy = ...task_id();
}
prev_out_seed_deps[5] = seed_dummy;           // = invalid()
// each consumer: if (seed_dummy.is_valid()) params_deps[c++] = seed_dummy;  // skipped

So seed_dummy = pl.system.task_dummy(deps=[]) is never submitted and never becomes a dependency edge — the barrier is a silent no-op. Confirmed three ways:

• Real run (no dep_gen, swimlane): all three seeds early-dispatch — barrier defeated.
• deps.json: the seed_dummy node + its 3 edges are absent (dummy-task count 4→2); kv_seed's producers drop from [copy_hidden, DUMMY] → [copy_hidden]. Faithful — the runtime built no such task/edge.
• Instrumented simpler (temporary LOG_INFO_V9 in the orchestrator): a dummy-submit marker fires for the two out_proj barriers only, never for seed_dummy; and the "explicit dep skipped because producer already reclaimed" path (dep_local_task_id < last_task_alive) fires 0 times — so the device-runtime last-task-alive logic is not involved.

A barrier dummy with no upstream deps fences nothing (ready at t=0, imposes no ordering), so pypto drops it. By depending on prev_out_tids, seed_dummy collects ≥1 dep → is actually submitted → becomes the unflagged fanin edge that disqualifies its consumers.

Implication. The barrier only exists as long as it carries a real dependency. seed_dummy's prev_out deps are what make it a submitted, gating task — a future refactor must not "simplify" it to deps=[]. (More robustly, a "never early-dispatch these seeds" intent could be a producer-side opt-out flag carried on the task itself, rather than a dependency-only dummy.)

Suggest a one-line comment at the seed_dummy definition noting its prev_out dependency is required — an empty-deps task_dummy is elided by codegen and gates nothing.

[ChaoWao]

#762