diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py b/examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py index 37c573e36e..5482a36f9c 100644 --- a/examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py @@ -14,7 +14,7 @@ import sys from pathlib import Path -from simpler_setup import SceneTestCase, scene_test +from simpler_setup import PinnedTorchAllocator, SceneTestCase, scene_test from simpler_setup.goldens.qwen3_14b_decode import compute_golden as _decode_golden from simpler_setup.goldens.qwen3_14b_decode import generate_inputs as _decode_generate_inputs @@ -64,6 +64,14 @@ def generate_args(self, params): n_layers=N_LAYERS, ) + def generate_args_for_worker(self, worker, params): + return _decode_generate_inputs( + seed=params.get("seed", 1234), + seq_len=params.get("seq_len", 3500), + n_layers=N_LAYERS, + allocator=PinnedTorchAllocator(worker), + ) + def compute_golden(self, args, params): _decode_golden(args, n_layers=N_LAYERS) diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 484c817640..bc54bdf3e1 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -3034,6 +3034,8 @@ NB_MODULE(_task_interface, m) { ) .def("malloc", &ChipWorker::malloc, nb::arg("size")) .def("free", &ChipWorker::free, nb::arg("ptr")) + .def("alloc_pinned_host", &ChipWorker::alloc_pinned_host, nb::arg("size")) + .def("free_pinned_host", &ChipWorker::free_pinned_host, nb::arg("ptr")) .def("copy_to", &ChipWorker::copy_to, nb::arg("dst"), nb::arg("src"), nb::arg("size")) .def("copy_from", &ChipWorker::copy_from, nb::arg("dst"), nb::arg("src"), nb::arg("size")) .def( diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 06772d983a..9c5b66b704 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -1603,6 +1603,14 @@ def free(self, ptr): """Free memory allocated by ``malloc()``.""" self._impl.free(int(ptr)) + def alloc_pinned_host(self, size): + """Allocate page-locked host memory. Returns a pointer (uint64).""" + return int(self._impl.alloc_pinned_host(int(size))) + + def free_pinned_host(self, ptr): + """Release memory allocated by ``alloc_pinned_host()``.""" + self._impl.free_pinned_host(int(ptr)) + def copy_to(self, dst, src, size): """Copy *size* bytes from host *src* to worker *dst*.""" self._impl.copy_to(int(dst), int(src), int(size)) diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 45777523b1..ee30dcc09d 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -4305,6 +4305,53 @@ def _on_cancel(_signum, _frame): # --------------------------------------------------------------------------- +class _PinnedHostAllocation: + """Lifetime token for one ChipWorker-owned page-locked host allocation.""" + + __slots__ = ("base", "nbytes", "_worker") + + def __init__(self, worker: ChipWorker, nbytes: int) -> None: + self._worker = worker + self.nbytes = int(nbytes) + self.base = worker.alloc_pinned_host(self.nbytes) + + def __del__(self) -> None: + base = self.base + if base == 0: + return + self.base = 0 + try: + self._worker.free_pinned_host(base) + except Exception as exc: # noqa: BLE001 -- ChipWorker.finalize already reclaims every still-live allocation + sys.stderr.write(f"PinnedHostBuffer cleanup: free_pinned_host failed (continuing best-effort): {exc}\n") + sys.stderr.flush() + + +class PinnedHostBuffer: + """A page-locked host byte span whose exported views keep its allocation alive. + + Use ``buffer`` with consumers of Python's buffer protocol, for example + ``torch.frombuffer(handle.buffer, dtype=...)``. The tensor retains the + ctypes exporter, which retains the allocation token, so dropping this + handle cannot release storage while a tensor view still exists. + """ + + __slots__ = ("base", "nbytes", "_buffer") + + def __init__(self, worker: ChipWorker, nbytes: int) -> None: + allocation = _PinnedHostAllocation(worker, nbytes) + raw = (ctypes.c_ubyte * allocation.nbytes).from_address(allocation.base) + raw._simpler_pinned_owner = allocation + self.base = allocation.base + self.nbytes = allocation.nbytes + self._buffer = raw + + @property + def buffer(self): + """Writable object implementing Python's contiguous buffer protocol.""" + return self._buffer + + class Worker: """Unified worker for all hierarchy levels. @@ -10259,6 +10306,23 @@ def copy_from(self, dst, src: Buffer) -> None: # Post-fork zero-copy host buffers # ------------------------------------------------------------------ + def alloc_pinned_host(self, nbytes: int) -> PinnedHostBuffer: + """Allocate page-locked host storage for direct L2 host/device copies. + + The returned byte span is local to this process and therefore only + valid on an L2 Worker. Build producer tensors directly over its + ``buffer``; copying an existing pageable tensor into it would be a + bounce buffer and defeats this API's purpose. + """ + if self.level != 2: + raise TypeError("alloc_pinned_host requires a level-2 Worker") + nbytes = int(nbytes) + if nbytes <= 0: + raise ValueError("alloc_pinned_host: nbytes must be positive") + with self._operation_lease("alloc_pinned_host"): + assert self._chip_worker is not None + return PinnedHostBuffer(self._chip_worker, nbytes) + def create_buffer(self, nbytes: int) -> Buffer: """Allocate a shared ``Buffer`` owned by this Worker (P1-B). diff --git a/simpler_setup/__init__.py b/simpler_setup/__init__.py index b6792c60f0..0403e16ad3 100644 --- a/simpler_setup/__init__.py +++ b/simpler_setup/__init__.py @@ -23,11 +23,12 @@ scene_level, scene_test, ) -from .torch_interop import make_chip_tensor_arg, torch_dtype_to_datatype +from .torch_interop import PinnedTorchAllocator, make_chip_tensor_arg, torch_dtype_to_datatype __all__ = [ "CallableNamespace", "KernelCompiler", + "PinnedTorchAllocator", "RuntimeBuilder", "Scalar", "SceneTestCase", diff --git a/simpler_setup/goldens/qwen3_14b_decode.py b/simpler_setup/goldens/qwen3_14b_decode.py index 5196306749..a10b6dcfeb 100644 --- a/simpler_setup/goldens/qwen3_14b_decode.py +++ b/simpler_setup/goldens/qwen3_14b_decode.py @@ -103,10 +103,11 @@ def _rope_half(vec: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch return torch.cat([lo * cos_lo - hi * sin_lo, hi * cos_hi + lo * sin_hi], dim=-1) -def _paged_block_table_slot_mapping(seq_lens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: +def _paged_block_table_slot_mapping(seq_lens: torch.Tensor, empty=torch.empty) -> tuple[torch.Tensor, torch.Tensor]: """Identity paging within a layer pool (same for every layer).""" - block_table = torch.arange(BATCH * MAX_BLOCKS_PER_SEQ, dtype=torch.int32) - slot_mapping = torch.empty(BATCH, dtype=torch.int32) + block_table = empty((BATCH * MAX_BLOCKS_PER_SEQ,), dtype=torch.int32) + torch.arange(BATCH * MAX_BLOCKS_PER_SEQ, dtype=torch.int32, out=block_table) + slot_mapping = empty((BATCH,), dtype=torch.int32) for b in range(BATCH): pos = int(seq_lens[b].item()) - 1 logical_block = pos // BLOCK_SIZE @@ -119,6 +120,7 @@ def generate_inputs( seed: int = 1234, seq_len: int = DEFAULT_SEQ_LEN, n_layers: int = N_LAYERS, + allocator=None, ) -> TaskArgsBuilder: """Deterministic fixture for decode_fwd_layers, stacked along dim 0. @@ -136,6 +138,9 @@ def generate_inputs( raise ValueError(f"n_layers must be positive, got {n_layers}") g = torch.Generator().manual_seed(seed) + if allocator is not None: + return _generate_direct_inputs(allocator, g, seq_len, n_layers) + def rn(shape, std=1.0, bias=0.0): return torch.empty(shape).normal_(0.0, std, generator=g) + bias @@ -177,6 +182,70 @@ def s0(t): # replicate along dim 0 (one slice per layer) return TaskArgsBuilder(*specs) +def _generate_direct_inputs(allocator, generator, seq_len: int, n_layers: int) -> TaskArgsBuilder: + """Generate every final Qwen arg directly in allocator-owned storage.""" + + def empty(shape, *, dtype=torch.float32): + return allocator.empty(shape, dtype=dtype) + + def normal(shape, std=1.0, bias=0.0, *, dtype=torch.float32): + result = empty(shape, dtype=dtype) + result.normal_(0.0, std, generator=generator) + if bias: + result.add_(bias) + return result + + def stacked(shape, std=1.0, bias=0.0, *, dtype=torch.float32): + rows = int(shape[0]) + result = empty((n_layers * rows, *shape[1:]), dtype=dtype) + first = result[:rows] + first.normal_(0.0, std, generator=generator) + if bias: + first.add_(bias) + for layer in range(1, n_layers): + result[layer * rows : (layer + 1) * rows].copy_(first) + return result + + seq_lens = empty((BATCH,), dtype=torch.int32).fill_(seq_len) + block_table, slot_mapping = _paged_block_table_slot_mapping(seq_lens, empty) + + posv = torch.arange(MAX_SEQ).float().unsqueeze(1) + inv_freq = 1.0 / (ROPE_THETA ** (torch.arange(0, HALF_DIM).float() / HALF_DIM)) + ang = posv * inv_freq.unsqueeze(0) + rope_cos = empty((MAX_SEQ, HEAD_DIM)) + rope_sin = empty((MAX_SEQ, HEAD_DIM)) + torch.cos(ang, out=rope_cos[:, :HALF_DIM]) + torch.sin(ang, out=rope_sin[:, :HALF_DIM]) + rope_cos[:, HALF_DIM:].copy_(rope_cos[:, :HALF_DIM]) + rope_sin[:, HALF_DIM:].copy_(rope_sin[:, :HALF_DIM]) + + tensors = { + "hidden_states": normal((BATCH, HIDDEN), dtype=torch.bfloat16), + "input_rms_weight": stacked((1, HIDDEN), 0.1, 1.0), + "wq": stacked((HIDDEN, HIDDEN), 0.02, dtype=torch.bfloat16), + "wk": stacked((HIDDEN, KV_HIDDEN), 0.02, dtype=torch.bfloat16), + "wv": stacked((HIDDEN, KV_HIDDEN), 0.02, dtype=torch.bfloat16), + "q_norm_weight": stacked((1, HEAD_DIM), 0.1, 1.0), + "k_norm_weight": stacked((1, HEAD_DIM), 0.1, 1.0), + "seq_lens": seq_lens, + "block_table": block_table, + "slot_mapping": slot_mapping, + "rope_cos": rope_cos, + "rope_sin": rope_sin, + "k_cache": stacked((CACHE_ROWS, HEAD_DIM), 0.01, dtype=torch.bfloat16), + "v_cache": stacked((CACHE_ROWS, HEAD_DIM), 0.02, 0.3, dtype=torch.bfloat16), + "wo": stacked((HIDDEN, HIDDEN), 0.0006, dtype=torch.bfloat16), + "w_gate": stacked((HIDDEN, INTERMEDIATE), 0.02, dtype=torch.bfloat16), + "w_up": stacked((HIDDEN, INTERMEDIATE), 0.02, dtype=torch.bfloat16), + "w_down": stacked((INTERMEDIATE, HIDDEN), 0.0004, dtype=torch.bfloat16), + "post_rms_weight": stacked((1, HIDDEN), 0.1, 1.0), + } + specs = [TensorArg(name, tensors[name]) for name in INPUT_NAMES] + output = empty((BATCH, HIDDEN), dtype=torch.bfloat16).zero_() + specs.append(TensorArg("out", output)) + return TaskArgsBuilder(*specs) + + def _one_layer(args, layer: int, x: torch.Tensor) -> torch.Tensor: """One decode layer of the chunk. ``x`` is the FP32 layer input; returns the FP32 residual-stream output (down + h1, NO bf16 round). Also writes the diff --git a/simpler_setup/scene_test.py b/simpler_setup/scene_test.py index ba5425edb1..6e3def8041 100644 --- a/simpler_setup/scene_test.py +++ b/simpler_setup/scene_test.py @@ -1445,6 +1445,15 @@ def generate_args(self, params) -> TaskArgsBuilder: """Return TaskArgsBuilder with ordered TensorArg/Scalar specs.""" raise NotImplementedError + def generate_args_for_worker(self, worker, params) -> TaskArgsBuilder: + """Build L2 args with access to worker-owned allocation APIs. + + Most cases use ordinary CPU tensors and inherit this forwarding + implementation. Cases whose producer writes directly into specialized + storage, such as pinned HBG args, override this hook. + """ + return self.generate_args(params) + def compute_golden(self, args: TaskArgsBuilder, params) -> None: """Compute expected outputs in-place on a cloned TaskArgsBuilder.""" raise NotImplementedError @@ -1650,7 +1659,7 @@ def _run_and_validate_l2( # noqa: PLR0913 -- threads CLI diagnostic flags + cas type(self)._st_l2_handle = handle # Build args - test_args = self.generate_args(params) + test_args = self.generate_args_for_worker(worker, params) chip_args, output_names = _build_l2_ref_args(test_args, orch_sig, worker) # Compute golden (unless skip_golden) diff --git a/simpler_setup/torch_interop.py b/simpler_setup/torch_interop.py index 8f981ac144..7924558892 100644 --- a/simpler_setup/torch_interop.py +++ b/simpler_setup/torch_interop.py @@ -30,6 +30,7 @@ from __future__ import annotations +import math from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -38,6 +39,30 @@ _TORCH_DTYPE_MAP = None +class PinnedTorchAllocator: + """Allocate CPU tensors directly on a Worker's page-locked host storage.""" + + def __init__(self, worker) -> None: + if int(worker.level) != 2: + raise TypeError("PinnedTorchAllocator requires a level-2 Worker") + self._worker = worker + + def empty(self, shape, *, dtype=None): + """Return a contiguous tensor whose storage comes from ``aclrtMallocHost``.""" + import torch # pyright: ignore[reportMissingImports] + + shape = tuple(int(dim) for dim in shape) + dtype = torch.get_default_dtype() if dtype is None else dtype + if any(dim < 0 for dim in shape): + raise ValueError(f"PinnedTorchAllocator.empty: shape dimensions must be non-negative, got {shape}") + numel = math.prod(shape) + if numel == 0: + return torch.empty(shape, dtype=dtype) + element_size = torch.empty((), dtype=dtype).element_size() + backing = self._worker.alloc_pinned_host(numel * element_size) + return torch.frombuffer(backing.buffer, dtype=dtype, count=numel).reshape(shape) + + def _ensure_torch_map(): global _TORCH_DTYPE_MAP if _TORCH_DTYPE_MAP is not None: diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index d68bbf269c..e5aaf841aa 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -32,6 +32,7 @@ #include "task_args.h" #include "native_run_context.h" +#include #include #include #include @@ -331,6 +332,17 @@ void device_free_ctx(DeviceContextHandle ctx, void *dev_ptr) { } catch (...) {} } +int alloc_pinned_host_ctx(DeviceContextHandle ctx, size_t size, void **host_ptr) { + if (ctx == NULL || size == 0 || host_ptr == NULL) return PTO_RUNTIME_ERR_INTERNAL; + *host_ptr = NULL; + return aclrtMallocHost(host_ptr, size); +} + +int free_pinned_host_ctx(DeviceContextHandle ctx, void *host_ptr) { + if (ctx == NULL || host_ptr == NULL) return PTO_RUNTIME_ERR_INTERNAL; + return aclrtFreeHost(host_ptr); +} + int copy_to_device_ctx(DeviceContextHandle ctx, void *dev_ptr, const void *host_ptr, size_t size) { if (ctx == NULL || dev_ptr == NULL || host_ptr == NULL) return PTO_RUNTIME_ERR_INTERNAL; try { diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index cfed968131..0ab142da0f 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -314,6 +314,18 @@ void device_free_ctx(DeviceContextHandle ctx, void *dev_ptr) { } catch (...) {} } +int alloc_pinned_host_ctx(DeviceContextHandle ctx, size_t size, void **host_ptr) { + if (ctx == NULL || size == 0 || host_ptr == NULL) return PTO_RUNTIME_ERR_INTERNAL; + *host_ptr = std::malloc(size); + return *host_ptr == NULL ? PTO_RUNTIME_ERR_INTERNAL : 0; +} + +int free_pinned_host_ctx(DeviceContextHandle ctx, void *host_ptr) { + if (ctx == NULL || host_ptr == NULL) return PTO_RUNTIME_ERR_INTERNAL; + std::free(host_ptr); + return 0; +} + int copy_to_device_ctx(DeviceContextHandle ctx, void *dev_ptr, const void *host_ptr, size_t size) { if (ctx == NULL || dev_ptr == NULL || host_ptr == NULL) return PTO_RUNTIME_ERR_INTERNAL; try { diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 2cacacd82b..30c23b3b08 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -206,6 +206,8 @@ void ChipWorker::init( destroy_device_context_fn_ = load_symbol(handle, "destroy_device_context"); device_malloc_ctx_fn_ = load_symbol(handle, "device_malloc_ctx"); device_free_ctx_fn_ = load_symbol(handle, "device_free_ctx"); + alloc_pinned_host_ctx_fn_ = load_symbol(handle, "alloc_pinned_host_ctx"); + free_pinned_host_ctx_fn_ = load_symbol(handle, "free_pinned_host_ctx"); device_committed_memory_fn_ = load_symbol(handle, "committed_device_memory_ctx"); copy_to_device_ctx_fn_ = load_symbol(handle, "copy_to_device_ctx"); copy_from_device_ctx_fn_ = load_symbol(handle, "copy_from_device_ctx"); @@ -332,6 +334,8 @@ void ChipWorker::init( destroy_device_context_fn_ = nullptr; device_malloc_ctx_fn_ = nullptr; device_free_ctx_fn_ = nullptr; + alloc_pinned_host_ctx_fn_ = nullptr; + free_pinned_host_ctx_fn_ = nullptr; device_committed_memory_fn_ = nullptr; copy_to_device_ctx_fn_ = nullptr; copy_from_device_ctx_fn_ = nullptr; @@ -383,6 +387,8 @@ void ChipWorker::init( destroy_device_context_fn_ = nullptr; device_malloc_ctx_fn_ = nullptr; device_free_ctx_fn_ = nullptr; + alloc_pinned_host_ctx_fn_ = nullptr; + free_pinned_host_ctx_fn_ = nullptr; device_committed_memory_fn_ = nullptr; copy_to_device_ctx_fn_ = nullptr; copy_from_device_ctx_fn_ = nullptr; @@ -492,6 +498,19 @@ void ChipWorker::finalize() { // communicator handles and streams before tearing down the device context. clear_comm_sessions(); + if (device_ctx_ != nullptr && free_pinned_host_ctx_fn_ != nullptr) { + for (void *ptr : pinned_host_allocations_) { + int rc = free_pinned_host_ctx_fn_(device_ctx_, ptr); + if (rc != 0) { + HostLogger::get_instance().log( + simpler::log::LogLevel::ERROR, __func__, + "free_pinned_host_ctx failed during finalization for %p: %d", ptr, rc + ); + } + } + } + pinned_host_allocations_.clear(); + if (device_ctx_ != nullptr && finalize_device_fn_ != nullptr && initialized_) { finalize_device_fn_(device_ctx_); } @@ -507,6 +526,8 @@ void ChipWorker::finalize() { destroy_device_context_fn_ = nullptr; device_malloc_ctx_fn_ = nullptr; device_free_ctx_fn_ = nullptr; + alloc_pinned_host_ctx_fn_ = nullptr; + free_pinned_host_ctx_fn_ = nullptr; device_committed_memory_fn_ = nullptr; copy_to_device_ctx_fn_ = nullptr; copy_from_device_ctx_fn_ = nullptr; @@ -1082,6 +1103,43 @@ void ChipWorker::free(uint64_t ptr) { device_free_ctx_fn_(device_ctx_, reinterpret_cast(ptr)); } +uint64_t ChipWorker::alloc_pinned_host(size_t size) { + if (!initialized_) { + throw std::runtime_error("ChipWorker not initialized; call init() first"); + } + if (size == 0) { + throw std::invalid_argument("alloc_pinned_host: size must be positive"); + } + void *ptr = nullptr; + int rc = alloc_pinned_host_ctx_fn_(device_ctx_, size, &ptr); + if (rc != 0 || ptr == nullptr) { + throw std::runtime_error( + "alloc_pinned_host(" + std::to_string(size) + ") failed with code " + std::to_string(rc) + ); + } + try { + pinned_host_allocations_.insert(ptr); + } catch (...) { + (void)free_pinned_host_ctx_fn_(device_ctx_, ptr); + throw; + } + return reinterpret_cast(ptr); +} + +void ChipWorker::free_pinned_host(uint64_t ptr) { + void *host_ptr = reinterpret_cast(ptr); + auto it = pinned_host_allocations_.find(host_ptr); + if (it == pinned_host_allocations_.end()) return; + if (!initialized_) { + throw std::runtime_error("ChipWorker not initialized; call init() first"); + } + int rc = free_pinned_host_ctx_fn_(device_ctx_, host_ptr); + if (rc != 0) { + throw std::runtime_error("free_pinned_host failed with code " + std::to_string(rc)); + } + pinned_host_allocations_.erase(it); +} + void ChipWorker::copy_to(uint64_t dst, uint64_t src, size_t size) { if (!initialized_) { throw std::runtime_error("ChipWorker not initialized; call init() first"); diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 3bfb343474..5c96235b0c 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -172,6 +172,8 @@ class ChipWorker { uint64_t malloc(size_t size); void free(uint64_t ptr); + uint64_t alloc_pinned_host(size_t size); + void free_pinned_host(uint64_t ptr); void copy_to(uint64_t dst, uint64_t src, size_t size); void copy_from(uint64_t dst, uint64_t src, size_t size); @@ -248,6 +250,8 @@ class ChipWorker { using DestroyDeviceContextFn = void (*)(void *); using DeviceMallocCtxFn = void *(*)(void *, size_t); using DeviceFreeCtxFn = void (*)(void *, void *); + using AllocPinnedHostCtxFn = int (*)(void *, size_t, void **); + using FreePinnedHostCtxFn = int (*)(void *, void *); using CopyToDeviceCtxFn = int (*)(void *, void *, const void *, size_t); using CopyFromDeviceCtxFn = int (*)(void *, void *, const void *, size_t); using GetRuntimeSizeFn = size_t (*)(); @@ -309,11 +313,14 @@ class ChipWorker { DestroyDeviceContextFn destroy_device_context_fn_ = nullptr; DeviceMallocCtxFn device_malloc_ctx_fn_ = nullptr; DeviceFreeCtxFn device_free_ctx_fn_ = nullptr; + AllocPinnedHostCtxFn alloc_pinned_host_ctx_fn_ = nullptr; + FreePinnedHostCtxFn free_pinned_host_ctx_fn_ = nullptr; CopyToDeviceCtxFn copy_to_device_ctx_fn_ = nullptr; CopyFromDeviceCtxFn copy_from_device_ctx_fn_ = nullptr; GetRuntimeSizeFn get_runtime_size_fn_ = nullptr; GetRuntimeAlignmentFn get_runtime_alignment_fn_ = nullptr; GetCommittedDeviceMemoryFn device_committed_memory_fn_ = nullptr; + std::unordered_set pinned_host_allocations_; SimplerInitFn simpler_init_fn_ = nullptr; SimplerRegisterCallableFn register_callable_fn_ = nullptr; SimplerRunFn run_fn_ = nullptr; diff --git a/src/common/worker/runtime_c_api.h b/src/common/worker/runtime_c_api.h index e0dfd44ed4..de38f3daad 100644 --- a/src/common/worker/runtime_c_api.h +++ b/src/common/worker/runtime_c_api.h @@ -242,6 +242,12 @@ void *device_malloc_ctx(DeviceContextHandle ctx, size_t size); /** Free device memory previously allocated in the given device context. */ void device_free_ctx(DeviceContextHandle ctx, void *dev_ptr); +/** Allocate page-locked host memory owned by the given device context. */ +int alloc_pinned_host_ctx(DeviceContextHandle ctx, size_t size, void **host_ptr); + +/** Release host memory returned by alloc_pinned_host_ctx. */ +int free_pinned_host_ctx(DeviceContextHandle ctx, void *host_ptr); + /** * Total device HBM (bytes) currently committed by this device context's * MemoryAllocator (user tensors + pooled arenas + Graph execution blocks + diff --git a/tests/ut/py/test_qwen3_pinned_inputs.py b/tests/ut/py/test_qwen3_pinned_inputs.py new file mode 100644 index 0000000000..1f47c5e2f8 --- /dev/null +++ b/tests/ut/py/test_qwen3_pinned_inputs.py @@ -0,0 +1,50 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Unit coverage for direct-to-final-storage Qwen input generation.""" + +import torch + +from simpler_setup.goldens import qwen3_14b_decode as qwen + + +class _RecordingAllocator: + def __init__(self): + self.tensors = [] + + def empty(self, shape, *, dtype=None): + tensor = torch.empty(shape, dtype=dtype) + self.tensors.append(tensor) + return tensor + + +def test_qwen_args_are_generated_in_allocator_storage(monkeypatch): + dimensions = { + "BATCH": 2, + "HIDDEN": 8, + "KV_HIDDEN": 4, + "HEAD_DIM": 4, + "HALF_DIM": 2, + "INTERMEDIATE": 6, + "MAX_SEQ": 8, + "BLOCK_SIZE": 2, + "MAX_BLOCKS_PER_SEQ": 4, + "CACHE_ROWS": 16, + } + for name, value in dimensions.items(): + monkeypatch.setattr(qwen, name, value) + + allocator = _RecordingAllocator() + args = qwen.generate_inputs(seed=7, seq_len=5, n_layers=2, allocator=allocator) + allocated_ptrs = {tensor.data_ptr() for tensor in allocator.tensors} + + assert args.tensor_names() == [*qwen.INPUT_NAMES, "out"] + assert all(spec.value.data_ptr() in allocated_ptrs for spec in args.specs) + assert args.wq.shape == (2 * dimensions["HIDDEN"], dimensions["HIDDEN"]) + assert torch.equal(args.wq[: dimensions["HIDDEN"]], args.wq[dimensions["HIDDEN"] :]) + assert torch.count_nonzero(args.out) == 0 diff --git a/tests/ut/py/test_scene_test_golden_hooks.py b/tests/ut/py/test_scene_test_golden_hooks.py index ca8d918b0c..c4ee6d464e 100644 --- a/tests/ut/py/test_scene_test_golden_hooks.py +++ b/tests/ut/py/test_scene_test_golden_hooks.py @@ -65,6 +65,24 @@ def test_default_case_computes_and_compares(): assert inst.compare_calls == [_OUTPUT_NAMES] +def test_l2_arg_generation_can_allocate_from_worker(): + class _WorkerAware(_Recorder): + def __init__(self): + super().__init__() + self.arg_worker = None + + def generate_args(self, params): + raise AssertionError("ordinary arg generation must not run") + + def generate_args_for_worker(self, worker, params): + self.arg_worker = worker + return TaskArgsBuilder(TensorArg("y", torch.zeros(4, dtype=torch.float32))) + + inst = _WorkerAware() + _drive(inst, {"name": "c"}) + assert inst.arg_worker is not None + + def test_per_case_skip_golden_bypasses_both_hooks(): inst = _Recorder() _drive(inst, {"name": "c", "skip_golden": True}) diff --git a/tests/ut/py/test_task_interface.py b/tests/ut/py/test_task_interface.py index 880b18ed88..5491b4e286 100644 --- a/tests/ut/py/test_task_interface.py +++ b/tests/ut/py/test_task_interface.py @@ -15,6 +15,7 @@ import struct import weakref from multiprocessing.shared_memory import SharedMemory +from typing import Any, cast import pytest import simpler.task_interface as task_interface_module @@ -200,6 +201,84 @@ def test_dtype_names(self, dtype, expected): class TestTorchInterop: + def test_pinned_torch_allocator_builds_tensor_over_worker_storage(self): + import torch # pyright: ignore[reportMissingImports] + from simpler.worker import PinnedHostBuffer + + from simpler_setup.torch_interop import PinnedTorchAllocator + + class FakeChipWorker: + def __init__(self): + self.blocks = {} + self.alloc_sizes = [] + self.freed = [] + + def alloc_pinned_host(self, size): + block = (ctypes.c_ubyte * size)() + ptr = ctypes.addressof(block) + self.blocks[ptr] = block + self.alloc_sizes.append(size) + return ptr + + def free_pinned_host(self, ptr): + self.freed.append(ptr) + self.blocks.pop(ptr) + + class FakeWorker: + level = 2 + + def __init__(self): + self.chip_worker = FakeChipWorker() + + def alloc_pinned_host(self, size): + return PinnedHostBuffer(cast(Any, self.chip_worker), size) + + worker = FakeWorker() + tensor = PinnedTorchAllocator(worker).empty((2, 3), dtype=torch.float32) + ptr = tensor.data_ptr() + + assert worker.chip_worker.alloc_sizes == [2 * 3 * 4] + assert ptr in worker.chip_worker.blocks + tensor.fill_(1.25) + assert tensor.tolist() == [[1.25, 1.25, 1.25], [1.25, 1.25, 1.25]] + + # torch keeps the ctypes exporter (and therefore its allocation token) + # alive after PinnedTorchAllocator.empty() drops its local handle. + gc.collect() + assert worker.chip_worker.freed == [] + del tensor + gc.collect() + assert worker.chip_worker.freed == [ptr] + + def test_pinned_torch_allocator_rejects_non_l2_worker(self): + from simpler_setup.torch_interop import PinnedTorchAllocator + + with pytest.raises(TypeError, match="level-2"): + PinnedTorchAllocator(type("Worker", (), {"level": 3})()) + + def test_pinned_host_cleanup_reports_free_failure(self, capsys): + from simpler.worker import PinnedHostBuffer + + class FailingChipWorker: + def __init__(self): + self.block = (ctypes.c_ubyte * 16)() + + def alloc_pinned_host(self, size): + assert size == len(self.block) + return ctypes.addressof(self.block) + + def free_pinned_host(self, ptr): + raise RuntimeError(f"injected failure for {ptr}") + + buffer = PinnedHostBuffer(cast(Any, FailingChipWorker()), 16) + ptr = buffer.base + del buffer + gc.collect() + + assert capsys.readouterr().err == ( + f"PinnedHostBuffer cleanup: free_pinned_host failed (continuing best-effort): injected failure for {ptr}\n" + ) + def test_torch_dtype_to_datatype(self): import torch # pyright: ignore[reportMissingImports]