remap_registered_memory

Applicability

Note: For Atlas A2 training product/Atlas A2 inference product, only the Atlas 800I A2 inference server and A200I A2 Box heterogeneous subrack are supported.

Function Description

During LLM inference, if a memory UCE occurs (that is, ACL_ERROR_RT_DEVICE_MEM_ERROR is returned), the upper-layer framework needs to first determine whether the faulty memory is a KV cache memory. If it is not a KV cache memory, refer to the torch_npu.npu.restart_device API description in Ascend Extension for PyTorch Custom API Reference to obtain and fix the incorrect virtual address of the memory UCE. If it is a KV cache memory, additionally call this API to restore the KV cache memory registered with the NIC.

Note: This API is reserved, and the function is not supported currently.

Prototype

1

remap_registered_memory(mem_infos: Union[MemInfo, list[MemInfo]]) -> None

Parameters

Example

1
2
3
4
5
6
7
8
9

from llm_datadist import *
role = LLMRole.PROMPT  # Role where the fault occurs.
cluster_id = 1 # ID of the LLM-DataDist instance where the fault occurs.
datadist = LLMDataDist(role, cluster_id)
cache_manager = datadist.cache_manager
addr = 12345678 # Address where the fault occurs.
size = 1024  # Size of the faulty memory region.
mem_info = MemInfo(Memtype.MEM_TYPE_DEVICE, addr, size)
cache_manager.remap_registered_memory(mem_info)

Returns

In normal cases, no value is returned.

If the input data type is incorrect, the TypeError or ValueError exception is reported.

Restrictions

Currently, only the repair of device memory is supported.