Fix Qwen3 attention mask NaNs

[superxf]

Summary

• Replace Qwen3-14B attention fillpad min masks with a finite negative sentinel before row-wise softmax.
• Add pypto-serving's Qwen3 serving guard to the cross-repo CI action.

Related Issues

Fixes #713

[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: 1e8f0850-47f6-447c-a991-a6010f602962

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

Replaces pl.fillpad(..., pad_value=pl.PadValue.min) with an explicit finite ATTN_MASK_NEG_INF = -3.0e38 constant for masking attention scores in Qwen3 14B prefill and decode attention paths, avoiding NaNs from fully invalid padded rows. Updates the CI serving-tests action to also run the Qwen3 serving test.

Changes

Attention Masking Fix

Layer / File(s)	Summary
Prefill attention score masking models/qwen3/14b/prefill_fwd.py	Adds ATTN_MASK_NEG_INF constant; scores tensor now built by filling with the constant then assembling in scaled valid scores, replacing fillpad(PadValue.min).
Decode attention score masking models/qwen3/14b/decode_layer.py	Adds same ATTN_MASK_NEG_INF constant; fused attention loop now builds a full FP32 score tile filled with the constant and assembles valid scores into it, replacing fillpad(PadValue.min).
CI serving test coverage .github/actions/pypto-serving-tests/action.yml	Pytest invocation now also runs tests/test_qwen3_serving.py alongside the existing accuracy test.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• hw-native-sys/pypto-lib#51: Both PRs modify the same Qwen3 decode attention softmax masking logic, changing how invalid/padded score entries are filled.

Poem

A rabbit hopped through -inf's shade,
where NaNs in padded rows were made.
With -3e38 held tight,
the softmax rows now stay just right. 🐇
No more glitches, hop hop hooray!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately reflects the main fix to Qwen3 attention mask NaNs.
Description check	✅ Passed	The description is clearly related to the changeset and matches the stated Qwen3 mask fix and CI update.
Linked Issues check	✅ Passed	The PR replaces PadValue.min masks with a finite sentinel in Qwen3 prefill and decode, addressing the NaN issue.
Out of Scope Changes check	✅ Passed	All changes map to the stated fix and CI guard; no unrelated edits are evident.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

---

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 adds a new serving test to the GitHub Actions workflow and updates the Qwen3 14B model's decode and prefill forward layers to replace pl.fillpad with pl.full and pl.assemble when preparing attention scores. The review feedback points out that because scores is now a scratch tensor (TensorType), it must be loaded back as a Vec tile using pl.slice before calling vector instructions like pl.row_max and pl.row_expand_sub.

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]

In PyPTO, vector instructions such as pl.row_max and pl.row_expand_sub require their arguments to be a Vec TileType, not a TensorType. Since scores is initialized via pl.full and updated via pl.assemble, it represents a TensorType (scratch tensor). To use it in vector operations, you must first load it back as a Vec tile using pl.slice.

Suggested change

	cur_mi = pl.row_max(scores)
	exp_scores = pl.exp(pl.row_expand_sub(scores, cur_mi))
	scores_tile = pl.slice(scores, [Q_HEAD_PAD, SEQ_TILE], [0, 0])
	cur_mi = pl.row_max(scores_tile)
	exp_scores = pl.exp(pl.row_expand_sub(scores_tile, cur_mi))

References

1. In PyPTO, vector instructions require a Vec TileType rather than a TensorType. When assembling to a scratch tensor, the value must be loaded back as a Vec tile before being consumed by vector operations.

[gemini-code-assist]

In PyPTO, vector instructions such as pl.row_max and pl.row_expand_sub require their arguments to be a Vec TileType, not a TensorType. Since scores is initialized via pl.full and updated via pl.assemble, it represents a TensorType (scratch tensor). To use it in vector operations, you must first load it back as a Vec tile using pl.slice.

Suggested change

	cur_mi = pl.row_max(scores)
	exp_scores = pl.exp(pl.row_expand_sub(scores, cur_mi))
	scores_tile = pl.slice(scores, [Q_HEAD_BATCH_PAD, SEQ_TILE], [0, 0])
	cur_mi = pl.row_max(scores_tile)
	exp_scores = pl.exp(pl.row_expand_sub(scores_tile, cur_mi))

References

1. In PyPTO, vector instructions require a Vec TileType rather than a TensorType. When assembling to a scratch tensor, the value must be loaded back as a Vec tile before being consumed by vector operations.

[coderabbitai]

🧹 Nitpick comments (1)

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

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

Constant duplicated across two files.

ATTN_MASK_NEG_INF is defined identically here and in decode_layer.py. Worth hoisting to a shared module to avoid future drift between the two sentinel values.

🤖 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/prefill_fwd.py` at line 121, The ATT N mask sentinel is
duplicated in both prefill_fwd.py and decode_layer.py, so hoist ATT
N_MASK_NEG_INF into a shared module and update both callers to import the same
constant. Use the existing ATT N_MASK_NEG_INF symbol in prefill_fwd.py and
decode_layer.py as the reference points, and ensure both paths read the value
from one source so the sentinel cannot drift.

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

Nitpick comments:
In `@models/qwen3/14b/prefill_fwd.py`:
- Line 121: The ATT N mask sentinel is duplicated in both prefill_fwd.py and
decode_layer.py, so hoist ATT N_MASK_NEG_INF into a shared module and update
both callers to import the same constant. Use the existing ATT N_MASK_NEG_INF
symbol in prefill_fwd.py and decode_layer.py as the reference points, and ensure
both paths read the value from one source so the sentinel cannot drift.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 4f908f12-0f7e-4738-bcf8-7f86dda05a18

📥 Commits

Reviewing files that changed from the base of the PR and between c159c32 and d87ceb1.

📒 Files selected for processing (3)

• .github/actions/pypto-serving-tests/action.yml
• models/qwen3/14b/decode_layer.py
• models/qwen3/14b/prefill_fwd.py

[superxf]

Update: I withdrew the model-level workaround from this PR after validating the CI failure locally.

Evidence:

• Original pypto-serving baseline passes: tests/test_qwen3_accuracy.py -> 1 passed on NPU.
• Replacing fillpad(PadValue.min) with finite full + assemble avoids the NaN path but changes Qwen3 baseline tokens in CI.
• Follow-up model-level attempts (cur_mi zeroing, exp_scores zeroing, and pl.maximum clamp after fillpad) either hit ptoas alignment issues or fail warmup with AICore 507018.

Conclusion: this should be fixed in PyPTO core fillpad/PadValue lowering, not with a Qwen3 model-side rewrite in pypto-lib. The branch is now an empty diff against main so it no longer carries the CI-breaking changes.

[superxf]

Update: re-tested the minimal clamp fix without PTO2_RING_HEAP, PTO2_RING_TASK_WINDOW, and PTO2_RING_DEP_POOL.

Current patch:

• keep fillpad(PadValue.min) unchanged;
• add scores = pl.maximum(scores, -3.4028234663852886e38) before row-wise softmax in Qwen3-14B prefill/decode.

Validation:

• tests/test_qwen3_accuracy.py: passed (1 passed, no AICore 507018) when run without the three PTO2_RING_* env overrides.
• tests/test_qwen3_serving.py::test_prefix_cache_reuses_prefix_and_preserves_output: pytest passed (1 passed). The task wrapper returned 137 afterwards because taskqueue reported daemon stop/restart, but the captured pytest result was successful.

This keeps the original fillpad valid-shape path and only prevents the -inf - -inf -> NaN case on padded rows.