Add: allocate Qwen args in pinned host memory

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)
 

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(

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))

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).
 

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",

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

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)

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:

src/common/platform/onboard/host/c_api_shared.cpp

@@ -32,6 +32,7 @@
 #include "task_args.h"
 #include "native_run_context.h"
 
+#include <acl/acl_rt.h>
 #include <dlfcn.h>
 #include <cstdlib>
 #include <cstdio>
@@ -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 {

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 {