create_npu_tensors

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

函数功能

通过一串Device地址创建PyTorch在NPU上的tensors。

主要使用场景是用于创建大模型中的kv cache tensors，所有kv cache的shape和dtype都一致。

LLM-DataDist是大模型分布式集群和数据加速组件，提供了集群KV数据管理能力，以支持全量图和增量图分离部署。

函数原型

def create_npu_tensors(shape: List[int], dtype: torch.dtype, addresses: List[int]) -> List[torch.Tensor]

参数说明

返回值说明

正常情况下，返回torch.Tensor的List。异常情况会抛出异常，创建tensor会失败。

约束说明

当前该接口仅限大模型分离部署场景下使用。

调用示例

1
2
3
4
5

import torch, torch_npu, torchair
kv_tensor_addrs = [1, 2, 3, 4]           # 需要是device返回出来的addr
shape = [4, 1024, 128]
dtype = torch.float16
kv_tensors = torchair.llm_datadist.create_npu_tensors(shape, dtype, kv_tensor_addrs)