diff --git a/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index 87b65c9..0e2b1d3 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,13 @@ 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) + + # 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: - # 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 +332,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 +364,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 +383,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 +423,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 +478,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 +509,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/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/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..c40e476 --- /dev/null +++ b/examples/model/qwen3_14b/runner/npu_executor_a8w8.py @@ -0,0 +1,577 @@ +# 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 +from collections.abc import Sequence +from pathlib import Path + +import torch + +from examples.model.qwen3_14b.runner.npu_runner_a8w8 import ( + _CompiledKernels, + _KernelLayerWeights, + _L3Callable, + 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 + + @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") + 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) + 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( + prefill_jit, + 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( + decode_jit, + 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") + 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() + + 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, + **l3_buffers, + ) + + 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, + ) -> _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 + 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, + ) -> _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 + 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], + ) -> _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) + 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, + ) + 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, + block_dim=_QWEN14B_BLOCK_DIM, + aicpu_thread_num=4, + ) + + @staticmethod + def _validate_kernel_runtime_contract( + model: RuntimeModel, + qwen3_prefill_fwd: object, + qwen3_decode_layer: object, + ) -> None: + """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( + 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 new file mode 100644 index 0000000..163254f --- /dev/null +++ b/examples/model/qwen3_14b/runner/npu_runner_a8w8.py @@ -0,0 +1,801 @@ +# 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 os +from dataclasses import dataclass +from typing import Any + +import torch + +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, + DecodeResult, + PrefillBatch, + PrefillResult, + RuntimeModel, +) +from python.profile import profile_span + +_QWEN14B_A8W8_PREFILL_CHUNK_LAYERS = 10 +_QWEN14B_LM_HEAD_CHUNK_ROWS = 8192 + + +def _kernel_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 _CompiledKernels: + """Compiled Qwen3-14B kernels and immutable runtime tensors.""" + + prefill: _L3Callable + decode: _L3Callable + final_norm_weight: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + padded_vocab: int + 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 +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 + + +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 | None, + 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._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: + 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 _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: 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 + def _validate_kv_cache_bounds( + model: RuntimeModel, + block_table: torch.Tensor, + slot_mapping: torch.Tensor, + cache: Any, + ) -> 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.""" + self._l3_log("run_prefill: start") + 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 = 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}") + 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) -> Any: + tensor = dw[name][start * rows_per_layer_ : (start + layers) * rows_per_layer_] + 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( + _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 + 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, + 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._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._device_tensor_view( + k_cache, + cache_row_start * model.config.head_dim, + (cache_rows, model.config.head_dim), + 1, + ), + self._device_tensor_view( + v_cache, + cache_row_start * model.config.head_dim, + (cache_rows, model.config.head_dim), + 1, + ), + 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._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()) + 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. + """ + 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 + 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]}" + ) + + 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. + 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_distributed_program( + compiled.decode, + hidden, + 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._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._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( + 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 + + def _device_tensor_view( + self, + tensor: Any, + element_offset: int, + shape: tuple[int, ...], + element_size: int, + ) -> 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 _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, + "block_dim": callable_spec.block_dim, + "aicpu_thread_num": callable_spec.aicpu_thread_num, + } + with profile_span( + _kernel_trace_name(callable_spec.name), + cat="kernel", + level="kernel", + args=span_args, + ): + 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 _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 _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 + + 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) + + 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 _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("L3 static tensor must be on CPU") + if not tensor.is_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()) + 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 _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: + """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 A8W8 L3 worker resources.""" + self.close_kv_cache() + worker = self._l3_worker + if worker is not None: + worker.close() + self._l3_worker = None + self._l3_static_tensors.clear() + self._l3_static_storage_tensors.clear() + + 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 + page_size = model.runtime.page_size + max_blocks = self._max_blocks_per_seq(model) + + 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) + 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): + 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, + seq_lens=seq_lens, + chunk_lens=chunk_lens, + chunk_offsets=chunk_offsets, + block_table=block_table, + slot_mapping=slot_mapping, + ) + + 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_(), + ) 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/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. diff --git a/tests/test_batching.py b/tests/test_batching.py index 850682e..f662ce3 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,120 @@ 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, + ) + 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, + ) + monkeypatch.setattr(runner, "_shared_l3_worker", lambda: _FakeWorker()) + 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 +650,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 +763,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 +842,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