Add: dynamic public batch (1..16) for Qwen3-14B decode - #841
Conversation
📝 WalkthroughWalkthroughQwen3-14B now separates dynamic public batch dimensions from fixed padded pipeline dimensions. Configuration, prefill, decode, sampling, attention, RoPE kernel wiring, regeneration tooling, and contract tests are updated to use ChangesQwen3 batch ABI transition
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (6)
models/qwen3/14b/decode_layer_a8w8.py (1)
308-317: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThis path assumes
seq_lensalways hasBATCH_PADrows.
batch(Line 132) is read from theseq_lensdescriptor, butfa_work_build,rms_recip, and the quant loops allunroll(BATCH_PAD)and indexseq_lens[wb]directly. Since the A8W8 decode entry never binds a dynamic batch, that holds today; a shortseq_lenswould read out of bounds. A short docstring/comment stating the static-batch contract (asgreedy_samplenow does) would keep this from drifting once the A8W8 path follows the BF16 one.🤖 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_layer_a8w8.py` around lines 308 - 317, Document the static-batch contract for the A8W8 decode path near the batch handling or shared setup: `seq_lens` must contain `BATCH_PAD` rows because `fa_work_build`, `rms_recip`, and the quantization loops unroll and index all padded batch entries. Keep the existing `BATCH_PAD` iteration behavior unchanged and follow the concise contract documentation style used by `greedy_sample`.models/qwen3/14b/rope_qkv_regen.py (3)
360-373: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMissing the standard
-d/--deviceflag for an executable model script.Repo convention is that model scripts accept
-p {a2a3,a2a3sim,a5,a5sim}and-d <device_id>. This one narrows-p(justified in the_backend_typecomment) and omits-dentirely. Since the program is compile-only,-dmay genuinely be unnecessary — if so, state that in the argparse description so the deviation is intentional and visible.As per path instructions: "Executable model scripts generally accept
-p {a2a3,a2a3sim,a5,a5sim}and-d <device_id>."🤖 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/rope_qkv_regen.py` around lines 360 - 373, Update the __main__ argparse setup to follow the executable model-script CLI convention by adding a -d/--device device identifier option, and broaden --platform choices to include a5 and a5sim unless the existing _backend_type constraint requires the narrower set. If this compile-only script intentionally does not need a device option, state that exception explicitly in the ArgumentParser description instead.Source: Path instructions
354-357: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_backend_typeignores itsplatformargument.Both accepted choices map to
Ascend910B, so-p a2a3simbehaves identically to-p a2a3and the parameter is dead. Either drop the parameter (keeping the constant plus the comment) or make the sim path meaningful.🤖 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/rope_qkv_regen.py` around lines 354 - 357, Update _backend_type so it no longer accepts the unused platform argument, while preserving the Ascend910B return value and its A2/A3-only comment; also adjust every caller to use the parameterless function.
313-322: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDead local and duplicated artifact lookup.
names(Line 315) is computed and never used, and Lines 385-388 re-implement_newest_artifact()verbatim. Alsore.search(...).group(1)on Line 317 will raise an opaqueAttributeErrorif the launcher shape changes — a clear error would help the next person regenerating.♻️ Proposed cleanup
- launcher = src[src.index('extern "C"'):] - names = [m.split(": ", 1)[1] for m in - re.findall(r"// (?:Unpack \w+|Extract dynamic dim): .*", launcher)] - call = re.search(r"rope_qkv\((.*?)\);", launcher, re.DOTALL).group(1) - call_args = [a.strip() for a in call.split(",")] + launcher = src[src.index('extern "C"'):] + call_match = re.search(r"rope_qkv\((.*?)\);", launcher, re.DOTALL) + if call_match is None: + raise SystemExit(f"no rope_qkv forwarding call found in {artifact}") + call_args = [a.strip() for a in call_match.group(1).split(",")]- out_root = Path(__file__).parent / "build_output" - builds = sorted(out_root.glob("_jit_rope_qkv_regen_*"), key=lambda p: p.stat().st_mtime) - if builds: - artifact = builds[-1] / "kernels" / "aiv" / "rope_qkv.cpp" + artifact = _newest_artifact() + if artifact.exists():Also applies to: 385-390
🤖 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/rope_qkv_regen.py` around lines 313 - 322, Remove the unused names computation near launcher parsing, and replace the duplicated artifact lookup around _newest_artifact() with a call to that existing helper. Validate the rope_qkv regex match before accessing its capture group, raising a clear regeneration error when the launcher shape is unexpected.models/qwen3/14b/decode_fwd.py (1)
1175-1208: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
decode_fwd_layerscopy_hidden has no valid-row guard, unlikedecode_fwd.
batchis derived fromseq_lens(Line 1175) andpa_max_blocksdivides by it, so this entry can be driven with a shortseq_lens(the--validate-fwdreference path does exactly that withsl[:UB]), yetcopy_hidden/copy_outslicehidden_states/outat fullBATCH_PADwith novalid_shape. Today the harness always passesBATCH_PAD-row host tensors so it is safe, but the mismatch is easy to trip on later. Worth either asserting the static contract in the docstring or mirroringdecode_fwd'sfillpad(..., valid_shape=[batch, ...]).Also applies to: 1227-1245
🤖 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 1175 - 1208, Update decode_fwd_layers so copy_hidden and copy_out handle the runtime batch derived from seq_lens instead of assuming BATCH_PAD valid rows. Mirror decode_fwd’s fillpad behavior by supplying valid_shape based on batch when staging hidden_states and out, or explicitly assert the full BATCH_PAD-row contract in the function documentation if that contract is intentional.models/qwen3/14b/paged_attention_cce.py (1)
61-74: 🗄️ Data Integrity & Integration | 🔵 TrivialAdd a contract assertion for the metadata slot invariant.
METADATA_BATCH_SLOTSis currently only documented to matchkernel/metadata_layout.h::kLengthArrayBytesandtiling/qwen_fai_runtime_tiler.hpp::kMaxBatch, but the C++ constants still sit outside any shared/runtime check. Add coverage that fails these constants from drift, mirroring the existing contract-banner pattern.[low_effort And_high_reward]
🤖 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/paged_attention_cce.py` around lines 61 - 74, The metadata slot invariant is documented but lacks an executable drift check. Add a contract assertion near METADATA_BATCH_SLOTS, following the existing contract-banner pattern, that validates it against the fixed 16-slot metadata layout and the corresponding qwen_fai_tiler::kMaxBatch value; ensure the check fails if either C++ contract changes.
🤖 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 1811-1813: In the logits comparison block, split the
semicolon-joined assignments for a and e into separate statements on distinct
lines to satisfy ruff E702; leave the subsequent argmax comparison logic
unchanged.
In `@models/qwen3/14b/rope_qkv_regen.py`:
- Around line 260-276: Validate the --batch value against BATCH_PAD before
constructing or compiling inputs in _dummy_inputs and the surrounding entry
flow. Reject batches greater than BATCH_PAD with a clear range error, while
preserving valid batch behavior and existing tensor shapes.
In `@tests/contract/test_qwen3_14b_contract.py`:
- Around line 281-300: Update the flag-balance assertions in the contract test
to validate complete feasible execution paths rather than requiring each guarded
block in _flag_balance to be self-contained. Cover the no-block,
first-block-only, and all-blocks paths, ensuring each path’s prologue and
executed-block sets are drained by the epilogue; alternatively regenerate the
kernel so each guard is independently balanced.
---
Nitpick comments:
In `@models/qwen3/14b/decode_fwd.py`:
- Around line 1175-1208: Update decode_fwd_layers so copy_hidden and copy_out
handle the runtime batch derived from seq_lens instead of assuming BATCH_PAD
valid rows. Mirror decode_fwd’s fillpad behavior by supplying valid_shape based
on batch when staging hidden_states and out, or explicitly assert the full
BATCH_PAD-row contract in the function documentation if that contract is
intentional.
In `@models/qwen3/14b/decode_layer_a8w8.py`:
- Around line 308-317: Document the static-batch contract for the A8W8 decode
path near the batch handling or shared setup: `seq_lens` must contain
`BATCH_PAD` rows because `fa_work_build`, `rms_recip`, and the quantization
loops unroll and index all padded batch entries. Keep the existing `BATCH_PAD`
iteration behavior unchanged and follow the concise contract documentation style
used by `greedy_sample`.
In `@models/qwen3/14b/paged_attention_cce.py`:
- Around line 61-74: The metadata slot invariant is documented but lacks an
executable drift check. Add a contract assertion near METADATA_BATCH_SLOTS,
following the existing contract-banner pattern, that validates it against the
fixed 16-slot metadata layout and the corresponding qwen_fai_tiler::kMaxBatch
value; ensure the check fails if either C++ contract changes.
In `@models/qwen3/14b/rope_qkv_regen.py`:
- Around line 360-373: Update the __main__ argparse setup to follow the
executable model-script CLI convention by adding a -d/--device device identifier
option, and broaden --platform choices to include a5 and a5sim unless the
existing _backend_type constraint requires the narrower set. If this
compile-only script intentionally does not need a device option, state that
exception explicitly in the ArgumentParser description instead.
- Around line 354-357: Update _backend_type so it no longer accepts the unused
platform argument, while preserving the Ascend910B return value and its
A2/A3-only comment; also adjust every caller to use the parameterless function.
- Around line 313-322: Remove the unused names computation near launcher
parsing, and replace the duplicated artifact lookup around _newest_artifact()
with a call to that existing helper. Validate the rope_qkv regex match before
accessing its capture group, raising a clear regeneration error when the
launcher shape is unexpected.
🪄 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 Plus
Run ID: 431e6930-c371-45e4-be69-b96b63bea61b
📒 Files selected for processing (15)
models/qwen3/14b/config.pymodels/qwen3/14b/constants.pymodels/qwen3/14b/contract.pymodels/qwen3/14b/decode_fwd.pymodels/qwen3/14b/decode_layer_a8w8.pymodels/qwen3/14b/greedy_sample.pymodels/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hppmodels/qwen3/14b/kernels/paged_attention_cce/kernel/rope_qkv_generated.hppmodels/qwen3/14b/paged_attention_cce.pymodels/qwen3/14b/prefill_fwd.pymodels/qwen3/14b/prefill_fwd_a8w8.pymodels/qwen3/14b/rms_lm_head.pymodels/qwen3/14b/rope_qkv_regen.pymodels/qwen3/14b/topk_select.pytests/contract/test_qwen3_14b_contract.py
3996cb8 to
110c906
Compare
decode_fwd was static batch-16: the public batch was pinned to the
padded pipeline width, so one compiled program could serve only a full
batch. The batch axes are now pl.dynamic and the live batch is read from
the seq_lens descriptor, matching what prefill_fwd already did.
- Separate the two conflated concepts. The padded pipeline width (the M
of every matmul) is BATCH_PAD; the public batch is BATCH_DYN /
`batch`. It is a pad width, not an architectural ceiling, so it is not
named MAX_BATCH. The contract now spells prefill's ("USER_BATCH",) and
decode's ("B",) as one ("BATCH",). paged_attention_cce's BATCH also
served two roles and is split into METADATA_BATCH_SLOTS (physical
metadata layout, must stay 16) and BATCH_PAD (standalone test-kernel
shapes).
- Regenerate rope_qkv_generated.hpp with a runtime batch. The embedded
body had batch=16 baked in as C++ literals and cannot be hand-edited:
ptoas CSEs constants by value, so 16 is both the batch divisor and the
Q-RMSNorm tile row count, and 128 is both HEAD_DIM and the item bound.
Add rope_qkv_regen.py as the regeneration source the header's
instruction referred to -- the scope it was generated from was deleted
in hw-native-sys#796, leaving that instruction unusable. Its --emit-header re-wraps
the artifact and reproduces the committed body byte for byte.
- fai_body reads the batch from the seq_lens descriptor, as the tiler
does, so the RoPE and attention phases cannot disagree about how many
sequences are live. A static_assert on the function-pointer type
guards the arity of the hand-written arg mapping.
- Relax the contract to 1 <= max_batch_size <= batch_pad and add
-b/--batch to the decode harness.
- Add a sync-safety invariant test over the generated body. Both guards
are unreachable at the padded batch (max item id 127 < 128), so the
skip path only goes live once the batch shrinks. The test simulates
every reachable path -- neither block, first only, both; a later item
cannot pass a guard the earlier one failed -- and requires that no
wait_flag runs without an outstanding set_flag and that the epilogue
drains every credit.
Validated on a2a3: batch-16 output is bit-identical to the pre-change
baseline; batch 1/2/3/5/8/15/16 all pass --validate-fwd; 4-layer runs
pass at batch 1 and 16. Single-layer makespan 981.3 -> 969.2 us over the
same 488 aicore tasks.
greedy_sample_fwd stays static batch and is documented as not composable
with a dynamic-batch decode; decode samples inline and does not use it.
1 participant
Summary
decode_fwdserved only a full batch: the public batch was pinned to the padded pipeline width. Its batch axes are nowpl.dynamicand the live batch is read from theseq_lensdescriptor, asprefill_fwdalready did, so one compiled program serves batch 1..16.BATCH_PAD; the public batch isBATCH_DYN/batch. It is a pad width rather than an architectural ceiling, so it is not namedMAX_BATCH. The contract now spells prefill's("USER_BATCH",)and decode's("B",)as one("BATCH",).paged_attention_cce'sBATCHalso served two roles and is split intoMETADATA_BATCH_SLOTS(physical metadata layout, must stay 16) andBATCH_PAD(standalone test-kernel shapes).rope_qkv_generated.hppwith a runtime batch. The embedded body hadbatch=16baked in as C++ literals and cannot be hand-edited: ptoas CSEs constants by value, so 16 is both the batch divisor and the Q-RMSNorm tile row count, and 128 is bothHEAD_DIMand the item bound.rope_qkv_regen.py, the regeneration source the header's instruction referred to — the scope it was generated from was deleted in perf(qwen3-14b decode): fold RoPE into the CANN FAI attention extern #796, leaving that instruction unusable.--emit-headerre-wraps the artifact and reproduces the committed body byte for byte.fai_bodyreads the batch from theseq_lensdescriptor, as the tiler does, so the RoPE and attention phases cannot disagree about how many sequences are live. Astatic_asserton the function-pointer type guards the arity of the hand-written arg mapping.1 <= max_batch_size <= batch_padand add-b/--batchto the decode harness.wait_flagruns without an outstandingset_flagand that the epilogue drains every credit.Validated on a2a3: batch-16 output is bit-identical to the pre-change baseline; batch 1/2/3/5/8/15/16 all pass
--validate-fwd; 4-layer runs pass at batch 1 and 16. Single-layer makespan 981.3 -> 969.2 us over the same 488 aicore tasks.greedy_sample_fwdstays static batch and is documented as not composable with a dynamic-batch decode; decode samples inline and does not use it.