From b2080dc85ff2ca9a903ee5bc66cfadc88d291613 Mon Sep 17 00:00:00 2001 From: ndleslx Date: Wed, 3 Jun 2026 17:02:49 +0800 Subject: [PATCH] Rewrite non-L3 kernels through L3 worker --- examples/model/qwen3_14b/runner/npu_runner.py | 318 ++++++++++++++---- python/runtime/worker.py | 139 +++++++- tests/test_npu_runner_l3_worker.py | 245 ++++++++++++++ 3 files changed, 623 insertions(+), 79 deletions(-) create mode 100644 tests/test_npu_runner_l3_worker.py diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 9445a5a..5798c2e 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -133,6 +133,26 @@ class _L2ProgramHandle: runtime_name: str +@dataclass(frozen=True) +class _DeferredChildTensor: + """CPU tensor that should be materialized as child memory inside an L3 run.""" + + runtime_name: str + tensor: torch.Tensor + upload: bool = True + refresh: bool = False + require_existing: bool = False + + +@dataclass(frozen=True) +class _PostSubmitCopy: + """Device-to-host copy to enqueue after an L3 child task has drained.""" + + dst: int + src: int + nbytes: int + + class Qwen314BModelRunner(ModelRunner): """Runtime wrapper for one Qwen3-14B model's compiled PyPTO kernels.""" @@ -227,7 +247,6 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: k_cache, v_cache = self.materialize_full_layer_cache( model.config.model_id, ) - refresh_kv_cache = model.config.model_id in self._l2_dirty_kv_models if decode_inputs.actual_batch > compiled.decode_logits_buffer.shape[0]: raise ValueError( @@ -249,8 +268,8 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: decode_inputs.slot_mapping, self._l2_child_tensor(compiled.decode.runtime_name, compiled.rope_cos), self._l2_child_tensor(compiled.decode.runtime_name, compiled.rope_sin), - self._l2_child_tensor(compiled.decode.runtime_name, k_cache, refresh=refresh_kv_cache), - self._l2_child_tensor(compiled.decode.runtime_name, v_cache, refresh=refresh_kv_cache), + k_cache, + v_cache, self._l2_child_tensor(compiled.decode.runtime_name, dw["decode_wo"]), self._l2_child_tensor(compiled.decode.runtime_name, dw["decode_post_rms_weight"]), self._l2_child_tensor(compiled.decode.runtime_name, dw["decode_w_gate"]), @@ -293,67 +312,97 @@ def _project_logits(self, model: RuntimeModel, hidden: torch.Tensor) -> torch.Te compiled.lm_head(normed, compiled.padded_lm_head_weight, logits_padded, config=None) return logits_padded[:actual_batch, :vocab_size].to(hidden.device) - normed: torch.Tensor | WorkerTensor - x_arg: torch.Tensor | WorkerTensor = x - worker: LlmWorker | None = None + normed = torch.zeros((_LOGITS_BATCH_TILE, hidden_size), dtype=torch.bfloat16).share_memory_() + normed_arg: torch.Tensor | _DeferredChildTensor = normed if compiled.final_rms.runtime_name == compiled.lm_head.runtime_name: - worker = self._worker_for_runtime(compiled.final_rms.runtime_name) - x_arg = worker.alloc_tensor(x.shape, x.dtype, init=x) - normed = worker.alloc_tensor(x.shape, x.dtype) - else: - normed = torch.zeros((_LOGITS_BATCH_TILE, hidden_size), dtype=torch.bfloat16).share_memory_() + normed_arg = self._l2_child_tensor(compiled.final_rms.runtime_name, normed, upload=False) - try: - self._run_l2_program( - compiled.final_rms, - x_arg, - self._l2_child_tensor(compiled.final_rms.runtime_name, compiled.final_norm_weight), - normed, - ) + self._run_l2_program( + compiled.final_rms, + x, + self._l2_child_tensor(compiled.final_rms.runtime_name, compiled.final_norm_weight), + normed_arg, + ) - logits_padded = torch.zeros((_LOGITS_BATCH_TILE, padded_vocab), dtype=torch.float32).share_memory_() - self._run_l2_program( - compiled.lm_head, + lm_head_input: torch.Tensor | _DeferredChildTensor = normed + if isinstance(normed_arg, _DeferredChildTensor): + lm_head_input = self._l2_child_tensor( + compiled.lm_head.runtime_name, normed, - self._l2_child_tensor(compiled.lm_head.runtime_name, compiled.padded_lm_head_weight), - logits_padded, + upload=False, + require_existing=True, ) - finally: - if isinstance(x_arg, WorkerTensor): - if worker is None: - raise RuntimeError("missing L2 worker for child-memory logits projection") - worker.free_tensor(x_arg) - if isinstance(normed, WorkerTensor): - if worker is None: - raise RuntimeError("missing L2 worker for child-memory logits projection") - worker.free_tensor(normed) + + logits_padded = torch.zeros((_LOGITS_BATCH_TILE, padded_vocab), dtype=torch.float32).share_memory_() + self._run_l2_program( + compiled.lm_head, + lm_head_input, + self._l2_child_tensor(compiled.lm_head.runtime_name, compiled.padded_lm_head_weight), + logits_padded, + ) return logits_padded[:actual_batch, :vocab_size].to(hidden.device) def _run_l2_program(self, callable_spec: _L2Callable, *args: Any) -> None: - """Run a compiled non-L3 program through the LLM Simpler worker.""" + """Run one compiled chip program through a wrapped L3 worker.""" from simpler.task_interface import CallConfig # noqa: PLC0415 + self._share_l2_host_args(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] - worker.run(handle.callable_id, orch_args, cfg) + post_submit_copies: list[_PostSubmitCopy] = [] + + def _orch_fn(orchestrator: Any, _args: Any, _config: Any) -> None: + orch_args, copies = self._build_l2_orch_args( + callable_spec, + args, + worker=worker, + orchestrator=orchestrator, + ) + post_submit_copies.extend(copies) + worker.submit_next_level( + handle.callable_id, + orch_args, + cfg, + worker_id=0, + orchestrator=orchestrator, + ) + + try: + worker.run(_orch_fn) + if post_submit_copies: + def _copy_back_orch_fn(orchestrator: Any, _args: Any, _config: Any) -> None: + for copy in post_submit_copies: + worker.copy_from( + copy.dst, + copy.src, + copy.nbytes, + worker_id=0, + orchestrator=orchestrator, + ) + + worker.run(_copy_back_orch_fn) + finally: + self._discard_l2_worker(handle.runtime_name) def _worker_for_runtime(self, runtime_name: str) -> LlmWorker: - """Return an initialized worker for ``runtime_name``.""" + """Return a wrapped L3 worker for ``runtime_name``.""" worker = self._l2_workers.get(runtime_name) if worker is not None: return worker + for active_runtime_name in list(self._l2_workers): + if active_runtime_name != runtime_name: + self._close_l2_worker(active_runtime_name) worker = LlmWorker( - level=2, + level=3, platform=self._platform, runtime=runtime_name, - device_id=self._device_id, - auto_init=True, + device_ids=[self._device_id], + num_sub_workers=0, ) self._l2_workers[runtime_name] = worker return worker @@ -366,7 +415,20 @@ def _ensure_l2_program(self, callable_spec: _L2Callable) -> _L2ProgramHandle: return cached worker = self._worker_for_runtime(callable_spec.runtime_name) + if not worker.initialized: + handle = self._register_l2_program(callable_spec, worker) + self._register_known_l2_programs_for_runtime(callable_spec.runtime_name, worker) + worker.init() + return handle + + return self._register_l2_program(callable_spec, worker) + def _register_l2_program(self, callable_spec: _L2Callable, worker: LlmWorker) -> _L2ProgramHandle: + """Register ``callable_spec`` on ``worker`` once and cache the callable id.""" + key = id(callable_spec) + cached = self._l2_programs.get(key) + if cached is not None: + return cached handle = _L2ProgramHandle( callable_id=worker.register(callable_spec.chip_callable), runtime_name=callable_spec.runtime_name, @@ -374,6 +436,20 @@ def _ensure_l2_program(self, callable_spec: _L2Callable) -> _L2ProgramHandle: self._l2_programs[key] = handle return handle + def _register_known_l2_programs_for_runtime(self, runtime_name: str, worker: LlmWorker) -> None: + """Pre-register known callables for a runtime before the L3 worker starts.""" + compiled = self._compiled + if not isinstance(compiled, _CompiledKernels): + return + for callable_spec in ( + compiled.prefill, + compiled.decode, + compiled.final_rms, + compiled.lm_head, + ): + if isinstance(callable_spec, _L2Callable) and callable_spec.runtime_name == runtime_name: + self._register_l2_program(callable_spec, worker) + def _l2_child_tensor( self, runtime_name: str, @@ -381,15 +457,32 @@ def _l2_child_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 - + require_existing: bool = False, + ) -> _DeferredChildTensor: + """Return a deferred worker-resident view for a CPU tensor backing storage.""" 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) + return _DeferredChildTensor( + runtime_name=runtime_name, + tensor=self._share_cpu_tensor(tensor), + upload=upload, + refresh=refresh, + require_existing=require_existing, + ) + + def _resolve_l3_child_tensor( + self, + deferred: _DeferredChildTensor, + *, + worker: LlmWorker, + orchestrator: Any, + ) -> tuple[WorkerTensor, _PostSubmitCopy | None]: + """Allocate or reuse a child-memory tensor inside an L3 orchestration scope.""" + from simpler_setup.torch_interop import torch_dtype_to_datatype # noqa: PLC0415 + + tensor = deferred.tensor storage = tensor.untyped_storage() storage_ptr = int(storage.data_ptr()) storage_nbytes = int(storage.nbytes()) @@ -397,26 +490,27 @@ def _l2_child_tensor( 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) + key = (deferred.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) + if deferred.require_existing: + raise RuntimeError("missing worker-resident tensor allocation") + dev_ptr = worker.malloc(storage_nbytes, worker_id=0, orchestrator=orchestrator) + if deferred.upload: + worker.copy_to(dev_ptr, storage_ptr, storage_nbytes, worker_id=0, orchestrator=orchestrator) 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) + elif deferred.upload and deferred.refresh: + worker.copy_to(alloc[0], storage_ptr, storage_nbytes, worker_id=0, orchestrator=orchestrator) dev_base, _ = alloc shape = tuple(int(dim) for dim in tensor.shape) - return WorkerTensor( + worker_tensor = WorkerTensor( data_ptr=dev_base + tensor_offset, shape=shape, dtype=torch_dtype_to_datatype(tensor.dtype), ) + return worker_tensor, _PostSubmitCopy(storage_ptr, dev_base, storage_nbytes) @staticmethod def _share_cpu_tensor(tensor: torch.Tensor) -> torch.Tensor: @@ -425,23 +519,65 @@ def _share_cpu_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor.share_memory_() return tensor + @classmethod + def _share_l2_host_args(cls, args: tuple[Any, ...]) -> None: + """Share host tensor mappings before L3 child workers are initialized.""" + for arg in args: + if isinstance(arg, _DeferredChildTensor): + cls._share_cpu_tensor(arg.tensor) + elif isinstance(arg, torch.Tensor): + cls._share_cpu_tensor(arg) + def close(self) -> None: """Release non-L3 child-memory allocations and L2 workers.""" - 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 runtime_name in list(self._l2_workers): + self._close_l2_worker(runtime_name) self._l2_dirty_kv_models.clear() - for worker in self._l2_workers.values(): + + def _close_l2_worker(self, runtime_name: str) -> None: + """Close one wrapped L3 worker and forget runtime-local child state.""" + worker = self._l2_workers.pop(runtime_name, None) + for key, (dev_ptr, _nbytes) in list(self._l2_child_allocs.items()): + alloc_runtime_name, _storage_ptr = key + if alloc_runtime_name != runtime_name: + continue + self._l2_child_allocs.pop(key, None) + if worker is not None and worker.initialized and worker.level < 3: + worker.free(dev_ptr) + for key, handle in list(self._l2_programs.items()): + if handle.runtime_name == runtime_name: + self._l2_programs.pop(key, None) + 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, ContinuousTensor, scalar_to_uint64 # noqa: PLC0415 + def _discard_l2_worker(self, runtime_name: str) -> None: + """Discard one-shot L3 worker state after a submitted chip task.""" + worker = self._l2_workers.pop(runtime_name, None) + for key in list(self._l2_child_allocs): + alloc_runtime_name, _storage_ptr = key + if alloc_runtime_name == runtime_name: + self._l2_child_allocs.pop(key, None) + for key, handle in list(self._l2_programs.items()): + if handle.runtime_name == runtime_name: + self._l2_programs.pop(key, None) + if worker is not None: + worker.discard_l3_children() + + def _build_l2_orch_args( + self, + callable_spec: _L2Callable, + args: tuple[Any, ...], + *, + worker: LlmWorker | None = None, + orchestrator: Any = None, + ): + """Build ``TaskArgs`` for submitting a compiled child chip program.""" + from simpler.task_interface import ( # noqa: PLC0415 + ContinuousTensor, + TaskArgs, + TensorArgType, + scalar_to_uint64, + ) from simpler_setup.torch_interop import make_tensor_arg # noqa: PLC0415 param_infos = callable_spec.param_infos @@ -451,18 +587,30 @@ def _build_l2_orch_args(callable_spec: _L2Callable, args: tuple[Any, ...]): 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): + orch_args = TaskArgs() + post_submit_copies: list[_PostSubmitCopy] = [] + for arg_index, (info, arg) in enumerate(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 + tag = self._l2_tensor_arg_type(info, arg, arg_index, len(args), TensorArgType) + if isinstance(arg, _DeferredChildTensor): + if worker is None or orchestrator is None: + raise RuntimeError("deferred child tensor requires L3 worker and orchestrator") + arg, copy = self._resolve_l3_child_tensor( + arg, + worker=worker, + orchestrator=orchestrator, + ) + if tag == TensorArgType.INOUT and copy is not None: + post_submit_copies.append(copy) if isinstance(arg, WorkerTensor): - orch_args.add_tensor(arg.to_continuous_tensor()) + orch_args.add_tensor(arg.to_continuous_tensor(), tag) continue if isinstance(arg, ContinuousTensor): - orch_args.add_tensor(arg) + orch_args.add_tensor(arg, tag) continue if not isinstance(arg, torch.Tensor): raise TypeError(f"tensor parameter {info.name!r} expects torch.Tensor, got {type(arg).__name__}") @@ -472,8 +620,36 @@ def _build_l2_orch_args(callable_spec: _L2Callable, args: tuple[Any, ...]): 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 + orch_args.add_tensor(make_tensor_arg(arg), tag) + return orch_args, post_submit_copies + + @staticmethod + def _l2_tensor_arg_type( + info: object, + arg: object, + arg_index: int, + arg_count: int, + tensor_arg_type: object, + ) -> object: + """Translate PyPTO parameter metadata into a Simpler tensor argument tag.""" + name = getattr(info, "name", "") + if name in {"k_cache", "v_cache"}: + return tensor_arg_type.INOUT + + direction = getattr(info, "direction", None) + direction_name = getattr(direction, "name", None) + if direction_name is None and direction is not None: + direction_name = str(direction).rsplit(".", maxsplit=1)[-1] + if direction_name == "InOut": + return tensor_arg_type.INOUT + if direction_name == "Out": + return tensor_arg_type.OUTPUT_EXISTING + + if name in {"out", "output"}: + return tensor_arg_type.OUTPUT_EXISTING + if isinstance(arg, torch.Tensor) and arg_index == arg_count - 1: + return tensor_arg_type.OUTPUT_EXISTING + return tensor_arg_type.INPUT # ── L3-wrapped generate: entire prefill + decode loop in one worker.run() ── diff --git a/python/runtime/worker.py b/python/runtime/worker.py index 5854cfe..a13beb8 100644 --- a/python/runtime/worker.py +++ b/python/runtime/worker.py @@ -18,6 +18,7 @@ from __future__ import annotations +import os from collections.abc import Sequence from dataclasses import dataclass from typing import Any @@ -198,6 +199,11 @@ def initialized(self) -> bool: """Whether the wrapped Simpler worker has been initialized.""" return self._initialized + @property + def chip_contexts(self) -> Any: + """Return L3 child chip contexts when exposed by the wrapped worker.""" + return getattr(self._worker, "chip_contexts", None) + def init(self) -> None: """Initialize runtime resources if they are not already live.""" if self._initialized: @@ -212,6 +218,68 @@ def close(self) -> None: self._worker.close() self._initialized = False + def discard_l3_children(self) -> None: + """Best-effort cleanup for one-shot L3 workers whose chip child has exited.""" + if self.level < 3: + self.close() + return + if not self._initialized: + return + + from simpler.worker import ( # noqa: PLC0415 + _OFF_STATE, + _SHUTDOWN, + _buffer_field_addr, + _mailbox_store_i32, + ) + + impl = self._worker + for shm in list(getattr(impl, "_sub_shms", [])) + list(getattr(impl, "_chip_shms", [])) + list( + getattr(impl, "_next_level_shms", []) + ): + try: + buf = shm.buf + if buf is not None: + _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _SHUTDOWN) + except Exception: # noqa: BLE001 + pass + + for pid in list(getattr(impl, "_sub_pids", [])) + list(getattr(impl, "_chip_pids", [])) + list( + getattr(impl, "_next_level_pids", []) + ): + try: + os.waitpid(int(pid), 0) + except ChildProcessError: + pass + except OSError: + pass + + for shm in list(getattr(impl, "_sub_shms", [])) + list(getattr(impl, "_chip_shms", [])) + list( + getattr(impl, "_next_level_shms", []) + ): + try: + shm.close() + except Exception: # noqa: BLE001 + pass + try: + shm.unlink() + except FileNotFoundError: + pass + except Exception: # noqa: BLE001 + pass + + for attr in ("_sub_pids", "_chip_pids", "_next_level_pids"): + getattr(impl, attr, []).clear() + for attr in ("_sub_shms", "_chip_shms", "_next_level_shms", "_next_level_workers"): + getattr(impl, attr, []).clear() + if hasattr(impl, "_worker"): + impl._worker = None # noqa: SLF001 + if hasattr(impl, "_orch"): + impl._orch = None # noqa: SLF001 + if hasattr(impl, "_hierarchical_started"): + impl._hierarchical_started = False # noqa: SLF001 + self._initialized = False + def _require_initialized(self, op: str) -> None: if not self._initialized: raise RuntimeError(f"Worker.{op} requires init() first") @@ -230,36 +298,91 @@ def prepare_callable(self, callable_id: int, callable_obj: Any) -> None: self._require_initialized("prepare_callable") self._worker.prepare_callable(int(callable_id), callable_obj) - def run(self, callable_obj: Any, args: Any = None, config: Any = None) -> None: + def run(self, callable_obj: Any, args: Any = None, config: Any = None) -> Any: """Run one L2 callable id or one L3 orchestration function.""" self._require_initialized("run") - self._worker.run(callable_obj, args=args, config=config) + return self._worker.run(callable_obj, args=args, config=config) + + def submit_next_level( + self, + callable_id: int, + args: Any, + config: Any = None, + *, + worker_id: int = 0, + orchestrator: Any, + ) -> Any: + """Submit a child chip callable from inside an L3 orchestration callback.""" + if self.level < 3: + raise RuntimeError("Worker.submit_next_level requires a level 3 worker") + return orchestrator.submit_next_level( + int(callable_id), + args, + config, + worker=int(worker_id), + ) + + def run_next_level( + self, + callable_id: int, + args: Any, + config: Any = None, + *, + worker_id: int = 0, + ) -> Any: + """Run one chip callable through an L3 orchestration scope.""" + if self.level < 3: + return self.run(int(callable_id), args=args, config=config) + + def _orch_fn(orchestrator: Any, _args: Any, _config: Any) -> None: + self.submit_next_level( + int(callable_id), + args, + config, + worker_id=worker_id, + orchestrator=orchestrator, + ) + + return self.run(_orch_fn) def run_prepared(self, callable_id: int, args: Any = None, config: Any = None) -> None: """Run a callable previously staged with :meth:`prepare_callable`.""" self._require_initialized("run_prepared") self._worker.run_prepared(int(callable_id), args=args, config=config) - def malloc(self, nbytes: int, *, worker_id: int = 0) -> int: + def malloc(self, nbytes: int, *, worker_id: int = 0, orchestrator: Any = None) -> int: """Allocate bytes on a worker child and return the device pointer.""" self._require_initialized("malloc") nbytes = _validate_positive_int("nbytes", nbytes) + if orchestrator is not None: + return int(orchestrator.malloc(worker_id=int(worker_id), size=nbytes)) return self._worker.malloc(nbytes, worker_id=int(worker_id)) - def free(self, ptr: int, *, worker_id: int = 0) -> None: + def free(self, ptr: int, *, worker_id: int = 0, orchestrator: Any = None) -> None: """Free a pointer previously returned by :meth:`malloc`.""" self._require_initialized("free") + if orchestrator is not None: + orchestrator.free(worker_id=int(worker_id), ptr=int(ptr)) + return self._worker.free(int(ptr), worker_id=int(worker_id)) - def copy_to(self, dst: int, src: int, nbytes: int, *, worker_id: int = 0) -> None: + def copy_to(self, dst: int, src: int, nbytes: int, *, worker_id: int = 0, orchestrator: Any = None) -> None: """Copy ``nbytes`` from a host pointer to worker device memory.""" self._require_initialized("copy_to") - self._worker.copy_to(int(dst), int(src), _validate_positive_int("nbytes", nbytes), worker_id=int(worker_id)) + nbytes = _validate_positive_int("nbytes", nbytes) + if orchestrator is not None: + orchestrator.copy_to(worker_id=int(worker_id), dst=int(dst), src=int(src), size=nbytes) + return + self._worker.copy_to(int(dst), int(src), nbytes, worker_id=int(worker_id)) - def copy_from(self, dst: int, src: int, nbytes: int, *, worker_id: int = 0) -> None: + def copy_from(self, dst: int, src: int, nbytes: int, *, worker_id: int = 0, orchestrator: Any = None) -> None: """Copy ``nbytes`` from worker device memory to a host pointer.""" self._require_initialized("copy_from") - self._worker.copy_from(int(dst), int(src), _validate_positive_int("nbytes", nbytes), worker_id=int(worker_id)) + nbytes = _validate_positive_int("nbytes", nbytes) + if orchestrator is not None: + orchestrator.copy_from(worker_id=int(worker_id), dst=int(dst), src=int(src), size=nbytes) + return + self._worker.copy_from(int(dst), int(src), nbytes, worker_id=int(worker_id)) def alloc_tensor( self, diff --git a/tests/test_npu_runner_l3_worker.py b/tests/test_npu_runner_l3_worker.py new file mode 100644 index 0000000..49edb1e --- /dev/null +++ b/tests/test_npu_runner_l3_worker.py @@ -0,0 +1,245 @@ +# 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 + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import torch +from simpler.task_interface import TensorArgType + +from examples.model.qwen3_14b.runner import npu_runner +from examples.model.qwen3_14b.runner.npu_runner import Qwen314BModelRunner +from examples.model.qwen3_14b.runner.npu_runner import _CompiledKernels +from examples.model.qwen3_14b.runner.npu_runner import _L2Callable + + +@dataclass +class _ParamInfo: + name: str + shape: tuple[int, ...] | None + direction: object | None = None + + +class _FakeOrchestrator: + def __init__(self, worker: "_FakeL3Worker") -> None: + self._worker = worker + self.ops: list[tuple[Any, ...]] = [] + + def malloc(self, worker_id: int, size: int) -> int: + ptr = self._worker.next_ptr + self._worker.next_ptr += int(size) + 64 + self.ops.append(("malloc", worker_id, size, ptr)) + return ptr + + def copy_to(self, worker_id: int, dst: int, src: int, size: int) -> None: + self.ops.append(("copy_to", worker_id, dst, src, size)) + + def copy_from(self, worker_id: int, dst: int, src: int, size: int) -> None: + self.ops.append(("copy_from", worker_id, dst, src, size)) + + def submit_next_level(self, callable_id: int, args: Any, config: Any, worker: int = -1) -> None: + self.ops.append(("submit_next_level", callable_id, args, config.block_dim, worker)) + + +class _FakeL3Worker: + instances: list["_FakeL3Worker"] = [] + + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + self.level = int(kwargs["level"]) + self.initialized = False + self.closed = False + self.registered: list[object] = [] + self.runs: list[list[tuple[Any, ...]]] = [] + self.freed: list[int] = [] + self.next_ptr = 0x100000 + _FakeL3Worker.instances.append(self) + + def register(self, callable_obj: object) -> int: + self.registered.append(callable_obj) + return len(self.registered) - 1 + + def init(self) -> None: + self.initialized = True + + def run(self, callable_obj: object, args: Any = None, config: Any = None) -> object: + orchestrator = _FakeOrchestrator(self) + callable_obj(orchestrator, args, config) + self.runs.append(orchestrator.ops) + return SimpleNamespace(host_wall_us=1, device_wall_us=0) + + def submit_next_level( + self, + callable_id: int, + args: Any, + config: Any = None, + *, + worker_id: int = 0, + orchestrator: _FakeOrchestrator, + ) -> None: + orchestrator.submit_next_level(callable_id, args, config, worker=worker_id) + + def malloc(self, nbytes: int, *, worker_id: int = 0, orchestrator: _FakeOrchestrator) -> int: + return orchestrator.malloc(worker_id=worker_id, size=nbytes) + + def copy_to( + self, + dst: int, + src: int, + nbytes: int, + *, + worker_id: int = 0, + orchestrator: _FakeOrchestrator, + ) -> None: + orchestrator.copy_to(worker_id=worker_id, dst=dst, src=src, size=nbytes) + + def copy_from( + self, + dst: int, + src: int, + nbytes: int, + *, + worker_id: int = 0, + orchestrator: _FakeOrchestrator, + ) -> None: + orchestrator.copy_from(worker_id=worker_id, dst=dst, src=src, size=nbytes) + + def free(self, _ptr: int) -> None: + self.freed.append(_ptr) + + def close(self) -> None: + self.initialized = False + self.closed = True + + def discard_l3_children(self) -> None: + self.initialized = False + self.closed = True + + +def _callable(runtime_name: str, params: tuple[_ParamInfo, ...]) -> _L2Callable: + return _L2Callable( + chip_callable=object(), + runtime_name=runtime_name, + block_dim=24, + aicpu_thread_num=1, + param_infos=params, + ) + + +def _compiled(callable_spec: _L2Callable) -> _CompiledKernels: + return _CompiledKernels( + prefill=callable_spec, + decode=callable_spec, + final_rms=None, + lm_head=None, + final_norm_weight=torch.ones(1, 4), + rope_cos=torch.zeros(1, 4), + rope_sin=torch.zeros(1, 4), + padded_vocab=8, + padded_lm_head_weight=torch.zeros(8, 4), + layers=[], + decode_weights={}, + decode_logits_buffer=torch.zeros(1, 8), + ) + + +def _runner(compiled: _CompiledKernels) -> Qwen314BModelRunner: + return Qwen314BModelRunner( + model_id="model", + compiled=compiled, + platform="a2a3", + device_id=3, + save_kernels_dir=None, + l3_trace=False, + ) + + +def test_non_l3_program_runs_through_l3_worker(monkeypatch): + _FakeL3Worker.instances.clear() + monkeypatch.setattr(npu_runner, "LlmWorker", _FakeL3Worker) + params = ( + _ParamInfo("weight", (2, 2)), + _ParamInfo("out", (2, 2)), + ) + callable_spec = _callable("prefill", params) + runner = _runner(_compiled(callable_spec)) + weight = torch.ones(2, 2, dtype=torch.float32) + out = torch.zeros(2, 2, dtype=torch.float32).share_memory_() + + runner._run_l2_program(callable_spec, runner._l2_child_tensor(callable_spec.runtime_name, weight), out) + + worker = _FakeL3Worker.instances[0] + assert worker.kwargs["level"] == 3 + assert worker.kwargs["device_ids"] == [3] + assert worker.kwargs["num_sub_workers"] == 0 + assert worker.closed + assert callable_spec.chip_callable in worker.registered + assert len(worker.runs) == 1 + assert [op[0] for op in worker.runs[0]] == ["malloc", "copy_to", "submit_next_level"] + assert worker.runs[0][-1][2].tag(0) == TensorArgType.INPUT + assert worker.runs[0][-1][2].tag(1) == TensorArgType.OUTPUT_EXISTING + + +def test_full_kv_host_tensors_are_inout_without_extra_copyback(monkeypatch): + _FakeL3Worker.instances.clear() + monkeypatch.setattr(npu_runner, "LlmWorker", _FakeL3Worker) + params = ( + _ParamInfo("k_cache", (4, 2)), + _ParamInfo("v_cache", (4, 2)), + _ParamInfo("out", (1, 4)), + ) + callable_spec = _callable("decode", params) + runner = _runner(_compiled(callable_spec)) + k_cache = torch.zeros(4, 2, dtype=torch.bfloat16) + v_cache = torch.zeros(4, 2, dtype=torch.bfloat16) + out = torch.zeros(1, 4, dtype=torch.float32).share_memory_() + + runner._run_l2_program( + callable_spec, + k_cache, + v_cache, + out, + ) + + worker = _FakeL3Worker.instances[0] + assert len(worker.runs) == 1 + submit_args = worker.runs[0][-1][2] + assert submit_args.tag(0) == TensorArgType.INOUT + assert submit_args.tag(1) == TensorArgType.INOUT + assert submit_args.tag(2) == TensorArgType.OUTPUT_EXISTING + + +def test_switching_runtime_closes_previous_l3_worker(monkeypatch): + _FakeL3Worker.instances.clear() + monkeypatch.setattr(npu_runner, "LlmWorker", _FakeL3Worker) + params = ( + _ParamInfo("weight", (2, 2)), + _ParamInfo("out", (2, 2)), + ) + prefill = _callable("prefill-rt", params) + decode = _callable("decode-rt", params) + compiled = _compiled(prefill) + compiled.decode = decode + runner = _runner(compiled) + weight = torch.ones(2, 2, dtype=torch.float32) + out = torch.zeros(2, 2, dtype=torch.float32).share_memory_() + + runner._run_l2_program(prefill, runner._l2_child_tensor("prefill-rt", weight), out) + prefill_worker = _FakeL3Worker.instances[0] + assert prefill_worker.closed + + runner._run_l2_program(decode, runner._l2_child_tensor("decode-rt", weight), out) + + assert prefill_worker.closed + assert not prefill_worker.freed + assert len(_FakeL3Worker.instances) == 2 + assert _FakeL3Worker.instances[1].kwargs["runtime"] == "decode-rt"