remap_registered_memory

Applicability

Note: For the Atlas A2 training products/Atlas A2 inference products, only the Atlas 800I A2 inference server, Atlas 300I A2 inference card, and A200I A2 Box heterogeneous components 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 erroneous 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 occurred
cluster_id = 1 # ID of the LLM-DataDist instance where the fault occurred
datadist = LLMDataDist(role, cluster_id)
cache_manager = datadist.cache_manager
addr = 12345678 # Address where the fault occurred
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, a TypeError or ValueError is thrown.

Restrictions

Currently, only the repair of device memory is supported.