Integrate Qwen3 CANN FAI decode

.github/workflows/ci.yml

@@ -174,7 +174,7 @@ jobs:
           PYPTO_QWEN3_MODEL_DIR: /data/l00955553/model/Qwen3-14B
           PTO2_RING_DEP_POOL: 16384
           PTO2_RING_TASK_WINDOW: 16384
-          PTO2_RING_HEAP: 1073741824
+          PTO2_RING_HEAP: 2147483648
         run: |
           source activate.sh
           marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}"
@@ -196,7 +196,7 @@ jobs:
           PYPTO_QWEN3_MODEL_DIR: /data/l00955553/model/Qwen3-14B
           PTO2_RING_DEP_POOL: 16384
           PTO2_RING_TASK_WINDOW: 16384
-          PTO2_RING_HEAP: 1073741824
+          PTO2_RING_HEAP: 2147483648
         run: |
           source activate.sh
           marker="pypto-serving-ci-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}"

examples/model/qwen3_14b/npu_generate.py

@@ -345,7 +345,7 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--model-dir", required=True, help="Local model directory, e.g. a Hugging Face snapshot.")
     parser.add_argument("--prompt", required=True, help="Prompt text.")
     parser.add_argument("--model-id", default="qwen3-14b-local")
-    parser.add_argument("--platform", default="a2a3", choices=["a2a3sim", "a2a3", "a5sim", "a5"])
+    parser.add_argument("--platform", default="a2a3", choices=["a2a3"])
     parser.add_argument("--device-id", type=int, default=0, help="Default NPU device id when --devices is unset.")
     parser.add_argument(
         "--devices",
@@ -374,9 +374,7 @@ def build_parser() -> argparse.ArgumentParser:
         "--max-num-batched-tokens",
         type=int,
         default=4096,
-        help="Total tokens per scheduling step (used by warmup). NOTE: the 40-layer "
-        "fused prefill deadlocks the single-die ring-heap above ~415 total tokens, "
-        "so set this low enough that the per-request count stays under the ceiling.",
+        help="Total tokens per scheduling step and production-scale warmup.",
     )
     parser.add_argument(
         "--npu-memory-utilization",

examples/model/qwen3_14b/runner/npu_executor.py

@@ -27,7 +27,7 @@
 from python.core._profiling import StageTimer
 from python.core.model_runner import ModelRunner
 from python.core.pypto_executor import PyptoExecutor as CorePyptoExecutor
-from python.core.types import RuntimeModel
+from python.core.types import GenerateConfig, ModelRecord, RuntimeModel
 from python.core.utils import rope_tables, round_up
 
 
@@ -104,11 +104,13 @@ def __init__(
         self,
         kv_cache_manager=None,
         *,
-        platform: str = "a2a3sim",
+        platform: str = "a2a3",
         device_ids: Sequence[int] = (0,),
         save_kernels_dir: str | None = None,
         l3_trace: bool = False,
     ) -> None:
+        if platform != "a2a3":
+            raise ValueError("Qwen3 direct CANN FAI decode supports only A2/A3 onboard")
         super().__init__(
             kv_cache_manager,
             platform=platform,
@@ -132,6 +134,21 @@ def supports_device_embedding(self) -> bool:
         """Qwen3 NPU decode embeds greedy token ids inside the device kernel."""
         return True
 
+    def validate_generate_batch(
+        self,
+        record: ModelRecord,
+        batch_size: int,
+        config: GenerateConfig,
+    ) -> None:
+        """Accept only the active B1 prefix or all B16 KV-cache rows."""
+        del config
+        kernel_batch = record.runtime.max_batch_size
+        if batch_size not in (1, kernel_batch):
+            raise ValueError(
+                "Qwen3 direct CANN FAI decode supports only batch 1 or the "
+                f"full batch capacity {kernel_batch}, got {batch_size}"
+            )
+
     def _create_runner(self, model_id: str, compiled: object) -> ModelRunner:
         """Create the Qwen3-14B runtime runner for compiled kernels."""
         if not isinstance(compiled, _CompiledKernels):
@@ -180,10 +197,10 @@ def _mark(label: str) -> None:
 
         if int(qwen3_decode_fwd.BATCH) != kernel_batch:
             raise ValueError(
-                "decode_fwd.decode_fwd is compiled for a fixed kernel BATCH of "
+                "decode_fwd.decode_fwd has an internal batch capacity of "
                 f"{int(qwen3_decode_fwd.BATCH)}, but runtime max_batch_size is "
-                f"{kernel_batch}; they must match (decode statically computes and "
-                "writes BATCH rows / BATCH logit rows)."
+                f"{kernel_batch}; they must match. Runtime dispatch supports only "
+                "batch 1 or the full batch capacity."
             )
         if int(model.config.num_hidden_layers) != int(qwen3_decode_fwd.NUM_LAYERS):
             raise ValueError(
@@ -455,12 +472,13 @@ def _compile_decode_fwd_callable(
     ) -> _L3Callable:
         """Compile the fused all-layer PAGED decode HOST wrapper into a distributed program.
 
-        Signature (21 args; PAGED KV via block_table + slot_mapping, same pool as
+        Signature (25 args; PAGED KV via block_table + slot_mapping, same pool as
         prefill):
           input_rms_weight, wq, wk, wv, q_norm_weight,
           k_norm_weight, seq_lens, block_table, slot_mapping, rope_cos, rope_sin,
           k_cache, v_cache, wo, w_gate, w_up, w_down, post_rms_weight,
-          final_norm_weight, lm_head_weight, out.
+          final_norm_weight, lm_head_weight, out, embed_weight,
+          sampled_ids_in, sampled_ids_out, next_hidden.
 
         k_cache/v_cache are the PAGED pool (rows = num_layers * batch *
         runtime_cache_blocks * num_kv_heads * page_size — identical to prefill);

examples/model/qwen3_14b/runner/npu_runner.py

@@ -134,14 +134,16 @@ class _DecodeInputs:
 
 @dataclass
 class _DecodeKernelInputs:
-    """Fixed-batch tensors passed to the fused decode kernel."""
+    """Active-prefix tensors passed to the fused decode kernel."""
 
     actual_batch: int
     token_ids: torch.Tensor
     seq_lens: torch.Tensor
     block_table: torch.Tensor
     slot_mapping: torch.Tensor
     logits: torch.Tensor
+    sampled_ids: torch.Tensor
+    next_hidden: torch.Tensor
 
 
 @dataclass
@@ -183,16 +185,16 @@ def __init__(
             self._share_static_kernel_tensors()
             self._static_args = self._build_static_kernel_args()
 
-    #: Scratch KV pages for the profile pass — slot=-1 means only page 0
-    #: is ever touched (reads via block_table=0, writes via slot clamp to 0).
+    #: Scratch KV pages for the profile pass. Warmup maps every read and write
+    #: to slot 0, so one page covers the full synthetic batch.
     _PROFILE_PAGES = 1
 
     def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConfig) -> int:
         """Create the L3 worker-resident cache before the first request.
 
         Order (vLLM-style): run a profile warmup FIRST so the simpler arena
         is allocated before the KV cache competes for HBM.  The profile uses
-        ``slot_mapping=-1`` / ``block_table=0`` so only a single dummy page
+        ``slot_mapping=0`` / ``block_table=0`` so only a single dummy page
         is needed.  The KV cache size is then computed by the estimation
         formula and allocated into the remaining space; if allocation fails
         the page count is halved and retried.
@@ -213,7 +215,7 @@ def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConf
             self._materialize_static_tensors()
 
         # -- phase 1: profile warmup → arena allocated ----------------------
-        # Uses slot_mapping=-1 so no real KV cache pages are needed; the
+        # Uses slot_mapping=0 so no real KV cache pages are needed; the
         # 1-page scratch is the dummy target for all reads/writes.
         logger.info(f"[init_kv_cache] profile warmup (scratch {self._PROFILE_PAGES} page) …")
         ModelRunner.init_kv_cache(self, model_id, config, runtime, num_pages=self._PROFILE_PAGES)
@@ -381,15 +383,12 @@ def warmup(self, model: RuntimeModel) -> None:
         self._warmup_dispatch(model.runtime)
 
     def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
-        """Production-scale prefill + decode warm-up with slot_mapping=-1.
+        """Production-scale prefill + decode warm-up mapped to scratch slot 0.
 
         Sizes the prefill to one serving scheduling step — total tokens =
         ``max_num_batched_tokens`` spread across ``max_batch`` requests.
-        This deliberately exercises the kernel at the configured capacity so
-        that a too-large ``max_num_batched_tokens`` (which would hit the
-        single-die attention heap ceiling around seq≈415 in the 40-layer
-        fused prefill) fails at startup rather than on the first real
-        request.
+        This exercises the configured capacity so ring sizing or allocation
+        failures surface at startup rather than on the first real request.
         """
         batch = runtime.max_batch_size
         max_seq = runtime.max_seq_len
@@ -400,7 +399,7 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
 
         logger.info(
             f"[warmup] starting (batch={batch}, max_num_batched_tokens={mnb}, "
-            f"max_seq={max_seq}, per_req={per_req}, total_tokens={total_tokens}, slot=-1)",
+            f"max_seq={max_seq}, per_req={per_req}, total_tokens={total_tokens}, slot=0)",
 
         )
         compiled = self._compiled
@@ -412,7 +411,7 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
         compiled.prefill_chunk_lens_buffer.zero_()
         compiled.prefill_chunk_offsets_buffer.zero_()
         compiled.prefill_block_table_buffer.fill_(0)    # all reads from page 0
-        compiled.prefill_slot_mapping_buffer.fill_(-1)  # all writes to page 0
+        compiled.prefill_slot_mapping_buffer.zero_()      # all writes to page 0
 
         token_offset = 0
         for b in range(batch):
@@ -442,11 +441,11 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
         )
         logger.info(f"[warmup] prefill done ({time.perf_counter() - t0:.2f} s)")
 
-        # -- decode (full fixed batch, minimal seq) -------------------------
+        # -- decode (full capacity batch, minimal seq) ----------------------
         compiled.decode_token_ids_buffer.zero_()
         compiled.decode_seq_lens_buffer.zero_()
         compiled.decode_block_table_buffer.fill_(0)     # all reads from page 0
-        compiled.decode_slot_mapping_buffer.fill_(-1)   # all writes to page 0
+        compiled.decode_slot_mapping_buffer.zero_()      # all writes to page 0
 
         for b in range(batch):
             compiled.decode_seq_lens_buffer[b] = min(per_req + 1, max_seq)
@@ -458,6 +457,8 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
             block_table=compiled.decode_block_table_buffer,
             slot_mapping=compiled.decode_slot_mapping_buffer,
             logits=compiled.decode_logits_buffer,
+            sampled_ids=compiled.decode_sampled_ids_buffer,
+            next_hidden=compiled.decode_next_hidden_buffer,
         )
 
         logger.info(f"[warmup] decode dispatch … (batch={batch}, seq_len={per_req + 1})")
@@ -614,12 +615,9 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult:
         KV is already in place with no bridge. KV is keyed by block_table page id, not by
         kernel row, so a request may occupy any row each step (no stable-slot shim).
 
-        The kernel is FIXED-BATCH (it computes all max_batch_size rows and writes
-        each row's current-token KV). Pad the active batch up to the kernel batch by
-        REPLICATING active row 0's inputs into the padding rows: those rows then
-        recompute row 0's K/V and write row 0's own slot with byte-identical values
-        (an idempotent, safe write), and their logits are trimmed off below. This
-        avoids padded rows clobbering an unrelated request's physical page.
+        The public decode batch is either one active row or all 16 independent
+        rows. Dynamic tensor descriptors carry only that active prefix into the
+        kernel; B1 never materializes or computes replicated KV rows.
         """
         compiled = self._compiled
         model_id = model.config.model_id
@@ -631,10 +629,8 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult:
         k_cache = kv_cache.key_pages
         v_cache = kv_cache.value_pages
 
-        kernel_inputs = self._pad_decode_inputs(model, decode_inputs)
+        kernel_inputs = self._stage_decode_inputs(model, decode_inputs)
 
-        # Padded block_table / slot_mapping only ever reference row 0's
-        # already-valid pages, so bound-check exactly what the kernel will read.
         self._validate_kv_cache_bounds(model, kernel_inputs.block_table, kernel_inputs.slot_mapping, k_cache)
 
         self._run_distributed_program(
@@ -644,7 +640,7 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult:
         for batch_idx, alloc in enumerate(batch.kv_allocations):
             alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item()))
         sampled_ids, next_hidden = self._integrated_sample_result(
-            compiled.decode_sampled_ids_buffer,
+            kernel_inputs.sampled_ids,
             # decode_fwd's next_hidden output is the embedding for sampled_ids_in
             # used by this decode step. The newly sampled token is embedded at the
             # start of the following decode_fwd call, so there is no next-step
@@ -655,9 +651,7 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult:
         )
         return DecodeResult(
             hidden_states=decode_inputs.hidden.float(),
-            logits=kernel_inputs.logits[: kernel_inputs.actual_batch, : model.config.vocab_size].to(
-                decode_inputs.hidden.device
-            ),
+            logits=kernel_inputs.logits[:, : model.config.vocab_size].to(decode_inputs.hidden.device),
             sampled_token_ids=sampled_ids,
             next_hidden_states=next_hidden,
         )
@@ -776,22 +770,23 @@ def _decode_kernel_args(
             inputs.logits,
             static.padded_embed_weight,
             inputs.token_ids,
-            self._compiled.decode_sampled_ids_buffer,
-            self._compiled.decode_next_hidden_buffer,
+            inputs.sampled_ids,
+            inputs.next_hidden,
         )
 
-    def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _DecodeKernelInputs:
-        """Pad active decode rows to the fixed kernel batch.
-
-        The fused decode kernel computes all ``max_batch_size`` rows. Inactive
-        rows replicate row 0 so their KV writes are idempotent instead of
-        targeting unrelated pages.
-        """
+    def _stage_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _DecodeKernelInputs:
+        """Copy active decode rows into shared buffers and return prefix views."""
         compiled = self._compiled
         actual_batch = inputs.actual_batch
         kernel_batch = model.runtime.max_batch_size
         max_blocks = self._max_blocks_per_seq(model)
 
+        if actual_batch not in (1, kernel_batch):
+            raise ValueError(
+                "Qwen3 direct CANN FAI decode supports only batch 1 or the "
+                f"full batch capacity {kernel_batch}, got {actual_batch}"
+            )
+
         if kernel_batch > compiled.decode_logits_buffer.shape[0]:
             raise ValueError(
                 f"kernel batch {kernel_batch} exceeds logits buffer batch "
@@ -808,36 +803,23 @@ def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _Dec
             )
         width = 1
         token_ids[:actual_batch, :width].copy_(active_token_ids[:, :width])
-        if actual_batch < kernel_batch:
-            token_ids[actual_batch:, :width].copy_(
-                active_token_ids[0:1, :width].expand(kernel_batch - actual_batch, width)
-            )
+
+        seq_lens = compiled.decode_seq_lens_buffer[:actual_batch]
+        block_table = compiled.decode_block_table_buffer[: actual_batch * max_blocks]
+        slot_mapping = compiled.decode_slot_mapping_buffer[:actual_batch]
+        seq_lens.copy_(inputs.seq_lens.reshape(actual_batch))
+        block_table.copy_(inputs.block_table.reshape(actual_batch * max_blocks))
+        slot_mapping.copy_(inputs.slot_mapping.reshape(actual_batch))
 
         return _DecodeKernelInputs(
             actual_batch=actual_batch,
-            token_ids=token_ids,
-            seq_lens=self._copy_replicated_rows(
-                compiled.decode_seq_lens_buffer,
-                inputs.seq_lens,
-                actual_batch,
-                kernel_batch,
-                rows_each=1,
-            ),
-            block_table=self._copy_replicated_rows(
-                compiled.decode_block_table_buffer,
-                inputs.block_table,
-                actual_batch,
-                kernel_batch,
-                rows_each=max_blocks,
-            ),
-            slot_mapping=self._copy_replicated_rows(
-                compiled.decode_slot_mapping_buffer,
-                inputs.slot_mapping,
-                actual_batch,
-                kernel_batch,
-                rows_each=1,
-            ),
-            logits=compiled.decode_logits_buffer,
+            token_ids=token_ids[:actual_batch],
+            seq_lens=seq_lens,
+            block_table=block_table,
+            slot_mapping=slot_mapping,
+            logits=compiled.decode_logits_buffer[:actual_batch],
+            sampled_ids=compiled.decode_sampled_ids_buffer[:actual_batch],
+            next_hidden=compiled.decode_next_hidden_buffer[:actual_batch],
         )
 
     def _run_distributed_program(self, callable_spec: _L3Callable, *args: Any) -> Any:
@@ -908,23 +890,6 @@ def _materialize_static_tensors(self) -> None:
         ):
             self._coerce_l3_arg(worker, arg)
 
-    @staticmethod
-    def _copy_replicated_rows(
-        dst: torch.Tensor,
-        active: torch.Tensor,
-        actual_batch: int,
-        kernel_batch: int,
-        *,
-        rows_each: int,
-    ) -> torch.Tensor:
-        """Copy active rows and fill inactive rows by replicating row 0."""
-        active_view = active.reshape(actual_batch, rows_each)
-        dst_view = dst.reshape(kernel_batch, rows_each)
-        dst_view[:actual_batch].copy_(active_view)
-        if actual_batch < kernel_batch:
-            dst_view[actual_batch:].copy_(active_view[0:1].expand(kernel_batch - actual_batch, rows_each))
-        return dst
-
     @staticmethod
     def _static_device_tensor(tensor: torch.Tensor) -> _StaticDeviceTensor:
         """Mark a CPU tensor for one-time upload to the shared worker."""

examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py

@@ -102,7 +102,7 @@ def qwen3_decode_host(
     sampled_ids: pl.Out[pl.Tensor],
     next_hidden: pl.Out[pl.Tensor],
 ) -> tuple[pl.Tensor, pl.Tensor, pl.Tensor]:
-    logits, sampled_ids, next_hidden = decode_fwd(
+    out, sampled_ids, next_hidden = decode_fwd(
         input_rms_weight,
         wq,
         wk,
@@ -129,7 +129,7 @@ def qwen3_decode_host(
         sampled_ids,
         next_hidden,
     )
-    return logits, sampled_ids, next_hidden
+    return out, sampled_ids, next_hidden
 
 
 @pl.jit.host