create_npu_tensors

功能说明

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

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

• LLM-DataDist是大模型分布式集群和数据加速组件，提供了集群KV数据管理能力，以支持全量图和增量图分离部署。
• 该接口目前适用如下产品：
  • Atlas 800I A2 推理产品/A200I A2 Box 异构组件
  • Atlas A3 训练系列产品/Atlas A3 推理系列产品

函数原型

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)