HSTU 生成式推荐模型在 CANN 平台的 NPU 推理适配与优化实战指南【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer本篇技术指南以 CANN 开源推理样例仓中的 HSTU 模型适配实现为核心系统介绍生成式推荐GR模型 HSTU 在 Atlas A2 系列昇腾 NPU 上的推理全流程从 Docker 镜像与 CANN 环境准备、RecSDK 算子编译安装到 kuairand-1k 数据集的预处理、单卡与多卡推理命令再到 KV Cache 多级缓存、Paged Attention、融合算子与 ACLGraph 等关键优化点的源码级原理。读完本文你将能够在本仓库 models/hstu 目录下完整复现 HSTU 的 NPU 推理并理解每一项 NPU 适配优化背后的设计意图与实现位置。HSTU 模型与生成式推荐背景Meta 论文《Actions Speak Louder Than Words》提出生成式推荐Generative Recommendation, GR范式将推荐问题建模为序列生成任务并引入 HSTUHierarchical Sequential Transduction Unit架构以应对推荐场景中高基数high-cardinality、非平稳non-stationary的数据流同时统一支持检索与排序两大任务。基于 HSTU 构建的超大规模模型1.5 万亿参数在线 A/B 指标提升 12.4%已在 Meta 多个平台部署。本仓库 models/hstu/README.md 中给出的示例即聚焦于 HSTU 在 CANN 平台的 NPU 推理实现其核心工作是在 NVIDIA recsys-examples 的 HSTU 推理代码基础上完成设备层适配CUDA → NPU、KV Cache 管理替换以及高性能融合算子接入。支持的产品型号与运行环境本样例支持Atlas A2 系列产品的单卡、多卡推理运行环境基于RecSDK-Torch镜像与CANN 8.5.0软件包。从源码结构看整个样例的入口脚本、模型实现与算子编译脚本均位于仓库的 models/hstu 目录下包括推理入口命令行解析、数据加载、模型构建与推理循环模型主体InferenceRankingGR端到端推理模型分页 HSTU 推理层单层注意力计算与 KV Cache 写入KV Cache 管理 与 KV Cache 实现算子一键编译脚本 与 适配补丁。环境准备镜像准备本样例基于 RecSDK-Torch 环境运行官方提供 Docker 镜像便于快速完成环境配置。流程如下从昇腾 AscendHub 社区下载rec_sdk-torch镜像对应昇腾开源开发套件 RecSDK-Torch使用如下命令拉起容器示例映射了 8 张 NPU 卡可依据实际硬件调整--device参数docker run -u root -itd --name rec_gr --ulimit nproc65535:65535 --ipchost \ --device/dev/davinci0 --device/dev/davinci1 \ --device/dev/davinci2 --device/dev/davinci3 \ --device/dev/davinci4 --device/dev/davinci5 \ --device/dev/davinci6 --device/dev/davinci7 \ --device/dev/davinci_manager --device/dev/devmm_svm \ --device/dev/hisi_hdc \ -v /home/:/home \ -v /data:/data \ -v /etc/localtime:/etc/localtime \ -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ -v /etc/ascend_install.info:/etc/ascend_install.info -v /var/log/npu/:/usr/slog \ -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /sys/fs/cgroup:/sys/fs/cgroup:ro \ -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/sbin:/usr/local/sbin \ -v /etc/hccn.conf:/etc/hccn.conf -v /root/.pip:/root/.pip -v /etc/hosts:/etc/hosts \ -v /usr/bin/hostname:/usr/bin/hostname \ --nethost \ --shm-size128g \ --privileged \ $REPOSITORY:TAG \ # 通过docker images命令查看REPOSITORY和TAG例如 swr.cn-south-1.myhuaweicloud.com/ascendhub/rec_sdk-torch:openeuler2203-arm /bin/bash其中-v /data:/data用于映射代码文件、权重文件和数据集文件。容器拉起后使用docker exec -it rec_gr bash进入容器内部并设置如下环境变量export PATH/usr/local/python3.11.0/bin:$PATH export LD_LIBRARY_PATH/usr/local/python3.11.0/lib:$LD_LIBRARY_PATH安装 CANN 8.5.0 软件包本样例的编译执行依赖 CANN 开发套件包与 CANN 二进制算子包支持的软件版本为CANN 8.5.0从昇腾社区软件包下载页面获取Ascend-cann-toolkit_${version}_linux-${arch}.run与Ascend-cann-910b-ops_${version}_linux-${arch}.run两个软件包并参考昇腾 CANN 安装文档进行安装${version}表示 CANN 包版本号如8.5.0${arch}表示 CPU 架构如aarch64、x86_64。安装完成后设置 CANN 环境变量source /usr/local/Ascend/cann-8.5.0/set_env.sh代码准备本仓库的 HSTU 示例依赖 recsys-examples 开源仓的代码按以下三步完成代码准备克隆 cann-recipes 推理仓mkdir -p /data/code cd /data/code git clone https://gitcode.com/cann/cann-recipes-infer.git克隆 recsys-examples 仓v25.11 分支获取所需文件并使用非覆盖模式替换文件-n表示不覆盖已存在的同名文件避免破坏本仓库已有的 NPU 适配代码git clone --branch v25.11 --depth 1 https://github.com/NVIDIA/recsys-examples.git cp -an recsys-examples/examples/hstu/* cann-recipes-infer/models/hstu/应用 patch 完成 NPU 适配cd cann-recipes-infer/models/hstu git apply hstu_cann.patch从补丁 hstu_cann.patch 的内容看其主要完成将torch.cuda.current_device()全部替换为torch.npu.current_device()移除对 TensorRT-LLM 绑定库的依赖改为本地KVCacheManagerImpl实现重写 Host 侧 KV 存储管理的按页读写逻辑去除训练相关的 ranking/retrieval 模型工厂代码等。相关依赖安装Docker 容器内需要安装 nnal 加速库与 RecSDK 算子组件具体如下。1. nnal 加速库安装本样例 KV Cache 管理功能依赖 NPU 算子_npu_reshape_and_cache使用该算子需要安装 nnal 加速库。从昇腾社区 CANN 8.5.0 下载页面获取Ascend-cann-nnal_8.5.0_linux-${arch}.run包并在 docker 环境中安装bash Ascend-cann-nnal_8.5.0_linux-${arch}.run --install安装完成后设置环境变量source /usr/local/Ascend/nnal/atb/set_env.sh2. RecSDK 算子库安装从昇腾 RecSDK 仓库克隆适配代码cd cann-recipes-infer/models/hstu/ git clone https://gitcode.com/Ascend/RecSDK.git部分算子依赖外部组件需要提前准备下载json-3.9.1.tar.gznlohmann/json 的 v3.9.1 版本并重命名为v3.9.1.tar.gz放置于RecSDK/cust_op/ascendc_op/build/scripts/onnx_plugin目录in_linear_silu算子依赖 CATLASS 源码验证过的版本为 catlass v1.3.0下载压缩包后解压。concat_nd_jagged自定义算子适配concat_2d_jagged算子原生只支持 2 个Jagged Tensor进行 concat。为支持多个Jagged Tensor的 concat 操作需要对 Torch 算子适配层代码进行修改代码路径为RecSDK/cust_op/framework/torch_plugin/torch_library/concat_2d_jagged/concat_jagged_tensor.cpp修改方式如下// 添加函数实现concat_nd_jagged_npu at::Tensor concat_nd_jagged_npu( const int64_t maxSeqlen, at::TensorList values, at::TensorList offsets, const bool isReplace false, const int64_t nPrefixFromRight 0) { TORCH_CHECK(values.size() 2, values list must have at least 2 tensors); TORCH_CHECK(values.size() offsets.size(), values list and offsets list must have same length); TORCH_CHECK(values[0].dim() 2, values[0] must be a 2D tensor); const auto v_dtype values[0].dtype(); const auto v_device values[0].device(); const int64_t resultCols values[0].size(1); TORCH_CHECK(offsets[0].dim() 1, offsets[0] must be a 1D tensor); const auto o_dtype offsets[0].dtype(); int64_t offsetlen offsets[0].size(0); int64_t resultRows 0; for (int64_t i 0; i (int64_t)values.size(); i){ TORCH_CHECK(values[i].dim() 2, values[, i, ] must be 2D); TORCH_CHECK(values[i].size(1) resultCols, values must have same hidden dim); TORCH_CHECK(values[i].dtype() v_dtype, values must have same dtype); TORCH_CHECK(values[i].device() v_device, values must be on the same device); TORCH_CHECK(offsets[i].dim() 1, offsets[, i, ] must be 1D); TORCH_CHECK(offsets[i].size(0) offsetlen, all offsets must have the same length); TORCH_CHECK(offsets[i].dtype() o_dtype, all offsets must have the same dtype); resultRows values[i].size(0); } int64_t jtNum (int64_t)values.size(); std::vectorat::Tensor offsets_vec; offsets_vec.reverse(jtNum); for (auto t : offsets) offsets_vec.emplace_back(t); at::Tensor offsets_cat at::cat(offsets_vec, 0); at::Tensor offsets_int64 offsets_cat.to(at::kLong); std::vectorint64_t offsetVector tensor_to_vectorint64_t(offsets_int64); auto offsetArray vec_to_intarray(offsetVector); auto result at::empty({resultRows, resultCols}, values[0].options()); EXEC_NPU_CMD(aclnnConcatJaggedTensor, values, offsetArray, offsetlen, jtNum, result); return result; } // 在npu命名空间里注册concat_2d_jagged。找到TORCH_LIBRARY_FRAGMENT(mxrec, m)添加下面m.def()内容 TORCH_LIBRARY_FRAGMENT(mxrec, m) { m.def(concat_nd_jagged_npu(SymInt maxSeqlen, Tensor[] values, Tensor[] offsets, bool isReplaceFalse, SymInt nPrefixFromRight0) - Tensor); } // 添加注册在TORCH_LIBRARY_IMPL(mxrec, PrivateUse1, m)中添加m.impl()内容 TORCH_LIBRARY_IMPL(mxrec, PrivateUse1, m) { m.impl(concat_nd_jagged_npu, concat_nd_jagged_npu) }设置 CATLASS 环境变量后执行 build_install_ops.sh 一键完成所需算子的编译与适配层安装export CATLASS_HOMEcatlass_home chmod x build_install_ops.sh bash build_install_ops.sh A2 ./RecSDK从脚本实现看build_install_ops.sh会依次完成四个步骤调用RecSDK/cust_op/ascendc_op/build/build_ai_core_op.sh构建 AI Core 算子解压Ascend-recsdk-npu-ops-A2-linux-aarch64.tar.gz产物安装mxrec_opp_hstu_dense_forward、mxrec_opp_concat_2d_jagged、mxrec_opp_hstu_dense_backward、mxrec_opp_split_embedding_codegen_forward_unweighted等算子运行包并额外安装in_linear_silu位于ai_core_op/in_linear_silu/v220目录最后通过torch_library/common/build_ops.sh构建 Torch 适配库。脚本也支持只传架构参数如bash build_install_ops.sh A2时自动向上查找 RecSDK 根目录。3. 其他库安装最后安装与导入相关 Python 库pip uninstall torchrec pip install torchrec1.1.0 pip install rich einops export LIB_FBGEMM_NPU_API_SO_PATH/usr/local/python3.11.0/lib/python3.11/site-packages/libfbgemm_npu_api.so其中LIB_FBGEMM_NPU_API_SO_PATH指向 fbgemm 的 NPU API 动态库推理入口脚本 与 HSTU 处理器 中均通过torch.ops.load_library(lib_fbgemm_npu_api_so_path)加载它用于asynchronous_complete_cumsum等稀疏计算原语。模型推理下载并处理数据集使用 preprocessor.py 下载并处理数据集示例数据集为 kuairand-1k即 Kuaishou 推荐数据的 1k 子集python3 ./preprocessor.py --dataset_name kuairand-1k --inference单机推理推理入口脚本 inference_gr_ranking.py 支持eval与simulate两种运行模式并支持 profiling 采集# eval 模式 python3 ./inference/inference_gr_ranking.py --gin_config_file ./inference/configs/kuairand_1k_inference_ranking.gin --mode eval --enable_fused_ops all # eval 模式 profiling 采集 python3 ./inference/inference_gr_ranking.py --gin_config_file ./inference/configs/kuairand_1k_inference_ranking.gin --mode eval --enable_fused_ops all --enable_profiler # simulate 模式 python3 ./inference/inference_gr_ranking.py --gin_config_file ./inference/configs/kuairand_1k_inference_ranking.gin --mode simulate --enable_fused_ops all各命令行参数的语义与 入口脚本 的 argparse 定义一致参数取值说明--gin_config_filegin 文件路径必填网络结构、数据集与任务超参数配置--modeeval/simulateeval 走离线数据集评测simulate 模拟用户时间线增量推理并计算 AUC--enable_fused_opsall/concat_nd_jagged/in_linear_silu使能融合算子可单独或全部开启--enable_profiler开关使用torch_npu.profiler采集 NPU 侧 profile仅 eval 模式生效--disable_auc开关关闭 AUC 指标计算仅 simulate 模式生效--disable_context开关禁用上下文特征仅 simulate 模式生效--embed_tp_size整数默认 1Embedding 表张量并行切分大小用于多卡推理--enable_fused_ops all会在脚本内设置builtins.ENABLED_FUSED_OPS {concat_nd_jagged, in_linear_silu}后续模型代码通过fused_enabled(name)判断是否走 NPU 融合算子分支。分布式推理分布式推理支持单机多卡示例为两卡通过torchrun启动并使用 HCCL 作为通信后端ASCEND_RT_VISIBLE_DEVICES0,1 torchrun --rdzv-backendc10d --rdzv-endpointlocalhost:6000 --nnodes1 --nproc-per-node2 ./inference/inference_gr_ranking.py --gin_config_file ./inference/configs/kuairand_1k_inference_ranking.gin --embed_tp_size 2 --mode eval --enable_fused_ops all当--embed_tp_size 1时入口脚本 会通过create_tp_process_group创建 HCCL 进程组并将组信息传入模型InferenceEmbedding 中对 vocab 规模不小于阈值默认 1000的 Embedding 表使用VocabParallelEmbedding按 TP 维度切分从而在 Embedding Table 占用显存较大的场景下支持更大的 dense 参数模型推理。网络与缓存关键配置推理使用的 gin 配置文件 定义如下核心参数DatasetArgs.dataset_name kuairand-1k DatasetArgs.max_sequence_length 4096 DatasetArgs.max_num_candidates 512 DatasetArgs.shuffle False NetworkArgs.num_layers 4 NetworkArgs.num_attention_heads 4 NetworkArgs.hidden_size 512 NetworkArgs.kv_channels 128 RankingArgs.prediction_head_arch [512, 8] RankingArgs.prediction_head_bias True RankingArgs.num_tasks 8KV Cache 相关配置在 inference_gr_ranking.py 的get_inference_hstu_model中硬编码构建包括blocks_in_primary_pool5120主池页数、page_size32每页 token 数、offload_chunksize1024Host offload 分块大小、max_batch_size1、max_seq_lenceil(total_max_seqlen/32)*32。位置编码配置为num_position_buckets8192、num_time_buckets2048、use_time_encodingFalse。kuairand-1k 的 Embedding 表配置中user_idvocab 为 1000video_idvocab 为 10,000,000其余特征表如user_active_degree8、action_weights233均为小表Embedding 维度统一等于hidden_size512。NPU 适配与优化点深度解析基于 recsys-examples 库本样例针对 NPU 做了两类基础适配并在此之上使能了六项关键优化。以下结合源码逐一展开。基础适配一TensorRT-LLM 库 KV Cache 管理 impl 方法的 NPU 适配替换原始 recsys-examples 的 HSTU 推理直接依赖 TensorRT-LLM 的KVCacheManager绑定。NPU 侧无法复用该 CUDA 实现因此本仓库在 kv_cache_manager_impl.py 中实现了纯 PyTorch 版KVCacheManagerImpl并在 hstu_cann.patch 中删除了对tensorrt_llm.bindings.internal.batch_manager.KVCacheManager的导入。从实现看该管理器具备以下能力主池primary_pool形状为[num_layers, 2, total_blocks, block_size]的 KV 缓存张量其中block_size tokens_per_block * num_heads * size_per_head2对应 K/V 两份LRU 淘汰维护_seq_lru_list_lru_touch将最近访问的序列移到队首空间不足时从队尾挑选 victim 序列淘汰按序列连续追加add_sequence_with_eviction校验start_pos必须与上次缓存末尾连续当前仅支持连续追加自淘汰self-evict当序列所需块数超过max_blocks_per_seq时优先淘汰自身前缀块_evict_blocks_prefix并将seq_cache_start_pos右移以维持缓存语义预留块reserved_blocks_in_primary_pool预留的页被 GPU 端管理器 用作 Host onload 的目标缓冲。基础适配二使用 torch 小算子替换 Triton 与 CUDA 算子原始 HSTU 推理中大量使用 Triton/CUDA 自定义算子如 jagged concat、attention 内核对齐等。本仓库在 hstu_processor.py 中提供了纯 torch 实现jagged_2d_tensor_concat/torch_split_2d_jagged通过searchsorted计算 batch 索引与行内位置配合index_copy_完成任意数量 Jagged Tensor 的拼接与切分作为融合算子未使能时的兜底实现位置编码、LayerNorm、MLP 等模块全部基于torch.nn原生算子可直接运行在 NPU 上。在 paged_hstu_infer_layer.py 中forward_naive路径使用F.silu(self._linear_uvqk(normed_input))加torch.split完成 u/v/q/k 投影并通过torch_npu._npu_reshape_and_cache完成 KV 写入。优化点一KV Cache 多级缓存、onload 与 offload 流程适配HSTU 面向超长用户行为序列单个用户的完整 KV 无法常驻 NPU 显存。本样例实现了NPU 显存主池 Host 内存pin memory 缓冲两级缓存HSTUGpuKVCacheManager 维护 NPU 侧主池并创建独立的_onload_stream、_offload_stream与_offload_start_event/_offload_end_eventHost 侧由HSTUHostKVStorageManager补丁后为纯 PyTorch 实现按用户维度存储 offload 的历史 KV 数据offload 流程offload_async根据 Host 已 offload 位置与 GPU 缓存位置之差计算新增待 offload 的 chunk通过index_select从主池取出对应页并copy_到 pin memory 的_offload_kvdata_host_buffersoffload_layer_async支持按层粒度在 attention 完成后立即 offload 该层配合_attn_done_events事件同步而无需等待整网结束onload 流程onload将 Host 数据非阻塞拷贝到主池末尾预留的 onload 缓冲get_onload_buffers并记录onload_history_kv_eventsHSTU 层在调用 Paged Attention 前会wait对应层事件确保 onload 完成后再做注意力计算。模型主体 的prepare_kv_cache中Host 中已缓存的页 ID 与本次新分配的页 ID 通过torch.ops.mxrec.concat_2d_jagged使能融合算子时或torch_concat_2d_jagged未使能时拼接后统一写入kv_indices随后触发 onload。优化点二高性能融合算子接入使能--enable_fused_ops all后以下三个 NPU 融合算子被激活由 RecSDK 提供经 build_install_ops.sh 编译安装hstu_pagedPaged Attention 融合算子。在 paged_hstu_infer_layer.py 中通过torch.ops.mxrec.hstu_paged(q, k, v, kv_cache, ...)调用一次完成历史 KV来自页表page_ids/page_offsets/last_page_len与当前 token 的注意力计算替代原始的 Triton attention 内核in_linear_silu线性投影 SiLU 激活融合。通过torch.ops.mxrec.distance_in_linear_silu_forward(normed_input, weight, bias, split_arg_list)一次返回 u/v/q/k 四路输出减少中间张量访存concat_nd_jagged多 Jagged Tensor 拼接算子。前文已给出其在 RecSDK 适配层 的调用方式jagged_2d_tensor_concat_mxrec用于上下文特征、历史序列与候选序列 Embedding 的合并。算子未使能时模型自动回退到 torch 小算子路径保证功能可用性。优化点三ACLGraph 使能为降低 Python 侧算子下发开销模型支持图执行模式。从 paged_hstu_infer_layer.py 看层前向被拆分为forward_inputLayerNorm u/v/q/k 投影 KV 写入与forward_outputPaged Attention 输出层配合 inference_ranking_gr.py 中的use_cudagraphTrue与预分配的_hidden_states、_jagged_metadata、_kvcache_metadata静态缓冲将整网执行路径固化为可重放的图避免逐算子解释执行的开销KV Cache 元数据通过copy_kvcache_metadata同步到静态缓冲保证每次推理的数据更新不破坏图结构。优化点四KV Cache layout 优化引入_npu_reshape_and_cache页式 KV Cache 的写入采用 NPU 原生算子_npu_reshape_and_cache。在 append_kvcache 中先通过pos2page/pos2entry查找表将 token 位置映射为页号与页内偏移再经kv_indptr/kv_indices的 CSR 结构定位物理页最终计算slot_id page_id * page_size entry_idx并以 int32 传入torch_npu._npu_reshape_and_cache(src_k, src_v, paged_k_cache, paged_v_cache, slot_i32)该算子直接在 NPU 侧完成 K/V 的按页写入与 layout 转换规避了 CPU 侧逐个 token 的索引拼接。同时 GPU 端管理器 校验page_size必须为 32 或 64当前 Paged Attention 内核仅支持这两种页大小并在主池页数不足以容纳单 batch 潜在最大页数时主动报错给出明确的配置约束。优化点五分布式推理支持更大 dense 参数模型针对 Embedding Table 占用显存大的场景通过--embed_tp_size对 Embedding 表做张量并行切分见上文分布式推理一节。大表如video_idvocab 1000 万切分到多卡小表保持全量复制从而在受限显存下支持更大的 dense 参数模型推理。优化点六单算子/图模式下 KV Offload 与 dense 层计算并行为掩盖 Host offload 的传输延迟本样例使能了通算掩盖在start_offload_plan阶段记录_offload_start_event后立即继续 dense 层计算每个 HSTU 层在 attention 完成后通过record_attn_done记录事件offload 流等待该事件后即可并行搬移该层 KV而无需阻塞后续层计算最后一层 offload 完成后记录_offload_end_event主循环在下一 batch 的allocate前wait该事件保证缓存一致性。在 Cudagraph图模式下KV offload 同样按层流水化执行offload_layer_async实现 offload 与计算在不同 NPU 流上的重叠。推理流程梳理从数据到 logits综合 inference_ranking_gr.py 的forward与各模块调用一次推理的主要链路为prepare_kv_cache为 batch 内用户分配页、查询 Host 存储并 onload 历史 KV拼接页表元数据Embedding lookupInferenceEmbedding→ 预处理hstu_processor.py 中完成 contextual/history/candidate 的拼接与位置编码逐层执行HSTUBlockInferenceLayerNorm → u/v/q/k 投影可融合in_linear_silu→append_kvcache写入 KV →hstu_paged注意力 → 输出层与残差_postprocessor切出候选 token 的表示并做 L2 归一化MLP输出各任务 logitsoffload_kv_cache_wait将本 batch 新增 KV 落盘到 Host 存储供下一轮同用户复用。总结本文完整梳理了 HSTU 生成式推荐模型在 CANN 平台上的 NPU 推理落地路径环境侧依赖 RecSDK-Torch 镜像与 CANN 8.5.0 nnal 加速库代码侧通过hstu_cann.patch完成设备替换与 KV Cache 管理去 TensorRT-LLM 化性能侧通过 Paged Attentionhstu_paged、融合算子in_linear_silu、concat_nd_jagged、_npu_reshape_and_cache、KV 多级缓存与 offload/计算并行等优化在 Atlas A2 系列产品上实现了单卡与多卡可运行的 HSTU 推理样例。读者可直接以 models/hstu/README.md 为操作入口结合本文给出的源码路径逐层验证各优化点的实现细节并将其中的 KV Cache 多级缓存与通算掩盖思路迁移到其他长序列推荐/大模型推理场景。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考