vllm.models.deepseek_v41.nvidia.ops.engram_fabric_table

Engram tables in cuMem, shared between processes by FABRIC handle.

Host tables: one copy in host-NUMA memory per node-local DP group, which GPUs map with 2M pages instead of /dev/shm's 64K pages, so row lookups miss the TLB less. Peer tables: rows split over the HBM of several engines' GPUs.

Classes:

• ExportedHostTable –

  What another rank needs to map a host table.

• PeerTableMappings –

  This process's mappings of every slice of a peer table, unmapped once

Functions:

• fabric_host_memory_supported –

  Whether host-NUMA cuMem allocations can export FABRIC handles here.

• share_host_table –

  Map one table on every rank of a node-local group, or None on all of them.

• share_peer_table –

  Allocate this GPU's slice of a table in its HBM and map every GPU's slice.

ExportedHostTable

Bases: NamedTuple

What another rank needs to map a host table.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

class ExportedHostTable(NamedTuple):
    """What another rank needs to map a host table."""

    numa_node: int
    padded_bytes: int
    fabric_handles: list[bytes]

PeerTableMappings

This process's mappings of every slice of a peer table, unmapped once neither the table nor any tensor over its slice holds them. A slice's memory is freed when every GPU has unmapped it.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

class PeerTableMappings:
    """This process's mappings of every slice of a peer table, unmapped once
    neither the table nor any tensor over its slice holds them. A slice's
    memory is freed when every GPU has unmapped it."""

    def __init__(self, size: int) -> None:
        self.addresses: list[int] = []
        finalizer = weakref.finalize(self, _unmap_slices, self.addresses, size)
        finalizer.atexit = False  # type: ignore[misc]

_call_driver(fn, *args)

Call a CUDA driver API and return its result, raising on failure.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

def _call_driver(fn, *args):
    """Call a CUDA driver API and return its result, raising on failure."""
    error, *result = fn(*args)
    if error != cu.CUresult.CUDA_SUCCESS:
        raise RuntimeError(f"{fn.__name__} failed: {error.name}")
    return result[0] if result else None

_create_table(num_bytes)

Allocate and export a host table on this GPU's NUMA node.

Stops once the node would drop below its reserve. The kernel compacts to find 2M pages and fails the chunk, rather than OOM-killing, if it cannot.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

def _create_table(num_bytes: int) -> tuple[list, ExportedHostTable]:
    """Allocate and export a host table on this GPU's NUMA node.

    Stops once the node would drop below its reserve. The kernel compacts to
    find 2M pages and fails the chunk, rather than OOM-killing, if it cannot.
    """
    numa_node = _gpu_numa_node()
    prop = _host_allocation_prop(numa_node)
    padded_bytes = round_up(num_bytes, _allocation_granularity(prop))
    reserve = get_memory_node_info(numa_node).total_memory * NUMA_RESERVE_FRACTION
    handles: list = []
    try:
        for offset, length in _chunk_ranges(padded_bytes):
            handles.append(_call_driver(cu.cuMemCreate, length, prop, 0))
            if get_memory_node_info(numa_node).available_memory < reserve:
                raise MemoryError(
                    f"NUMA node {numa_node} fell below its "
                    f"{reserve / 2**30:.1f} GiB reserve after "
                    f"{(offset + length) / 2**30:.1f} of "
                    f"{padded_bytes / 2**30:.1f} GiB"
                )
        exported = [_export_fabric_handle(handle) for handle in handles]
    except BaseException:
        _release(handles)
        raise
    return handles, ExportedHostTable(numa_node, padded_bytes, exported)

_gpu_view(address, num_bytes, owner=None)

A uint8 tensor over mapped memory; torch retains owner with it.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

def _gpu_view(address: int, num_bytes: int, owner: object = None) -> torch.Tensor:
    """A uint8 tensor over mapped memory; torch retains `owner` with it."""
    # Name the device: an index-less "cuda" makes torch copy the memory.
    return torch.as_tensor(
        SimpleNamespace(
            __cuda_array_interface__={
                "shape": (num_bytes,),
                "typestr": "|u1",
                "data": (address, False),
                "version": 3,
            },
            owner=owner,
        ),
        device=torch.device("cuda", torch.accelerator.current_device_index()),
    )

_map_table(table, num_bytes, handles)

Map a host table, importing its handles unless given, and release them.

Returns CPU and GPU uint8 views; either keeps the mapping alive.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

