[Feature] Expose simpler's committed NPU memory (rtMalloc pool is invisible to aclrtGetMemInfo)

[sunghajung6688]

Summary

Downstream runtimes (we run pypto-serving, a vLLM-style Ascend serving stack on top of simpler) size their paged KV cache the way vLLM does: after loading weights and warming up, they compute the non-KV peak as total_HBM − free_HBM via torch.npu.mem_get_info (which wraps aclrtGetMemInfo) and hand the rest of HBM to the KV cache.

That scheme is only correct if free_HBM reflects all device memory already committed — including the ring-heap arena that simpler allocates itself via rtMalloc (aclrtMalloc). It does not appear to. As a result total − free under-reports real committed memory, the KV cache is over-provisioned, and the device runs out of HBM at runtime. We need either a simpler API to query its committed device memory, or confirmation of how the arena is reserved so we can account for it.

How we size the KV cache (consumer side)

In pypto_serving/model/qwen/npu_runner.py, after the L3 worker is created, weights are uploaded, and a profile warm-up has run (so simpler's arena is already allocated):

free_bytes, total_bytes = torch.npu.mem_get_info(f"npu:{device_id}")
peak_non_kv = total_bytes - free_bytes                 # ← the assumption lives here
kv_budget   = int(total_bytes * utilization - peak_non_kv)   # utilization default 0.90
num_pages   = max(kv_budget // bytes_per_page, 1)

We deliberately order everything (worker creation, static tensors, decode scratch, warm-up) before this measurement so that peak_non_kv = total − free is meant to capture weights + simpler arena + compiled buffers + scratch. The KV cache is then allocated into kv_budget; on alloc OOM we halve the page count and retry.

The entire estimate hinges on one assumption: total − free is a true measure of committed HBM.

The problem

simpler's worker arena is its own rtMalloc (aclrtMalloc) allocation, sized on the order of PTO2_RING_HEAP × 4 + 128 MB (≈ 1.1 GB at the default PTO2_RING_HEAP = 256 MB). torch and simpler are two independent NPU allocators. torch.npu.mem_get_info (aclrtGetMemInfo) does not appear to subtract simpler's rtMalloc pool from free, so:

• free is over-reported → peak_non_kv is under-reported → kv_budget is too large → the KV cache is allocated into HBM that simpler has already reserved but torch still believes is free.

This produces two distinct failure shapes:

1. Alloc-time (partly mitigated). The torch KV allocation itself OOMs and our halve-and-retry shrinks the cache. This only catches the instantaneous torch-allocator failure.
2. Runtime (NOT mitigated). The KV allocation succeeds (torch thought the memory was free), then simpler uses / grows its arena during real serving and collides with the already-resident KV cache → a runtime OOM or device fault that the halve-retry cannot catch.

A second-order uncertainty makes this worse: is the arena fully reserved at worker init, or grown lazily on demand? If it is lazy, the free we measure after warm-up reflects only a fraction of the arena's eventual footprint, so the over-provisioning is larger and the runtime collision in (2) is more likely.

How to reproduce / confirm

Run this probe on the device and report the free₁ − free₂ delta:

import torch, simpler   # simpler worker / allocator handle

free1, total = torch.npu.mem_get_info("npu:0")

# Trigger a known-size simpler rtMalloc allocation (e.g. init a worker,
# or allocate one arena slab of S bytes).
S = ...                      # bytes simpler is asked to commit
...                          # simpler allocation here

free2, _ = torch.npu.mem_get_info("npu:0")
print("delta:", free1 - free2, "expected:", S)

• If free₁ − free₂ ≈ S → mem_get_info sees simpler's rtMalloc pool (no problem from our side; we'd look elsewhere).
• If free₁ − free₂ ≈ 0 → simpler's pool is invisible to aclrtGetMemInfo and this issue is the root cause.

We can run this and attach numbers, but the allocation internals are simpler's, so we'd like the simpler side to (a) confirm the expected behavior and (b) tell us whether rtMalloc here is plain aclrtMalloc or the virtual-memory path (aclrtReserveMemAddress + aclrtMallocPhysical), which aclrtGetMemInfo is known not to count.

What we need from simpler

Any one of the following unblocks us; the first is the clean fix:

1. An API to query simpler's committed device memory — the sum of rtMalloc-backed pools / the resident arena size — so consumers can subtract it explicitly: peak_non_kv = (total − free) + simpler_committed_bytes.
2. Confirmation of whether aclrtGetMemInfo is expected to reflect simpler's allocations, and which CANN allocation API the ring heap uses (aclrtMalloc vs. the virtual-memory reserve/map path).
3. Arena lifecycle guidance: is PTO2_RING_HEAP × 4 + 128 MB reserved up-front at worker init, or grown lazily? This determines whether a warm-up-time mem_get_info measurement is ever safe to rely on.

Environment

• Device: Ascend 910B/C
• CANN: <fill in>
• torch_npu: <fill in>
• simpler / pypto-lib: <commit>
• PTO2_RING_HEAP: default 256 MB (arena ≈ 1.1 GB)

Workaround we can adopt in the meantime

Stop relying on mem_get_info for the simpler portion and subtract the arena explicitly:

arena = int(os.environ.get("PTO2_RING_HEAP", 256 * 1024 * 1024)) * 4 + 128 * 1024 * 1024
kv_budget = int(total_bytes * utilization - peak_non_kv) - arena

This works only if the arena is reserved up-front; it also duplicates simpler's internal sizing formula and will silently drift whenever simpler changes how it sizes the arena — which is exactly why a real query API (item 1 above) would be far more robust.

Motivation / Use Case

"aclrtGetMemInfo / torch.npu.mem_get_info does not reflect simpler's rtMalloc ring-heap arena; please expose simpler's committed device memory".

Proposed API / Behavior

No response

Alternatives Considered

No response

Additional Context

No response