Initial C++ aclnn EAGER inference for Qwen3-235B-A22B MoE on Ascend 910 × 16 NPU

Pure-C++ inference runtime built directly on aclnn single-op API (no graph
compilation, no PyTorch, no ggml). Targets Qwen3-235B-A22B-Instruct-2507 BF16
with TP=16 HCCL tensor parallelism.

Quality-preserving throughput:
- Untuned baseline: 12 t/s
- Recommended (HCCL env + Fused RoPE + small ops): ~27 t/s (all prompts)
- PLD with degeneration guard: 29-45 t/s (structured long-form text)

Key components:
- 12 headers + 6 sources implementing attention/MoE/Runner/HCCL/RoPE
- Fused RoPE via aclnnApplyRotaryPosEmbV2 (layout=1, "half")
- PLD (Prompt Lookup Decoding) with degeneration guard:
low-distinct + tail-echo heuristics block loop-amplifying drafts
- bench_pld_safe.sh classifies each run as OK / LOOP_N / LOW_DIVERSITY
and separates TG stats accordingly (honest performance reporting)
- 19 unit / integration tests + end-to-end smoke test

HCCL environment (applied by tp_launch.sh):
HCCL_OP_EXPANSION_MODE=AIV + HCCL_OP_BASE_FFTS_MODE_ENABLE=1 +
TASK_QUEUE_ENABLE=2 contributes +89% TG vs default ring-only.

Known limitations:
- Does not exceed cann-recipes-infer GE graph baseline of ~54 t/s
- PLD on factual/code prompts is unreliable (disable or use bench_pld_safe.sh)
- Requires Ascend 910 initial-gen × 16 NPU and CANN 8.5.1

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,48 @@
+# Build artifacts
+/build/
+*.o
+*.obj
+*.a
+*.so
+*.exe
+
+# CMake
+CMakeCache.txt
+CMakeFiles/
+cmake_install.cmake
+Makefile
+compile_commands.json
+
+# Tokenizer output (regenerated from HF model)
+/tokenizer_data/
+*.bin
+!tests/**/*.bin
+
+# Reference data (regenerated by scripts/gen_*.py; too large to commit)
+/tests/attn_data/
+/tests/moe_data/
+/tests/mm_data/
+/tests/rms_norm_data/
+/tests/poc_data/
+
+# Runtime state
+/tmp/
+*.log
+/tp_rank_*.log
+hccl_root_info.bin
+/tmp/hccl_root_info.bin
+
+# Editor / IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+.DS_Store
+
+# Python
+__pycache__/
+*.pyc
+.ipynb_checkpoints/
+
+# Benchmark output
+bench_result*.log

CMakeLists.txt

@@ -0,0 +1,110 @@
+cmake_minimum_required(VERSION 3.16)
+project(qwen3-moe-aclnn CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+
+if(NOT CMAKE_BUILD_TYPE)
+    set(CMAKE_BUILD_TYPE Release)
+endif()
+
+set(CMAKE_CXX_FLAGS_RELEASE "-O2 -g")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-unused-function")
+
+# CANN paths
+if(NOT DEFINED CANN_INSTALL_DIR)
+    if(DEFINED ENV{ASCEND_TOOLKIT_HOME})
+        set(CANN_INSTALL_DIR $ENV{ASCEND_TOOLKIT_HOME})
+    else()
+        set(CANN_INSTALL_DIR /usr/local/Ascend/ascend-toolkit/latest)
+    endif()
+endif()
+message(STATUS "CANN_INSTALL_DIR: ${CANN_INSTALL_DIR}")
+
+include_directories(
+    ${CANN_INSTALL_DIR}/include
+    ${CANN_INSTALL_DIR}/include/aclnn
+    ${CMAKE_SOURCE_DIR}/include
+    ${CMAKE_SOURCE_DIR}/external
+)
+
+link_directories(${CANN_INSTALL_DIR}/lib64)
+
+set(CANN_LIBS ascendcl nnopbase opapi opapi_transformer acl_op_compiler hccl)
+
+# HCCL headers live under include/ but we need explicit include dir for <hccl/hccl.h>.
+include_directories(${CANN_INSTALL_DIR}/include/hccl)
+
+# ---- Library: qwen3-moe-aclnn core ----
+set(LCA_SOURCES
+    src/safetensors_loader.cpp
+    src/model_config.cpp
+    src/tokenizer.cpp
+    src/device_weights.cpp
+    src/runner.cpp
+)
+add_library(qwen3-moe-aclnn-core STATIC ${LCA_SOURCES})
+target_link_libraries(qwen3-moe-aclnn-core PUBLIC ${CANN_LIBS})
+
+# ---- Binaries ----
+add_executable(hello_acl tests/hello_acl.cpp)
+target_link_libraries(hello_acl qwen3-moe-aclnn-core)
+
+add_executable(test_safetensors tests/test_safetensors.cpp)
+target_link_libraries(test_safetensors qwen3-moe-aclnn-core)
+
+add_executable(test_model_config tests/test_model_config.cpp)
+target_link_libraries(test_model_config qwen3-moe-aclnn-core)
+
+add_executable(test_tokenizer tests/test_tokenizer.cpp)
+target_link_libraries(test_tokenizer qwen3-moe-aclnn-core)
+
+add_executable(test_rms_norm tests/test_rms_norm.cpp)
+target_link_libraries(test_rms_norm qwen3-moe-aclnn-core)
+
+add_executable(test_weight_load tests/test_weight_load.cpp)
+target_link_libraries(test_weight_load qwen3-moe-aclnn-core)
+
+add_executable(test_linear_hf tests/test_linear_hf.cpp)
+target_link_libraries(test_linear_hf qwen3-moe-aclnn-core)
+
+add_executable(test_rope tests/test_rope.cpp)
+target_link_libraries(test_rope qwen3-moe-aclnn-core)
+
+add_executable(test_rope_manual tests/test_rope_manual.cpp)
+target_link_libraries(test_rope_manual qwen3-moe-aclnn-core)
+
+add_executable(test_attention_layer tests/test_attention_layer.cpp)
+target_link_libraries(test_attention_layer qwen3-moe-aclnn-core)
+
+add_executable(test_moe_layer tests/test_moe_layer.cpp)
+target_link_libraries(test_moe_layer qwen3-moe-aclnn-core)
+
+add_executable(test_attention_decode tests/test_attention_decode.cpp)
+target_link_libraries(test_attention_decode qwen3-moe-aclnn-core)
+
+add_executable(test_engine_smoke tests/test_engine_smoke.cpp)
+target_link_libraries(test_engine_smoke qwen3-moe-aclnn-core)
+
+add_executable(test_layer_forward tests/test_layer_forward.cpp)
+target_link_libraries(test_layer_forward qwen3-moe-aclnn-core)
+
+add_executable(test_runner tests/test_runner.cpp)
+target_link_libraries(test_runner qwen3-moe-aclnn-core)
+
+# ---- Main CLI ----
+add_executable(qwen3-moe-aclnn src/main_cli.cpp)
+target_link_libraries(qwen3-moe-aclnn qwen3-moe-aclnn-core)
+
+add_executable(test_op_support tests/test_op_support.cpp)
+target_link_libraries(test_op_support qwen3-moe-aclnn-core)
+
+add_executable(test_rope_fused tests/test_rope_fused.cpp)
+target_link_libraries(test_rope_fused qwen3-moe-aclnn-core)
+
+add_executable(test_batch_decode tests/test_batch_decode.cpp)
+target_link_libraries(test_batch_decode qwen3-moe-aclnn-core)
+
+add_executable(test_batch_correctness tests/test_batch_correctness.cpp)
+target_link_libraries(test_batch_correctness qwen3-moe-aclnn-core)

LICENSE

@@ -0,0 +1,176 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for describing the origin of the Work and
+      reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Support. While redistributing the Work or
+      Derivative Works thereof, You may choose to offer, and charge a
+      fee for, acceptance of support, warranty, indemnity, or other
+      liability obligations and/or rights consistent with this License.
+      However, in accepting such obligations, You may act only on Your
+      own behalf and on Your sole responsibility, not on behalf of any
+      other Contributor, and only if You agree to indemnify, defend,
+      and hold each Contributor harmless for any liability incurred by,
+      or claims asserted against, such Contributor by reason of your
+      accepting any such warranty or support.
+
+   END OF TERMS AND CONDITIONS

README.md

@@ -0,0 +1,338 @@
+# qwen3-moe-aclnn
+
+Pure C++ inference of **Qwen3-235B-A22B-Instruct** BF16 on **Ascend 910 × 16 NPU**, built directly on the aclnn EAGER API (no graph compilation, no PyTorch, no ggml).
+
+---
+
+## Performance
+
+Measured on Ascend 910 initial-gen × 16 NPU (TP=16) with Qwen3-235B-A22B-Instruct-2507 BF16 weights.
+All numbers are **quality-preserving TG** (output was manually verified); greedy `temperature=0`.
+
+| Configuration | TG | Applicable prompts |
+|---|---|---|
+| Untuned baseline | 12 t/s | All |
+| **Default recommended** (no PLD) | **~27 t/s** | **All prompts, stable output** |
+| PLD with degeneration guard | 29-45 t/s | Structured text (essays, long-form answers) |
+| PLD on creative prompts | 25-40 t/s | Stories / varied generation |
+| PLD on factual / code prompts | unstable (21-95 t/s, high variance) | Not recommended |
+
+Reference: `cann-recipes-infer` GE graph baseline reports ~54 t/s on the same hardware. **This project does not exceed that baseline** — it trades some peak speed for (a) no graph compilation, (b) no PyTorch dependency, (c) full control over operator scheduling.
+
+### Key optimizations that contributed (in order of magnitude)
+
+| Rank | Optimization | Gain | Where |
+|---|---|---|---|
+| 🥇 | HCCL env tuning (`AIV` + `FFTS` + `TASK_QUEUE=2`) | +89% (12→23 t/s) | `scripts/tp_launch.sh` |
+| 🥈 | Fused RoPE via `aclnnApplyRotaryPosEmbV2` | +17% (23→27 t/s) | `include/rope.h` |
+| 🥉 | Prompt Lookup Decoding (PLD) w/ degeneration guard | +10-60% on applicable prompts | `src/main_cli.cpp` |
+| ○ | Device-side topk-w normalize, MoE argsort, cos/sin cache | ~+15% cumulative | `include/engine.h` |
+| ○ | WorkspacePool (thread-local + retain-old) | reduces alloc overhead | `include/workspace_pool.h` |
+
+---
+
+## Architecture
+
+**Model**: Qwen3-235B-A22B, 94 layers, 128 experts (top-k=8), GQA (64 Q heads, 4 KV heads), BF16.
+
+**Parallelism**: TP=16 via HCCL ring AllReduce. KV heads sharded 1-per-rank (since 4 KV heads < 16 ranks, Q heads 0-3 on each rank share KV head 0).
+
+**Execution**: aclnn EAGER mode — every op goes through `aclnn*` single-op API with workspace pool; no graph capture, no GE IR. Async stream execution with `TASK_QUEUE_ENABLE=2` for kernel submission overlap.
+
+**Tokenizer**: Uses HuggingFace `transformers` via a Python subprocess for encoding; vocab decode is pure C++ from an exported `vocab.bin`.
+
+### Per-layer forward flow
+
+```
+x_in [S, D=4096]
+  ↓
+┌── Attention branch (TP: Q_DIM=512=4h×128, KV_DIM=128=1h×128) ──┐
+│  RmsNorm(input_layernorm)
+│  linear_hf q_proj / k_proj / v_proj        → q, k, v
+│  Per-head RmsNorm q_norm, k_norm
+│  Fused RoPE: aclnnApplyRotaryPosEmbV2 (layout=1, "half")
+│  Append K, V to per-layer KV cache
+│  Mask selection:
+│    prefill:      2048×2048 causal + sparse_mode=3
+│    decode S=1:   nullptr + sparse_mode=0
+│    batch decode: [1,1,S,past+S] custom bool mask + sparse_mode=0
+│  FIAS (aclnnFusedInferAttentionScore)
+│  o_proj linear_hf → partial per-rank
+│  HCCL AllReduce (ring + AIV + FFTS) → full
+└─────────┘
+  ↓ residual add
+┌── MoE branch ──┐
+│  RmsNorm(post_attention_layernorm)
+│  router linear_hf → logits [S, 128]
+│  moe_gating_topk_softmax → topk_w[S,8], topk_idx[S,8]
+│  Device-side normalize (reduce_sum + adds + cast + div)
+│  moe_init_routing_v3 → expanded_x, expanded_ri, tokens_per_expert
+│  grouped_matmul_v4 gate/up/down (SwiGLU activation)
+│  Device-side argsort × 2 → fwd permutation (avoids host sync)
+│  IndexSelect → packed
+│  Broadcast-mul by topk_w + ReduceSum axis=1
+│  HCCL AllReduce → full
+└─────────┘
+  ↓ residual add
+x_out
+```
+
+---
+
+## Model weights
+
+This project targets **Qwen3-235B-A22B-Instruct-2507** (BF16). About **470 GB** of safetensors shards.
+
+**Download sources**:
+- HuggingFace: https://huggingface.co/Qwen/Qwen3-235B-A22B-Instruct-2507
+- ModelScope: https://www.modelscope.cn/models/Qwen/Qwen3-235B-A22B-Instruct-2507
+
+Download via `huggingface-cli` or `modelscope` CLI:
+```bash
+# HuggingFace
+huggingface-cli download Qwen/Qwen3-235B-A22B-Instruct-2507 --local-dir /path/to/Qwen3-235B-A22B-Instruct-2507-BF16
+
+# ModelScope
+modelscope download --model Qwen/Qwen3-235B-A22B-Instruct-2507 --local_dir /path/to/Qwen3-235B-A22B-Instruct-2507-BF16
+```
+
+**Weights format**: the binary reads HuggingFace `.safetensors` shards (multi-shard mmap), `config.json`, and `tokenizer.json` directly from the model directory. No conversion step is needed — point `--model-dir` at the downloaded directory.
+
+**Expected directory contents**:
+```
+Qwen3-235B-A22B-Instruct-2507-BF16/
+├── config.json
+├── tokenizer.json
+├── tokenizer_config.json
+├── model-00001-of-000XX.safetensors
+├── ...
+└── model.safetensors.index.json
+```
+
+---
+
+## Build
+
+```bash
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+cmake -B build
+cmake --build build -j8 --target qwen3-moe-aclnn
+```
+
+**Requires**:
+- CANN 8.5.1 or compatible
+- Python 3 + `transformers` + `torch_npu` (for tokenizer subprocess and reference-data generation only)
+- C++17 compiler
+- Ascend 910 × 16 NPU
+- nlohmann/json (bundled as `external/json.hpp`)
+
+**Python environment setup** — the tokenizer calls a Python subprocess. Override the activation command via `QWEN3_PYENV_INIT` if your conda / venv layout differs from the default:
+```bash
+export QWEN3_PYENV_INIT="source /opt/my_conda/etc/profile.d/conda.sh && conda activate my_env && "
+```
+If unset, the default tries `${HOME}/miniconda3` with env `qwen3` and auto-sources the Ascend toolkit.
+
+---
+
+## Quick-start inference
+
+```bash
+# 1. Export tokenizer vocab to binary (one-time setup)
+python3 scripts/export_vocab.py /path/to/Qwen3-235B-A22B-Instruct-2507-BF16
+
+# 2. Run inference (TP=16)
+./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn \
+    --model-dir /path/to/Qwen3-235B-A22B-Instruct-2507-BF16 \
+    --prompt "The capital of France is" \
+    --n-predict 100 \
+    --temperature 0 \
+    --vocab tokenizer_data/vocab.bin
+```
+
+Expected: ~27 t/s, coherent output.
+
+### Recommended flags by use case
+
+**Universal default (stable, any prompt)** — no PLD:
+```bash
+./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... --temperature 0 --no-stream
+```
+
+**Structured / long-form (essays, explanations)** — PLD with guard gives +60-90%:
+```bash
+./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... --pld --temperature 0 --no-stream
+```
+
+**Interactive REPL (multi-turn chat)**:
+```bash
+./scripts/tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... \
+    --interactive --chat --temperature 0.7 --top-p 0.8
+```
+
+---
+
+## PLD degeneration guard
+
+Prompt Lookup Decoding speeds up generation by having the model verify a batch of "draft" tokens in a single forward pass. The drafts are copied from the generation history via n-gram match.
+
+**Known failure mode**: on prompts the model tends to repeat on (factual Q&A, code generation), the n-gram match feeds the model's own repetition back as drafts, creating a positive feedback loop that accelerates degenerate output. Early versions of this project reported misleading peak TG numbers driven by this loop.
+
+**This project's guard** blocks suspect drafts with two heuristics:
+
+1. **low-distinct**: draft's distinct-token count < threshold → reject
+2. **tail-echo**: all of last N hist tokens equal draft[0] → reject
+
+Rejected drafts fall back to single-token decode. A `[warn]` line is emitted once if the generated tail shows 8 consecutive identical tokens.
+
+Flags:
+```
+--pld                  enable PLD (opt-in)
+--pld-k N              draft window size (default: 10)
+--pld-ngram N          n-gram match size (default: 1, with multi-level fallback)
+--pld-min-hist N       skip PLD until history >= N tokens (default: 20)
+--pld-no-guard         disable the degeneration guard (dangerous: can produce dead loops)
+--pld-guard-distinct N minimum distinct tokens in draft (default: 3)
+--pld-guard-tail N     tail-echo window (default: 6)
+--pld-loop-warn N      emit warning on N consecutive identical tokens (default: 8)
+```
+
+**Honest benchmarking**: use `scripts/bench_pld_safe.sh`, which classifies each run's output as OK / LOOP_N / LOW_DIVERSITY and separates TG statistics for OK-only vs degraded runs.
+
+---
+
+## Correctness verification
+
+15+ unit / integration tests checked against Python (HuggingFace Transformers) reference:
+
+```bash
+./build/test_attention_layer    # rel=4.9e-4 vs Python prefill
+./build/test_attention_decode   # rel=0  (bit-exact)
+./build/test_moe_layer          # rel=3.6e-3
+./build/test_layer_forward      # full single layer
+./build/test_runner             # multi-layer runner
+./build/test_rope_fused         # aclnnApplyRotaryPosEmbV2 vs manual HF rotate_half
+./build/test_batch_decode       # S=1..8 timing
+./build/test_batch_correctness  # argmax consistency
+./build/test_op_support         # 910-specific op availability
+# Integration smoke:
+./tests/test_chat_flow.sh       # 7/7 PASS
+```
+
+Tests expect reference data under `tests/<name>_data/` generated by `scripts/gen_*_reference.py`. See each script's docstring.
+
+---
+
+## Environment tuning (auto-applied by `tp_launch.sh`)
+
+```bash
+HCCL_WHITELIST_DISABLE=1
+HCCL_ALGO=level0:ring                  # ring, not fullmesh (fullmesh causes garbled output)
+HCCL_BUFFSIZE=200                       # sweet spot; 100 and 400 both slower
+HCCL_OP_EXPANSION_MODE=AIV              # key: AI Vector cores participate in reduce scheduling
+HCCL_OP_BASE_FFTS_MODE_ENABLE=1         # key: Fast Frequently-used Transfer Scheduling
+TASK_QUEUE_ENABLE=2                     # key: aggressive async task submission
+```
+
+Removing any of the three "key" env vars drops TG by 20-40%.
+
+---
+
+## Directory layout
+
+```
+include/
+├── acl_common.h           RAII wrappers, DeviceBuffer, make_contig_tensor
+├── aclnn_ops.h            single-op wrappers + WorkspacePool integration
+├── acl_runtime.h          AclRuntime (device + stream management)
+├── device_weights.h       safetensors → device loading + TP sharding
+├── engine.h               attention_forward + moe_forward + RopeCache
+├── hccl_comm.h            HCCL init + allreduce + broadcast
+├── model_config.h         Qwen3 hyperparameters + compute_derived
+├── rope.h                 apply_rope_fused (aclnnApplyRotaryPosEmbV2 wrapper)
+├── runner.h               Runner class (prefill/decode/decode_batch/rewind/profile)
+├── safetensors_loader.h   multi-shard safetensors mmap parser
+├── tokenizer.h            vocab decode + Python subprocess encode
+└── workspace_pool.h       thread-local aclnn workspace pool (retain-old)
+
+src/
+├── device_weights.cpp     load_attention (GQA fix), load_moe (permute sync fix)
+├── main_cli.cpp           CLI entry + PLD main loop + degeneration guard + multi-turn
+├── model_config.cpp       compute_derived (GQA KV sharding)
+├── runner.cpp             Runner (build_batch_decode_mask_ etc.)
+├── safetensors_loader.cpp
+└── tokenizer.cpp
+
+scripts/
+├── tp_launch.sh           production launcher (auto-applies HCCL env)
+├── bench_tg.sh            stable N-run TG measurement
+├── bench_pld_safe.sh      PLD benchmark with output-correctness classifier
+├── bench_hccl[_adv].sh    HCCL parameter sweep
+├── bench_pld[_k].sh       PLD K × ngram sweep (legacy, prefer bench_pld_safe.sh)
+├── export_vocab.py        vocab.bin exporter from HF tokenizer
+└── gen_*_reference.py     per-op Python reference data generators
+
+tests/
+├── test_attention_*       attention correctness (prefill / decode)
+├── test_moe_layer         MoE correctness
+├── test_layer_forward     full single layer
+├── test_runner            multi-layer Runner
+├── test_rope_fused        fused RoPE vs manual HF
+├── test_batch_*           batch decode timing + correctness
+├── test_op_support        910-specific op availability probe
+└── test_chat_flow.sh      end-to-end integration smoke
+```
+
+---
+
+## CLI reference
+
+```
+--model-dir <path>         (required) HF safetensors directory
+--prompt "<text>"          prompt text
+--prompt-file FILE         read prompt from file (avoids shell-escape issues)
+--n-predict N              maximum tokens to generate
+--tp-size N                tensor parallelism (or set TP_SIZE env)
+--max-seq N                KV cache + context cap (default: 512)
+--temperature F            0 = greedy; typical 0.7
+--top-k N                  0 = disabled
+--top-p F                  1.0 = disabled
+--seed N                   0 = time-based
+--chat                     apply Qwen3 chat template
+--system "<text>"          system role text (with --chat)
+--interactive, -i          REPL mode (multi-turn memory with --chat)
+--reset                    force stateless REPL (reset KV between turns)
+--no-stream                batch-print final text instead of per-token streaming
+--vocab <path>             vocab.bin path (default: tokenizer_data/vocab.bin)
+--pld*                     see "PLD degeneration guard" section
+```
+
+---
+
+## Known limitations
+
+- **Not yet reaching cann-recipes GE graph 54 t/s baseline** (currently ~27 t/s stable / up to ~45 t/s PLD).
+  Closing the gap requires one of: (a) real graph compilation, (b) fused collectives (`MatmulAllReduce`, `GroupedMatmulAllReduce`) which are absent on 910 initial-gen, (c) migration to 910B/A2/A3.
+- **Only `tp_size` ∈ {1, 2, 4, 8, 16}** supported. Values that don't evenly divide 64 Q heads will error.
+- **PLD on factual/code prompts is unreliable** — either produces baseline TG (guard rejects most drafts) or enters partial degeneration the classifier may not catch at low-severity. Use `bench_pld_safe.sh` to evaluate honestly.
+- **Tokenizer requires Python subprocess** — adds ~1s startup for first encode. Override via `QWEN3_PYENV_INIT` env if default conda path doesn't match.
+- **NPU performance has high run-to-run variance** (up to 4× in some configurations) due to BF16 + MoE intrinsic non-determinism and shared hardware resources. Report medians over ≥5 runs.
+
+---
+
+## Future directions (prioritized)
+
+1. **Draft Model Speculative Decoding** with Qwen3-0.6B — more stable accept rate than n-gram PLD, expected +60-100% TG across prompt types (1-2 week implementation).
+2. **HCCL AllReduce / compute overlap** — ~+10-15% in theory, limited by EAGER path serial dependencies.
+3. **KV cache INT8 quantization** — reduces memory-bandwidth pressure, ~+15-25% on long contexts (pending 910-initial-gen op support verification).
+4. **W8 weight quantization** — ~+10-20% if aclnn quantization kernels exist on 910 initial-gen.
+
+Not recommended:
+- `aclmdlRI` stream-capture-style graph recording (POC proved 1.13× ceiling, not worth the engineering cost).
+- Custom AscendC fused ops (high maintenance cost unless dedicated kernel engineer).
+- torchair / torch.compile migration (breaks pure-C++ design).
+
+---
+
+## License
+
+Apache License 2.0 — see `LICENSE`.

include/acl_common.h

@@ -0,0 +1,106 @@
+#pragma once
+#include <acl/acl.h>
+#include <aclnn/acl_meta.h>
+#include <cstdio>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#define ACL_CHECK(x) do { \
+    aclError __e = (x); \
+    if (__e != ACL_ERROR_NONE) { \
+        fprintf(stderr, "ACL error %d at %s:%d : %s\n", __e, __FILE__, __LINE__, #x); \
+        std::abort(); \
+    } \
+} while(0)
+
+#define ACLNN_CHECK(x) do { \
+    aclnnStatus __e = (x); \
+    if (__e != 0) { \
+        const char* __msg = aclGetRecentErrMsg(); \
+        fprintf(stderr, "aclnn error %d at %s:%d : %s\n  msg: %s\n", (int)__e, __FILE__, __LINE__, #x, __msg ? __msg : "(null)"); \
+        std::abort(); \
+    } \
+} while(0)
+
+// RAII wrapper for aclTensor: call aclDestroyTensor on dtor
+struct AclTensorDel { void operator()(aclTensor* t) const { if (t) aclDestroyTensor(t); } };
+using AclTensorPtr = std::unique_ptr<aclTensor, AclTensorDel>;
+
+struct AclTensorListDel { void operator()(aclTensorList* t) const { if (t) aclDestroyTensorList(t); } };
+using AclTensorListPtr = std::unique_ptr<aclTensorList, AclTensorListDel>;
+
+struct AclIntArrayDel { void operator()(aclIntArray* a) const { if (a) aclDestroyIntArray(a); } };
+using AclIntArrayPtr = std::unique_ptr<aclIntArray, AclIntArrayDel>;
+
+// Create ACL tensor with explicit row-major shape (outermost leftmost) and element strides.
+// NOTE: stride is in ELEMENTS, not bytes.
+inline AclTensorPtr make_acl_tensor(void* data, aclDataType dt,
+                                    const std::vector<int64_t>& shape,
+                                    const std::vector<int64_t>& stride_elems,
+                                    aclFormat fmt = ACL_FORMAT_ND) {
+    int64_t n = (int64_t)shape.size();
+    int64_t storage_len = 1;
+    for (int i = 0; i < n; i++) storage_len += (shape[i] - 1) * stride_elems[i];
+    aclTensor* t = aclCreateTensor(
+        shape.data(), (uint64_t)n, dt,
+        stride_elems.data(), 0, fmt,
+        &storage_len, 1, data);
+    return AclTensorPtr(t);
+}
+
+// Default contiguous strides for row-major tensor: stride[i] = product of shape[i+1..n-1]
+inline std::vector<int64_t> contiguous_strides(const std::vector<int64_t>& shape) {
+    int n = (int)shape.size();
+    std::vector<int64_t> s(n);
+    int64_t acc = 1;
+    for (int i = n - 1; i >= 0; --i) {
+        s[i] = acc;
+        acc *= shape[i];
+    }
+    return s;
+}
+
+inline AclTensorPtr make_contig_tensor(void* data, aclDataType dt,
+                                       const std::vector<int64_t>& shape,
+                                       aclFormat fmt = ACL_FORMAT_ND) {
+    return make_acl_tensor(data, dt, shape, contiguous_strides(shape), fmt);
+}
+
+inline size_t dtype_size(aclDataType dt) {
+    switch (dt) {
+        case ACL_FLOAT:     return 4;
+        case ACL_FLOAT16:   return 2;
+        case ACL_BF16:      return 2;
+        case ACL_INT8:      return 1;
+        case ACL_INT32:     return 4;
+        case ACL_INT64:     return 8;
+        default: return 0;
+    }
+}
+
+// Device buffer RAII: allocates via aclrtMalloc, frees in dtor
+struct DeviceBuffer {
+    void*  ptr  = nullptr;
+    size_t size = 0;
+
+    DeviceBuffer() = default;
+    explicit DeviceBuffer(size_t bytes) { alloc(bytes); }
+    ~DeviceBuffer() { if (ptr) aclrtFree(ptr); }
+    DeviceBuffer(const DeviceBuffer&) = delete;
+    DeviceBuffer& operator=(const DeviceBuffer&) = delete;
+    DeviceBuffer(DeviceBuffer&& o) noexcept : ptr(o.ptr), size(o.size) { o.ptr = nullptr; o.size = 0; }
+    DeviceBuffer& operator=(DeviceBuffer&& o) noexcept {
+        if (this != &o) { if (ptr) aclrtFree(ptr); ptr = o.ptr; size = o.size; o.ptr = nullptr; o.size = 0; }
+        return *this;
+    }
+
+    void alloc(size_t bytes) {
+        if (ptr) aclrtFree(ptr);
+        ACL_CHECK(aclrtMalloc(&ptr, bytes, ACL_MEM_MALLOC_HUGE_FIRST));
+        size = bytes;
+    }
+    void* get() { return ptr; }
+    const void* get() const { return ptr; }
+};

include/acl_runtime.h

@@ -0,0 +1,41 @@
+// acl_runtime.h — per-rank ACL runtime init/teardown.
+#pragma once
+#include "acl_common.h"
+#include <cstdio>
+
+class AclRuntime {
+public:
+    AclRuntime() = default;
+    ~AclRuntime() { shutdown(); }
+
+    bool init(int device_id) {
+        if (initialized_) return true;
+        device_id_ = device_id;
+        ACL_CHECK(aclInit(nullptr));
+        ACL_CHECK(aclrtSetDevice(device_id));
+        ACL_CHECK(aclrtCreateContext(&ctx_, device_id));
+        ACL_CHECK(aclrtCreateStream(&stream_));
+        initialized_ = true;
+        return true;
+    }
+
+    void shutdown() {
+        if (!initialized_) return;
+        if (stream_) { aclrtDestroyStream(stream_); stream_ = nullptr; }
+        if (ctx_)    { aclrtDestroyContext(ctx_); ctx_ = nullptr; }
+        aclrtResetDevice(device_id_);
+        aclFinalize();
+        initialized_ = false;
+    }
+
+    void sync() { if (stream_) ACL_CHECK(aclrtSynchronizeStream(stream_)); }
+
+    aclrtStream stream() const { return stream_; }
+    int device_id() const { return device_id_; }
+
+private:
+    bool initialized_ = false;
+    int device_id_ = 0;
+    aclrtContext ctx_ = nullptr;
+    aclrtStream  stream_ = nullptr;
+};

include/aclnn_ops.h

@@ -0,0 +1,345 @@
+// aclnn_ops.h — thin wrappers around common aclnn operators used in forward pass.
+// Each wrapper does GetWorkspaceSize + op call on the provided stream.
+//
+// All tensors are passed as raw aclTensor* (caller owns them).
+// Workspace allocation uses DeviceBuffer (RAII).
+#pragma once
+#include "acl_common.h"
+#include "workspace_pool.h"
+
+// Thread-local shared workspace pool for all aclnn wrappers below. Single-threaded stream
+// means we can safely reuse one buffer across serial op calls. Set via `GGML_CANN_WP=0` is
+// not supported here — if truly needed, we'd wire a flag.
+inline WorkspacePool& _lca_pool() {
+    thread_local WorkspacePool pool;
+    return pool;
+}
+
+#include <aclnnop/aclnn_add.h>
+#include <aclnnop/aclnn_addcmul.h>
+#include <aclnnop/aclnn_grouped_matmul_v4.h>
+#include <aclnnop/aclnn_moe_finalize_routing.h>
+#include <aclnnop/aclnn_moe_finalize_routing_v2.h>
+#include <aclnnop/aclnn_moe_gating_top_k_softmax.h>
+#include <aclnnop/aclnn_moe_init_routing_v3.h>
+#include <aclnnop/aclnn_cast.h>
+#include <aclnnop/aclnn_copy.h>
+#include <aclnnop/aclnn_div.h>
+#include <aclnnop/aclnn_fused_infer_attention_score.h>
+#include <aclnnop/aclnn_index_select.h>
+#include <aclnnop/aclnn_matmul.h>
+#include <aclnnop/aclnn_mul.h>
+#include <aclnnop/aclnn_neg.h>
+#include <aclnnop/aclnn_reduce_sum.h>
+#include <aclnnop/aclnn_silu.h>
+
+// ---- RmsNorm ----
+// Signature (based on ggml-cann usage): aclnnRmsNorm(x, gamma, eps, y, rstd)
+// where rstd (rsqrt of mean-square) is an extra output we usually discard.
+
+// Forward declare header; include happens in impl file to keep this header light.
+extern "C" {
+#include <aclnnop/aclnn_rms_norm.h>
+}
+
+inline void rms_norm(aclrtStream stream,
+                     aclTensor* x,         // [N, D] BF16/FP16
+                     aclTensor* gamma,     // [D] same dtype as x
+                     double     eps,
+                     aclTensor* y,         // [N, D]
+                     aclTensor* rstd       // [N] fp32 (required output)
+) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnRmsNormGetWorkspaceSize(x, gamma, eps, y, rstd, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnRmsNorm(wp, ws, exec, stream));
+}
+
+// ---- Silu ----
+inline void silu(aclrtStream stream, aclTensor* x, aclTensor* y) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnSiluGetWorkspaceSize(x, y, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnSilu(wp, ws, exec, stream));
+}
+
+// ---- Mul (element-wise) ----
+inline void mul(aclrtStream stream, aclTensor* a, aclTensor* b, aclTensor* out) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnMulGetWorkspaceSize(a, b, out, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnMul(wp, ws, exec, stream));
+}
+
+// ---- Cast ----
+inline void cast(aclrtStream stream, aclTensor* x, aclDataType dst_dtype, aclTensor* y) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnCastGetWorkspaceSize(x, dst_dtype, y, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnCast(wp, ws, exec, stream));
+}
+
+// ---- InplaceCopy: copy src (possibly non-contiguous via strides) into contiguous dst ----
+inline void inplace_copy(aclrtStream stream, aclTensor* dst, aclTensor* src) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnInplaceCopyGetWorkspaceSize(dst, src, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnInplaceCopy(wp, ws, exec, stream));
+}
+
+// ---- Matmul: out = a @ b ----
+// cube_math_type:
+//   0 = KEEP_DTYPE, 1 = ALLOW_FP32_DOWN_PRECISION, 2 = USE_FP16, 3 = USE_HF32
+inline void matmul(aclrtStream stream,
+                   aclTensor* a, aclTensor* b, aclTensor* out,
+                   int8_t cube_math_type = 1) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnMatmulGetWorkspaceSize(a, b, out, cube_math_type, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnMatmul(wp, ws, exec, stream));
+}
+
+// ---- Neg ----
+inline void neg(aclrtStream stream, aclTensor* x, aclTensor* y) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnNegGetWorkspaceSize(x, y, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnNeg(wp, ws, exec, stream));
+}
+
+// ---- Addcmul: self = self + value * (tensor1 * tensor2) ----
+inline void addcmul(aclrtStream stream, aclTensor* self_io, aclTensor* t1, aclTensor* t2, float value) {
+    aclScalar* v = aclCreateScalar(&value, ACL_FLOAT);
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnAddcmulGetWorkspaceSize(self_io, t1, t2, v, self_io, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnAddcmul(wp, ws, exec, stream));
+    aclDestroyScalar(v);
+}
+
+// ---- MoE Gating TopK Softmax ----
+// x [N, E] → y [N, K] (top-K softmax probs), expert_idx [N, K] int32, row_idx [N, K] int32
+inline void moe_gating_topk_softmax(aclrtStream stream,
+                                    aclTensor* x, int64_t k,
+                                    aclTensor* y_out, aclTensor* idx_out, aclTensor* row_idx_out) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnMoeGatingTopKSoftmaxGetWorkspaceSize(x, nullptr, k, y_out, idx_out, row_idx_out, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnMoeGatingTopKSoftmax(wp, ws, exec, stream));
+}
+
+// ---- MoE Init Routing V3 ----
+// x [N, D], expert_idx [N, K] int32 → expanded_x [N*K, D], expanded_row_idx [N*K] int32,
+// tokens_per_expert [E] int64
+inline void moe_init_routing_v3(aclrtStream stream,
+                                aclTensor* x, aclTensor* expert_idx,
+                                int64_t n_experts, int64_t active_num,
+                                aclTensor* expanded_x, aclTensor* expanded_row_idx,
+                                aclTensor* tokens_per_expert)
+{
+    int64_t range[2] = {0, n_experts};
+    aclIntArray* r = aclCreateIntArray(range, 2);
+    // scale_out_optional we dummy since quant_mode=-1 (no quant) still requires pass a placeholder?
+    // Per our POC test earlier: pass a real tensor for scale_out works.
+    // For simplicity here, we'll allocate a dummy [active_num] float tensor.
+    DeviceBuffer dummy(active_num * 4);
+    auto t_dummy = make_contig_tensor(dummy.get(), ACL_FLOAT, {active_num});
+
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    // rowIdxType=1: expanded_row_idx[i] = sorted_position p for i-th original (n,k) flat index.
+    // This lets us use expanded_row_idx directly as the gather index (forward permutation).
+    ACLNN_CHECK(aclnnMoeInitRoutingV3GetWorkspaceSize(
+        x, expert_idx, nullptr, nullptr,
+        active_num, 0, n_experts, 0, 1, true, -1,
+        r, 1,
+        expanded_x, expanded_row_idx, tokens_per_expert, t_dummy.get(),
+        &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnMoeInitRoutingV3(wp, ws, exec, stream));
+    aclDestroyIntArray(r);
+}
+
+// ---- GroupedMatmulV4 (single-in single-out, M-axis split) ----
+// x [T, K_in], w [E, K_in, N_out] contiguous row-major, group_list [E] int64 → y [T, N_out]
+// group_list_type: 0=cumsum, 1=counts (V4 doc)
+inline void grouped_matmul_v4(aclrtStream stream,
+                              aclTensor* x, aclTensor* w, aclTensor* group_list, aclTensor* y,
+                              int64_t group_list_type = 1)
+{
+    aclTensor* xa[] = {x}; aclTensorList* x_list = aclCreateTensorList(xa, 1);
+    aclTensor* wa[] = {w}; aclTensorList* w_list = aclCreateTensorList(wa, 1);
+    aclTensor* ya[] = {y}; aclTensorList* y_list = aclCreateTensorList(ya, 1);
+
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnGroupedMatmulV4GetWorkspaceSize(
+        x_list, w_list,
+        nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
+        group_list,
+        nullptr, nullptr, nullptr,
+        3, 0, group_list_type, 0,
+        y_list, nullptr, nullptr,
+        &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnGroupedMatmulV4(wp, ws, exec, stream));
+    // NOTE: TensorList takes ownership of the raw tensors. Destroying the list frees them,
+    // which would cause double-free in the caller's AclTensorPtr. Leak the list (small cost).
+    // A cleaner API would accept (ptr, shape, dtype) triples and build tensors internally.
+    // TODO(M6): refactor for long-running use.
+}
+
+// ---- MoE Finalize Routing V2: out = x1 + weighted_sum of top-K outputs ----
+// V2 has all inputs optional except expandedX/expandedRowIdx/out; pass nullptr for x1 to
+// skip the residual add, or pass the residual to fuse it into this op.
+inline void moe_finalize_routing(aclrtStream stream,
+                                 aclTensor* expanded_x,
+                                 aclTensor* x1_skip,               // [N, D] added to output (nullable)
+                                 aclTensor* scales,                // weights [N, K]
+                                 aclTensor* expanded_row_idx,
+                                 aclTensor* expert_idx,             // [N, K] topk expert indices (nullable)
+                                 aclTensor* out)
+{
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnMoeFinalizeRoutingV2GetWorkspaceSize(
+        expanded_x,
+        expanded_row_idx,
+        x1_skip,        // x1Optional
+        nullptr,        // x2Optional
+        nullptr,        // biasOptional
+        scales,         // scalesOptional
+        expert_idx,     // expertIdxOptional (needed for correct routing)
+        0,              // dropPadMode (0 = dropless, which matches our pipeline)
+        out,
+        &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnMoeFinalizeRoutingV2(wp, ws, exec, stream));
+}
+
+// ---- Div: self / other (broadcast supported) ----
+inline void div_tensor(aclrtStream stream, aclTensor* self, aclTensor* other, aclTensor* out) {
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnDivGetWorkspaceSize(self, other, out, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnDiv(wp, ws, exec, stream));
+}
+
+// ---- In-place scalar add: self += scalar ----
+#include <aclnnop/aclnn_add.h>
+#include <aclnnop/aclnn_argsort.h>
+
+// ---- Argsort: indices that would sort self along dim (returns INT64) ----
+inline void argsort(aclrtStream stream, aclTensor* self, int64_t dim, bool descending,
+                    aclTensor* indices_out) {
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnArgsortGetWorkspaceSize(self, dim, descending, indices_out, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnArgsort(wp, ws, exec, stream));
+}
+
+inline void inplace_adds(aclrtStream stream, aclTensor* self, double value) {
+    float v = (float)value;
+    aclScalar* s = aclCreateScalar(&v, ACL_FLOAT);
+    float alpha_v = 1.0f;
+    aclScalar* al = aclCreateScalar(&alpha_v, ACL_FLOAT);
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnInplaceAddsGetWorkspaceSize(self, s, al, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnInplaceAdds(wp, ws, exec, stream));
+    aclDestroyScalar(s);
+    aclDestroyScalar(al);
+}
+
+// ---- ReduceSum over specified dims ----
+inline void reduce_sum(aclrtStream stream, aclTensor* self, const std::vector<int64_t>& dims,
+                       bool keep_dims, aclDataType out_dtype, aclTensor* out) {
+    aclIntArray* d = aclCreateIntArray(dims.data(), dims.size());
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnReduceSumGetWorkspaceSize(self, d, keep_dims, out_dtype, out, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnReduceSum(wp, ws, exec, stream));
+    aclDestroyIntArray(d);
+}
+
+// ---- IndexSelect: out[j] = self[index[j], ...] ----
+inline void index_select(aclrtStream stream, aclTensor* self, int64_t dim, aclTensor* index, aclTensor* out) {
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnIndexSelectGetWorkspaceSize(self, dim, index, out, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnIndexSelect(wp, ws, exec, stream));
+}
+
+// ---- FusedInferAttentionScore (simplified wrapper for prefill/decode without quant, BSH layout).
+// Caller owns q/k/v/mask/out; k/v are single-tensor lists.
+inline void fused_infer_attention_score(
+    aclrtStream stream,
+    aclTensor* q,                       // [B, S, Hq*Dh] BF16
+    aclTensor* k,                       // [B, S, Hkv*Dh] BF16
+    aclTensor* v,                       // [B, S, Hkv*Dh] BF16
+    aclTensor* atten_mask,              // [1, 1, M, M] bool, sparse_mode=3 needs M=2048
+    std::vector<int64_t> actual_seq_lens,
+    std::vector<int64_t> actual_seq_lens_kv,
+    int64_t num_heads, int64_t num_kv_heads,
+    double scale, int64_t sparse_mode,
+    aclTensor* out)                     // [B, S, Hq*Dh]
+{
+    aclTensor* k_arr[] = {k};
+    aclTensor* v_arr[] = {v};
+    aclTensorList* k_list = aclCreateTensorList(k_arr, 1);
+    aclTensorList* v_list = aclCreateTensorList(v_arr, 1);
+    aclIntArray* sq   = aclCreateIntArray(actual_seq_lens.data(),    (uint64_t)actual_seq_lens.size());
+    aclIntArray* skv  = aclCreateIntArray(actual_seq_lens_kv.data(), (uint64_t)actual_seq_lens_kv.size());
+
+    uint64_t ws = 0;
+    aclOpExecutor* exec = nullptr;
+    ACLNN_CHECK(aclnnFusedInferAttentionScoreGetWorkspaceSize(
+        q, k_list, v_list,
+        nullptr,            // pseShift
+        atten_mask,
+        sq, skv,
+        nullptr, nullptr, nullptr, nullptr, nullptr,    // dequant/quant scales
+        nullptr, nullptr,                                // antiquant
+        nullptr, nullptr, nullptr,                       // block_table, q_padding, kv_padding
+        num_heads,
+        scale,
+        2147483647, 2147483647,                         // pre/next tokens (no limit)
+        (char*)"BSH",
+        num_kv_heads,
+        sparse_mode,
+        0,                                               // inner_precise
+        0, 0,                                            // block_size, antiquant_mode
+        false,                                           // softmax_lse_flag
+        out, nullptr,
+        &ws, &exec));
+
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnFusedInferAttentionScore(wp, ws, exec, stream));
+    // See note on grouped_matmul_v4 — intentionally leak lists to avoid double-free with caller RAII.
+    (void)k_list; (void)v_list;
+    aclDestroyIntArray(sq);
+    aclDestroyIntArray(skv);
+}
+
+// ---- "Linear" helper: y = x @ W.T where W is stored as [out_features, in_features] (HF convention).
+// Achieved by viewing W as [in_features, out_features] with stride [1, in_features] (elements).
+// Returns y [N, out_features].
+// Caller allocates y.
+inline void linear_hf(aclrtStream stream,
+                      aclTensor* x,                       // [N, in_features]
+                      void* W_data, aclDataType dtype,
+                      int64_t out_features, int64_t in_features,
+                      aclTensor* y_out)                   // [N, out_features]
+{
+    auto W_view = make_acl_tensor(W_data, dtype,
+                                  {in_features, out_features},
+                                  {1, in_features});  // strides: d0=1 elem, d1=in_features elems
+    matmul(stream, x, W_view.get(), y_out);
+}

include/device_weights.h

@@ -0,0 +1,82 @@
+// device_weights.h — load safetensors weights to device memory with proper TP shard.
+//
+// For M3 (attention only): loads attention + norm weights. MoE expert weights come in M4.
+//
+#pragma once
+#include "acl_common.h"
+#include "model_config.h"
+#include "safetensors_loader.h"
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+// Per-layer MoE weights on device (BF16).
+// After loading: weights are in GMM-ready layout [E, K_in, N_out] row-major contiguous.
+// For gate/up:  K_in=D,   N_out=I_per_rank
+// For down:     K_in=I,   N_out=D
+struct LayerMoEWeights {
+    DeviceBuffer router;       // [E, D]            BF16 replicated
+    DeviceBuffer gate_exps;    // [E, D, I_per_rank]  (permuted from HF [E, I, D])
+    DeviceBuffer up_exps;      // [E, D, I_per_rank]
+    DeviceBuffer down_exps;    // [E, I_per_rank, D]  (permuted from HF [E, D, I])
+};
+
+// Per-layer attention weights on device (BF16 unless noted).
+struct LayerAttnWeights {
+    DeviceBuffer input_layernorm;          // [D]            BF16
+    DeviceBuffer post_attention_layernorm; // [D]            BF16
+    // Q/K/V/O projections. HF stores as [out, in] BF16.
+    // For M3 we keep HF layout as-is; matmul wrappers handle the transpose via aclnnMm semantics.
+    DeviceBuffer q_proj;   // [Q_full,  D] on rank, but physical stored as [Q_rank, D] (sliced by head)
+    DeviceBuffer k_proj;   // [KV, D]       (replicated if tp_size > num_kv_heads)
+    DeviceBuffer v_proj;   // [KV, D]
+    DeviceBuffer o_proj;   // [D, Q_rank]   (row-parallel on Q dim)
+    DeviceBuffer q_norm;   // [head_dim]    BF16 (Qwen3 per-head norm)
+    DeviceBuffer k_norm;   // [head_dim]    BF16
+};
+
+// Shared model weights (replicated across ranks).
+struct SharedWeights {
+    DeviceBuffer embed_tokens;   // [vocab, D]
+    DeviceBuffer lm_head;        // [vocab, D]
+    DeviceBuffer final_norm;     // [D]
+};
+
+class DeviceWeightsLoader {
+public:
+    DeviceWeightsLoader(SafetensorsLoader& st, const ModelConfig& cfg)
+        : st_(st), cfg_(cfg) {}
+
+    // Load shared (embed, norm, lm_head). Replicated on every rank.
+    bool load_shared(SharedWeights& out);
+
+    // Load ONE attention layer's weights with TP sharding.
+    bool load_attention(int layer_idx, LayerAttnWeights& out);
+
+    // Load ONE MoE layer's weights. Stacks 128 experts and permutes to GMM-ready layout.
+    // stream: ACL stream for the permute op (aclnnInplaceCopy).
+    bool load_moe(int layer_idx, aclrtStream stream, LayerMoEWeights& out);
+
+    // Expose underlying safetensors for direct access (diagnostic use).
+    SafetensorsLoader& st() { return st_; }
+
+private:
+    SafetensorsLoader& st_;
+    const ModelConfig& cfg_;
+
+    // Helper: load HF tensor (full shape) into device buffer (simple H2D).
+    bool load_tensor_full_(const std::string& name, DeviceBuffer& buf);
+
+    // Helper: load HF tensor and keep only [row_lo, row_hi) of first dim (TP shard by "out" dim).
+    // HF format: tensor has shape [D0, D1, ...] stored row-major. We take rows [lo, hi) to form
+    // a sharded tensor of shape [hi-lo, D1, ...].
+    bool load_tensor_row_slice_(const std::string& name,
+                                 int64_t row_lo, int64_t row_hi,
+                                 DeviceBuffer& buf);
+
+    // TP shard by "in" dim (second axis for 2D, etc.) — used for o_proj (row-parallel).
+    bool load_tensor_col_slice_(const std::string& name,
+                                 int64_t col_lo, int64_t col_hi,
+                                 DeviceBuffer& buf);
+};

include/engine.h

@@ -0,0 +1,354 @@
+// engine.h — single-layer forward functions for attention and MoE.
+//
+// Both functions operate on device tensors. The caller owns all buffers (input, output, weights,
+// KV cache slots, scratch). They take RoPE cos/sin tables and act as pure forward kernels.
+//
+// Design goals:
+//   - Zero allocations per call (all scratch is passed in)
+//   - Same signature works for prefill (S>=1) and decode (S=1); caller picks sparse_mode.
+//   - Residual connection is NOT included (caller decides when to add residual).
+#pragma once
+#include "acl_common.h"
+#include "aclnn_ops.h"
+#include "device_weights.h"
+#include "hccl_comm.h"
+#include "model_config.h"
+#include "rope.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <tuple>
+#include <vector>
+
+// Bf16 conversion helpers used by fill_cos_sin.
+static inline uint16_t _engine_f2bf16(float x) {
+    uint32_t u; std::memcpy(&u, &x, 4);
+    return (uint16_t)((u + 0x7FFF + ((u >> 16) & 1)) >> 16);
+}
+
+// Fill cos/sin tables for positions [p0, p0+L) with HF half-half layout. Returns
+// contiguous [L*Dh] BF16 in provided host vectors (caller uploads to device).
+inline void fill_cos_sin_hf(std::vector<uint16_t>& cos_h, std::vector<uint16_t>& sin_h,
+                            int64_t p0, int64_t L, int64_t Dh, float theta) {
+    cos_h.resize(L * Dh);
+    sin_h.resize(L * Dh);
+    int64_t half = Dh / 2;
+    for (int64_t s = 0; s < L; s++) {
+        for (int64_t d = 0; d < Dh; d++) {
+            int64_t pair = (d < half) ? d : (d - half);
+            float theta_pair = 1.0f / std::pow(theta, (2.0f * pair) / Dh);
+            float angle = (float)(p0 + s) * theta_pair;
+            cos_h[s * Dh + d] = _engine_f2bf16(std::cos(angle));
+            sin_h[s * Dh + d] = _engine_f2bf16(std::sin(angle));
+        }
+    }
+}
+
+// Precomputed RoPE cos/sin table: BF16 [max_seq, Dh]. One-time cost per runtime.
+struct RopeCache {
+    DeviceBuffer cos;   // [max_seq, Dh] BF16
+    DeviceBuffer sin;   // [max_seq, Dh] BF16
+    int64_t max_seq = 0;
+    int64_t head_dim = 0;
+    float theta = 0.0f;
+};
+
+inline bool rope_cache_build(RopeCache& rc, int64_t max_seq, int64_t head_dim, float theta) {
+    std::vector<uint16_t> cos_h, sin_h;
+    fill_cos_sin_hf(cos_h, sin_h, /*p0=*/0, max_seq, head_dim, theta);
+    rc.cos.alloc(max_seq * head_dim * 2);
+    rc.sin.alloc(max_seq * head_dim * 2);
+    ACL_CHECK(aclrtMemcpy(rc.cos.get(), cos_h.size() * 2, cos_h.data(), cos_h.size() * 2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(rc.sin.get(), sin_h.size() * 2, sin_h.data(), sin_h.size() * 2, ACL_MEMCPY_HOST_TO_DEVICE));
+    rc.max_seq = max_seq; rc.head_dim = head_dim; rc.theta = theta;
+    return true;
+}
+
+// Attention forward for a single layer.
+//
+//   x_in  [S, D]    (hidden state, pre input_layernorm)
+//   x_out [S, D]    (attention output — NOT residual-added)
+//
+// K cache / V cache are contiguous [MAX_LEN, KV_DIM] BF16 buffers. This call writes new
+// positions at [past_len, past_len+S) and then runs FIAS over [0, past_len+S).
+//
+// Scratch requirements:
+//   q_scratch:    S * Q_DIM * 2 bytes
+//   k_scratch:    S * KV_DIM * 2 bytes
+//   v_scratch:    S * KV_DIM * 2 bytes
+//   xn_scratch:   S * D * 2 bytes
+//   rstd_scratch: S * 4 bytes (RmsNorm rstd output)
+//   rope_scratch: S * Hq * Dh * 2 bytes
+//
+// mask: [1, 1, 2048, 2048] bool for prefill; ignored (pass nullptr) for decode.
+inline void attention_forward(
+    aclrtStream stream,
+    const ModelConfig& cfg,
+    LayerAttnWeights& w,
+    void* x_in,                 // [S, D] BF16
+    int64_t S,
+    int64_t past_len,           // prior KV positions
+    void* k_cache, void* v_cache, int64_t max_len,
+    aclTensor* mask_tensor,     // may be nullptr for decode
+    void* q_scratch, void* k_scratch, void* v_scratch,
+    void* xn_scratch, void* rstd_scratch, void* rope_scratch,
+    void* attn_out_scratch,     // S * Q_DIM * 2 bytes (FIAS output before o_proj)
+    void* x_out,                // [S, D] BF16
+    HcclCtx* hccl_ctx = nullptr, // if tp_size > 1, AllReduce x_out after o_proj
+    const RopeCache* rope_cache = nullptr,  // if provided, use cached cos/sin table; avoids per-call H2D
+    int64_t sparse_mode = -1    // -1=auto (3 for prefill, 0 for decode); explicit 0/3 overrides
+) {
+    const int64_t D = cfg.hidden_size;
+    const int64_t Hq = cfg.n_heads_per_rank;
+    const int64_t Hkv = cfg.n_kv_heads_per_rank;
+    const int64_t Dh = cfg.head_dim;
+    const int64_t Q_DIM = Hq * Dh;
+    const int64_t KV_DIM = Hkv * Dh;
+    const double scale = 1.0 / std::sqrt((double)Dh);
+    const double eps = cfg.rms_norm_eps;
+    const float theta = cfg.rope_theta;
+
+    // 1. Input layernorm: xn = rmsnorm(x_in, input_layernorm_weight)
+    auto t_x   = make_contig_tensor(x_in,         ACL_BF16, {S, D});
+    auto t_xn  = make_contig_tensor(xn_scratch,   ACL_BF16, {S, D});
+    auto t_lnw = make_contig_tensor(w.input_layernorm.get(), ACL_BF16, {D});
+    auto t_rstd = make_contig_tensor(rstd_scratch, ACL_FLOAT, {S});
+    rms_norm(stream, t_x.get(), t_lnw.get(), eps, t_xn.get(), t_rstd.get());
+
+    // 2. Q/K/V projection
+    auto t_q = make_contig_tensor(q_scratch, ACL_BF16, {S, Q_DIM});
+    auto t_k = make_contig_tensor(k_scratch, ACL_BF16, {S, KV_DIM});
+    auto t_v = make_contig_tensor(v_scratch, ACL_BF16, {S, KV_DIM});
+    linear_hf(stream, t_xn.get(), w.q_proj.get(), ACL_BF16, Q_DIM,  D, t_q.get());
+    linear_hf(stream, t_xn.get(), w.k_proj.get(), ACL_BF16, KV_DIM, D, t_k.get());
+    linear_hf(stream, t_xn.get(), w.v_proj.get(), ACL_BF16, KV_DIM, D, t_v.get());
+
+    // 3. Per-head q_norm, k_norm
+    auto t_q_4d = make_contig_tensor(q_scratch, ACL_BF16, {1, S, Hq,  Dh});
+    auto t_k_4d = make_contig_tensor(k_scratch, ACL_BF16, {1, S, Hkv, Dh});
+    auto t_qn_w = make_contig_tensor(w.q_norm.get(), ACL_BF16, {Dh});
+    auto t_kn_w = make_contig_tensor(w.k_norm.get(), ACL_BF16, {Dh});
+    // reuse rstd_scratch split or allocate? reuse xn_scratch's first S*Hq*4 bytes.
+    // Simpler: require rstd_scratch to have max(S, S*max(Hq,Hkv)) * 4 bytes.
+    // For single-rank attention tests we pass enough.
+    auto t_rstd_q = make_contig_tensor(rstd_scratch, ACL_FLOAT, {1, S, Hq});
+    auto t_rstd_k = make_contig_tensor(rstd_scratch, ACL_FLOAT, {1, S, Hkv});
+    rms_norm(stream, t_q_4d.get(), t_qn_w.get(), eps, t_q_4d.get(), t_rstd_q.get());
+    rms_norm(stream, t_k_4d.get(), t_kn_w.get(), eps, t_k_4d.get(), t_rstd_k.get());
+
+    // 4. RoPE: positions [past_len, past_len + S). Fused aclnnApplyRotaryPosEmbV2 is 1 op
+    // vs 8-op manual version — saves ~7 kernel launches/layer × 94 layers = 658/token.
+    if (rope_cache && rope_cache->cos.get() && past_len + S <= rope_cache->max_seq) {
+        void* cos_ptr = (char*)rope_cache->cos.get() + past_len * Dh * 2;
+        void* sin_ptr = (char*)rope_cache->sin.get() + past_len * Dh * 2;
+        apply_rope_fused(stream, q_scratch, 1, S, Hq, Dh, k_scratch, Hkv, cos_ptr, sin_ptr);
+    } else {
+        std::vector<uint16_t> cos_h, sin_h;
+        fill_cos_sin_hf(cos_h, sin_h, past_len, S, Dh, theta);
+        DeviceBuffer cos_dev(S * Dh * 2), sin_dev(S * Dh * 2);
+        ACL_CHECK(aclrtMemcpy(cos_dev.get(), S*Dh*2, cos_h.data(), S*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+        ACL_CHECK(aclrtMemcpy(sin_dev.get(), S*Dh*2, sin_h.data(), S*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+        apply_rope_manual(stream, q_scratch, 1, S, Hq, Dh, k_scratch, Hkv,
+                          cos_dev.get(), sin_dev.get(), rope_scratch);
+        // Local DeviceBuffers would be freed on return while async kernels still read them.
+        ACL_CHECK(aclrtSynchronizeStream(stream));
+    }
+
+    // 5. Append K, V to cache at [past_len, past_len + S)
+    ACL_CHECK(aclrtMemcpyAsync((char*)k_cache + past_len * KV_DIM * 2, S * KV_DIM * 2,
+                               k_scratch, S * KV_DIM * 2, ACL_MEMCPY_DEVICE_TO_DEVICE, stream));
+    ACL_CHECK(aclrtMemcpyAsync((char*)v_cache + past_len * KV_DIM * 2, S * KV_DIM * 2,
+                               v_scratch, S * KV_DIM * 2, ACL_MEMCPY_DEVICE_TO_DEVICE, stream));
+
+    // 6. FIAS: q [1, S, Q_DIM], k/v [1, kv_len, KV_DIM] from cache
+    int64_t kv_len = past_len + S;
+    auto t_q_bsh = make_contig_tensor(q_scratch,  ACL_BF16, {1, S,      Q_DIM});
+    auto t_k_bsh = make_contig_tensor(k_cache,    ACL_BF16, {1, kv_len, KV_DIM});
+    auto t_v_bsh = make_contig_tensor(v_cache,    ACL_BF16, {1, kv_len, KV_DIM});
+    // FIAS writes to a separate buffer (attn_out_scratch) — aliasing q→out is unsafe.
+    auto t_attn_out_bsh = make_contig_tensor(attn_out_scratch, ACL_BF16, {1, S, Q_DIM});
+    // sparse_mode selection:
+    //   3 = left-top causal (prefill, q.S == kv.S with 2048 mask)
+    //   0 = user mask (decode with cache, batch verify)
+    //  -1 (sentinel) = auto: 3 if mask given & past_len==0 & S>1 (prefill), else 0
+    int64_t sparse = sparse_mode;
+    if (sparse < 0) {
+        sparse = (mask_tensor != nullptr && past_len == 0 && S > 1) ? 3 : 0;
+    }
+    fused_infer_attention_score(
+        stream, t_q_bsh.get(), t_k_bsh.get(), t_v_bsh.get(),
+        mask_tensor, {S}, {kv_len},
+        Hq, Hkv, scale, sparse, t_attn_out_bsh.get());
+
+    // 7. O projection: y = attn_out @ o_proj.T → [S, D]
+    auto t_attn_2d = make_contig_tensor(attn_out_scratch, ACL_BF16, {S, Q_DIM});
+    auto t_out     = make_contig_tensor(x_out,     ACL_BF16, {S, D});
+    linear_hf(stream, t_attn_2d.get(), w.o_proj.get(), ACL_BF16, D, Q_DIM, t_out.get());
+
+    // 8. TP AllReduce on x_out (row-parallel o_proj → SUM across ranks)
+    if (hccl_ctx && hccl_ctx->tp_size > 1) {
+        hccl_allreduce_bf16(*hccl_ctx, x_out, S * D, stream);
+    }
+}
+
+// MoE forward for a single layer. Residual NOT applied here.
+//
+//   x_in  [S, D]    (hidden state, pre post_attention_layernorm)
+//   x_out [S, D]    (MoE output)
+//
+// Scratch:
+//   xn_scratch:         S * D * 2
+//   rstd_scratch:       S * 4
+//   logits_scratch:     S * E * 2
+//   topk_w_scratch:     S * K * 2
+//   topk_idx_scratch:   S * K * 4
+//   row_idx_scratch:    S * K * 4  (gating output unused)
+//   expanded_x_scratch: TOTAL * D * 2
+//   expanded_ri_scratch:TOTAL * 4
+//   tpe_scratch:        E * 8
+//   fwd_dev:            TOTAL * 8
+//   gate_out_scratch:   TOTAL * I * 2
+//   up_out_scratch:     TOTAL * I * 2
+//   down_out_scratch:   TOTAL * D * 2
+//   packed_scratch:     TOTAL * D * 2
+//   weighted_scratch:   S * K * D * 2
+//
+// where TOTAL = S * K, I = cfg.i_per_rank, E = cfg.num_experts, K = cfg.num_experts_per_tok.
+//
+// IMPORTANT: post_attention_layernorm weight in `attn_w` (not in LayerMoEWeights).
+inline void moe_forward(
+    aclrtStream stream,
+    const ModelConfig& cfg,
+    LayerAttnWeights& attn_w,            // for post_attention_layernorm
+    LayerMoEWeights& w,
+    void* x_in, int64_t S,
+    void* xn_scratch, void* rstd_scratch,
+    void* logits_scratch,
+    void* topk_w_scratch, void* topk_idx_scratch, void* row_idx_scratch,
+    void* expanded_x_scratch, void* expanded_ri_scratch, void* tpe_scratch,
+    void* fwd_scratch,
+    void* gate_out_scratch, void* up_out_scratch, void* down_out_scratch,
+    void* packed_scratch, void* weighted_scratch,
+    void* x_out,
+    HcclCtx* hccl_ctx = nullptr, // if tp_size > 1, AllReduce after reduce_sum
+    void* norm_sum_scratch = nullptr  // S * 2 bytes — persistent buffer for topk_w normalize
+) {
+    const int64_t D = cfg.hidden_size;
+    const int64_t I = cfg.i_per_rank;
+    const int64_t E = cfg.num_experts;
+    const int64_t K = cfg.num_experts_per_tok;
+    const double eps = cfg.rms_norm_eps;
+    const int64_t TOTAL = S * K;
+
+    // 1. post_attention_layernorm
+    auto t_x    = make_contig_tensor(x_in,         ACL_BF16, {S, D});
+    auto t_xn   = make_contig_tensor(xn_scratch,   ACL_BF16, {S, D});
+    auto t_lnw  = make_contig_tensor(attn_w.post_attention_layernorm.get(), ACL_BF16, {D});
+    auto t_rstd = make_contig_tensor(rstd_scratch, ACL_FLOAT, {S});
+    rms_norm(stream, t_x.get(), t_lnw.get(), eps, t_xn.get(), t_rstd.get());
+
+    // 2. Router linear: logits = xn @ router.T → [S, E]
+    auto t_logits = make_contig_tensor(logits_scratch, ACL_BF16, {S, E});
+    linear_hf(stream, t_xn.get(), w.router.get(), ACL_BF16, E, D, t_logits.get());
+
+    // 3. TopK softmax
+    auto t_topk_w   = make_contig_tensor(topk_w_scratch,   ACL_BF16,  {S, K});
+    auto t_topk_idx = make_contig_tensor(topk_idx_scratch, ACL_INT32, {S, K});
+    auto t_row_idx  = make_contig_tensor(row_idx_scratch,  ACL_INT32, {S, K});
+    moe_gating_topk_softmax(stream, t_logits.get(), K, t_topk_w.get(), t_topk_idx.get(), t_row_idx.get());
+
+    // 4. Device-side normalize topk weights (Qwen3 norm_topk_prob=true).
+    //   sum = reduce_sum(topk_w, dim=-1, keepdim=true)   # [S, 1]  F32 in rstd_scratch
+    //   sum += 1e-20
+    //   sum_bf16 = cast(sum, BF16)                        # [S, 1]  in norm_sum_scratch (caller-owned)
+    //   topk_w /= sum_bf16                                # broadcast divide
+    // No per-layer syncs — all scratch buffers persist across layers.
+    if (norm_sum_scratch) {
+        auto t_sum      = make_contig_tensor(rstd_scratch,      ACL_FLOAT, {S, 1});
+        auto t_sum_bf16 = make_contig_tensor(norm_sum_scratch,  ACL_BF16,  {S, 1});
+        reduce_sum(stream, t_topk_w.get(), {-1}, /*keep_dims=*/true, ACL_FLOAT, t_sum.get());
+        inplace_adds(stream, t_sum.get(), 1e-20);
+        cast(stream, t_sum.get(), ACL_BF16, t_sum_bf16.get());
+        div_tensor(stream, t_topk_w.get(), t_sum_bf16.get(), t_topk_w.get());
+    } else {
+        // Fallback: host-side normalize (for callers that didn't provide scratch).
+        ACL_CHECK(aclrtSynchronizeStream(stream));
+        std::vector<uint16_t> h_tw(S * K);
+        ACL_CHECK(aclrtMemcpy(h_tw.data(), S*K*2, topk_w_scratch, S*K*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        for (int s = 0; s < S; s++) {
+            float sum = 0;
+            for (int k = 0; k < K; k++) {
+                uint32_t u = (uint32_t)h_tw[s*K + k] << 16;
+                float v; std::memcpy(&v, &u, 4);
+                sum += v;
+            }
+            sum += 1e-20f;
+            for (int k = 0; k < K; k++) {
+                uint32_t u = (uint32_t)h_tw[s*K + k] << 16;
+                float v; std::memcpy(&v, &u, 4);
+                v /= sum;
+                std::memcpy(&u, &v, 4);
+                h_tw[s*K + k] = (uint16_t)((u + 0x7FFF + ((u >> 16) & 1)) >> 16);
+            }
+        }
+        ACL_CHECK(aclrtMemcpy(topk_w_scratch, S*K*2, h_tw.data(), S*K*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    }
+
+    // 5. MoE init routing
+    auto t_ex_x  = make_contig_tensor(expanded_x_scratch,  ACL_BF16,  {TOTAL, D});
+    auto t_ex_ri = make_contig_tensor(expanded_ri_scratch, ACL_INT32, {TOTAL});
+    auto t_tpe   = make_contig_tensor(tpe_scratch,         ACL_INT64, {E});
+    moe_init_routing_v3(stream, t_xn.get(), t_topk_idx.get(),
+                        E, TOTAL, t_ex_x.get(), t_ex_ri.get(), t_tpe.get());
+
+    // 6. GMM gate + up
+    auto t_gate_out = make_contig_tensor(gate_out_scratch, ACL_BF16, {TOTAL, I});
+    auto t_up_out   = make_contig_tensor(up_out_scratch,   ACL_BF16, {TOTAL, I});
+    auto t_w_gate   = make_contig_tensor(w.gate_exps.get(), ACL_BF16, {E, D, I});
+    auto t_w_up     = make_contig_tensor(w.up_exps.get(),   ACL_BF16, {E, D, I});
+    grouped_matmul_v4(stream, t_ex_x.get(), t_w_gate.get(), t_tpe.get(), t_gate_out.get(), 1);
+    grouped_matmul_v4(stream, t_ex_x.get(), t_w_up.get(),   t_tpe.get(), t_up_out.get(),   1);
+
+    // 7. SwiGLU: gate_out = silu(gate_out) * up_out
+    silu(stream, t_gate_out.get(), t_gate_out.get());
+    mul(stream, t_gate_out.get(), t_up_out.get(), t_gate_out.get());
+
+    // 8. GMM down
+    auto t_down_out = make_contig_tensor(down_out_scratch, ACL_BF16, {TOTAL, D});
+    auto t_w_down   = make_contig_tensor(w.down_exps.get(), ACL_BF16, {E, I, D});
+    grouped_matmul_v4(stream, t_gate_out.get(), t_w_down.get(), t_tpe.get(), t_down_out.get(), 1);
+
+    // 9. Device-side finalize: build forward perm via two consecutive argsorts on topk_idx.
+    // No host sync — safe for graph capture.
+    //   inv_fwd = argsort(topk_idx.flat)   // each (n,k) → sorted position (primary key: expert_id)
+    //   fwd     = argsort(inv_fwd)          // inverse perm — what IndexSelect needs
+    // Stability: aclnnArgsort preserves input order for equal keys; flat index = n*K + k orders
+    // ties by n-then-k, matching our previous manual sort convention.
+    //
+    // Scratch for inv_fwd: reuse first TOTAL*8 bytes of weighted_scratch (gets overwritten
+    // by the subsequent mul op, so aliasing is safe).
+    {
+        auto t_topk_idx_flat = make_contig_tensor(topk_idx_scratch,  ACL_INT32, {TOTAL});
+        auto t_inv_fwd       = make_contig_tensor(weighted_scratch,  ACL_INT64, {TOTAL});
+        auto t_fwd_64        = make_contig_tensor(fwd_scratch,       ACL_INT64, {TOTAL});
+        argsort(stream, t_topk_idx_flat.get(), /*dim=*/0, /*descending=*/false, t_inv_fwd.get());
+        argsort(stream, t_inv_fwd.get(),       /*dim=*/0, /*descending=*/false, t_fwd_64.get());
+    }
+    auto t_fwd = make_contig_tensor(fwd_scratch, ACL_INT64, {TOTAL});
+    auto t_packed = make_contig_tensor(packed_scratch, ACL_BF16, {TOTAL, D});
+    index_select(stream, t_down_out.get(), 0, t_fwd.get(), t_packed.get());
+
+    auto t_packed_3d = make_contig_tensor(packed_scratch, ACL_BF16, {S, K, D});
+    auto t_topk_w_3d = make_contig_tensor(topk_w_scratch, ACL_BF16, {S, K, 1});
+    auto t_weighted  = make_contig_tensor(weighted_scratch, ACL_BF16, {S, K, D});
+    mul(stream, t_packed_3d.get(), t_topk_w_3d.get(), t_weighted.get());
+
+    auto t_out = make_contig_tensor(x_out, ACL_BF16, {S, D});
+    reduce_sum(stream, t_weighted.get(), {1}, false, ACL_BF16, t_out.get());
+
+    // TP AllReduce on MoE output (column-parallel experts → SUM partial intermediate outputs)
+    if (hccl_ctx && hccl_ctx->tp_size > 1) {
+        hccl_allreduce_bf16(*hccl_ctx, x_out, S * D, stream);
+    }
+}

include/hccl_comm.h

@@ -0,0 +1,106 @@
+// hccl_comm.h — minimal HCCL wrapper for TP=N AllReduce.
+//
+// Multi-process mode (each rank is a separate process, device 0 each):
+//   - Rank 0 calls HcclGetRootInfo, writes to /tmp/hccl_root_info.bin
+//   - Rank 1..N-1 wait for that file, read it
+//   - All ranks call HcclCommInitRootInfo → shared HcclComm
+//   - allreduce() does in-place HcclAllReduce with SUM op
+//
+// Launcher sets HCCL_WHITELIST_DISABLE=1, ASCEND_RT_VISIBLE_DEVICES=<rank>, etc.
+#pragma once
+#include <hccl/hccl.h>
+#include <hccl/hccl_types.h>
+#include <acl/acl.h>
+
+#include <chrono>
+#include <cstdio>
+#include <cstring>
+#include <string>
+#include <thread>
+
+#define HCCL_ROOT_INFO_PATH "/tmp/hccl_root_info.bin"
+
+struct HcclCtx {
+    HcclComm comm = nullptr;
+    int tp_size = 1;
+    int tp_rank = 0;
+    bool initialized = false;
+};
+
+inline bool hccl_init(HcclCtx& ctx, int tp_size, int tp_rank) {
+    if (tp_size <= 1) { ctx.tp_size = 1; ctx.tp_rank = 0; ctx.initialized = true; return true; }
+    ctx.tp_size = tp_size;
+    ctx.tp_rank = tp_rank;
+
+    HcclRootInfo rootInfo;
+    std::memset(&rootInfo, 0, sizeof(rootInfo));
+
+    if (tp_rank == 0) {
+        if (HcclGetRootInfo(&rootInfo) != HCCL_SUCCESS) {
+            fprintf(stderr, "[HCCL] HcclGetRootInfo failed\n"); return false;
+        }
+        FILE* f = fopen(HCCL_ROOT_INFO_PATH, "wb");
+        if (!f) { fprintf(stderr, "[HCCL] cannot write %s\n", HCCL_ROOT_INFO_PATH); return false; }
+        fwrite(&rootInfo, sizeof(rootInfo), 1, f);
+        fclose(f);
+    } else {
+        bool found = false;
+        for (int r = 0; r < 600; r++) {   // 60s timeout
+            FILE* f = fopen(HCCL_ROOT_INFO_PATH, "rb");
+            if (f) {
+                size_t rd = fread(&rootInfo, 1, sizeof(rootInfo), f);
+                fclose(f);
+                if (rd == sizeof(rootInfo)) { found = true; break; }
+            }
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
+        if (!found) { fprintf(stderr, "[HCCL] rank %d timeout waiting for root info\n", tp_rank); return false; }
+    }
+
+    HcclResult r = HcclCommInitRootInfo((uint32_t)tp_size, &rootInfo, (uint32_t)tp_rank, &ctx.comm);
+    if (r != HCCL_SUCCESS) {
+        fprintf(stderr, "[HCCL] HcclCommInitRootInfo failed: %d (rank=%d)\n", (int)r, tp_rank);
+        return false;
+    }
+    ctx.initialized = true;
+    fprintf(stderr, "[HCCL] rank %d/%d comm OK\n", tp_rank, tp_size);
+    return true;
+}
+
+// In-place AllReduce SUM on BF16 tensor. dtype = HCCL_DATA_TYPE_BFP16.
+inline bool hccl_allreduce_bf16(const HcclCtx& ctx, void* data, int64_t count, aclrtStream stream) {
+    if (!ctx.initialized) return false;
+    if (ctx.tp_size <= 1) return true;   // no-op
+
+    HcclResult r = HcclAllReduce(data, data, (uint64_t)count,
+                                  HCCL_DATA_TYPE_BFP16, HCCL_REDUCE_SUM,
+                                  ctx.comm, stream);
+    if (r != HCCL_SUCCESS) {
+        fprintf(stderr, "[HCCL] AllReduce failed: %d\n", (int)r);
+        return false;
+    }
+    return true;
+}
+
+// Broadcast buffer from root (rank 0) to all ranks. Used to share prompt tokens across ranks.
+// `data_dev` must be device memory. dtype generic (e.g., HCCL_DATA_TYPE_INT32).
+inline bool hccl_broadcast(const HcclCtx& ctx, void* data_dev, int64_t count,
+                           HcclDataType dtype, uint32_t root, aclrtStream stream) {
+    if (!ctx.initialized) return false;
+    if (ctx.tp_size <= 1) return true;
+
+    HcclResult r = HcclBroadcast(data_dev, (uint64_t)count, dtype, root, ctx.comm, stream);
+    if (r != HCCL_SUCCESS) {
+        fprintf(stderr, "[HCCL] Broadcast failed: %d\n", (int)r);
+        return false;
+    }
+    return true;
+}
+
+inline void hccl_shutdown(HcclCtx& ctx) {
+    if (ctx.comm) {
+        HcclCommDestroy(ctx.comm);
+        ctx.comm = nullptr;
+    }
+    ctx.initialized = false;
+}

include/model_config.h

@@ -0,0 +1,52 @@
+// model_config.h — Qwen3 hparams loaded from HF config.json, plus TP-derived per-rank sizes.
+#pragma once
+#include <cstdint>
+#include <string>
+
+struct ModelConfig {
+    // ---- Raw hparams from config.json ----
+    int64_t vocab_size               = 0;
+    int64_t hidden_size              = 0;   // D
+    int64_t intermediate_size        = 0;   // dense FFN (not used for MoE layers; kept for completeness)
+    int64_t moe_intermediate_size    = 0;   // I per expert
+    int64_t num_hidden_layers        = 0;   // = 94 for Qwen3-235B
+    int64_t num_attention_heads      = 0;   // = 64
+    int64_t num_key_value_heads      = 0;   // = 4 (GQA)
+    int64_t head_dim                 = 0;   // = 128
+    int64_t num_experts              = 0;   // = 128
+    int64_t num_experts_per_tok      = 0;   // top_k = 8
+    int64_t max_position_embeddings  = 0;
+    float   rope_theta               = 0.0f;
+    float   rms_norm_eps             = 1e-6f;
+    bool    norm_topk_prob           = true;
+    bool    tie_word_embeddings      = false;
+    int64_t bos_token_id             = 0;
+    int64_t eos_token_id             = 0;
+
+    // ---- TP configuration ----
+    int tp_size = 1;
+    int tp_rank = 0;
+
+    // ---- Derived per-rank sizes ----
+    // Attention Q: split along num_heads (head-parallel)
+    //    n_heads_per_rank = num_attention_heads / tp_size
+    //    q_dim_per_rank   = n_heads_per_rank * head_dim
+    int64_t n_heads_per_rank     = 0;
+    int64_t q_dim_per_rank       = 0;
+
+    // Attention KV: GQA with num_kv_heads < tp_size needs special handling.
+    // For Qwen3-235B: num_kv_heads = 4, tp_size = 16 → each KV head is replicated 4× across ranks.
+    // Simple scheme: each rank computes ALL kv heads (small, 4 × 128 = 512 features)
+    //                then slices attention output for its own q heads.
+    // Alternative: split KV heads if tp_size <= num_kv_heads.
+    int64_t n_kv_heads_per_rank  = 0;
+    int64_t kv_dim_per_rank      = 0;
+
+    // MoE: intermediate dim split. Each rank holds 1/tp_size of experts' intermediate_size.
+    //    i_per_rank = moe_intermediate_size / tp_size
+    int64_t i_per_rank           = 0;
+
+    bool load_from_json(const std::string& path);
+    void compute_derived(int tp_size, int tp_rank);
+    std::string describe() const;
+};

include/rope.h

@@ -0,0 +1,94 @@
+// rope.h — Manual HF-style RoPE using basic aclnn ops.
+//
+// Formula: q_out = q * cos + rotate_half(q) * sin
+//   where rotate_half(q) = concat(-q[..., d/2:], q[..., :d/2], dim=-1)
+//
+// Tensor layout: q/k [B, S, N, Dh] BF16, cos/sin [1, S, Dh] BF16
+//   (cos/sin are broadcast across B and N dims)
+//
+#pragma once
+#include "acl_common.h"
+#include "aclnn_ops.h"
+#include <aclnnop/aclnn_apply_rotary_pos_emb_v2.h>
+
+// Fused RoPE via aclnnApplyRotaryPosEmbV2 — replaces the 8-op manual version with a single
+// op, saving ~7 launches per layer × 94 layers = ~658 launches/token. Validated on 910 initial:
+// layout=1 + rotaryMode="half" matches HF rotate_half semantics (rel=1.24e-3 vs manual).
+//
+// q_data: [B, S, Nq, Dh] BF16 (modified in place)
+// k_data: [B, S, Nk, Dh] BF16 (modified in place)
+// cos_data / sin_data: [1, S, 1, Dh] BF16  (single contiguous buffer slice from RopeCache)
+inline void apply_rope_fused(aclrtStream stream,
+                             void* q_data, int64_t B, int64_t S, int64_t Nq, int64_t Dh,
+                             void* k_data, int64_t Nk,
+                             void* cos_data, void* sin_data) {
+    const aclDataType dt = ACL_BF16;
+    auto t_q   = make_contig_tensor(q_data, dt, {B, S, Nq, Dh});
+    auto t_k   = make_contig_tensor(k_data, dt, {B, S, Nk, Dh});
+    auto t_cos = make_contig_tensor(cos_data, dt, {1, S, 1, Dh});
+    auto t_sin = make_contig_tensor(sin_data, dt, {1, S, 1, Dh});
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    char mode[] = "half";
+    ACLNN_CHECK(aclnnApplyRotaryPosEmbV2GetWorkspaceSize(
+        t_q.get(), t_k.get(), t_cos.get(), t_sin.get(),
+        /*layout=*/1, mode, &ws, &exec));
+    void* wp = (ws > 0) ? _lca_pool().alloc(ws) : nullptr;
+    ACLNN_CHECK(aclnnApplyRotaryPosEmbV2(wp, ws, exec, stream));
+}
+
+// Apply RoPE in-place to q and k.
+// q_data: pointer to [B, S, Nq, Dh] BF16 (modified in place)
+// k_data: pointer to [B, S, Nk, Dh] BF16 (modified in place)
+// cos_data, sin_data: [1, S, Dh] BF16
+// scratch_data: pointer to contiguous [B, S, max(Nq,Nk), Dh] BF16 scratch buffer for rotate_half
+inline void apply_rope_manual(aclrtStream stream,
+                              void* q_data, int64_t B, int64_t S, int64_t Nq, int64_t Dh,
+                              void* k_data, int64_t Nk,
+                              void* cos_data, void* sin_data,
+                              void* scratch_data) {
+    const aclDataType dt = ACL_BF16;
+    const size_t elem = 2;
+    const int64_t halfDh = Dh / 2;
+
+    auto process = [&](void* x_data, int64_t N) {
+        // Strides in elements (row-major [B, S, N, Dh]):
+        // stride = [S*N*Dh, N*Dh, Dh, 1]
+        const std::vector<int64_t> full_shape   = {B, S, N, Dh};
+        const std::vector<int64_t> full_stride  = {S*N*Dh, N*Dh, Dh, 1};
+        const std::vector<int64_t> half_shape   = {B, S, N, halfDh};
+        const std::vector<int64_t> half_stride  = full_stride;  // same leading 3 strides
+
+        // View x as full
+        auto t_x = make_acl_tensor(x_data, dt, full_shape, full_stride);
+
+        // View of x left half and right half (shifted pointers, same layout, last dim half)
+        auto t_x_left  = make_acl_tensor(x_data,                           dt, half_shape, half_stride);
+        auto t_x_right = make_acl_tensor((char*)x_data + halfDh*elem,     dt, half_shape, half_stride);
+
+        // rotate_half buffer view (contiguous [B, S, N, Dh])
+        const std::vector<int64_t> rh_stride = {S*N*Dh, N*Dh, Dh, 1};
+        auto t_rh       = make_acl_tensor(scratch_data,                    dt, full_shape, rh_stride);
+        auto t_rh_left  = make_acl_tensor(scratch_data,                    dt, half_shape, rh_stride);
+        auto t_rh_right = make_acl_tensor((char*)scratch_data + halfDh*elem, dt, half_shape, rh_stride);
+
+        // rh[..., :Dh/2] = -x[..., Dh/2:]
+        neg(stream, t_x_right.get(), t_rh_left.get());
+        // rh[..., Dh/2:] = x[..., :Dh/2]
+        inplace_copy(stream, t_rh_right.get(), t_x_left.get());
+
+        // cos/sin views broadcastable to [B, S, N, Dh]
+        // Original storage: [1, S, Dh]. For broadcast, use shape [1, S, 1, Dh] with strides [0, Dh, 0, 1].
+        auto t_cos = make_acl_tensor(cos_data, dt, {1, S, 1, Dh}, {0, Dh, 0, 1});
+        auto t_sin = make_acl_tensor(sin_data, dt, {1, S, 1, Dh}, {0, Dh, 0, 1});
+
+        // q_rot = q * cos + rh * sin  (use addcmul: q *= cos, then q += rh * sin)
+        // Compute tmp = q * cos (fresh buffer needed; use scratch_data is occupied by rh)
+        // Better: multiply x in place: x *= cos, then x += rh * sin
+        // aclnnMul with x as both in and out is inplace.
+        mul(stream, t_x.get(), t_cos.get(), t_x.get());         // x = x * cos
+        addcmul(stream, t_x.get(), t_rh.get(), t_sin.get(), 1); // x += 1 * (rh * sin)
+    };
+
+    process(q_data, Nq);
+    process(k_data, Nk);
+}

include/runner.h

@@ -0,0 +1,128 @@
+// runner.h — multi-layer transformer Runner for Qwen3-235B-A22B.
+//
+// Owns: shared weights, per-layer attention + MoE weights, KV cache, scratch buffers.
+// Provides: prefill(tokens) and decode(new_token) methods returning logits [vocab] on device.
+//
+// Memory budget at TP=1 for testing a SUBSET of layers (num_layers_to_load <= 94). Full 94-layer
+// inference requires TP=16 where per-rank MoE fits ~28GB.
+#pragma once
+#include "acl_common.h"
+#include "acl_runtime.h"
+#include "aclnn_ops.h"
+#include "device_weights.h"
+#include "engine.h"
+#include "hccl_comm.h"
+#include "model_config.h"
+#include "safetensors_loader.h"
+
+#include <vector>
+
+class Runner {
+public:
+    Runner() = default;
+    ~Runner() = default;
+    Runner(const Runner&) = delete;
+    Runner& operator=(const Runner&) = delete;
+
+    // Initialize runtime, open safetensors, load shared weights. tp_size/tp_rank configure
+    // MoE + attention sharding. num_layers is how many transformer blocks to load (1..94).
+    // max_seq is the maximum sequence length (for KV cache allocation).
+    bool init(const std::string& model_dir, int tp_size, int tp_rank,
+              int num_layers_to_load, int64_t max_seq, int device_id = 0);
+
+    // Prefill: ingest S>=1 tokens, produces logits [vocab] for the LAST position. Populates KV
+    // cache starting at position 0. `hidden_out` optionally returns the final hidden state [S, D].
+    bool prefill(const int32_t* tokens, int64_t S, DeviceBuffer& logits_out);
+
+    // Decode: take 1 new token, produce logits [vocab] from the new position.
+    bool decode(int32_t token, DeviceBuffer& logits_out);
+
+    // Batched decode: take S tokens as "candidate verify batch" at positions [past_len..past_len+S),
+    // produce logits [S, vocab]. Uses causal-with-past mask (token i sees past+tokens[0..i]).
+    // Foundation for speculative decoding / PLD.
+    //   tokens: [S] int32
+    //   S: 1 .. 16
+    //   all_logits_out: will hold S * vocab_size * 2 bytes BF16, row-major [S, V]
+    // Updates past_len by +S on success.
+    bool decode_batch(const int32_t* tokens, int64_t S, DeviceBuffer& all_logits_out);
+
+    // Warmup: run N dummy decode() calls (resetting cache) to pre-compile aclnn executors,
+    // warm HCCL collective buffers, and stabilize NPU thermals. Improves first-N-token latency
+    // by ~1 s (especially noticeable on short generations or REPL cold start).
+    // Call after init(); safe to call multiple times. Does NOT affect past_len.
+    void warmup(int iterations = 3);
+
+    // Accessors
+    const ModelConfig& cfg() const { return cfg_; }
+    aclrtStream stream() { return rt_.stream(); }
+    int64_t past_len() const { return past_len_; }
+    void reset_cache() { past_len_ = 0; }
+    // Rewind past_len by n. Used by speculative decoding to discard rejected draft tokens'
+    // KV cache entries (they'll be overwritten by subsequent writes).
+    void rewind_cache(int64_t n) { if (n > 0 && n <= past_len_) past_len_ -= n; }
+    HcclCtx& hccl_ctx() { return hccl_ctx_; }
+
+    // Profiling: set via LCA_PROFILE=1 env in main_cli. If enabled, decode() accumulates
+    // per-phase wall-clock ms into the timer accumulators below.
+    bool profile_enabled = false;
+    double t_embed_ms = 0, t_layers_ms = 0, t_final_ms = 0;
+    int64_t profile_calls = 0;
+    void print_profile_summary() const;
+
+private:
+    // One-layer forward: x_in [S, D] → x_out [S, D] via attention + residual + MoE + residual.
+    // Uses this layer's KV cache starting at past_len; caller updates past_len after each call.
+    // batch_decode_mode: true for S>1 at past_len>0 (spec decoding) — uses custom causal mask
+    //                    with past instead of the 2048×2048 prefill mask.
+    void layer_forward_(int layer_idx, int64_t S, void* x_in, void* x_out,
+                        bool batch_decode_mode = false);
+
+    // Build causal-with-past mask in batch_mask_dev_ for decode_batch at current past_len.
+    // Shape [1, 1, S, past_len+S] bool, mask[i, j] = 1 iff j > past_len+i.
+    void build_batch_decode_mask_(int64_t S);
+
+    // Final: final_norm + lm_head on last position → logits [vocab].
+    void final_logits_(void* hidden_last /*[1, D]*/, DeviceBuffer& logits_out);
+
+    // Batched final: final_norm + lm_head on [S, D] → logits [S, V].
+    void final_logits_batch_(void* hidden /*[S, D]*/, int64_t S, DeviceBuffer& logits_out);
+
+    AclRuntime rt_;
+    SafetensorsLoader st_;
+    ModelConfig cfg_;
+    HcclCtx hccl_ctx_;
+    int num_layers_ = 0;
+    int64_t max_seq_ = 0;
+
+    SharedWeights                 shared_;
+    std::vector<LayerAttnWeights> attn_;
+    std::vector<LayerMoEWeights>  moe_;
+
+    // Per-layer KV cache
+    std::vector<DeviceBuffer> k_cache_;
+    std::vector<DeviceBuffer> v_cache_;
+
+    // Scratch (reallocated per-call sized by current S)
+    DeviceBuffer q_sc_, k_sc_, v_sc_, xn_sc_, rstd_sc_, rope_sc_, attn_fias_sc_, attn_out_sc_;
+    DeviceBuffer moe_xn_, moe_rstd_, moe_logits_;
+    DeviceBuffer moe_topk_w_, moe_topk_idx_, moe_row_idx_;
+    DeviceBuffer moe_ex_x_, moe_ex_ri_, moe_tpe_;
+    DeviceBuffer moe_fwd_;
+    DeviceBuffer moe_gate_, moe_up_, moe_down_;
+    DeviceBuffer moe_packed_, moe_weighted_, moe_out_;
+    DeviceBuffer moe_norm_sum_;       // BF16 [S, 1] for on-device topk_w normalize
+    DeviceBuffer x_buf_a_, x_buf_b_;   // ping-pong for residual chain
+
+    // Causal mask for prefill (2048 x 2048 bool); decode uses nullptr
+    DeviceBuffer prefill_mask_dev_;
+
+    // Batch decode mask: S_MAX × KV_MAX bool, where mask[i, j] = 1 (masked out) if
+    // j > past_len + i. Built on-demand per-call (past_len changes).
+    DeviceBuffer batch_mask_dev_;
+
+    // Pre-computed RoPE cos/sin table (sized for max_seq_)
+    RopeCache rope_cache_;
+
+    int64_t past_len_ = 0;
+    int64_t cur_S_capacity_ = 0;   // scratch sized for this many tokens
+};

include/safetensors_loader.h

@@ -0,0 +1,78 @@
+// safetensors_loader.h — lazy multi-shard safetensors reader.
+//
+// Usage:
+//     SafetensorsLoader loader;
+//     loader.open("/path/to/model_dir");   // parses index.json + all shard headers
+//     auto meta = loader.get("model.layers.0.self_attn.q_proj.weight");
+//     const void* host_ptr = loader.data_ptr(meta);   // mmap-backed, host memory
+//     // copy to device: aclrtMemcpy(d_ptr, n, host_ptr, n, ACL_MEMCPY_HOST_TO_DEVICE);
+//
+// Files are mmap'd on first access and unmapped at destruction.
+//
+#pragma once
+#include <cstdint>
+#include <map>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+struct TensorMeta {
+    std::string name;
+    std::string dtype;      // "BF16", "F16", "F32", "I32", "I64"
+    std::vector<int64_t> shape;
+    int shard_id = -1;      // index into SafetensorsLoader::shards_
+    size_t offset = 0;       // byte offset within shard (after 8B header_len + JSON header)
+    size_t nbytes = 0;
+};
+
+struct ShardFile {
+    std::string path;
+    int fd = -1;
+    void* mmap_ptr = nullptr;
+    size_t mmap_size = 0;
+    size_t data_base = 0;    // byte offset to first tensor data within file
+};
+
+class SafetensorsLoader {
+public:
+    SafetensorsLoader();
+    ~SafetensorsLoader();
+
+    // Opens a HuggingFace model directory. Returns false on failure.
+    // Expects: <dir>/model.safetensors.index.json + model-XXXXX-of-YYYYY.safetensors
+    bool open(const std::string& model_dir);
+
+    // Get tensor metadata. Returns nullptr if name not found.
+    const TensorMeta* get(const std::string& name) const;
+
+    // Return host pointer to tensor's raw bytes (mmap-backed). Null if not found or mmap failed.
+    const void* data_ptr(const TensorMeta& m);
+    const void* data_ptr(const std::string& name);
+
+    // Enumerate all tensor names (stable order = lexicographic).
+    std::vector<std::string> list_tensor_names() const;
+
+    // Stats
+    size_t tensor_count() const { return tensors_.size(); }
+    size_t shard_count()  const { return shards_.size(); }
+    size_t total_bytes() const;
+
+private:
+    bool parse_shard_header_(int shard_id);
+    bool mmap_shard_(int shard_id);
+
+    std::string model_dir_;
+    std::vector<ShardFile> shards_;
+    std::map<std::string, TensorMeta> tensors_;  // ordered for determinism
+};
+
+// ---- Helpers ----
+
+// Convert safetensors dtype string to element byte size.
+inline size_t sdtype_size(const std::string& s) {
+    if (s == "F32" || s == "I32") return 4;
+    if (s == "F16" || s == "BF16" || s == "I16") return 2;
+    if (s == "F64" || s == "I64") return 8;
+    if (s == "I8"  || s == "U8"  || s == "BOOL") return 1;
+    return 0;
+}

include/tokenizer.h

@@ -0,0 +1,38 @@
+// tokenizer.h — minimal Qwen3 tokenizer.
+//
+// M2-phase1: decode() is native C++ (simple vocab lookup). encode() is a Python subprocess
+// (one-time cost at prompt setup). Native BPE encode is a future item.
+//
+#pragma once
+#include <string>
+#include <vector>
+#include <cstdint>
+
+class Tokenizer {
+public:
+    bool load(const std::string& vocab_bin_path);
+
+    // Decode a single token id to UTF-8 string.
+    std::string decode(int token_id) const;
+
+    // Decode list of token ids to concatenated UTF-8 string.
+    std::string decode(const std::vector<int>& token_ids) const;
+
+    // Encode prompt to token ids. Uses a Python subprocess since Qwen3 needs proper BPE.
+    // The subprocess call takes ~200ms but is only invoked once per prompt.
+    std::vector<int> encode_via_python(const std::string& model_dir,
+                                       const std::string& prompt,
+                                       bool apply_chat_template = false) const;
+
+    // Encode a multi-turn conversation by applying the model's chat template. Each pair is
+    // (role, content) — typical roles: "system", "user", "assistant". Uses Python subprocess.
+    std::vector<int> encode_conversation_via_python(
+        const std::string& model_dir,
+        const std::vector<std::pair<std::string, std::string>>& conversation,
+        bool add_generation_prompt = true) const;
+
+    size_t size() const { return id_to_bytes_.size(); }
+
+private:
+    std::vector<std::string> id_to_bytes_;  // id -> raw utf-8 bytes
+};

include/workspace_pool.h

@@ -0,0 +1,84 @@
+// workspace_pool.h — reusable aclnn workspace buffer pool.
+//
+// Problem: every aclnn op does `aclrtMalloc(workspace)` + `aclrtFree`. For decode at 94 layers
+// × ~30 ops = 2820 mallocs/frees per token, this is significant overhead.
+//
+// Solution: pool of DeviceBuffers, grow-only. Pool returns a pointer >= requested size.
+// Most ops reuse the SAME buffer since they don't overlap on-stream (serial execution).
+//
+// Thread safety: not thread-safe. One pool per Runner (one thread).
+#pragma once
+#include "acl_common.h"
+#include <algorithm>
+#include <vector>
+
+class WorkspacePool {
+public:
+    WorkspacePool() = default;
+    ~WorkspacePool() = default;
+    WorkspacePool(const WorkspacePool&) = delete;
+    WorkspacePool& operator=(const WorkspacePool&) = delete;
+
+    // Return a device pointer of at least `bytes`. Reuses the current buffer
+    // if it's big enough; otherwise grows by allocating a new one and
+    // **retaining old buffers** (async kernels may still be reading them —
+    // freeing too early would corrupt in-flight workspaces).
+    //
+    // Periodically call `reset_after_sync()` when the stream is idle to
+    // reclaim all-but-largest buffers and reset grow count.
+    void* alloc(size_t bytes) {
+        if (bytes == 0) return nullptr;
+        if (current_size_ < bytes) {
+            // Keep old buffer alive (don't free!) — aclnn kernels may still use it.
+            old_bufs_.push_back(std::move(buf_));
+            buf_.alloc(bytes);
+            current_size_ = bytes;
+            grow_count_++;
+        }
+        return buf_.get();
+    }
+
+    size_t current_size() const { return current_size_; }
+    size_t grow_count() const { return grow_count_; }
+    size_t retained_count() const { return old_bufs_.size(); }
+
+    // Call only when the stream is guaranteed idle (e.g., after aclrtSynchronizeStream).
+    // Drops all retained older buffers, freeing device memory. Current active buffer kept.
+    void reset_after_sync() {
+        old_bufs_.clear();
+    }
+
+    void clear() {
+        old_bufs_.clear();
+        buf_ = DeviceBuffer();
+        current_size_ = 0;
+        grow_count_ = 0;
+    }
+
+private:
+    DeviceBuffer buf_;                       // current active (largest so far)
+    std::vector<DeviceBuffer> old_bufs_;     // older, smaller — still live until stream sync
+    size_t current_size_ = 0;
+    size_t grow_count_ = 0;
+};
+
+// Convenience: per-stream RAII guard that acts like a `DeviceBuffer` but draws from pool.
+// Used in aclnn_ops.h wrappers as a drop-in replacement for the local DeviceBuffer.
+class PoolBuffer {
+public:
+    // Fallback mode: if pool is nullptr, allocate own buffer (current behavior).
+    // Pool mode: return pool's shared pointer.
+    PoolBuffer(WorkspacePool* pool, size_t bytes) {
+        if (pool) {
+            ptr_ = pool->alloc(bytes);
+        } else if (bytes > 0) {
+            local_.alloc(bytes);
+            ptr_ = local_.get();
+        }
+    }
+    void* get() { return ptr_; }
+
+private:
+    DeviceBuffer local_;   // only used when pool is null
+    void*        ptr_ = nullptr;
+};

scripts/bench_hccl.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+# bench_hccl.sh — HCCL 参数矩阵 benchmark for TG
+#
+# 遍历 HCCL_ALGO × HCCL_BUFFSIZE 组合，每组 N_RUNS 次取中位数，找最佳配置。
+# 固定 prompt + seed=0 + n_predict=200 保证可比性。
+set -u
+cd "$(dirname "$0")/.."
+
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+BIN="./build/qwen3-moe-aclnn"
+LAUNCH="./scripts/tp_launch.sh"
+TP="${TP_SIZE:-16}"
+N_PREDICT="${N_PREDICT:-150}"
+N_RUNS="${N_RUNS:-2}"
+PROMPT="${PROMPT:-The history of artificial intelligence spans several decades and}"
+VOCAB="tokenizer_data/vocab.bin"
+
+OUT=/tmp/bench_hccl_results.csv
+echo "algo,buffsize,runs,best_tgs" > $OUT
+
+run_one() {
+    local algo="$1" buf="$2"
+    local tgs=()
+    for r in $(seq 1 $N_RUNS); do
+        export HCCL_ALGO="$algo" HCCL_BUFFSIZE="$buf"
+        local out
+        out=$(${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+                --prompt "$PROMPT" --n-predict $N_PREDICT \
+                --vocab "$VOCAB" --seed 0 2>&1 | grep "decode :" | awk '{print $(NF-2)}')
+        tgs+=("${out:-0}")
+    done
+    local sorted=($(printf '%s\n' "${tgs[@]}" | sort -n))
+    local best="${sorted[-1]}"
+    local csv="$algo,$buf,${tgs[*]},$best"
+    echo "$csv" | sed 's/ /|/g' >> $OUT
+    printf "  %-22s  buf=%-4s  %s  best=%s\n" \
+        "${algo:-(auto)}" "$buf" "${tgs[*]}" "$best"
+}
+
+# Matrix
+ALGOS=("" "level0:ring" "level0:fullmesh")
+BUFSIZES=("100" "200" "400")
+
+echo "HCCL matrix: ${#ALGOS[@]} algos × ${#BUFSIZES[@]} buffsizes × ${N_RUNS} runs each"
+echo "Results → $OUT"
+echo ""
+
+for algo in "${ALGOS[@]}"; do
+    for buf in "${BUFSIZES[@]}"; do
+        run_one "$algo" "$buf"
+    done
+done
+
+echo ""
+echo "====== Summary (sorted by best TG) ======"
+(head -1 $OUT; tail -n +2 $OUT | sort -t, -k4 -gr) | column -t -s,

scripts/bench_hccl_adv.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+# bench_hccl_adv.sh — 进阶 HCCL 参数调优
+# 在已确定的 ring:200 baseline 上加入 OP_EXPANSION_MODE=AIV 等 knob
+set -u
+cd "$(dirname "$0")/.."
+
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+BIN="./build/qwen3-moe-aclnn"
+LAUNCH="./scripts/tp_launch.sh"
+TP=16
+N_PREDICT=200
+N_RUNS=2
+PROMPT="The history of artificial intelligence spans several decades and"
+VOCAB="tokenizer_data/vocab.bin"
+
+OUT=/tmp/bench_hccl_adv.csv
+echo "config,run1,run2,best,median" > $OUT
+
+run_one() {
+    local name="$1"; shift
+    # remaining args are env assignments: KEY=VALUE ...
+    local tgs=()
+    for r in $(seq 1 $N_RUNS); do
+        local out
+        # set env vars for this run
+        local env_cmd=""
+        for a in "$@"; do env_cmd="$env_cmd $a"; done
+        out=$(env HCCL_ALGO=level0:ring HCCL_BUFFSIZE=200 $@ \
+              ${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+              --prompt "$PROMPT" --n-predict $N_PREDICT \
+              --vocab "$VOCAB" --seed 0 2>&1 | grep "decode :" | awk '{print $(NF-2)}')
+        tgs+=("${out:-0}")
+    done
+    local sorted=($(printf '%s\n' "${tgs[@]}" | sort -n))
+    local best="${sorted[-1]}"
+    local median="${sorted[$((${#sorted[@]}/2))]}"
+    echo "$name,${tgs[0]},${tgs[1]},$best,$median" >> $OUT
+    printf "  %-40s  %s  best=%s median=%s\n" "$name" "${tgs[*]}" "$best" "$median"
+}
+
+echo "Adv HCCL bench: baseline ring:200 + additional knobs"
+echo "Results → $OUT"
+echo ""
+
+run_one "baseline (ring+200 only)"
+
+run_one "+ OP_EXPANSION_MODE=AIV"                        HCCL_OP_EXPANSION_MODE=AIV
+run_one "+ OP_BASE_FFTS_MODE=1"                          HCCL_OP_BASE_FFTS_MODE_ENABLE=1
+run_one "+ OP_EXPANSION=AIV + FFTS=1"                    HCCL_OP_EXPANSION_MODE=AIV HCCL_OP_BASE_FFTS_MODE_ENABLE=1
+run_one "+ OP_EXPANSION=AIV + BUF=256"                   HCCL_OP_EXPANSION_MODE=AIV HCCL_BUFFSIZE=256
+run_one "+ OP_EXPANSION=AIV + BUF=512"                   HCCL_OP_EXPANSION_MODE=AIV HCCL_BUFFSIZE=512
+run_one "+ OP_EXPANSION=AIV + ALGO=fullmesh"             HCCL_OP_EXPANSION_MODE=AIV HCCL_ALGO=level0:fullmesh
+
+echo ""
+echo "====== Sorted by best TG ======"
+(head -1 $OUT; tail -n +2 $OUT | sort -t, -k4 -gr) | column -t -s,

scripts/bench_hccl_adv2.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+# bench_hccl_adv2.sh — layer 2 env knob exploration on top of AIV+FFTS=1 baseline.
+# Target: break past 25 t/s MUST barrier.
+set -u
+cd "$(dirname "$0")/.."
+
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+BIN="./build/qwen3-moe-aclnn"
+LAUNCH="./scripts/tp_launch.sh"
+TP=16
+N_PREDICT=200
+N_RUNS=3
+LONG_PROMPT="Write a very long detailed essay about artificial intelligence, machine learning, deep learning and their applications in modern society. Include historical context, current state of the art, and future predictions."
+VOCAB="tokenizer_data/vocab.bin"
+
+OUT=/tmp/bench_hccl_adv2.csv
+echo "config,runs,best,median" > $OUT
+
+run_one() {
+    local name="$1"; shift
+    local tgs=()
+    for r in $(seq 1 $N_RUNS); do
+        local out
+        out=$(env HCCL_ALGO=level0:ring HCCL_BUFFSIZE=200 \
+                  HCCL_OP_EXPANSION_MODE=AIV HCCL_OP_BASE_FFTS_MODE_ENABLE=1 \
+                  "$@" \
+              ${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+              --prompt "$LONG_PROMPT" --n-predict $N_PREDICT \
+              --vocab "$VOCAB" --seed 0 --no-stream 2>&1 \
+              | grep "decode :" | awk '{print $(NF-2)}')
+        tgs+=("${out:-0}")
+    done
+    local sorted=($(printf '%s\n' "${tgs[@]}" | sort -n))
+    local best="${sorted[-1]}"
+    local median="${sorted[$((${#sorted[@]}/2))]}"
+    echo "$name,${tgs[*]},$best,$median" | tr ' ' '|' | sed 's/|/,/' | sed 's/|/ /g' >> $OUT
+    printf "  %-40s  %s  best=%s median=%s\n" "$name" "${tgs[*]}" "$best" "$median"
+}
+
+echo "Bench: AIV+FFTS baseline + single additional knob"
+echo "$N_RUNS runs × $N_PREDICT tokens"
+echo ""
+
+run_one "baseline (AIV + FFTS)"
+
+run_one "+ TASK_QUEUE_ENABLE=1"          TASK_QUEUE_ENABLE=1
+run_one "+ TASK_QUEUE_ENABLE=2"          TASK_QUEUE_ENABLE=2
+run_one "+ HCCL_BUFFSIZE=256"            HCCL_BUFFSIZE=256
+run_one "+ HCCL_DETERMINISTIC=false"     HCCL_DETERMINISTIC=false
+run_one "+ HCCL_INTRA_ROCE_ENABLE=1"     HCCL_INTRA_ROCE_ENABLE=1
+run_one "+ HCCL_CLUSTER_TIMEOUT=600"     HCCL_CLUSTER_TIMEOUT=600
+run_one "+ ASCEND_LAUNCH_BLOCKING=0"     ASCEND_LAUNCH_BLOCKING=0
+
+echo ""
+echo "====== Sorted by best TG ======"
+(head -1 $OUT; tail -n +2 $OUT | sort -t, -k3 -gr) | column -t -s,

scripts/bench_pld.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# bench_pld.sh — sweep K × n-gram with corrected causal-with-past mask.
+# Measures TG + accept rate stability across N_RUNS per config.
+set -u
+cd "$(dirname "$0")/.."
+
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+BIN="./build/qwen3-moe-aclnn"
+LAUNCH="./scripts/tp_launch.sh"
+TP=16
+N_PREDICT=200
+N_RUNS="${N_RUNS:-3}"
+PROMPT="${PROMPT:-Write a long Python function that computes the Fibonacci sequence with memoization, extensive comments, and type hints.}"
+VOCAB="tokenizer_data/vocab.bin"
+
+OUT=/tmp/bench_pld.csv
+echo "k,ngram,run_tgs,best,median,avg_accept" > $OUT
+
+run_one() {
+    local k="$1" ng="$2"
+    local tgs=() accs=()
+    for r in $(seq 1 $N_RUNS); do
+        local output
+        output=$(${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+                --prompt "$PROMPT" --n-predict $N_PREDICT --max-seq 512 \
+                --vocab "$VOCAB" --seed 0 --no-stream \
+                --pld --pld-k $k --pld-ngram $ng 2>&1)
+        local tg
+        tg=$(echo "$output" | grep "decode :" | awk '{print $(NF-2)}')
+        local acc
+        acc=$(echo "$output" | grep "\[pld\]" | grep -oE "avg=[0-9.]+" | cut -d= -f2)
+        tgs+=("${tg:-0}")
+        accs+=("${acc:-0}")
+    done
+    local sorted=($(printf '%s\n' "${tgs[@]}" | sort -n))
+    local n=${#sorted[@]}
+    local best="${sorted[-1]}"
+    local median="${sorted[$((n/2))]}"
+    local accs_avg=$(printf '%s\n' "${accs[@]}" | awk '{s+=$1} END {printf "%.2f", s/NR}')
+    echo "$k,$ng,$(IFS=/; echo "${tgs[*]}"),$best,$median,$accs_avg" >> $OUT
+    printf "  K=%-2d ng=%-1d  runs=[%s]  best=%s  median=%s  accept_avg=%s\n" \
+        "$k" "$ng" "${tgs[*]}" "$best" "$median" "$accs_avg"
+}
+
+echo "PLD sweep on '$PROMPT' ($N_RUNS runs × $N_PREDICT tokens)"
+echo ""
+
+for k in 2 4 6 8 12; do
+    for ng in 1 2 3; do
+        run_one $k $ng
+    done
+done
+
+# Baseline for reference
+echo ""
+echo "Baseline (no PLD):"
+tgs=()
+for r in $(seq 1 $N_RUNS); do
+    tg=$(${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+         --prompt "$PROMPT" --n-predict $N_PREDICT --max-seq 512 \
+         --vocab "$VOCAB" --seed 0 --no-stream 2>&1 | grep "decode :" | awk '{print $(NF-2)}')
+    tgs+=("${tg:-0}")
+done
+sorted=($(printf '%s\n' "${tgs[@]}" | sort -n))
+echo "  baseline:  ${tgs[*]}  median=${sorted[$((${#sorted[@]}/2))]}"
+
+echo ""
+echo "====== Sorted by median TG ======"
+(head -1 $OUT; tail -n +2 $OUT | sort -t, -k5 -gr) | column -t -s,

scripts/bench_pld_k.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# bench_pld_k.sh — isolated K sweep with FIXED K (no adaptive) to characterize raw K effect.
+# Larger K = more draft candidates per verify. Peak observed accept=7.38 suggests K=8 not saturated.
+set -u
+cd "$(dirname "$0")/.."
+
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+BIN="./build/qwen3-moe-aclnn"
+LAUNCH="./scripts/tp_launch.sh"
+TP=16
+N_PREDICT=200
+N_RUNS=3
+PROMPT="Write a long Python function that computes the Fibonacci sequence with memoization, extensive comments, and type hints."
+VOCAB="tokenizer_data/vocab.bin"
+
+OUT=/tmp/bench_pld_k.csv
+echo "k,runs,median,max,avg_accept" > $OUT
+
+for K in 4 6 8 10 12 16; do
+    tgs=() accs=()
+    for r in $(seq 1 $N_RUNS); do
+        out=$(${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+             --prompt "$PROMPT" --n-predict $N_PREDICT --max-seq 512 \
+             --vocab "$VOCAB" --seed 0 --no-stream \
+             --pld --pld-k $K --pld-ngram 1 --pld-fixed-k 2>&1)
+        tg=$(echo "$out" | grep "decode :" | awk '{print $(NF-2)}')
+        acc=$(echo "$out" | grep "\[pld\]" | grep -oE "avg=[0-9.]+" | cut -d= -f2)
+        tgs+=("${tg:-0}")
+        accs+=("${acc:-0}")
+    done
+    sorted=($(printf '%s\n' "${tgs[@]}" | sort -n))
+    median="${sorted[$((${#sorted[@]}/2))]}"
+    max="${sorted[-1]}"
+    accs_avg=$(printf '%s\n' "${accs[@]}" | awk '{s+=$1} END {printf "%.2f", s/NR}')
+    echo "$K,$(IFS=/; echo "${tgs[*]}"),$median,$max,$accs_avg" >> $OUT
+    printf "  K=%-2d  runs=[%s]  median=%s  max=%s  accept=%s\n" "$K" "${tgs[*]}" "$median" "$max" "$accs_avg"
+done
+
+echo ""
+echo "====== Sorted by median ======"
+(head -1 $OUT; tail -n +2 $OUT | sort -t, -k3 -gr) | column -t -s,

scripts/bench_pld_safe.sh

@@ -0,0 +1,154 @@
+#!/usr/bin/env bash
+# bench_pld_safe.sh — PLD benchmark with output correctness check.
+# Unlike bench_tg.sh (which only reports TG numbers), this wrapper also inspects the
+# generated text for degeneration signals (consecutive identical tokens / very low
+# distinct-token ratio in the tail) and flags runs whose high TG came from dead-loop
+# output rather than real acceleration.
+#
+# Usage:  ./scripts/bench_pld_safe.sh [N_RUNS] [PROMPT_FILE]
+#         Prompts with "|" separator: "tag|prompt text"
+#         Default: tests multiple prompt classes and reports which ones PLD helps safely.
+set -u
+cd "$(dirname "$0")/.."
+
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+BIN="./build/qwen3-moe-aclnn"
+N_RUNS="${1:-3}"
+N_PREDICT="${N_PREDICT:-120}"
+VOCAB="tokenizer_data/vocab.bin"
+
+# Default prompt suite: one per class. Override via PROMPTS env or arg 2 (file with "tag|prompt" per line).
+default_prompts=(
+    "story|Once upon a time, in a small village,"
+    "factual|The capital of France is"
+    "code|Write a Python function that computes Fibonacci."
+    "essay|The history of artificial intelligence spans several decades and"
+)
+
+if [ "${2:-}" != "" ] && [ -f "${2:-}" ]; then
+    mapfile -t prompts < "$2"
+else
+    prompts=("${default_prompts[@]}")
+fi
+
+# ----- Correctness classifier -----
+# Reads generated text from stdin, returns:
+#   OK           — no loop signals
+#   LOOP_N       — N+ consecutive identical non-space words detected
+#   LOW_DIVERSITY — tail 40 words have < 10 distinct words (heavy repetition)
+classify_output() {
+    awk '
+    {
+        # Tokenize on whitespace; strip punct at edges for comparison.
+        n = split($0, w, /[[:space:]]+/);
+        for (i = 1; i <= n; i++) {
+            gsub(/^[[:punct:]]+|[[:punct:]]+$/, "", w[i]);
+            if (w[i] == "") continue;
+            words[++nw] = tolower(w[i]);
+        }
+    }
+    END {
+        if (nw < 5) { print "OK"; exit }
+        # consecutive-same detection
+        run = 1; max_run = 1;
+        for (i = 2; i <= nw; i++) {
+            if (words[i] == words[i-1]) { run++; if (run > max_run) max_run = run; }
+            else run = 1;
+        }
+        if (max_run >= 6) { printf "LOOP_%d\n", max_run; exit }
+
+        # tail diversity: last 40 words
+        tail_start = nw - 39; if (tail_start < 1) tail_start = 1;
+        delete seen;
+        distinct = 0;
+        for (i = tail_start; i <= nw; i++) {
+            if (!(words[i] in seen)) { seen[words[i]] = 1; distinct++; }
+        }
+        tail_n = nw - tail_start + 1;
+        if (tail_n >= 20 && distinct < 10) {
+            printf "LOW_DIVERSITY_%d/%d\n", distinct, tail_n;
+            exit;
+        }
+        print "OK";
+    }'
+}
+
+run_once() {
+    local prompt="$1"
+    local extra_flags="$2"
+    # Launch. The binary prints to stdout: rank/cli headers, runner loading lines,
+    # generated text (--no-stream), then perf lines. pld/warn go to stderr.
+    local stdout_file=$(mktemp)
+    local stderr_file=$(mktemp)
+    # Ensure no lockfile leftover.
+    ssh_cleanup_lockfile
+    ./scripts/tp_launch.sh 16 $BIN --model-dir "$MODEL" \
+        --prompt "$prompt" --n-predict $N_PREDICT \
+        --vocab "$VOCAB" --seed 0 --no-stream --temperature 0 \
+        $extra_flags 1>"$stdout_file" 2>"$stderr_file"
+    # TG lives on stdout (from printf in binary).
+    local tg=$(grep "\[perf\] decode" "$stdout_file" | awk '{print $(NF-2)}')
+    # Generated text: the line that begins with the prompt (--no-stream echoes prompt+text).
+    local gen_text=$(grep -F -- "$prompt" "$stdout_file" | grep -v '^\[' | tail -1)
+    local stripped="${gen_text#$prompt}"
+    local verdict=$(echo "$stripped" | classify_output)
+    local has_warn=""
+    if grep -q "\[warn\]" "$stderr_file"; then has_warn="WARN"; fi
+    local pld_line=$(grep "\[pld\]" "$stderr_file" | tail -1 | sed 's/^\[pld\] //')
+    rm -f "$stdout_file" "$stderr_file"
+    echo "${tg:-0}|${verdict}|${has_warn}|${pld_line}"
+}
+
+ssh_cleanup_lockfile() {
+    rm -f /tmp/hccl_root_info.bin 2>/dev/null || true
+}
+
+bench_prompt() {
+    local tag="$1"; local prompt="$2"; local flags="$3"
+    echo ""
+    echo "=== [$tag] $(echo "$prompt" | head -c 50)... (flags: ${flags:-none}) ==="
+    local tgs=() verdicts=() warns=() plds=()
+    for r in $(seq 1 $N_RUNS); do
+        result=$(run_once "$prompt" "$flags")
+        IFS='|' read -r tg verdict warn pld <<< "$result"
+        printf "  run %d: TG=%s verdict=%s %s\n" "$r" "$tg" "$verdict" "$warn"
+        [ -n "$pld" ] && printf "         %s\n" "$pld"
+        tgs+=("${tg:-0}"); verdicts+=("$verdict"); warns+=("$warn")
+        rm -f /tmp/hccl_root_info.bin
+    done
+    # Split good vs degraded
+    local good_tgs=() bad_tgs=()
+    for i in "${!tgs[@]}"; do
+        if [ "${verdicts[$i]}" = "OK" ]; then good_tgs+=("${tgs[$i]}"); else bad_tgs+=("${tgs[$i]}"); fi
+    done
+    local n_good=${#good_tgs[@]}
+    local n_bad=${#bad_tgs[@]}
+    echo "  → $n_good/$N_RUNS OK, $n_bad/$N_RUNS degraded"
+    if [ $n_good -gt 0 ]; then
+        local mean=$(printf '%s\n' "${good_tgs[@]}" | awk '{s+=$1} END {printf "%.2f", s/NR}')
+        echo "  → OK mean TG: $mean t/s  (values: ${good_tgs[*]})"
+    fi
+    if [ $n_bad -gt 0 ]; then
+        local bad_mean=$(printf '%s\n' "${bad_tgs[@]}" | awk '{s+=$1} END {printf "%.2f", s/NR}')
+        echo "  → degraded mean TG: $bad_mean t/s (DO NOT REPORT as speedup)  (values: ${bad_tgs[*]})"
+    fi
+}
+
+echo "bench_pld_safe: $N_RUNS runs × $N_PREDICT tokens per prompt; comparing [no-pld, pld+guard, pld+no-guard]"
+
+for entry in "${prompts[@]}"; do
+    tag="${entry%%|*}"
+    prompt="${entry#*|}"
+    bench_prompt "$tag/base"      "$prompt" ""
+    bench_prompt "$tag/pld+guard" "$prompt" "--pld"
+    bench_prompt "$tag/pld-raw"   "$prompt" "--pld --pld-no-guard"
+done
+
+echo ""
+echo "=========================================================="
+echo "Interpretation:"
+echo "  OK mean TG is the only honest number to report."
+echo "  Any 'degraded' result with high TG is a dead-loop artifact."
+echo "  Expected: pld+guard matches or beats base on creative/story prompts,"
+echo "  matches base on factual/code prompts (drafts rejected → fallback to single decode)."
+echo "  pld-raw (no guard) on repetitive prompts produces 'degraded' with high TG."

scripts/bench_tg.sh

@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+# bench_tg.sh — stable TG measurement: N runs × 200 tokens, drop cold-starts, report median.
+#
+# Usage:  ./scripts/bench_tg.sh [N_RUNS]      (default 5)
+#         LCA_WARMUP=3 ./scripts/bench_tg.sh  (with warmup enabled)
+set -u
+cd "$(dirname "$0")/.."
+
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+BIN="./build/qwen3-moe-aclnn"
+N_RUNS="${1:-5}"
+N_PREDICT="${N_PREDICT:-200}"
+PROMPT="The history of artificial intelligence spans several decades and"
+VOCAB="tokenizer_data/vocab.bin"
+
+echo "bench_tg: $N_RUNS runs × $N_PREDICT tokens (LCA_WARMUP=${LCA_WARMUP:-0})"
+tgs=()
+for r in $(seq 1 $N_RUNS); do
+    local_out=$(./scripts/tp_launch.sh 16 $BIN --model-dir "$MODEL" \
+                --prompt "$PROMPT" --n-predict $N_PREDICT \
+                --vocab "$VOCAB" --seed 0 2>&1 | grep "decode :" | awk '{print $(NF-2)}')
+    printf "  run %d: %s t/s\n" "$r" "$local_out"
+    tgs+=("${local_out:-0}")
+done
+
+echo ""
+echo "====== Summary ======"
+sorted=($(printf '%s\n' "${tgs[@]}" | sort -n))
+n=${#sorted[@]}
+mid=$((n / 2))
+median="${sorted[$mid]}"
+min="${sorted[0]}"
+max="${sorted[-1]}"
+mean=$(printf '%s\n' "${tgs[@]}" | awk '{s+=$1} END {printf "%.2f", s/NR}')
+
+echo "  all     : ${tgs[*]}"
+echo "  min     : $min t/s"
+echo "  median  : $median t/s"
+echo "  mean    : $mean t/s"
+echo "  max     : $max t/s"

scripts/export_vocab.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""Export Qwen3 tokenizer vocab to a simple binary format.
+
+Format (little-endian):
+  u32 num_tokens
+  for each id in [0, num_tokens):
+    u32 byte_length
+    u8[byte_length] utf8_bytes
+
+Also emits special_tokens.txt with id + content pairs for reference.
+"""
+import json, sys, struct, os
+
+model_dir = sys.argv[1] if len(sys.argv) > 1 else '/path/to/Qwen3-235B-A22B-Instruct-2507-BF16'
+out_dir   = sys.argv[2] if len(sys.argv) > 2 else 'tokenizer_data'
+os.makedirs(out_dir, exist_ok=True)
+
+with open(os.path.join(model_dir, 'tokenizer.json'), 'r') as f:
+    tok = json.load(f)
+
+# Byte-level decoder map: HF Qwen uses byte-level BPE like GPT-2
+# Each non-ASCII vocab token is a mapping of U+0100..U+017F etc back to raw bytes.
+# For decode we just need the reverse map from printable chars to raw bytes.
+def build_byte_decoder():
+    bs = list(range(ord('!'), ord('~')+1)) + list(range(ord('¡'), ord('¬')+1)) + list(range(ord('®'), ord('ÿ')+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    return {chr(c): bytes([b]) for b, c in zip(bs, cs)}
+
+byte_decoder = build_byte_decoder()
+
+# Merge vocab + added_tokens into id -> utf8_bytes lookup
+vocab = tok['model']['vocab']  # {token_str: id}
+added = tok.get('added_tokens', [])  # list of {id, content, ...}
+
+id_to_bytes = {}
+for token, tid in vocab.items():
+    # Decode byte-level encoding back to raw utf8 bytes
+    raw = b''
+    for ch in token:
+        if ch in byte_decoder:
+            raw += byte_decoder[ch]
+        else:
+            raw += ch.encode('utf-8')
+    id_to_bytes[int(tid)] = raw
+
+for a in added:
+    # Special tokens stored as raw utf8
+    id_to_bytes[int(a['id'])] = a['content'].encode('utf-8')
+
+max_id = max(id_to_bytes.keys())
+num = max_id + 1
+print(f"max_id = {max_id}, num_tokens = {num}")
+print(f"num_special_tokens = {len(added)}")
+
+# Write vocab.bin
+vocab_path = os.path.join(out_dir, 'vocab.bin')
+with open(vocab_path, 'wb') as f:
+    f.write(struct.pack('<I', num))
+    for i in range(num):
+        b = id_to_bytes.get(i, b'')
+        f.write(struct.pack('<I', len(b)))
+        f.write(b)
+print(f"Wrote {vocab_path} ({os.path.getsize(vocab_path)} bytes)")
+
+# Write special tokens
+with open(os.path.join(out_dir, 'special_tokens.txt'), 'w') as f:
+    for a in added:
+        f.write(f"{a['id']}\t{a['content']}\n")
+print(f"Wrote special_tokens.txt")
+
+# Verify via a known prompt
+from transformers import AutoTokenizer
+atok = AutoTokenizer.from_pretrained(model_dir)
+test = "The capital of France is"
+ids = atok.encode(test)
+print(f"\nTest encode '{test}' -> {ids}")
+decoded = ''.join(id_to_bytes.get(i, b'?').decode('utf-8', errors='replace') for i in ids)
+print(f"Our decode:   '{decoded}'")
+print(f"HF  decode:   '{atok.decode(ids)}'")

scripts/gen_attention_reference.py

@@ -0,0 +1,179 @@
+#!/usr/bin/env python3
+"""Generate a single-layer attention forward reference for Qwen3-235B layer 0.
+
+Input: token ids (representing "The capital of France is")
+Output: hidden_states after layer 0 attention (residual already added).
+Also dumps all intermediate tensors for step-wise debugging.
+"""
+import os, json, math, struct
+import torch
+import torch_npu
+from safetensors.torch import load_file
+
+torch.npu.set_device(0)
+torch.set_grad_enabled(False)
+
+MODEL_DIR = '/path/to/Qwen3-235B-A22B-Instruct-2507-BF16'
+OUT_DIR   = 'tests/attn_data'
+os.makedirs(OUT_DIR, exist_ok=True)
+
+cfg = json.load(open(os.path.join(MODEL_DIR, 'config.json')))
+D   = cfg['hidden_size']            # 4096
+Hq  = cfg['num_attention_heads']    # 64
+Hkv = cfg['num_key_value_heads']    # 4
+Dh  = cfg['head_dim']               # 128
+Q_DIM  = Hq * Dh                    # 8192
+KV_DIM = Hkv * Dh                   # 512
+eps    = cfg['rms_norm_eps']
+theta  = cfg['rope_theta']          # 5e6 for Qwen3-235B
+
+# ---- Find which safetensors shard contains layer 0 attention + input_layernorm ----
+idx = json.load(open(os.path.join(MODEL_DIR, 'model.safetensors.index.json')))
+wm  = idx['weight_map']
+
+needed = [
+    'model.embed_tokens.weight',
+    'model.layers.0.input_layernorm.weight',
+    'model.layers.0.self_attn.q_proj.weight',
+    'model.layers.0.self_attn.k_proj.weight',
+    'model.layers.0.self_attn.v_proj.weight',
+    'model.layers.0.self_attn.o_proj.weight',
+    'model.layers.0.self_attn.q_norm.weight',
+    'model.layers.0.self_attn.k_norm.weight',
+]
+shards = sorted({wm[n] for n in needed})
+print("Need to load shards:", shards)
+
+weights = {}
+for sh in shards:
+    t = load_file(os.path.join(MODEL_DIR, sh))
+    for n in needed:
+        if n in t:
+            weights[n] = t[n].to('npu')
+print("loaded:", list(weights.keys()))
+
+# ---- Forward ----
+# Input tokens (from tokenizer: "The capital of France is")
+token_ids = torch.tensor([785, 6722, 315, 9625, 374], dtype=torch.long).npu()
+S = token_ids.shape[0]
+print(f"S = {S}")
+
+# Embedding lookup
+x = weights['model.embed_tokens.weight'][token_ids]   # [S, D]
+x = x.unsqueeze(0)                                     # [1, S, D]
+print("embed x:", x.shape, x.dtype)
+
+# Residual
+residual = x
+
+# Input layernorm (RMSNorm)
+ln = weights['model.layers.0.input_layernorm.weight']
+xn, _ = torch_npu.npu_rms_norm(x, ln, epsilon=eps)
+print("after_input_norm xn:", xn.shape)
+
+# Q/K/V projections
+Wq = weights['model.layers.0.self_attn.q_proj.weight']
+Wk = weights['model.layers.0.self_attn.k_proj.weight']
+Wv = weights['model.layers.0.self_attn.v_proj.weight']
+q = torch.matmul(xn, Wq.t())   # [1, S, Q_DIM]
+k = torch.matmul(xn, Wk.t())   # [1, S, KV_DIM]
+v = torch.matmul(xn, Wv.t())
+
+# Reshape to heads
+q = q.view(1, S, Hq,  Dh)
+k = k.view(1, S, Hkv, Dh)
+v = v.view(1, S, Hkv, Dh)
+
+# Per-head RMSNorm on head_dim (Qwen3 specific)
+qn_w = weights['model.layers.0.self_attn.q_norm.weight']  # [Dh]
+kn_w = weights['model.layers.0.self_attn.k_norm.weight']
+q_normed, _ = torch_npu.npu_rms_norm(q, qn_w, epsilon=eps)
+k_normed, _ = torch_npu.npu_rms_norm(k, kn_w, epsilon=eps)
+
+# RoPE: compute cos/sin for positions [0, S)
+position_ids = torch.arange(S, device='npu').unsqueeze(0)  # [1, S]
+inv_freq = 1.0 / (theta ** (torch.arange(0, Dh, 2, dtype=torch.float32).npu() / Dh))
+freqs = position_ids.float().unsqueeze(-1) * inv_freq.unsqueeze(0).unsqueeze(0)  # [1, S, Dh/2]
+# Concat (half, half) to get [1, S, Dh]
+emb = torch.cat([freqs, freqs], dim=-1)
+cos = emb.cos().to(torch.bfloat16)                         # [1, S, Dh]
+sin = emb.sin().to(torch.bfloat16)
+
+# Apply RoPE — npu_apply_rotary_pos_emb expects BSND layout
+# cos/sin shape: [1, S, 1, Dh] for broadcast over heads
+cos_b = cos.unsqueeze(2)
+sin_b = sin.unsqueeze(2)
+q_roped, k_roped = torch_npu.npu_apply_rotary_pos_emb(q_normed, k_normed, cos_b, sin_b)
+
+# Flatten for FIAS (BSH layout)
+q_bsh = q_roped.reshape(1, S, Q_DIM)
+k_bsh = k_roped.reshape(1, S, KV_DIM)
+v_bsh = v.reshape(1, S, KV_DIM)
+
+# FIAS with causal mask for prefill
+scale = 1.0 / math.sqrt(Dh)
+# sparse_mode=3 requires fixed 2048×2048 mask
+MASK_SIZE = 2048
+mask = torch.triu(torch.ones(MASK_SIZE, MASK_SIZE, dtype=torch.bool, device='npu'), diagonal=1)
+mask = mask.view(1, 1, MASK_SIZE, MASK_SIZE)
+attn_out, _ = torch_npu.npu_fused_infer_attention_score(
+    q_bsh, k_bsh, v_bsh,
+    num_heads=Hq,
+    num_key_value_heads=Hkv,
+    scale=scale,
+    input_layout="BSH",
+    sparse_mode=3,
+    atten_mask=mask,
+    actual_seq_lengths=[S],
+    actual_seq_lengths_kv=[S],
+)
+print("attn_out:", attn_out.shape)  # [1, S, Q_DIM]
+
+# Output projection
+Wo = weights['model.layers.0.self_attn.o_proj.weight']
+o = torch.matmul(attn_out, Wo.t())   # [1, S, D]
+
+# Residual add
+out = residual + o
+print("out:", out.shape, out[0, 0, :4].float().tolist())
+
+# ---- Dump ----
+def dump(name, t):
+    p = os.path.join(OUT_DIR, name + '.bin')
+    a = t.contiguous().cpu().view(torch.int16).numpy().astype('int16')
+    open(p, 'wb').write(a.tobytes())
+
+# Save token_ids
+with open(os.path.join(OUT_DIR, 'token_ids.bin'), 'wb') as f:
+    f.write(struct.pack('<i', S))
+    for tid in token_ids.cpu().tolist():
+        f.write(struct.pack('<i', tid))
+
+# Save inputs
+dump('x_input',    x)                                          # embed result
+dump('x_normed',   xn)
+dump('q_normed',   q_normed)
+dump('k_normed',   k_normed)
+dump('q_roped',    q_roped)
+dump('k_roped',    k_roped)
+dump('cos',        cos)
+dump('sin',        sin)
+dump('attn_out',   attn_out)
+dump('final_out',  out)
+# Save weights used (dtype=BF16)
+for name, path_name in [
+    ('model.layers.0.input_layernorm.weight', 'w_input_norm'),
+    ('model.layers.0.self_attn.q_proj.weight', 'w_q_proj'),
+    ('model.layers.0.self_attn.k_proj.weight', 'w_k_proj'),
+    ('model.layers.0.self_attn.v_proj.weight', 'w_v_proj'),
+    ('model.layers.0.self_attn.o_proj.weight', 'w_o_proj'),
+    ('model.layers.0.self_attn.q_norm.weight', 'w_q_norm'),
+    ('model.layers.0.self_attn.k_norm.weight', 'w_k_norm'),
+]:
+    dump(path_name, weights[name])
+
+with open(os.path.join(OUT_DIR, 'shape.txt'), 'w') as f:
+    f.write(f"S={S}\nD={D}\nHq={Hq}\nHkv={Hkv}\nDh={Dh}\nQ_DIM={Q_DIM}\nKV_DIM={KV_DIM}\neps={eps}\ntheta={theta}\n")
+
+print("\nAll dumps in:", OUT_DIR)
+print("Final output first 4:", out[0, 0, :4].float().cpu().tolist())

scripts/gen_gmm_reference.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Generate a GMM reference case using torch_npu.npu_grouped_matmul.
+
+Dumps: x, w (unpermuted ggml-like layout), group_list, and y_ref to binary files
+for the C++ POC to load and validate against.
+"""
+import os
+import sys
+import numpy as np
+import struct
+
+# Enable torch_npu
+os.environ.setdefault('LD_LIBRARY_PATH', '')
+import torch
+import torch_npu
+
+torch.npu.set_device(0)
+torch.manual_seed(42)
+
+# Toy Qwen3-like MoE shape: D=64 hidden, I=32 intermediate, E=8 experts, N*K=16 expanded tokens
+# (small enough to eyeball; large enough to catch layout bugs)
+D, I, E, TOTAL = 64, 32, 8, 16
+
+# Input x: [TOTAL, D] BF16 — expanded routed tokens
+x = torch.randn(TOTAL, D, dtype=torch.bfloat16).npu()
+
+# Weight w: per-expert [I, D] BF16 — gate/up has this shape in HF
+# We will stack into [E, I, D] and also provide [E, D, I] permuted for comparison
+w_per_expert = [torch.randn(I, D, dtype=torch.bfloat16).npu() for _ in range(E)]
+w_stacked_IDL = torch.stack(w_per_expert, dim=0)  # [E, I, D]
+
+# group_list: counts of tokens per expert, sum = TOTAL
+group_list = torch.tensor([3, 2, 1, 2, 1, 3, 2, 2], dtype=torch.int64).npu()
+assert group_list.sum().item() == TOTAL
+
+# Reference: use torch_npu.npu_grouped_matmul
+# Per cann-recipes: weight needs to be in [E, D, I] for matmul y = x @ w (y shape [total, I])
+# i.e. per-expert w is transposed from HF's [I, D] to [D, I]
+w_transposed = w_stacked_IDL.transpose(1, 2).contiguous()  # [E, D, I]
+
+# Call GMM: y = x @ w, result [TOTAL, I]
+y_ref = torch_npu.npu_grouped_matmul(
+    [x],                                    # x list
+    [w_transposed],                         # weight list (transposed)
+    group_list=group_list,
+    group_type=0,
+    group_list_type=1,                      # counts
+    split_item=3                            # single-in single-out
+)[0]  # unwrap tensor list
+
+print("x shape:", x.shape, x.dtype)
+print("w_stacked_IDL shape:", w_stacked_IDL.shape, w_stacked_IDL.dtype)
+print("w_transposed shape:", w_transposed.shape)
+print("group_list:", group_list.cpu().tolist())
+print("y_ref shape:", y_ref.shape)
+print("y_ref[0, 0:4]:", y_ref[0, 0:4].cpu().float().tolist())
+
+# Save binary dumps
+out_dir = 'tests/poc_data'
+os.makedirs(out_dir, exist_ok=True)
+
+def dump_bf16(name, tensor):
+    path = os.path.join(out_dir, name + '.bin')
+    arr = tensor.contiguous().cpu().view(torch.int16).numpy().astype('int16')
+    with open(path, 'wb') as f:
+        f.write(arr.tobytes())
+    print(f"  wrote {name}.bin: {arr.shape} int16 = BF16 raw, {arr.nbytes} bytes")
+
+def dump_int64(name, tensor):
+    path = os.path.join(out_dir, name + '.bin')
+    arr = tensor.contiguous().cpu().numpy().astype('int64')
+    with open(path, 'wb') as f:
+        f.write(arr.tobytes())
+    print(f"  wrote {name}.bin: {arr.shape} int64, {arr.nbytes} bytes")
+
+# HF-style weight layout (ggml stores similar): [E, I, D] = what C++ gets from safetensors after stack
+dump_bf16('x', x)
+dump_bf16('w_hf_EID', w_stacked_IDL)          # C++ input weight (HF layout)
+dump_bf16('w_ref_EDI', w_transposed)          # Already-permuted reference (for debug)
+dump_int64('group_list', group_list)
+dump_bf16('y_ref', y_ref)
+
+# Also dump shapes header
+with open(os.path.join(out_dir, 'shapes.txt'), 'w') as f:
+    f.write(f"D={D}\nI={I}\nE={E}\nTOTAL={TOTAL}\n")
+
+print("\nAll dumps in:", out_dir)
+print("\nTo validate: C++ loads w_hf_EID, permutes [0,2,1] to [E,D,I], NZ-casts, calls GMMV4, "
+      "compares output to y_ref.")

scripts/gen_mm_reference.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+"""Generate a linear (y = x @ W.T) reference for a realistic Qwen3 attention shape."""
+import os, struct, torch, torch_npu
+torch.npu.set_device(0)
+torch.manual_seed(7)
+
+N, D, OUT = 5, 4096, 8192  # prompt len, hidden, q_dim
+
+x = torch.randn(N, D, dtype=torch.bfloat16).npu()
+W = torch.randn(OUT, D, dtype=torch.bfloat16).npu()  # HF layout [out, in]
+# y = x @ W.T, shape [N, OUT]
+y_ref = torch.matmul(x, W.t())
+
+out_dir = 'tests/mm_data'
+os.makedirs(out_dir, exist_ok=True)
+def dump(name, t):
+    p = os.path.join(out_dir, name + '.bin')
+    a = t.contiguous().cpu().view(torch.int16).numpy().astype('int16')
+    open(p, 'wb').write(a.tobytes())
+dump('x', x); dump('W', W); dump('y_ref', y_ref)
+with open(os.path.join(out_dir, 'shape.txt'), 'w') as f:
+    f.write(f"N={N}\nD={D}\nOUT={OUT}\n")
+print(f"N={N} D={D} OUT={OUT}, y_ref[0, :4] = {y_ref[0, :4].float().cpu().tolist()}")

scripts/gen_moe_reference.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Generate MoE layer forward reference for Qwen3-235B layer 0.
+
+Input: hidden_states from attention output (use attn_data/final_out.bin as input — realistic).
+Output: hidden_states after MoE + residual.
+"""
+import os, json, math, torch, torch_npu
+from safetensors.torch import load_file
+
+torch.npu.set_device(0)
+torch.set_grad_enabled(False)
+
+MODEL_DIR = '/path/to/Qwen3-235B-A22B-Instruct-2507-BF16'
+OUT_DIR   = 'tests/moe_data'
+os.makedirs(OUT_DIR, exist_ok=True)
+
+cfg = json.load(open(os.path.join(MODEL_DIR, 'config.json')))
+D    = cfg['hidden_size']                  # 4096
+I    = cfg['moe_intermediate_size']        # 1536
+E    = cfg['num_experts']                  # 128
+TK   = cfg['num_experts_per_tok']          # 8
+eps  = cfg['rms_norm_eps']
+norm_topk = cfg.get('norm_topk_prob', True)
+
+# Use attention output as input (more realistic than random)
+attn_out_raw = open('tests/attn_data/final_out.bin', 'rb').read()
+S = 5
+x_in = torch.frombuffer(bytearray(attn_out_raw), dtype=torch.int16).view(1, S, D).view(torch.bfloat16).npu()
+print(f"x_in: {x_in.shape}")
+
+# Load required weights for layer 0
+idx = json.load(open(os.path.join(MODEL_DIR, 'model.safetensors.index.json')))
+wm = idx['weight_map']
+
+needed = [f'model.layers.0.post_attention_layernorm.weight',
+          f'model.layers.0.mlp.gate.weight']
+for e in range(E):
+    for p in ['gate_proj', 'up_proj', 'down_proj']:
+        needed.append(f'model.layers.0.mlp.experts.{e}.{p}.weight')
+
+shards = sorted({wm[n] for n in needed})
+weights = {}
+for sh in shards:
+    t = load_file(os.path.join(MODEL_DIR, sh))
+    for n in needed:
+        if n in t:
+            weights[n] = t[n].to('npu')
+print("loaded %d tensors from %d shards" % (len(weights), len(shards)))
+
+# Residual = input
+residual = x_in
+
+# Post-attention RmsNorm
+xn, _ = torch_npu.npu_rms_norm(x_in, weights['model.layers.0.post_attention_layernorm.weight'], epsilon=eps)
+xn_flat = xn.view(S, D)  # flatten batch
+
+# Router: logits [S, E]
+W_router = weights['model.layers.0.mlp.gate.weight']                # [E, D]
+logits = xn_flat @ W_router.t()                                      # [S, E]
+
+# Top-k softmax
+topk_logits, topk_idx = logits.topk(TK, dim=-1)                      # both [S, TK]
+topk_weights = torch.softmax(topk_logits.float(), dim=-1)             # [S, TK] F32
+if norm_topk:
+    topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20)
+topk_weights = topk_weights.to(torch.bfloat16)
+topk_idx = topk_idx.to(torch.int32)
+
+print(f"topk_idx[0]: {topk_idx[0].cpu().tolist()}")
+print(f"topk_weights[0]: {topk_weights[0].cpu().float().tolist()}")
+
+# MoE forward — loop over tokens (simple reference, not optimized)
+out_flat = torch.zeros(S, D, dtype=torch.bfloat16, device='npu')
+for s in range(S):
+    token = xn_flat[s]                                                # [D]
+    acc = torch.zeros(D, dtype=torch.bfloat16, device='npu')
+    for k in range(TK):
+        e = int(topk_idx[s, k].item())
+        w = topk_weights[s, k]
+        Wg = weights[f'model.layers.0.mlp.experts.{e}.gate_proj.weight']   # [I, D]
+        Wu = weights[f'model.layers.0.mlp.experts.{e}.up_proj.weight']     # [I, D]
+        Wd = weights[f'model.layers.0.mlp.experts.{e}.down_proj.weight']   # [D, I]
+        gate = token @ Wg.t()      # [I]
+        up = token @ Wu.t()
+        act = torch.nn.functional.silu(gate) * up
+        down = act @ Wd.t()        # [D]
+        acc = acc + w * down
+    out_flat[s] = acc
+
+moe_out = out_flat.view(1, S, D)
+final_out = residual + moe_out
+print(f"final_out[0,0,:4] = {final_out[0,0,:4].float().cpu().tolist()}")
+
+# Dump
+def dump(name, t):
+    p = os.path.join(OUT_DIR, name + '.bin')
+    a = t.contiguous().cpu().view(torch.int16).numpy().astype('int16')
+    open(p, 'wb').write(a.tobytes())
+
+dump('x_in', x_in)
+dump('final_out', final_out)
+dump('moe_out', moe_out)
+dump('router', W_router)
+dump('xn', xn)
+dump('topk_w', topk_weights)      # [S, TK] BF16 (normalized)
+dump('out_flat', out_flat)        # [S, D] BF16 — moe contrib before residual
+
+# expert_idx as int32 dump (raw bytes)
+topk_idx_bytes = topk_idx.contiguous().cpu().numpy().astype('int32').tobytes()
+open(os.path.join(OUT_DIR, 'topk_idx.bin'), 'wb').write(topk_idx_bytes)
+
+with open(os.path.join(OUT_DIR, 'shape.txt'), 'w') as f:
+    f.write(f"S={S}\nD={D}\nI={I}\nE={E}\nTK={TK}\n")
+
+print(f"\nDumps in {OUT_DIR}")

scripts/gen_rms_norm_reference.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+"""Generate a RmsNorm reference using PyTorch."""
+import os, struct
+import numpy as np
+import torch
+import torch_npu
+
+torch.npu.set_device(0)
+torch.manual_seed(123)
+
+N, D = 5, 4096  # 5 tokens, Qwen3 hidden_size
+eps = 1e-6
+
+x = torch.randn(N, D, dtype=torch.bfloat16).npu()
+gamma = torch.randn(D, dtype=torch.bfloat16).npu() * 0.1 + 1.0
+
+# Use torch_npu's npu_rms_norm if available, else do it manually
+y_ref, _ = torch_npu.npu_rms_norm(x, gamma, epsilon=eps)
+
+out_dir = 'tests/rms_norm_data'
+os.makedirs(out_dir, exist_ok=True)
+
+def dump_bf16(name, t):
+    path = os.path.join(out_dir, name + '.bin')
+    a = t.contiguous().cpu().view(torch.int16).numpy().astype('int16')
+    with open(path, 'wb') as f:
+        f.write(a.tobytes())
+    return path
+
+dump_bf16('x',     x)
+dump_bf16('gamma', gamma)
+dump_bf16('y_ref', y_ref)
+
+with open(os.path.join(out_dir, 'shape.txt'), 'w') as f:
+    f.write(f"N={N}\nD={D}\neps={eps}\n")
+
+print(f"x shape: {x.shape}, gamma: {gamma.shape}, y_ref: {y_ref.shape}")
+print("y_ref[0, :8]:", y_ref[0, :8].float().cpu().tolist())
+print("saved in", out_dir)

scripts/regen_rope_reference.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+"""Re-generate RoPE reference using explicit HF formula (not torch_npu.npu_apply_rotary_pos_emb)."""
+import os, math, torch, torch_npu
+torch.npu.set_device(0)
+torch.manual_seed(42)
+
+S = 5; Hq = 64; Hkv = 4; Dh = 128
+theta = 5e6
+data = 'tests/attn_data'
+
+def load_bf16(name, shape):
+    raw = open(os.path.join(data, name + '.bin'), 'rb').read()
+    a = torch.frombuffer(bytearray(raw), dtype=torch.int16).view(*shape).view(torch.bfloat16)
+    return a.npu()
+
+q = load_bf16('q_normed', [1, S, Hq, Dh])
+k = load_bf16('k_normed', [1, S, Hkv, Dh])
+
+# Compute cos/sin identical to HF (rope_theta=5e6, 0..S positions)
+inv_freq = 1.0 / (theta ** (torch.arange(0, Dh, 2, dtype=torch.float32).npu() / Dh))
+pos = torch.arange(S, device='npu').float().unsqueeze(-1)
+freqs = pos * inv_freq
+emb = torch.cat([freqs, freqs], dim=-1)    # [S, Dh]
+cos = emb.cos().to(torch.bfloat16)          # [S, Dh]
+sin = emb.sin().to(torch.bfloat16)
+
+# HF (Qwen3) style RoPE: q_rot = q * cos + rotate_half(q) * sin
+def rotate_half(x):
+    h = x.shape[-1] // 2
+    x1 = x[..., :h]
+    x2 = x[..., h:]
+    return torch.cat([-x2, x1], dim=-1)
+
+# Broadcast cos/sin from [S, Dh] to [1, S, 1, Dh]
+cos_b = cos.unsqueeze(0).unsqueeze(2)
+sin_b = sin.unsqueeze(0).unsqueeze(2)
+
+q_roped_hf = q * cos_b + rotate_half(q) * sin_b
+k_roped_hf = k * cos_b + rotate_half(k) * sin_b
+
+print("HF-style q_roped[0,0,:4]:", q_roped_hf[0,0,0,:4].float().cpu().tolist())
+print("cos[0,:4]:", cos[0,:4].float().cpu().tolist())
+print("sin[0,:4]:", sin[0,:4].float().cpu().tolist())
+print("cos[1,:4]:", cos[1,:4].float().cpu().tolist())
+
+# Compare with existing q_roped (from torch_npu.npu_apply_rotary_pos_emb)
+old_q_roped = load_bf16('q_roped', [1, S, Hq, Dh])
+diff = (q_roped_hf - old_q_roped).float().abs().max().item()
+print(f"\nDiff between HF formula and npu_apply: max={diff:.4f}")
+
+# Save HF version as ground truth
+def dump(name, t):
+    p = os.path.join(data, name + '.bin')
+    a = t.contiguous().cpu().view(torch.int16).numpy().astype('int16')
+    open(p, 'wb').write(a.tobytes())
+dump('q_roped', q_roped_hf)
+dump('k_roped', k_roped_hf)
+# Overwrite cos, sin to [1, S, Dh] layout
+dump('cos', cos.unsqueeze(0))  # [1, S, Dh]
+dump('sin', sin.unsqueeze(0))
+
+print("\nOverwrote q_roped, k_roped, cos, sin with HF-formula ground truth.")

scripts/tp_launch.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+# tp_launch.sh — launcher for TP>1 multi-process qwen3-moe-aclnn.
+#
+# Usage: ./tp_launch.sh <tp_size> <bin> [args...]
+#   e.g. ./tp_launch.sh 16 ./build/qwen3-moe-aclnn --model-dir ... --prompt "..." --n-predict 20
+#
+# Each rank runs as a separate process with:
+#   ASCEND_RT_VISIBLE_DEVICES=<rank>
+#   TP_RANK=<rank>  TP_SIZE=<tp_size>
+#   HCCL_WHITELIST_DISABLE=1
+# rank 0 creates /tmp/hccl_root_info.bin; other ranks wait for it.
+set -euo pipefail
+
+TP_SIZE="${1:?tp_size required}"; shift
+BIN="${1:?binary required}"; shift
+
+# Clean any stale HCCL coordination file
+rm -f /tmp/hccl_root_info.bin
+
+export HCCL_WHITELIST_DISABLE=1
+# Benchmark-tuned defaults (bench_hccl_adv.sh 2026-04-21):
+#   ring:200 + OP_EXPANSION_MODE=AIV + OP_BASE_FFTS_MODE_ENABLE=1  →  ~18.8 t/s median
+#   vs baseline (auto) ~12 t/s.  +54% from HCCL env knobs alone.
+export HCCL_ALGO="${HCCL_ALGO:-level0:ring}"
+export HCCL_BUFFSIZE="${HCCL_BUFFSIZE:-200}"
+export HCCL_OP_EXPANSION_MODE="${HCCL_OP_EXPANSION_MODE:-AIV}"
+export HCCL_OP_BASE_FFTS_MODE_ENABLE="${HCCL_OP_BASE_FFTS_MODE_ENABLE:-1}"
+# TASK_QUEUE_ENABLE=2: aggressive async task queueing (marginal gain on top of AIV+FFTS)
+export TASK_QUEUE_ENABLE="${TASK_QUEUE_ENABLE:-2}"
+
+# Launch ranks 1..N-1 in background with stdin/stdout redirected to /dev/null / logfile.
+# Launch rank 0 LAST in foreground, inheriting the terminal stdin/stdout — so --interactive works.
+pids=()
+for rank in $(seq 1 $((TP_SIZE - 1))); do
+    logfile="/tmp/tp_rank_${rank}.log"
+    env ASCEND_RT_VISIBLE_DEVICES=${rank} \
+        TP_RANK=${rank} \
+        TP_SIZE=${TP_SIZE} \
+        "${BIN}" "$@" < /dev/null > "${logfile}" 2>&1 &
+    pids+=($!)
+    echo "[tp_launch] rank ${rank} pid=$! log=${logfile}"
+done
+
+# Give ranks 1..N-1 a moment to reach HcclCommInitRootInfo's file-wait before rank 0 writes it.
+sleep 1
+
+# Rank 0 in foreground — terminal stdin/stdout passthrough for REPL.
+env ASCEND_RT_VISIBLE_DEVICES=0 \
+    TP_RANK=0 \
+    TP_SIZE=${TP_SIZE} \
+    "${BIN}" "$@"
+ec=$?
+
+# Wait for background ranks to finish (rank 0 exit signals end-of-work, but they may take a bit).
+for i in "${!pids[@]}"; do
+    wait "${pids[$i]}" || true
+done
+exit $ec

src/device_weights.cpp

@@ -0,0 +1,221 @@
+#include "device_weights.h"
+#include "aclnn_ops.h"
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+bool DeviceWeightsLoader::load_tensor_full_(const std::string& name, DeviceBuffer& buf) {
+    const auto* m = st_.get(name);
+    if (!m) { fprintf(stderr, "load_tensor_full_: missing %s\n", name.c_str()); return false; }
+    const void* host = st_.data_ptr(*m);
+    if (!host) { fprintf(stderr, "load_tensor_full_: null host ptr %s\n", name.c_str()); return false; }
+    buf.alloc(m->nbytes);
+    ACL_CHECK(aclrtMemcpy(buf.get(), m->nbytes, host, m->nbytes, ACL_MEMCPY_HOST_TO_DEVICE));
+    return true;
+}
+
+bool DeviceWeightsLoader::load_tensor_row_slice_(const std::string& name,
+                                                 int64_t row_lo, int64_t row_hi,
+                                                 DeviceBuffer& buf) {
+    const auto* m = st_.get(name);
+    if (!m) { fprintf(stderr, "load_tensor_row_slice_: missing %s\n", name.c_str()); return false; }
+    if (m->shape.empty()) { fprintf(stderr, "%s: empty shape\n", name.c_str()); return false; }
+    int64_t D0 = m->shape[0];
+    if (row_hi > D0 || row_lo < 0 || row_hi <= row_lo) {
+        fprintf(stderr, "load_tensor_row_slice_: %s bad range [%ld,%ld) vs D0=%ld\n",
+                name.c_str(), row_lo, row_hi, D0);
+        return false;
+    }
+    size_t elem = sdtype_size(m->dtype);
+    size_t inner = 1;
+    for (size_t i = 1; i < m->shape.size(); i++) inner *= m->shape[i];
+    size_t row_bytes = inner * elem;
+    size_t slice_bytes = (row_hi - row_lo) * row_bytes;
+
+    const auto* host = (const char*)st_.data_ptr(*m);
+    buf.alloc(slice_bytes);
+    ACL_CHECK(aclrtMemcpy(buf.get(), slice_bytes,
+                          host + row_lo * row_bytes, slice_bytes,
+                          ACL_MEMCPY_HOST_TO_DEVICE));
+    return true;
+}
+
+bool DeviceWeightsLoader::load_tensor_col_slice_(const std::string& name,
+                                                 int64_t col_lo, int64_t col_hi,
+                                                 DeviceBuffer& buf) {
+    const auto* m = st_.get(name);
+    if (!m || m->shape.size() < 2) {
+        fprintf(stderr, "load_tensor_col_slice_: bad shape %s\n", name.c_str()); return false;
+    }
+    int64_t D0 = m->shape[0];
+    int64_t D1 = m->shape[1];
+    if (col_hi > D1 || col_lo < 0 || col_hi <= col_lo) {
+        fprintf(stderr, "load_tensor_col_slice_: bad range %ld-%ld D1=%ld\n",
+                col_lo, col_hi, D1); return false;
+    }
+    size_t elem = sdtype_size(m->dtype);
+    int64_t new_cols = col_hi - col_lo;
+    size_t slice_bytes = D0 * new_cols * elem;
+    buf.alloc(slice_bytes);
+
+    // Need to copy row-by-row since source has stride D1 but dest has stride new_cols.
+    const auto* host = (const char*)st_.data_ptr(*m);
+    std::vector<char> staging(slice_bytes);
+    size_t src_row = D1 * elem;
+    size_t dst_row = new_cols * elem;
+    size_t col_off = col_lo * elem;
+    for (int64_t r = 0; r < D0; r++) {
+        std::memcpy(staging.data() + r * dst_row, host + r * src_row + col_off, dst_row);
+    }
+    ACL_CHECK(aclrtMemcpy(buf.get(), slice_bytes, staging.data(), slice_bytes,
+                          ACL_MEMCPY_HOST_TO_DEVICE));
+    return true;
+}
+
+bool DeviceWeightsLoader::load_shared(SharedWeights& out) {
+    if (!load_tensor_full_("model.embed_tokens.weight", out.embed_tokens)) return false;
+    if (!load_tensor_full_("lm_head.weight",            out.lm_head))       return false;
+    if (!load_tensor_full_("model.norm.weight",         out.final_norm))    return false;
+    return true;
+}
+
+bool DeviceWeightsLoader::load_moe(int L, aclrtStream stream, LayerMoEWeights& out) {
+    const int64_t E = cfg_.num_experts;
+    const int64_t D = cfg_.hidden_size;
+    const int64_t I_full = cfg_.moe_intermediate_size;
+    const int64_t I_rank = cfg_.i_per_rank;
+    const size_t elem = 2;  // BF16
+
+    auto base = "model.layers." + std::to_string(L);
+
+    // 1. Router [E, D] — small, fully replicated
+    if (!load_tensor_full_(base + ".mlp.gate.weight", out.router)) return false;
+
+    // 2. MoE expert weights: need to stack 128 experts + TP slice + permute
+    // HF gate/up: each expert [I_full, D] → TP slice rows to [I_rank, D]
+    // HF down:    each expert [D, I_full] → TP slice cols to [D, I_rank]
+
+    auto load_and_stack = [&](const std::string& proj_name,
+                              bool is_down, DeviceBuffer& final_buf) -> bool {
+        // HF shape for gate/up: [I_full, D]; for down: [D, I_full]
+        // After TP slice: gate/up rows [I_rank, D]; down cols [D, I_rank]
+        // Stacked:
+        //   gate/up: [E, I_rank, D] → permute to [E, D, I_rank]
+        //   down:    [E, D, I_rank] → permute to [E, I_rank, D]
+
+        int64_t K_in, N_out;
+        bool row_slice;
+        if (!is_down) {
+            K_in = I_rank;            // HF first dim after row-slice
+            N_out = D;
+            row_slice = true;
+        } else {
+            K_in = D;
+            N_out = I_rank;
+            row_slice = false;        // col slice
+        }
+
+        // Stage: stacked HF-layout [E, K_in, N_out] on device (before permute)
+        size_t elem_stack = K_in * N_out * elem;
+        DeviceBuffer stacked_hf(E * elem_stack);
+
+        // For each expert, load + TP slice + memcpy to stacked_hf[e]
+        // We use the existing row_slice/col_slice helpers on a per-expert basis.
+        DeviceBuffer tmp;
+        for (int e = 0; e < E; e++) {
+            std::string name = base + ".mlp.experts." + std::to_string(e) + "." + proj_name + ".weight";
+            if (row_slice) {
+                int64_t lo = cfg_.tp_rank * I_rank;
+                int64_t hi = lo + I_rank;
+                if (!load_tensor_row_slice_(name, lo, hi, tmp)) return false;
+            } else {
+                int64_t lo = cfg_.tp_rank * I_rank;
+                int64_t hi = lo + I_rank;
+                if (!load_tensor_col_slice_(name, lo, hi, tmp)) return false;
+            }
+            if (tmp.size != elem_stack) {
+                fprintf(stderr, "load_moe: expert %d %s slice size %zu != expected %zu\n",
+                        e, name.c_str(), tmp.size, elem_stack);
+                return false;
+            }
+            // Synchronous D2D: tmp is about to be reallocated in the next iteration,
+            // so we cannot enqueue an async copy that would still reference it.
+            ACL_CHECK(aclrtMemcpy(
+                (char*)stacked_hf.get() + e * elem_stack, elem_stack,
+                tmp.get(), elem_stack,
+                ACL_MEMCPY_DEVICE_TO_DEVICE));
+        }
+
+        // Now permute stacked_hf [E, K_in, N_out] → final [E, N_out, K_in] row-major
+        // (swap last two dims)
+        final_buf.alloc(E * elem_stack);
+        const int64_t d0 = E, d1 = K_in, d2 = N_out;
+        // View stacked_hf with permuted strides pointing to same data:
+        // logical shape [E, N_out, K_in], strides [K_in*N_out, 1, N_out]
+        // (since physical is [E, K_in, N_out] row-major with strides [K_in*N_out, N_out, 1])
+        auto t_src = make_acl_tensor(stacked_hf.get(), ACL_BF16,
+                                     {d0, d2, d1},  // [E, N_out, K_in]
+                                     {d1 * d2, 1, d2});
+        auto t_dst = make_contig_tensor(final_buf.get(), ACL_BF16, {d0, d2, d1});
+        inplace_copy(stream, t_dst.get(), t_src.get());
+        // Must sync before stacked_hf goes out of scope — the inplace_copy is async and
+        // reads from stacked_hf's memory. If we return without syncing, DeviceBuffer's
+        // destructor frees stacked_hf while the permute kernel is still running, producing
+        // garbage in final_buf.
+        ACL_CHECK(aclrtSynchronizeStream(stream));
+        return true;
+    };
+
+    if (!load_and_stack("gate_proj", false, out.gate_exps)) return false;
+    if (!load_and_stack("up_proj",   false, out.up_exps))   return false;
+    if (!load_and_stack("down_proj", true,  out.down_exps)) return false;
+
+    return true;
+}
+
+bool DeviceWeightsLoader::load_attention(int L, LayerAttnWeights& out) {
+    auto base = "model.layers." + std::to_string(L);
+
+    if (!load_tensor_full_(base + ".input_layernorm.weight",          out.input_layernorm))          return false;
+    if (!load_tensor_full_(base + ".post_attention_layernorm.weight", out.post_attention_layernorm)) return false;
+    if (!load_tensor_full_(base + ".self_attn.q_norm.weight", out.q_norm)) return false;
+    if (!load_tensor_full_(base + ".self_attn.k_norm.weight", out.k_norm)) return false;
+
+    const int64_t head_dim = cfg_.head_dim;
+    const int64_t q_full   = cfg_.num_attention_heads * head_dim;    // 64 * 128 = 8192
+
+    // q_proj: [q_full, D], shard rows by head. Each rank gets n_heads_per_rank heads.
+    int64_t q_rows_per_rank = cfg_.n_heads_per_rank * head_dim;
+    int64_t q_row_lo = cfg_.tp_rank * q_rows_per_rank;
+    int64_t q_row_hi = q_row_lo + q_rows_per_rank;
+    if (!load_tensor_row_slice_(base + ".self_attn.q_proj.weight",
+                                 q_row_lo, q_row_hi, out.q_proj)) return false;
+
+    // k_proj, v_proj: HF shape [num_kv * head_dim, D].
+    //   Case A (tp <= n_kv): split rows across ranks, each rank gets n_kv/tp KV heads.
+    //   Case B (tp > n_kv):  each rank gets exactly ONE KV head; group of (tp/n_kv) ranks share it.
+    //       kv_head_idx = tp_rank / (tp_size / n_kv)
+    if (cfg_.tp_size <= cfg_.num_key_value_heads) {
+        int64_t kv_rows_per_rank = cfg_.n_kv_heads_per_rank * head_dim;
+        int64_t kv_row_lo = cfg_.tp_rank * kv_rows_per_rank;
+        int64_t kv_row_hi = kv_row_lo + kv_rows_per_rank;
+        if (!load_tensor_row_slice_(base + ".self_attn.k_proj.weight", kv_row_lo, kv_row_hi, out.k_proj)) return false;
+        if (!load_tensor_row_slice_(base + ".self_attn.v_proj.weight", kv_row_lo, kv_row_hi, out.v_proj)) return false;
+    } else {
+        // GQA replicated-group mode: 1 KV head per rank, selected by group.
+        int64_t ranks_per_kv = cfg_.tp_size / cfg_.num_key_value_heads;
+        int64_t kv_head_idx = cfg_.tp_rank / ranks_per_kv;
+        int64_t kv_row_lo = kv_head_idx * head_dim;
+        int64_t kv_row_hi = kv_row_lo + head_dim;
+        if (!load_tensor_row_slice_(base + ".self_attn.k_proj.weight", kv_row_lo, kv_row_hi, out.k_proj)) return false;
+        if (!load_tensor_row_slice_(base + ".self_attn.v_proj.weight", kv_row_lo, kv_row_hi, out.v_proj)) return false;
+    }
+
+    // o_proj: [D, q_full], row-parallel → shard cols (input dim) by head.
+    int64_t o_col_lo = q_row_lo;  // same slicing as q rows
+    int64_t o_col_hi = q_row_hi;
+    if (!load_tensor_col_slice_(base + ".self_attn.o_proj.weight",
+                                 o_col_lo, o_col_hi, out.o_proj)) return false;
+
+    return true;
+}

src/main_cli.cpp

@@ -0,0 +1,816 @@
+// main_cli.cpp — qwen3-moe-aclnn entry point.
+//
+// Usage:
+//   qwen3-moe-aclnn --model-dir <path> --prompt "<text>" --n-predict <N>
+//                   [--tp-size 1|16] [--vocab <path>] [--max-seq N] [--num-layers N]
+//                   [--chat] [--temperature 0.7] [--top-k 20] [--top-p 0.8] [--seed N]
+//                   [--no-stream]
+//
+// At TP>1 each rank is a separate process (env TP_RANK=<i>, TP_SIZE=<n>) launched by
+// scripts/tp_launch.sh. Only rank 0 prints text output.
+#include "runner.h"
+#include "tokenizer.h"
+
+// Escape hatch for HCCL broadcast from within CLI (defined in runner.cpp)
+HcclCtx* runner_hccl_ctx_shim(Runner& r);
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <random>
+#include <string>
+#include <vector>
+
+static float bf16_to_float(uint16_t x) {
+    uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f;
+}
+
+// Truncate a string to the last complete UTF-8 character boundary. If the last 1-3 bytes
+// form an incomplete multi-byte sequence (e.g., assistant response cut mid-codepoint at
+// n_predict limit), drop them so the JSON encoder downstream sees only valid UTF-8.
+static std::string utf8_trim_incomplete(const std::string& s) {
+    if (s.empty()) return s;
+    size_t n = s.size();
+    // Walk back up to 4 bytes looking for the start of a UTF-8 sequence.
+    for (size_t back = 0; back < 4 && back < n; back++) {
+        size_t i = n - 1 - back;
+        unsigned char c = (unsigned char)s[i];
+        if ((c & 0x80) == 0)        { return s; }                    // ASCII: already complete
+        if ((c & 0xC0) == 0x80)     { continue; }                    // continuation byte: keep going
+        // Start byte: 110xxxxx (2-byte), 1110xxxx (3-byte), 11110xxx (4-byte)
+        size_t need = 0;
+        if      ((c & 0xE0) == 0xC0) need = 2;
+        else if ((c & 0xF0) == 0xE0) need = 3;
+        else if ((c & 0xF8) == 0xF0) need = 4;
+        else                          return s.substr(0, i);         // invalid start — drop
+        size_t have = back + 1;
+        return (have >= need) ? s : s.substr(0, i);                  // trim incomplete trailing sequence
+    }
+    // Should not reach here; return as-is.
+    return s;
+}
+
+struct Args {
+    std::string model_dir;
+    std::string prompt = "The capital of France is";
+    std::string vocab_path = "tokenizer_data/vocab.bin";
+    int n_predict = 100;
+    int tp_size = 1;
+    int tp_rank = 0;
+    int num_layers = 0;   // 0 = auto
+    int max_seq = 512;
+    int device_id = 0;
+    bool chat_template = false;
+    bool stream = true;
+    bool interactive = false;
+    bool reset_each_turn = false;   // if true, REPL clears KV cache between turns (stateless)
+    std::string system_prompt;       // optional system role for chat mode
+    std::string prompt_file;         // read prompt from file (avoids shell escaping)
+    bool pld_enabled = false;        // prompt lookup decoding
+    int  pld_k = 10;                 // bench_pld_k.sh: K=10 median 105 t/s (3/3 runs 100+), K=8 was 35
+    int  pld_ngram = 1;              // n-gram match size — 1 with multi-level fallback best
+    bool pld_adaptive = false;       // fixed K=10 is simpler and mean-optimal; adaptive --pld-adaptive
+    int  pld_min_hist = 20;          // skip PLD until history >= this (avoid early-token false matches)
+    // PLD degeneration guard (on by default): prevents PLD from amplifying repetition loops.
+    bool pld_guard = true;           // --pld-no-guard disables
+    int  pld_guard_distinct = 3;     // reject draft if distinct tokens < this (≥K/3 heuristic)
+    int  pld_guard_tail = 6;         // reject if draft[0] matches all last N hist tokens
+    int  pld_loop_warn = 8;          // warn once when N consecutive identical tokens emitted
+    float temperature = 0.0f;  // 0 = greedy
+    int   top_k = 0;           // 0 = disabled
+    float top_p = 1.0f;        // 1.0 = disabled
+    uint64_t seed = 0;          // 0 = use time
+    // Qwen3 EOS tokens (from generation_config.json)
+    std::vector<int> eos_ids = {151645, 151643};
+};
+
+static bool parse_args(int argc, char** argv, Args& a) {
+    for (int i = 1; i < argc; i++) {
+        std::string s = argv[i];
+        auto next = [&](const char* f)->const char* {
+            if (i + 1 >= argc) { fprintf(stderr, "missing value for %s\n", f); return nullptr; }
+            return argv[++i];
+        };
+        if      (s == "--model-dir")   { auto v = next(s.c_str()); if (!v) return false; a.model_dir = v; }
+        else if (s == "--prompt")      { auto v = next(s.c_str()); if (!v) return false; a.prompt = v; }
+        else if (s == "--vocab")       { auto v = next(s.c_str()); if (!v) return false; a.vocab_path = v; }
+        else if (s == "--n-predict")   { auto v = next(s.c_str()); if (!v) return false; a.n_predict = std::atoi(v); }
+        else if (s == "--tp-size")     { auto v = next(s.c_str()); if (!v) return false; a.tp_size = std::atoi(v); }
+        else if (s == "--num-layers")  { auto v = next(s.c_str()); if (!v) return false; a.num_layers = std::atoi(v); }
+        else if (s == "--max-seq")     { auto v = next(s.c_str()); if (!v) return false; a.max_seq = std::atoi(v); }
+        else if (s == "--device")      { auto v = next(s.c_str()); if (!v) return false; a.device_id = std::atoi(v); }
+        else if (s == "--temperature") { auto v = next(s.c_str()); if (!v) return false; a.temperature = (float)std::atof(v); }
+        else if (s == "--top-k")       { auto v = next(s.c_str()); if (!v) return false; a.top_k = std::atoi(v); }
+        else if (s == "--top-p")       { auto v = next(s.c_str()); if (!v) return false; a.top_p = (float)std::atof(v); }
+        else if (s == "--seed")        { auto v = next(s.c_str()); if (!v) return false; a.seed = (uint64_t)std::atoll(v); }
+        else if (s == "--chat")        { a.chat_template = true; }
+        else if (s == "--no-stream")   { a.stream = false; }
+        else if (s == "--interactive" || s == "-i") { a.interactive = true; }
+        else if (s == "--reset")       { a.reset_each_turn = true; }
+        else if (s == "--system")      { auto v = next(s.c_str()); if (!v) return false; a.system_prompt = v; }
+        else if (s == "--prompt-file") { auto v = next(s.c_str()); if (!v) return false; a.prompt_file = v; }
+        else if (s == "--pld")         { a.pld_enabled = true; }
+        else if (s == "--pld-k")       { auto v = next(s.c_str()); if (!v) return false; a.pld_k = std::atoi(v); }
+        else if (s == "--pld-ngram")   { auto v = next(s.c_str()); if (!v) return false; a.pld_ngram = std::atoi(v); }
+        else if (s == "--pld-adaptive"){ a.pld_adaptive = true; }
+        else if (s == "--pld-fixed-k") { a.pld_adaptive = false; }   // opt out of adaptive
+        else if (s == "--pld-min-hist"){ auto v = next(s.c_str()); if (!v) return false; a.pld_min_hist = std::atoi(v); }
+        else if (s == "--pld-no-guard"){ a.pld_guard = false; }
+        else if (s == "--pld-guard-distinct"){ auto v = next(s.c_str()); if (!v) return false; a.pld_guard_distinct = std::atoi(v); }
+        else if (s == "--pld-guard-tail"){ auto v = next(s.c_str()); if (!v) return false; a.pld_guard_tail = std::atoi(v); }
+        else if (s == "--pld-loop-warn"){ auto v = next(s.c_str()); if (!v) return false; a.pld_loop_warn = std::atoi(v); }
+        else if (s == "--help" || s == "-h") {
+            printf("Usage: %s --model-dir <path> [options]\n", argv[0]);
+            printf("  --prompt \"text\"     prompt text (default: \"%s\")\n", a.prompt.c_str());
+            printf("  --prompt-file FILE   read prompt from file (overrides --prompt)\n");
+            printf("  --n-predict N        max tokens to generate (default: %d)\n", a.n_predict);
+            printf("  --tp-size N          tensor parallelism (default: 1; or TP_SIZE env)\n");
+            printf("  --num-layers N       limit layers, testing only (default: all)\n");
+            printf("  --max-seq N          KV cache + context cap (default: %d)\n", a.max_seq);
+            printf("  --chat               apply Qwen3 chat template\n");
+            printf("  --system \"text\"    system role for chat\n");
+            printf("  --temperature F      0 = greedy; typical 0.7\n");
+            printf("  --top-k N            0 = disabled\n");
+            printf("  --top-p F            1.0 = disabled; typical 0.8\n");
+            printf("  --seed N             0 = time-based (default)\n");
+            printf("  --no-stream          batch-print final text\n");
+            printf("  -i, --interactive    REPL (multi-turn memory when --chat)\n");
+            printf("  --reset              force stateless REPL (reset each turn)\n");
+            printf("  --pld                enable Prompt Lookup Decoding (greedy only)\n");
+            printf("  --pld-k N            draft window size (default: 4)\n");
+            printf("  --pld-ngram N        match n-gram size (default: 2; multi-level fallback)\n");
+            printf("  --pld-adaptive       adjust K based on recent accept rate\n");
+            printf("  --pld-min-hist N     skip PLD until history >= N tokens (default: 20)\n");
+            printf("  --pld-no-guard       disable degeneration guard (dangerous: can amplify loops)\n");
+            printf("  --pld-guard-distinct N  reject draft with distinct tokens < N (default: 3)\n");
+            printf("  --pld-guard-tail N   reject draft if draft[0] matches all last N hist (default: 6)\n");
+            printf("  --pld-loop-warn N    warn once on N consecutive identical emitted tokens (default: 8)\n");
+            return false;
+        }
+        else { fprintf(stderr, "unknown arg: %s\n", s.c_str()); return false; }
+    }
+    if (a.model_dir.empty()) { fprintf(stderr, "--model-dir required\n"); return false; }
+    if (const char* r = std::getenv("TP_RANK")) a.tp_rank = std::atoi(r);
+    if (const char* s = std::getenv("TP_SIZE")) a.tp_size = std::atoi(s);
+    return true;
+}
+
+// Sample next token from logits. temperature=0 → greedy argmax. Otherwise top-k / top-p.
+static int sample_token(const std::vector<uint16_t>& logits_bf16, int64_t V,
+                        float temperature, int top_k, float top_p, std::mt19937& rng) {
+    if (temperature <= 0.0f) {
+        int best = 0;
+        float bv = bf16_to_float(logits_bf16[0]);
+        for (int64_t i = 1; i < V; i++) {
+            float v = bf16_to_float(logits_bf16[i]);
+            if (v > bv) { bv = v; best = (int)i; }
+        }
+        return best;
+    }
+
+    // Build (logit, id) list as float
+    std::vector<std::pair<float, int>> scored;
+    scored.reserve(V);
+    for (int64_t i = 0; i < V; i++) {
+        scored.emplace_back(bf16_to_float(logits_bf16[i]) / temperature, (int)i);
+    }
+
+    // Top-k: keep highest k entries (partial sort)
+    if (top_k > 0 && top_k < (int)scored.size()) {
+        std::nth_element(scored.begin(), scored.begin() + top_k, scored.end(),
+                         [](const auto& a, const auto& b){ return a.first > b.first; });
+        scored.resize(top_k);
+    }
+
+    // Sort descending for top-p
+    std::sort(scored.begin(), scored.end(),
+              [](const auto& a, const auto& b){ return a.first > b.first; });
+
+    // Softmax (numerically stable)
+    float maxv = scored[0].first;
+    double sum = 0;
+    for (auto& p : scored) { p.first = std::exp(p.first - maxv); sum += p.first; }
+    for (auto& p : scored) p.first /= (float)sum;
+
+    // Top-p nucleus
+    if (top_p > 0.0f && top_p < 1.0f) {
+        double cum = 0;
+        size_t cutoff = scored.size();
+        for (size_t i = 0; i < scored.size(); i++) {
+            cum += scored[i].first;
+            if (cum >= top_p) { cutoff = i + 1; break; }
+        }
+        scored.resize(cutoff);
+        // re-normalize
+        double s = 0; for (auto& p : scored) s += p.first;
+        for (auto& p : scored) p.first /= (float)s;
+    }
+
+    // Sample
+    std::uniform_real_distribution<float> U(0.0f, 1.0f);
+    float r = U(rng), acc = 0.0f;
+    for (auto& p : scored) {
+        acc += p.first;
+        if (r <= acc) return p.second;
+    }
+    return scored.back().second;
+}
+
+// Broadcast a prompt's token_ids from rank 0 to all ranks. For TP>1 the non-master ranks need
+// the tokens before prefill. We use HCCL broadcast: rank 0 provides the count, then the ids.
+// Uses a pre-allocated device buffer (must be large enough for max_seq tokens).
+static bool broadcast_token_ids(Runner& runner, std::vector<int32_t>& ids,
+                                int64_t max_seq, bool is_master) {
+    const ModelConfig& cfg = runner.cfg();
+    if (cfg.tp_size <= 1) return true;
+
+    // Step 1: broadcast count (as int32 on device)
+    DeviceBuffer cnt_dev(4);
+    int32_t cnt = is_master ? (int32_t)ids.size() : 0;
+    ACL_CHECK(aclrtMemcpy(cnt_dev.get(), 4, &cnt, 4, ACL_MEMCPY_HOST_TO_DEVICE));
+    // Access Runner's HCCL context via stream (exposed) and rely on the fact that ctx.comm is owned.
+    // Since hccl_broadcast needs HcclCtx, we need access. Cheapest: friend access via a shim member.
+    // For now, Runner has a stream() accessor; HCCL ctx is private. We'll accept that and broadcast
+    // via a direct call on the comm — but ctx is hidden. Workaround: expose hccl_ctx() on Runner.
+    // ... (see Runner::hccl_ctx() accessor added below)
+    extern HcclCtx* runner_hccl_ctx_shim(Runner& r);   // forward from runner.cpp
+    HcclCtx* ctx = runner_hccl_ctx_shim(runner);
+    if (!ctx) return false;
+
+    if (!hccl_broadcast(*ctx, cnt_dev.get(), 1, HCCL_DATA_TYPE_INT32, 0, runner.stream())) return false;
+    ACL_CHECK(aclrtSynchronizeStream(runner.stream()));
+    ACL_CHECK(aclrtMemcpy(&cnt, 4, cnt_dev.get(), 4, ACL_MEMCPY_DEVICE_TO_HOST));
+    if (cnt <= 0 || cnt > (int32_t)max_seq) {
+        fprintf(stderr, "[rank %d] broadcast: bad count %d\n", cfg.tp_rank, cnt);
+        return false;
+    }
+
+    // Step 2: broadcast the id buffer
+    DeviceBuffer ids_dev(cnt * 4);
+    if (is_master) {
+        ACL_CHECK(aclrtMemcpy(ids_dev.get(), cnt*4, ids.data(), cnt*4, ACL_MEMCPY_HOST_TO_DEVICE));
+    }
+    if (!hccl_broadcast(*ctx, ids_dev.get(), cnt, HCCL_DATA_TYPE_INT32, 0, runner.stream())) return false;
+    ACL_CHECK(aclrtSynchronizeStream(runner.stream()));
+    if (!is_master) {
+        ids.resize(cnt);
+        ACL_CHECK(aclrtMemcpy(ids.data(), cnt*4, ids_dev.get(), cnt*4, ACL_MEMCPY_DEVICE_TO_HOST));
+    }
+    return true;
+}
+
+// Run one generation turn. Assumes KV cache is reset. Returns perf summary.
+struct TurnStats {
+    double prefill_ms = 0; double decode_ms = 0;
+    int n_prompt = 0; int decoded = 0; bool hit_eos = false;
+};
+static TurnStats run_turn(Runner& runner, Tokenizer& tokenizer, const Args& args,
+                          const std::string& prompt, std::mt19937& rng, bool is_master) {
+    TurnStats st;
+
+    // --- Tokenize (on master; broadcast for TP>1) ---
+    std::vector<int32_t> input_ids;
+    if (is_master) {
+        auto raw = tokenizer.encode_via_python(args.model_dir, prompt, args.chat_template);
+        if (raw.empty()) return st;
+        input_ids.reserve(raw.size());
+        for (int v : raw) input_ids.push_back((int32_t)v);
+    }
+    if (args.tp_size > 1) {
+        if (!broadcast_token_ids(runner, input_ids, args.max_seq, is_master)) return st;
+    }
+    if (input_ids.empty()) return st;
+
+    const int64_t V = runner.cfg().vocab_size;
+    std::vector<uint16_t> logits_h(V);
+    auto load_logits = [&](DeviceBuffer& buf) {
+        ACL_CHECK(aclrtMemcpy(logits_h.data(), V*2, buf.get(), V*2, ACL_MEMCPY_DEVICE_TO_HOST));
+    };
+    auto is_eos = [&](int id) {
+        for (int e : args.eos_ids) if (id == e) return true;
+        return false;
+    };
+
+    // --- Prefill ---
+    st.n_prompt = (int)input_ids.size();
+    auto t0 = std::chrono::steady_clock::now();
+    DeviceBuffer logits;
+    if (!runner.prefill(input_ids.data(), (int64_t)input_ids.size(), logits)) return st;
+    auto t1 = std::chrono::steady_clock::now();
+    st.prefill_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
+
+    load_logits(logits);
+    int next_id = sample_token(logits_h, V, args.temperature, args.top_k, args.top_p, rng);
+
+    if (is_master && args.stream) {
+        if (!args.chat_template) printf("%s", prompt.c_str());
+        printf("%s", tokenizer.decode(next_id).c_str());
+        fflush(stdout);
+    }
+
+    std::vector<int> generated = { next_id };
+    st.hit_eos = is_eos(next_id);
+
+    // All tokens (prompt + generated) for PLD n-gram lookup. Non-master ranks still need to
+    // track consistent length for HCCL broadcast of draft proposals.
+    std::vector<int32_t> hist;
+    hist.reserve(input_ids.size() + args.n_predict + 16);
+    for (auto x : input_ids) hist.push_back(x);
+    hist.push_back(next_id);
+
+    // PLD n-gram lookup: search for suffix match ending at end-of-hist; return K tokens following match.
+    // Longer matches = more reliable drafts. Multi-level: try n, fall back to smaller n if no match.
+    auto lookup_one = [&](int ngram, int K) -> std::vector<int32_t> {
+        int hs = (int)hist.size();
+        if (hs < ngram + 1 || K <= 0) return {};
+        for (int start = hs - ngram - 1; start >= 0; start--) {
+            bool match = true;
+            for (int k = 0; k < ngram; k++) {
+                if (hist[start + k] != hist[hs - ngram + k]) { match = false; break; }
+            }
+            if (match) {
+                int after = start + ngram;
+                std::vector<int32_t> d;
+                for (int k = 0; k < K && after + k < hs; k++) {
+                    d.push_back(hist[after + k]);
+                    if (is_eos(hist[after + k])) break;
+                }
+                if (!d.empty()) return d;
+            }
+        }
+        return {};
+    };
+    // Multi-level: try configured n first, then n-1, then n-2 (down to 1).
+    auto lookup_draft = [&](int ngram, int K) -> std::vector<int32_t> {
+        for (int n = ngram; n >= 1; n--) {
+            auto d = lookup_one(n, K);
+            if (!d.empty()) return d;
+        }
+        return {};
+    };
+
+    // Degeneration guard: classify a draft as repetition-induced so we can fall back to single
+    // decode (and avoid PLD amplifying model's own repetition loop into a runaway "W W W …" mess).
+    // Returns nullptr if draft is OK, else a short reason string for stats.
+    auto draft_degenerate = [&](const std::vector<int32_t>& d) -> const char* {
+        if (!args.pld_guard || d.empty()) return nullptr;
+        // (1) distinct-token count: a draft of K tokens with < args.pld_guard_distinct distinct
+        //     values means n-gram is echoing a loop. Only apply when draft is long enough.
+        if ((int)d.size() >= 3) {
+            int distinct = 0;
+            for (int i = 0; i < (int)d.size(); i++) {
+                bool seen = false;
+                for (int j = 0; j < i; j++) { if (d[j] == d[i]) { seen = true; break; } }
+                if (!seen) distinct++;
+            }
+            if (distinct < args.pld_guard_distinct) return "low-distinct";
+        }
+        // (2) tail echo: if the last N hist tokens are all equal to draft[0], the model is already
+        //     in a short loop — accepting the draft will just confirm the loop at batch speed.
+        int tail_n = std::min(args.pld_guard_tail, (int)hist.size());
+        if (tail_n >= 3) {
+            int matches = 0;
+            for (int i = (int)hist.size() - tail_n; i < (int)hist.size(); i++) {
+                if (hist[i] == d[0]) matches++;
+            }
+            if (matches == tail_n) return "tail-echo";
+        }
+        return nullptr;
+    };
+
+    // --- Decode loop ---
+    auto t2 = std::chrono::steady_clock::now();
+    int pld_verifies = 0, pld_accepted = 0;
+    int pld_rej_lowdist = 0, pld_rej_tailecho = 0;   // guard rejection counters
+    bool loop_warned = false;                         // warn-once state
+
+    // Adaptive K state: recent accept counts for moving-average decisions
+    const int ADAPT_WINDOW = 8;
+    std::vector<int> recent_accepts;
+    int current_k = args.pld_k;
+    bool pld_disabled_adapt = false;   // set true when recent accept rate is too low to benefit
+
+    while (st.decoded < args.n_predict - 1 && !st.hit_eos) {
+        // Adaptive K: scale K with recent accept rate.
+        // No auto-disable: since S=K+1 forward ≈ S=1 forward (latency-bound), even accept=0.1
+        // still nets slightly positive — PLD doesn't "hurt" as long as ngram lookup is cheap.
+        if (args.pld_adaptive && (int)recent_accepts.size() >= ADAPT_WINDOW) {
+            double avg = 0;
+            for (int a : recent_accepts) avg += a;
+            avg /= recent_accepts.size();
+            // Aim: K = 2*avg + 4 (generous window to catch upswings). Clamp [4, 12].
+            current_k = std::max(4, std::min(12, (int)std::round(2.0 * avg + 4.0)));
+        }
+
+        // Try PLD speculation path — skip until enough history accumulated
+        std::vector<int32_t> draft;
+        if (args.pld_enabled && (int)hist.size() >= args.pld_min_hist && is_master) {
+            draft = lookup_draft(args.pld_ngram, current_k);
+            // Degeneration guard: if draft looks like repetition-loop echo, drop it so this
+            // iteration falls through to normal single decode. This does NOT stop a loop the model
+            // is already in (greedy is deterministic), but it prevents PLD from running the loop
+            // at batch speed while masquerading as a speedup.
+            if (!draft.empty()) {
+                const char* reason = draft_degenerate(draft);
+                if (reason) {
+                    if (reason[0] == 'l') pld_rej_lowdist++;
+                    else                  pld_rej_tailecho++;
+                    draft.clear();
+                }
+            }
+        }
+        // For TP>1, broadcast draft across ranks. Only broadcast if master has a non-empty draft;
+        // otherwise all ranks take the no-draft path (normal decode).
+        bool has_draft = is_master ? !draft.empty() : false;
+        // Broadcast the "has_draft" flag (using a 1-element count: 1 = yes, 0 = no)
+        if (args.tp_size > 1) {
+            extern HcclCtx* runner_hccl_ctx_shim(Runner&);
+            HcclCtx* ctx = runner_hccl_ctx_shim(runner);
+            DeviceBuffer flag(4);
+            int32_t f = has_draft ? 1 : 0;
+            ACL_CHECK(aclrtMemcpy(flag.get(), 4, &f, 4, ACL_MEMCPY_HOST_TO_DEVICE));
+            hccl_broadcast(*ctx, flag.get(), 1, HCCL_DATA_TYPE_INT32, 0, runner.stream());
+            ACL_CHECK(aclrtSynchronizeStream(runner.stream()));
+            ACL_CHECK(aclrtMemcpy(&f, 4, flag.get(), 4, ACL_MEMCPY_DEVICE_TO_HOST));
+            has_draft = (f != 0);
+            if (has_draft) {
+                std::vector<int32_t> d = draft;
+                broadcast_token_ids(runner, d, args.max_seq, is_master);
+                draft = d;
+            } else {
+                draft.clear();
+            }
+        }
+
+        if (args.pld_enabled && (int)draft.size() >= 1 && args.temperature == 0.0f) {
+            // Batch verify: input = [next_id, draft[0], ..., draft[K-1]]
+            std::vector<int32_t> batch_input = { next_id };
+            for (auto d : draft) batch_input.push_back(d);
+            int S = (int)batch_input.size();
+            DeviceBuffer batch_logits;
+            if (!runner.decode_batch(batch_input.data(), S, batch_logits)) break;
+
+            std::vector<uint16_t> blh(S * V);
+            if (is_master) ACL_CHECK(aclrtMemcpy(blh.data(), S*V*2, batch_logits.get(), S*V*2, ACL_MEMCPY_DEVICE_TO_HOST));
+
+            // Accept longest prefix: draft[i] is "candidate" for position past+i+1.
+            // blh row i predicts position past+i+1 (follows batch_input[i]).
+            // Verify: blh[0].argmax == draft[0]? (i.e., does model agree with draft's first proposal)
+            int accept = 0, new_next = next_id;
+            if (is_master) {
+                for (int i = 0; i < S - 1; i++) {
+                    int pred = 0; float bv = bf16_to_float(blh[i * V]);
+                    for (int k = 1; k < V; k++) { float v = bf16_to_float(blh[i*V + k]); if (v > bv) { bv = v; pred = k; } }
+                    if (pred == (int)draft[i]) accept++;
+                    else { new_next = pred; break; }
+                }
+                if (accept == S - 1) {
+                    // All draft accepted, bonus from last row
+                    int pred = 0; float bv = bf16_to_float(blh[(S-1) * V]);
+                    for (int k = 1; k < V; k++) { float v = bf16_to_float(blh[(S-1)*V + k]); if (v > bv) { bv = v; pred = k; } }
+                    new_next = pred;
+                }
+            }
+
+            // Broadcast accept count + new_next across TP ranks
+            if (args.tp_size > 1) {
+                int32_t packed[2] = { (int32_t)accept, (int32_t)new_next };
+                std::vector<int32_t> p(packed, packed + 2);
+                broadcast_token_ids(runner, p, args.max_seq, is_master);
+                if (p.size() == 2) { accept = p[0]; new_next = p[1]; }
+            }
+
+            // Rewind KV for rejected drafts
+            int64_t rewind = (int64_t)(S - 1 - accept);  // drafts not accepted (excluding bonus)
+            if (rewind > 0) runner.rewind_cache(rewind);
+
+            // Commit accepted drafts + bonus to hist and emit
+            for (int i = 0; i < accept; i++) {
+                int tok = (int)draft[i];
+                hist.push_back(tok);
+                generated.push_back(tok);
+                st.decoded++;
+                if (is_master && args.stream) { printf("%s", tokenizer.decode(tok).c_str()); fflush(stdout); }
+                if (is_eos(tok)) { st.hit_eos = true; break; }
+            }
+            pld_verifies++; pld_accepted += accept;
+            // Track recent accept for adaptive K
+            if (args.pld_adaptive) {
+                recent_accepts.push_back(accept);
+                if ((int)recent_accepts.size() > ADAPT_WINDOW) recent_accepts.erase(recent_accepts.begin());
+            }
+            if (st.hit_eos) break;
+
+            // Bonus token (new_next) is also committed
+            hist.push_back(new_next);
+            generated.push_back(new_next);
+            st.decoded++;
+            if (is_master && args.stream) { printf("%s", tokenizer.decode(new_next).c_str()); fflush(stdout); }
+            if (is_eos(new_next)) { st.hit_eos = true; break; }
+            next_id = new_next;
+        } else {
+            // Normal decode
+            DeviceBuffer logits2;
+            if (!runner.decode((int32_t)next_id, logits2)) break;
+            load_logits(logits2);
+            next_id = sample_token(logits_h, V, args.temperature, args.top_k, args.top_p, rng);
+            hist.push_back(next_id);
+            generated.push_back(next_id);
+            st.decoded++;
+            if (is_master && args.stream) { printf("%s", tokenizer.decode(next_id).c_str()); fflush(stdout); }
+            if (is_eos(next_id)) { st.hit_eos = true; break; }
+        }
+        // Loop-warn: emit a one-shot warning to stderr if the tail of generated is all-same-token.
+        // Does not stop generation (user may want to see what happens) — just flags output quality.
+        if (is_master && !loop_warned && args.pld_loop_warn > 0 &&
+            (int)generated.size() >= args.pld_loop_warn) {
+            int tail = args.pld_loop_warn;
+            int anchor = generated[(int)generated.size() - tail];
+            bool all_same = true;
+            for (int i = (int)generated.size() - tail + 1; i < (int)generated.size(); i++) {
+                if (generated[i] != anchor) { all_same = false; break; }
+            }
+            if (all_same) {
+                fprintf(stderr, "\n[warn] %d consecutive identical tokens — likely degeneration loop; output after this point is suspect\n", tail);
+                loop_warned = true;
+            }
+        }
+    }
+    auto t3 = std::chrono::steady_clock::now();
+    st.decode_ms = std::chrono::duration<double, std::milli>(t3 - t2).count();
+    if (is_master && args.pld_enabled) {
+        if (pld_verifies > 0) {
+            fprintf(stderr, "\n[pld] %d verifies, %d drafts accepted, avg=%.2f",
+                    pld_verifies, pld_accepted, (double)pld_accepted / pld_verifies);
+        } else {
+            fprintf(stderr, "\n[pld] 0 verifies (all drafts blocked or none found)");
+        }
+        if (args.pld_guard && (pld_rej_lowdist + pld_rej_tailecho) > 0) {
+            fprintf(stderr, "; guard rejections: low-distinct=%d tail-echo=%d",
+                    pld_rej_lowdist, pld_rej_tailecho);
+        }
+        fprintf(stderr, "\n");
+    }
+
+    if (is_master) {
+        if (args.stream) { printf("\n"); fflush(stdout); }
+        else {
+            std::string text = tokenizer.decode(generated);
+            printf("%s%s\n", args.chat_template ? "" : prompt.c_str(), text.c_str());
+        }
+    }
+    return st;
+}
+
+static bool load_file(const std::string& path, std::string& out) {
+    FILE* f = fopen(path.c_str(), "rb");
+    if (!f) { fprintf(stderr, "[cli] cannot open %s\n", path.c_str()); return false; }
+    fseek(f, 0, SEEK_END); long sz = ftell(f); fseek(f, 0, SEEK_SET);
+    out.resize(sz);
+    size_t n = fread(out.data(), 1, sz, f);
+    fclose(f);
+    if ((long)n != sz) { fprintf(stderr, "[cli] short read from %s\n", path.c_str()); return false; }
+    // Strip a single trailing newline (common in text files)
+    if (!out.empty() && out.back() == '\n') out.pop_back();
+    return true;
+}
+
+int main(int argc, char** argv) {
+    Args args;
+    if (!parse_args(argc, argv, args)) return 1;
+
+    // --prompt-file overrides --prompt
+    if (!args.prompt_file.empty()) {
+        if (!load_file(args.prompt_file, args.prompt)) return 1;
+    }
+
+    const bool is_master = (args.tp_rank == 0);
+    std::mt19937 rng(args.seed ? args.seed :
+                     (uint64_t)std::chrono::steady_clock::now().time_since_epoch().count());
+
+    if (is_master) {
+        printf("[cli] model=%s\n", args.model_dir.c_str());
+        printf("[cli] tp=%d   n_predict=%d   temp=%.2f top_k=%d top_p=%.2f   chat=%d   interactive=%d\n",
+               args.tp_size, args.n_predict, args.temperature, args.top_k, args.top_p,
+               args.chat_template, args.interactive);
+        fflush(stdout);
+    }
+
+    Tokenizer tokenizer;
+    if (!tokenizer.load(args.vocab_path)) {
+        fprintf(stderr, "[cli] failed to load vocab %s\n", args.vocab_path.c_str()); return 1;
+    }
+
+    Runner runner;
+    int num_layers = args.num_layers;
+    if (num_layers == 0) {
+        ModelConfig probe;
+        if (!probe.load_from_json(args.model_dir + "/config.json")) return 1;
+        num_layers = (int)probe.num_hidden_layers;
+    }
+    if (!runner.init(args.model_dir, args.tp_size, args.tp_rank,
+                     num_layers, args.max_seq, args.device_id)) return 1;
+    if (const char* p = std::getenv("LCA_PROFILE"); p && std::atoi(p) != 0) {
+        runner.profile_enabled = true;
+    }
+    // Warmup: cut cold-start latency. Controlled via LCA_WARMUP env (default 0 to keep behavior).
+    if (const char* w = std::getenv("LCA_WARMUP"); w) {
+        int n = std::atoi(w);
+        if (n > 0) runner.warmup(n);
+    }
+
+    if (args.interactive) {
+        const bool multi_turn = args.chat_template && !args.reset_each_turn;
+        if (is_master) {
+            printf("\n[cli] === interactive mode ===\n");
+            if (multi_turn) {
+                printf("[cli] multi-turn chat (KV cache preserved). Commands: 'quit', 'reset'.\n");
+                if (!args.system_prompt.empty()) {
+                    printf("[cli] system: %s\n", args.system_prompt.c_str());
+                }
+            } else {
+                printf("[cli] stateless mode (KV cache reset each turn). Command: 'quit'.\n");
+                if (!args.chat_template) {
+                    printf("[cli] (hint: add --chat for multi-turn conversational memory)\n");
+                }
+            }
+            fflush(stdout);
+        }
+
+        // Conversation history: accumulated (role, content) pairs. System prompt seeded if present.
+        std::vector<std::pair<std::string, std::string>> conversation;
+        if (multi_turn && !args.system_prompt.empty()) {
+            conversation.emplace_back("system", args.system_prompt);
+        }
+
+        auto* hccl_ctx = runner_hccl_ctx_shim(runner);
+
+        // Signal types (broadcast as int32):  0 = normal turn, 1 = quit, 2 = reset.
+        auto broadcast_signal = [&](int32_t sig)->int32_t {
+            if (args.tp_size <= 1) return sig;
+            DeviceBuffer s(4);
+            ACL_CHECK(aclrtMemcpy(s.get(), 4, &sig, 4, ACL_MEMCPY_HOST_TO_DEVICE));
+            hccl_broadcast(*hccl_ctx, s.get(), 1, HCCL_DATA_TYPE_INT32, 0, runner.stream());
+            ACL_CHECK(aclrtSynchronizeStream(runner.stream()));
+            int32_t r; ACL_CHECK(aclrtMemcpy(&r, 4, s.get(), 4, ACL_MEMCPY_DEVICE_TO_HOST));
+            return r;
+        };
+
+        while (true) {
+            std::string prompt;
+            int32_t sig = 0;
+            if (is_master) {
+                printf("\n> "); fflush(stdout);
+                if (!std::getline(std::cin, prompt)) sig = 1;
+                else if (prompt == "quit" || prompt == "exit") sig = 1;
+                else if (prompt == "reset")                    sig = 2;
+                else if (prompt.empty())                        sig = 3;   // skip
+            }
+            sig = broadcast_signal(sig);
+            if (sig == 1) break;
+            if (sig == 2) {
+                runner.reset_cache();
+                conversation.clear();
+                if (multi_turn && !args.system_prompt.empty())
+                    conversation.emplace_back("system", args.system_prompt);
+                if (is_master) { printf("[cli] (cache + conversation reset)\n"); fflush(stdout); }
+                continue;
+            }
+            if (sig == 3) continue;
+
+            TurnStats st;
+            if (multi_turn) {
+                // Append user message and tokenize full conversation. Prefill DELTA only.
+                if (is_master) conversation.emplace_back("user", prompt);
+                // Also ranks 1..N-1 need to track conversation (needed for correct delta count on
+                // subsequent turns if TP ever tokenizes per-rank — currently rank 0 tokenizes).
+                std::vector<int32_t> full_ids;
+                if (is_master) {
+                    auto raw = tokenizer.encode_conversation_via_python(args.model_dir, conversation, /*gen_prompt=*/true);
+                    full_ids.reserve(raw.size());
+                    for (int v : raw) full_ids.push_back((int32_t)v);
+                }
+                // Broadcast full_ids (variable-length). Use the same shim as broadcast_token_ids.
+                if (args.tp_size > 1) {
+                    if (!broadcast_token_ids(runner, full_ids, args.max_seq, is_master)) break;
+                }
+                if (full_ids.empty()) { if (is_master) printf("[cli] tokenize failed\n"); continue; }
+
+                int64_t past = runner.past_len();
+                if ((int64_t)full_ids.size() < past) { runner.reset_cache(); past = 0; }
+                std::vector<int32_t> delta(full_ids.begin() + past, full_ids.end());
+                if (delta.empty()) {
+                    if (is_master) printf("[cli] (no new tokens)\n");
+                    continue;
+                }
+                // Overflow check — simple policy: warn + auto-reset if the turn + generation
+                // would exceed max_seq. Conversation history is cleared (except --system) so
+                // the user's current prompt still fits.
+                if ((int64_t)(past + delta.size()) + args.n_predict > args.max_seq) {
+                    if (is_master) {
+                        fprintf(stderr, "[cli] context %ld + gen %d > max_seq %d — auto-resetting\n",
+                                (long)(past + delta.size()), args.n_predict, args.max_seq);
+                    }
+                    runner.reset_cache();
+                    // Rebuild conversation: keep only system + current user turn.
+                    if (is_master) {
+                        std::vector<std::pair<std::string, std::string>> fresh;
+                        for (auto& m : conversation) if (m.first == "system") fresh.push_back(m);
+                        if (!conversation.empty() && conversation.back().first == "user") {
+                            fresh.push_back(conversation.back());
+                        }
+                        conversation = std::move(fresh);
+                        auto raw = tokenizer.encode_conversation_via_python(args.model_dir, conversation, true);
+                        full_ids.clear();
+                        for (int v : raw) full_ids.push_back((int32_t)v);
+                    }
+                    if (args.tp_size > 1) {
+                        if (!broadcast_token_ids(runner, full_ids, args.max_seq, is_master)) break;
+                    }
+                    delta.assign(full_ids.begin(), full_ids.end());
+                    past = 0;
+                }
+
+                // --- Prefill the delta ---
+                st.n_prompt = (int)delta.size();
+                auto t0 = std::chrono::steady_clock::now();
+                DeviceBuffer logits;
+                if (!runner.prefill(delta.data(), (int64_t)delta.size(), logits)) break;
+                auto t1 = std::chrono::steady_clock::now();
+                st.prefill_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
+
+                const int64_t V = runner.cfg().vocab_size;
+                std::vector<uint16_t> logits_h(V);
+                auto load_logits = [&](DeviceBuffer& buf) {
+                    ACL_CHECK(aclrtMemcpy(logits_h.data(), V*2, buf.get(), V*2, ACL_MEMCPY_DEVICE_TO_HOST));
+                };
+                auto is_eos = [&](int id) {
+                    for (int e : args.eos_ids) if (id == e) return true;
+                    return false;
+                };
+                load_logits(logits);
+                int next_id = sample_token(logits_h, V, args.temperature, args.top_k, args.top_p, rng);
+
+                std::vector<int> assistant_ids = { next_id };
+                if (is_master) { printf("%s", tokenizer.decode(next_id).c_str()); fflush(stdout); }
+                st.hit_eos = is_eos(next_id);
+
+                auto t2 = std::chrono::steady_clock::now();
+                for (int step = 1; step < args.n_predict && !st.hit_eos; step++) {
+                    DeviceBuffer logits2;
+                    if (!runner.decode((int32_t)next_id, logits2)) break;
+                    load_logits(logits2);
+                    next_id = sample_token(logits_h, V, args.temperature, args.top_k, args.top_p, rng);
+                    assistant_ids.push_back(next_id);
+                    st.decoded++;
+                    if (is_master) { printf("%s", tokenizer.decode(next_id).c_str()); fflush(stdout); }
+                    if (is_eos(next_id)) { st.hit_eos = true; break; }
+                }
+                auto t3 = std::chrono::steady_clock::now();
+                st.decode_ms = std::chrono::duration<double, std::milli>(t3 - t2).count();
+                if (is_master) { printf("\n"); fflush(stdout); }
+
+                // Record assistant reply in conversation (strip trailing EOS before decode,
+                // and trim incomplete UTF-8 tail if generation was cut mid-codepoint).
+                if (is_master) {
+                    std::vector<int> content_ids;
+                    for (int id : assistant_ids) { if (is_eos(id)) break; content_ids.push_back(id); }
+                    conversation.emplace_back("assistant", utf8_trim_incomplete(tokenizer.decode(content_ids)));
+                }
+            } else {
+                // Stateless: reset cache, one-shot prompt
+                runner.reset_cache();
+                st = run_turn(runner, tokenizer, args, prompt, rng, is_master);
+            }
+
+            if (is_master) {
+                double tgs = (st.decode_ms > 0) ? (st.decoded * 1000.0 / st.decode_ms) : 0.0;
+                printf("[perf] prefill %d tok %.0fms   decode %d tok %.0fms = %.2f t/s%s   past_len=%ld\n",
+                       st.n_prompt, st.prefill_ms, st.decoded, st.decode_ms, tgs,
+                       st.hit_eos ? "  (EOS)" : "", runner.past_len());
+                fflush(stdout);
+            }
+        }
+        if (is_master) printf("[cli] bye\n");
+        return 0;
+    }
+
+    // One-shot mode
+    TurnStats st = run_turn(runner, tokenizer, args, args.prompt, rng, is_master);
+    if (is_master) runner.print_profile_summary();
+    if (is_master) {
+        if (st.hit_eos) printf("[cli] (hit EOS)\n");
+        printf("\n[perf] prefill: %.1fms for %d tokens = %.2f t/s\n",
+               st.prefill_ms, st.n_prompt,
+               (st.prefill_ms > 0) ? (st.n_prompt * 1000.0 / st.prefill_ms) : 0.0);
+        if (st.decoded > 0) {
+            printf("[perf] decode : %.1fms for %d tokens = %.2f t/s (TG)\n",
+                   st.decode_ms, st.decoded, (st.decoded * 1000.0) / st.decode_ms);
+        }
+    }
+    return 0;
+}

src/model_config.cpp

@@ -0,0 +1,115 @@
+#include "model_config.h"
+
+#include <cstdio>
+#include <fstream>
+#include <sstream>
+
+#include "json.hpp"
+using json = nlohmann::json;
+
+bool ModelConfig::load_from_json(const std::string& path) {
+    std::ifstream f(path);
+    if (!f) {
+        fprintf(stderr, "ModelConfig: cannot open %s\n", path.c_str());
+        return false;
+    }
+    json j;
+    try { f >> j; } catch (std::exception& e) {
+        fprintf(stderr, "ModelConfig: bad json: %s\n", e.what());
+        return false;
+    }
+
+    auto get = [&](const char* k, auto def) {
+        if (j.contains(k) && !j[k].is_null()) return j[k].get<decltype(def)>();
+        return def;
+    };
+
+    vocab_size               = get("vocab_size",              (int64_t)0);
+    hidden_size              = get("hidden_size",             (int64_t)0);
+    intermediate_size        = get("intermediate_size",       (int64_t)0);
+    moe_intermediate_size    = get("moe_intermediate_size",   (int64_t)0);
+    num_hidden_layers        = get("num_hidden_layers",       (int64_t)0);
+    num_attention_heads      = get("num_attention_heads",     (int64_t)0);
+    num_key_value_heads      = get("num_key_value_heads",     (int64_t)0);
+    head_dim                 = get("head_dim",                (int64_t)0);
+    num_experts              = get("num_experts",             (int64_t)0);
+    num_experts_per_tok      = get("num_experts_per_tok",     (int64_t)0);
+    max_position_embeddings  = get("max_position_embeddings", (int64_t)0);
+    rope_theta               = (float)get("rope_theta",       (double)10000.0);
+    rms_norm_eps             = (float)get("rms_norm_eps",     (double)1e-6);
+    norm_topk_prob           = get("norm_topk_prob",          true);
+    tie_word_embeddings      = get("tie_word_embeddings",     false);
+    bos_token_id             = get("bos_token_id",            (int64_t)0);
+    eos_token_id             = get("eos_token_id",            (int64_t)0);
+
+    // Sanity
+    if (num_attention_heads == 0 || head_dim == 0 || hidden_size == 0) {
+        fprintf(stderr, "ModelConfig: required fields missing\n");
+        return false;
+    }
+    return true;
+}
+
+void ModelConfig::compute_derived(int tps, int tpr) {
+    tp_size = tps;
+    tp_rank = tpr;
+
+    // Attention Q: split by head
+    if (num_attention_heads % tp_size != 0) {
+        fprintf(stderr, "WARN: num_attention_heads=%ld not divisible by tp_size=%d\n",
+                num_attention_heads, tp_size);
+    }
+    n_heads_per_rank = num_attention_heads / tp_size;
+    q_dim_per_rank   = n_heads_per_rank * head_dim;
+
+    // Attention KV: GQA sharding.
+    //   Case A (tp_size <= num_kv_heads): split KV heads across ranks.
+    //     n_kv_heads_per_rank = num_kv_heads / tp_size
+    //   Case B (tp_size > num_kv_heads): each rank gets ONE kv head shared by multiple ranks.
+    //     Ranks in the same "group" share one kv head (ratio = tp_size / num_kv_heads).
+    //     n_kv_heads_per_rank = 1
+    //     kv_head_idx_for_rank = tp_rank / (tp_size / num_kv_heads)
+    //   This matches the GQA semantics: each group of (num_q_heads / num_kv_heads) Q heads
+    //   shares one KV head. FIAS is given matched Hq (rank-local Q heads) and Hkv=1.
+    if (tp_size <= num_key_value_heads && num_key_value_heads % tp_size == 0) {
+        n_kv_heads_per_rank = num_key_value_heads / tp_size;
+    } else if (tp_size % num_key_value_heads == 0) {
+        n_kv_heads_per_rank = 1;
+    } else {
+        fprintf(stderr, "WARN: non-standard TP/KV head ratio: tp=%d kv=%ld — falling back to replicate-all\n",
+                tp_size, num_key_value_heads);
+        n_kv_heads_per_rank = num_key_value_heads;
+    }
+    kv_dim_per_rank = n_kv_heads_per_rank * head_dim;
+
+    // MoE intermediate dim split
+    if (moe_intermediate_size % tp_size != 0) {
+        fprintf(stderr, "WARN: moe_intermediate_size=%ld not divisible by tp_size=%d\n",
+                moe_intermediate_size, tp_size);
+    }
+    i_per_rank = moe_intermediate_size / tp_size;
+}
+
+std::string ModelConfig::describe() const {
+    std::ostringstream os;
+    os << "Qwen3MoE config:\n"
+       << "  vocab_size            = " << vocab_size              << "\n"
+       << "  hidden_size           = " << hidden_size             << "\n"
+       << "  num_hidden_layers     = " << num_hidden_layers       << "\n"
+       << "  num_attention_heads   = " << num_attention_heads     << "\n"
+       << "  num_key_value_heads   = " << num_key_value_heads     << "\n"
+       << "  head_dim              = " << head_dim                << "\n"
+       << "  num_experts           = " << num_experts             << "\n"
+       << "  num_experts_per_tok   = " << num_experts_per_tok     << "\n"
+       << "  moe_intermediate_size = " << moe_intermediate_size   << "\n"
+       << "  rope_theta            = " << rope_theta              << "\n"
+       << "  rms_norm_eps          = " << rms_norm_eps            << "\n"
+       << "  max_pos_embeddings    = " << max_position_embeddings << "\n"
+       << "TP rank " << tp_rank << " / " << tp_size << " derived:\n"
+       << "  n_heads_per_rank      = " << n_heads_per_rank        << "\n"
+       << "  q_dim_per_rank        = " << q_dim_per_rank          << "\n"
+       << "  n_kv_heads_per_rank   = " << n_kv_heads_per_rank     << "\n"
+       << "  kv_dim_per_rank       = " << kv_dim_per_rank         << "\n"
+       << "  i_per_rank            = " << i_per_rank              << "\n";
+    return os.str();
+}

src/runner.cpp

@@ -0,0 +1,428 @@
+#include "runner.h"
+
+#include <chrono>
+#include <cstdio>
+#include <cstring>
+
+// Expose HCCL context for the CLI broadcast helper.
+HcclCtx* runner_hccl_ctx_shim(Runner& r) { return &r.hccl_ctx(); }
+
+bool Runner::init(const std::string& model_dir, int tp_size, int tp_rank,
+                  int num_layers_to_load, int64_t max_seq, int device_id) {
+    if (!cfg_.load_from_json(model_dir + "/config.json")) return false;
+    cfg_.compute_derived(tp_size, tp_rank);
+    if (num_layers_to_load < 1 || num_layers_to_load > (int)cfg_.num_hidden_layers) {
+        fprintf(stderr, "runner: invalid num_layers %d (max %ld)\n",
+                num_layers_to_load, cfg_.num_hidden_layers);
+        return false;
+    }
+    num_layers_ = num_layers_to_load;
+    max_seq_ = max_seq;
+
+    if (!st_.open(model_dir)) return false;
+    rt_.init(device_id);
+
+    // HCCL init (no-op if tp_size == 1)
+    if (!hccl_init(hccl_ctx_, tp_size, tp_rank)) {
+        fprintf(stderr, "runner: HCCL init failed\n");
+        return false;
+    }
+
+    DeviceWeightsLoader dw(st_, cfg_);
+    printf("runner: loading shared weights (embed, lm_head, final_norm)...\n");
+    if (!dw.load_shared(shared_)) return false;
+
+    attn_.resize(num_layers_);
+    moe_.resize(num_layers_);
+    k_cache_.resize(num_layers_);
+    v_cache_.resize(num_layers_);
+
+    const int64_t KV_DIM = cfg_.n_kv_heads_per_rank * cfg_.head_dim;
+    for (int L = 0; L < num_layers_; L++) {
+        printf("runner: loading layer %d/%d...\n", L + 1, num_layers_);
+        if (!dw.load_attention(L, attn_[L])) return false;
+        if (!dw.load_moe(L, rt_.stream(), moe_[L])) return false;
+        k_cache_[L].alloc(max_seq_ * KV_DIM * 2);
+        v_cache_[L].alloc(max_seq_ * KV_DIM * 2);
+    }
+    rt_.sync();
+
+    // Prefill mask (2048x2048 bool causal)
+    const int64_t MASK = 2048;
+    std::vector<uint8_t> mh(MASK * MASK, 0);
+    for (int i = 0; i < MASK; i++)
+        for (int j = i+1; j < MASK; j++) mh[i*MASK + j] = 1;
+    prefill_mask_dev_.alloc(MASK * MASK);
+    ACL_CHECK(aclrtMemcpy(prefill_mask_dev_.get(), MASK*MASK, mh.data(), MASK*MASK, ACL_MEMCPY_HOST_TO_DEVICE));
+
+    // Pre-compute RoPE cos/sin table once (covers all positions up to max_seq_)
+    rope_cache_build(rope_cache_, max_seq_, cfg_.head_dim, cfg_.rope_theta);
+
+    past_len_ = 0;
+    cur_S_capacity_ = 0;
+    return true;
+}
+
+static void ensure_sc_(DeviceBuffer& buf, size_t needed) {
+    if (buf.size < needed) buf.alloc(needed);
+}
+
+static void ensure_all_scratch_(Runner* self, int64_t S, const ModelConfig& cfg,
+                                 DeviceBuffer& q_sc, DeviceBuffer& k_sc, DeviceBuffer& v_sc,
+                                 DeviceBuffer& xn_sc, DeviceBuffer& rstd_sc, DeviceBuffer& rope_sc,
+                                 DeviceBuffer& attn_fias_sc, DeviceBuffer& attn_out_sc,
+                                 DeviceBuffer& moe_xn, DeviceBuffer& moe_rstd, DeviceBuffer& moe_logits,
+                                 DeviceBuffer& moe_topk_w, DeviceBuffer& moe_topk_idx, DeviceBuffer& moe_row_idx,
+                                 DeviceBuffer& moe_ex_x, DeviceBuffer& moe_ex_ri, DeviceBuffer& moe_tpe,
+                                 DeviceBuffer& moe_fwd,
+                                 DeviceBuffer& moe_gate, DeviceBuffer& moe_up, DeviceBuffer& moe_down,
+                                 DeviceBuffer& moe_packed, DeviceBuffer& moe_weighted, DeviceBuffer& moe_out,
+                                 DeviceBuffer& moe_norm_sum,
+                                 DeviceBuffer& x_buf_a, DeviceBuffer& x_buf_b) {
+    (void)self;
+    const int64_t D = cfg.hidden_size;
+    const int64_t Hq = cfg.n_heads_per_rank, Hkv = cfg.n_kv_heads_per_rank;
+    const int64_t Dh = cfg.head_dim;
+    const int64_t Q_DIM = Hq * Dh;
+    const int64_t KV_DIM = Hkv * Dh;
+    const int64_t I = cfg.i_per_rank, E = cfg.num_experts, K = cfg.num_experts_per_tok;
+    const int64_t TOTAL = S * K;
+
+    ensure_sc_(q_sc,          S * Q_DIM  * 2);
+    ensure_sc_(k_sc,          S * KV_DIM * 2);
+    ensure_sc_(v_sc,          S * KV_DIM * 2);
+    ensure_sc_(xn_sc,         S * D * 2);
+    ensure_sc_(rstd_sc,       S * std::max(Hq, Hkv) * 4);
+    ensure_sc_(rope_sc,       1 * S * Hq * Dh * 2);
+    ensure_sc_(attn_fias_sc,  S * Q_DIM * 2);
+    ensure_sc_(attn_out_sc,   S * D * 2);
+
+    ensure_sc_(moe_xn,        S * D * 2);
+    ensure_sc_(moe_rstd,      S * 4);
+    ensure_sc_(moe_logits,    S * E * 2);
+    ensure_sc_(moe_topk_w,    S * K * 2);
+    ensure_sc_(moe_topk_idx,  S * K * 4);
+    ensure_sc_(moe_row_idx,   S * K * 4);
+    ensure_sc_(moe_ex_x,      TOTAL * D * 2);
+    ensure_sc_(moe_ex_ri,     TOTAL * 4);
+    ensure_sc_(moe_tpe,       E * 8);
+    ensure_sc_(moe_fwd,       TOTAL * 8);
+    ensure_sc_(moe_gate,      TOTAL * I * 2);
+    ensure_sc_(moe_up,        TOTAL * I * 2);
+    ensure_sc_(moe_down,      TOTAL * D * 2);
+    ensure_sc_(moe_packed,    TOTAL * D * 2);
+    ensure_sc_(moe_weighted,  S * K * D * 2);
+    ensure_sc_(moe_out,       S * D * 2);
+    ensure_sc_(moe_norm_sum,  S * 2);
+
+    ensure_sc_(x_buf_a,       S * D * 2);
+    ensure_sc_(x_buf_b,       S * D * 2);
+}
+
+void Runner::layer_forward_(int layer_idx, int64_t S, void* x_in, void* x_out, bool batch_decode_mode) {
+    const int64_t D = cfg_.hidden_size;
+
+    // Attention mask selection:
+    //   prefill (S>1, past=0):       2048×2048 upper-tri + sparse_mode=3 (FIAS internal causal)
+    //   decode  (S==1):              mask=nullptr + sparse_mode=0 (single query sees all cache)
+    //   batch decode (S>1, past>0):  S × (past+S) causal-with-past + sparse_mode=0
+    aclTensor* mask = nullptr;
+    int64_t sparse_mode = -1;  // auto
+    AclTensorPtr t_mask_ptr;
+    if (batch_decode_mode) {
+        build_batch_decode_mask_(S);
+        int64_t kv_len = past_len_ + S;
+        t_mask_ptr = make_contig_tensor(batch_mask_dev_.get(), ACL_BOOL, {1, 1, S, kv_len});
+        mask = t_mask_ptr.get();
+        sparse_mode = 0;
+    } else if (S > 1) {
+        // Pure prefill from past=0
+        t_mask_ptr = make_contig_tensor(prefill_mask_dev_.get(), ACL_BOOL, {1, 1, 2048, 2048});
+        mask = t_mask_ptr.get();
+        sparse_mode = 3;
+    }
+    // else: S=1 decode, mask=nullptr, sparse_mode=0 (auto)
+
+    attention_forward(
+        rt_.stream(), cfg_, attn_[layer_idx],
+        x_in, S, past_len_,
+        k_cache_[layer_idx].get(), v_cache_[layer_idx].get(), max_seq_,
+        mask,
+        q_sc_.get(), k_sc_.get(), v_sc_.get(),
+        xn_sc_.get(), rstd_sc_.get(), rope_sc_.get(),
+        attn_fias_sc_.get(),
+        attn_out_sc_.get(),
+        (hccl_ctx_.tp_size > 1) ? &hccl_ctx_ : nullptr,
+        &rope_cache_,
+        sparse_mode);
+
+    // x1 = x_in + attn_out (residual)
+    auto t_x_in    = make_contig_tensor(x_in,               ACL_BF16, {S, D});
+    auto t_attn_out= make_contig_tensor(attn_out_sc_.get(), ACL_BF16, {S, D});
+    auto t_x1      = make_contig_tensor(x_buf_a_.get(),     ACL_BF16, {S, D});
+    {
+        float a = 1.0f; aclScalar* al = aclCreateScalar(&a, ACL_FLOAT);
+        uint64_t ws = 0; aclOpExecutor* e = nullptr;
+        ACLNN_CHECK(aclnnAddGetWorkspaceSize(t_x_in.get(), t_attn_out.get(), al, t_x1.get(), &ws, &e));
+        DeviceBuffer wb; if (ws > 0) wb.alloc(ws);
+        ACLNN_CHECK(aclnnAdd(wb.get(), ws, e, rt_.stream()));
+        aclDestroyScalar(al);
+    }
+
+    // MoE
+    moe_forward(
+        rt_.stream(), cfg_, attn_[layer_idx], moe_[layer_idx],
+        x_buf_a_.get(), S,
+        moe_xn_.get(), moe_rstd_.get(),
+        moe_logits_.get(),
+        moe_topk_w_.get(), moe_topk_idx_.get(), moe_row_idx_.get(),
+        moe_ex_x_.get(), moe_ex_ri_.get(), moe_tpe_.get(),
+        moe_fwd_.get(),
+        moe_gate_.get(), moe_up_.get(), moe_down_.get(),
+        moe_packed_.get(), moe_weighted_.get(),
+        moe_out_.get(),
+        (hccl_ctx_.tp_size > 1) ? &hccl_ctx_ : nullptr,
+        moe_norm_sum_.get());
+
+    // x_out = x1 + moe_out (residual)
+    auto t_moe_out = make_contig_tensor(moe_out_.get(), ACL_BF16, {S, D});
+    auto t_out     = make_contig_tensor(x_out,          ACL_BF16, {S, D});
+    {
+        float a = 1.0f; aclScalar* al = aclCreateScalar(&a, ACL_FLOAT);
+        uint64_t ws = 0; aclOpExecutor* e = nullptr;
+        ACLNN_CHECK(aclnnAddGetWorkspaceSize(t_x1.get(), t_moe_out.get(), al, t_out.get(), &ws, &e));
+        DeviceBuffer wb; if (ws > 0) wb.alloc(ws);
+        ACLNN_CHECK(aclnnAdd(wb.get(), ws, e, rt_.stream()));
+        aclDestroyScalar(al);
+    }
+}
+
+void Runner::final_logits_(void* hidden_last, DeviceBuffer& logits_out) {
+    // Single-position variant: hidden_last is [1, D], output [1, V].
+    final_logits_batch_(hidden_last, 1, logits_out);
+}
+
+void Runner::final_logits_batch_(void* hidden, int64_t S, DeviceBuffer& logits_out) {
+    const int64_t D = cfg_.hidden_size;
+    const int64_t V = cfg_.vocab_size;
+
+    DeviceBuffer hn(S * D * 2), rstd(S * 4);
+    auto t_h   = make_contig_tensor(hidden,   ACL_BF16, {S, D});
+    auto t_hn  = make_contig_tensor(hn.get(), ACL_BF16, {S, D});
+    auto t_lnw = make_contig_tensor(shared_.final_norm.get(), ACL_BF16, {D});
+    auto t_rstd = make_contig_tensor(rstd.get(), ACL_FLOAT, {S});
+    rms_norm(rt_.stream(), t_h.get(), t_lnw.get(), cfg_.rms_norm_eps, t_hn.get(), t_rstd.get());
+
+    logits_out.alloc(S * V * 2);
+    auto t_logits = make_contig_tensor(logits_out.get(), ACL_BF16, {S, V});
+    linear_hf(rt_.stream(), t_hn.get(), shared_.lm_head.get(), ACL_BF16, V, D, t_logits.get());
+}
+
+bool Runner::decode_batch(const int32_t* tokens, int64_t S, DeviceBuffer& all_logits_out) {
+    if (S < 1) return false;
+    if (past_len_ + S > max_seq_) {
+        fprintf(stderr, "runner: decode_batch exceeds max_seq (%ld + %ld > %ld)\n",
+                past_len_, S, max_seq_);
+        return false;
+    }
+    const int64_t D = cfg_.hidden_size;
+
+    ensure_all_scratch_(this, S, cfg_,
+        q_sc_, k_sc_, v_sc_, xn_sc_, rstd_sc_, rope_sc_, attn_fias_sc_, attn_out_sc_,
+        moe_xn_, moe_rstd_, moe_logits_,
+        moe_topk_w_, moe_topk_idx_, moe_row_idx_,
+        moe_ex_x_, moe_ex_ri_, moe_tpe_,
+        moe_fwd_,
+        moe_gate_, moe_up_, moe_down_,
+        moe_packed_, moe_weighted_, moe_out_,
+        moe_norm_sum_,
+        x_buf_a_, x_buf_b_);
+
+    // Embed S tokens
+    DeviceBuffer tok_dev(S * 4);
+    ACL_CHECK(aclrtMemcpy(tok_dev.get(), S*4, tokens, S*4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tok = make_contig_tensor(tok_dev.get(), ACL_INT32, {S});
+    auto t_embed_w = make_contig_tensor(shared_.embed_tokens.get(), ACL_BF16, {cfg_.vocab_size, D});
+
+    DeviceBuffer x0(S * D * 2);
+    auto t_x0 = make_contig_tensor(x0.get(), ACL_BF16, {S, D});
+    index_select(rt_.stream(), t_embed_w.get(), 0, t_tok.get(), t_x0.get());
+
+    DeviceBuffer xping(S * D * 2), xpong(S * D * 2);
+    ACL_CHECK(aclrtMemcpyAsync(xping.get(), S*D*2, x0.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_DEVICE, rt_.stream()));
+    void* cur_in  = xping.get();
+    void* cur_out = xpong.get();
+    // batch_decode_mode=true uses proper causal-with-past mask (S × past+S, sparse_mode=0).
+    for (int L = 0; L < num_layers_; L++) {
+        layer_forward_(L, S, cur_in, cur_out, /*batch_decode_mode=*/past_len_ > 0);
+        std::swap(cur_in, cur_out);
+    }
+    rt_.sync();
+
+    // Get logits for ALL S positions (not just last)
+    final_logits_batch_(cur_in, S, all_logits_out);
+    rt_.sync();
+
+    past_len_ += S;
+    return true;
+}
+
+bool Runner::prefill(const int32_t* tokens, int64_t S, DeviceBuffer& logits_out) {
+    if (S < 1) return false;
+    if (past_len_ + S > max_seq_) {
+        fprintf(stderr, "runner: prefill exceeds max_seq (%ld + %ld > %ld)\n",
+                past_len_, S, max_seq_);
+        return false;
+    }
+
+    const int64_t D = cfg_.hidden_size;
+    ensure_all_scratch_(this, S, cfg_,
+        q_sc_, k_sc_, v_sc_, xn_sc_, rstd_sc_, rope_sc_, attn_fias_sc_, attn_out_sc_,
+        moe_xn_, moe_rstd_, moe_logits_,
+        moe_topk_w_, moe_topk_idx_, moe_row_idx_,
+        moe_ex_x_, moe_ex_ri_, moe_tpe_,
+        moe_fwd_,
+        moe_gate_, moe_up_, moe_down_,
+        moe_packed_, moe_weighted_, moe_out_,
+        moe_norm_sum_,
+        x_buf_a_, x_buf_b_);
+
+    // Embed
+    DeviceBuffer tok_dev(S * 4);
+    ACL_CHECK(aclrtMemcpy(tok_dev.get(), S*4, tokens, S*4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tok = make_contig_tensor(tok_dev.get(), ACL_INT32, {S});
+    auto t_embed_w = make_contig_tensor(shared_.embed_tokens.get(), ACL_BF16, {cfg_.vocab_size, D});
+
+    DeviceBuffer x0(S * D * 2);
+    auto t_x0 = make_contig_tensor(x0.get(), ACL_BF16, {S, D});
+    index_select(rt_.stream(), t_embed_w.get(), 0, t_tok.get(), t_x0.get());
+
+    // Layer chain: ping-pong between two buffers
+    DeviceBuffer xping(S * D * 2), xpong(S * D * 2);
+    ACL_CHECK(aclrtMemcpyAsync(xping.get(), S*D*2, x0.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_DEVICE, rt_.stream()));
+
+    void* cur_in  = xping.get();
+    void* cur_out = xpong.get();
+    for (int L = 0; L < num_layers_; L++) {
+        layer_forward_(L, S, cur_in, cur_out);
+        std::swap(cur_in, cur_out);
+    }
+    rt_.sync();
+
+    // Take last position's hidden → final_logits
+    DeviceBuffer last(1 * D * 2);
+    ACL_CHECK(aclrtMemcpy(last.get(), 1*D*2,
+                          (char*)cur_in + (S - 1) * D * 2, 1*D*2,
+                          ACL_MEMCPY_DEVICE_TO_DEVICE));
+    final_logits_(last.get(), logits_out);
+    rt_.sync();
+
+    past_len_ += S;
+    return true;
+}
+
+bool Runner::decode(int32_t token, DeviceBuffer& logits_out) {
+    const int64_t D = cfg_.hidden_size;
+    if (past_len_ + 1 > max_seq_) {
+        fprintf(stderr, "runner: decode exceeds max_seq\n");
+        return false;
+    }
+
+    const int64_t S = 1;
+    ensure_all_scratch_(this, S, cfg_,
+        q_sc_, k_sc_, v_sc_, xn_sc_, rstd_sc_, rope_sc_, attn_fias_sc_, attn_out_sc_,
+        moe_xn_, moe_rstd_, moe_logits_,
+        moe_topk_w_, moe_topk_idx_, moe_row_idx_,
+        moe_ex_x_, moe_ex_ri_, moe_tpe_,
+        moe_fwd_,
+        moe_gate_, moe_up_, moe_down_,
+        moe_packed_, moe_weighted_, moe_out_,
+        moe_norm_sum_,
+        x_buf_a_, x_buf_b_);
+
+    DeviceBuffer tok_dev(1 * 4);
+    ACL_CHECK(aclrtMemcpy(tok_dev.get(), 4, &token, 4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tok = make_contig_tensor(tok_dev.get(), ACL_INT32, {1});
+    auto t_embed_w = make_contig_tensor(shared_.embed_tokens.get(), ACL_BF16, {cfg_.vocab_size, D});
+
+    auto t0 = std::chrono::steady_clock::now();
+
+    DeviceBuffer x0(1 * D * 2);
+    auto t_x0 = make_contig_tensor(x0.get(), ACL_BF16, {1, D});
+    index_select(rt_.stream(), t_embed_w.get(), 0, t_tok.get(), t_x0.get());
+
+    DeviceBuffer xping(1 * D * 2), xpong(1 * D * 2);
+    ACL_CHECK(aclrtMemcpyAsync(xping.get(), 1*D*2, x0.get(), 1*D*2, ACL_MEMCPY_DEVICE_TO_DEVICE, rt_.stream()));
+    if (profile_enabled) { ACL_CHECK(aclrtSynchronizeStream(rt_.stream())); }
+    auto t1 = std::chrono::steady_clock::now();
+
+    void* cur_in  = xping.get();
+    void* cur_out = xpong.get();
+    for (int L = 0; L < num_layers_; L++) {
+        layer_forward_(L, 1, cur_in, cur_out);
+        std::swap(cur_in, cur_out);
+    }
+    rt_.sync();
+    auto t2 = std::chrono::steady_clock::now();
+
+    final_logits_(cur_in, logits_out);
+    rt_.sync();
+    auto t3 = std::chrono::steady_clock::now();
+
+    if (profile_enabled) {
+        using ms = std::chrono::duration<double, std::milli>;
+        t_embed_ms  += ms(t1 - t0).count();
+        t_layers_ms += ms(t2 - t1).count();
+        t_final_ms  += ms(t3 - t2).count();
+        profile_calls++;
+    }
+
+    past_len_ += 1;
+    return true;
+}
+
+void Runner::build_batch_decode_mask_(int64_t S) {
+    int64_t kv_len = past_len_ + S;
+    size_t bytes = (size_t)S * kv_len;   // bool = 1 byte
+    if (batch_mask_dev_.size < bytes) batch_mask_dev_.alloc(bytes);
+    std::vector<uint8_t> h_mask(bytes, 0);
+    for (int64_t i = 0; i < S; i++) {
+        // Row i: positions j ≤ past_len_+i are visible (0), j > past_len_+i are masked (1).
+        for (int64_t j = past_len_ + i + 1; j < kv_len; j++) {
+            h_mask[i * kv_len + j] = 1;
+        }
+    }
+    ACL_CHECK(aclrtMemcpy(batch_mask_dev_.get(), bytes, h_mask.data(), bytes,
+                          ACL_MEMCPY_HOST_TO_DEVICE));
+}
+
+void Runner::warmup(int iterations) {
+    if (num_layers_ == 0) return;
+    int64_t saved_past = past_len_;
+    past_len_ = 0;
+    int32_t dummy_tok = 0;      // token id 0, valid for Qwen3 (bos)
+    DeviceBuffer dummy_logits;
+    for (int i = 0; i < iterations; i++) {
+        past_len_ = 0;
+        if (!decode(dummy_tok, dummy_logits)) break;
+    }
+    past_len_ = saved_past;
+    fprintf(stderr, "[runner] warmup: %d iterations done\n", iterations);
+}
+
+void Runner::print_profile_summary() const {
+    if (!profile_enabled || profile_calls == 0) return;
+    double total = t_embed_ms + t_layers_ms + t_final_ms;
+    fprintf(stderr, "\n=== Runner profile (%ld decode calls) ===\n", profile_calls);
+    fprintf(stderr, "  phase        total_ms    avg_ms/call   pct\n");
+    fprintf(stderr, "  embed        %8.1f    %10.3f    %5.1f%%\n",
+            t_embed_ms,  t_embed_ms  / profile_calls, 100.0 * t_embed_ms  / total);
+    fprintf(stderr, "  layers (x%d) %8.1f    %10.3f    %5.1f%%  → %.3f ms/layer/call\n",
+            num_layers_, t_layers_ms, t_layers_ms / profile_calls,
+            100.0 * t_layers_ms / total,
+            t_layers_ms / profile_calls / num_layers_);
+    fprintf(stderr, "  final+lm_hd  %8.1f    %10.3f    %5.1f%%\n",
+            t_final_ms,  t_final_ms  / profile_calls, 100.0 * t_final_ms  / total);
+    fprintf(stderr, "  total        %8.1f    %10.3f   100.0%%\n",
+            total, total / profile_calls);
+}

src/safetensors_loader.cpp

@@ -0,0 +1,172 @@
+#include "safetensors_loader.h"
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <sstream>
+
+#include "json.hpp"
+
+using json = nlohmann::json;
+
+SafetensorsLoader::SafetensorsLoader() = default;
+
+SafetensorsLoader::~SafetensorsLoader() {
+    for (auto& s : shards_) {
+        if (s.mmap_ptr) munmap(s.mmap_ptr, s.mmap_size);
+        if (s.fd >= 0) close(s.fd);
+    }
+}
+
+bool SafetensorsLoader::open(const std::string& dir) {
+    model_dir_ = dir;
+
+    // 1. Parse index.json to discover shard files
+    std::string idx_path = dir + "/model.safetensors.index.json";
+    std::ifstream idx_file(idx_path);
+    if (!idx_file) {
+        // Fallback: single-file model
+        std::string single = dir + "/model.safetensors";
+        std::ifstream f(single);
+        if (!f) {
+            fprintf(stderr, "SafetensorsLoader: neither index.json nor model.safetensors found in %s\n", dir.c_str());
+            return false;
+        }
+        shards_.push_back({single});
+        return parse_shard_header_(0);
+    }
+
+    json idx;
+    try { idx_file >> idx; } catch (std::exception& e) {
+        fprintf(stderr, "SafetensorsLoader: bad index.json: %s\n", e.what());
+        return false;
+    }
+    if (!idx.contains("weight_map")) {
+        fprintf(stderr, "SafetensorsLoader: index.json missing weight_map\n");
+        return false;
+    }
+
+    // Collect unique shard filenames (preserving discovery order).
+    std::map<std::string, int> shard_name_to_id;
+    for (auto& [name, file] : idx["weight_map"].items()) {
+        std::string shard_name = file.get<std::string>();
+        if (shard_name_to_id.count(shard_name) == 0) {
+            int id = (int)shards_.size();
+            shard_name_to_id[shard_name] = id;
+            shards_.push_back({dir + "/" + shard_name});
+        }
+    }
+
+    // 2. Parse header of each shard to discover tensor offsets
+    for (int i = 0; i < (int)shards_.size(); i++) {
+        if (!parse_shard_header_(i)) {
+            fprintf(stderr, "SafetensorsLoader: failed to parse shard %s\n", shards_[i].path.c_str());
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool SafetensorsLoader::parse_shard_header_(int shard_id) {
+    ShardFile& sh = shards_[shard_id];
+    std::ifstream f(sh.path, std::ios::binary);
+    if (!f) return false;
+
+    // Read 8-byte little-endian header length
+    uint64_t header_len = 0;
+    f.read((char*)&header_len, 8);
+    if (!f) return false;
+
+    std::string header(header_len, '\0');
+    f.read(header.data(), header_len);
+    if (!f) return false;
+
+    sh.data_base = 8 + header_len;
+
+    json j;
+    try { j = json::parse(header); } catch (std::exception& e) {
+        fprintf(stderr, "SafetensorsLoader: bad shard header JSON in %s: %s\n", sh.path.c_str(), e.what());
+        return false;
+    }
+
+    for (auto it = j.begin(); it != j.end(); ++it) {
+        const std::string& name = it.key();
+        if (name == "__metadata__") continue;
+        const auto& entry = it.value();
+
+        TensorMeta m;
+        m.name  = name;
+        m.dtype = entry["dtype"].get<std::string>();
+        for (auto& d : entry["shape"]) m.shape.push_back(d.get<int64_t>());
+        const auto& offs = entry["data_offsets"];
+        size_t begin = offs[0].get<size_t>();
+        size_t end   = offs[1].get<size_t>();
+        m.offset = sh.data_base + begin;
+        m.nbytes = end - begin;
+        m.shard_id = shard_id;
+
+        tensors_[name] = std::move(m);
+    }
+
+    return true;
+}
+
+bool SafetensorsLoader::mmap_shard_(int shard_id) {
+    ShardFile& sh = shards_[shard_id];
+    if (sh.mmap_ptr) return true;
+
+    sh.fd = ::open(sh.path.c_str(), O_RDONLY);
+    if (sh.fd < 0) {
+        perror("open");
+        return false;
+    }
+    struct stat st;
+    if (fstat(sh.fd, &st) != 0) return false;
+    sh.mmap_size = st.st_size;
+
+    sh.mmap_ptr = mmap(nullptr, sh.mmap_size, PROT_READ, MAP_PRIVATE, sh.fd, 0);
+    if (sh.mmap_ptr == MAP_FAILED) {
+        perror("mmap");
+        sh.mmap_ptr = nullptr;
+        return false;
+    }
+    return true;
+}
+
+const TensorMeta* SafetensorsLoader::get(const std::string& name) const {
+    auto it = tensors_.find(name);
+    if (it == tensors_.end()) return nullptr;
+    return &it->second;
+}
+
+const void* SafetensorsLoader::data_ptr(const TensorMeta& m) {
+    if (m.shard_id < 0 || (size_t)m.shard_id >= shards_.size()) return nullptr;
+    if (!mmap_shard_(m.shard_id)) return nullptr;
+    ShardFile& sh = shards_[m.shard_id];
+    return (const char*)sh.mmap_ptr + m.offset;
+}
+
+const void* SafetensorsLoader::data_ptr(const std::string& name) {
+    const auto* m = get(name);
+    if (!m) return nullptr;
+    return data_ptr(*m);
+}
+
+std::vector<std::string> SafetensorsLoader::list_tensor_names() const {
+    std::vector<std::string> out;
+    out.reserve(tensors_.size());
+    for (auto& [k, v] : tensors_) out.push_back(k);
+    return out;
+}
+
+size_t SafetensorsLoader::total_bytes() const {
+    size_t sum = 0;
+    for (auto& [k, v] : tensors_) sum += v.nbytes;
+    return sum;
+}

src/tokenizer.cpp

@@ -0,0 +1,176 @@
+#include "tokenizer.h"
+
+#include <array>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <memory>
+#include <sstream>
+#include <unistd.h>
+
+bool Tokenizer::load(const std::string& vocab_bin_path) {
+    std::ifstream f(vocab_bin_path, std::ios::binary);
+    if (!f) {
+        fprintf(stderr, "Tokenizer: cannot open %s\n", vocab_bin_path.c_str());
+        return false;
+    }
+    uint32_t num;
+    f.read((char*)&num, 4);
+    if (!f) return false;
+    id_to_bytes_.resize(num);
+    for (uint32_t i = 0; i < num; i++) {
+        uint32_t len;
+        f.read((char*)&len, 4);
+        if (!f) return false;
+        id_to_bytes_[i].resize(len);
+        if (len > 0) f.read(id_to_bytes_[i].data(), len);
+    }
+    return true;
+}
+
+std::string Tokenizer::decode(int id) const {
+    if (id < 0 || (size_t)id >= id_to_bytes_.size()) return "";
+    return id_to_bytes_[id];
+}
+
+std::string Tokenizer::decode(const std::vector<int>& ids) const {
+    std::string out;
+    for (int id : ids) out += decode(id);
+    return out;
+}
+
+std::vector<int> Tokenizer::encode_via_python(const std::string& model_dir,
+                                              const std::string& prompt,
+                                              bool apply_chat_template) const {
+    // Call python subprocess to tokenize. Embed prompt via stdin to avoid shell-escape bugs.
+    std::string cmd;
+    // Set QWEN3_PYENV_INIT to override the Python env activation sequence (e.g., "source /opt/my_env/activate && ")
+    // Default assumes conda at ~/miniconda3 with env 'qwen3' and Ascend toolkit installed.
+    if (const char* init = std::getenv("QWEN3_PYENV_INIT")) {
+        cmd += init;
+    } else {
+        cmd += "source ${HOME}/miniconda3/etc/profile.d/conda.sh 2>/dev/null && ";
+        cmd += "conda activate qwen3 2>/dev/null || true; ";
+        cmd += "source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || true; ";
+    }
+    cmd += "python3 -c \"";
+    cmd += "import sys, json;";
+    cmd += "from transformers import AutoTokenizer;";
+    cmd += "t = AutoTokenizer.from_pretrained('" + model_dir + "');";
+    cmd += "p = sys.stdin.read();";
+    if (apply_chat_template) {
+        cmd += "msg = [{'role': 'user', 'content': p}];";
+        cmd += "ids = t.apply_chat_template(msg, add_generation_prompt=True);";
+    } else {
+        cmd += "ids = t.encode(p);";
+    }
+    cmd += "print(' '.join(str(i) for i in ids));";
+    cmd += "\"";
+
+    // popen with stdin: use the two-pipe dance via temp file for safety
+    char tmpl[] = "/tmp/lca_prompt_XXXXXX";
+    int fd = mkstemp(tmpl);
+    if (fd < 0) { perror("mkstemp"); return {}; }
+    write(fd, prompt.data(), prompt.size());
+    close(fd);
+
+    std::string full = cmd + " < " + tmpl + " 2>/dev/null";
+    FILE* pipe = popen(full.c_str(), "r");
+    if (!pipe) { perror("popen"); unlink(tmpl); return {}; }
+
+    std::string out;
+    char buf[4096];
+    while (size_t n = fread(buf, 1, sizeof(buf), pipe)) out.append(buf, n);
+    pclose(pipe);
+    unlink(tmpl);
+
+    std::vector<int> ids;
+    std::istringstream iss(out);
+    int x;
+    while (iss >> x) ids.push_back(x);
+    return ids;
+}
+
+// Shell-quote a string for embedding in a JSON string (escape ", \, control chars).
+static std::string json_escape(const std::string& s) {
+    std::string out;
+    out.reserve(s.size() + 8);
+    for (char c : s) {
+        switch (c) {
+            case '"':  out += "\\\""; break;
+            case '\\': out += "\\\\"; break;
+            case '\n': out += "\\n";  break;
+            case '\r': out += "\\r";  break;
+            case '\t': out += "\\t";  break;
+            default:
+                if ((unsigned char)c < 0x20) {
+                    char buf[8];
+                    snprintf(buf, sizeof(buf), "\\u%04x", (unsigned char)c);
+                    out += buf;
+                } else {
+                    out += c;
+                }
+        }
+    }
+    return out;
+}
+
+std::vector<int> Tokenizer::encode_conversation_via_python(
+    const std::string& model_dir,
+    const std::vector<std::pair<std::string, std::string>>& conversation,
+    bool add_generation_prompt) const
+{
+    // Build JSON array of messages. Pass via stdin to avoid shell-escape issues.
+    std::string json_msgs = "[";
+    for (size_t i = 0; i < conversation.size(); i++) {
+        if (i > 0) json_msgs += ",";
+        json_msgs += "{\"role\":\"" + json_escape(conversation[i].first) + "\",";
+        json_msgs += "\"content\":\"" + json_escape(conversation[i].second) + "\"}";
+    }
+    json_msgs += "]";
+
+    std::string cmd;
+    // Set QWEN3_PYENV_INIT to override the Python env activation sequence (e.g., "source /opt/my_env/activate && ")
+    // Default assumes conda at ~/miniconda3 with env 'qwen3' and Ascend toolkit installed.
+    if (const char* init = std::getenv("QWEN3_PYENV_INIT")) {
+        cmd += init;
+    } else {
+        cmd += "source ${HOME}/miniconda3/etc/profile.d/conda.sh 2>/dev/null && ";
+        cmd += "conda activate qwen3 2>/dev/null || true; ";
+        cmd += "source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || true; ";
+    }
+    cmd += "python3 -c \"";
+    cmd += "import sys, json;";
+    cmd += "from transformers import AutoTokenizer;";
+    cmd += "t = AutoTokenizer.from_pretrained('" + model_dir + "');";
+    cmd += "msgs = json.loads(sys.stdin.read());";
+    cmd += "ids = t.apply_chat_template(msgs, add_generation_prompt=";
+    cmd += add_generation_prompt ? "True" : "False";
+    cmd += ");";
+    cmd += "print(' '.join(str(i) for i in ids));";
+    cmd += "\"";
+
+    char tmpl[] = "/tmp/lca_conv_XXXXXX";
+    int fd = mkstemp(tmpl);
+    if (fd < 0) { perror("mkstemp"); return {}; }
+    write(fd, json_msgs.data(), json_msgs.size());
+    close(fd);
+
+    std::string full = cmd + " < " + tmpl + " 2>/dev/null";
+    FILE* pipe = popen(full.c_str(), "r");
+    if (!pipe) { perror("popen"); unlink(tmpl); return {}; }
+
+    std::string out;
+    char buf[4096];
+    while (size_t n = fread(buf, 1, sizeof(buf), pipe)) out.append(buf, n);
+    pclose(pipe);
+    unlink(tmpl);
+
+    std::vector<int> ids;
+    std::istringstream iss(out);
+    int x;
+    while (iss >> x) ids.push_back(x);
+    return ids;
+}

tests/hello_acl.cpp

@@ -0,0 +1,62 @@
+// hello_acl.cpp — smoke test: aclInit + device + stream + simple tensor + aclnnAdd
+#include "acl_common.h"
+#include <aclnnop/aclnn_add.h>
+#include <cstdio>
+#include <vector>
+
+int main() {
+    ACL_CHECK(aclInit(nullptr));
+    ACL_CHECK(aclrtSetDevice(0));
+    aclrtContext ctx;
+    ACL_CHECK(aclrtCreateContext(&ctx, 0));
+    aclrtStream stream;
+    ACL_CHECK(aclrtCreateStream(&stream));
+
+    // Tiny test: a = [1, 2, 3, 4] f32, b = [10, 20, 30, 40] f32, out = a + b
+    const int64_t N = 4;
+    std::vector<float> a_host = {1.0f, 2.0f, 3.0f, 4.0f};
+    std::vector<float> b_host = {10.0f, 20.0f, 30.0f, 40.0f};
+    std::vector<float> out_host(N, 0.0f);
+
+    DeviceBuffer a_dev(N * sizeof(float));
+    DeviceBuffer b_dev(N * sizeof(float));
+    DeviceBuffer out_dev(N * sizeof(float));
+
+    ACL_CHECK(aclrtMemcpy(a_dev.get(), N * 4, a_host.data(), N * 4, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(b_dev.get(), N * 4, b_host.data(), N * 4, ACL_MEMCPY_HOST_TO_DEVICE));
+
+    auto a_t   = make_contig_tensor(a_dev.get(),   ACL_FLOAT, {N});
+    auto b_t   = make_contig_tensor(b_dev.get(),   ACL_FLOAT, {N});
+    auto out_t = make_contig_tensor(out_dev.get(), ACL_FLOAT, {N});
+
+    // aclnnAdd: out = a + alpha * b
+    float alpha_val = 1.0f;
+    aclScalar* alpha = aclCreateScalar(&alpha_val, ACL_FLOAT);
+
+    uint64_t ws_size = 0;
+    aclOpExecutor* executor = nullptr;
+    ACLNN_CHECK(aclnnAddGetWorkspaceSize(a_t.get(), b_t.get(), alpha, out_t.get(), &ws_size, &executor));
+
+    DeviceBuffer ws;
+    if (ws_size > 0) ws.alloc(ws_size);
+    ACLNN_CHECK(aclnnAdd(ws.get(), ws_size, executor, stream));
+
+    ACL_CHECK(aclrtSynchronizeStream(stream));
+
+    ACL_CHECK(aclrtMemcpy(out_host.data(), N * 4, out_dev.get(), N * 4, ACL_MEMCPY_DEVICE_TO_HOST));
+
+    printf("hello_acl: ");
+    for (int i = 0; i < N; i++) printf("%.1f ", out_host[i]);
+    printf("\n");
+
+    bool ok = (out_host[0] == 11.0f && out_host[1] == 22.0f &&
+               out_host[2] == 33.0f && out_host[3] == 44.0f);
+    printf(ok ? "PASS\n" : "FAIL\n");
+
+    aclDestroyScalar(alpha);
+    ACL_CHECK(aclrtDestroyStream(stream));
+    ACL_CHECK(aclrtDestroyContext(ctx));
+    ACL_CHECK(aclrtResetDevice(0));
+    aclFinalize();
+    return ok ? 0 : 1;
+}

tests/test_attention_decode.cpp

@@ -0,0 +1,319 @@
+// test_attention_decode.cpp — validates single-layer attention with KV cache.
+//
+// Strategy: compare two paths yielding the same pos-5 attention output:
+//   Path A (reference): prefill 6 tokens in one shot → attn_out[5]
+//   Path B (decode):    prefill 5 tokens → K/V cache; decode 6th token via cache → attn_out_decode[0]
+//
+// The two outputs should match within BF16 precision.
+#include "acl_common.h"
+#include "acl_runtime.h"
+#include "aclnn_ops.h"
+#include "device_weights.h"
+#include "model_config.h"
+#include "rope.h"
+#include "safetensors_loader.h"
+
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <vector>
+
+static float bf16_to_float(uint16_t x) {
+    uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f;
+}
+static uint16_t float_to_bf16(float x) {
+    uint32_t u; std::memcpy(&u, &x, 4);
+    return (uint16_t)((u + 0x7FFF + ((u >> 16) & 1)) >> 16);
+}
+static std::vector<uint8_t> read_file(const std::string& p) {
+    std::ifstream f(p, std::ios::binary | std::ios::ate); size_t s = f.tellg();
+    f.seekg(0); std::vector<uint8_t> v(s); f.read((char*)v.data(), s); return v;
+}
+
+// Fill cos/sin tables for a range of positions [p0, p0+L). HF layout: half-half.
+static void fill_cos_sin(std::vector<uint16_t>& cos_h, std::vector<uint16_t>& sin_h,
+                         int64_t p0, int64_t L, int64_t Dh, float theta) {
+    cos_h.resize(L * Dh); sin_h.resize(L * Dh);
+    int64_t half = Dh / 2;
+    for (int64_t s = 0; s < L; s++) {
+        for (int64_t d = 0; d < Dh; d++) {
+            int64_t pair = (d < half) ? d : (d - half);
+            float theta_pair = 1.0f / std::pow(theta, (2.0f * pair) / Dh);
+            float angle = (float)(p0 + s) * theta_pair;
+            cos_h[s * Dh + d] = float_to_bf16(std::cos(angle));
+            sin_h[s * Dh + d] = float_to_bf16(std::sin(angle));
+        }
+    }
+}
+
+int main() {
+    const std::string model_dir = "/path/to/Qwen3-235B-A22B-Instruct-2507-BF16";
+    const std::string data_dir  = "tests/attn_data";
+
+    ModelConfig cfg;
+    if (!cfg.load_from_json(model_dir + "/config.json")) return 1;
+    cfg.compute_derived(1, 0);
+    const int64_t D      = cfg.hidden_size;
+    const int64_t Hq     = cfg.num_attention_heads;
+    const int64_t Hkv    = cfg.num_key_value_heads;
+    const int64_t Dh     = cfg.head_dim;
+    const int64_t Q_DIM  = Hq  * Dh;
+    const int64_t KV_DIM = Hkv * Dh;
+    const double  scale  = 1.0 / std::sqrt((double)Dh);
+    const double  eps    = cfg.rms_norm_eps;
+    const float   theta  = cfg.rope_theta;
+
+    SafetensorsLoader st;
+    if (!st.open(model_dir)) return 1;
+    AclRuntime rt;
+    rt.init(0);
+
+    DeviceWeightsLoader dw(st, cfg);
+    SharedWeights shared;
+    LayerAttnWeights attn;
+    printf("Loading weights...\n");
+    if (!dw.load_shared(shared))     return 1;
+    if (!dw.load_attention(0, attn)) return 1;
+
+    // ---- Load 5 prefill tokens + use token[5]=random as "6th" decoded token ----
+    auto tok_raw = read_file(data_dir + "/token_ids.bin");
+    int32_t S_prefill = *(int32_t*)tok_raw.data();
+    if (S_prefill < 5) { fprintf(stderr, "need >=5 tokens\n"); return 1; }
+    std::vector<int32_t> tokens(S_prefill);
+    std::memcpy(tokens.data(), tok_raw.data() + 4, S_prefill * 4);
+
+    // Build 6-token sequence (reuse first 5; pick a 6th token id — use token 0 as a simple choice)
+    const int64_t S6 = 6;
+    const int64_t S5 = 5;
+    std::vector<int32_t> tok6(S6);
+    for (int i = 0; i < S5; i++) tok6[i] = tokens[i];
+    tok6[5] = tokens[0];  // any token works for cross-consistency test
+    printf("tokens6=["); for (auto t : tok6) printf("%d,", t); printf("]\n");
+
+    // ---- Causal mask (2048x2048, sparse_mode=3) shared across both paths ----
+    const int64_t MASK = 2048;
+    DeviceBuffer mask_dev(MASK * MASK);
+    std::vector<uint8_t> mask_host(MASK * MASK, 0);
+    for (int i = 0; i < MASK; i++)
+        for (int j = i+1; j < MASK; j++)
+            mask_host[i*MASK + j] = 1;
+    ACL_CHECK(aclrtMemcpy(mask_dev.get(), MASK*MASK, mask_host.data(), MASK*MASK, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_mask = make_contig_tensor(mask_dev.get(), ACL_BOOL, {1, 1, MASK, MASK});
+
+    // =========================================================================
+    // PATH A: 6-token prefill (reference)
+    // =========================================================================
+    printf("\n[Path A] 6-token prefill reference\n");
+
+    DeviceBuffer tokA_dev(S6 * 4);
+    ACL_CHECK(aclrtMemcpy(tokA_dev.get(), S6*4, tok6.data(), S6*4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tokA = make_contig_tensor(tokA_dev.get(), ACL_INT32, {S6});
+    auto t_embed_w = make_contig_tensor(shared.embed_tokens.get(), ACL_BF16, {cfg.vocab_size, D});
+
+    DeviceBuffer xA_dev(S6 * D * 2);
+    auto t_xA = make_contig_tensor(xA_dev.get(), ACL_BF16, {S6, D});
+    index_select(rt.stream(), t_embed_w.get(), 0, t_tokA.get(), t_xA.get());
+    rt.sync();
+
+    DeviceBuffer xnA_dev(S6 * D * 2);
+    DeviceBuffer rstdA_dev(S6 * 4);
+    auto t_xnA   = make_contig_tensor(xnA_dev.get(),   ACL_BF16, {S6, D});
+    auto t_ln_w  = make_contig_tensor(attn.input_layernorm.get(), ACL_BF16, {D});
+    auto t_rstdA = make_contig_tensor(rstdA_dev.get(), ACL_FLOAT, {S6});
+    rms_norm(rt.stream(), t_xA.get(), t_ln_w.get(), eps, t_xnA.get(), t_rstdA.get());
+
+    DeviceBuffer qA_dev(S6 * Q_DIM  * 2);
+    DeviceBuffer kA_dev(S6 * KV_DIM * 2);
+    DeviceBuffer vA_dev(S6 * KV_DIM * 2);
+    auto t_qA = make_contig_tensor(qA_dev.get(), ACL_BF16, {S6, Q_DIM});
+    auto t_kA = make_contig_tensor(kA_dev.get(), ACL_BF16, {S6, KV_DIM});
+    auto t_vA = make_contig_tensor(vA_dev.get(), ACL_BF16, {S6, KV_DIM});
+    linear_hf(rt.stream(), t_xnA.get(), attn.q_proj.get(), ACL_BF16, Q_DIM,  D, t_qA.get());
+    linear_hf(rt.stream(), t_xnA.get(), attn.k_proj.get(), ACL_BF16, KV_DIM, D, t_kA.get());
+    linear_hf(rt.stream(), t_xnA.get(), attn.v_proj.get(), ACL_BF16, KV_DIM, D, t_vA.get());
+
+    // Per-head norm
+    auto t_qA_4d  = make_contig_tensor(qA_dev.get(), ACL_BF16, {1, S6, Hq,  Dh});
+    auto t_kA_4d  = make_contig_tensor(kA_dev.get(), ACL_BF16, {1, S6, Hkv, Dh});
+    auto t_qn_w   = make_contig_tensor(attn.q_norm.get(), ACL_BF16, {Dh});
+    auto t_kn_w   = make_contig_tensor(attn.k_norm.get(), ACL_BF16, {Dh});
+    DeviceBuffer rstd_qA(S6 * Hq  * 4), rstd_kA(S6 * Hkv * 4);
+    auto t_rstd_qA = make_contig_tensor(rstd_qA.get(), ACL_FLOAT, {1, S6, Hq});
+    auto t_rstd_kA = make_contig_tensor(rstd_kA.get(), ACL_FLOAT, {1, S6, Hkv});
+    rms_norm(rt.stream(), t_qA_4d.get(), t_qn_w.get(), eps, t_qA_4d.get(), t_rstd_qA.get());
+    rms_norm(rt.stream(), t_kA_4d.get(), t_kn_w.get(), eps, t_kA_4d.get(), t_rstd_kA.get());
+
+    // RoPE for positions 0..5
+    std::vector<uint16_t> cosA_h, sinA_h;
+    fill_cos_sin(cosA_h, sinA_h, 0, S6, Dh, theta);
+    DeviceBuffer cosA_dev(S6 * Dh * 2), sinA_dev(S6 * Dh * 2);
+    ACL_CHECK(aclrtMemcpy(cosA_dev.get(), S6*Dh*2, cosA_h.data(), S6*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(sinA_dev.get(), S6*Dh*2, sinA_h.data(), S6*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    DeviceBuffer ropeA_scratch(1 * S6 * Hq * Dh * 2);
+    apply_rope_manual(rt.stream(), qA_dev.get(), 1, S6, Hq, Dh, kA_dev.get(), Hkv,
+                      cosA_dev.get(), sinA_dev.get(), ropeA_scratch.get());
+
+    auto t_qA_bsh = make_contig_tensor(qA_dev.get(), ACL_BF16, {1, S6, Q_DIM});
+    auto t_kA_bsh = make_contig_tensor(kA_dev.get(), ACL_BF16, {1, S6, KV_DIM});
+    auto t_vA_bsh = make_contig_tensor(vA_dev.get(), ACL_BF16, {1, S6, KV_DIM});
+
+    DeviceBuffer attnA_out(1 * S6 * Q_DIM * 2);
+    auto t_attnA_out = make_contig_tensor(attnA_out.get(), ACL_BF16, {1, S6, Q_DIM});
+    fused_infer_attention_score(
+        rt.stream(), t_qA_bsh.get(), t_kA_bsh.get(), t_vA_bsh.get(),
+        t_mask.get(), {S6}, {S6}, Hq, Hkv, scale, 3, t_attnA_out.get());
+    rt.sync();
+
+    // Extract attnA_out[pos=5] into [1, 1, Q_DIM] for comparison
+    std::vector<uint16_t> refA_host(Q_DIM);
+    ACL_CHECK(aclrtMemcpy(refA_host.data(), Q_DIM*2,
+                          (char*)attnA_out.get() + 5 * Q_DIM * 2, Q_DIM*2,
+                          ACL_MEMCPY_DEVICE_TO_HOST));
+    printf("  attnA_out[5, :4] = %.5f %.5f %.5f %.5f\n",
+           bf16_to_float(refA_host[0]), bf16_to_float(refA_host[1]),
+           bf16_to_float(refA_host[2]), bf16_to_float(refA_host[3]));
+
+    // =========================================================================
+    // PATH B: 5-token prefill + KV cache → 1-token decode
+    // =========================================================================
+    printf("\n[Path B] 5-prefill + 1-decode via KV cache\n");
+
+    const int64_t MAX_LEN = 128;  // small cache for test
+    DeviceBuffer k_cache(MAX_LEN * KV_DIM * 2);
+    DeviceBuffer v_cache(MAX_LEN * KV_DIM * 2);
+    // Zero-init unused slots (not strictly needed, FIAS uses actual_seq_lens).
+
+    // ---- Prefill 5 tokens ----
+    DeviceBuffer tokB_dev(S5 * 4);
+    ACL_CHECK(aclrtMemcpy(tokB_dev.get(), S5*4, tok6.data(), S5*4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tokB = make_contig_tensor(tokB_dev.get(), ACL_INT32, {S5});
+    DeviceBuffer xB_dev(S5 * D * 2);
+    auto t_xB = make_contig_tensor(xB_dev.get(), ACL_BF16, {S5, D});
+    index_select(rt.stream(), t_embed_w.get(), 0, t_tokB.get(), t_xB.get());
+    rt.sync();
+
+    DeviceBuffer xnB_dev(S5 * D * 2);
+    DeviceBuffer rstdB_dev(S5 * 4);
+    auto t_xnB   = make_contig_tensor(xnB_dev.get(),   ACL_BF16, {S5, D});
+    auto t_rstdB = make_contig_tensor(rstdB_dev.get(), ACL_FLOAT, {S5});
+    rms_norm(rt.stream(), t_xB.get(), t_ln_w.get(), eps, t_xnB.get(), t_rstdB.get());
+
+    DeviceBuffer qB_dev(S5 * Q_DIM  * 2);
+    DeviceBuffer kB_dev(S5 * KV_DIM * 2);
+    DeviceBuffer vB_dev(S5 * KV_DIM * 2);
+    auto t_qB = make_contig_tensor(qB_dev.get(), ACL_BF16, {S5, Q_DIM});
+    auto t_kB = make_contig_tensor(kB_dev.get(), ACL_BF16, {S5, KV_DIM});
+    auto t_vB = make_contig_tensor(vB_dev.get(), ACL_BF16, {S5, KV_DIM});
+    linear_hf(rt.stream(), t_xnB.get(), attn.q_proj.get(), ACL_BF16, Q_DIM,  D, t_qB.get());
+    linear_hf(rt.stream(), t_xnB.get(), attn.k_proj.get(), ACL_BF16, KV_DIM, D, t_kB.get());
+    linear_hf(rt.stream(), t_xnB.get(), attn.v_proj.get(), ACL_BF16, KV_DIM, D, t_vB.get());
+
+    auto t_qB_4d  = make_contig_tensor(qB_dev.get(), ACL_BF16, {1, S5, Hq,  Dh});
+    auto t_kB_4d  = make_contig_tensor(kB_dev.get(), ACL_BF16, {1, S5, Hkv, Dh});
+    DeviceBuffer rstd_qB(S5 * Hq  * 4), rstd_kB(S5 * Hkv * 4);
+    auto t_rstd_qB = make_contig_tensor(rstd_qB.get(), ACL_FLOAT, {1, S5, Hq});
+    auto t_rstd_kB = make_contig_tensor(rstd_kB.get(), ACL_FLOAT, {1, S5, Hkv});
+    rms_norm(rt.stream(), t_qB_4d.get(), t_qn_w.get(), eps, t_qB_4d.get(), t_rstd_qB.get());
+    rms_norm(rt.stream(), t_kB_4d.get(), t_kn_w.get(), eps, t_kB_4d.get(), t_rstd_kB.get());
+
+    std::vector<uint16_t> cosB_h, sinB_h;
+    fill_cos_sin(cosB_h, sinB_h, 0, S5, Dh, theta);
+    DeviceBuffer cosB_dev(S5 * Dh * 2), sinB_dev(S5 * Dh * 2);
+    ACL_CHECK(aclrtMemcpy(cosB_dev.get(), S5*Dh*2, cosB_h.data(), S5*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(sinB_dev.get(), S5*Dh*2, sinB_h.data(), S5*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    DeviceBuffer ropeB_scratch(1 * S5 * Hq * Dh * 2);
+    apply_rope_manual(rt.stream(), qB_dev.get(), 1, S5, Hq, Dh, kB_dev.get(), Hkv,
+                      cosB_dev.get(), sinB_dev.get(), ropeB_scratch.get());
+    rt.sync();
+
+    // Append K, V to cache at positions 0..4.
+    ACL_CHECK(aclrtMemcpy(k_cache.get(), S5 * KV_DIM * 2,
+                          kB_dev.get(),  S5 * KV_DIM * 2, ACL_MEMCPY_DEVICE_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(v_cache.get(), S5 * KV_DIM * 2,
+                          vB_dev.get(),  S5 * KV_DIM * 2, ACL_MEMCPY_DEVICE_TO_DEVICE));
+    printf("  cached K/V at positions 0..%ld\n", S5 - 1);
+
+    // ---- Decode 1 token (position = 5) ----
+    DeviceBuffer tokD_dev(1 * 4);
+    int32_t tok_dec = tok6[5];
+    ACL_CHECK(aclrtMemcpy(tokD_dev.get(), 4, &tok_dec, 4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tokD = make_contig_tensor(tokD_dev.get(), ACL_INT32, {1});
+    DeviceBuffer xD_dev(1 * D * 2);
+    auto t_xD = make_contig_tensor(xD_dev.get(), ACL_BF16, {1, D});
+    index_select(rt.stream(), t_embed_w.get(), 0, t_tokD.get(), t_xD.get());
+
+    DeviceBuffer xnD_dev(1 * D * 2), rstdD_dev(1 * 4);
+    auto t_xnD   = make_contig_tensor(xnD_dev.get(), ACL_BF16, {1, D});
+    auto t_rstdD = make_contig_tensor(rstdD_dev.get(), ACL_FLOAT, {1});
+    rms_norm(rt.stream(), t_xD.get(), t_ln_w.get(), eps, t_xnD.get(), t_rstdD.get());
+
+    DeviceBuffer qD_dev(1 * Q_DIM  * 2), kD_dev(1 * KV_DIM * 2), vD_dev(1 * KV_DIM * 2);
+    auto t_qD = make_contig_tensor(qD_dev.get(), ACL_BF16, {1, Q_DIM});
+    auto t_kD = make_contig_tensor(kD_dev.get(), ACL_BF16, {1, KV_DIM});
+    auto t_vD = make_contig_tensor(vD_dev.get(), ACL_BF16, {1, KV_DIM});
+    linear_hf(rt.stream(), t_xnD.get(), attn.q_proj.get(), ACL_BF16, Q_DIM,  D, t_qD.get());
+    linear_hf(rt.stream(), t_xnD.get(), attn.k_proj.get(), ACL_BF16, KV_DIM, D, t_kD.get());
+    linear_hf(rt.stream(), t_xnD.get(), attn.v_proj.get(), ACL_BF16, KV_DIM, D, t_vD.get());
+
+    auto t_qD_4d  = make_contig_tensor(qD_dev.get(), ACL_BF16, {1, 1, Hq,  Dh});
+    auto t_kD_4d  = make_contig_tensor(kD_dev.get(), ACL_BF16, {1, 1, Hkv, Dh});
+    DeviceBuffer rstd_qD(1 * Hq  * 4), rstd_kD(1 * Hkv * 4);
+    auto t_rstd_qD = make_contig_tensor(rstd_qD.get(), ACL_FLOAT, {1, 1, Hq});
+    auto t_rstd_kD = make_contig_tensor(rstd_kD.get(), ACL_FLOAT, {1, 1, Hkv});
+    rms_norm(rt.stream(), t_qD_4d.get(), t_qn_w.get(), eps, t_qD_4d.get(), t_rstd_qD.get());
+    rms_norm(rt.stream(), t_kD_4d.get(), t_kn_w.get(), eps, t_kD_4d.get(), t_rstd_kD.get());
+
+    // RoPE for position 5 only
+    std::vector<uint16_t> cosD_h, sinD_h;
+    fill_cos_sin(cosD_h, sinD_h, /*p0=*/5, /*L=*/1, Dh, theta);
+    DeviceBuffer cosD_dev(1 * Dh * 2), sinD_dev(1 * Dh * 2);
+    ACL_CHECK(aclrtMemcpy(cosD_dev.get(), Dh*2, cosD_h.data(), Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(sinD_dev.get(), Dh*2, sinD_h.data(), Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    DeviceBuffer ropeD_scratch(1 * 1 * Hq * Dh * 2);
+    apply_rope_manual(rt.stream(), qD_dev.get(), 1, 1, Hq, Dh, kD_dev.get(), Hkv,
+                      cosD_dev.get(), sinD_dev.get(), ropeD_scratch.get());
+    rt.sync();
+
+    // Append K, V to cache at position 5.
+    ACL_CHECK(aclrtMemcpy((char*)k_cache.get() + S5 * KV_DIM * 2, KV_DIM * 2,
+                          kD_dev.get(), KV_DIM * 2, ACL_MEMCPY_DEVICE_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy((char*)v_cache.get() + S5 * KV_DIM * 2, KV_DIM * 2,
+                          vD_dev.get(), KV_DIM * 2, ACL_MEMCPY_DEVICE_TO_DEVICE));
+
+    // FIAS decode: q [1, 1, Q_DIM], k/v [1, 6, KV_DIM] from cache.
+    auto t_qD_bsh = make_contig_tensor(qD_dev.get(),  ACL_BF16, {1, 1,  Q_DIM});
+    auto t_kC_bsh = make_contig_tensor(k_cache.get(), ACL_BF16, {1, S6, KV_DIM});
+    auto t_vC_bsh = make_contig_tensor(v_cache.get(), ACL_BF16, {1, S6, KV_DIM});
+
+    DeviceBuffer attnD_out(1 * 1 * Q_DIM * 2);
+    auto t_attnD_out = make_contig_tensor(attnD_out.get(), ACL_BF16, {1, 1, Q_DIM});
+    // Decode: q has 1 token, k/v have 6 tokens. Use sparse_mode=0 with no mask — the single q
+    // at the end can attend to all cached positions; there's no causal constraint on it.
+    fused_infer_attention_score(
+        rt.stream(), t_qD_bsh.get(), t_kC_bsh.get(), t_vC_bsh.get(),
+        nullptr, {1}, {S6},
+        Hq, Hkv, scale, 0, t_attnD_out.get());
+    rt.sync();
+
+    std::vector<uint16_t> decB_host(Q_DIM);
+    ACL_CHECK(aclrtMemcpy(decB_host.data(), Q_DIM*2, attnD_out.get(), Q_DIM*2, ACL_MEMCPY_DEVICE_TO_HOST));
+
+    // ---- Compare Path A vs Path B ----
+    printf("\n  attnB_decode[:4] = %.5f %.5f %.5f %.5f\n",
+           bf16_to_float(decB_host[0]), bf16_to_float(decB_host[1]),
+           bf16_to_float(decB_host[2]), bf16_to_float(decB_host[3]));
+
+    double l2d = 0, l2r = 0, maxd = 0;
+    for (int i = 0; i < Q_DIM; i++) {
+        float a = bf16_to_float(decB_host[i]), b = bf16_to_float(refA_host[i]);
+        l2d += (a-b)*(a-b); l2r += b*b;
+        if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+    }
+    double rel = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+    printf("\nDecode vs 6-prefill comparison: rel=%.4e max_abs=%.4f\n", rel, maxd);
+
+    bool pass = rel < 5e-2;
+    printf("\n%s\n", pass ? "=== test_attention_decode PASS ===" : "=== test_attention_decode FAIL ===");
+    return pass ? 0 : 1;
+}

tests/test_attention_layer.cpp

@@ -0,0 +1,219 @@
+// test_attention_layer.cpp — full single-layer attention forward (Qwen3-235B layer 0), TP=1.
+// Validates C++ output against Python HF-style reference (attn_data/final_out.bin).
+#include "acl_common.h"
+#include "acl_runtime.h"
+#include "aclnn_ops.h"
+#include "device_weights.h"
+#include "model_config.h"
+#include "rope.h"
+#include "safetensors_loader.h"
+
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <vector>
+
+static float bf16_to_float(uint16_t x) {
+    uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f;
+}
+static uint16_t float_to_bf16(float x) {
+    uint32_t u; std::memcpy(&u, &x, 4);
+    return (uint16_t)((u + 0x7FFF + ((u >> 16) & 1)) >> 16);
+}
+
+static std::vector<uint8_t> read_file(const std::string& p) {
+    std::ifstream f(p, std::ios::binary | std::ios::ate); size_t s = f.tellg();
+    f.seekg(0); std::vector<uint8_t> v(s); f.read((char*)v.data(), s); return v;
+}
+
+int main() {
+    const std::string model_dir = "/path/to/Qwen3-235B-A22B-Instruct-2507-BF16";
+    const std::string data_dir  = "tests/attn_data";
+
+    ModelConfig cfg;
+    if (!cfg.load_from_json(model_dir + "/config.json")) return 1;
+    cfg.compute_derived(/*tp_size=*/1, /*tp_rank=*/0);   // single rank for correctness test
+    const int64_t D   = cfg.hidden_size;
+    const int64_t Hq  = cfg.num_attention_heads;
+    const int64_t Hkv = cfg.num_key_value_heads;
+    const int64_t Dh  = cfg.head_dim;
+    const int64_t Q_DIM  = Hq  * Dh;
+    const int64_t KV_DIM = Hkv * Dh;
+    const double  scale = 1.0 / std::sqrt((double)Dh);
+    const double  eps   = cfg.rms_norm_eps;
+    const float   theta = cfg.rope_theta;
+
+    SafetensorsLoader st;
+    if (!st.open(model_dir)) return 1;
+
+    AclRuntime rt;
+    rt.init(0);
+
+    // ---- Load weights (layer 0 attention + embed) ----
+    DeviceWeightsLoader dw(st, cfg);
+    SharedWeights shared;
+    LayerAttnWeights attn;
+    printf("Loading weights...\n");
+    if (!dw.load_shared(shared))         return 1;
+    if (!dw.load_attention(0, attn))     return 1;
+    printf("  shared.embed %.0fMB, attn total ~140MB\n", shared.embed_tokens.size / 1e6);
+
+    // ---- Load token ids (5 tokens: "The capital of France is") ----
+    auto tok_raw = read_file(data_dir + "/token_ids.bin");
+    int32_t S = *(int32_t*)tok_raw.data();
+    std::vector<int32_t> tokens(S);
+    std::memcpy(tokens.data(), tok_raw.data() + 4, S * 4);
+    printf("S=%d tokens=[", S); for (auto t : tokens) printf("%d,", t); printf("]\n");
+
+    // ---- Embed lookup: [S, D] ----
+    DeviceBuffer tok_dev(S * 4);
+    ACL_CHECK(aclrtMemcpy(tok_dev.get(), S * 4, tokens.data(), S * 4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tok = make_contig_tensor(tok_dev.get(), ACL_INT32, {S});
+
+    // embed weight shape [vocab, D]
+    auto t_embed_w = make_contig_tensor(shared.embed_tokens.get(), ACL_BF16, {cfg.vocab_size, D});
+
+    DeviceBuffer x_dev(S * D * 2);
+    auto t_x = make_contig_tensor(x_dev.get(), ACL_BF16, {S, D});
+    index_select(rt.stream(), t_embed_w.get(), 0, t_tok.get(), t_x.get());
+    rt.sync();
+
+    // ---- Residual snapshot (copy x) ----
+    DeviceBuffer residual_dev(S * D * 2);
+    ACL_CHECK(aclrtMemcpyAsync(residual_dev.get(), S*D*2, x_dev.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_DEVICE, rt.stream()));
+
+    // ---- Input layernorm ----
+    DeviceBuffer xn_dev(S * D * 2);
+    DeviceBuffer rstd_dev(S * 4);
+    auto t_xn   = make_contig_tensor(xn_dev.get(),   ACL_BF16, {S, D});
+    auto t_ln_w = make_contig_tensor(attn.input_layernorm.get(), ACL_BF16, {D});
+    auto t_rstd = make_contig_tensor(rstd_dev.get(), ACL_FLOAT, {S});
+    rms_norm(rt.stream(), t_x.get(), t_ln_w.get(), eps, t_xn.get(), t_rstd.get());
+
+    // ---- Q/K/V projections (linear_hf: y = x @ W.T, W stored as [out, in]) ----
+    DeviceBuffer q_dev(S * Q_DIM  * 2);
+    DeviceBuffer k_dev(S * KV_DIM * 2);
+    DeviceBuffer v_dev(S * KV_DIM * 2);
+    auto t_q = make_contig_tensor(q_dev.get(), ACL_BF16, {S, Q_DIM});
+    auto t_k = make_contig_tensor(k_dev.get(), ACL_BF16, {S, KV_DIM});
+    auto t_v = make_contig_tensor(v_dev.get(), ACL_BF16, {S, KV_DIM});
+    linear_hf(rt.stream(), t_xn.get(), attn.q_proj.get(), ACL_BF16, Q_DIM,  D, t_q.get());
+    linear_hf(rt.stream(), t_xn.get(), attn.k_proj.get(), ACL_BF16, KV_DIM, D, t_k.get());
+    linear_hf(rt.stream(), t_xn.get(), attn.v_proj.get(), ACL_BF16, KV_DIM, D, t_v.get());
+
+    // ---- Reshape Q, K as [B=1, S, N, Dh] for q_norm/k_norm + RoPE ----
+    // Same memory; just new views.
+    // q_dev has S * Q_DIM = S * Hq * Dh BF16
+    auto t_q_4d = make_contig_tensor(q_dev.get(), ACL_BF16, {1, S, Hq,  Dh});
+    auto t_k_4d = make_contig_tensor(k_dev.get(), ACL_BF16, {1, S, Hkv, Dh});
+
+    // Per-head RmsNorm on last dim (gamma shape [Dh])
+    auto t_qn_w = make_contig_tensor(attn.q_norm.get(), ACL_BF16, {Dh});
+    auto t_kn_w = make_contig_tensor(attn.k_norm.get(), ACL_BF16, {Dh});
+    DeviceBuffer rstd_q_dev(S * Hq  * 4);  // rstd shape = q's all-but-last dims
+    DeviceBuffer rstd_k_dev(S * Hkv * 4);
+    auto t_rstd_q = make_contig_tensor(rstd_q_dev.get(), ACL_FLOAT, {1, S, Hq});
+    auto t_rstd_k = make_contig_tensor(rstd_k_dev.get(), ACL_FLOAT, {1, S, Hkv});
+    // RmsNorm in place on q/k
+    rms_norm(rt.stream(), t_q_4d.get(), t_qn_w.get(), eps, t_q_4d.get(), t_rstd_q.get());
+    rms_norm(rt.stream(), t_k_4d.get(), t_kn_w.get(), eps, t_k_4d.get(), t_rstd_k.get());
+
+    // ---- Compute cos/sin on device ----
+    // cos/sin shape [1, S, Dh] BF16
+    std::vector<uint16_t> cos_host(S * Dh), sin_host(S * Dh);
+    for (int s = 0; s < S; s++) {
+        for (int64_t d = 0; d < Dh; d++) {
+            // freq index: for half-half layout, index d corresponds to pair index (d % (Dh/2))
+            int64_t half = Dh / 2;
+            int64_t pair = (d < half) ? d : (d - half);
+            float theta_pair = 1.0f / std::pow(theta, (2.0f * pair) / Dh);
+            float angle = (float)s * theta_pair;
+            cos_host[s * Dh + d] = float_to_bf16(std::cos(angle));
+            sin_host[s * Dh + d] = float_to_bf16(std::sin(angle));
+        }
+    }
+    DeviceBuffer cos_dev(S * Dh * 2);
+    DeviceBuffer sin_dev(S * Dh * 2);
+    ACL_CHECK(aclrtMemcpy(cos_dev.get(), S*Dh*2, cos_host.data(), S*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(sin_dev.get(), S*Dh*2, sin_host.data(), S*Dh*2, ACL_MEMCPY_HOST_TO_DEVICE));
+
+    // ---- RoPE ----
+    DeviceBuffer rope_scratch(1 * S * Hq * Dh * 2);
+    apply_rope_manual(rt.stream(),
+                      q_dev.get(), 1, S, Hq, Dh,
+                      k_dev.get(), Hkv,
+                      cos_dev.get(), sin_dev.get(),
+                      rope_scratch.get());
+
+    // ---- FIAS ----
+    // q/k/v are reshaped back to BSH [1, S, Hq*Dh or Hkv*Dh]
+    auto t_q_bsh = make_contig_tensor(q_dev.get(), ACL_BF16, {1, S, Q_DIM});
+    auto t_k_bsh = make_contig_tensor(k_dev.get(), ACL_BF16, {1, S, KV_DIM});
+    auto t_v_bsh = make_contig_tensor(v_dev.get(), ACL_BF16, {1, S, KV_DIM});
+
+    // Causal mask 2048x2048 (sparse_mode=3 requires fixed size)
+    const int64_t MASK = 2048;
+    DeviceBuffer mask_dev(MASK * MASK);  // bool = 1 byte
+    std::vector<uint8_t> mask_host(MASK * MASK, 0);
+    for (int i = 0; i < MASK; i++)
+        for (int j = i+1; j < MASK; j++)
+            mask_host[i*MASK + j] = 1;  // upper triangular = True
+    ACL_CHECK(aclrtMemcpy(mask_dev.get(), MASK*MASK, mask_host.data(), MASK*MASK, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_mask = make_contig_tensor(mask_dev.get(), ACL_BOOL, {1, 1, MASK, MASK});
+
+    DeviceBuffer attn_out_dev(1 * S * Q_DIM * 2);
+    auto t_attn_out = make_contig_tensor(attn_out_dev.get(), ACL_BF16, {1, S, Q_DIM});
+
+    fused_infer_attention_score(
+        rt.stream(),
+        t_q_bsh.get(), t_k_bsh.get(), t_v_bsh.get(),
+        t_mask.get(),
+        {S}, {S},
+        Hq, Hkv,
+        scale,
+        3,                  // sparse_mode = causal
+        t_attn_out.get());
+
+    // ---- O projection ----
+    auto t_attn_out_2d = make_contig_tensor(attn_out_dev.get(), ACL_BF16, {S, Q_DIM});
+    DeviceBuffer o_dev(S * D * 2);
+    auto t_o = make_contig_tensor(o_dev.get(), ACL_BF16, {S, D});
+    linear_hf(rt.stream(), t_attn_out_2d.get(), attn.o_proj.get(), ACL_BF16, D, Q_DIM, t_o.get());
+
+    // ---- Residual add: out = residual + o ----
+    auto t_res = make_contig_tensor(residual_dev.get(), ACL_BF16, {S, D});
+    float alpha_v = 1.0f;
+    aclScalar* alpha = aclCreateScalar(&alpha_v, ACL_FLOAT);
+    DeviceBuffer out_dev(S * D * 2);
+    auto t_out = make_contig_tensor(out_dev.get(), ACL_BF16, {S, D});
+    {
+        uint64_t ws = 0; aclOpExecutor* e = nullptr;
+        ACLNN_CHECK(aclnnAddGetWorkspaceSize(t_res.get(), t_o.get(), alpha, t_out.get(), &ws, &e));
+        DeviceBuffer wb; if (ws > 0) wb.alloc(ws);
+        ACLNN_CHECK(aclnnAdd(wb.get(), ws, e, rt.stream()));
+    }
+    aclDestroyScalar(alpha);
+    rt.sync();
+
+    // ---- Compare with Python reference ----
+    auto ref_h = read_file(data_dir + "/final_out.bin");
+    std::vector<uint16_t> cxx(S * D);
+    ACL_CHECK(aclrtMemcpy(cxx.data(), S*D*2, out_dev.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+    auto* ref = (const uint16_t*)ref_h.data();
+
+    double l2d = 0, l2r = 0, maxd = 0;
+    for (int i = 0; i < S * D; i++) {
+        float a = bf16_to_float(cxx[i]), b = bf16_to_float(ref[i]);
+        l2d += (a-b)*(a-b); l2r += b*b;
+        if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+    }
+    double rel = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+    printf("\nAttention layer output compare: rel=%.4e max_abs=%.4f\n", rel, maxd);
+    printf("  cxx[0, :4]: "); for (int i = 0; i < 4; i++) printf("%.6f ", bf16_to_float(cxx[i]));
+    printf("\n  ref[0, :4]: "); for (int i = 0; i < 4; i++) printf("%.6f ", bf16_to_float(ref[i])); printf("\n");
+
+    bool pass = rel < 5e-2;   // BF16 accumulation across 5+ ops loses ~1-2% per step
+    printf("\n%s\n", pass ? "=== test_attention_layer PASS ===" : "=== test_attention_layer FAIL ===");
+    return pass ? 0 : 1;
+}

tests/test_batch_correctness.cpp

@@ -0,0 +1,98 @@
+// test_batch_correctness.cpp — verify that forward with S>1 at past_len>0 produces the
+// same logits at each position as sequential S=1 decodes.
+//
+// This is the foundation for speculative decoding / PLD: the main model must predict logits
+// for each of K candidate positions in one batched forward pass matching sequential behavior.
+#include "runner.h"
+
+#include <cstdio>
+#include <cstring>
+#include <vector>
+#include <cmath>
+
+static float bf16_to_float(uint16_t x) {
+    uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f;
+}
+
+int main() {
+    const std::string model_dir = "/path/to/Qwen3-235B-A22B-Instruct-2507-BF16";
+    int tp_rank = 0, tp_size = 1;
+    if (const char* v = std::getenv("TP_RANK")) tp_rank = std::atoi(v);
+    if (const char* v = std::getenv("TP_SIZE")) tp_size = std::atoi(v);
+    bool is_master = tp_rank == 0;
+
+    Runner r;
+    if (!r.init(model_dir, tp_size, tp_rank, 94, 512)) return 1;
+    const int64_t V = r.cfg().vocab_size;
+
+    // Prefix
+    std::vector<int32_t> prompt = {785, 6722, 315, 9625, 374};
+    DeviceBuffer logits0;
+    r.prefill(prompt.data(), prompt.size(), logits0);
+    std::vector<uint16_t> h_last0(V);
+    if (is_master) ACL_CHECK(aclrtMemcpy(h_last0.data(), V*2, logits0.get(), V*2, ACL_MEMCPY_DEVICE_TO_HOST));
+    int next0 = 0;
+    if (is_master) {
+        float best = -1e30; for (int i = 0; i < V; i++) { float v = bf16_to_float(h_last0[i]); if (v > best) { best = v; next0 = i; } }
+    }
+    // Broadcast next0 (simple: let rank 0 decide and non-master ranks independently too)
+    int32_t token_seq[4];
+    if (is_master) token_seq[0] = next0;
+
+    // --- Path A: sequential S=1 decode × 4 times ---
+    std::vector<std::vector<uint16_t>> seq_logits(4);
+    for (int i = 0; i < 4; i++) seq_logits[i].resize(V);
+
+    // first decode: takes prompt's last logit argmax
+    // Here we need identical approach on all ranks. Use random token id for consistency.
+    std::vector<int32_t> seq_tokens = {next0, 100, 200, 300};  // deterministic for test
+
+    for (int i = 0; i < 4; i++) {
+        DeviceBuffer out;
+        r.decode(seq_tokens[i], out);
+        if (is_master) ACL_CHECK(aclrtMemcpy(seq_logits[i].data(), V*2, out.get(), V*2, ACL_MEMCPY_DEVICE_TO_HOST));
+    }
+    int64_t past_after_seq = r.past_len();
+
+    // --- Path B: reset, re-prefill, then ONE batch forward with S=4 ---
+    r.reset_cache();
+    DeviceBuffer logits_reprefill;
+    r.prefill(prompt.data(), prompt.size(), logits_reprefill);
+
+    DeviceBuffer batch_logits;
+    r.prefill(seq_tokens.data(), 4, batch_logits);
+    // prefill returns logits for LAST position only (S=4 gives [1, V], not [4, V]).
+    // Hmm — that's a limitation. To do PLD we need logits for all 4 positions.
+    // For now, just compare the LAST one (position 4 after prefix).
+
+    std::vector<uint16_t> batch_last(V);
+    if (is_master) ACL_CHECK(aclrtMemcpy(batch_last.data(), V*2, batch_logits.get(), V*2, ACL_MEMCPY_DEVICE_TO_HOST));
+
+    if (is_master) {
+        printf("\n=== Batched vs Sequential Decode Correctness ===\n");
+        double l2d=0, l2r=0, maxd=0;
+        for (int i = 0; i < V; i++) {
+            float a = bf16_to_float(batch_last[i]), b = bf16_to_float(seq_logits[3][i]);
+            l2d += (a-b)*(a-b); l2r += b*b;
+            if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+        }
+        double rel = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+        printf("Last-position logits:\n");
+        printf("  seq[3]  argmax = "); {
+            int b = 0; float bv = bf16_to_float(seq_logits[3][0]);
+            for (int i = 1; i < V; i++) if (bf16_to_float(seq_logits[3][i]) > bv) { bv = bf16_to_float(seq_logits[3][i]); b = i; }
+            printf("%d (%.3f)\n", b, bv);
+        }
+        printf("  batch   argmax = "); {
+            int b = 0; float bv = bf16_to_float(batch_last[0]);
+            for (int i = 1; i < V; i++) if (bf16_to_float(batch_last[i]) > bv) { bv = bf16_to_float(batch_last[i]); b = i; }
+            printf("%d (%.3f)\n", b, bv);
+        }
+        printf("  rel=%.4e  max=%.4f\n", rel, maxd);
+        printf("  %s\n", rel < 5e-2 ? "PASS" : "FAIL (batch forward diverges from sequential)");
+        printf("\nNote: current Runner.prefill() returns ONLY last-position logits. For PLD\n");
+        printf("we need all-position logits: requires extending prefill to optionally output\n");
+        printf("[S, V] logits tensor.\n");
+    }
+    return 0;
+}

tests/test_batch_decode.cpp

@@ -0,0 +1,85 @@
+// test_batch_decode.cpp — benchmark decode with different batch sizes S = 1, 2, 4, 8.
+//
+// Purpose: quantify the cost of "batched decode" (a.k.a. the ingredient speculative decoding
+// relies on). If Runner.prefill(S=K) forward-pass is only a small overhead over S=1, then
+// spec-decoding with K draft tokens gives ~K× speedup at high accept rate.
+//
+// Per-token amortized cost:
+//   cost(S) / S
+// Speculative decoding benefit:
+//   expected_accept_rate * K = effective tokens per forward
+//   → TG = expected / (cost(S=K+1) / 1_sec)
+#include "runner.h"
+
+#include <chrono>
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+int main() {
+    const std::string model_dir = "/path/to/Qwen3-235B-A22B-Instruct-2507-BF16";
+    Runner r;
+    int tp_rank = 0, tp_size = 1;
+    if (const char* v = std::getenv("TP_RANK")) tp_rank = std::atoi(v);
+    if (const char* v = std::getenv("TP_SIZE")) tp_size = std::atoi(v);
+    bool is_master = tp_rank == 0;
+
+    if (!r.init(model_dir, tp_size, tp_rank, /*num_layers=*/94, /*max_seq=*/512)) return 1;
+
+    // Prefill a short context so decode has some KV cache
+    std::vector<int32_t> prompt = {785, 6722, 315, 9625, 374};  // "The capital of France is"
+    DeviceBuffer logits;
+    r.prefill(prompt.data(), prompt.size(), logits);
+
+    auto now = []() { return std::chrono::steady_clock::now(); };
+    auto ms = [](auto t0, auto t1) { return std::chrono::duration<double, std::milli>(t1 - t0).count(); };
+
+    std::vector<int> batch_sizes = {1, 2, 4, 8};
+    int N_ITERS = 20;
+
+    if (is_master) {
+        printf("\n=== Batched decode forward benchmark (94 layers, TP=%d) ===\n", tp_size);
+        printf("Each row: forward with S=K new tokens after prefill\n");
+        printf("%-5s %-12s %-18s %-18s %s\n",
+               "S", "ms/forward", "ms/token (amort)", "tokens/sec", "vs S=1 efficiency");
+    }
+
+    double base_per_token = 0;
+    for (int S : batch_sizes) {
+        // Reset cache between measurements to keep cache size fair (same position for each)
+        // Actually we want to simulate: after some past_len, do 1 forward with S new tokens.
+        // Use prefill() which accepts S>=1.
+
+        std::vector<double> times;
+        for (int iter = 0; iter < N_ITERS + 3; iter++) {  // +3 for warmup
+            r.reset_cache();
+            r.prefill(prompt.data(), prompt.size(), logits);  // re-prefill
+
+            // New forward with S tokens (as if doing speculative verify)
+            std::vector<int32_t> new_tokens(S, 100);  // dummy token ids
+            auto t0 = now();
+            DeviceBuffer logits2;
+            r.prefill(new_tokens.data(), S, logits2);
+            auto t1 = now();
+            if (iter >= 3) times.push_back(ms(t0, t1));
+        }
+        std::sort(times.begin(), times.end());
+        double median_ms = times[times.size() / 2];
+        double per_token = median_ms / S;
+        double tok_per_sec = 1000.0 / per_token;
+        if (S == 1) base_per_token = per_token;
+        double efficiency = base_per_token / per_token * 100.0;
+
+        if (is_master) {
+            printf("%-5d %-12.2f %-18.2f %-18.2f %.1f%%\n",
+                   S, median_ms, per_token, tok_per_sec, efficiency);
+        }
+    }
+
+    if (is_master) {
+        printf("\n=== Interpretation ===\n");
+        printf("If S=4 forward ~ S=1 (efficiency high), spec decoding with accept_rate=70%%\n");
+        printf("gives TG = 0.7*4 / cost(S=5) vs baseline 1 / cost(S=1) = up to 2.8× speedup.\n");
+    }
+    return 0;
+}

tests/test_chat_flow.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+# test_chat_flow.sh — end-to-end integration smoke test for the CLI.
+#
+# Exercises:
+#   - --prompt-file
+#   - Multi-turn --chat memory (remembers Alice's name in turn 2)
+#   - --reset command in REPL
+#   - --system prompt
+#   - EOS detection at <|im_end|>
+#
+# Requires TP=16 Ascend 910 setup. Run from the repo root.
+#
+# Exit: 0 on all-pass, nonzero with reason.
+set -u
+BIN="./build/qwen3-moe-aclnn"
+MODEL="${MODEL_DIR:-/path/to/Qwen3-235B-A22B-Instruct-2507-BF16}"
+LAUNCH="./scripts/tp_launch.sh"
+TP="${TP_SIZE:-16}"
+VOCAB="tokenizer_data/vocab.bin"
+
+[ -x "$BIN" ] || { echo "FAIL: $BIN not built"; exit 1; }
+[ -x "$LAUNCH" ] || { echo "FAIL: $LAUNCH not found"; exit 1; }
+
+pass=0; fail=0
+check() {
+    local name="$1"; shift
+    local out="$1"; shift
+    local needle="$1"; shift
+    if echo "$out" | grep -qiF "$needle"; then
+        echo "  [PASS] $name (found: '$needle')"; pass=$((pass+1))
+    else
+        echo "  [FAIL] $name (did NOT find: '$needle')"; fail=$((fail+1))
+        echo "  ---- output ----"; echo "$out" | tail -20; echo "  ---- end ----"
+    fi
+}
+
+echo "===== Test 1: --prompt-file + EOS ====="
+echo "What is the capital of Japan?" > /tmp/chat_test_prompt.txt
+OUT=$(${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+    --prompt-file /tmp/chat_test_prompt.txt \
+    --chat --n-predict 50 --temperature 0 --vocab "$VOCAB" 2>&1)
+check "prompt-file loaded"        "$OUT" "capital of Japan"
+check "answer mentions Tokyo"     "$OUT" "Tokyo"
+check "hit EOS"                   "$OUT" "hit EOS"
+
+echo ""
+echo "===== Test 2: multi-turn memory (remembers name) ====="
+OUT=$(printf "My name is Alice.\nWhat is my name?\nquit\n" | \
+    ${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+    --interactive --chat \
+    --system "You are a concise assistant. Answer in one short sentence." \
+    --temperature 0 --n-predict 40 --max-seq 512 \
+    --vocab "$VOCAB" 2>&1)
+check "recalls Alice"             "$OUT" "Alice"
+check "has 2 turns"               "$OUT" "past_len="
+
+echo ""
+echo "===== Test 3: reset command clears memory ====="
+OUT=$(printf "My name is Bob.\nreset\nWhat is my name?\nquit\n" | \
+    ${LAUNCH} ${TP} ${BIN} --model-dir "$MODEL" \
+    --interactive --chat \
+    --system "Answer truthfully in one sentence." \
+    --temperature 0 --n-predict 40 --max-seq 512 \
+    --vocab "$VOCAB" 2>&1)
+check "reset acknowledged"        "$OUT" "cache + conversation reset"
+# After reset, model should NOT know the name is Bob (probably says "don't know" or asks)
+# We can't reliably check negation, so just check that the reset ran and turn 3 produced output
+check "turn 3 ran"                "$OUT" "bye"
+
+echo ""
+echo "===== Summary: $pass passed, $fail failed ====="
+exit $fail

tests/test_engine_smoke.cpp

@@ -0,0 +1,8 @@
+// test_engine_smoke.cpp — just verify engine.h compiles and links.
+#include "engine.h"
+
+int main() {
+    // No-op — all engine functions take many parameters and need real runtime. This test
+    // only validates that the header compiles and the core lib links.
+    return 0;
+}

tests/test_layer_forward.cpp

@@ -0,0 +1,192 @@
+// test_layer_forward.cpp — integration test for one full transformer layer via engine.h.
+//
+// Chain: embed_5_tokens → attention_forward (prefill, past=0) → +residual → moe_forward → +residual
+// Expected: final output matches moe_data/final_out.bin within BF16 precision (rel < 5e-2).
+#include "acl_common.h"
+#include "acl_runtime.h"
+#include "aclnn_ops.h"
+#include "device_weights.h"
+#include "engine.h"
+#include "model_config.h"
+#include "safetensors_loader.h"
+
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <vector>
+
+static float bf16_to_float(uint16_t x) {
+    uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f;
+}
+static std::vector<uint8_t> read_file(const std::string& p) {
+    std::ifstream f(p, std::ios::binary | std::ios::ate); size_t s = f.tellg();
+    f.seekg(0); std::vector<uint8_t> v(s); f.read((char*)v.data(), s); return v;
+}
+
+// Add: out = a + b (BF16).
+static void bf16_add(aclrtStream stream, aclTensor* a, aclTensor* b, aclTensor* out) {
+    float alpha = 1.0f; aclScalar* al = aclCreateScalar(&alpha, ACL_FLOAT);
+    uint64_t ws = 0; aclOpExecutor* e = nullptr;
+    ACLNN_CHECK(aclnnAddGetWorkspaceSize(a, b, al, out, &ws, &e));
+    DeviceBuffer wb; if (ws > 0) wb.alloc(ws);
+    ACLNN_CHECK(aclnnAdd(wb.get(), ws, e, stream));
+    aclDestroyScalar(al);
+}
+
+int main() {
+    const std::string model_dir = "/path/to/Qwen3-235B-A22B-Instruct-2507-BF16";
+    const std::string attn_data = "tests/attn_data";
+    const std::string moe_data  = "tests/moe_data";
+
+    ModelConfig cfg;
+    if (!cfg.load_from_json(model_dir + "/config.json")) return 1;
+    cfg.compute_derived(1, 0);
+    const int64_t D      = cfg.hidden_size;
+    const int64_t Hq     = cfg.n_heads_per_rank;
+    const int64_t Hkv    = cfg.n_kv_heads_per_rank;
+    const int64_t Dh     = cfg.head_dim;
+    const int64_t Q_DIM  = Hq * Dh;
+    const int64_t KV_DIM = Hkv * Dh;
+    const int64_t I      = cfg.i_per_rank;
+    const int64_t E      = cfg.num_experts;
+    const int64_t K      = cfg.num_experts_per_tok;
+    printf("Dims: D=%ld Q_DIM=%ld KV_DIM=%ld I=%ld E=%ld K=%ld\n", D, Q_DIM, KV_DIM, I, E, K);
+
+    SafetensorsLoader st;
+    if (!st.open(model_dir)) return 1;
+    AclRuntime rt;
+    rt.init(0);
+
+    DeviceWeightsLoader dw(st, cfg);
+    SharedWeights shared;
+    LayerAttnWeights attn;
+    LayerMoEWeights  moe;
+    printf("Loading weights...\n");
+    if (!dw.load_shared(shared))          return 1;
+    if (!dw.load_attention(0, attn))      return 1;
+    if (!dw.load_moe(0, rt.stream(), moe)) return 1;
+    rt.sync();
+
+    // ---- Load 5 prefill tokens ----
+    auto tok_raw = read_file(attn_data + "/token_ids.bin");
+    int32_t S = *(int32_t*)tok_raw.data();
+    std::vector<int32_t> tokens(S);
+    std::memcpy(tokens.data(), tok_raw.data() + 4, S * 4);
+    printf("S=%d tokens=[", S); for (auto t : tokens) printf("%d,", t); printf("]\n");
+
+    // ---- Embed ----
+    DeviceBuffer tok_dev(S * 4);
+    ACL_CHECK(aclrtMemcpy(tok_dev.get(), S * 4, tokens.data(), S * 4, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_tok = make_contig_tensor(tok_dev.get(), ACL_INT32, {S});
+    auto t_embed_w = make_contig_tensor(shared.embed_tokens.get(), ACL_BF16, {cfg.vocab_size, D});
+
+    DeviceBuffer x_dev(S * D * 2);          // residual / input to layer
+    auto t_x = make_contig_tensor(x_dev.get(), ACL_BF16, {S, D});
+    index_select(rt.stream(), t_embed_w.get(), 0, t_tok.get(), t_x.get());
+    rt.sync();
+
+    // ---- Scratch buffers for attention_forward ----
+    const int64_t MAX_LEN = 128;
+    DeviceBuffer k_cache(MAX_LEN * KV_DIM * 2), v_cache(MAX_LEN * KV_DIM * 2);
+    DeviceBuffer q_sc(S * Q_DIM * 2), k_sc(S * KV_DIM * 2), v_sc(S * KV_DIM * 2);
+    DeviceBuffer xn_sc(S * D * 2), rstd_sc(S * std::max(Hq, Hkv) * 4);
+    DeviceBuffer rope_sc(1 * S * Hq * Dh * 2);
+    DeviceBuffer attn_fias_sc(S * Q_DIM * 2);  // FIAS output buffer (before o_proj)
+    DeviceBuffer attn_out_dev(S * D * 2);
+
+    // ---- Causal mask (2048x2048) for prefill ----
+    const int64_t MASK = 2048;
+    DeviceBuffer mask_dev(MASK * MASK);
+    std::vector<uint8_t> mh(MASK * MASK, 0);
+    for (int i = 0; i < MASK; i++)
+        for (int j = i+1; j < MASK; j++) mh[i*MASK + j] = 1;
+    ACL_CHECK(aclrtMemcpy(mask_dev.get(), MASK*MASK, mh.data(), MASK*MASK, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto t_mask = make_contig_tensor(mask_dev.get(), ACL_BOOL, {1, 1, MASK, MASK});
+
+    // ---- Attention forward ----
+    attention_forward(
+        rt.stream(), cfg, attn,
+        x_dev.get(), S,
+        /*past_len=*/0, k_cache.get(), v_cache.get(), MAX_LEN,
+        t_mask.get(),
+        q_sc.get(), k_sc.get(), v_sc.get(),
+        xn_sc.get(), rstd_sc.get(), rope_sc.get(),
+        attn_fias_sc.get(),
+        attn_out_dev.get());
+    rt.sync();
+
+    // ---- x1 = x + attn_out (residual) — should match attn_data/final_out.bin ----
+    DeviceBuffer x1_dev(S * D * 2);
+    auto t_attn_out = make_contig_tensor(attn_out_dev.get(), ACL_BF16, {S, D});
+    auto t_x1       = make_contig_tensor(x1_dev.get(),       ACL_BF16, {S, D});
+    bf16_add(rt.stream(), t_x.get(), t_attn_out.get(), t_x1.get());
+    rt.sync();
+
+    auto attn_ref_h = read_file(attn_data + "/final_out.bin");
+    std::vector<uint16_t> x1_host(S * D);
+    ACL_CHECK(aclrtMemcpy(x1_host.data(), S*D*2, x1_dev.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+    auto* ar = (const uint16_t*)attn_ref_h.data();
+    double al2d=0, al2r=0, amaxd=0;
+    for (int i = 0; i < S*D; i++) {
+        float a = bf16_to_float(x1_host[i]), b = bf16_to_float(ar[i]);
+        al2d += (a-b)*(a-b); al2r += b*b;
+        if (std::abs(a-b) > amaxd) amaxd = std::abs(a-b);
+    }
+    double arel = std::sqrt(al2d) / (std::sqrt(al2r) + 1e-10);
+    printf("  [attn] x + attn_out vs attn_data/final_out.bin:  rel=%.4e max=%.4f\n", arel, amaxd);
+
+    // ---- MoE scratch buffers ----
+    const int64_t TOTAL = S * K;
+    DeviceBuffer moe_xn(S * D * 2), moe_rstd(S * 4);
+    DeviceBuffer moe_logits(S * E * 2);
+    DeviceBuffer moe_topk_w(S * K * 2), moe_topk_idx(S * K * 4), moe_row_idx(S * K * 4);
+    DeviceBuffer moe_ex_x(TOTAL * D * 2), moe_ex_ri(TOTAL * 4), moe_tpe(E * 8);
+    DeviceBuffer moe_fwd(TOTAL * 8);
+    DeviceBuffer moe_gate(TOTAL * I * 2), moe_up(TOTAL * I * 2), moe_down(TOTAL * D * 2);
+    DeviceBuffer moe_packed(TOTAL * D * 2), moe_weighted(S * K * D * 2);
+    DeviceBuffer moe_out_dev(S * D * 2);
+
+    moe_forward(rt.stream(), cfg, attn, moe,
+                x1_dev.get(), S,
+                moe_xn.get(), moe_rstd.get(),
+                moe_logits.get(),
+                moe_topk_w.get(), moe_topk_idx.get(), moe_row_idx.get(),
+                moe_ex_x.get(), moe_ex_ri.get(), moe_tpe.get(),
+                moe_fwd.get(),
+                moe_gate.get(), moe_up.get(), moe_down.get(),
+                moe_packed.get(), moe_weighted.get(),
+                moe_out_dev.get());
+    rt.sync();
+
+    // ---- x2 = x1 + moe_out (residual) — should match moe_data/final_out.bin ----
+    DeviceBuffer x2_dev(S * D * 2);
+    auto t_moe_out = make_contig_tensor(moe_out_dev.get(), ACL_BF16, {S, D});
+    auto t_x2      = make_contig_tensor(x2_dev.get(),      ACL_BF16, {S, D});
+    bf16_add(rt.stream(), t_x1.get(), t_moe_out.get(), t_x2.get());
+    rt.sync();
+
+    auto moe_ref_h = read_file(moe_data + "/final_out.bin");
+    std::vector<uint16_t> x2_host(S * D);
+    ACL_CHECK(aclrtMemcpy(x2_host.data(), S*D*2, x2_dev.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+    auto* mr = (const uint16_t*)moe_ref_h.data();
+    double ml2d=0, ml2r=0, mmaxd=0;
+    for (int i = 0; i < S*D; i++) {
+        float a = bf16_to_float(x2_host[i]), b = bf16_to_float(mr[i]);
+        ml2d += (a-b)*(a-b); ml2r += b*b;
+        if (std::abs(a-b) > mmaxd) mmaxd = std::abs(a-b);
+    }
+    double mrel = std::sqrt(ml2d) / (std::sqrt(ml2r) + 1e-10);
+    printf("  [full] x1 + moe_out vs moe_data/final_out.bin:   rel=%.4e max=%.4f\n", mrel, mmaxd);
+    printf("    x2[0, :4]:  %.5f %.5f %.5f %.5f\n",
+           bf16_to_float(x2_host[0]), bf16_to_float(x2_host[1]), bf16_to_float(x2_host[2]), bf16_to_float(x2_host[3]));
+    printf("    ref[0, :4]: %.5f %.5f %.5f %.5f\n",
+           bf16_to_float(mr[0]), bf16_to_float(mr[1]), bf16_to_float(mr[2]), bf16_to_float(mr[3]));
+
+    // Tolerance: attn chain 5e-3 (tight, only linear ops); full layer 1e-1 (MoE's discrete topk
+    // routing amplifies BF16 noise — tiny input changes flip expert selection, magnifying output
+    // delta. End-to-end CLI correctness is validated by test_chat_flow.sh separately.)
+    bool pass = (arel < 5e-3) && (mrel < 1e-1);
+    printf("\n%s\n", pass ? "=== test_layer_forward PASS ===" : "=== test_layer_forward FAIL ===");
+    return pass ? 0 : 1;
+}

tests/test_linear_hf.cpp

@@ -0,0 +1,73 @@
+// test_linear_hf.cpp — verify linear_hf (y = x @ W.T with HF [out, in] layout).
+#include "acl_common.h"
+#include "acl_runtime.h"
+#include "aclnn_ops.h"
+
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <vector>
+
+static float bf16_to_float(uint16_t x) {
+    uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f;
+}
+
+int main() {
+    const std::string data = "tests/mm_data";
+    int64_t N = 0, D = 0, OUT = 0;
+    {
+        std::ifstream f(data + "/shape.txt"); std::string line;
+        while (std::getline(f, line)) {
+            auto eq = line.find('='); if (eq == std::string::npos) continue;
+            auto k = line.substr(0, eq); auto v = std::atoll(line.c_str() + eq + 1);
+            if (k == "N") N = v; else if (k == "D") D = v; else if (k == "OUT") OUT = v;
+        }
+    }
+    printf("N=%ld D=%ld OUT=%ld\n", N, D, OUT);
+
+    AclRuntime rt;
+    rt.init(0);
+
+    auto read_all = [&](const std::string& p) {
+        std::ifstream f(p, std::ios::binary | std::ios::ate); size_t sz = f.tellg();
+        f.seekg(0); std::vector<uint8_t> v(sz); f.read((char*)v.data(), sz); return v;
+    };
+    auto x_h  = read_all(data + "/x.bin");
+    auto W_h  = read_all(data + "/W.bin");
+    auto yr_h = read_all(data + "/y_ref.bin");
+
+    DeviceBuffer x_d(N * D * 2);
+    DeviceBuffer W_d(OUT * D * 2);
+    DeviceBuffer y_d(N * OUT * 2);
+    ACL_CHECK(aclrtMemcpy(x_d.get(), x_h.size(), x_h.data(), x_h.size(), ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(W_d.get(), W_h.size(), W_h.data(), W_h.size(), ACL_MEMCPY_HOST_TO_DEVICE));
+
+    auto t_x = make_contig_tensor(x_d.get(), ACL_BF16, {N, D});
+    auto t_y = make_contig_tensor(y_d.get(), ACL_BF16, {N, OUT});
+
+    linear_hf(rt.stream(), t_x.get(), W_d.get(), ACL_BF16, OUT, D, t_y.get());
+    rt.sync();
+
+    std::vector<uint16_t> y_cxx(N * OUT);
+    ACL_CHECK(aclrtMemcpy(y_cxx.data(), N * OUT * 2, y_d.get(), N * OUT * 2, ACL_MEMCPY_DEVICE_TO_HOST));
+    auto* y_ref = (const uint16_t*)yr_h.data();
+
+    double l2d = 0, l2r = 0, maxd = 0;
+    for (int i = 0; i < N * OUT; i++) {
+        float a = bf16_to_float(y_cxx[i]);
+        float b = bf16_to_float(y_ref[i]);
+        l2d += (a-b)*(a-b); l2r += b*b;
+        if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+    }
+    double rel = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+    printf("L2 diff=%.4f ref=%.4f relative=%.4e max_abs=%.4f\n",
+           std::sqrt(l2d), std::sqrt(l2r), rel, maxd);
+    printf("y_cxx[0..3]: "); for (int i = 0; i < 4; i++) printf("%.3f ", bf16_to_float(y_cxx[i])); printf("\n");
+    printf("y_ref[0..3]: "); for (int i = 0; i < 4; i++) printf("%.3f ", bf16_to_float(y_ref[i])); printf("\n");
+
+    // BF16 matmul has more precision loss than RmsNorm. Allow 1% relative error.
+    bool ok = rel < 1e-2;
+    printf("\n%s\n", ok ? "=== test_linear_hf PASS ===" : "=== test_linear_hf FAIL ===");
+    return ok ? 0 : 1;
+}

tests/test_model_config.cpp

@@ -0,0 +1,106 @@
+// test_model_config.cpp — load config.json, derive TP shard sizes, verify all expected
+// HF tensors exist in safetensors for Qwen3-235B.
+#include "model_config.h"
+#include "safetensors_loader.h"
+
+#include <cstdio>
+#include <sstream>
+#include <string>
+
+int main(int argc, char** argv) {
+    std::string dir = argc > 1 ? argv[1]
+                               : "/path/to/Qwen3-235B-A22B-Instruct-2507-BF16";
+    int tp_size = argc > 2 ? std::atoi(argv[2]) : 16;
+    int tp_rank = argc > 3 ? std::atoi(argv[3]) : 0;
+
+    ModelConfig cfg;
+    if (!cfg.load_from_json(dir + "/config.json")) return 1;
+    cfg.compute_derived(tp_size, tp_rank);
+    printf("%s\n", cfg.describe().c_str());
+
+    SafetensorsLoader loader;
+    if (!loader.open(dir)) return 1;
+
+    // Verify all expected tensor names & shapes match cfg.
+    int missing = 0, shape_mismatch = 0;
+    auto check_shape = [&](const std::string& name, const std::vector<int64_t>& expected) {
+        auto* m = loader.get(name);
+        if (!m) {
+            printf("  MISSING: %s\n", name.c_str());
+            missing++;
+            return;
+        }
+        if (m->shape != expected) {
+            printf("  SHAPE MISMATCH: %s got=[", name.c_str());
+            for (size_t i = 0; i < m->shape.size(); i++) printf("%s%ld", i ? "," : "", m->shape[i]);
+            printf("] want=[");
+            for (size_t i = 0; i < expected.size(); i++) printf("%s%ld", i ? "," : "", expected[i]);
+            printf("]\n");
+            shape_mismatch++;
+        }
+    };
+
+    // embed/head
+    check_shape("model.embed_tokens.weight", {cfg.vocab_size, cfg.hidden_size});
+    check_shape("lm_head.weight",            {cfg.vocab_size, cfg.hidden_size});
+    check_shape("model.norm.weight",         {cfg.hidden_size});
+
+    // Attention weights (HF stores as [out, in])
+    int64_t q_full = cfg.num_attention_heads * cfg.head_dim;
+    int64_t kv_full = cfg.num_key_value_heads * cfg.head_dim;
+    for (int L = 0; L < cfg.num_hidden_layers; L++) {
+        auto base = "model.layers." + std::to_string(L);
+        check_shape(base + ".input_layernorm.weight",                   {cfg.hidden_size});
+        check_shape(base + ".post_attention_layernorm.weight",          {cfg.hidden_size});
+        check_shape(base + ".self_attn.q_proj.weight",                  {q_full,  cfg.hidden_size});
+        check_shape(base + ".self_attn.k_proj.weight",                  {kv_full, cfg.hidden_size});
+        check_shape(base + ".self_attn.v_proj.weight",                  {kv_full, cfg.hidden_size});
+        check_shape(base + ".self_attn.o_proj.weight",                  {cfg.hidden_size, q_full});
+        // Qwen3 uses q_norm / k_norm (norm per head) — check existence
+        check_shape(base + ".self_attn.q_norm.weight",                  {cfg.head_dim});
+        check_shape(base + ".self_attn.k_norm.weight",                  {cfg.head_dim});
+        // MoE router
+        check_shape(base + ".mlp.gate.weight",                          {cfg.num_experts, cfg.hidden_size});
+        // Spot-check few experts (full enumeration is 94*384=36096 lines)
+        for (int e : {0, 1, 63, 127}) {
+            auto ebase = base + ".mlp.experts." + std::to_string(e);
+            check_shape(ebase + ".gate_proj.weight", {cfg.moe_intermediate_size, cfg.hidden_size});
+            check_shape(ebase + ".up_proj.weight",   {cfg.moe_intermediate_size, cfg.hidden_size});
+            check_shape(ebase + ".down_proj.weight", {cfg.hidden_size, cfg.moe_intermediate_size});
+        }
+    }
+
+    // Print TP memory estimate
+    int64_t attn_bytes_per_rank = 0;
+    attn_bytes_per_rank += cfg.q_dim_per_rank     * cfg.hidden_size   * 2;   // q_proj
+    attn_bytes_per_rank += cfg.kv_dim_per_rank    * cfg.hidden_size   * 2;   // k_proj
+    attn_bytes_per_rank += cfg.kv_dim_per_rank    * cfg.hidden_size   * 2;   // v_proj
+    attn_bytes_per_rank += cfg.hidden_size        * cfg.q_dim_per_rank * 2;  // o_proj
+    attn_bytes_per_rank *= cfg.num_hidden_layers;
+
+    int64_t moe_bytes_per_rank = 0;
+    // gate_exps + up_exps: [E, I_per_rank, D]
+    moe_bytes_per_rank += 2 * cfg.num_experts * cfg.i_per_rank * cfg.hidden_size * 2;
+    // down_exps: [E, D, I_per_rank]
+    moe_bytes_per_rank += cfg.num_experts * cfg.hidden_size * cfg.i_per_rank * 2;
+    moe_bytes_per_rank *= cfg.num_hidden_layers;
+
+    int64_t embed_bytes = cfg.vocab_size * cfg.hidden_size * 2 * 2;  // embed + lm_head
+    int64_t router_bytes = cfg.num_experts * cfg.hidden_size * 2 * cfg.num_hidden_layers;
+    int64_t norm_bytes = cfg.hidden_size * 2 * (2 * cfg.num_hidden_layers + 1);
+    int64_t total_per_rank = attn_bytes_per_rank + moe_bytes_per_rank + embed_bytes + router_bytes + norm_bytes;
+
+    printf("\nPer-rank weight memory estimate (BF16, TP=%d):\n", tp_size);
+    printf("  attention: %.2f GB\n", attn_bytes_per_rank / 1e9);
+    printf("  MoE exps:  %.2f GB\n", moe_bytes_per_rank / 1e9);
+    printf("  embed+head: %.2f GB (replicated)\n", embed_bytes / 1e9);
+    printf("  router:    %.2f MB (replicated)\n", router_bytes / 1e6);
+    printf("  norms:     %.2f MB (replicated)\n", norm_bytes / 1e6);
+    printf("  TOTAL:     %.2f GB\n", total_per_rank / 1e9);
+
+    int errors = missing + shape_mismatch;
+    printf("\nMissing: %d, Shape mismatch: %d\n", missing, shape_mismatch);
+    printf("%s\n", errors == 0 ? "=== test_model_config PASS ==="
+                               : "=== test_model_config FAIL ===");
+    return errors == 0 ? 0 : 1;
+}

tests/test_moe_layer.cpp

@@ -0,0 +1,676 @@
+// test_moe_layer.cpp — Full MoE layer forward (Qwen3-235B layer 0), TP=1.
+//
+// Pipeline:
+//   1. Post-attention RmsNorm (input from attn_data/final_out.bin)
+//   2. Router: xn @ W_router.T → logits [S, E]
+//   3. TopK softmax → weights [S, K], expert_ids [S, K]
+//   4. Host-normalize top_k weights (Qwen3 norm_topk_prob)
+//   5. MoeInitRoutingV3 → expanded_x [S*K, D], expanded_row_idx, tokens_per_expert
+//   6. GMM gate: expanded_x × gate_exps → [S*K, I]
+//   7. GMM up:   same → [S*K, I]
+//   8. silu(gate) * up → [S*K, I]
+//   9. GMM down: act × down_exps → [S*K, D]
+//  10. MoeFinalizeRouting (weighted sum) → [S, D]
+//  11. + residual
+#include "acl_common.h"
+#include "acl_runtime.h"
+#include "aclnn_ops.h"
+#include "device_weights.h"
+#include "model_config.h"
+#include "safetensors_loader.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+#include <fstream>
+#include <tuple>
+#include <vector>
+
+static float bf16_to_float(uint16_t x) {
+    uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f;
+}
+static uint16_t float_to_bf16(float x) {
+    uint32_t u; std::memcpy(&u, &x, 4);
+    return (uint16_t)((u + 0x7FFF + ((u >> 16) & 1)) >> 16);
+}
+static std::vector<uint8_t> read_file(const std::string& p) {
+    std::ifstream f(p, std::ios::binary | std::ios::ate); size_t s = f.tellg();
+    f.seekg(0); std::vector<uint8_t> v(s); f.read((char*)v.data(), s); return v;
+}
+
+int main() {
+    const std::string model_dir = "/path/to/Qwen3-235B-A22B-Instruct-2507-BF16";
+    const std::string data_dir  = "tests/moe_data";
+
+    ModelConfig cfg;
+    if (!cfg.load_from_json(model_dir + "/config.json")) return 1;
+    cfg.compute_derived(1, 0);  // TP=1
+    const int64_t D = cfg.hidden_size;
+    const int64_t I = cfg.moe_intermediate_size;
+    const int64_t E = cfg.num_experts;
+    const int64_t K = cfg.num_experts_per_tok;
+    const double eps = cfg.rms_norm_eps;
+
+    AclRuntime rt;
+    rt.init(0);
+    printf("[dbg] rt init ok\n"); fflush(stdout);
+
+    SafetensorsLoader st;
+    if (!st.open(model_dir)) return 1;
+
+    // ---- Load weights ----
+    printf("Loading layer 0 attention weights (for post_attention_layernorm)...\n");
+    DeviceWeightsLoader dw(st, cfg);
+    LayerAttnWeights attn;
+    if (!dw.load_attention(0, attn)) return 1;
+
+    printf("Loading layer 0 MoE weights (128 experts × 3 projections, stacking + permute)...\n"); fflush(stdout);
+    LayerMoEWeights moe;
+    if (!dw.load_moe(0, rt.stream(), moe)) return 1;
+    rt.sync();
+    printf("[dbg] moe load ok\n"); fflush(stdout);
+    printf("  router     %.1f MB  gate_exps %.0f MB  up_exps %.0f MB  down_exps %.0f MB\n",
+           moe.router.size / 1e6, moe.gate_exps.size / 1e6, moe.up_exps.size / 1e6, moe.down_exps.size / 1e6);
+
+    // ---- Load input & Python reference ----
+    int S = 5;
+    auto x_in_host = read_file(data_dir + "/x_in.bin");
+    auto ref_out_host = read_file(data_dir + "/final_out.bin");
+    DeviceBuffer x_dev(S * D * 2);
+    ACL_CHECK(aclrtMemcpy(x_dev.get(), x_in_host.size(), x_in_host.data(), x_in_host.size(), ACL_MEMCPY_HOST_TO_DEVICE));
+
+    // Residual snapshot
+    DeviceBuffer residual_dev(S * D * 2);
+    ACL_CHECK(aclrtMemcpy(residual_dev.get(), S*D*2, x_dev.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_DEVICE));
+
+    printf("[dbg] loaded data and residual ok, TOTAL=%ld\n", S * K); fflush(stdout);
+
+    // ---- Step 1: Post-attention RmsNorm ----
+    DeviceBuffer xn_dev(S * D * 2);
+    DeviceBuffer rstd_dev(S * 4);
+    auto t_x  = make_contig_tensor(x_dev.get(),  ACL_BF16, {S, D});
+    auto t_xn = make_contig_tensor(xn_dev.get(), ACL_BF16, {S, D});
+    auto t_ln = make_contig_tensor(attn.post_attention_layernorm.get(), ACL_BF16, {D});
+    auto t_rstd = make_contig_tensor(rstd_dev.get(), ACL_FLOAT, {S});
+    rms_norm(rt.stream(), t_x.get(), t_ln.get(), eps, t_xn.get(), t_rstd.get());
+    rt.sync();
+    printf("[dbg] rms_norm ok\n"); fflush(stdout);
+
+    // ---- Step 2: Router (gate matmul) ----
+    DeviceBuffer logits_dev(S * E * 2);
+    auto t_logits = make_contig_tensor(logits_dev.get(), ACL_BF16, {S, E});
+    // router is [E, D] (HF). logits = xn @ router.T
+    linear_hf(rt.stream(), t_xn.get(), moe.router.get(), ACL_BF16, E, D, t_logits.get());
+    rt.sync();
+    printf("[dbg] router linear ok\n"); fflush(stdout);
+
+    // ---- Step 3: TopK softmax ----
+    DeviceBuffer topk_w_dev(S * K * 2);          // BF16
+    DeviceBuffer topk_idx_dev(S * K * 4);        // int32
+    DeviceBuffer row_idx_dev(S * K * 4);         // int32 (from gating op, unused for our routing)
+    auto t_topk_w   = make_contig_tensor(topk_w_dev.get(),   ACL_BF16,  {S, K});
+    auto t_topk_idx = make_contig_tensor(topk_idx_dev.get(), ACL_INT32, {S, K});
+    auto t_row_idx  = make_contig_tensor(row_idx_dev.get(),  ACL_INT32, {S, K});
+    moe_gating_topk_softmax(rt.stream(), t_logits.get(), K, t_topk_w.get(), t_topk_idx.get(), t_row_idx.get());
+    rt.sync();
+    printf("[dbg] topk_softmax ok\n"); fflush(stdout);
+
+    // ---- Step 4: Host-normalize top_k weights (norm_topk_prob=true) ----
+    std::vector<uint16_t> tw_bf(S * K);
+    ACL_CHECK(aclrtMemcpy(tw_bf.data(), S*K*2, topk_w_dev.get(), S*K*2, ACL_MEMCPY_DEVICE_TO_HOST));
+    for (int s = 0; s < S; s++) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) sum += bf16_to_float(tw_bf[s*K + k]);
+        sum += 1e-20f;
+        for (int k = 0; k < K; k++) {
+            float v = bf16_to_float(tw_bf[s*K + k]) / sum;
+            tw_bf[s*K + k] = float_to_bf16(v);
+        }
+    }
+    ACL_CHECK(aclrtMemcpy(topk_w_dev.get(), S*K*2, tw_bf.data(), S*K*2, ACL_MEMCPY_HOST_TO_DEVICE));
+
+    // ---- Step 5: MoE init routing ----
+    int64_t TOTAL = S * K;
+    DeviceBuffer expanded_x_dev(TOTAL * D * 2);
+    DeviceBuffer expanded_row_idx_dev(TOTAL * 4);
+    DeviceBuffer tokens_per_expert_dev(E * 8);
+
+    auto t_ex_x   = make_contig_tensor(expanded_x_dev.get(),        ACL_BF16,  {TOTAL, D});
+    auto t_ex_ri  = make_contig_tensor(expanded_row_idx_dev.get(),  ACL_INT32, {TOTAL});
+    auto t_tpe    = make_contig_tensor(tokens_per_expert_dev.get(), ACL_INT64, {E});
+
+    moe_init_routing_v3(rt.stream(),
+                        t_xn.get(), t_topk_idx.get(),
+                        E, TOTAL,
+                        t_ex_x.get(), t_ex_ri.get(), t_tpe.get());
+    rt.sync();
+    printf("[dbg] moe_init_routing ok\n"); fflush(stdout);
+
+    // Convert tokens_per_expert from counts to cumsum (on host) for GMM groupListType=0.
+    DeviceBuffer tpe_cumsum_dev(E * 8);
+    {
+        std::vector<int64_t> h_counts(E), h_cum(E);
+        ACL_CHECK(aclrtMemcpy(h_counts.data(), E*8, tokens_per_expert_dev.get(), E*8, ACL_MEMCPY_DEVICE_TO_HOST));
+        int64_t acc = 0;
+        for (int i = 0; i < E; i++) { acc += h_counts[i]; h_cum[i] = acc; }
+        ACL_CHECK(aclrtMemcpy(tpe_cumsum_dev.get(), E*8, h_cum.data(), E*8, ACL_MEMCPY_HOST_TO_DEVICE));
+    }
+    auto t_tpe_cum = make_contig_tensor(tpe_cumsum_dev.get(), ACL_INT64, {E});
+
+    // ---- Step 6/7: GMM gate and up ----
+    DeviceBuffer gate_out_dev(TOTAL * I * 2);
+    DeviceBuffer up_out_dev(TOTAL * I * 2);
+    auto t_gate_out = make_contig_tensor(gate_out_dev.get(), ACL_BF16, {TOTAL, I});
+    auto t_up_out   = make_contig_tensor(up_out_dev.get(),   ACL_BF16, {TOTAL, I});
+    // gate/up_exps loaded as [E, D, I] row-major
+    auto t_w_gate = make_contig_tensor(moe.gate_exps.get(), ACL_BF16, {E, D, I});
+    auto t_w_up   = make_contig_tensor(moe.up_exps.get(),   ACL_BF16, {E, D, I});
+    // Use cumsum group_list (groupListType=0): empirically more reliable with many zero-count experts.
+    grouped_matmul_v4(rt.stream(), t_ex_x.get(), t_w_gate.get(), t_tpe_cum.get(), t_gate_out.get(), 0);
+    rt.sync();
+    printf("[dbg] gmm gate ok\n"); fflush(stdout);
+    grouped_matmul_v4(rt.stream(), t_ex_x.get(), t_w_up.get(),   t_tpe_cum.get(), t_up_out.get(), 0);
+    rt.sync();
+    printf("[dbg] gmm up ok\n"); fflush(stdout);
+
+    // ---- Step 8: SwiGLU ----
+    // act = silu(gate) * up  (inplace on gate_out)
+    silu(rt.stream(), t_gate_out.get(), t_gate_out.get());
+    rt.sync(); printf("[dbg] silu ok\n"); fflush(stdout);
+    mul(rt.stream(), t_gate_out.get(), t_up_out.get(), t_gate_out.get());
+    rt.sync(); printf("[dbg] mul ok\n"); fflush(stdout);
+    // now gate_out_dev contains the activated intermediate
+
+    // ---- Step 9: GMM down ----
+    DeviceBuffer down_out_dev(TOTAL * D * 2);
+    auto t_down_out = make_contig_tensor(down_out_dev.get(), ACL_BF16, {TOTAL, D});
+    auto t_w_down   = make_contig_tensor(moe.down_exps.get(), ACL_BF16, {E, I, D});
+    grouped_matmul_v4(rt.stream(), t_gate_out.get(), t_w_down.get(), t_tpe_cum.get(), t_down_out.get(), 0);
+    rt.sync();
+    printf("[dbg] gmm down ok\n"); fflush(stdout);
+
+    // ---- Step 10: Device-side manual finalize (replacement for buggy MoeFinalizeRoutingV2) ----
+    // Compute forward permutation fwd[n*K + k] = p where token n's k-th expert's output is at
+    // expanded position p. We use tokens_per_expert (cumsum) + topk_idx to resolve this correctly,
+    // regardless of the exact rowIdxType semantics returned by MoeInitRoutingV3.
+    DeviceBuffer fwd_dev(TOTAL * 8);
+    {
+        std::vector<int64_t> h_tpe2(E);
+        std::vector<int32_t> h_tidx3(S * K);
+        ACL_CHECK(aclrtMemcpy(h_tpe2.data(),  E*8,   tokens_per_expert_dev.get(), E*8,   ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_tidx3.data(), S*K*4, topk_idx_dev.get(),          S*K*4, ACL_MEMCPY_DEVICE_TO_HOST));
+
+        // Sort (n, k) pairs by expert ascending (stable). For each expert in order, tokens
+        // appear in ascending token index (since MoeInitRoutingV3 is stable by s).
+        // Specifically: expanded positions 0..tpe[0]-1 are for expert 0 (tokens picking e=0, in n-ascending order),
+        // next tpe[1] are for expert 1, etc.
+        //
+        // To build fwd: for each (n, k), expert e = topk_idx[n, k]. Position p is the base of expert e's
+        // block plus the rank of n within tokens picking e.
+        std::vector<int64_t> expert_base(E + 1, 0);
+        for (int e = 0; e < E; e++) expert_base[e + 1] = expert_base[e] + h_tpe2[e];
+
+        std::vector<int> expert_slot(E, 0);  // next available slot per expert
+        std::vector<int64_t> fwd(TOTAL);
+        // Iterate in token-ascending, k-ascending order — match MoeInitRoutingV3's stable sort convention.
+        // For each (n, k) sorted by (expert[n,k], n), assign p.
+        // Simpler: pre-collect (e, n, k) triples, sort by (e, n), then p is the rank.
+        std::vector<std::tuple<int, int, int>> triples;
+        triples.reserve(TOTAL);
+        for (int n = 0; n < S; n++) for (int k = 0; k < K; k++) {
+            triples.emplace_back(h_tidx3[n * K + k], n, k);
+        }
+        std::sort(triples.begin(), triples.end(), [](const auto& a, const auto& b){
+            if (std::get<0>(a) != std::get<0>(b)) return std::get<0>(a) < std::get<0>(b);
+            return std::get<1>(a) < std::get<1>(b);
+        });
+        for (int64_t p = 0; p < TOTAL; p++) {
+            auto [e, n, k] = triples[p];
+            fwd[n * K + k] = p;
+        }
+        ACL_CHECK(aclrtMemcpy(fwd_dev.get(), TOTAL*8, fwd.data(), TOTAL*8, ACL_MEMCPY_HOST_TO_DEVICE));
+    }
+    auto t_fwd = make_contig_tensor(fwd_dev.get(), ACL_INT64, {TOTAL});
+
+    // Gather: packed [S*K, D] = down_out[fwd, :]
+    DeviceBuffer packed_dev(TOTAL * D * 2);
+    auto t_packed = make_contig_tensor(packed_dev.get(), ACL_BF16, {TOTAL, D});
+    index_select(rt.stream(), t_down_out.get(), 0, t_fwd.get(), t_packed.get());
+    rt.sync();
+
+    // Broadcast-multiply by topk_w: view packed as [S, K, D], topk_w as [S, K, 1].
+    auto t_packed_3d = make_contig_tensor(packed_dev.get(), ACL_BF16, {S, K, D});
+    auto t_topk_w_3d = make_contig_tensor(topk_w_dev.get(), ACL_BF16, {S, K, 1});
+    DeviceBuffer weighted_dev(S * K * D * 2);
+    auto t_weighted = make_contig_tensor(weighted_dev.get(), ACL_BF16, {S, K, D});
+    mul(rt.stream(), t_packed_3d.get(), t_topk_w_3d.get(), t_weighted.get());
+    rt.sync();
+
+    // Verify broadcast mul + sum by dumping all k entries and summing on host.
+    {
+        std::vector<uint16_t> h_pk_all(S * K * D);
+        std::vector<uint16_t> h_wt_all(S * K * D);
+        std::vector<uint16_t> h_tw_all(S * K);
+        ACL_CHECK(aclrtMemcpy(h_pk_all.data(), S*K*D*2, packed_dev.get(), S*K*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_wt_all.data(), S*K*D*2, weighted_dev.get(), S*K*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_tw_all.data(), S*K*2, topk_w_dev.get(), S*K*2, ACL_MEMCPY_DEVICE_TO_HOST));
+
+        printf("  verify weighted[0, k, 0] = packed[0, k, 0] * topk_w[0, k] for all k:\n");
+        float host_sum = 0;
+        for (int k = 0; k < K; k++) {
+            float p = bf16_to_float(h_pk_all[k * D]);          // packed[0, k, 0] = offset s*K*D + k*D + 0 = k*D (for s=0)
+            float w = bf16_to_float(h_tw_all[k]);               // topk_w[0, k]
+            float wt = bf16_to_float(h_wt_all[k * D]);          // weighted[0, k, 0]
+            host_sum += p * w;
+            printf("    k=%d: packed=%.5f * topk_w=%.5f = expect=%.5f  dev=%.5f\n",
+                   k, p, w, p*w, wt);
+        }
+        printf("  host_sum_of_weighted[0, :, 0] = %.5f  (expected moe_out[0,0] = -0.02466)\n", host_sum);
+    }
+
+    // ReduceSum over K axis → [S, D]
+    DeviceBuffer moe_out_dev(S * D * 2);
+    auto t_moe_out = make_contig_tensor(moe_out_dev.get(), ACL_BF16, {S, D});
+    reduce_sum(rt.stream(), t_weighted.get(), {1}, /*keep_dims=*/false, ACL_BF16, t_moe_out.get());
+    rt.sync();
+    printf("[dbg] device-side finalize (gather+mul+reduce) ok\n"); fflush(stdout);
+
+    // Residual add to produce final_out
+    float alpha_v = 1.0f; aclScalar* alpha = aclCreateScalar(&alpha_v, ACL_FLOAT);
+    DeviceBuffer final_dev(S * D * 2);
+    auto t_final = make_contig_tensor(final_dev.get(), ACL_BF16, {S, D});
+    auto t_res   = make_contig_tensor(residual_dev.get(), ACL_BF16, {S, D});
+    {
+        uint64_t ws = 0; aclOpExecutor* e = nullptr;
+        ACLNN_CHECK(aclnnAddGetWorkspaceSize(t_res.get(), t_moe_out.get(), alpha, t_final.get(), &ws, &e));
+        DeviceBuffer wb; if (ws > 0) wb.alloc(ws);
+        ACLNN_CHECK(aclnnAdd(wb.get(), ws, e, rt.stream()));
+    }
+    aclDestroyScalar(alpha);
+    rt.sync();
+
+    // ---- Compare (intermediate + final) ----
+    auto compare_bf16 = [&](const char* label, void* dev_ptr, int64_t nelem,
+                             const std::string& ref_file) {
+        std::vector<uint16_t> cxx(nelem);
+        ACL_CHECK(aclrtMemcpy(cxx.data(), nelem*2, dev_ptr, nelem*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        auto refbuf = read_file(data_dir + "/" + ref_file);
+        auto* ref = (const uint16_t*)refbuf.data();
+        double l2d = 0, l2r = 0, maxd = 0;
+        for (int64_t i = 0; i < nelem; i++) {
+            float a = bf16_to_float(cxx[i]), b = bf16_to_float(ref[i]);
+            l2d += (a-b)*(a-b); l2r += b*b;
+            if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+        }
+        double rel = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+        printf("  [cmp] %-12s rel=%.4e max_abs=%.4f   cxx[:4]=%.5f %.5f %.5f %.5f  ref[:4]=%.5f %.5f %.5f %.5f\n",
+               label, rel, maxd,
+               bf16_to_float(cxx[0]), bf16_to_float(cxx[1]), bf16_to_float(cxx[2]), bf16_to_float(cxx[3]),
+               bf16_to_float(ref[0]), bf16_to_float(ref[1]), bf16_to_float(ref[2]), bf16_to_float(ref[3]));
+        return rel;
+    };
+
+    printf("\n=== Intermediate diagnostics ===\n");
+    compare_bf16("xn",       xn_dev.get(),         S * D,   "xn.bin");
+    compare_bf16("topk_w",   topk_w_dev.get(),     S * K,   "topk_w.bin");
+
+    // Dump topk_idx (int32) to compare
+    {
+        std::vector<int32_t> cxx_idx(S*K);
+        ACL_CHECK(aclrtMemcpy(cxx_idx.data(), S*K*4, topk_idx_dev.get(), S*K*4, ACL_MEMCPY_DEVICE_TO_HOST));
+        auto refbuf = read_file(data_dir + "/topk_idx.bin");
+        auto* ref = (const int32_t*)refbuf.data();
+        int mismatches = 0;
+        for (int i = 0; i < S*K; i++) if (cxx_idx[i] != ref[i]) mismatches++;
+        printf("  [cmp] topk_idx    mismatches=%d/%d   cxx[0,:4]=%d %d %d %d  ref[0,:4]=%d %d %d %d\n",
+               mismatches, S*K,
+               cxx_idx[0], cxx_idx[1], cxx_idx[2], cxx_idx[3],
+               ref[0], ref[1], ref[2], ref[3]);
+    }
+
+    printf("\n=== MoE-only (before residual) ===\n");
+    compare_bf16("moe_out",  moe_out_dev.get(),    S * D,   "out_flat.bin");
+
+    // Manual host-side finalize: verify what down_out + expanded_row_idx + topk_w produce.
+    {
+        std::vector<uint16_t> h_down(TOTAL * D);
+        std::vector<int32_t>  h_ri(TOTAL);
+        std::vector<uint16_t> h_tw(S * K);
+        ACL_CHECK(aclrtMemcpy(h_down.data(), TOTAL*D*2, down_out_dev.get(), TOTAL*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_ri.data(),   TOTAL*4,   expanded_row_idx_dev.get(), TOTAL*4, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_tw.data(),   S*K*2,     topk_w_dev.get(),           S*K*2,   ACL_MEMCPY_DEVICE_TO_HOST));
+
+        printf("  expanded_row_idx (all %ld):\n  ", TOTAL);
+        for (int i = 0; i < TOTAL; i++) {
+            printf("%d ", h_ri[i]);
+            if ((i+1) % 10 == 0) printf("\n  ");
+        }
+        printf("\n");
+        // count unique and check bijection
+        std::vector<int> count(TOTAL, 0);
+        int out_of_range = 0;
+        for (int i = 0; i < TOTAL; i++) {
+            int v = h_ri[i];
+            if (v >= 0 && v < TOTAL) count[v]++;
+            else out_of_range++;
+        }
+        int bijection_ok = (out_of_range == 0);
+        for (int i = 0; i < TOTAL && bijection_ok; i++) if (count[i] != 1) bijection_ok = 0;
+        printf("  bijection=%s out_of_range=%d\n", bijection_ok ? "YES" : "NO", out_of_range);
+
+        // Also dump tokens_per_expert (int64) — should sum to TOTAL
+        std::vector<int64_t> h_tpe(E);
+        ACL_CHECK(aclrtMemcpy(h_tpe.data(), E*8, tokens_per_expert_dev.get(), E*8, ACL_MEMCPY_DEVICE_TO_HOST));
+        int64_t tpe_sum = 0, nonzero = 0;
+        int64_t tpe_max = 0;
+        for (int i = 0; i < E; i++) { tpe_sum += h_tpe[i]; if (h_tpe[i]>0) nonzero++; if (h_tpe[i]>tpe_max) tpe_max=h_tpe[i]; }
+        printf("  tokens_per_expert: sum=%ld nonzero=%ld max=%ld (expected sum=%ld if counts, or last=%ld if cumsum)\n",
+               tpe_sum, nonzero, tpe_max, TOTAL, TOTAL);
+        printf("  tpe[last 4]: %ld %ld %ld %ld\n", h_tpe[E-4], h_tpe[E-3], h_tpe[E-2], h_tpe[E-1]);
+
+        std::vector<float> manual(S * D, 0.0f);
+        for (int64_t p = 0; p < TOTAL; p++) {
+            int32_t src = h_ri[p];
+            int s = src / K;
+            int k = src % K;
+            if (s < 0 || s >= S || k < 0 || k >= K) { printf("  bad idx p=%ld src=%d\n", p, src); continue; }
+            float w = bf16_to_float(h_tw[s * K + k]);
+            for (int d = 0; d < D; d++) {
+                manual[s * D + d] += w * bf16_to_float(h_down[p * D + d]);
+            }
+        }
+        // Convert to bf16 and compare to Python out_flat
+        auto refbuf = read_file(data_dir + "/out_flat.bin");
+        auto* ref = (const uint16_t*)refbuf.data();
+        double l2d=0, l2r=0, maxd=0;
+        for (int64_t i = 0; i < S*D; i++) {
+            float a = manual[i], b = bf16_to_float(ref[i]);
+            l2d += (a-b)*(a-b); l2r += b*b;
+            if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+        }
+        double rel_manual = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+        printf("  [cmp] MANUAL(row_idx=src→flat)  rel=%.4e max_abs=%.4f  m[:4]=%.5f %.5f %.5f %.5f  r[:4]=%.5f %.5f %.5f %.5f\n",
+               rel_manual, maxd,
+               manual[0], manual[1], manual[2], manual[3],
+               bf16_to_float(ref[0]), bf16_to_float(ref[1]), bf16_to_float(ref[2]), bf16_to_float(ref[3]));
+
+        // Alternative semantic: row_idx[p] = destination position
+        // In that case: p=src_row, dst=h_ri[p]
+        std::vector<float> manual2(S * D, 0.0f);
+        for (int64_t p = 0; p < TOTAL; p++) {
+            int32_t dst = h_ri[p];
+            int s = dst / K;
+            int k = dst % K;
+            if (s < 0 || s >= S || k < 0 || k >= K) continue;
+            float w = bf16_to_float(h_tw[s * K + k]);
+            for (int d = 0; d < D; d++) {
+                manual2[s * D + d] += w * bf16_to_float(h_down[p * D + d]);
+            }
+        }
+        double l2d2=0, l2r2=0, maxd2=0;
+        for (int64_t i = 0; i < S*D; i++) {
+            float a = manual2[i], b = bf16_to_float(ref[i]);
+            l2d2 += (a-b)*(a-b); l2r2 += b*b;
+            if (std::abs(a-b) > maxd2) maxd2 = std::abs(a-b);
+        }
+        double rel_manual2 = std::sqrt(l2d2) / (std::sqrt(l2r2) + 1e-10);
+        printf("  [cmp] MANUAL(row_idx=p→dst_flat) rel=%.4e max_abs=%.4f  m[:4]=%.5f %.5f %.5f %.5f\n",
+               rel_manual2, maxd2,
+               manual2[0], manual2[1], manual2[2], manual2[3]);
+    }
+
+    // Manual finalize using cumsum (semantics-independent):
+    // For each (n, k), find p such that actual_s(p)=n AND expert(p)=topk_idx[n,k], then
+    // out[n] += topk_w[n,k] * down_out[p].
+    {
+        std::vector<uint16_t> h_down(TOTAL * D);
+        std::vector<int64_t>  h_tpe(E);
+        std::vector<int32_t>  h_tidx(S * K);
+        std::vector<uint16_t> h_tw(S * K);
+        std::vector<uint16_t> h_xn_all(S * D);
+        std::vector<uint16_t> h_ex_all(TOTAL * D);
+        ACL_CHECK(aclrtMemcpy(h_down.data(),   TOTAL*D*2, down_out_dev.get(), TOTAL*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_tpe.data(),    E*8,       tokens_per_expert_dev.get(), E*8, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_tidx.data(),   S*K*4,     topk_idx_dev.get(), S*K*4, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_tw.data(),     S*K*2,     topk_w_dev.get(),   S*K*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_xn_all.data(), S*D*2,     xn_dev.get(),       S*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_ex_all.data(), TOTAL*D*2, expanded_x_dev.get(), TOTAL*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+
+        // Build p → (actual_s, actual_expert).
+        // actual_s: find s with xn[s,0] == expanded_x[p,0]
+        // actual_expert: find e such that cumsum_tpe[e-1] <= p < cumsum_tpe[e]
+        std::vector<int> p_to_s(TOTAL), p_to_e(TOTAL);
+        int64_t cum = 0;
+        int cursor_e = 0;
+        for (int64_t p = 0; p < TOTAL; p++) {
+            while (cursor_e < E && p >= cum + h_tpe[cursor_e]) { cum += h_tpe[cursor_e]; cursor_e++; }
+            p_to_e[p] = cursor_e;
+            float ev = bf16_to_float(h_ex_all[p * D]);
+            int best = -1; float bd = 1e30f;
+            for (int s = 0; s < S; s++) {
+                float df = std::abs(bf16_to_float(h_xn_all[s * D]) - ev);
+                if (df < bd) { bd = df; best = s; }
+            }
+            p_to_s[p] = best;
+        }
+
+        // Build (n, k) → p lookup via (n, expert) → p
+        std::vector<float> manual_cum(S * D, 0.0f);
+        int found_count = 0;
+        for (int n = 0; n < S; n++) {
+            for (int k = 0; k < K; k++) {
+                int e = h_tidx[n * K + k];
+                float w = bf16_to_float(h_tw[n * K + k]);
+                // search p with p_to_s[p]==n and p_to_e[p]==e
+                int found_p = -1;
+                for (int64_t p = 0; p < TOTAL; p++) {
+                    if (p_to_s[p] == n && p_to_e[p] == e) { found_p = p; break; }
+                }
+                if (found_p < 0) {
+                    printf("  [!!!] not found: n=%d k=%d expert=%d\n", n, k, e);
+                    continue;
+                }
+                found_count++;
+                for (int d = 0; d < D; d++)
+                    manual_cum[n * D + d] += w * bf16_to_float(h_down[found_p * D + d]);
+            }
+        }
+        auto refbuf = read_file(data_dir + "/out_flat.bin");
+        auto* ref = (const uint16_t*)refbuf.data();
+        double l2d=0, l2r=0, maxd=0;
+        for (int64_t i = 0; i < S*D; i++) {
+            float a = manual_cum[i], b = bf16_to_float(ref[i]);
+            l2d += (a-b)*(a-b); l2r += b*b;
+            if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+        }
+        double rel_cum = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+        printf("  [cmp] MANUAL_CUMSUM (p via expert cumsum) rel=%.4e max=%.4f found=%d/40  m[:4]=%.5f %.5f %.5f %.5f\n",
+               rel_cum, maxd, found_count, manual_cum[0], manual_cum[1], manual_cum[2], manual_cum[3]);
+    }
+
+    // Dump all expanded_x[p, 0] and all xn[s, 0] to determine the mapping.
+    {
+        std::vector<uint16_t> h_xn_all(S * D);
+        ACL_CHECK(aclrtMemcpy(h_xn_all.data(), S*D*2, xn_dev.get(), S*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        std::vector<uint16_t> h_ex_all(TOTAL * D);
+        ACL_CHECK(aclrtMemcpy(h_ex_all.data(), TOTAL*D*2, expanded_x_dev.get(), TOTAL*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        printf("  xn[s, 0]: ");
+        for (int s = 0; s < S; s++) printf("%.5f ", bf16_to_float(h_xn_all[s * D]));
+        printf("\n  expanded_x[p, 0]: ");
+        for (int p = 0; p < TOTAL; p++) printf("%.5f ", bf16_to_float(h_ex_all[p * D]));
+        printf("\n  mapping p→s (by matching expanded_x[p,0] to xn[s,0]): ");
+        for (int p = 0; p < TOTAL; p++) {
+            float e = bf16_to_float(h_ex_all[p * D]);
+            int match = -1; float best = 1e30f;
+            for (int s = 0; s < S; s++) {
+                float df = std::abs(bf16_to_float(h_xn_all[s * D]) - e);
+                if (df < best) { best = df; match = s; }
+            }
+            printf("%d ", match);
+        }
+        printf("\n");
+    }
+
+    // Dump gate_out[p=4, :8] — gate activation of xn[0] via expert 10
+    {
+        std::vector<uint16_t> h_gate(I);
+        // NOTE: gate_out_dev was overwritten by silu+mul. So we need to reload from scratch.
+        // Instead just show down_out[4, :4].
+        std::vector<uint16_t> h_d(D);
+        ACL_CHECK(aclrtMemcpy(h_d.data(), D*2, (char*)down_out_dev.get() + 4*D*2, D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        printf("  down_out[p=4, :4] (s=0, k=0, expert=10): %.5f %.5f %.5f %.5f\n",
+               bf16_to_float(h_d[0]), bf16_to_float(h_d[1]), bf16_to_float(h_d[2]), bf16_to_float(h_d[3]));
+        // If GMM is correct, down_out[4] ~ ref[0] / topk_w[0,0]. ref[0,:4]=[-0.025, -0.007, 0.005, -0.008] / 0.224 ~ [-0.113, -0.031, 0.024, -0.036].
+        // But it's just ONE contribution so hard to compare directly.
+    }
+
+    // Single-expert verification using linear_hf: compute gate/up/down for (xn[0], expert=10)
+    // and compare with GMM's down_out at the corresponding position.
+    // linear_hf expects HF-layout weight [out_features, in_features]; our stacked gate_exps/up_exps
+    // are [E, D, I] — meaning per-expert shape is [D, I] (K, N) NOT HF [I, D]. So we can NOT directly
+    // linear_hf from gate_exps. Instead, load the expert-10 weight fresh and use linear_hf.
+    {
+        std::vector<int32_t> h_tidx_local(S * K);
+        ACL_CHECK(aclrtMemcpy(h_tidx_local.data(), S*K*4, topk_idx_dev.get(), S*K*4, ACL_MEMCPY_DEVICE_TO_HOST));
+        int target_expert = h_tidx_local[0 * K + 0];  // topk_idx[0, 0] should be 10 from Python ref
+        printf("\n  === Single-expert linear_hf vs GMM sanity (token 0, expert %d) ===\n", target_expert);
+
+        // Recompute p_to_s and p_to_e from host data (scoped locally).
+        std::vector<int64_t> h_tpe2(E);
+        std::vector<uint16_t> h_xn_all2(S * D);
+        std::vector<uint16_t> h_ex_all2(TOTAL * D);
+        ACL_CHECK(aclrtMemcpy(h_tpe2.data(),    E*8,       tokens_per_expert_dev.get(), E*8, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_xn_all2.data(), S*D*2,     xn_dev.get(),       S*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        ACL_CHECK(aclrtMemcpy(h_ex_all2.data(), TOTAL*D*2, expanded_x_dev.get(), TOTAL*D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+        std::vector<int> p_to_s(TOTAL), p_to_e(TOTAL);
+        {
+            int64_t cum = 0; int ce = 0;
+            for (int64_t p = 0; p < TOTAL; p++) {
+                while (ce < E && p >= cum + h_tpe2[ce]) { cum += h_tpe2[ce]; ce++; }
+                p_to_e[p] = ce;
+                float ev = bf16_to_float(h_ex_all2[p * D]);
+                int best = -1; float bd = 1e30f;
+                for (int s = 0; s < S; s++) {
+                    float df = std::abs(bf16_to_float(h_xn_all2[s * D]) - ev);
+                    if (df < bd) { bd = df; best = s; }
+                }
+                p_to_s[p] = best;
+            }
+        }
+
+        DeviceBuffer g_w, u_w, d_w;
+        char ename[256];
+        snprintf(ename, sizeof(ename), "model.layers.0.mlp.experts.%d.gate_proj.weight", target_expert);
+        if (!dw.st().get(ename)) { printf("  missing %s\n", ename); goto after_sanity; }
+
+        // Load full per-expert weight using public helpers (indirectly via loader).
+        // Easiest: use load_tensor_full_ via friend access... Instead, use st_ directly.
+        {
+            auto* m_gate = dw.st().get(ename);
+            DeviceBuffer gw_buf(m_gate->nbytes);
+            ACL_CHECK(aclrtMemcpy(gw_buf.get(), m_gate->nbytes, dw.st().data_ptr(*m_gate), m_gate->nbytes, ACL_MEMCPY_HOST_TO_DEVICE));
+            g_w = std::move(gw_buf);
+
+            snprintf(ename, sizeof(ename), "model.layers.0.mlp.experts.%d.up_proj.weight", target_expert);
+            auto* m_up = dw.st().get(ename);
+            DeviceBuffer uw_buf(m_up->nbytes);
+            ACL_CHECK(aclrtMemcpy(uw_buf.get(), m_up->nbytes, dw.st().data_ptr(*m_up), m_up->nbytes, ACL_MEMCPY_HOST_TO_DEVICE));
+            u_w = std::move(uw_buf);
+
+            snprintf(ename, sizeof(ename), "model.layers.0.mlp.experts.%d.down_proj.weight", target_expert);
+            auto* m_down = dw.st().get(ename);
+            DeviceBuffer dw_buf(m_down->nbytes);
+            ACL_CHECK(aclrtMemcpy(dw_buf.get(), m_down->nbytes, dw.st().data_ptr(*m_down), m_down->nbytes, ACL_MEMCPY_HOST_TO_DEVICE));
+            d_w = std::move(dw_buf);
+        }
+
+        // Compute gate = xn[0] @ gate_w.T → [I]; up = xn[0] @ up_w.T → [I]; act; down = act @ down_w.T → [D]
+        DeviceBuffer xn0_dev(D * 2);
+        ACL_CHECK(aclrtMemcpy(xn0_dev.get(), D*2, xn_dev.get(), D*2, ACL_MEMCPY_DEVICE_TO_DEVICE));
+
+        DeviceBuffer gate_v(I * 2), up_v(I * 2), act_v(I * 2), down_v(D * 2);
+        auto t_xn0   = make_contig_tensor(xn0_dev.get(), ACL_BF16, {1, D});
+        auto t_gate  = make_contig_tensor(gate_v.get(),  ACL_BF16, {1, I});
+        auto t_up    = make_contig_tensor(up_v.get(),    ACL_BF16, {1, I});
+        auto t_act   = make_contig_tensor(act_v.get(),   ACL_BF16, {1, I});
+        auto t_down  = make_contig_tensor(down_v.get(),  ACL_BF16, {1, D});
+        linear_hf(rt.stream(), t_xn0.get(), g_w.get(), ACL_BF16, I, D, t_gate.get());  // gate_proj HF [I, D]
+        linear_hf(rt.stream(), t_xn0.get(), u_w.get(), ACL_BF16, I, D, t_up.get());
+        rt.sync();
+        silu(rt.stream(), t_gate.get(), t_act.get());
+        mul(rt.stream(),  t_act.get(),  t_up.get(), t_act.get());
+        rt.sync();
+        linear_hf(rt.stream(), t_act.get(), d_w.get(), ACL_BF16, D, I, t_down.get());  // down_proj HF [D, I]
+        rt.sync();
+
+        std::vector<uint16_t> h_down_lin(D);
+        ACL_CHECK(aclrtMemcpy(h_down_lin.data(), D*2, down_v.get(), D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+
+        // Find the p in GMM output that corresponds to (s=0, expert=target_expert)
+        int found_p = -1;
+        for (int64_t p = 0; p < TOTAL; p++) {
+            if (p_to_s[p] == 0 && p_to_e[p] == target_expert) { found_p = p; break; }
+        }
+        if (found_p >= 0) {
+            std::vector<uint16_t> h_down_gmm(D);
+            ACL_CHECK(aclrtMemcpy(h_down_gmm.data(), D*2, (char*)down_out_dev.get() + found_p*D*2, D*2, ACL_MEMCPY_DEVICE_TO_HOST));
+            double l2d=0, l2r=0, maxd=0;
+            for (int i = 0; i < D; i++) {
+                float a = bf16_to_float(h_down_gmm[i]), b = bf16_to_float(h_down_lin[i]);
+                l2d += (a-b)*(a-b); l2r += b*b;
+                if (std::abs(a-b) > maxd) maxd = std::abs(a-b);
+            }
+            double rel = std::sqrt(l2d) / (std::sqrt(l2r) + 1e-10);
+            printf("  GMM down_out[p=%d]  vs  linear_hf down:  rel=%.4e max=%.4f\n", found_p, rel, maxd);
+            printf("    GMM[:4]:    %.5f %.5f %.5f %.5f\n",
+                   bf16_to_float(h_down_gmm[0]), bf16_to_float(h_down_gmm[1]), bf16_to_float(h_down_gmm[2]), bf16_to_float(h_down_gmm[3]));
+            printf("    linear[:4]: %.5f %.5f %.5f %.5f\n",
+                   bf16_to_float(h_down_lin[0]), bf16_to_float(h_down_lin[1]), bf16_to_float(h_down_lin[2]), bf16_to_float(h_down_lin[3]));
+        } else {
+            printf("  not found p for (s=0, expert=%d)\n", target_expert);
+        }
+    }
+    after_sanity:;
+
+    // Direct verification: gate_exps[expert_10, :4, :4] vs HF gate_proj_10 (transposed).
+    {
+        int expert_id = 10;
+        std::vector<uint16_t> h_stacked(4 * 4);
+        // gate_exps shape [E, D, I]. Expert 10 starts at offset expert_id * D * I * 2.
+        // Read the first 4 rows (d=0..3), first 4 cols (i=0..3). Row stride = I * 2 bytes.
+        for (int d = 0; d < 4; d++) {
+            ACL_CHECK(aclrtMemcpy(h_stacked.data() + d*4, 8,
+                (char*)moe.gate_exps.get() + (expert_id * D * I + d * I) * 2, 8,
+                ACL_MEMCPY_DEVICE_TO_HOST));
+        }
+        char ename[256];
+        snprintf(ename, sizeof(ename), "model.layers.0.mlp.experts.%d.gate_proj.weight", expert_id);
+        auto* m = dw.st().get(ename);
+        // HF gate_proj [I, D] row-major. Element at (i, d) is at offset (i*D + d)*2.
+        // Expected gate_exps[10, d, i] == HF_gate_proj[10][i, d].
+        // So for d in 0..3, i in 0..3: expected is HF[i, d].
+        std::vector<uint16_t> h_expected(4 * 4);
+        auto* hf = (const uint16_t*)dw.st().data_ptr(*m);
+        for (int d = 0; d < 4; d++) {
+            for (int i = 0; i < 4; i++) {
+                h_expected[d*4 + i] = hf[i * D + d];  // HF[i, d]
+            }
+        }
+        printf("\n  === gate_exps[10, :4, :4] layout check ===\n");
+        printf("  stacked: ");
+        for (int i = 0; i < 16; i++) printf("%.5f ", bf16_to_float(h_stacked[i]));
+        printf("\n  expected: ");
+        for (int i = 0; i < 16; i++) printf("%.5f ", bf16_to_float(h_expected[i]));
+        printf("\n");
+        int mism = 0;
+        for (int i = 0; i < 16; i++) if (h_stacked[i] != h_expected[i]) mism++;
+        printf("  mismatches: %d / 16\n", mism);
+    }
+
+    printf("\n=== Final (with residual) ===\n");
+    double rel = compare_bf16("final_out", final_dev.get(), S * D, "final_out.bin");
+    bool pass = rel < 5e-2;
+    printf("\n%s\n", pass ? "=== test_moe_layer PASS ===" : "=== test_moe_layer FAIL ===");
+    return pass ? 0 : 1;
+}

tests/test_op_support.cpp

@@ -0,0 +1,190 @@
+// test_op_support.cpp — smoke test which aclnn ops actually RUN on 910 初代.
+// Just call each candidate op with small tensors; report SUCCESS/FAILURE.
+// Guides optimization feasibility analysis.
+#include "acl_common.h"
+#include "acl_runtime.h"
+#include "aclnn_ops.h"
+#include <acl/acl.h>
+#include <aclnnop/aclnn_add_rms_norm.h>
+#include <aclnnop/aclnn_npu_format_cast.h>
+#include <aclnnop/aclnn_matmul.h>
+
+#include <cstdio>
+#include <cstring>
+#include <vector>
+
+static float bf16_to_float(uint16_t x) { uint32_t u = (uint32_t)x << 16; float f; std::memcpy(&f, &u, 4); return f; }
+static uint16_t f_to_bf16(float f) { uint32_t u; std::memcpy(&u, &f, 4); return (uint16_t)((u + 0x7FFF + ((u >> 16) & 1)) >> 16); }
+
+static const char* test_add_rms_norm(AclRuntime& rt) {
+    // Inputs: x1 [1, 16], x2 [1, 16] BF16; gamma [16] BF16
+    const int64_t D = 16;
+    std::vector<uint16_t> h_x1(D, f_to_bf16(0.5f));
+    std::vector<uint16_t> h_x2(D, f_to_bf16(0.3f));
+    std::vector<uint16_t> h_gamma(D, f_to_bf16(1.0f));
+    DeviceBuffer x1(D*2), x2(D*2), g(D*2), y(D*2), rstd(1*4), x_out(D*2);
+    ACL_CHECK(aclrtMemcpy(x1.get(), D*2, h_x1.data(), D*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(x2.get(), D*2, h_x2.data(), D*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(g.get(),  D*2, h_gamma.data(), D*2, ACL_MEMCPY_HOST_TO_DEVICE));
+
+    auto tx1 = make_contig_tensor(x1.get(),    ACL_BF16, {1, D});
+    auto tx2 = make_contig_tensor(x2.get(),    ACL_BF16, {1, D});
+    auto tg  = make_contig_tensor(g.get(),     ACL_BF16, {D});
+    auto ty  = make_contig_tensor(y.get(),     ACL_BF16, {1, D});
+    auto trs = make_contig_tensor(rstd.get(),  ACL_FLOAT, {1});
+    auto tout= make_contig_tensor(x_out.get(), ACL_BF16, {1, D});
+
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    aclnnStatus s = aclnnAddRmsNormGetWorkspaceSize(tx1.get(), tx2.get(), tg.get(), 1e-6,
+                                                     ty.get(), trs.get(), tout.get(), &ws, &exec);
+    if (s != 0) return "GetWorkspaceSize FAILED";
+    DeviceBuffer ws_buf;
+    if (ws > 0) ws_buf.alloc(ws);
+    s = aclnnAddRmsNorm(ws_buf.get(), ws, exec, rt.stream());
+    if (s != 0) return "aclnnAddRmsNorm FAILED (kernel not available on 910?)";
+    if (aclrtSynchronizeStream(rt.stream()) != 0) return "sync FAILED";
+    return "OK";
+}
+
+static const char* test_npu_format_cast_nz(AclRuntime& rt) {
+    // Transform a small [16, 16] BF16 tensor from ND to NZ format.
+    const int64_t H = 16, W = 16;
+    std::vector<uint16_t> h(H * W, f_to_bf16(1.0f));
+    DeviceBuffer src(H * W * 2);
+    ACL_CHECK(aclrtMemcpy(src.get(), H*W*2, h.data(), H*W*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    auto tsrc = make_contig_tensor(src.get(), ACL_BF16, {H, W});
+
+    // Step 1: calculate NZ shape
+    int64_t* dst_shape = nullptr;
+    uint64_t dst_shape_size = 0;
+    int actual_fmt = 0;
+    aclnnStatus s = aclnnNpuFormatCastCalculateSizeAndFormat(
+        tsrc.get(), /*dstFormat=*/29 /* FRACTAL_NZ */,
+        /*additionalDtype=*/27 /* BF16 */,
+        &dst_shape, &dst_shape_size, &actual_fmt);
+    if (s != 0) return "CalculateSizeAndFormat FAILED";
+
+    // Step 2: alloc dst and call cast
+    int64_t total = 1;
+    std::vector<int64_t> shape_vec(dst_shape, dst_shape + dst_shape_size);
+    for (auto d : shape_vec) total *= d;
+    DeviceBuffer dst(total * 2);
+    auto tdst = make_acl_tensor(dst.get(), ACL_BF16, shape_vec, {}, (aclFormat)actual_fmt);
+
+    uint64_t ws = 0; aclOpExecutor* exec = nullptr;
+    s = aclnnNpuFormatCastGetWorkspaceSize(tsrc.get(), tdst.get(), &ws, &exec);
+    if (s != 0) return "FormatCast GetWorkspaceSize FAILED";
+    DeviceBuffer ws_buf; if (ws > 0) ws_buf.alloc(ws);
+    s = aclnnNpuFormatCast(ws_buf.get(), ws, exec, rt.stream());
+    if (s != 0) return "aclnnNpuFormatCast FAILED (NZ not supported on 910?)";
+    if (aclrtSynchronizeStream(rt.stream()) != 0) return "sync FAILED";
+    return "OK";
+}
+
+static const char* test_matmul_nz(AclRuntime& rt) {
+    // Try a MatMul with NZ-format weight.
+    const int64_t M = 16, K = 32, N = 16;
+    std::vector<uint16_t> h_x(M * K, f_to_bf16(0.1f));
+    std::vector<uint16_t> h_w(K * N, f_to_bf16(0.1f));
+    DeviceBuffer x(M*K*2), w(K*N*2), y(M*N*2);
+    ACL_CHECK(aclrtMemcpy(x.get(), M*K*2, h_x.data(), M*K*2, ACL_MEMCPY_HOST_TO_DEVICE));
+    ACL_CHECK(aclrtMemcpy(w.get(), K*N*2, h_w.data(), K*N*2, ACL_MEMCPY_HOST_TO_DEVICE));
+
+    auto tx = make_contig_tensor(x.get(), ACL_BF16, {M, K});
+    auto tw_nd = make_contig_tensor(w.get(), ACL_BF16, {K, N});
+
+    // Convert W to NZ
+    int64_t* dst_shape = nullptr; uint64_t dst_size = 0; int fmt = 0;
+    if (aclnnNpuFormatCastCalculateSizeAndFormat(tw_nd.get(), 29, 27, &dst_shape, &dst_size, &fmt) != 0)
+        return "calc NZ FAILED";
+    int64_t total = 1;
+    std::vector<int64_t> sh(dst_shape, dst_shape + dst_size);
+    for (auto d : sh) total *= d;
+    DeviceBuffer w_nz(total * 2);
+    auto tw_nz = make_acl_tensor(w_nz.get(), ACL_BF16, sh, {}, (aclFormat)fmt);
+
+    uint64_t ws = 0; aclOpExecutor* e = nullptr;
+    aclnnStatus s = aclnnNpuFormatCastGetWorkspaceSize(tw_nd.get(), tw_nz.get(), &ws, &e);
+    if (s != 0) return "NZ cast ws FAILED";
+    DeviceBuffer wb; if (ws > 0) wb.alloc(ws);
+    if (aclnnNpuFormatCast(wb.get(), ws, e, rt.stream()) != 0) return "NZ cast EXEC FAILED";
+    if (aclrtSynchronizeStream(rt.stream()) != 0) return "NZ cast sync FAILED";
+
+    // Now try MatMul with x (ND) × w_nz (NZ)
+    auto ty = make_contig_tensor(y.get(), ACL_BF16, {M, N});
+    ws = 0; e = nullptr;
+    s = aclnnMatmulGetWorkspaceSize(tx.get(), tw_nz.get(), ty.get(), 0 /*trans*/, &ws, &e);
+    if (s != 0) return "MatMul NZ GetWorkspaceSize FAILED";
+    DeviceBuffer mwb; if (ws > 0) mwb.alloc(ws);
+    if (aclnnMatmul(mwb.get(), ws, e, rt.stream()) != 0) return "MatMul NZ EXEC FAILED (MatMul doesn't accept NZ on 910?)";
+    if (aclrtSynchronizeStream(rt.stream()) != 0) return "MatMul NZ sync FAILED";
+    return "OK";
+}
+
+static const char* test_multi_stream(AclRuntime& rt) {
+    // Allocate a SECOND stream and check it works.
+    aclrtStream s2 = nullptr;
+    if (aclrtCreateStream(&s2) != 0) return "aclrtCreateStream FAILED";
+    // Simple dummy op on s2
+    DeviceBuffer x(16 * 2);
+    std::vector<uint16_t> hx(16, 0);
+    if (aclrtMemcpyAsync(x.get(), 16*2, hx.data(), 16*2, ACL_MEMCPY_HOST_TO_DEVICE, s2) != 0) return "memcpy on s2 FAILED";
+    if (aclrtSynchronizeStream(s2) != 0) return "sync s2 FAILED";
+    aclrtDestroyStream(s2);
+    return "OK";
+}
+
+int main() {
+    AclRuntime rt;
+    rt.init(0);
+
+    printf("=== 910 op support smoke test ===\n");
+
+    const char* r1 = test_add_rms_norm(rt);
+    printf("  aclnnAddRmsNorm (fused Add+RmsNorm):     %s\n", r1);
+
+    const char* r2 = test_npu_format_cast_nz(rt);
+    printf("  aclnnNpuFormatCast → FRACTAL_NZ:         %s\n", r2);
+
+    const char* r3 = test_matmul_nz(rt);
+    printf("  aclnnMatmul with NZ weight:              %s\n", r3);
+
+    const char* r4 = test_multi_stream(rt);
+    printf("  Multi-stream (compute/comm overlap):     %s\n", r4);
+
+    // More candidates
+    printf("\n=== Additional 910 op candidates ===\n");
+
+    // InplaceAddRmsNorm
+    #include <aclnnop/aclnn_inplace_add_rms_norm.h>
+    {
+        const int64_t D = 16;
+        std::vector<uint16_t> h(D, f_to_bf16(0.5f)), hg(D, f_to_bf16(1.0f));
+        DeviceBuffer x1(D*2), x2(D*2), g(D*2), rstd(4);
+        ACL_CHECK(aclrtMemcpy(x1.get(), D*2, h.data(), D*2, ACL_MEMCPY_HOST_TO_DEVICE));
+        ACL_CHECK(aclrtMemcpy(x2.get(), D*2, h.data(), D*2, ACL_MEMCPY_HOST_TO_DEVICE));
+        ACL_CHECK(aclrtMemcpy(g.get(),  D*2, hg.data(), D*2, ACL_MEMCPY_HOST_TO_DEVICE));
+        auto tx1 = make_contig_tensor(x1.get(), ACL_BF16, {1, D});
+        auto tx2 = make_contig_tensor(x2.get(), ACL_BF16, {1, D});
+        auto tg  = make_contig_tensor(g.get(), ACL_BF16, {D});
+        auto tr  = make_contig_tensor(rstd.get(), ACL_FLOAT, {1});
+        uint64_t ws = 0; aclOpExecutor* e = nullptr;
+        aclnnStatus s = aclnnInplaceAddRmsNormGetWorkspaceSize(tx1.get(), tx2.get(), tg.get(), 1e-6,
+                                                                tr.get(), &ws, &e);
+        printf("  aclnnInplaceAddRmsNorm:  %s\n", s == 0 ? "GetWS OK" : "GetWS FAILED");
+        if (s == 0) {
+            DeviceBuffer wb; if (ws > 0) wb.alloc(ws);
+            s = aclnnInplaceAddRmsNorm(wb.get(), ws, e, rt.stream());
+            printf("    exec: %s\n", s == 0 ? "OK" : "FAILED");
+        }
+    }
+
+    // Test HCCL AllReduce on a separate stream
+    printf("  HCCL AllReduce on stream2:  requires TP>1, skipped in this smoke test\n");
+
+    printf("\n=== FINAL Feasibility Summary ===\n");
+    printf("  Optimization A (FRACTAL_NZ):       INFEASIBLE (910 不支持)\n");
+    printf("  Optimization B (multi-stream):     FEASIBLE\n");
+    printf("  Optimization C (Add+RmsNorm):      INFEASIBLE (910 无 kernel)\n");
+    return 0;
+}