def _map_table(
    table: ExportedHostTable, num_bytes: int, handles: list
) -> tuple[torch.Tensor, torch.Tensor]:
    """Map a host table, importing its handles unless given, and release them.

    Returns CPU and GPU uint8 views; either keeps the mapping alive.
    """
    device = torch.accelerator.current_device_index()
    address = None
    mapped: list[tuple[int, int]] = []
    try:
        if not handles:
            for exported in table.fabric_handles:
                handles.append(
                    _call_driver(
                        cu.cuMemImportFromShareableHandle, exported, FABRIC_HANDLE
                    )
                )
        address = int(_call_driver(cu.cuMemAddressReserve, table.padded_bytes, 0, 0, 0))
        for (offset, length), handle in zip(_chunk_ranges(table.padded_bytes), handles):
            _call_driver(cu.cuMemMap, address + offset, length, 0, handle, 0)
            mapped.append((offset, length))
        location = cu.CUmemLocationType
        access = [
            _read_write_access(location.CU_MEM_LOCATION_TYPE_DEVICE, device),
            _read_write_access(
                location.CU_MEM_LOCATION_TYPE_HOST_NUMA, table.numa_node
            ),
        ]
        _call_driver(
            cu.cuMemSetAccess, address, table.padded_bytes, access, len(access)
        )
    except BaseException:
        if address is not None:
            _unmap(address, table.padded_bytes, mapped)
        raise
    finally:
        # The mappings keep the allocations alive.
        _release(handles)

    host_buffer = (ctypes.c_uint8 * num_bytes).from_address(address)
    finalizer = weakref.finalize(
        host_buffer, _unmap, address, table.padded_bytes, mapped
    )
    finalizer.atexit = False  # type: ignore[misc]
    gpu_view = _gpu_view(address, num_bytes, owner=host_buffer)
    return torch.frombuffer(host_buffer, dtype=torch.uint8), gpu_view

fabric_host_memory_supported() cached

Whether host-NUMA cuMem allocations can export FABRIC handles here.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

@cache
def fabric_host_memory_supported() -> bool:
    """Whether host-NUMA cuMem allocations can export FABRIC handles here."""
    try:
        prop = _host_allocation_prop(_gpu_numa_node())
        handle = _call_driver(cu.cuMemCreate, _allocation_granularity(prop), prop, 0)
        try:
            _export_fabric_handle(handle)
        finally:
            cu.cuMemRelease(handle)
    except Exception as exc:
        logger.info_once(
            "Host-NUMA cuMem FABRIC handles are unavailable (%s); Engram DP "
            "replicas share tables through /dev/shm.",
            exc,
        )
        return False
    return True

share_host_table(num_bytes, group)

Map one table on every rank of a node-local group, or None on all of them.

Rank 0 creates the table and broadcasts its handles; the others import them. Returns CPU and GPU uint8 views of the table.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

def share_host_table(
    num_bytes: int, group: GroupCoordinator
) -> tuple[torch.Tensor, torch.Tensor] | None:
    """Map one table on every rank of a node-local group, or None on all of them.

    Rank 0 creates the table and broadcasts its handles; the others import them.
    Returns CPU and GPU uint8 views of the table.
    """
    handles: list = []
    table = error = None
    if group.rank_in_group == 0 and fabric_host_memory_supported():
        try:
            handles, table = _create_table(num_bytes)
        except Exception as exc:
            error = describe_error(exc)
    table, error = group.broadcast_object((table, error))
    if table is not None:
        views = None
        try:
            views = _map_table(table, num_bytes, handles)
        except Exception as exc:
            error = describe_error(exc)
        if not (error := gather_errors(group, error)):
            return views
    if error and group.rank_in_group == 0:
        logger.warning(
            "Engram host table cannot use host-NUMA cuMem (%s); falling back "
            "to /dev/shm, whose 64K GPU pages make lookups slower.",
            error,
        )
    return None

share_peer_table(num_bytes, group)

Allocate this GPU's slice of a table in its HBM and map every GPU's slice.

Every slice holds num_bytes. Returns this GPU's slice and the mappings, whose addresses are in rank order.

Source code in vllm/models/deepseek_v41/nvidia/ops/engram_fabric_table.py

def share_peer_table(
    num_bytes: int, group: EngramPeerGroup
) -> tuple[torch.Tensor, PeerTableMappings]:
    """Allocate this GPU's slice of a table in its HBM and map every GPU's slice.

    Every slice holds `num_bytes`. Returns this GPU's slice and the mappings,
    whose addresses are in rank order.
    """
    device = torch.accelerator.current_device_index()
    prop = cu.CUmemAllocationProp()
    prop.type = cu.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
    prop.location.type = cu.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
    prop.location.id = device
    prop.requestedHandleTypes = FABRIC_HANDLE
    size = round_up(num_bytes, _allocation_granularity(prop))
    own = _call_driver(cu.cuMemCreate, size, prop, 0)
    exported = group.all_gather(_export_fabric_handle(own))
    access = _read_write_access(
        cu.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE, device
    )
    mappings = PeerTableMappings(size)
    for rank, peer in enumerate(exported):
        handle = (
            own
            if rank == group.rank_in_group
            else _call_driver(cu.cuMemImportFromShareableHandle, peer, FABRIC_HANDLE)
        )
        address = int(_call_driver(cu.cuMemAddressReserve, size, 0, 0, 0))
        mappings.addresses.append(address)
        _call_driver(cu.cuMemMap, address, size, 0, handle, 0)
        _call_driver(cu.cuMemSetAccess, address, size, [access], 1)
        # The mapping keeps the slice alive.
        _release([handle])
    address = mappings.addresses[group.rank_in_group]
    return _gpu_view(address, num_bytes, owner=mappings), mappings