From db411a7f54cc2b1ca91a2b284f51ad9be425ddc6 Mon Sep 17 00:00:00 2001 From: vegetabledoww Date: Thu, 9 Jul 2026 18:09:33 -0700 Subject: [PATCH 1/4] feat(qwen3-a8w8): add loader and L2 runtime support Add the compressed-tensors A8W8 loader for Qwen3-14B checkpoints and extend the common PyPTO runtime plumbing for PTO ISA pinning and L2 child-memory tensors. --- .../model/qwen3_14b/runner/a8w8_loader.py | 187 ++++++++++++++++++ python/core/pypto_executor.py | 3 + python/runtime/worker.py | 8 + 3 files changed, 198 insertions(+) create mode 100644 examples/model/qwen3_14b/runner/a8w8_loader.py diff --git a/examples/model/qwen3_14b/runner/a8w8_loader.py b/examples/model/qwen3_14b/runner/a8w8_loader.py new file mode 100644 index 0000000..43b1436 --- /dev/null +++ b/examples/model/qwen3_14b/runner/a8w8_loader.py @@ -0,0 +1,187 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import json +import struct +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + +import torch + +from python.core.model_loader import ( + ModelLoadRequest, + _build_layer_specs, + _build_model_config, +) +from python.core.tokenizer import TransformersTokenizerAdapter +from python.core.types import LoadedModel, RuntimeConfig, RuntimeModel + + +_DTYPE_MAP = { + "I8": torch.int8, + "F32": torch.float32, + "BF16": torch.bfloat16, +} + + +@dataclass(frozen=True) +class _TensorRecord: + shard: Path + dtype: str + shape: tuple[int, ...] + data_start: int + offsets: tuple[int, int] + + +class _SafeTensorIndex: + """Minimal safetensors reader for Qwen compressed-tensors checkpoints.""" + + def __init__(self, root: Path) -> None: + self.root = root + self.records: dict[str, _TensorRecord] = {} + for shard in _safetensor_shards(root): + with shard.open("rb") as f: + header_len = struct.unpack(" torch.Tensor: + rec = self.records[key] + dtype = _DTYPE_MAP[rec.dtype] + begin, end = rec.offsets + with rec.shard.open("rb") as f: + f.seek(rec.data_start + begin) + raw = bytearray(f.read(end - begin)) + return torch.frombuffer(raw, dtype=dtype).clone().reshape(rec.shape) + + +def _safetensor_shards(root: Path) -> list[Path]: + index_path = root / "model.safetensors.index.json" + if index_path.exists(): + index_data = json.loads(index_path.read_text()) + return [root / filename for filename in sorted(set(index_data["weight_map"].values()))] + return sorted(root.glob("*.safetensors")) + + +def _hf_linear_to_kernel_i8(index: _SafeTensorIndex, key: str) -> torch.Tensor: + return index.load(key).t().contiguous() + + +def _hf_scale_to_kernel(index: _SafeTensorIndex, key: str) -> torch.Tensor: + return index.load(key).reshape(1, -1).float().contiguous() + + +def _hf_linear_to_bf16(index: _SafeTensorIndex, weight_key: str, scale_key: str) -> torch.Tensor: + weight = index.load(weight_key).float() + scale = index.load(scale_key).float().reshape(-1, 1) + return (weight * scale).t().contiguous().to(torch.bfloat16) + + +def _layer_prefix(layer_idx: int) -> str: + return f"model.layers.{layer_idx}" + + +class Qwen3A8W8DirectoryLoader: + """Loader for local Qwen3-14B compressed-tensors W8A8 checkpoints.""" + + format_names = ("qwen3-a8w8", "qwen3-w8a8", "a8w8") + + def supports_format(self, model_format: str) -> bool: + return model_format.lower() in self.format_names + + def can_load(self, model_path: Path) -> bool: + if not (model_path / "config.json").exists(): + return False + return bool(_safetensor_shards(model_path)) + + def load(self, request: ModelLoadRequest) -> LoadedModel: + model_path = Path(request.model_dir) + config_data = json.loads((model_path / "config.json").read_text()) + trust_remote_code = bool(request.loader_options.get("trust_remote_code", False)) + tokenizer = TransformersTokenizerAdapter.from_pretrained( + str(model_path), + trust_remote_code=trust_remote_code, + ) + config = _build_model_config(request.model_id, config_data, tokenizer) + runtime = request.runtime_config or RuntimeConfig(max_seq_len=config.max_position_embeddings) + decode_backend = str(request.loader_options.get("decode_backend", "a8w8")).lower() + if decode_backend != "a8w8": + raise ValueError(f"unsupported qwen3-a8w8 decode_backend: {decode_backend!r}") + layer_specs = _build_layer_specs(config) + index = _SafeTensorIndex(model_path) + + embed_tokens = index.load("model.embed_tokens.weight").to(torch.bfloat16).contiguous() + final_norm_weight = index.load("model.norm.weight").float().contiguous() + lm_head = index.load("lm_head.weight").to(torch.bfloat16).contiguous() + + layers = [] + for spec in layer_specs: + prefix = _layer_prefix(spec.layer_idx) + layer = SimpleNamespace( + quantization="a8w8", + input_rms_weight=index.load(f"{prefix}.input_layernorm.weight").reshape(1, -1).float(), + wq=_hf_linear_to_kernel_i8(index, f"{prefix}.self_attn.q_proj.weight"), + wk=_hf_linear_to_kernel_i8(index, f"{prefix}.self_attn.k_proj.weight"), + wv=_hf_linear_to_kernel_i8(index, f"{prefix}.self_attn.v_proj.weight"), + wq_scale=_hf_scale_to_kernel(index, f"{prefix}.self_attn.q_proj.weight_scale"), + wk_scale=_hf_scale_to_kernel(index, f"{prefix}.self_attn.k_proj.weight_scale"), + wv_scale=_hf_scale_to_kernel(index, f"{prefix}.self_attn.v_proj.weight_scale"), + q_norm_weight=index.load(f"{prefix}.self_attn.q_norm.weight").reshape(1, -1).float(), + k_norm_weight=index.load(f"{prefix}.self_attn.k_norm.weight").reshape(1, -1).float(), + wo=_hf_linear_to_kernel_i8(index, f"{prefix}.self_attn.o_proj.weight"), + wo_scale=_hf_scale_to_kernel(index, f"{prefix}.self_attn.o_proj.weight_scale"), + post_rms_weight=index.load(f"{prefix}.post_attention_layernorm.weight").reshape(1, -1).float(), + w_gate=_hf_linear_to_bf16( + index, + f"{prefix}.mlp.gate_proj.weight", + f"{prefix}.mlp.gate_proj.weight_scale", + ), + w_up=_hf_linear_to_bf16( + index, + f"{prefix}.mlp.up_proj.weight", + f"{prefix}.mlp.up_proj.weight_scale", + ), + w_down=_hf_linear_to_bf16( + index, + f"{prefix}.mlp.down_proj.weight", + f"{prefix}.mlp.down_proj.weight_scale", + ), + ) + layers.append(layer) + + runtime_model = RuntimeModel( + config=config, + runtime=runtime, + embed_tokens=embed_tokens, + final_norm_weight=final_norm_weight, + lm_head=lm_head, + layers=layers, + ) + return LoadedModel( + model_id=request.model_id, + model_dir=str(model_path), + config=config, + tokenizer=tokenizer, + layer_specs=layer_specs, + runtime_model=runtime_model, + ) diff --git a/python/core/pypto_executor.py b/python/core/pypto_executor.py index 026d8a1..3c6b571 100644 --- a/python/core/pypto_executor.py +++ b/python/core/pypto_executor.py @@ -42,6 +42,7 @@ def __init__( platform: str = "a2a3sim", device_ids: Sequence[int] = (0,), save_kernels_dir: str | None = None, + pto_isa_commit: str | None = None, ) -> None: """Initialize common PyPTO runtime options and model registries.""" super().__init__(kv_cache_manager) @@ -50,6 +51,7 @@ def __init__( if not self._device_ids: raise ValueError("device_ids must contain at least one device id") self._save_kernels_dir = save_kernels_dir + self._pto_isa_commit = pto_isa_commit self._runners: dict[str, ModelRunner] = {} self._compiled: dict[str, object] = {} @@ -118,6 +120,7 @@ def _run_config(self, *, codegen_only: bool): codegen_only=codegen_only, save_kernels=self._save_kernels_dir is not None, save_kernels_dir=self._save_kernels_dir, + pto_isa_commit=self._pto_isa_commit, ) @abstractmethod diff --git a/python/runtime/worker.py b/python/runtime/worker.py index 32609a1..c6d61ff 100644 --- a/python/runtime/worker.py +++ b/python/runtime/worker.py @@ -24,6 +24,10 @@ import torch from simpler.task_interface import DataType +try: + from simpler.task_interface import ContinuousTensor +except ImportError: # Newer simpler exposes the same child-memory wrapper as Tensor. + from simpler.task_interface import Tensor as ContinuousTensor from simpler.worker import Worker as SimplerWorker @@ -119,6 +123,10 @@ def torch_dtype(self) -> torch.dtype: """Return the corresponding ``torch.dtype``.""" return _to_torch_dtype(self.dtype) + def to_continuous_tensor(self) -> ContinuousTensor: + """Return a Simpler tensor view over existing worker child memory.""" + return ContinuousTensor.make(self.data_ptr, self.shape, self.dtype, child_memory=True) + class Worker: """Manage a Simpler L2 or L3 worker for LLM execution. From ec286f6dd8c36356af5c40bd63e1b6f603b06a36 Mon Sep 17 00:00:00 2001 From: vegetabledoww Date: Thu, 9 Jul 2026 18:09:40 -0700 Subject: [PATCH 2/4] feat(qwen3-a8w8): add A8W8 L2 model runner Add the Qwen3-14B A8W8 L2 runner responsible for paged KV cache management, child-memory argument preparation, prefill dispatch, decode dispatch, and logits collection. --- .../model/qwen3_14b/runner/npu_runner_a8w8.py | 773 ++++++++++++++++++ 1 file changed, 773 insertions(+) create mode 100644 examples/model/qwen3_14b/runner/npu_runner_a8w8.py diff --git a/examples/model/qwen3_14b/runner/npu_runner_a8w8.py b/examples/model/qwen3_14b/runner/npu_runner_a8w8.py new file mode 100644 index 0000000..578b6fb --- /dev/null +++ b/examples/model/qwen3_14b/runner/npu_runner_a8w8.py @@ -0,0 +1,773 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import ctypes +from dataclasses import dataclass +from typing import Any + +import torch + +from examples.model.qwen3_14b.runner.npu_runner import Qwen314BModelRunner, _add_run_timing_args +from python.core.model_runner import ModelRunner +from python.core.types import ( + DecodeBatch, + DecodeResult, + PrefillBatch, + PrefillResult, + RuntimeModel, +) +from python.profile import profile_span +from python.runtime.worker import Worker as LlmWorker +from python.runtime.worker import WorkerTensor + +_QWEN14B_A8W8_PREFILL_CHUNK_LAYERS = 10 +_QWEN14B_LM_HEAD_CHUNK_ROWS = 8192 + + +def _l2_trace_name(kernel_name: str) -> str: + if "prefill" in kernel_name: + return "kernel.prefill_fwd" + if "decode" in kernel_name: + return "kernel.decode_fwd" + if "final_rms" in kernel_name: + return "kernel.final_rms" + if "lm_head" in kernel_name: + return "kernel.lm_head" + return f"kernel.{kernel_name}" + + +@dataclass +class _KernelLayerWeights: + """Kernel-ready weights for one transformer layer.""" + + input_rms_weight: torch.Tensor + wq: torch.Tensor + wk: torch.Tensor + wv: torch.Tensor + q_norm_weight: torch.Tensor + k_norm_weight: torch.Tensor + wo: torch.Tensor + post_rms_weight: torch.Tensor + w_gate: torch.Tensor + w_up: torch.Tensor + w_down: torch.Tensor + wq_scale: torch.Tensor | None = None + wk_scale: torch.Tensor | None = None + wv_scale: torch.Tensor | None = None + wo_scale: torch.Tensor | None = None + + +@dataclass +class _L2Callable: + """Assembled non-L3 callable and launch metadata.""" + + chip_callable: object + name: str + runtime_name: str + block_dim: int + aicpu_thread_num: int + param_infos: tuple[object, ...] + + +@dataclass +class _CompiledKernels: + """Compiled Qwen3-14B kernels and immutable runtime tensors.""" + + prefill: _L2Callable + decode: _L2Callable + final_norm_weight: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + padded_vocab: int + padded_lm_head_weight: torch.Tensor + layers: list[_KernelLayerWeights] + decode_weights: dict[str, torch.Tensor] + decode_logits_buffer: torch.Tensor + + +@dataclass +class _PrefillInputs: + """Host tensors passed to the prefill kernel.""" + + actual_batch: int + hidden: torch.Tensor + seq_lens: torch.Tensor + chunk_lens: torch.Tensor + chunk_offsets: torch.Tensor + block_table: torch.Tensor + slot_mapping: torch.Tensor + + +@dataclass +class _DecodeInputs: + """Padded host tensors passed to the decode kernel.""" + + actual_batch: int + hidden: torch.Tensor + seq_lens: torch.Tensor + block_table: torch.Tensor + slot_mapping: torch.Tensor + + +@dataclass +class _L2ProgramHandle: + """L2 callable registration state for one runner process.""" + + callable_id: int + runtime_name: str + + +class Qwen314BA8W8ModelRunner(ModelRunner): + """Runtime wrapper for one Qwen3-14B model's compiled PyPTO kernels.""" + + _compute_slot_mapping = staticmethod(Qwen314BModelRunner._compute_slot_mapping) + _write_block_table_row = staticmethod(Qwen314BModelRunner._write_block_table_row) + _validate_batch_size = staticmethod(Qwen314BModelRunner._validate_batch_size) + _max_blocks_per_seq = staticmethod(Qwen314BModelRunner._max_blocks_per_seq) + + def __init__( + self, + *, + compiled: _CompiledKernels, + model_id: str = "test-model", + platform: str = "a2a3sim", + device_id: int = 0, + save_kernels_dir: str | None = None, + ) -> None: + super().__init__() + self._model_id = model_id + self._compiled = compiled + self._platform = platform + self._device_id = device_id + self._save_kernels_dir = save_kernels_dir + self._l2_workers: dict[str, LlmWorker] = {} + self._l2_programs: dict[int, _L2ProgramHandle] = {} + self._l2_child_allocs: dict[tuple[str, int], tuple[int, int]] = {} + self._kv_scale_caches: dict[str, tuple[WorkerTensor, WorkerTensor]] = {} + + def init_kv_cache(self, model_id: str, config, runtime) -> int: + """Create the runner-owned KV cache, plus INT8 scale pages for A8W8.""" + num_pages = super().init_kv_cache(model_id, config, runtime) + if model_id in self._kv_scale_caches: + return num_pages + cache_rows = config.num_hidden_layers * num_pages * config.num_key_value_heads * runtime.page_size + key_scale = self._alloc_kv_cache_tensor((cache_rows, 8), torch.float32) + try: + value_scale = self._alloc_kv_cache_tensor((cache_rows, 8), torch.float32) + except Exception: + self._free_kv_cache_tensor(key_scale) + raise + self._kv_scale_caches[model_id] = (key_scale, value_scale) + return num_pages + + def close_kv_cache(self) -> None: + for key_scale, value_scale in list(self._kv_scale_caches.values()): + self._free_kv_cache_tensor(key_scale) + self._free_kv_cache_tensor(value_scale) + self._kv_scale_caches.clear() + super().close_kv_cache() + + def _kv_cache_runtime_name(self) -> str: + if self._compiled.prefill.runtime_name != self._compiled.decode.runtime_name: + raise ValueError( + "device-side KV cache requires prefill and decode to use the same L2 runtime: " + f"{self._compiled.prefill.runtime_name!r} != {self._compiled.decode.runtime_name!r}" + ) + return self._compiled.prefill.runtime_name + + def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> WorkerTensor: + """Allocate one KV cache tensor on the L2 NPU worker.""" + worker = self._worker_for_runtime(self._kv_cache_runtime_name()) + return worker.alloc_tensor(shape, dtype) + + def _free_kv_cache_tensor(self, tensor: WorkerTensor) -> None: + """Free one KV cache tensor from the L2 NPU worker.""" + worker = self._l2_workers.get(self._kv_cache_runtime_name()) + if worker is not None and worker.initialized: + worker.free_tensor(tensor) + + @staticmethod + def _validate_kv_cache_bounds( + model: RuntimeModel, + block_table: torch.Tensor, + slot_mapping: torch.Tensor, + cache: WorkerTensor, + ) -> None: + """Fail on host before an invalid KV page id reaches the NPU kernel.""" + valid_blocks = block_table[block_table >= 0] + valid_slots = slot_mapping[slot_mapping >= 0] + if valid_blocks.numel() == 0 and valid_slots.numel() == 0: + return + max_block_id = int(valid_blocks.max().item()) if valid_blocks.numel() else -1 + max_slot_block = int(valid_slots.max().item()) // model.runtime.page_size if valid_slots.numel() else -1 + max_page_id = max(max_block_id, max_slot_block) + rows_per_layer = cache.shape[0] // model.config.num_hidden_layers + max_pages = rows_per_layer // (model.config.num_key_value_heads * model.runtime.page_size) + if max_page_id >= max_pages: + raise RuntimeError( + "KV cache page id exceeds runner device cache capacity: " + f"max_page_id={max_page_id}, max_pages={max_pages}, " + f"cache_shape={cache.shape}, block_table_shape={tuple(block_table.shape)}, " + f"slot_mapping_shape={tuple(slot_mapping.shape)}" + ) + + def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult: + """Run the JIT all-layer prefill kernel and return next-token logits.""" + compiled = self._compiled + prefill_inputs = self._prepare_prefill_inputs(model, batch) + dw = compiled.decode_weights + + kv_cache = self._kv_caches.get(model.config.model_id) + if kv_cache is None: + raise RuntimeError(f"KV cache for model {model.config.model_id!r} is not initialized") + k_cache = kv_cache.key_pages + v_cache = kv_cache.value_pages + self._validate_kv_cache_bounds(model, prefill_inputs.block_table, prefill_inputs.slot_mapping, k_cache) + kv_scales = self._kv_scale_caches.get(model.config.model_id) + logits_padded = torch.zeros( + (prefill_inputs.actual_batch, compiled.padded_vocab), + dtype=torch.float32, + ).share_memory_() + + if kv_scales is None: + raise RuntimeError(f"missing A8W8 KV scale cache for model {model.config.model_id!r}") + k_cache_scale, v_cache_scale = kv_scales + rows_per_layer = k_cache.shape[0] // model.config.num_hidden_layers + hidden = prefill_inputs.hidden + + def weight_slice(name: str, start: int, layers: int, rows_per_layer_: int = 1) -> WorkerTensor: + tensor = dw[name][start * rows_per_layer_ : (start + layers) * rows_per_layer_] + return self._l2_child_tensor(compiled.prefill.runtime_name, tensor) + + for layer_start in range(0, model.config.num_hidden_layers, _QWEN14B_A8W8_PREFILL_CHUNK_LAYERS): + layer_count = min( + _QWEN14B_A8W8_PREFILL_CHUNK_LAYERS, + model.config.num_hidden_layers - layer_start, + ) + cache_row_start = layer_start * rows_per_layer + cache_rows = layer_count * rows_per_layer + hidden_out = torch.empty_like(hidden).share_memory_() + scratch_logits = torch.empty_like(logits_padded).share_memory_() + self._run_l2_program( + compiled.prefill, + hidden, + prefill_inputs.seq_lens, + prefill_inputs.chunk_lens, + prefill_inputs.chunk_offsets, + weight_slice("decode_input_rms_weight", layer_start, layer_count), + weight_slice("decode_wq", layer_start, layer_count, model.config.hidden_size), + weight_slice("decode_wk", layer_start, layer_count, model.config.hidden_size), + weight_slice("decode_wv", layer_start, layer_count, model.config.hidden_size), + weight_slice("decode_wq_scale", layer_start, layer_count), + weight_slice("decode_wk_scale", layer_start, layer_count), + weight_slice("decode_wv_scale", layer_start, layer_count), + weight_slice("decode_q_norm_weight", layer_start, layer_count), + weight_slice("decode_k_norm_weight", layer_start, layer_count), + self._l2_child_tensor(compiled.prefill.runtime_name, compiled.rope_cos), + self._l2_child_tensor(compiled.prefill.runtime_name, compiled.rope_sin), + prefill_inputs.block_table, + prefill_inputs.slot_mapping, + self._worker_tensor_view( + k_cache, + cache_row_start * model.config.head_dim, + (cache_rows, model.config.head_dim), + 1, + ), + self._worker_tensor_view( + v_cache, + cache_row_start * model.config.head_dim, + (cache_rows, model.config.head_dim), + 1, + ), + self._worker_tensor_view(k_cache_scale, cache_row_start * 8, (cache_rows, 8), 4), + self._worker_tensor_view(v_cache_scale, cache_row_start * 8, (cache_rows, 8), 4), + weight_slice("decode_wo", layer_start, layer_count, model.config.hidden_size), + weight_slice("decode_wo_scale", layer_start, layer_count), + weight_slice("decode_post_rms_weight", layer_start, layer_count), + weight_slice("decode_w_gate", layer_start, layer_count, model.config.hidden_size), + weight_slice("decode_w_up", layer_start, layer_count, model.config.hidden_size), + weight_slice("decode_w_down", layer_start, layer_count, model.config.intermediate_size), + self._l2_child_tensor(compiled.prefill.runtime_name, compiled.final_norm_weight), + self._l2_child_tensor(compiled.prefill.runtime_name, compiled.padded_lm_head_weight), + scratch_logits, + hidden_out, + ) + hidden = hidden_out + logits_padded = self._project_logits_host(model, compiled, prefill_inputs, hidden) + + for batch_idx, alloc in enumerate(batch.kv_allocations): + seq_len = int(batch.seq_lens[batch_idx].item()) + alloc.tokens_used = max(alloc.tokens_used, seq_len) + return PrefillResult( + last_hidden=None, + logits=logits_padded[:, : model.config.vocab_size], + ) + + def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: + """Run the fused all-layer PAGED ``decode_layer.decode_fwd`` and return logits. + + ``decode_fwd`` runs all NUM_LAYERS + the LM head in one dispatch over the + PAGED KV pool, addressing KV via ``block_table`` + ``slot_mapping`` — the + SAME device-resident KV pool prefill writes (``self._kv_caches``), so prompt + 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. + """ + compiled = self._compiled + model_id = model.config.model_id + decode_inputs = self._prepare_decode_inputs(model, batch) + actual_batch = decode_inputs.actual_batch + dw = compiled.decode_weights + rt = compiled.decode.runtime_name + kernel_batch = model.runtime.max_batch_size + max_blocks = self._max_blocks_per_seq(model) + + kv_cache = self._kv_caches.get(model_id) + if kv_cache is None: + raise RuntimeError(f"KV cache for model {model_id!r} is not initialized") + k_cache = kv_cache.key_pages + v_cache = kv_cache.value_pages + kv_scales = self._kv_scale_caches.get(model_id) + + if kernel_batch > compiled.decode_logits_buffer.shape[0]: + raise ValueError( + f"kernel batch {kernel_batch} exceeds logits buffer batch " + f"{compiled.decode_logits_buffer.shape[0]}" + ) + + # Pad active inputs up to the fixed kernel batch by replicating row 0. + def _pad_rows(active: torch.Tensor, rows_each: int) -> torch.Tensor: + view = active.reshape(actual_batch, rows_each) + padded = view[0:1].expand(kernel_batch - actual_batch, rows_each) + return torch.cat([view, padded], dim=0).reshape(-1).contiguous() + + hidden = torch.zeros((kernel_batch, model.config.hidden_size), dtype=torch.bfloat16) + hidden[:actual_batch] = decode_inputs.hidden + hidden[actual_batch:] = decode_inputs.hidden[0:1] + hidden = hidden.share_memory_() + seq_lens = _pad_rows(decode_inputs.seq_lens, 1).to(torch.int32).share_memory_() + block_table = _pad_rows(decode_inputs.block_table, max_blocks).to(torch.int32).share_memory_() + slot_mapping = _pad_rows(decode_inputs.slot_mapping, 1).to(torch.int32).share_memory_() + + # 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, block_table, slot_mapping, k_cache) + + logits_padded = compiled.decode_logits_buffer # full [kernel_batch, vocab]; trimmed below + if kv_scales is None: + raise RuntimeError(f"missing A8W8 KV scale cache for model {model_id!r}") + k_cache_scale, v_cache_scale = kv_scales + self._run_l2_program( + compiled.decode, + hidden, + self._l2_child_tensor(rt, dw["decode_input_rms_weight"]), + self._l2_child_tensor(rt, dw["decode_wq"]), + self._l2_child_tensor(rt, dw["decode_wk"]), + self._l2_child_tensor(rt, dw["decode_wv"]), + self._l2_child_tensor(rt, dw["decode_wq_scale"]), + self._l2_child_tensor(rt, dw["decode_wk_scale"]), + self._l2_child_tensor(rt, dw["decode_wv_scale"]), + self._l2_child_tensor(rt, dw["decode_q_norm_weight"]), + self._l2_child_tensor(rt, dw["decode_k_norm_weight"]), + seq_lens, + block_table, + slot_mapping, + self._l2_child_tensor(rt, compiled.rope_cos), + self._l2_child_tensor(rt, compiled.rope_sin), + k_cache, + v_cache, + k_cache_scale, + v_cache_scale, + self._l2_child_tensor(rt, dw["decode_wo"]), + self._l2_child_tensor(rt, dw["decode_wo_scale"]), + self._l2_child_tensor(rt, dw["decode_w_gate"]), + self._l2_child_tensor(rt, dw["decode_w_up"]), + self._l2_child_tensor(rt, dw["decode_w_down"]), + self._l2_child_tensor(rt, dw["decode_post_rms_weight"]), + self._l2_child_tensor(rt, compiled.final_norm_weight), + self._l2_child_tensor(rt, compiled.padded_lm_head_weight), + logits_padded, + ) + for batch_idx, alloc in enumerate(batch.kv_allocations): + alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item())) + return DecodeResult( + hidden_states=decode_inputs.hidden.float(), + logits=logits_padded[:actual_batch, : model.config.vocab_size].to(decode_inputs.hidden.device), + ) + + def _project_logits_host( + self, + model: RuntimeModel, + compiled: _CompiledKernels, + prefill_inputs: _PrefillInputs, + hidden: torch.Tensor, + ) -> torch.Tensor: + """Project final prefill hidden states on host for A8W8 chunked prefill. + + The A8W8 prefill kernel is currently chunked by transformer layers and + returns the final hidden states. Keep the final RMSNorm + LM head here as + an explicit fallback until the prefill kernel owns that projection end to + end, so the performance tradeoff is visible at the serving boundary. + """ + final_hidden = torch.zeros((prefill_inputs.actual_batch, model.config.hidden_size), dtype=torch.bfloat16) + for batch_idx in range(prefill_inputs.actual_batch): + chunk_offset = int(prefill_inputs.chunk_offsets[batch_idx].item()) + chunk_len = int(prefill_inputs.chunk_lens[batch_idx].item()) + final_hidden[batch_idx] = hidden[chunk_offset + chunk_len - 1] + + gamma = compiled.final_norm_weight.view(1, -1).float() + x = final_hidden.float() + normed = (x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + model.config.rms_norm_eps) * gamma).to( + torch.bfloat16 + ) + logits = torch.empty((prefill_inputs.actual_batch, compiled.padded_vocab), dtype=torch.float32).share_memory_() + lm_head = compiled.padded_lm_head_weight + for row_start in range(0, compiled.padded_vocab, _QWEN14B_LM_HEAD_CHUNK_ROWS): + row_end = min(row_start + _QWEN14B_LM_HEAD_CHUNK_ROWS, compiled.padded_vocab) + logits[:, row_start:row_end] = normed.float() @ lm_head[row_start:row_end].float().T + return logits + + @staticmethod + def _worker_tensor_view( + tensor: WorkerTensor, + element_offset: int, + shape: tuple[int, ...], + element_size: int, + ) -> WorkerTensor: + """Return a contiguous WorkerTensor view starting at an element offset.""" + return WorkerTensor( + data_ptr=tensor.data_ptr + element_offset * element_size, + shape=shape, + dtype=tensor.dtype, + ) + + def _run_l2_program(self, callable_spec: _L2Callable, *args: Any) -> Any: + """Run a compiled non-L3 program through the LLM Simpler worker.""" + from simpler.task_interface import CallConfig # noqa: PLC0415 + + span_args = { + "kernel": callable_spec.name, + "runtime": callable_spec.runtime_name, + "block_dim": callable_spec.block_dim, + "aicpu_thread_num": callable_spec.aicpu_thread_num, + } + with profile_span( + _l2_trace_name(callable_spec.name), + cat="kernel", + level="kernel", + args=span_args, + ): + handle = self._ensure_l2_program(callable_spec) + orch_args = self._build_l2_orch_args(callable_spec, args) + + cfg = CallConfig() + cfg.block_dim = callable_spec.block_dim + cfg.aicpu_thread_num = callable_spec.aicpu_thread_num + + worker = self._l2_workers[handle.runtime_name] + timing = worker.run(handle.callable_id, orch_args, cfg) + _add_run_timing_args(span_args, timing) + return timing + + def _worker_for_runtime(self, runtime_name: str) -> LlmWorker: + """Return an initialized worker for ``runtime_name``.""" + worker = self._l2_workers.get(runtime_name) + if worker is not None: + return worker + worker = LlmWorker( + level=2, + platform=self._platform, + runtime=runtime_name, + device_id=self._device_id, + auto_init=True, + ) + self._l2_workers[runtime_name] = worker + return worker + + def _ensure_l2_program(self, callable_spec: _L2Callable) -> _L2ProgramHandle: + """Register and cache one executor-assembled non-L3 callable.""" + key = id(callable_spec) + cached = self._l2_programs.get(key) + if cached is not None: + return cached + + worker = self._worker_for_runtime(callable_spec.runtime_name) + + handle = _L2ProgramHandle( + callable_id=worker.register(callable_spec.chip_callable), + runtime_name=callable_spec.runtime_name, + ) + self._l2_programs[key] = handle + return handle + + def _l2_child_tensor( + self, + runtime_name: str, + tensor: torch.Tensor, + *, + upload: bool = True, + refresh: bool = False, + ) -> WorkerTensor: + """Return a worker-resident view for a CPU tensor's backing storage.""" + from simpler_setup.torch_interop import torch_dtype_to_datatype # noqa: PLC0415 + + if tensor.device.type != "cpu": + raise ValueError("child-memory tensor must be on CPU") + if not tensor.is_contiguous(): + raise ValueError("child-memory tensor must be contiguous") + tensor = self._share_cpu_tensor(tensor) + storage = tensor.untyped_storage() + storage_ptr = int(storage.data_ptr()) + storage_nbytes = int(storage.nbytes()) + tensor_offset = int(tensor.data_ptr()) - storage_ptr + if tensor_offset < 0 or tensor_offset + int(tensor.nbytes) > storage_nbytes: + raise ValueError("tensor view is outside its backing storage") + + key = (runtime_name, storage_ptr) + alloc = self._l2_child_allocs.get(key) + if alloc is None: + worker = self._worker_for_runtime(runtime_name) + dev_ptr = worker.malloc(storage_nbytes) + if upload: + worker.copy_to(dev_ptr, storage_ptr, storage_nbytes) + alloc = (dev_ptr, storage_nbytes) + self._l2_child_allocs[key] = alloc + elif upload and refresh: + worker = self._worker_for_runtime(runtime_name) + worker.copy_to(alloc[0], storage_ptr, storage_nbytes) + + dev_base, _ = alloc + shape = tuple(int(dim) for dim in tensor.shape) + return WorkerTensor( + data_ptr=dev_base + tensor_offset, + shape=shape, + dtype=torch_dtype_to_datatype(tensor.dtype), + ) + + def _release_l2_child_allocs(self, runtime_name: str) -> None: + """Free cached child-memory allocations for one L2 runtime.""" + worker = self._l2_workers.get(runtime_name) + for key, (dev_ptr, _nbytes) in list(self._l2_child_allocs.items()): + key_runtime, _storage_ptr = key + if key_runtime != runtime_name: + continue + if worker is not None and worker.initialized: + worker.free(dev_ptr) + self._l2_child_allocs.pop(key, None) + + @staticmethod + def _share_cpu_tensor(tensor: torch.Tensor) -> torch.Tensor: + """Move a CPU tensor's storage to shared memory if needed.""" + if tensor.device.type == "cpu" and not tensor.is_shared(): + return tensor.share_memory_() + return tensor + + def close(self) -> None: + """Release non-L3 child-memory allocations and L2 workers.""" + self.close_kv_cache() + for (runtime_name, _), (dev_ptr, _nbytes) in list(self._l2_child_allocs.items()): + worker = self._l2_workers.get(runtime_name) + if worker is not None and worker.initialized: + worker.free(dev_ptr) + self._l2_child_allocs.clear() + self._l2_programs.clear() + for worker in self._l2_workers.values(): + worker.close() + self._l2_workers.clear() + + @staticmethod + def _build_l2_orch_args(callable_spec: _L2Callable, args: tuple[Any, ...]): + """Build ``ChipStorageTaskArgs`` for a compiled L2 program call.""" + from simpler.task_interface import ChipStorageTaskArgs, scalar_to_uint64 # noqa: PLC0415 + try: + from simpler.task_interface import ContinuousTensor # noqa: PLC0415 + except ImportError: + from simpler.task_interface import Tensor as ContinuousTensor # noqa: PLC0415 + from simpler_setup.torch_interop import make_tensor_arg # noqa: PLC0415 + + param_infos = callable_spec.param_infos + if len(args) != len(param_infos): + names = [p.name for p in param_infos] + raise TypeError( + f"compiled program expects {len(param_infos)} arguments, got {len(args)}. Parameters: {names}" + ) + + orch_args = ChipStorageTaskArgs() + for info, arg in zip(param_infos, args, strict=True): + if info.shape is None: + if not isinstance(arg, ctypes._SimpleCData): + raise TypeError(f"scalar parameter {info.name!r} must be passed as a ctypes scalar") + orch_args.add_scalar(scalar_to_uint64(arg)) + continue + if isinstance(arg, WorkerTensor): + orch_args.add_tensor(arg.to_continuous_tensor()) + continue + if isinstance(arg, ContinuousTensor): + orch_args.add_tensor(arg) + continue + if not isinstance(arg, torch.Tensor): + raise TypeError(f"tensor parameter {info.name!r} expects torch.Tensor, got {type(arg).__name__}") + if arg.device.type != "cpu": + raise ValueError(f"tensor parameter {info.name!r} must be on CPU for Simpler L2 dispatch") + if not arg.is_contiguous(): + raise ValueError(f"tensor parameter {info.name!r} must be contiguous") + if not arg.is_shared(): + arg.share_memory_() + orch_args.add_tensor(make_tensor_arg(arg)) + return orch_args + + def _prepare_prefill_inputs( + self, + model: RuntimeModel, + batch: PrefillBatch, + ) -> _PrefillInputs: + """Pack variable-length prefill requests into kernel input tensors.""" + batch_count = len(batch.kv_allocations) if batch.kv_allocations else int(batch.seq_lens.shape[0]) + actual_batch = self._validate_batch_size(model, batch_count) + max_seq = model.runtime.max_seq_len + hidden_size = model.config.hidden_size + page_size = model.runtime.page_size + max_blocks = self._max_blocks_per_seq(model) + + seq_lens = torch.empty((actual_batch,), dtype=torch.int32) + chunk_lens = torch.empty((actual_batch,), dtype=torch.int32) + chunk_offsets = torch.empty((actual_batch,), dtype=torch.int32) + block_table = torch.full((actual_batch * max_blocks,), -1, dtype=torch.int32) + seq_len_values = [int(batch.seq_lens[idx].item()) for idx in range(actual_batch)] + chunk_len_values: list[int] = [] + chunk_start_values: list[int] = [] + for batch_idx, seq_len in enumerate(seq_len_values): + if batch.positions is not None: + row_positions = batch.positions[batch_idx].detach().cpu() + valid_positions = row_positions[row_positions >= 0] + if valid_positions.numel() == 0: + raise ValueError("prefill positions must include at least one chunk token") + chunk_start = int(valid_positions[0].item()) + chunk_len = int(valid_positions.numel()) + expected_positions = torch.arange( + chunk_start, + chunk_start + chunk_len, + dtype=valid_positions.dtype, + ) + if not torch.equal(valid_positions, expected_positions): + raise ValueError( + "prefill batch.positions must form one contiguous chunk: " + f"chunk_start={chunk_start}, chunk_len={chunk_len}, seq_len={seq_len}" + ) + else: + chunk_len = seq_len + chunk_start = 0 + if chunk_len <= 0: + raise ValueError("prefill chunk_lens must be positive") + if chunk_start + chunk_len != seq_len: + raise ValueError( + "prefill chunk must end at seq_len: " + f"chunk_start={chunk_start}, chunk_len={chunk_len}, seq_len={seq_len}" + ) + chunk_len_values.append(chunk_len) + chunk_start_values.append(chunk_start) + total_tokens = sum(chunk_len_values) + hidden = torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16) + slot_mapping = torch.empty((total_tokens,), dtype=torch.int32) + + token_offset = 0 + for batch_idx in range(actual_batch): + alloc = batch.kv_allocations[batch_idx] if batch_idx < len(batch.kv_allocations) else None + seq_len = seq_len_values[batch_idx] + if seq_len <= 0: + raise ValueError("prefill seq_lens must be positive") + if seq_len > max_seq: + raise ValueError(f"prefill seq_len {seq_len} exceeds max_seq_len {max_seq}") + seq_lens[batch_idx] = seq_len + chunk_len = chunk_len_values[batch_idx] + chunk_start = chunk_start_values[batch_idx] + chunk_lens[batch_idx] = chunk_len + chunk_offsets[batch_idx] = token_offset + embeddings = batch.input_embeddings[batch_idx, :chunk_len, :].to(torch.bfloat16).cpu() + hidden[token_offset : token_offset + chunk_len, :] = embeddings + + if alloc is not None: + page_ids = alloc.page_ids + elif batch_idx < len(batch.block_ids): + page_ids = batch.block_ids[batch_idx] + else: + page_ids = [] + self._write_block_table_row(block_table, batch_idx, max_blocks, page_ids) + + slot_row = self._compute_slot_mapping(page_ids, chunk_len, page_size, start_pos=chunk_start) + slot_mapping[token_offset : token_offset + chunk_len] = slot_row + token_offset += chunk_len + + return _PrefillInputs( + actual_batch=actual_batch, + hidden=hidden.share_memory_(), + seq_lens=seq_lens.share_memory_(), + chunk_lens=chunk_lens.share_memory_(), + chunk_offsets=chunk_offsets.share_memory_(), + block_table=block_table.share_memory_(), + slot_mapping=slot_mapping.share_memory_(), + ) + + def _prepare_decode_inputs( + self, + model: RuntimeModel, + batch: DecodeBatch, + ) -> _DecodeInputs: + """Pack active decode requests into fused decode-kernel inputs.""" + batch_count = len(batch.kv_allocations) if batch.kv_allocations else int(batch.seq_lens.shape[0]) + actual_batch = self._validate_batch_size(model, batch_count) + hidden_size = model.config.hidden_size + page_size = model.runtime.page_size + max_blocks = self._max_blocks_per_seq(model) + + hidden = torch.zeros((actual_batch, hidden_size), dtype=torch.bfloat16) + seq_lens = torch.empty((actual_batch,), dtype=torch.int32) + block_table = torch.full((actual_batch * max_blocks,), -1, dtype=torch.int32) + slot_mapping = torch.empty((actual_batch,), dtype=torch.int32) + + for batch_idx in range(actual_batch): + alloc = batch.kv_allocations[batch_idx] if batch_idx < len(batch.kv_allocations) else None + seq_len = int(batch.seq_lens[batch_idx].item()) + if seq_len <= 0: + raise ValueError("decode seq_lens must be positive") + if seq_len > model.runtime.max_seq_len: + raise ValueError( + f"decode seq_len {seq_len} exceeds max_seq_len {model.runtime.max_seq_len}" + ) + hidden[batch_idx, :] = batch.hidden_states[batch_idx].to(torch.bfloat16).cpu() + seq_lens[batch_idx] = seq_len + + if alloc is not None: + page_ids = alloc.page_ids + elif batch_idx < len(batch.block_ids): + page_ids = batch.block_ids[batch_idx] + else: + page_ids = [] + self._write_block_table_row(block_table, batch_idx, max_blocks, page_ids) + + tokens_used = seq_len - 1 + page_idx = tokens_used // page_size + offset = tokens_used % page_size + slot_mapping[batch_idx] = page_ids[page_idx] * page_size + offset + + return _DecodeInputs( + actual_batch=actual_batch, + hidden=hidden.share_memory_(), + seq_lens=seq_lens.share_memory_(), + block_table=block_table.share_memory_(), + slot_mapping=slot_mapping.share_memory_(), + ) From 039e7385a96bc2b09267c3010bd92824291424cd Mon Sep 17 00:00:00 2001 From: vegetabledoww Date: Thu, 9 Jul 2026 18:10:14 -0700 Subject: [PATCH 3/4] feat(qwen3-a8w8): compile kernels and wire generation CLI Add the A8W8 PyPTO executor, connect npu_generate.py to the qwen3-a8w8 model format, keep BF16 and A8W8 runtime settings isolated, and add regression coverage for loader, runner, and CLI behavior. --- examples/model/qwen3_14b/npu_generate.py | 159 ++++-- .../qwen3_14b/runner/npu_executor_a8w8.py | 534 ++++++++++++++++++ .../model/qwen3_14b/runner/npu_runner_a8w8.py | 1 - tests/test_batching.py | 200 ++++++- 4 files changed, 834 insertions(+), 60 deletions(-) create mode 100644 examples/model/qwen3_14b/runner/npu_executor_a8w8.py diff --git a/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index 87b65c9..a9727b0 100644 --- a/examples/model/qwen3_14b/npu_generate.py +++ b/examples/model/qwen3_14b/npu_generate.py @@ -10,6 +10,7 @@ from __future__ import annotations import argparse +import dataclasses import statistics import sys import time @@ -33,11 +34,17 @@ def _bootstrap_package_root() -> None: from python.core import GenerateConfig, LLMEngine, RuntimeConfig from python.core.kv_cache import KvCacheManager +from python.core.model_loader import ModelLoader from python.core.parallel import ParallelConfig, parse_device_ids from python.profile import get_profiler, merge_profile, profile_span -from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor as PyptoExecutor +from examples.model.qwen3_14b.runner.a8w8_loader import Qwen3A8W8DirectoryLoader +from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor +from examples.model.qwen3_14b.runner.npu_executor_a8w8 import Qwen314BA8W8PyptoExecutor from python.core.types import LoadedModel -import dataclasses + + +_QWEN3_BF16_FORMAT = "qwen3-14b" +_QWEN3_A8W8_FORMAT = "qwen3-a8w8" # ----------------------------------------------------------------------------- @@ -70,24 +77,6 @@ def TimePhase(self, name: str): finally: self.phases[name] = self.phases.get(name, 0.0) + (time.perf_counter() - t0) - def WrapKernel(self, fn, name: str, *, group_by_decode_step: bool = False): - """Return a wrapper that records every call's duration under `name`.""" - - def wrapper(*args, **kwargs): - t0 = time.perf_counter() - try: - return fn(*args, **kwargs) - finally: - dt = time.perf_counter() - t0 - self.kernel_times[name].append(dt) - if group_by_decode_step and self._decode_step_idx >= 0: - bucket = self.kernel_per_decode_step[name] - while len(bucket) <= self._decode_step_idx: - bucket.append([]) - bucket[self._decode_step_idx].append(dt) - - return wrapper - def BeginDecodeStep(self) -> None: self._decode_step_idx += 1 @@ -145,24 +134,23 @@ def InstallProfiling(engine: LLMEngine, model_id: str, collector: _TimingCollect """ executor = engine._executor # type: ignore[attr-defined] compiled = executor._compiled[model_id] # type: ignore[attr-defined] + runner = executor._runners[model_id] # type: ignore[attr-defined] + kernel_names = { + id(compiled.prefill): ("kernel.prefill_fwd", False), + id(compiled.decode): ("kernel.decode_layer", True), + } - # Kernels are dispatched by Qwen314BModelRunner. - if hasattr(compiled.prefill, "chip_callable") or hasattr(compiled.prefill, "compiled"): - runner = executor._runners[model_id] # type: ignore[attr-defined] - orig_run_program = runner._run_distributed_program # type: ignore[attr-defined] - kernel_names = { - id(compiled.prefill): ("kernel.prefill_fwd", False), - id(compiled.decode): ("kernel.decode_layer", True), - } + def install_runner_kernel_timing(method_name: str) -> None: + orig_run = getattr(runner, method_name) - def timed_run_program(callable_spec, *args, **kwargs): + def timed_run(callable_spec, *args, **kwargs): kernel_info = kernel_names.get(id(callable_spec)) if kernel_info is None: - return orig_run_program(callable_spec, *args, **kwargs) + return orig_run(callable_spec, *args, **kwargs) name, group_by_decode_step = kernel_info t0 = time.perf_counter() try: - timing = orig_run_program(callable_spec, *args, **kwargs) + timing = orig_run(callable_spec, *args, **kwargs) finally: dt = time.perf_counter() - t0 collector.kernel_times[name].append(dt) @@ -174,14 +162,16 @@ def timed_run_program(callable_spec, *args, **kwargs): collector.RecordRunTiming(name, timing) return timing - runner._run_distributed_program = timed_run_program # type: ignore[attr-defined] + setattr(runner, method_name, timed_run) + + # A8W8 kernels are dispatched as L2 callables by Qwen314BA8W8ModelRunner. + if hasattr(compiled.prefill, "chip_callable"): + install_runner_kernel_timing("_run_l2_program") + # Original Qwen3-14B kernels are dispatched by the L3 model runner. + elif hasattr(compiled.prefill, "compiled"): + install_runner_kernel_timing("_run_distributed_program") else: - # Per-layer kernel wrappers. compiled.prefill / compiled.decode are invoked - # once per transformer layer inside run_prefill / run_decode respectively. - compiled.prefill = collector.WrapKernel(compiled.prefill, "kernel.prefill_fwd") - compiled.decode = collector.WrapKernel( - compiled.decode, "kernel.decode_layer", group_by_decode_step=True - ) + raise TypeError("unsupported compiled kernel wrapper for profiling") # Top-level executor API wrappers. orig_prefill = executor.run_prefill @@ -345,6 +335,15 @@ 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( + "--model-format", + default=_QWEN3_BF16_FORMAT, + choices=[_QWEN3_BF16_FORMAT, _QWEN3_A8W8_FORMAT], + help=( + "Qwen3-14B weight format. Use qwen3-14b for the original BF16/L3 path " + "or qwen3-a8w8 for compressed-tensors W8A8 checkpoints." + ), + ) parser.add_argument("--platform", default="a2a3", choices=["a2a3sim", "a2a3", "a5sim", "a5"]) parser.add_argument("--device-id", type=int, default=0, help="Default NPU device id when --devices is unset.") parser.add_argument( @@ -368,6 +367,7 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--max-seq-len", type=int, default=4096) parser.add_argument("--max-new-tokens", type=int, default=32) + parser.add_argument("--max-batch-size", type=int, default=16) parser.add_argument("--max-num-seqs", type=int, default=16, help="Max batch size / concurrent requests.") parser.add_argument("--block-size", type=int, default=128, help="KV cache page size.") parser.add_argument( @@ -386,18 +386,31 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--dtype", default="bfloat16", help="Weight data type.") parser.add_argument("--kv-cache-dtype", default="bfloat16", help="KV cache data type.") + parser.add_argument( + "--decode-backend", + default="a8w8", + choices=["a8w8"], + help="For qwen3-a8w8 only: run the A8W8 prefill/decode backend.", + ) parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--top-p", type=float, default=1.0) parser.add_argument("--top-k", type=int, default=None) parser.add_argument("--stream", action="store_true", default=False) parser.add_argument("--save-kernels-dir", default=None) + parser.add_argument( + "--pto-isa-commit", + default=None, + help="For qwen3-a8w8 only: pin PyPTO compile/assemble to the installed runtime's pto-isa revision.", + ) parser.add_argument( "--num-layers-override", type=int, default=None, - help="Validation knob: truncate the loaded model to N transformer " + help="BF16/L3 validation knob: truncate the loaded model to N transformer " "layers before compile/dispatch. Used to reduce HBM footprint " - "while validating Qwen3-14B kernels on memory-constrained devices.", + "while validating Qwen3-14B kernels on memory-constrained devices. " + "Unsupported for qwen3-a8w8 because its fused decode kernel is " + "compiled for the full fixed layer count.", ) parser.add_argument( "--profile", @@ -413,12 +426,43 @@ def build_parser() -> argparse.ArgumentParser: return parser +def _model_loader_for_format(model_format: str) -> ModelLoader | None: + if model_format != _QWEN3_A8W8_FORMAT: + return None + model_loader = ModelLoader() + model_loader.register(Qwen3A8W8DirectoryLoader()) + return model_loader + + +def _executor_class_for_format(model_format: str): + if model_format == _QWEN3_A8W8_FORMAT: + return Qwen314BA8W8PyptoExecutor + if model_format == _QWEN3_BF16_FORMAT: + return Qwen314BPyptoExecutor + raise ValueError(f"unsupported model_format: {model_format!r}") + + +def _validate_generation_args(args: argparse.Namespace) -> None: + if args.model_format != _QWEN3_A8W8_FORMAT: + return + if args.num_layers_override is not None: + raise ValueError("--num-layers-override is not supported for qwen3-a8w8 fused decode kernels") + if args.tensor_parallel_size != 1: + raise ValueError("qwen3-a8w8 currently requires --tp 1") + + +def _validate_a8w8_device_group(model_format: str, device_ids: tuple[int, ...]) -> None: + if model_format == _QWEN3_A8W8_FORMAT and len(device_ids) != 1: + raise ValueError("qwen3-a8w8 currently requires exactly one device") + + def main() -> None: args = build_parser().parse_args() import logging logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") for _n in ("simpler_setup", "pypto", "simpler"): logging.getLogger(_n).setLevel(logging.WARNING) + _validate_generation_args(args) get_profiler(process_name="npu_generate") model_dir = Path(args.model_dir).resolve() if not model_dir.is_dir(): @@ -437,16 +481,25 @@ def main() -> None: "serving replicas; data_parallel_size must be 1" ) device_ids = parallel_config.replica_device_groups[0] + _validate_a8w8_device_group(args.model_format, device_ids) kv_cache_manager = KvCacheManager() - executor = PyptoExecutor( + executor_cls = _executor_class_for_format(args.model_format) + executor_kwargs = { + "platform": args.platform, + "device_ids": device_ids, + "save_kernels_dir": args.save_kernels_dir, + "l3_trace": args.profile_verbose, + } + if args.model_format == _QWEN3_A8W8_FORMAT: + executor_kwargs["pto_isa_commit"] = args.pto_isa_commit + executor = executor_cls( kv_cache_manager, - platform=args.platform, - device_ids=device_ids, - save_kernels_dir=args.save_kernels_dir, - l3_trace=args.profile_verbose, + **executor_kwargs, ) + model_loader = _model_loader_for_format(args.model_format) engine = LLMEngine( + model_loader=model_loader, kv_cache_manager=kv_cache_manager, executor=executor, ) @@ -459,21 +512,19 @@ def main() -> None: engine.init_model( model_id=args.model_id, model_dir=str(model_dir), - model_format="huggingface", + model_format=_QWEN3_A8W8_FORMAT if args.model_format == _QWEN3_A8W8_FORMAT else "huggingface", + decode_backend=args.decode_backend, runtime_config=RuntimeConfig( - page_size=args.block_size, - max_batch_size=args.max_num_seqs, + page_size=128 if args.model_format == _QWEN3_A8W8_FORMAT else args.block_size, + max_batch_size=args.max_batch_size if args.model_format == _QWEN3_A8W8_FORMAT else args.max_num_seqs, max_seq_len=args.max_seq_len, max_new_tokens=args.max_new_tokens, device="cpu", - kv_dtype=args.kv_cache_dtype, - weight_dtype=args.dtype, + kv_dtype="int8" if args.model_format == _QWEN3_A8W8_FORMAT else args.kv_cache_dtype, + weight_dtype="bfloat16" if args.model_format == _QWEN3_A8W8_FORMAT else args.dtype, npu_memory_utilization=args.npu_memory_utilization, max_num_batched_tokens=args.max_num_batched_tokens, - # Conservative default — the decode kernel is compiled - # with this baked-in shape and cannot be resized later. - # 200 pages x 128 tokens = 25 600 tokens total capacity. - total_kv_pages=200, + total_kv_pages=None if args.model_format == _QWEN3_A8W8_FORMAT else 200, ), ) if collector is not None: diff --git a/examples/model/qwen3_14b/runner/npu_executor_a8w8.py b/examples/model/qwen3_14b/runner/npu_executor_a8w8.py new file mode 100644 index 0000000..32c4b1f --- /dev/null +++ b/examples/model/qwen3_14b/runner/npu_executor_a8w8.py @@ -0,0 +1,534 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import importlib.util +import sys +import tempfile +from collections.abc import Sequence +from pathlib import Path + +import torch + +from examples.model.qwen3_14b.runner.npu_runner_a8w8 import ( + _CompiledKernels, + _KernelLayerWeights, + _L2Callable, + Qwen314BA8W8ModelRunner, +) +from examples.model.qwen3_14b.runner.npu_executor import ( + _PYPTO_LIB_QWEN14B_DIR, + _QWEN14B_BLOCK_DIM, + _VOCAB_PAD_MULTIPLE, + Qwen314BPyptoExecutor, +) +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.utils import rope_tables, round_up + + +_QWEN14B_A8W8_PREFILL_CHUNK_LAYERS = 10 + + +def _patch_aicore_bitcast_helpers(work_dir: Path) -> None: + """Mark ptoas-generated bitcast helpers as AICore-callable for A8W8 JIT builds.""" + needle = "static inline To ptoas_bitcast(From from) {" + replacement = "static __aicore__ inline To ptoas_bitcast(From from) {" + patched = 0 + for cpp in work_dir.rglob("*.cpp"): + try: + text = cpp.read_text() + except UnicodeDecodeError: + continue + if needle not in text: + continue + cpp.write_text(text.replace(needle, replacement)) + patched += 1 + if patched: + print(f"[RUN] patched {patched} ptoas_bitcast helper(s) for aicore compilation", flush=True) + + +def _ensure_pypto_task_id_alias() -> None: + """Install the legacy TASK_ID alias required by older pypto-lib kernels.""" + import pypto.language as pl_mod # noqa: PLC0415 + + if not hasattr(pl_mod, "TASK_ID"): + # Keep this process-wide: kernel functions resolve pl.TASK_ID later when + # JIT compilation runs, not only while this module is imported. + pl_mod.TASK_ID = pl_mod.INDEX + + +def _load_pypto_lib_qwen14b_module(module_name: str) -> object: + """Load a Qwen3-14B kernel module from the pypto-lib submodule.""" + module_path = _PYPTO_LIB_QWEN14B_DIR / f"qwen3_14b_{module_name}.py" + if not module_path.is_file(): + module_path = _PYPTO_LIB_QWEN14B_DIR / f"{module_name}.py" + if not module_path.is_file(): + raise FileNotFoundError( + f"Missing pypto-lib Qwen3-14B kernel module: {module_path}. " + "Run `git submodule update --init --recursive`." + ) + spec = importlib.util.spec_from_file_location( + f"_pypto_lib_qwen3_14b_{module_name}", + module_path, + ) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load pypto-lib kernel module from {module_path}") + module = importlib.util.module_from_spec(spec) + sys.path.insert(0, str(_PYPTO_LIB_QWEN14B_DIR)) + try: + _ensure_pypto_task_id_alias() + spec.loader.exec_module(module) + finally: + try: + sys.path.remove(str(_PYPTO_LIB_QWEN14B_DIR)) + except ValueError: + pass + return module + + +class Qwen314BA8W8PyptoExecutor(CorePyptoExecutor): + """PyPTO executor that compiles and registers the Qwen3-14B kernels.""" + + _shared_tensor = staticmethod(Qwen314BPyptoExecutor._shared_tensor) + _validate_supported_shape = staticmethod(Qwen314BPyptoExecutor._validate_supported_shape) + _validate_total_kv_pages = staticmethod(Qwen314BPyptoExecutor._validate_total_kv_pages) + + def __init__( + self, + kv_cache_manager=None, + *, + platform: str = "a2a3sim", + device_ids: Sequence[int] = (0,), + save_kernels_dir: str | None = None, + pto_isa_commit: str | None = None, + l3_trace: bool = False, + ) -> None: + super().__init__( + kv_cache_manager, + platform=platform, + device_ids=device_ids, + save_kernels_dir=save_kernels_dir, + pto_isa_commit=pto_isa_commit, + ) + self._l3_trace = l3_trace + self._l2_compile_root: Path | None = None + + @property + def profile_verbose(self) -> bool: + """Return whether compile and L3 execution timing logs are enabled.""" + return self._l3_trace + + def _create_runner(self, model_id: str, compiled: object) -> ModelRunner: + """Create the Qwen3-14B runtime runner for compiled kernels.""" + if not isinstance(compiled, _CompiledKernels): + raise TypeError("Qwen314BA8W8PyptoExecutor requires Qwen3-14B compiled kernels.") + return Qwen314BA8W8ModelRunner( + model_id=model_id, + compiled=compiled, + platform=self._platform, + device_id=self._device_ids[0], + save_kernels_dir=self._save_kernels_dir, + ) + + def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: + """Compile Qwen3-14B PyPTO kernels and pack runtime artifacts.""" + timer = StageTimer( + enabled=self._l3_trace, + prefix="compile-breakdown", + title="_compile_model stage timings", + ) + + def _mark(label: str) -> None: + timer.mark(label) + + quantization = self._model_quantization(model) + if quantization != "a8w8": + raise ValueError("Qwen314BA8W8PyptoExecutor only supports qwen3-a8w8 loaded models.") + qwen3_prefill_fwd = _load_pypto_lib_qwen14b_module("prefill_fwd_a8w8") + # The fused all-layer decode lives in decode_layer.decode_fwd (the + # standalone decode_fwd.py module was removed in pypto-lib). It is now + # PAGED: it consumes block_table + slot_mapping and reads/writes the SAME + # device-resident paged KV pool prefill writes (self._kv_caches), so no + # contiguous bridge / MAX_SEQ env is needed. + qwen3_decode_layer = _load_pypto_lib_qwen14b_module("decode_layer_a8w8") + _mark("imports") + + self._validate_supported_shape(model) + kernel_batch = model.runtime.max_batch_size + if int(qwen3_decode_layer.BATCH) != kernel_batch: + raise ValueError( + "decode_layer.decode_fwd is compiled for a fixed kernel BATCH of " + f"{int(qwen3_decode_layer.BATCH)}, but runtime max_batch_size is " + f"{kernel_batch}; they must match (decode statically computes and " + "writes BATCH rows / BATCH logit rows)." + ) + if int(model.config.num_hidden_layers) != int(qwen3_decode_layer.NUM_LAYERS): + raise ValueError( + "decode_layer.decode_fwd fuses a FIXED " + f"NUM_LAYERS={int(qwen3_decode_layer.NUM_LAYERS)} loop (the layer count " + "is a kernel constant, not derived from the weight tensors), but the " + f"model has {model.config.num_hidden_layers} layers. The fused decode " + "does not support --num-layers-override; run the full model." + ) + self._validate_kernel_runtime_contract(model, qwen3_prefill_fwd, qwen3_decode_layer) + self._validate_total_kv_pages(model, kernel_batch) + + padded_vocab = round_up(model.config.vocab_size, _VOCAB_PAD_MULTIPLE) + kernel_vocab = getattr(qwen3_decode_layer, "VOCAB", padded_vocab) + if padded_vocab != int(kernel_vocab): + raise ValueError( + f"decode_layer.decode_fwd hard-codes VOCAB={int(kernel_vocab)} " + f"(config.VOCAB) for its fused LM head, but the runtime padded vocab is " + f"{padded_vocab} (round_up({model.config.vocab_size}, {_VOCAB_PAD_MULTIPLE})); " + "they must match for the decode logits buffer / lm_head weight to line up." + ) + page_size = model.runtime.page_size + max_blocks_per_seq = (model.runtime.max_seq_len + page_size - 1) // page_size + prefill = self._compile_prefill_fwd_callable_a8w8( + qwen3_prefill_fwd.prefill_hidden_a8w8, + batch=kernel_batch, + max_seq=model.runtime.max_seq_len, + hidden_size=model.config.hidden_size, + intermediate_size=model.config.intermediate_size, + num_heads=model.config.num_attention_heads, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + num_layers=min(model.config.num_hidden_layers, _QWEN14B_A8W8_PREFILL_CHUNK_LAYERS), + vocab_size=padded_vocab, + block_table_stride=max_blocks_per_seq, + page_size=page_size, + ) + _mark("compile_prefill") + decode = self._compile_decode_fwd_callable_a8w8( + qwen3_decode_layer.decode_fwd, + batch=kernel_batch, + max_seq=model.runtime.max_seq_len, + block_table_stride=max_blocks_per_seq, + hidden_size=model.config.hidden_size, + intermediate_size=model.config.intermediate_size, + num_heads=model.config.num_attention_heads, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + num_layers=model.config.num_hidden_layers, + vocab_size=padded_vocab, + page_size=page_size, + ) + _mark("compile_decode") + + rope_cos_raw, rope_sin_raw = rope_tables( + model.runtime.max_seq_len, + model.config.head_dim, + model.config.rope_theta, + ) + rope_cos = self._shared_tensor(rope_cos_raw) + rope_sin = self._shared_tensor(rope_sin_raw) + + _mark("rope_tables") + + lm_head_weight = model.lm_head + if padded_vocab != lm_head_weight.shape[0]: + pad_rows = padded_vocab - lm_head_weight.shape[0] + padding = torch.zeros( + (pad_rows, lm_head_weight.shape[1]), + dtype=lm_head_weight.dtype, + device=lm_head_weight.device, + ) + lm_head_weight = torch.cat([lm_head_weight, padding], dim=0) + padded_lm_head_weight = self._shared_tensor(lm_head_weight.to(torch.bfloat16).contiguous().cpu()) + _mark("pad_lm_head") + layers = [] + for layer in model.layers: + layers.append(self._kernel_layer_weights(layer)) + self._release_layer_weights(layer) + final_norm_weight = self._shared_tensor(model.final_norm_weight.view(1, -1).float().cpu()) + _mark("kernel_layer_weights") + + decode_weights = self._stack_decode_weights(layers) + layers.clear() + _mark("stack_decode_weights") + decode_logits_buffer = torch.empty( + (kernel_batch, padded_vocab), + dtype=torch.float32, + ).share_memory_() + _mark("decode_logits_buffer") + + timer.report() + + return _CompiledKernels( + prefill=prefill, + decode=decode, + final_norm_weight=final_norm_weight, + rope_cos=rope_cos, + rope_sin=rope_sin, + padded_vocab=padded_vocab, + padded_lm_head_weight=padded_lm_head_weight, + decode_weights=decode_weights, + decode_logits_buffer=decode_logits_buffer, + ) + + def _compile_prefill_fwd_callable_a8w8( + self, + jit_fn: object, + *, + batch: int, + max_seq: int, + block_table_stride: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_layers: int, + vocab_size: int, + page_size: int, + ) -> _L2Callable: + """Compile the A8W8 all-layer prefill kernel into an L2 callable.""" + kv_hidden = num_kv_heads * head_dim + total_tokens = batch * max_seq + runtime_cache_blocks = (max_seq + page_size - 1) // page_size + cache_rows = batch * runtime_cache_blocks * num_layers * num_kv_heads * page_size + dummy_args = [ + torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16), + torch.empty((batch,), dtype=torch.int32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.int8), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.int8), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.int8), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers, kv_hidden), dtype=torch.float32), + torch.empty((num_layers, kv_hidden), dtype=torch.float32), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((batch * block_table_stride,), dtype=torch.int32), + torch.empty((total_tokens,), dtype=torch.int32), + torch.empty((cache_rows, head_dim), dtype=torch.int8), + torch.empty((cache_rows, head_dim), dtype=torch.int8), + torch.empty((cache_rows, 8), dtype=torch.float32), + torch.empty((cache_rows, 8), dtype=torch.float32), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.int8), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), + torch.empty((1, hidden_size), dtype=torch.float32), + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), + torch.empty((batch, vocab_size), dtype=torch.float32), + torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16), + ] + return self._compile_jit_fwd_callable("prefill_hidden_a8w8", jit_fn, dummy_args) + + def _compile_decode_fwd_callable_a8w8( + self, + jit_fn: object, + *, + batch: int, + max_seq: int, + block_table_stride: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_layers: int, + vocab_size: int, + page_size: int, + ) -> _L2Callable: + """Compile the A8W8 fused all-layer paged decode kernel.""" + kv_hidden = num_kv_heads * head_dim + runtime_cache_blocks = (max_seq + page_size - 1) // page_size + cache_rows = num_layers * batch * runtime_cache_blocks * num_kv_heads * page_size + dummy_args = [ + torch.empty((batch, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.int8), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.int8), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.int8), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers, kv_hidden), dtype=torch.float32), + torch.empty((num_layers, kv_hidden), dtype=torch.float32), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((batch * block_table_stride,), dtype=torch.int32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((cache_rows, head_dim), dtype=torch.int8), + torch.empty((cache_rows, head_dim), dtype=torch.int8), + torch.empty((cache_rows, 8), dtype=torch.float32), + torch.empty((cache_rows, 8), dtype=torch.float32), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.int8), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((1, hidden_size), dtype=torch.float32), + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), + torch.empty((batch, vocab_size), dtype=torch.float32), + ] + return self._compile_jit_fwd_callable("decode_fwd_a8w8", jit_fn, dummy_args) + + def _compile_jit_fwd_callable( + self, + name: str, + jit_fn: object, + dummy_args: list[torch.Tensor], + ) -> _L2Callable: + """Compile a top-level ``@pl.jit`` kernel into an L2 callable.""" + from pypto.runtime.device_runner import compile_and_assemble # noqa: PLC0415 + from pypto.runtime.runner import _patch_orchestration_headers # noqa: PLC0415 + + config = self._run_config(codegen_only=True) + compiled = jit_fn.compile(*dummy_args, config=config) + work_dir = Path(compiled.output_dir) + _patch_orchestration_headers(work_dir) + _patch_aicore_bitcast_helpers(work_dir) + assembled = compile_and_assemble( + work_dir, + self._platform, + pto_isa_commit=config.pto_isa_commit, + ) + if len(assembled) == 2: + chip_callable, runtime_name = assembled + runtime_config = {} + else: + chip_callable, runtime_name, runtime_config = assembled + runtime_config = runtime_config or {} + param_infos, _, _ = compiled._get_metadata() + return _L2Callable( + chip_callable=chip_callable, + name=name, + runtime_name=runtime_name, + block_dim=int(runtime_config.get("block_dim", _QWEN14B_BLOCK_DIM)), + aicpu_thread_num=int(runtime_config.get("aicpu_thread_num", 4)), + param_infos=tuple(param_infos), + ) + + def _l2_work_dir(self, name: str) -> Path: + """Return a dedicated compile directory for one non-L3 program.""" + if self._save_kernels_dir is not None: + root = Path(self._save_kernels_dir) + else: + if self._l2_compile_root is None: + self._l2_compile_root = Path(tempfile.mkdtemp(prefix="qwen3_14b_l2_")) + root = self._l2_compile_root + work_dir = root / name + work_dir.mkdir(parents=True, exist_ok=True) + return work_dir + + @staticmethod + def _validate_kernel_runtime_contract( + model: RuntimeModel, + qwen3_prefill_fwd: object, + qwen3_decode_layer: object, + ) -> None: + """Validate runtime dimensions baked into the A8W8 L2 kernels.""" + expected_page_size = int(qwen3_decode_layer.BLOCK_SIZE) + if int(model.runtime.page_size) != expected_page_size: + raise ValueError( + f"qwen3-a8w8 requires page_size={expected_page_size} " + f"to match decode BLOCK_SIZE, got {model.runtime.page_size}" + ) + max_supported_seq = int(qwen3_prefill_fwd.MAX_SEQ) + if int(model.runtime.max_seq_len) > max_supported_seq: + raise ValueError( + f"qwen3-a8w8 supports max_seq_len <= {max_supported_seq}, " + f"got {model.runtime.max_seq_len}" + ) + + @classmethod + def _stack_decode_weights(cls, layers: list[_KernelLayerWeights]) -> dict[str, torch.Tensor]: + """Stack per-layer weights into fused decode-kernel tensors and release sources.""" + # Stack from already-prepared per-layer kernel weights. Each + # _KernelLayerWeights field is already in the kernel-ready shape/dtype + # (transposed bf16 cpu for projections, [1, N] float cpu for norms), + # so a plain cat along dim 0 is all that's left. Reading from the + # original model.layers here would crash because _release_layer_weights + # has already replaced those tensors with torch.empty(0). + def cat(attr: str) -> torch.Tensor: + stacked = cls._shared_tensor(torch.cat([getattr(l, attr) for l in layers], dim=0).contiguous()) + for layer in layers: + setattr(layer, attr, torch.empty(0)) + return stacked + + weights = { + "decode_input_rms_weight": cat("input_rms_weight"), + "decode_wq": cat("wq"), + "decode_wk": cat("wk"), + "decode_wv": cat("wv"), + "decode_q_norm_weight": cat("q_norm_weight"), + "decode_k_norm_weight": cat("k_norm_weight"), + "decode_wo": cat("wo"), + "decode_post_rms_weight": cat("post_rms_weight"), + "decode_w_gate": cat("w_gate"), + "decode_w_up": cat("w_up"), + "decode_w_down": cat("w_down"), + } + if layers and layers[0].wq_scale is not None: + weights.update( + { + "decode_wq_scale": cat("wq_scale"), + "decode_wk_scale": cat("wk_scale"), + "decode_wv_scale": cat("wv_scale"), + "decode_wo_scale": cat("wo_scale"), + } + ) + return weights + + @classmethod + def _kernel_layer_weights(cls, layer) -> _KernelLayerWeights: + """Convert one Hugging Face layer into kernel-ready weight tensors.""" + return _KernelLayerWeights( + input_rms_weight=cls._shared_tensor(layer.input_rms_weight.float().cpu()), + wq=cls._shared_tensor(layer.wq.contiguous().cpu()), + wk=cls._shared_tensor(layer.wk.contiguous().cpu()), + wv=cls._shared_tensor(layer.wv.contiguous().cpu()), + q_norm_weight=cls._shared_tensor(layer.q_norm_weight.float().cpu()), + k_norm_weight=cls._shared_tensor(layer.k_norm_weight.float().cpu()), + wo=cls._shared_tensor(layer.wo.contiguous().cpu()), + post_rms_weight=cls._shared_tensor(layer.post_rms_weight.float().cpu()), + w_gate=cls._shared_tensor(layer.w_gate.contiguous().cpu()), + w_up=cls._shared_tensor(layer.w_up.contiguous().cpu()), + w_down=cls._shared_tensor(layer.w_down.contiguous().cpu()), + wq_scale=cls._shared_tensor(layer.wq_scale.float().cpu()), + wk_scale=cls._shared_tensor(layer.wk_scale.float().cpu()), + wv_scale=cls._shared_tensor(layer.wv_scale.float().cpu()), + wo_scale=cls._shared_tensor(layer.wo_scale.float().cpu()), + ) + + @staticmethod + def _release_layer_weights(layer) -> None: + """Drop original layer tensors after kernel-ready copies are built.""" + Qwen314BPyptoExecutor._release_layer_weights(layer) + if hasattr(layer, "wq_scale"): + empty = torch.empty(0) + layer.wq_scale = empty + layer.wk_scale = empty + layer.wv_scale = empty + layer.wo_scale = empty + + @staticmethod + def _model_quantization(model: RuntimeModel) -> str: + """Return the quantization mode declared by loaded layer weights.""" + if not model.layers: + return "bf16" + quantization = getattr(model.layers[0], "quantization", None) + if quantization == "a8w8": + return str(quantization) + return "bf16" diff --git a/examples/model/qwen3_14b/runner/npu_runner_a8w8.py b/examples/model/qwen3_14b/runner/npu_runner_a8w8.py index 578b6fb..11d918a 100644 --- a/examples/model/qwen3_14b/runner/npu_runner_a8w8.py +++ b/examples/model/qwen3_14b/runner/npu_runner_a8w8.py @@ -88,7 +88,6 @@ class _CompiledKernels: rope_sin: torch.Tensor padded_vocab: int padded_lm_head_weight: torch.Tensor - layers: list[_KernelLayerWeights] decode_weights: dict[str, torch.Tensor] decode_logits_buffer: torch.Tensor diff --git a/tests/test_batching.py b/tests/test_batching.py index 850682e..e9014b1 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -8,6 +8,9 @@ # ----------------------------------------------------------------------------------------------------------- import asyncio +import argparse +import json +import struct from pathlib import Path from types import SimpleNamespace @@ -33,6 +36,10 @@ ) from python.core.scheduler import Request, ScheduledRequest from python.core.serving_worker import WorkerProcess +from examples.model.qwen3_14b.npu_generate import _validate_generation_args +from examples.model.qwen3_14b.runner.a8w8_loader import Qwen3A8W8DirectoryLoader +from examples.model.qwen3_14b.runner.a8w8_loader import _SafeTensorIndex +from examples.model.qwen3_14b.runner.npu_executor_a8w8 import Qwen314BA8W8PyptoExecutor from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor as PyptoExecutor from examples.model.qwen3_14b.runner.npu_runner import Qwen314BModelRunner as ModelRunner from examples.model.qwen3_14b.runner.npu_runner import _CompiledKernels @@ -40,6 +47,8 @@ from examples.model.qwen3_14b.runner.npu_runner import _add_run_timing_args from examples.model.qwen3_14b.runner.npu_runner import _kernel_trace_name from examples.model.qwen3_14b.runner.npu_runner import _run_timing_us +from examples.model.qwen3_14b.runner.npu_runner_a8w8 import _KernelLayerWeights +from examples.model.qwen3_14b.runner.npu_runner_a8w8 import Qwen314BA8W8ModelRunner from python.runtime.worker import WorkerTensor @@ -48,6 +57,24 @@ QWEN3_KERNEL_DIR = ROOT / "pypto-lib" / "models" / "qwen3" / "14b" +def _write_safetensor(path: Path, name: str, tensor: torch.Tensor) -> None: + raw = tensor.contiguous().numpy().tobytes() + dtype_names = { + torch.int8: "I8", + torch.float32: "F32", + torch.bfloat16: "BF16", + } + header = { + name: { + "dtype": dtype_names[tensor.dtype], + "shape": list(tensor.shape), + "data_offsets": [0, len(raw)], + } + } + header_bytes = json.dumps(header).encode("utf-8") + path.write_bytes(struct.pack(" list[int]: return [max(1, len(text))] @@ -336,6 +363,119 @@ def test_decode_inputs_use_actual_user_batch_without_padding_lanes(): assert prepared.slot_mapping.tolist() == [manager.slot_mapping_for_request(alloc)] +def test_a8w8_decode_inputs_use_actual_user_batch_without_padding_lanes(): + model = _model(max_batch_size=16) + manager = KvCacheManager() + manager.register_model(model.config.model_id, model.config, model.runtime) + runner = Qwen314BA8W8ModelRunner( + compiled=None, # type: ignore[arg-type] + ) + alloc = manager.allocate_for_prompt(model.config.model_id, "req-0", 1) + hidden_states = torch.ones(1, model.config.hidden_size) + + prepared = runner._prepare_decode_inputs( + model, + DecodeBatch( + request_ids=[alloc.request_id], + token_ids=torch.zeros(1, 1, dtype=torch.long), + hidden_states=hidden_states, + seq_lens=torch.tensor([1], dtype=torch.int32), + kv_allocations=[alloc], + ), + ) + + assert prepared.actual_batch == 1 + assert prepared.hidden.shape == (1, model.config.hidden_size) + assert prepared.seq_lens.tolist() == [1] + assert prepared.block_table.shape == (2,) + assert prepared.block_table[0].item() == alloc.page_ids[0] + assert prepared.slot_mapping.tolist() == [manager.slot_mapping_for_request(alloc)] + + +def test_a8w8_init_kv_cache_returns_page_count_for_first_and_repeated_init(monkeypatch): + model = _model(max_batch_size=16, max_seq_len=128, page_size=64) + runner = Qwen314BA8W8ModelRunner( + compiled=None, # type: ignore[arg-type] + ) + allocated_shapes = [] + + def alloc_kv_cache_tensor(shape, dtype): + allocated_shapes.append((shape, dtype)) + return WorkerTensor( + data_ptr=len(allocated_shapes), + shape=shape, + dtype=_FakeWorker._DTYPES[dtype], + ) + + monkeypatch.setattr(runner, "_alloc_kv_cache_tensor", alloc_kv_cache_tensor) + monkeypatch.setattr(runner, "_free_kv_cache_tensor", lambda tensor: None) + + first_pages = runner.init_kv_cache(model.config.model_id, model.config, model.runtime) + second_pages = runner.init_kv_cache(model.config.model_id, model.config, model.runtime) + + assert first_pages == 32 + assert second_pages == 32 + cache_rows = model.config.num_hidden_layers * 32 * model.config.num_key_value_heads * model.runtime.page_size + assert allocated_shapes == [ + ((cache_rows, model.config.head_dim), torch.bfloat16), + ((cache_rows, model.config.head_dim), torch.bfloat16), + ((cache_rows, 8), torch.float32), + ((cache_rows, 8), torch.float32), + ] + + +def test_a8w8_runtime_contract_rejects_mismatched_page_size(): + model = _model(max_batch_size=16, max_seq_len=128, page_size=64) + + with pytest.raises(ValueError, match="page_size=128"): + Qwen314BA8W8PyptoExecutor._validate_kernel_runtime_contract( + model, + SimpleNamespace(MAX_SEQ=128), + SimpleNamespace(BLOCK_SIZE=128), + ) + + +def test_a8w8_runtime_contract_rejects_excessive_max_seq_len(): + model = _model(max_batch_size=16, max_seq_len=256, page_size=128) + + with pytest.raises(ValueError, match="max_seq_len <= 128"): + Qwen314BA8W8PyptoExecutor._validate_kernel_runtime_contract( + model, + SimpleNamespace(MAX_SEQ=128), + SimpleNamespace(BLOCK_SIZE=128), + ) + + +def test_a8w8_stack_decode_weights_releases_per_layer_sources(): + def layer(value: float) -> _KernelLayerWeights: + return _KernelLayerWeights( + input_rms_weight=torch.full((1, 2), value), + wq=torch.full((2, 2), value), + wk=torch.full((2, 1), value), + wv=torch.full((2, 1), value), + q_norm_weight=torch.full((1, 1), value), + k_norm_weight=torch.full((1, 1), value), + wo=torch.full((2, 2), value), + post_rms_weight=torch.full((1, 2), value), + w_gate=torch.full((2, 3), value), + w_up=torch.full((2, 3), value), + w_down=torch.full((3, 2), value), + wq_scale=torch.full((1, 2), value), + wk_scale=torch.full((1, 1), value), + wv_scale=torch.full((1, 1), value), + wo_scale=torch.full((1, 2), value), + ) + + layers = [layer(1.0), layer(2.0)] + + weights = Qwen314BA8W8PyptoExecutor._stack_decode_weights(layers) + + assert weights["decode_wq"].shape == (4, 2) + assert weights["decode_wq_scale"].shape == (2, 2) + assert all(layer_weights.wq.numel() == 0 for layer_weights in layers) + assert all(layer_weights.wq_scale is not None and layer_weights.wq_scale.numel() == 0 for layer_weights in layers) + + def test_decode_kernel_inputs_reject_multi_token_rows(): model = _model(max_batch_size=2) runner = ModelRunner(compiled=_compiled_kernels(model)) @@ -509,6 +649,54 @@ def test_engine_ignores_device_sampled_tokens_for_non_greedy_config(): assert sampler.sample_calls == 1 +def test_a8w8_loader_detects_unindexed_safetensors(tmp_path): + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + _write_safetensor(tmp_path / "weights.safetensors", "tensor", torch.tensor([[1, 2]], dtype=torch.int8)) + + loader = Qwen3A8W8DirectoryLoader() + index = _SafeTensorIndex(tmp_path) + + assert loader.can_load(tmp_path) + assert torch.equal(index.load("tensor"), torch.tensor([[1, 2]], dtype=torch.int8)) + + +def test_a8w8_loader_uses_index_shard_names(tmp_path): + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + _write_safetensor(tmp_path / "custom-shard.safetensors", "tensor", torch.tensor([3], dtype=torch.int8)) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"tensor": "custom-shard.safetensors"}}), + encoding="utf-8", + ) + + loader = Qwen3A8W8DirectoryLoader() + index = _SafeTensorIndex(tmp_path) + + assert loader.can_load(tmp_path) + assert torch.equal(index.load("tensor"), torch.tensor([3], dtype=torch.int8)) + + +def test_a8w8_num_layers_override_fails_fast(): + with pytest.raises(ValueError, match="num-layers-override"): + _validate_generation_args( + argparse.Namespace( + model_format="qwen3-a8w8", + num_layers_override=1, + tensor_parallel_size=1, + ) + ) + + +def test_a8w8_tensor_parallel_fails_fast(): + with pytest.raises(ValueError, match="requires --tp 1"): + _validate_generation_args( + argparse.Namespace( + model_format="qwen3-a8w8", + num_layers_override=None, + tensor_parallel_size=2, + ) + ) + + def test_serving_worker_skips_decode_host_embedding_when_executor_embeds_on_device(): model = _model(max_batch_size=1, eos_token_id=0) manager = KvCacheManager() @@ -574,7 +762,7 @@ def test_pypto_executor_uses_cached_kernel_weights_after_registration(monkeypatc compiled=compiled, ) monkeypatch.setattr(runner, "_shared_l3_worker", lambda: _FakeWorker()) - monkeypatch.setattr(runner, "_compute_kv_cache_pages", lambda config, runtime, device_id=0: 1) + monkeypatch.setattr(runner, "_compute_kv_cache_pages", lambda config, runtime, device_id=None: 1) monkeypatch.setattr(runner, "_print_memory_breakdown", lambda *a, **kw: None) runner.init_kv_cache(model.config.model_id, model.config, model.runtime) monkeypatch.setattr(runner, "_static_device_tensor", lambda tensor: tensor) @@ -653,10 +841,12 @@ def test_decode_host_inlines_embedding_and_sampling_into_decode_fwd(): if not QWEN3_KERNEL_DIR.is_dir(): pytest.skip("pypto-lib submodule is not checked out") - decode_path = QWEN3_KERNEL_DIR / "decode_layer.py" - if not decode_path.is_file(): - decode_path = QWEN3_KERNEL_DIR / "decode_fwd.py" - decode_source = decode_path.read_text(encoding="utf-8") + decode_kernel = QWEN3_KERNEL_DIR / "decode_layer.py" + if not decode_kernel.is_file(): + decode_kernel = QWEN3_KERNEL_DIR / "decode_fwd.py" + if not decode_kernel.is_file(): + pytest.skip("pypto-lib decode kernel source is not checked out") + decode_source = decode_kernel.read_text(encoding="utf-8") assert 'name_hint="token_embed"' in decode_source assert 'name_hint="greedy_sample"' in decode_source From 91aae6894ccf872d7c7b1d22cfdad07d88bab465 Mon Sep 17 00:00:00 2001 From: vegetabledoww Date: Sun, 12 Jul 2026 19:09:34 -0700 Subject: [PATCH 4/4] Use L3 dispatch for Qwen3 A8W8 Route Qwen3-14B A8W8 prefill and decode through HOST-level L3 wrappers, preallocate shared IO buffers for the L3 runner, and remove the A8W8 L2 callable fallback from the serving path. Consolidate A8W8 HOST wrappers into the existing Qwen3 L3 dispatch module. Use a closure-based factory to keep A8W8 kernel dependencies isolated from the BF16 dispatch globals. Validation: Python compile and pre-commit checks passed; 8-token A8W8 generation produced expected text at 55.1 ms/token, and prior 48-token L3-only validation succeeded. --- examples/model/qwen3_14b/npu_generate.py | 7 +- .../qwen3_14b/runner/npu_executor_a8w8.py | 135 +++-- .../model/qwen3_14b/runner/npu_runner_a8w8.py | 535 +++++++++--------- .../qwen3_14b/runner/qwen3_l3_dispatch.py | 136 +++++ tests/test_batching.py | 5 +- 5 files changed, 512 insertions(+), 306 deletions(-) diff --git a/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index a9727b0..0e2b1d3 100644 --- a/examples/model/qwen3_14b/npu_generate.py +++ b/examples/model/qwen3_14b/npu_generate.py @@ -164,11 +164,8 @@ def timed_run(callable_spec, *args, **kwargs): setattr(runner, method_name, timed_run) - # A8W8 kernels are dispatched as L2 callables by Qwen314BA8W8ModelRunner. - if hasattr(compiled.prefill, "chip_callable"): - install_runner_kernel_timing("_run_l2_program") - # Original Qwen3-14B kernels are dispatched by the L3 model runner. - elif hasattr(compiled.prefill, "compiled"): + # Current Qwen3-14B kernels, including A8W8, are dispatched by the L3 model runner. + if hasattr(compiled.prefill, "compiled"): install_runner_kernel_timing("_run_distributed_program") else: raise TypeError("unsupported compiled kernel wrapper for profiling") diff --git a/examples/model/qwen3_14b/runner/npu_executor_a8w8.py b/examples/model/qwen3_14b/runner/npu_executor_a8w8.py index 32c4b1f..c40e476 100644 --- a/examples/model/qwen3_14b/runner/npu_executor_a8w8.py +++ b/examples/model/qwen3_14b/runner/npu_executor_a8w8.py @@ -11,7 +11,6 @@ import importlib.util import sys -import tempfile from collections.abc import Sequence from pathlib import Path @@ -20,7 +19,7 @@ from examples.model.qwen3_14b.runner.npu_runner_a8w8 import ( _CompiledKernels, _KernelLayerWeights, - _L2Callable, + _L3Callable, Qwen314BA8W8ModelRunner, ) from examples.model.qwen3_14b.runner.npu_executor import ( @@ -121,7 +120,6 @@ def __init__( pto_isa_commit=pto_isa_commit, ) self._l3_trace = l3_trace - self._l2_compile_root: Path | None = None @property def profile_verbose(self) -> bool: @@ -161,6 +159,12 @@ def _mark(label: str) -> None: # device-resident paged KV pool prefill writes (self._kv_caches), so no # contiguous bridge / MAX_SEQ env is needed. qwen3_decode_layer = _load_pypto_lib_qwen14b_module("decode_layer_a8w8") + from examples.model.qwen3_14b.runner import qwen3_l3_dispatch # noqa: PLC0415 + + prefill_jit, decode_jit = qwen3_l3_dispatch.create_qwen3_a8w8_dispatch( + qwen3_prefill_fwd.prefill_hidden_a8w8, + qwen3_decode_layer.decode_fwd, + ) _mark("imports") self._validate_supported_shape(model) @@ -195,7 +199,7 @@ def _mark(label: str) -> None: page_size = model.runtime.page_size max_blocks_per_seq = (model.runtime.max_seq_len + page_size - 1) // page_size prefill = self._compile_prefill_fwd_callable_a8w8( - qwen3_prefill_fwd.prefill_hidden_a8w8, + prefill_jit, batch=kernel_batch, max_seq=model.runtime.max_seq_len, hidden_size=model.config.hidden_size, @@ -210,7 +214,7 @@ def _mark(label: str) -> None: ) _mark("compile_prefill") decode = self._compile_decode_fwd_callable_a8w8( - qwen3_decode_layer.decode_fwd, + decode_jit, batch=kernel_batch, max_seq=model.runtime.max_seq_len, block_table_stride=max_blocks_per_seq, @@ -261,6 +265,37 @@ def _mark(label: str) -> None: dtype=torch.float32, ).share_memory_() _mark("decode_logits_buffer") + max_prefill_tokens = kernel_batch * model.runtime.max_seq_len + l3_buffers = { + "prefill_hidden_buffer": torch.empty( + (max_prefill_tokens, model.config.hidden_size), + dtype=torch.bfloat16, + ).share_memory_(), + "prefill_next_hidden_buffer": torch.empty( + (max_prefill_tokens, model.config.hidden_size), + dtype=torch.bfloat16, + ).share_memory_(), + "prefill_seq_lens_buffer": torch.empty((kernel_batch,), dtype=torch.int32).share_memory_(), + "prefill_chunk_lens_buffer": torch.empty((kernel_batch,), dtype=torch.int32).share_memory_(), + "prefill_chunk_offsets_buffer": torch.empty((kernel_batch,), dtype=torch.int32).share_memory_(), + "prefill_block_table_buffer": torch.empty( + (kernel_batch * max_blocks_per_seq,), + dtype=torch.int32, + ).share_memory_(), + "prefill_slot_mapping_buffer": torch.empty((max_prefill_tokens,), dtype=torch.int32).share_memory_(), + "prefill_logits_buffer": torch.empty((kernel_batch, padded_vocab), dtype=torch.float32).share_memory_(), + "decode_hidden_buffer": torch.empty( + (kernel_batch, model.config.hidden_size), + dtype=torch.bfloat16, + ).share_memory_(), + "decode_seq_lens_buffer": torch.empty((kernel_batch,), dtype=torch.int32).share_memory_(), + "decode_block_table_buffer": torch.empty( + (kernel_batch * max_blocks_per_seq,), + dtype=torch.int32, + ).share_memory_(), + "decode_slot_mapping_buffer": torch.empty((kernel_batch,), dtype=torch.int32).share_memory_(), + } + _mark("l3_buffers") timer.report() @@ -274,6 +309,7 @@ def _mark(label: str) -> None: padded_lm_head_weight=padded_lm_head_weight, decode_weights=decode_weights, decode_logits_buffer=decode_logits_buffer, + **l3_buffers, ) def _compile_prefill_fwd_callable_a8w8( @@ -291,8 +327,8 @@ def _compile_prefill_fwd_callable_a8w8( num_layers: int, vocab_size: int, page_size: int, - ) -> _L2Callable: - """Compile the A8W8 all-layer prefill kernel into an L2 callable.""" + ) -> _L3Callable: + """Compile the A8W8 all-layer prefill host wrapper into an L3 callable.""" kv_hidden = num_kv_heads * head_dim total_tokens = batch * max_seq runtime_cache_blocks = (max_seq + page_size - 1) // page_size @@ -347,8 +383,8 @@ def _compile_decode_fwd_callable_a8w8( num_layers: int, vocab_size: int, page_size: int, - ) -> _L2Callable: - """Compile the A8W8 fused all-layer paged decode kernel.""" + ) -> _L3Callable: + """Compile the A8W8 fused all-layer paged decode host wrapper.""" kv_hidden = num_kv_heads * head_dim runtime_cache_blocks = (max_seq + page_size - 1) // page_size cache_rows = num_layers * batch * runtime_cache_blocks * num_kv_heads * page_size @@ -389,56 +425,63 @@ def _compile_jit_fwd_callable( name: str, jit_fn: object, dummy_args: list[torch.Tensor], - ) -> _L2Callable: - """Compile a top-level ``@pl.jit`` kernel into an L2 callable.""" - from pypto.runtime.device_runner import compile_and_assemble # noqa: PLC0415 - from pypto.runtime.runner import _patch_orchestration_headers # noqa: PLC0415 + ) -> _L3Callable: + """Compile a top-level A8W8 HOST wrapper into an L3 callable.""" + return self._compile_l3_jit_fwd_callable(name, jit_fn, dummy_args) + + def _compile_l3_jit_fwd_callable( + self, + name: str, + jit_fn: object, + dummy_args: list[torch.Tensor], + ) -> _L3Callable: + """Compile a HOST wrapper into a PyPTO DistributedCompiledProgram.""" + from pypto.ir.distributed_compiled_program import DistributedCompiledProgram # noqa: PLC0415 + from pypto.ir.distributed_compiled_program import DistributedConfig # noqa: PLC0415 + from pypto.runtime import RunConfig # noqa: PLC0415 config = self._run_config(codegen_only=True) - compiled = jit_fn.compile(*dummy_args, config=config) - work_dir = Path(compiled.output_dir) - _patch_orchestration_headers(work_dir) - _patch_aicore_bitcast_helpers(work_dir) - assembled = compile_and_assemble( - work_dir, - self._platform, + distributed_config = DistributedConfig( + device_ids=list(self._device_ids), + num_sub_workers=0, + block_dim=_QWEN14B_BLOCK_DIM, + aicpu_thread_num=4, + ) + run_config = RunConfig( + platform=config.platform, + device_id=config.device_id, + backend_type=config.backend_type, + strategy=config.strategy, + dump_passes=config.dump_passes, + save_kernels=config.save_kernels, + save_kernels_dir=config.save_kernels_dir, + codegen_only=True, pto_isa_commit=config.pto_isa_commit, + diagnostic_phase=config.diagnostic_phase, + disabled_diagnostics=config.disabled_diagnostics, + compile_profiling=config.compile_profiling, + distributed_config=distributed_config, ) - if len(assembled) == 2: - chip_callable, runtime_name = assembled - runtime_config = {} - else: - chip_callable, runtime_name, runtime_config = assembled - runtime_config = runtime_config or {} - param_infos, _, _ = compiled._get_metadata() - return _L2Callable( - chip_callable=chip_callable, + compiled = jit_fn.compile(*dummy_args, config=run_config) + _patch_aicore_bitcast_helpers(Path(compiled.output_dir)) + if not isinstance(compiled, DistributedCompiledProgram): + raise TypeError( + f"{name} did not compile to DistributedCompiledProgram; got {type(compiled).__name__}" + ) + return _L3Callable( + compiled=compiled, name=name, - runtime_name=runtime_name, - block_dim=int(runtime_config.get("block_dim", _QWEN14B_BLOCK_DIM)), - aicpu_thread_num=int(runtime_config.get("aicpu_thread_num", 4)), - param_infos=tuple(param_infos), + block_dim=_QWEN14B_BLOCK_DIM, + aicpu_thread_num=4, ) - def _l2_work_dir(self, name: str) -> Path: - """Return a dedicated compile directory for one non-L3 program.""" - if self._save_kernels_dir is not None: - root = Path(self._save_kernels_dir) - else: - if self._l2_compile_root is None: - self._l2_compile_root = Path(tempfile.mkdtemp(prefix="qwen3_14b_l2_")) - root = self._l2_compile_root - work_dir = root / name - work_dir.mkdir(parents=True, exist_ok=True) - return work_dir - @staticmethod def _validate_kernel_runtime_contract( model: RuntimeModel, qwen3_prefill_fwd: object, qwen3_decode_layer: object, ) -> None: - """Validate runtime dimensions baked into the A8W8 L2 kernels.""" + """Validate runtime dimensions baked into the A8W8 kernels.""" expected_page_size = int(qwen3_decode_layer.BLOCK_SIZE) if int(model.runtime.page_size) != expected_page_size: raise ValueError( diff --git a/examples/model/qwen3_14b/runner/npu_runner_a8w8.py b/examples/model/qwen3_14b/runner/npu_runner_a8w8.py index 11d918a..163254f 100644 --- a/examples/model/qwen3_14b/runner/npu_runner_a8w8.py +++ b/examples/model/qwen3_14b/runner/npu_runner_a8w8.py @@ -9,13 +9,18 @@ from __future__ import annotations -import ctypes +import os from dataclasses import dataclass from typing import Any import torch -from examples.model.qwen3_14b.runner.npu_runner import Qwen314BModelRunner, _add_run_timing_args +from examples.model.qwen3_14b.runner.npu_runner import ( + Qwen314BModelRunner, + _L3Callable, + _StaticDeviceTensor, + _add_run_timing_args, +) from python.core.model_runner import ModelRunner from python.core.types import ( DecodeBatch, @@ -25,14 +30,12 @@ RuntimeModel, ) from python.profile import profile_span -from python.runtime.worker import Worker as LlmWorker -from python.runtime.worker import WorkerTensor _QWEN14B_A8W8_PREFILL_CHUNK_LAYERS = 10 _QWEN14B_LM_HEAD_CHUNK_ROWS = 8192 -def _l2_trace_name(kernel_name: str) -> str: +def _kernel_trace_name(kernel_name: str) -> str: if "prefill" in kernel_name: return "kernel.prefill_fwd" if "decode" in kernel_name: @@ -65,24 +68,12 @@ class _KernelLayerWeights: wo_scale: torch.Tensor | None = None -@dataclass -class _L2Callable: - """Assembled non-L3 callable and launch metadata.""" - - chip_callable: object - name: str - runtime_name: str - block_dim: int - aicpu_thread_num: int - param_infos: tuple[object, ...] - - @dataclass class _CompiledKernels: """Compiled Qwen3-14B kernels and immutable runtime tensors.""" - prefill: _L2Callable - decode: _L2Callable + prefill: _L3Callable + decode: _L3Callable final_norm_weight: torch.Tensor rope_cos: torch.Tensor rope_sin: torch.Tensor @@ -90,6 +81,18 @@ class _CompiledKernels: padded_lm_head_weight: torch.Tensor decode_weights: dict[str, torch.Tensor] decode_logits_buffer: torch.Tensor + prefill_hidden_buffer: torch.Tensor + prefill_next_hidden_buffer: torch.Tensor + prefill_seq_lens_buffer: torch.Tensor + prefill_chunk_lens_buffer: torch.Tensor + prefill_chunk_offsets_buffer: torch.Tensor + prefill_block_table_buffer: torch.Tensor + prefill_slot_mapping_buffer: torch.Tensor + prefill_logits_buffer: torch.Tensor + decode_hidden_buffer: torch.Tensor + decode_seq_lens_buffer: torch.Tensor + decode_block_table_buffer: torch.Tensor + decode_slot_mapping_buffer: torch.Tensor @dataclass @@ -116,14 +119,6 @@ class _DecodeInputs: slot_mapping: torch.Tensor -@dataclass -class _L2ProgramHandle: - """L2 callable registration state for one runner process.""" - - callable_id: int - runtime_name: str - - class Qwen314BA8W8ModelRunner(ModelRunner): """Runtime wrapper for one Qwen3-14B model's compiled PyPTO kernels.""" @@ -135,7 +130,7 @@ class Qwen314BA8W8ModelRunner(ModelRunner): def __init__( self, *, - compiled: _CompiledKernels, + compiled: _CompiledKernels | None, model_id: str = "test-model", platform: str = "a2a3sim", device_id: int = 0, @@ -147,24 +142,38 @@ def __init__( self._platform = platform self._device_id = device_id self._save_kernels_dir = save_kernels_dir - self._l2_workers: dict[str, LlmWorker] = {} - self._l2_programs: dict[int, _L2ProgramHandle] = {} - self._l2_child_allocs: dict[tuple[str, int], tuple[int, int]] = {} - self._kv_scale_caches: dict[str, tuple[WorkerTensor, WorkerTensor]] = {} + self._l3_worker: Any | None = None + self._l3_static_tensors: dict[tuple[int, tuple[int, ...], torch.dtype], object] = {} + self._l3_static_storage_tensors: dict[int, object] = {} + self._l3_static_host_tensors: dict[int, torch.Tensor] = {} + self._kv_scale_caches: dict[str, tuple[Any, Any]] = {} + if compiled is not None: + self._register_l3_static_host_tensors() def init_kv_cache(self, model_id: str, config, runtime) -> int: """Create the runner-owned KV cache, plus INT8 scale pages for A8W8.""" + self._l3_log("init_kv_cache: preparing DistributedWorker") + self._shared_l3_worker() + self._l3_log("init_kv_cache: DistributedWorker ready") + self._l3_log("init_kv_cache: allocating KV cache") num_pages = super().init_kv_cache(model_id, config, runtime) + self._l3_log(f"init_kv_cache: KV cache pages={num_pages}") if model_id in self._kv_scale_caches: return num_pages cache_rows = config.num_hidden_layers * num_pages * config.num_key_value_heads * runtime.page_size + self._l3_log("init_kv_cache: allocating KV scale cache") + self._l3_log("init_kv_cache: allocating key scale") key_scale = self._alloc_kv_cache_tensor((cache_rows, 8), torch.float32) + self._l3_log("init_kv_cache: key scale allocated") try: + self._l3_log("init_kv_cache: allocating value scale") value_scale = self._alloc_kv_cache_tensor((cache_rows, 8), torch.float32) + self._l3_log("init_kv_cache: value scale allocated") except Exception: self._free_kv_cache_tensor(key_scale) raise self._kv_scale_caches[model_id] = (key_scale, value_scale) + self._l3_log("init_kv_cache: scale cache ready") return num_pages def close_kv_cache(self) -> None: @@ -174,23 +183,14 @@ def close_kv_cache(self) -> None: self._kv_scale_caches.clear() super().close_kv_cache() - def _kv_cache_runtime_name(self) -> str: - if self._compiled.prefill.runtime_name != self._compiled.decode.runtime_name: - raise ValueError( - "device-side KV cache requires prefill and decode to use the same L2 runtime: " - f"{self._compiled.prefill.runtime_name!r} != {self._compiled.decode.runtime_name!r}" - ) - return self._compiled.prefill.runtime_name - - def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> WorkerTensor: - """Allocate one KV cache tensor on the L2 NPU worker.""" - worker = self._worker_for_runtime(self._kv_cache_runtime_name()) - return worker.alloc_tensor(shape, dtype) + def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> Any: + """Allocate one KV cache tensor on the active NPU worker.""" + return self._shared_l3_worker().alloc_tensor(shape, dtype) - def _free_kv_cache_tensor(self, tensor: WorkerTensor) -> None: - """Free one KV cache tensor from the L2 NPU worker.""" - worker = self._l2_workers.get(self._kv_cache_runtime_name()) - if worker is not None and worker.initialized: + def _free_kv_cache_tensor(self, tensor: Any) -> None: + """Free one KV cache tensor from the active NPU worker.""" + worker = self._l3_worker + if worker is not None: worker.free_tensor(tensor) @staticmethod @@ -198,7 +198,7 @@ def _validate_kv_cache_bounds( model: RuntimeModel, block_table: torch.Tensor, slot_mapping: torch.Tensor, - cache: WorkerTensor, + cache: Any, ) -> None: """Fail on host before an invalid KV page id reaches the NPU kernel.""" valid_blocks = block_table[block_table >= 0] @@ -220,6 +220,7 @@ def _validate_kv_cache_bounds( def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult: """Run the JIT all-layer prefill kernel and return next-token logits.""" + self._l3_log("run_prefill: start") compiled = self._compiled prefill_inputs = self._prepare_prefill_inputs(model, batch) dw = compiled.decode_weights @@ -231,10 +232,7 @@ def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult v_cache = kv_cache.value_pages self._validate_kv_cache_bounds(model, prefill_inputs.block_table, prefill_inputs.slot_mapping, k_cache) kv_scales = self._kv_scale_caches.get(model.config.model_id) - logits_padded = torch.zeros( - (prefill_inputs.actual_batch, compiled.padded_vocab), - dtype=torch.float32, - ).share_memory_() + logits_padded = compiled.prefill_logits_buffer[: prefill_inputs.actual_batch] if kv_scales is None: raise RuntimeError(f"missing A8W8 KV scale cache for model {model.config.model_id!r}") @@ -242,9 +240,9 @@ def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult rows_per_layer = k_cache.shape[0] // model.config.num_hidden_layers hidden = prefill_inputs.hidden - def weight_slice(name: str, start: int, layers: int, rows_per_layer_: int = 1) -> WorkerTensor: + def weight_slice(name: str, start: int, layers: int, rows_per_layer_: int = 1) -> Any: tensor = dw[name][start * rows_per_layer_ : (start + layers) * rows_per_layer_] - return self._l2_child_tensor(compiled.prefill.runtime_name, tensor) + return self._kernel_static_tensor(compiled.prefill, tensor) for layer_start in range(0, model.config.num_hidden_layers, _QWEN14B_A8W8_PREFILL_CHUNK_LAYERS): layer_count = min( @@ -253,9 +251,16 @@ def weight_slice(name: str, start: int, layers: int, rows_per_layer_: int = 1) - ) cache_row_start = layer_start * rows_per_layer cache_rows = layer_count * rows_per_layer - hidden_out = torch.empty_like(hidden).share_memory_() - scratch_logits = torch.empty_like(logits_padded).share_memory_() - self._run_l2_program( + base_hidden = compiled.prefill_hidden_buffer + next_hidden = compiled.prefill_next_hidden_buffer + hidden_out = ( + next_hidden[: hidden.shape[0]] + if hidden.data_ptr() == base_hidden.data_ptr() + else base_hidden[: hidden.shape[0]] + ) + scratch_logits = compiled.prefill_logits_buffer[: prefill_inputs.actual_batch] + self._l3_log(f"run_prefill: dispatch layers {layer_start}-{layer_start + layer_count - 1}") + self._run_distributed_program( compiled.prefill, hidden, prefill_inputs.seq_lens, @@ -270,37 +275,39 @@ def weight_slice(name: str, start: int, layers: int, rows_per_layer_: int = 1) - weight_slice("decode_wv_scale", layer_start, layer_count), weight_slice("decode_q_norm_weight", layer_start, layer_count), weight_slice("decode_k_norm_weight", layer_start, layer_count), - self._l2_child_tensor(compiled.prefill.runtime_name, compiled.rope_cos), - self._l2_child_tensor(compiled.prefill.runtime_name, compiled.rope_sin), + self._kernel_static_tensor(compiled.prefill, compiled.rope_cos), + self._kernel_static_tensor(compiled.prefill, compiled.rope_sin), prefill_inputs.block_table, prefill_inputs.slot_mapping, - self._worker_tensor_view( + self._device_tensor_view( k_cache, cache_row_start * model.config.head_dim, (cache_rows, model.config.head_dim), 1, ), - self._worker_tensor_view( + self._device_tensor_view( v_cache, cache_row_start * model.config.head_dim, (cache_rows, model.config.head_dim), 1, ), - self._worker_tensor_view(k_cache_scale, cache_row_start * 8, (cache_rows, 8), 4), - self._worker_tensor_view(v_cache_scale, cache_row_start * 8, (cache_rows, 8), 4), + self._device_tensor_view(k_cache_scale, cache_row_start * 8, (cache_rows, 8), 4), + self._device_tensor_view(v_cache_scale, cache_row_start * 8, (cache_rows, 8), 4), weight_slice("decode_wo", layer_start, layer_count, model.config.hidden_size), weight_slice("decode_wo_scale", layer_start, layer_count), weight_slice("decode_post_rms_weight", layer_start, layer_count), weight_slice("decode_w_gate", layer_start, layer_count, model.config.hidden_size), weight_slice("decode_w_up", layer_start, layer_count, model.config.hidden_size), weight_slice("decode_w_down", layer_start, layer_count, model.config.intermediate_size), - self._l2_child_tensor(compiled.prefill.runtime_name, compiled.final_norm_weight), - self._l2_child_tensor(compiled.prefill.runtime_name, compiled.padded_lm_head_weight), + self._kernel_static_tensor(compiled.prefill, compiled.final_norm_weight), + self._kernel_static_tensor(compiled.prefill, compiled.padded_lm_head_weight), scratch_logits, hidden_out, ) + self._l3_log(f"run_prefill: layers {layer_start}-{layer_start + layer_count - 1} done") hidden = hidden_out logits_padded = self._project_logits_host(model, compiled, prefill_inputs, hidden) + self._l3_log("run_prefill: done") for batch_idx, alloc in enumerate(batch.kv_allocations): seq_len = int(batch.seq_lens[batch_idx].item()) @@ -326,12 +333,12 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: (an idempotent, safe write), and their logits are trimmed off below. This avoids padded rows clobbering an unrelated request's physical page. """ + self._l3_log("run_decode: start") compiled = self._compiled model_id = model.config.model_id decode_inputs = self._prepare_decode_inputs(model, batch) actual_batch = decode_inputs.actual_batch dw = compiled.decode_weights - rt = compiled.decode.runtime_name kernel_batch = model.runtime.max_batch_size max_blocks = self._max_blocks_per_seq(model) @@ -348,19 +355,30 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: f"{compiled.decode_logits_buffer.shape[0]}" ) - # Pad active inputs up to the fixed kernel batch by replicating row 0. - def _pad_rows(active: torch.Tensor, rows_each: int) -> torch.Tensor: - view = active.reshape(actual_batch, rows_each) - padded = view[0:1].expand(kernel_batch - actual_batch, rows_each) - return torch.cat([view, padded], dim=0).reshape(-1).contiguous() - - hidden = torch.zeros((kernel_batch, model.config.hidden_size), dtype=torch.bfloat16) - hidden[:actual_batch] = decode_inputs.hidden - hidden[actual_batch:] = decode_inputs.hidden[0:1] - hidden = hidden.share_memory_() - seq_lens = _pad_rows(decode_inputs.seq_lens, 1).to(torch.int32).share_memory_() - block_table = _pad_rows(decode_inputs.block_table, max_blocks).to(torch.int32).share_memory_() - slot_mapping = _pad_rows(decode_inputs.slot_mapping, 1).to(torch.int32).share_memory_() + hidden = compiled.decode_hidden_buffer + hidden[:actual_batch].copy_(decode_inputs.hidden) + hidden[actual_batch:].copy_(decode_inputs.hidden[0:1]) + seq_lens = self._copy_replicated_rows( + compiled.decode_seq_lens_buffer, + decode_inputs.seq_lens, + actual_batch, + kernel_batch, + rows_each=1, + ) + block_table = self._copy_replicated_rows( + compiled.decode_block_table_buffer, + decode_inputs.block_table, + actual_batch, + kernel_batch, + rows_each=max_blocks, + ) + slot_mapping = self._copy_replicated_rows( + compiled.decode_slot_mapping_buffer, + decode_inputs.slot_mapping, + actual_batch, + kernel_batch, + rows_each=1, + ) # Padded block_table / slot_mapping only ever reference row 0's # already-valid pages, so bound-check exactly what the kernel will read. @@ -370,37 +388,38 @@ def _pad_rows(active: torch.Tensor, rows_each: int) -> torch.Tensor: if kv_scales is None: raise RuntimeError(f"missing A8W8 KV scale cache for model {model_id!r}") k_cache_scale, v_cache_scale = kv_scales - self._run_l2_program( + self._run_distributed_program( compiled.decode, hidden, - self._l2_child_tensor(rt, dw["decode_input_rms_weight"]), - self._l2_child_tensor(rt, dw["decode_wq"]), - self._l2_child_tensor(rt, dw["decode_wk"]), - self._l2_child_tensor(rt, dw["decode_wv"]), - self._l2_child_tensor(rt, dw["decode_wq_scale"]), - self._l2_child_tensor(rt, dw["decode_wk_scale"]), - self._l2_child_tensor(rt, dw["decode_wv_scale"]), - self._l2_child_tensor(rt, dw["decode_q_norm_weight"]), - self._l2_child_tensor(rt, dw["decode_k_norm_weight"]), + self._kernel_static_tensor(compiled.decode, dw["decode_input_rms_weight"]), + self._kernel_static_tensor(compiled.decode, dw["decode_wq"]), + self._kernel_static_tensor(compiled.decode, dw["decode_wk"]), + self._kernel_static_tensor(compiled.decode, dw["decode_wv"]), + self._kernel_static_tensor(compiled.decode, dw["decode_wq_scale"]), + self._kernel_static_tensor(compiled.decode, dw["decode_wk_scale"]), + self._kernel_static_tensor(compiled.decode, dw["decode_wv_scale"]), + self._kernel_static_tensor(compiled.decode, dw["decode_q_norm_weight"]), + self._kernel_static_tensor(compiled.decode, dw["decode_k_norm_weight"]), seq_lens, block_table, slot_mapping, - self._l2_child_tensor(rt, compiled.rope_cos), - self._l2_child_tensor(rt, compiled.rope_sin), + self._kernel_static_tensor(compiled.decode, compiled.rope_cos), + self._kernel_static_tensor(compiled.decode, compiled.rope_sin), k_cache, v_cache, k_cache_scale, v_cache_scale, - self._l2_child_tensor(rt, dw["decode_wo"]), - self._l2_child_tensor(rt, dw["decode_wo_scale"]), - self._l2_child_tensor(rt, dw["decode_w_gate"]), - self._l2_child_tensor(rt, dw["decode_w_up"]), - self._l2_child_tensor(rt, dw["decode_w_down"]), - self._l2_child_tensor(rt, dw["decode_post_rms_weight"]), - self._l2_child_tensor(rt, compiled.final_norm_weight), - self._l2_child_tensor(rt, compiled.padded_lm_head_weight), + self._kernel_static_tensor(compiled.decode, dw["decode_wo"]), + self._kernel_static_tensor(compiled.decode, dw["decode_wo_scale"]), + self._kernel_static_tensor(compiled.decode, dw["decode_w_gate"]), + self._kernel_static_tensor(compiled.decode, dw["decode_w_up"]), + self._kernel_static_tensor(compiled.decode, dw["decode_w_down"]), + self._kernel_static_tensor(compiled.decode, dw["decode_post_rms_weight"]), + self._kernel_static_tensor(compiled.decode, compiled.final_norm_weight), + self._kernel_static_tensor(compiled.decode, compiled.padded_lm_head_weight), logits_padded, ) + self._l3_log("run_decode: done") for batch_idx, alloc in enumerate(batch.kv_allocations): alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item())) return DecodeResult( @@ -440,133 +459,180 @@ def _project_logits_host( logits[:, row_start:row_end] = normed.float() @ lm_head[row_start:row_end].float().T return logits - @staticmethod - def _worker_tensor_view( - tensor: WorkerTensor, + def _device_tensor_view( + self, + tensor: Any, element_offset: int, shape: tuple[int, ...], element_size: int, - ) -> WorkerTensor: - """Return a contiguous WorkerTensor view starting at an element offset.""" - return WorkerTensor( + ) -> Any: + """Return a contiguous DeviceTensor view for the L3 worker.""" + from pypto.runtime import DeviceTensor # noqa: PLC0415 + + return DeviceTensor( data_ptr=tensor.data_ptr + element_offset * element_size, shape=shape, dtype=tensor.dtype, ) - def _run_l2_program(self, callable_spec: _L2Callable, *args: Any) -> Any: - """Run a compiled non-L3 program through the LLM Simpler worker.""" - from simpler.task_interface import CallConfig # noqa: PLC0415 + def _kernel_static_tensor(self, callable_spec: _L3Callable, tensor: torch.Tensor) -> Any: + """Return a backend-resident static tensor argument.""" + return self._l3_static_storage_view(tensor) + def _run_distributed_program(self, callable_spec: _L3Callable, *args: Any) -> Any: + """Run a compiled HOST wrapper through the shared PyPTO L3 worker.""" span_args = { "kernel": callable_spec.name, - "runtime": callable_spec.runtime_name, "block_dim": callable_spec.block_dim, "aicpu_thread_num": callable_spec.aicpu_thread_num, } with profile_span( - _l2_trace_name(callable_spec.name), + _kernel_trace_name(callable_spec.name), cat="kernel", level="kernel", args=span_args, ): - handle = self._ensure_l2_program(callable_spec) - orch_args = self._build_l2_orch_args(callable_spec, args) - - cfg = CallConfig() - cfg.block_dim = callable_spec.block_dim - cfg.aicpu_thread_num = callable_spec.aicpu_thread_num - - worker = self._l2_workers[handle.runtime_name] - timing = worker.run(handle.callable_id, orch_args, cfg) + worker = self._shared_l3_worker() + l3_args = callable_spec.dispatch_args + tuple(self._coerce_l3_arg(worker, arg) for arg in args) + timing = worker.run(callable_spec.compiled, *l3_args) _add_run_timing_args(span_args, timing) return timing - def _worker_for_runtime(self, runtime_name: str) -> LlmWorker: - """Return an initialized worker for ``runtime_name``.""" - worker = self._l2_workers.get(runtime_name) - if worker is not None: - return worker - worker = LlmWorker( - level=2, - platform=self._platform, - runtime=runtime_name, - device_id=self._device_id, - auto_init=True, - ) - self._l2_workers[runtime_name] = worker + def _shared_l3_worker(self) -> Any: + """Return the L3 worker shared by A8W8 prefill/decode.""" + worker = self._l3_worker + if worker is None: + from pypto.runtime import DistributedWorker # noqa: PLC0415 + + worker = DistributedWorker([ + self._compiled.prefill.compiled, + self._compiled.decode.compiled, + ]) + self._l3_worker = worker + self._l3_log("DistributedWorker constructed") return worker - def _ensure_l2_program(self, callable_spec: _L2Callable) -> _L2ProgramHandle: - """Register and cache one executor-assembled non-L3 callable.""" - key = id(callable_spec) - cached = self._l2_programs.get(key) + def _coerce_l3_arg(self, worker: Any, arg: Any) -> Any: + """Convert static upload markers to worker-resident tensors.""" + if not isinstance(arg, _StaticDeviceTensor): + return arg + tensor = arg.tensor + key = (tensor.data_ptr(), tuple(tensor.shape), tensor.dtype) + cached = self._l3_static_tensors.get(key) if cached is not None: return cached + dev = worker.alloc_tensor(tensor.shape, tensor.dtype, init=tensor) + self._l3_static_tensors[key] = dev + return dev - worker = self._worker_for_runtime(callable_spec.runtime_name) + def _register_l3_static_host_tensors(self) -> None: + """Register full host storages that L3 static views can share.""" + compiled = self._compiled + for tensor in ( + compiled.final_norm_weight, + compiled.rope_cos, + compiled.rope_sin, + compiled.padded_lm_head_weight, + *compiled.decode_weights.values(), + ): + self._register_l3_static_host_tensor(tensor) - handle = _L2ProgramHandle( - callable_id=worker.register(callable_spec.chip_callable), - runtime_name=callable_spec.runtime_name, - ) - self._l2_programs[key] = handle - return handle + def _register_l3_static_host_tensor(self, tensor: torch.Tensor) -> None: + """Remember a full shared host tensor by storage pointer.""" + if tensor.device.type != "cpu": + raise ValueError("L3 static host tensor must be on CPU") + if not tensor.is_contiguous(): + raise ValueError("L3 static host tensor must be contiguous") + tensor = self._share_cpu_tensor(tensor) + self._l3_static_host_tensors[int(tensor.untyped_storage().data_ptr())] = tensor - def _l2_child_tensor( - self, - runtime_name: str, - tensor: torch.Tensor, - *, - upload: bool = True, - refresh: bool = False, - ) -> WorkerTensor: - """Return a worker-resident view for a CPU tensor's backing storage.""" - from simpler_setup.torch_interop import torch_dtype_to_datatype # noqa: PLC0415 + def _l3_static_storage_view(self, tensor: torch.Tensor) -> Any: + """Upload a full static storage once and return a DeviceTensor view.""" + from pypto.runtime import DeviceTensor # noqa: PLC0415 if tensor.device.type != "cpu": - raise ValueError("child-memory tensor must be on CPU") + raise ValueError("L3 static tensor must be on CPU") if not tensor.is_contiguous(): - raise ValueError("child-memory tensor must be contiguous") + raise ValueError("L3 static tensor view must be contiguous") tensor = self._share_cpu_tensor(tensor) storage = tensor.untyped_storage() storage_ptr = int(storage.data_ptr()) - storage_nbytes = int(storage.nbytes()) - tensor_offset = int(tensor.data_ptr()) - storage_ptr - if tensor_offset < 0 or tensor_offset + int(tensor.nbytes) > storage_nbytes: - raise ValueError("tensor view is outside its backing storage") - - key = (runtime_name, storage_ptr) - alloc = self._l2_child_allocs.get(key) - if alloc is None: - worker = self._worker_for_runtime(runtime_name) - dev_ptr = worker.malloc(storage_nbytes) - if upload: - worker.copy_to(dev_ptr, storage_ptr, storage_nbytes) - alloc = (dev_ptr, storage_nbytes) - self._l2_child_allocs[key] = alloc - elif upload and refresh: - worker = self._worker_for_runtime(runtime_name) - worker.copy_to(alloc[0], storage_ptr, storage_nbytes) - - dev_base, _ = alloc - shape = tuple(int(dim) for dim in tensor.shape) - return WorkerTensor( - data_ptr=dev_base + tensor_offset, - shape=shape, - dtype=torch_dtype_to_datatype(tensor.dtype), + host_full = self._l3_static_host_tensors.get(storage_ptr) + if host_full is None: + self._register_l3_static_host_tensor(tensor) + host_full = tensor + + dev_full = self._l3_static_storage_tensors.get(storage_ptr) + if dev_full is None: + worker = self._shared_l3_worker() + dev_full = worker.alloc_tensor(host_full.shape, host_full.dtype, init=host_full) + self._l3_static_storage_tensors[storage_ptr] = dev_full + + byte_offset = int(tensor.data_ptr()) - storage_ptr + if byte_offset < 0 or byte_offset + int(tensor.nbytes) > int(storage.nbytes()): + raise ValueError("L3 static tensor view is outside its backing storage") + return DeviceTensor( + data_ptr=dev_full.data_ptr + byte_offset, + shape=tuple(int(dim) for dim in tensor.shape), + dtype=tensor.dtype, ) - def _release_l2_child_allocs(self, runtime_name: str) -> None: - """Free cached child-memory allocations for one L2 runtime.""" - worker = self._l2_workers.get(runtime_name) - for key, (dev_ptr, _nbytes) in list(self._l2_child_allocs.items()): - key_runtime, _storage_ptr = key - if key_runtime != runtime_name: - continue - if worker is not None and worker.initialized: - worker.free(dev_ptr) - self._l2_child_allocs.pop(key, None) + def _materialize_static_tensors(self) -> None: + """Upload immutable model tensors into the shared L3 worker once.""" + worker = self._shared_l3_worker() + compiled = self._compiled + for arg in ( + self._static_device_tensor(compiled.final_norm_weight), + self._static_device_tensor(compiled.rope_cos), + self._static_device_tensor(compiled.rope_sin), + self._static_device_tensor(compiled.padded_lm_head_weight), + *(self._static_device_tensor(tensor) for tensor in compiled.decode_weights.values()), + ): + self._coerce_l3_arg(worker, arg) + + @staticmethod + def _static_device_tensor(tensor: torch.Tensor) -> _StaticDeviceTensor: + """Mark a CPU tensor for one-time upload to the shared L3 worker.""" + if tensor.device.type != "cpu": + raise ValueError("worker-resident tensor must be on CPU") + if not tensor.is_contiguous(): + raise ValueError("worker-resident tensor must be contiguous") + return _StaticDeviceTensor(Qwen314BModelRunner._share_cpu_tensor(tensor)) + + def _clear_l3_static_tensors(self) -> None: + """Release L3 static uploads, mainly prefill layer-slice weights.""" + worker = self._l3_worker + if worker is None: + self._l3_static_tensors.clear() + return + for dev in list(self._l3_static_tensors.values()): + worker.free_tensor(dev) + self._l3_static_tensors.clear() + self._l3_static_storage_tensors.clear() + self._l3_log("static tensor uploads released") + + @staticmethod + def _l3_log(message: str) -> None: + """Print L3 diagnostics only when explicitly requested.""" + if os.environ.get("QWEN_A8W8_L3_DEBUG") == "1": + print(f"[a8w8-l3] {message}", flush=True) + + @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 _share_cpu_tensor(tensor: torch.Tensor) -> torch.Tensor: @@ -576,58 +642,14 @@ def _share_cpu_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor def close(self) -> None: - """Release non-L3 child-memory allocations and L2 workers.""" + """Release A8W8 L3 worker resources.""" self.close_kv_cache() - for (runtime_name, _), (dev_ptr, _nbytes) in list(self._l2_child_allocs.items()): - worker = self._l2_workers.get(runtime_name) - if worker is not None and worker.initialized: - worker.free(dev_ptr) - self._l2_child_allocs.clear() - self._l2_programs.clear() - for worker in self._l2_workers.values(): + worker = self._l3_worker + if worker is not None: worker.close() - self._l2_workers.clear() - - @staticmethod - def _build_l2_orch_args(callable_spec: _L2Callable, args: tuple[Any, ...]): - """Build ``ChipStorageTaskArgs`` for a compiled L2 program call.""" - from simpler.task_interface import ChipStorageTaskArgs, scalar_to_uint64 # noqa: PLC0415 - try: - from simpler.task_interface import ContinuousTensor # noqa: PLC0415 - except ImportError: - from simpler.task_interface import Tensor as ContinuousTensor # noqa: PLC0415 - from simpler_setup.torch_interop import make_tensor_arg # noqa: PLC0415 - - param_infos = callable_spec.param_infos - if len(args) != len(param_infos): - names = [p.name for p in param_infos] - raise TypeError( - f"compiled program expects {len(param_infos)} arguments, got {len(args)}. Parameters: {names}" - ) - - orch_args = ChipStorageTaskArgs() - for info, arg in zip(param_infos, args, strict=True): - if info.shape is None: - if not isinstance(arg, ctypes._SimpleCData): - raise TypeError(f"scalar parameter {info.name!r} must be passed as a ctypes scalar") - orch_args.add_scalar(scalar_to_uint64(arg)) - continue - if isinstance(arg, WorkerTensor): - orch_args.add_tensor(arg.to_continuous_tensor()) - continue - if isinstance(arg, ContinuousTensor): - orch_args.add_tensor(arg) - continue - if not isinstance(arg, torch.Tensor): - raise TypeError(f"tensor parameter {info.name!r} expects torch.Tensor, got {type(arg).__name__}") - if arg.device.type != "cpu": - raise ValueError(f"tensor parameter {info.name!r} must be on CPU for Simpler L2 dispatch") - if not arg.is_contiguous(): - raise ValueError(f"tensor parameter {info.name!r} must be contiguous") - if not arg.is_shared(): - arg.share_memory_() - orch_args.add_tensor(make_tensor_arg(arg)) - return orch_args + self._l3_worker = None + self._l3_static_tensors.clear() + self._l3_static_storage_tensors.clear() def _prepare_prefill_inputs( self, @@ -638,14 +660,9 @@ def _prepare_prefill_inputs( batch_count = len(batch.kv_allocations) if batch.kv_allocations else int(batch.seq_lens.shape[0]) actual_batch = self._validate_batch_size(model, batch_count) max_seq = model.runtime.max_seq_len - hidden_size = model.config.hidden_size page_size = model.runtime.page_size max_blocks = self._max_blocks_per_seq(model) - seq_lens = torch.empty((actual_batch,), dtype=torch.int32) - chunk_lens = torch.empty((actual_batch,), dtype=torch.int32) - chunk_offsets = torch.empty((actual_batch,), dtype=torch.int32) - block_table = torch.full((actual_batch * max_blocks,), -1, dtype=torch.int32) seq_len_values = [int(batch.seq_lens[idx].item()) for idx in range(actual_batch)] chunk_len_values: list[int] = [] chunk_start_values: list[int] = [] @@ -680,8 +697,20 @@ def _prepare_prefill_inputs( chunk_len_values.append(chunk_len) chunk_start_values.append(chunk_start) total_tokens = sum(chunk_len_values) - hidden = torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16) - slot_mapping = torch.empty((total_tokens,), dtype=torch.int32) + compiled = self._compiled + hidden_buffer = compiled.prefill_hidden_buffer + if total_tokens > hidden_buffer.shape[0]: + raise ValueError(f"prefill total tokens {total_tokens} exceeds L3 buffer {hidden_buffer.shape[0]}") + hidden = hidden_buffer[:total_tokens] + seq_lens = compiled.prefill_seq_lens_buffer[:actual_batch] + chunk_lens = compiled.prefill_chunk_lens_buffer[:actual_batch] + chunk_offsets = compiled.prefill_chunk_offsets_buffer[:actual_batch] + block_table = compiled.prefill_block_table_buffer[: actual_batch * max_blocks] + slot_mapping = compiled.prefill_slot_mapping_buffer[:total_tokens] + seq_lens.zero_() + chunk_lens.zero_() + chunk_offsets.zero_() + block_table.fill_(-1) token_offset = 0 for batch_idx in range(actual_batch): @@ -713,12 +742,12 @@ def _prepare_prefill_inputs( return _PrefillInputs( actual_batch=actual_batch, - hidden=hidden.share_memory_(), - seq_lens=seq_lens.share_memory_(), - chunk_lens=chunk_lens.share_memory_(), - chunk_offsets=chunk_offsets.share_memory_(), - block_table=block_table.share_memory_(), - slot_mapping=slot_mapping.share_memory_(), + hidden=hidden, + seq_lens=seq_lens, + chunk_lens=chunk_lens, + chunk_offsets=chunk_offsets, + block_table=block_table, + slot_mapping=slot_mapping, ) def _prepare_decode_inputs( diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index 781a01c..130f9f3 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -141,3 +141,139 @@ def qwen3_greedy_sample_host( logits, sampled_ids, ) + + +def create_qwen3_a8w8_dispatch(prefill_hidden_a8w8, decode_fwd): + """Create A8W8 HOST wrappers with format-local kernel dependencies.""" + + @pl.jit.host + def qwen3_a8w8_prefill_host( + hidden_states: pl.Tensor, + seq_lens: pl.Tensor, + chunk_lens: pl.Tensor, + chunk_offsets: pl.Tensor, + input_rms_weight: pl.Tensor, + wq: pl.Tensor, + wk: pl.Tensor, + wv: pl.Tensor, + wq_scale: pl.Tensor, + wk_scale: pl.Tensor, + wv_scale: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + block_table: pl.Tensor, + slot_mapping: pl.Tensor, + k_cache: pl.Tensor, + v_cache: pl.Tensor, + k_cache_scale: pl.Tensor, + v_cache_scale: pl.Tensor, + wo: pl.Tensor, + wo_scale: pl.Tensor, + post_rms_weight: pl.Tensor, + w_gate: pl.Tensor, + w_up: pl.Tensor, + w_down: pl.Tensor, + final_norm_weight: pl.Tensor, + lm_head_weight: pl.Tensor, + out: pl.Out[pl.Tensor], + hidden_out: pl.Out[pl.Tensor], + ) -> pl.Tensor: + return prefill_hidden_a8w8( + hidden_states, + seq_lens, + chunk_lens, + chunk_offsets, + input_rms_weight, + wq, + wk, + wv, + wq_scale, + wk_scale, + wv_scale, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + block_table, + slot_mapping, + k_cache, + v_cache, + k_cache_scale, + v_cache_scale, + wo, + wo_scale, + post_rms_weight, + w_gate, + w_up, + w_down, + final_norm_weight, + lm_head_weight, + out, + hidden_out, + ) + + @pl.jit.host + def qwen3_a8w8_decode_host( + hidden_states: pl.Tensor, + input_rms_weight: pl.Tensor, + wq: pl.Tensor, + wk: pl.Tensor, + wv: pl.Tensor, + wq_scale: pl.Tensor, + wk_scale: pl.Tensor, + wv_scale: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + seq_lens: pl.Tensor, + block_table: pl.Tensor, + slot_mapping: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + k_cache: pl.Tensor, + v_cache: pl.Tensor, + k_cache_scale: pl.Tensor, + v_cache_scale: pl.Tensor, + wo: pl.Tensor, + wo_scale: pl.Tensor, + w_gate: pl.Tensor, + w_up: pl.Tensor, + w_down: pl.Tensor, + post_rms_weight: pl.Tensor, + final_norm_weight: pl.Tensor, + lm_head_weight: pl.Tensor, + out: pl.Out[pl.Tensor], + ) -> pl.Tensor: + return decode_fwd( + hidden_states, + input_rms_weight, + wq, + wk, + wv, + wq_scale, + wk_scale, + wv_scale, + q_norm_weight, + k_norm_weight, + seq_lens, + block_table, + slot_mapping, + rope_cos, + rope_sin, + k_cache, + v_cache, + k_cache_scale, + v_cache_scale, + wo, + wo_scale, + w_gate, + w_up, + w_down, + post_rms_weight, + final_norm_weight, + lm_head_weight, + out, + ) + + return qwen3_a8w8_prefill_host, qwen3_a8w8_decode_host diff --git a/tests/test_batching.py b/tests/test_batching.py index e9014b1..f662ce3 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -368,7 +368,7 @@ def test_a8w8_decode_inputs_use_actual_user_batch_without_padding_lanes(): manager = KvCacheManager() manager.register_model(model.config.model_id, model.config, model.runtime) runner = Qwen314BA8W8ModelRunner( - compiled=None, # type: ignore[arg-type] + compiled=None, ) alloc = manager.allocate_for_prompt(model.config.model_id, "req-0", 1) hidden_states = torch.ones(1, model.config.hidden_size) @@ -395,8 +395,9 @@ def test_a8w8_decode_inputs_use_actual_user_batch_without_padding_lanes(): def test_a8w8_init_kv_cache_returns_page_count_for_first_and_repeated_init(monkeypatch): model = _model(max_batch_size=16, max_seq_len=128, page_size=64) runner = Qwen314BA8W8ModelRunner( - compiled=None, # type: ignore[arg-type] + compiled=None, ) + monkeypatch.setattr(runner, "_shared_l3_worker", lambda: _FakeWorker()) allocated_shapes = [] def alloc_kv_cache_tensor(shape, dtype):