Skip to content

[Feature] Expose device memory information on DistributedWorker #2413

Description

@high-cloud

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

  1. Let serving access the private underlying Simpler Worker.

    This breaks the DistributedWorker abstraction and couples consumers to its internal structure.

  2. Query ACL directly from serving.

    This duplicates device-context and lifecycle handling outside the runtime layer.

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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions