Summary
Expose device-wide HBM information through the PyPTO DistributedWorker API:
free_bytes, total_bytes = worker.device_memory_info(worker_id=0)
The method should forward the query to the underlying Simpler Worker and return free and total device memory in bytes.
Depends on: hw-native-sys/simpler#1878
Motivation / Use Case
Serving integrations need device-wide free memory to size KV caches dynamically. They currently call torch.npu.mem_get_info directly, which requires a large torch-npu dependency for a single query.
DistributedWorker already exposes worker-scoped runtime information such as committed_device_memory(worker_id). Providing the device memory query through the same facade allows consumers to:
- Query the logical worker that owns the target device.
- Avoid managing ACL device contexts in the serving process.
- Avoid depending on Simpler implementation details.
- Remove
torch-npu when it is only used for memory inspection.
The new method must remain semantically distinct from committed_device_memory:
committed_device_memory is the amount committed by the Simpler allocator.
device_memory_info is a device-wide snapshot of free and total HBM.
Proposed API / Behavior
Add the following method to DistributedWorker:
def device_memory_info(self, worker_id: int = 0) -> tuple[int, int]:
"""Return free and total device memory in bytes."""
Expected behavior:
- Require the
DistributedWorker to be open.
- Forward the logical
worker_id unchanged to the underlying Simpler Worker.
- Normalize both returned values to Python integers.
- Preserve errors from the underlying Worker, including unsupported simulator backends.
- Do not silently return
(0, 0) when the query fails.
Example:
with DistributedWorker(...) as worker:
free_bytes, total_bytes = worker.device_memory_info(0)
Alternatives Considered
-
Let serving access the private underlying Simpler Worker.
This breaks the DistributedWorker abstraction and couples consumers to its internal structure.
-
Query ACL directly from serving.
This duplicates device-context and lifecycle handling outside the runtime layer.
-
Continue using torch.npu.mem_get_info.
This retains torch-npu as a dependency even when no other torch-npu functionality is required.
Additional Context
Suggested validation:
- Verify that the logical Worker ID is forwarded unchanged.
- Verify that the returned values are exposed as Python integers.
- Verify that calls after
DistributedWorker.close() fail consistently with other methods.
- Preserve and test propagation of unsupported-backend and runtime errors.
Summary
Expose device-wide HBM information through the PyPTO
DistributedWorkerAPI:The method should forward the query to the underlying Simpler Worker and return free and total device memory in bytes.
Depends on: hw-native-sys/simpler#1878
Motivation / Use Case
Serving integrations need device-wide free memory to size KV caches dynamically. They currently call
torch.npu.mem_get_infodirectly, which requires a largetorch-npudependency for a single query.DistributedWorkeralready exposes worker-scoped runtime information such ascommitted_device_memory(worker_id). Providing the device memory query through the same facade allows consumers to:torch-npuwhen it is only used for memory inspection.The new method must remain semantically distinct from
committed_device_memory:committed_device_memoryis the amount committed by the Simpler allocator.device_memory_infois a device-wide snapshot of free and total HBM.Proposed API / Behavior
Add the following method to
DistributedWorker:Expected behavior:
DistributedWorkerto be open.worker_idunchanged to the underlying Simpler Worker.(0, 0)when the query fails.Example:
Alternatives Considered
Let serving access the private underlying Simpler Worker.
This breaks the
DistributedWorkerabstraction and couples consumers to its internal structure.Query ACL directly from serving.
This duplicates device-context and lifecycle handling outside the runtime layer.
Continue using
torch.npu.mem_get_info.This retains
torch-npuas a dependency even when no other torch-npu functionality is required.Additional Context
Suggested validation:
DistributedWorker.close()fail consistently with other methods.