Add: host_build_graph copy of the qwen3_14b_decode example - #1666
Open
noabauma wants to merge 1 commit into
Open
Add: host_build_graph copy of the qwen3_14b_decode example#1666noabauma wants to merge 1 commit into
noabauma wants to merge 1 commit into
Conversation
The 40-layer fused Qwen3-14B decode existed only for tensormap_and_ringbuffer. pto_orchestration_api.h is identical between the two runtimes and compile_orchestration() selects the include dirs of whichever runtime @scene_test names, so no C++ changes: this adds a test class and a README, and points every kernel source at the existing example's kernels/ via QWEN_KERNELS. Sharing the sources rather than copying them keeps the two runtimes' compute bit-for-bit comparable and keeps the vendored FusedInferAttentionScore tree in the repo exactly once. The heap must hold all 40 layers' intermediates simultaneously: host_build_graph builds the whole graph before the device schedules anything, so no task has completed while the graph is being built and the heap tail never advances off 0. The 256 MiB default is exhausted in the last layers (Task Allocator Deadlock - Heap Exhausted, ~254 of 256 MiB used); 512 MiB carries it, set per case via runtime_env. The task window is not the constraint: the graph is ~10.6K tasks, inside the 16384 default. tensormap_and_ringbuffer needs no sizing at all, because each layer's scope frees its intermediates as orchestration proceeds. Validated on a2a3 with golden checking across all 40 layers' KV writes and the hidden output, on three devices (36.8 / 36.8 / 36.9 ms device_wall).
📝 WalkthroughWalkthroughThis pull request adds a new example under ChangesQwen3-14B Decode Example
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Test as TestQwen314BDecodeHostBuildGraph
participant Kernels as AIC/AIV Kernel Registry
participant KVCache as INOUT KV Cache
participant Golden as Torch Golden Reference
Test->>Kernels: configure fused decode kernels
Test->>Test: generate_args(seed, seq_len)
Test->>Golden: compute_golden(args, params)
Test->>Kernels: run fused dispatch across 40 layers
Kernels->>KVCache: read/write KV cache (INOUT)
Kernels-->>Test: return decode output
Test->>Golden: compare output against golden reference
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 2
🧹 Nitpick comments (1)
examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py (1)
95-95: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: check the repo convention for RUF012 on scene tests.
Ruff flags
CALLABLEandCASESas mutable class attribute defaults. The base class insimpler_setup/scene_test.pydeclares them withoutClassVar, so existing scene tests probably share this pattern. If other scene tests silence or annotate the rule, follow the same approach here for consistency.#!/bin/bash # How do existing scene tests handle RUF012 on CALLABLE / CASES? rg -n 'RUF012' --glob '*.toml' --glob '*.py' --glob '*.cfg' rg -n -C 2 '^\s*(CALLABLE|CASES)\s*[:=]' --glob '*/test_*.py' | head -60Also applies to: 432-432
🤖 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 `@examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py` at line 95, Check existing scene tests and repository Ruff configuration for the established handling of RUF012 on CALLABLE and CASES. Update the CALLABLE and CASES declarations in the affected test class to match that convention, using ClassVar annotations or the project-approved suppression as appropriate, while preserving their existing values and behavior.Source: Linters/SAST tools
🤖 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 `@examples/a2a3/host_build_graph/qwen3_14b_decode/README.md`:
- Around line 145-148: Clarify the regeneration instructions in the README so
each source has an explicit destination: copy OUT/orchestration/ and
OUT/kernels/ to their intended locations, and copy only
models/qwen3/14b/kernels/paged_attention_cce/ into kernels/vendor/. Preserve the
existing codegen_only and CALLABLE re-transcription guidance.
In `@examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py`:
- Around line 49-54: Correct the vendored FusedInferAttentionScore path in the
comment near the CANN devkit header description to use
kernels/vendor/paged_attention_cce/ instead of
kernels/paged_attention_cce/vendor/. Keep the surrounding explanation unchanged.
---
Nitpick comments:
In `@examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py`:
- Line 95: Check existing scene tests and repository Ruff configuration for the
established handling of RUF012 on CALLABLE and CASES. Update the CALLABLE and
CASES declarations in the affected test class to match that convention, using
ClassVar annotations or the project-approved suppression as appropriate, while
preserving their existing values and behavior.
🪄 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: d8b299f6-9a74-4ded-90d9-353a99daf4d8
📒 Files selected for processing (2)
examples/a2a3/host_build_graph/qwen3_14b_decode/README.mdexamples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py
Comment on lines
+145
to
+148
| `codegen_only` needs no device. Then copy `OUT/orchestration/`, `OUT/kernels/` | ||
| and `models/qwen3/14b/kernels/paged_attention_cce/` into `kernels/vendor/` here, and | ||
| re-transcribe `CALLABLE` from `OUT/kernel_config.py` (which already records | ||
| `func_id`, `core_type`, per-kernel `signature`, and `extra_include_dirs`). |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Clarify the copy destinations in the regeneration steps.
The sentence lists three sources and one destination, so it reads as if OUT/orchestration/ and OUT/kernels/ also go into kernels/vendor/. Only the paged_attention_cce/ tree belongs there. State one destination per source.
📝 Proposed wording fix
-`codegen_only` needs no device. Then copy `OUT/orchestration/`, `OUT/kernels/`
-and `models/qwen3/14b/kernels/paged_attention_cce/` into `kernels/vendor/` here, and
-re-transcribe `CALLABLE` from `OUT/kernel_config.py` (which already records
-`func_id`, `core_type`, per-kernel `signature`, and `extra_include_dirs`).
+`codegen_only` needs no device. Then copy `OUT/orchestration/` and `OUT/kernels/`
+into `kernels/` here, copy `models/qwen3/14b/kernels/paged_attention_cce/` into
+`kernels/vendor/`, and re-transcribe `CALLABLE` from `OUT/kernel_config.py`
+(which already records `func_id`, `core_type`, per-kernel `signature`, and
+`extra_include_dirs`).📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| `codegen_only` needs no device. Then copy `OUT/orchestration/`, `OUT/kernels/` | |
| and `models/qwen3/14b/kernels/paged_attention_cce/` into `kernels/vendor/` here, and | |
| re-transcribe `CALLABLE` from `OUT/kernel_config.py` (which already records | |
| `func_id`, `core_type`, per-kernel `signature`, and `extra_include_dirs`). | |
| `codegen_only` needs no device. Then copy `OUT/orchestration/` and `OUT/kernels/` | |
| into `kernels/` here, copy `models/qwen3/14b/kernels/paged_attention_cce/` into | |
| `kernels/vendor/`, and re-transcribe `CALLABLE` from `OUT/kernel_config.py` | |
| (which already records `func_id`, `core_type`, per-kernel `signature`, and | |
| `extra_include_dirs`). |
🤖 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 `@examples/a2a3/host_build_graph/qwen3_14b_decode/README.md` around lines 145 -
148, Clarify the regeneration instructions in the README so each source has an
explicit destination: copy OUT/orchestration/ and OUT/kernels/ to their intended
locations, and copy only models/qwen3/14b/kernels/paged_attention_cce/ into
kernels/vendor/. Preserve the existing codegen_only and CALLABLE
re-transcription guidance.
Comment on lines
+49
to
+54
| # CANN devkit headers for the attention extern, which builds on AscendC and the | ||
| # vendored FusedInferAttentionScore under kernels/paged_attention_cce/vendor/. | ||
| # `vendor/.../attn_infra/base_defs.hpp` selects its AscendC entry header under | ||
| # `#if ASC_DEVKIT_MAJOR >= 9`, which ccec predefines from the installed devkit, | ||
| # so a CANN 9 box must be able to resolve `basic_api/kernel_basic_intf.h` from | ||
| # one of these. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the vendored tree path in this comment.
The registry loads the extern from {QWEN_KERNELS}/vendor/paged_attention_cce/... (Lines 134, 212, 238). This comment states kernels/paged_attention_cce/vendor/, which inverts the two directory levels. The README also documents the tree as kernels/vendor/paged_attention_cce/.
📝 Proposed comment fix
# CANN devkit headers for the attention extern, which builds on AscendC and the
-# vendored FusedInferAttentionScore under kernels/paged_attention_cce/vendor/.
-# `vendor/.../attn_infra/base_defs.hpp` selects its AscendC entry header under
+# vendored FusedInferAttentionScore under kernels/vendor/paged_attention_cce/.
+# `.../attn_infra/base_defs.hpp` selects its AscendC entry header under📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| # CANN devkit headers for the attention extern, which builds on AscendC and the | |
| # vendored FusedInferAttentionScore under kernels/paged_attention_cce/vendor/. | |
| # `vendor/.../attn_infra/base_defs.hpp` selects its AscendC entry header under | |
| # `#if ASC_DEVKIT_MAJOR >= 9`, which ccec predefines from the installed devkit, | |
| # so a CANN 9 box must be able to resolve `basic_api/kernel_basic_intf.h` from | |
| # one of these. | |
| # CANN devkit headers for the attention extern, which builds on AscendC and the | |
| # vendored FusedInferAttentionScore under kernels/vendor/paged_attention_cce/. | |
| # `.../attn_infra/base_defs.hpp` selects its AscendC entry header under | |
| # `#if ASC_DEVKIT_MAJOR >= 9`, which ccec predefines from the installed devkit, | |
| # so a CANN 9 box must be able to resolve `basic_api/kernel_basic_intf.h` from | |
| # one of these. |
🤖 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 `@examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py`
around lines 49 - 54, Correct the vendored FusedInferAttentionScore path in the
comment near the CANN devkit header description to use
kernels/vendor/paged_attention_cce/ instead of
kernels/paged_attention_cce/vendor/. Keep the surrounding explanation unchanged.
This was referenced Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The 40-layer fused Qwen3-14B decode existed only for tensormap_and_ringbuffer. pto_orchestration_api.h is identical between the two runtimes and compile_orchestration() selects the include dirs of whichever runtime @scene_test names, so no C++ changes: this adds a test class and a README, and points every kernel source at the existing example's kernels/ via QWEN_KERNELS. Sharing the sources rather than copying them keeps the two runtimes' compute bit-for-bit comparable and keeps the vendored FusedInferAttentionScore tree in the repo exactly once.
The heap must hold all 40 layers' intermediates simultaneously: host_build_graph builds the whole graph before the device schedules anything, so no task has completed while the graph is being built and the heap tail never advances off 0. The 256 MiB default is exhausted in the last layers (Task Allocator Deadlock - Heap Exhausted, ~254 of 256 MiB used); 512 MiB carries it, set per case via runtime_env. The task window is not the constraint: the graph is ~10.6K tasks, inside the 16384 default. tensormap_and_ringbuffer needs no sizing at all, because each layer's scope frees its intermediates as orchestration proceeds.
Validated on a2a3 with golden checking across all 40 layers' KV writes and the hidden output, on three devices (36.8 / 36.8 / 36.9 ms device_wall).