DeepInfer 源码拆解（三十三）：Platform 抽象与 CUDA/Ascend 双后端——一套引擎如何跨两种加速器

enum：给分支一个稳定的名字

PlatformEnum 只有 GPU 与 NPU 两个值。is_gpu 与 is_npu 都写成 enum is PlatformEnum.GPU 这样的身份比较，而不是字符串比较。这样做的直接好处是：调用方不需要知道我具体是 PlatformGPU 还是未来某个 PlatformROCm，只需要回答“是不是 GPU 语义”。如果以后新增第三种加速器，只要它不属于这两类，现有分支不会被误判。

反过来说，这也意味着引擎在语义上只承认两类执行模型：CUDA 风格的 GPU 与 Ascend 风格的 NPU。这不是一个通用的“任意加速器”抽象，而是一个针对两种已支持硬件的二值契约。ROCm 或其它后端若要接入，要么复用 GPU 语义，要么扩展枚举并逐处审查分支。

tilelang_ops_backend：kernel 目录名

1
class PlatformGPU(Platform):
2
    tilelang_ops_backend = "cuda"
3

4
class PlatformNPU(Platform):
5
    tilelang_ops_backend = "asc"

这个字段只有一个用途，但用途非常关键。tilelang_ops/ops.py 直接用它拼出包名：

1
_pkg = importlib.import_module(f"{__package__}.{PLATFORM.tilelang_ops_backend}")
2
for _info in pkgutil.iter_modules(_pkg.__path__):
3
    _mod = importlib.import_module(f"{_pkg.__name__}.{_info.name}")
4
    for _name, _obj in vars(_mod).items():
5
        if _name.startswith("_") or isinstance(_obj, ModuleType) or not _name.endswith("_tl"):
6
            continue
7
        globals()[_name] = _obj

也就是说，tilelang_ops/cuda/ 与 tilelang_ops/asc/ 不是两套被硬编码的 import 列表，而是两个同构的 Python 包。分发器只做三件事：按平台选包、遍历包内所有模块、把所有以 _tl 结尾的工厂函数重导出到 ops 命名空间。新增一个 kernel，只要在对应目录里实现同名工厂，调用方不用改任何 import。

这种设计把“平台差异”从调用点移到了目录结构上。代价是：两个目录必须保持接口同构，缺失的工厂只能在运行时以 KernelUnavailableError 暴露，编译期没有任何检查。

distributed_backend：NCCL 还是 HCCL

1
distributed_backend = "nccl"   # PlatformGPU
2
distributed_backend = "hccl"   # PlatformNPU

parallel_state.py 在创建进程组时直接读取这个字段：

1
backend = PLATFORM.distributed_backend
2
options = PLATFORM.get_distributed_options()
3

4
self.world = _split_group([ranks.tolist()], backend=backend, pg_options=options)
5
self.local = _split_group(ranks.unsqueeze(-1).tolist(), backend=backend, pg_options=options)
6
...
7
self.cpu_dp_group = _split_group(dp_ranks.tolist(), backend="gloo")

注意最后一行：CPU 侧的 DP 组固定用 gloo，不随平台变化。这说明引擎把“集合通信”拆成了两层：设备间通信走 NCCL/HCCL，主机侧控制面走 Gloo。Ascend 上 Gloo 依然可用，因为它是 CPU 进程间的 TCP 通信，与 NPU 运行时无关。

get_distributed_options() 则返回各自后端的 ProcessGroupOptions：

1
# GPU
2
options = dist.ProcessGroupNCCL.Options()
3
options.config.max_ctas = envs.DEEPINFER_NCCL_NUM_SMS
4
options.config.min_ctas = envs.DEEPINFER_NCCL_NUM_SMS
5

6
# NPU
7
options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()

GPU 侧把 NCCL 可用的 SM 数钉在 DEEPINFER_NCCL_NUM_SMS（默认 32）。这是一个很工程化的决定：NCCL kernel 与计算 kernel 争抢 SM，如果不限制，通信会在大 batch decode 时挤占计算资源；固定 32 个 CTA 相当于给通信留一条固定宽度的车道。NPU 侧没有等价的 max_ctas 配置，HCCL 的选项对象直接返回，资源划分交给 CANN/HCCL 自己处理。

fp8_sf_type：UE8M0 scale 的容器宽度

这是四个字段里最不直观的一个。CUDA 上是 torch.int32，Ascend 上是 torch.int16。它不表示计算精度，而表示 FP8/FP4 量化时，若干个 1 字节的 UE8M0 scale 被打包进什么类型的容器。第九节会完整推导它对张量形状和 kernel 实参的影响，这里先记住结论：CUDA 上 4 个 scale 打成一个 int32，Ascend 上 2 个 scale 打成一个 int16。

三个方法：pin、unpin、分布式选项

pin_host_memory 与 unpin_host_memory 处理的是宿主内存锁定。基类默认是 no-op，注释写得很清楚：不支持 host register 的平台会退回 pageable copy。GPU 子类调用 cudaHostRegister，NPU 子类通过 ctypes 调 aclrtHostRegisterV2。第八节会拆这条路径。

get_distributed_options 在基类里直接 raise NotImplementedError()，说明它不是一个“可选优化”，而是每个后端必须实现的契约。

把四个字段与三个方法放在一起，可以得到一张很短的对照表：

这张表也解释了为什么 Platform 只有 227 行：它不负责“怎么算”，只负责“用哪套规则算”。

延迟初始化：为什么 PLATFORM 不能是一个普通模块变量

deepinfer/__init__.py 的开头有一句很重的注释：

1
# 严禁在本文件内进行任何形式的 import torch 或初始化 device 的行为，会导致子进程的故障

于是 PLATFORM 被实现成一个模块级 __getattr__：

1
def __getattr__(name: str) -> Any:
2
    if name != "PLATFORM":
3
        raise AttributeError(f"module 'deepinfer' has no attribute {name!r}")
4

5
    # NPU initialization patches torch.cuda and initializes the ACL runtime, so
6
    # it must happen only after ASCEND_RT_VISIBLE_DEVICES has been configured.
7
    from deepinfer.platform import get_platform
8

9
    platform = get_platform()
10
    globals()[name] = platform
11
    return platform

这段注释点出了顺序问题的本质：PlatformNPU.__init__ 会 import torch_npu，而 torch_npu 的 transfer_to_npu 会把 torch.cuda.* 重定向到 torch.npu.*，并初始化 ACL 运行时。如果在 ASCEND_RT_VISIBLE_DEVICES 还没设置好之前就完成初始化，后续再改环境变量就无效了——设备可见性在运行时初始化那一刻就被冻结。

所以引擎把平台对象的创建推迟到第一次访问 deepinfer.PLATFORM。而真正的设备可见性由启动脚本/入口点先设置，再由 deep_cluster/hardware 的探测函数读取。get_platform() 本身用 functools.cache 缓存，保证整个进程只构造一次：

1
@functools.cache
2
def get_platform() -> Platform:
3
    from deep_cluster.hardware import accelerator_kind
4

5
    platform_class = PlatformNPU if accelerator_kind() == "ascend" else PlatformGPU
6
    return platform_class()

注意默认分支：只要不是 ascend，就返回 PlatformGPU。accelerator_kind() 可能返回 "cpu"，这时仍然会构造 PlatformGPU。这不是 bug，而是一个明确的降级策略——在没有加速卡的机器上，引擎仍然按 CUDA 语义初始化，后续在真正调用设备 API 时失败。对于只做导入、解析配置或跑 CPU 侧工具的场景，这个默认值让代码不必到处处理 "cpu" 分支。

同一段 __init__.py 里还有 _so_platform_tag()：

1
def _so_platform_tag() -> str:
2
    accelerator = envs.DEEPINFER_ACCELERATOR
3
    if not accelerator:
4
        try:
5
            from deep_cluster.hardware import accelerator_kind
6
            accelerator = accelerator_kind()
7
        except Exception:
8
            accelerator = ""
9
    if accelerator not in ("nvidia", "ascend"):
10
        accelerator = "nvidia"
11
    return f"{accelerator}-{os.uname().machine}"

这段代码只使用环境变量与 deep_cluster.hardware，刻意不 import torch。它生成的标签形如 nvidia-x86_64 或 ascend-aarch64，被用来把 deepinfer/_so/<tag>/ 加入包搜索路径。第六节会看到构建脚本如何产出这个目录。

（图为依据 deep_cluster/hardware/、deepinfer/platform.py 与各消费方源码整理）

探测顺序：Ascend 优先，然后 NVIDIA，最后 CPU

1
@functools.cache
2
def accelerator_kind(*, _root: Path = Path("/")) -> AcceleratorKind:
3
    if ascend.is_available(_root=_root):
4
        return "ascend"
5
    if nvidia.is_available(_root=_root):
6
        return "nvidia"
7
    return "cpu"

顺序是 Ascend 在前。原因不在于 Ascend 更“优先”，而在于两类设备的探测依据完全不同，同时存在的概率极低；把 Ascend 放前面可以让 Ascend 机器在只装了 NPU 驱动时快速命中，避免先去读不存在的 /proc/driver/nvidia。

functools.cache 还意味着探测结果在进程生命周期内不会变化。这对推理服务是合理的：设备热插拔不在支持范围内，缓存可以避免每次取设备数都扫一遍 sysfs。

NVIDIA：procfs、设备节点与 NVML

nvidia.py 的探测路径非常克制：

1
def _proc_entries(*, _root: Path = Path("/")) -> tuple[Path, ...]:
2
    gpu_dir = _path(_root, "/proc/driver/nvidia/gpus")
3
    ...
4

5
def _device_node_indices(*, _root: Path = Path("/")) -> tuple[int, ...]:
6
    device_dir = _path(_root, "/dev")
7
    ...  # 匹配 nvidia0, nvidia1, ...

它先看 /proc/driver/nvidia/gpus 下有没有按 PCI bus id 命名的目录；如果没有，再退回到 /dev/nvidia* 设备节点。两套依据都不需要打开 CUDA 上下文。

更细的一层是计算能力探测：

1
@functools.cache
2
def compute_capability(*, index: int = 0) -> tuple[int, int] | None:
3
    """Compute capability (major, minor) read via NVML; no CUDA context is created."""
4
    library = _nvml()
5
    ...
6
    library.nvmlDeviceGetCudaComputeCapability(handle, byref(major), byref(minor))

它通过 ctypes 加载 libnvidia-ml，调用 nvmlDeviceGetCudaComputeCapability。这条路径只初始化 NVML，不创建 CUDA context。为什么这件事重要？因为一旦创建 CUDA context，再改 CUDA_VISIBLE_DEVICES 就不生效了。引擎需要在“还没决定用哪几张卡”的阶段就知道卡的架构，例如判断是不是 SM100（is_sm_major_10()），以便选择不同的 BF16 峰值算力常数。

deepinfer/__init__.py 里还设置了：

1
os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"

这行让 PyTorch 的 torch.cuda.is_available() 走 NVML 检查而不是创建 context。注释写得很明白：导入依赖时可能调用 torch.cuda.is_available，提前初始化 CUDA 后再设置 CUDA_VISIBLE_DEVICES 会失效。这是一个典型的“库的隐式初始化”问题——你无法控制第三方库在 import 时做什么，只能通过环境变量改变它的检查方式。

NUMA 节点则从 PCI 路径读：

1
bus_id = entries[index].name
2
candidates = (
3
    _path(_root, f"/sys/bus/pci/devices/{bus_id}/numa_node"),
4
    _path(_root, f"/sys/bus/pci/drivers/nvidia/{bus_id}/numa_node"),
5
)

这段逻辑在第二十四篇 NUMA 绑核里已经用过：GPU 的 NUMA 节点决定 CPU 亲和与内存放置。平台层只负责把节点号读出来，怎么用是 deep_cluster 上层的事。

Ascend：ubcore、davinci 与 npu-smi

Ascend 的探测没有 /proc/driver/nvidia 这样的统一入口，代码走的是两条路径：

1
def _entries(*, _root: Path = Path("/")) -> tuple[Path, ...]:
2
    ubcore_dir = _path(_root, "/sys/class/ubcore")
3
    ...  # 匹配 udmaN 或 udmaCcDdEe
4

5
def _device_node_indices(*, _root: Path = Path("/")) -> tuple[int, ...]:
6
    device_dir = _path(_root, "/dev")
7
    ...  # 匹配 davinciN

ubcore 是 Ascend 的统一总线控制器目录，udma 条目对应 DMA 通道；每个条目下的 device/ub_totalues 如果大于 0，说明这个条目真的挂了可用单元。/dev/davinci* 则是设备节点。is_available() 先看 /dev/davinci_manager 是否存在，再看 ubcore 条目，最后退回到设备节点：

1
def is_available(*, _root: Path = Path("/")) -> bool:
2
    if _path(_root, "/dev/davinci_manager").exists() or _device_node_indices(_root=_root):
3
        return True
4
    for entry in _entries(_root=_root):
5
        total_ues = entry / "device/ub_totalues"
6
        if total_ues.is_file() and _read_int(total_ues) > 0:
7
            return True
8
    return False

discover_npu_devices() 返回 (entry_name, physical_id, numa_node) 三元组。它优先用 ubcore 条目，若没有则用 /dev/davinciN 的索引，NUMA 节点缺失时记为 -1。

NUMA 映射有一段值得注意的公式：

1
raw_numa = devices[index][2]
2
if raw_numa < 0:
3
    raise RuntimeError(f"NUMA metadata is unavailable for ascend device {index}")
4
return (raw_numa // 2) * 2 + index % 2

当 sysfs 里读到的原始 NUMA 号不可用时，代码用设备索引做奇偶映射：偶数卡落在一个节点，奇数卡落在另一个节点。这是一个 fallback，而不是通用规则。真正的 NUMA 信息优先来自 npu-smi info -t topo -i 0：

1
@functools.lru_cache(maxsize=1)
2
def _npu_smi_numa_nodes() -> dict[int, int]:
3
    output = _run(["npu-smi", "info", "-t", "topo", "-i", "0"])
4
    cpus_by_numa = physical_core_cpus_by_numa()
5
    ...

它解析 npu-smi 的拓扑输出，把每个 NPU 的 CPU 亲和列表与 lscpu 解析出的物理核集合求交集，取交集最大的 NUMA 节点。这里用 lscpu -p=CPU,NODE,SOCKET,CORE,ONLINE 并只保留每个物理核的最小逻辑 CPU，避免把超线程算成两个核。

（图源：华为昇腾官网 Ascend 950 产品页，图片仅供参考）

可见设备：两套环境变量必须一致

device.py 里定义了 visible_device_env_name()：Ascend 返回 ASCEND_RT_VISIBLE_DEVICES，NVIDIA 返回 CUDA_VISIBLE_DEVICES。visible_device_ids() 负责把逗号分隔的字符串拆成 tuple，并处理空串与 -1：

1
if value.strip() in {"", "-1"}:
2
    return ()

assign_visible_devices() 则把一个 rank 映射到一组连续设备，并写回环境变量。这是多进程启动时按 local rank 切分卡的入口。

真正有意思的是 PlatformNPU.__init__ 里的一致性检查：

1
cuda_devices = os.getenv("CUDA_VISIBLE_DEVICES")
2
ascend_devices = os.getenv("ASCEND_RT_VISIBLE_DEVICES")
3
if cuda_devices is None and ascend_devices is not None:
4
    os.environ["CUDA_VISIBLE_DEVICES"] = ascend_devices
5
    warnings.warn(...)
6
elif cuda_devices != ascend_devices:
7
    raise RuntimeError(
8
        f"CUDA_VISIBLE_DEVICES ({cuda_devices!r}) must match "
9
        f"ASCEND_RT_VISIBLE_DEVICES ({ascend_devices!r})"
10
    )

为什么 NPU 平台要管 CUDA_VISIBLE_DEVICES？因为 torch_npu 的 transfer_to_npu 把 torch.cuda 映射到 torch.npu，而框架里仍有大量代码读 CUDA_VISIBLE_DEVICES（或通过 torch.cuda.device_count() 间接依赖它）。如果两个变量不一致，一部分代码看到 4 张卡、另一部分看到 8 张卡，rank 到物理设备的映射就会错位。这里的策略是：缺一个就补上并告警，两个都有但不相等就直接拒绝启动。

峰值算力常数：432 与 2250/2500

device.py 里还有一个看似突兀的函数：

1
def get_bf16_tflops_limit() -> int:
2
    if accelerator_kind() == "ascend":
3
        return 432
4
    if is_sm_major_10():
5
        return 2500 if platform.machine() in {"aarch64", "arm64"} else 2250
6
    raise RuntimeError("unknown GPU type")

它被 llm_engine.py 用来构造 StatsLogger：

1
tflops_limit=get_bf16_tflops_limit() * 2,  # FP8-equivalent peak: 2x BF16.

也就是说，这个常数不是用来限制计算的，而是 MFU 计算的分母：引擎假定 FP8 峰值是 BF16 的两倍。Ascend 950 取 432 TFLOPS，SM100 在 x86 上取 2250、在 ARM 上取 2500。第十九篇已经拆过 FlopsEstimator 如何按调度长度估算 FLOPs，这里补上分母的来源：它同样是平台相关常量，无法从设备查询里通用地拿到，只能按已知硬件型号硬编码。

硬件探测这一层的设计原则可以总结成三句话：

• 不 import torch：所有探测只读文件系统、跑 npu-smi/lscpu、加载 NVML；
• 不创建运行时：NVML 的计算能力查询与 ACL 的设备节点读取都不初始化 context；
• 可注入根目录：_root 参数让探测逻辑可以在假 sysfs 树上做单元测试。

这三条合起来，才使得“先探测、后初始化”的顺序成为可能。

第一层：torch.cuda 整体重定向

transfer_to_npu 做的事情不是“提供一个相似 API”，而是把 torch.cuda 这个模块对象整体映射到 torch.npu。效果是：torch.cuda.current_stream()、torch.cuda.empty_cache()、torch.cuda.CUDAGraph、torch.cuda.graph(...) 这些调用在 NPU 机器上都能跑，但语义已经换成了 NPU 的流、内存池与图捕获。

这带来一个非常隐蔽的后果：torch.cuda.is_available() 在纯 NPU 机器上返回 True。注释里点出了连锁反应——TileLang 的目标自动探测正是用这个调用判断“当前是不是 CUDA 环境”。如果不管，TileLang 会把每个没有显式 target= 的 kernel 都按 CUDA 编译，然后在 Ascend 上失败。

引擎的应对是把默认目标钉住：

1
os.environ.setdefault("TILELANG_DEFAULT_TARGET", "ascend")

用 setdefault 而不是直接赋值，是为了让用户显式设置的 TILELANG_DEFAULT_TARGET 仍然优先。这是一个很小的细节，但它体现了平台初始化的一个原则：默认值可以改，显式配置不能被覆盖。

第二层：.cuda_stream 属性

1
torch.npu.Stream.cuda_stream = property(lambda self: self.npu_stream)

很多 CUDA 时代的代码会从流对象上取 .cuda_stream 拿到裸指针，再传给 C++ 扩展或第三方库。NPU 的流对象有 npu_stream，但没有 cuda_stream。给类动态补一个 property，比逐个改调用点便宜得多。

这个补丁只解决“读属性”的场景，不解决“把裸指针传给 CUDA API”的场景。后者必须靠平台分支或 C++ 侧的 ACL 实现来处理。

第三层：CUDA Graph 的参数差异

cuda_graph_runner.py 里有一段很典型的“同名 API、不同能力”：

1
# keep_graph=True keeps the raw cudaGraph_t alive so cuda_graph_timer can
2
# match its event-record nodes for ping-pong buffer.
3
# NPU only has instantiated graph and does not support parameter keep_graph.
4
graph = torch.cuda.CUDAGraph(keep_graph=True) if PLATFORM.is_gpu else torch.cuda.CUDAGraph()
5

6
with torch.cuda.graph(graph, pool=self.graph_memory_pool, stream=self.graph_stream,
7
                      capture_error_mode="thread_local"):
8
    ...
9

10
if PLATFORM.is_gpu:
11
    # NPU only has instantiated graph, no need to instantiate.
12
    graph.instantiate()

在 GPU 上，keep_graph=True 保留原始 cudaGraph_t，供图内计时器按事件记录节点做匹配；在 NPU 上，图对象只有“已实例化”这一种形态，既不支持 keep_graph，也不需要显式 instantiate()。代码用 is_gpu 把这两处差异包住，但 torch.cuda.graph(...) 这个上下文管理器本身仍然走别名层。

这也是别名层最危险的地方：调用看起来一样，能力集却不一样。如果某段代码只在 GPU 上测试过，搬到 NPU 上可能直到运行时才因为某个不支持的参数失败。

图内计时器：能力探测而不是别名

timer.py 的 CUDAGraphTimer.supported() 直接返回 PLATFORM.is_gpu：

1
@staticmethod
2
@functools.cache
3
def supported() -> bool:
4
    """Whether the platform backend supports graph probes."""
5
    return PLATFORM.is_gpu

而 C++ 侧的 graph_timer_asc.cpp 更直接——每个方法都抛异常：

1
EventHandle::EventHandle() {
2
    throw std::runtime_error("NPU graph timer primitives are not implemented yet");
3
}

这说明“别名层能跑”不等于“功能可用”。平台层用两种方式表达不可用：Python 侧的能力探测（supported()）与 C++ 侧的显式异常。前者让调用方提前分支，后者保证即使误调用也不会静默返回错误数据。

其它被别名层影响的路径

eplb/__init__.py 的单机 shuffle 是一个好例子：

1
# TODO(NPU): NPU doesn't support symm_mem yet.
2
self._intranode_shuffle = PG.is_within_nvl_domain(ep_group) and PLATFORM.is_gpu

torch.cuda.empty_cache()、symm_mem.empty(...)、symm_mem.rendezvous(...) 在 NPU 上通过别名层未必报错，但 symm_mem 本身没有 NPU 实现。所以这里不能只靠别名，必须显式加 is_gpu。

attn.py 里也有类似的分支：

1
self.use_fp8_einsum_output = PLATFORM.is_gpu and gran == 32 and self.einsum_out_clamp_value is None

FP8 einsum 输出是 CUDA 侧的优化，NPU 上直接关闭。mlp.py 里的 use_mega_gate、fuse_shared 也都以 is_gpu 为前提。

把这几处放在一起看，可以总结出别名层的使用边界：

这张表也解释了为什么仓库里仍有 47 处 PLATFORM.is_npu / is_gpu：别名层降低了移植成本，但没有消除语义差异。

设备可见性：两个变量必须对齐

前面已经看过 CUDA_VISIBLE_DEVICES 与 ASCEND_RT_VISIBLE_DEVICES 的一致性检查。这里补充一个细节：当只有 ASCEND_RT_VISIBLE_DEVICES 时，代码会把它的值复制给 CUDA_VISIBLE_DEVICES 并发出 UserWarning。这个告警不是噪音，而是在提示“有代码正在用 CUDA 变量做设备选择，请确认它看到的是同一组卡”。

内存分配器：expandable_segments

1
os.environ["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:True"

这是 NPU 侧对应 PyTorch PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 的设置。推理引擎的 KV 池、专家权重、通信缓冲尺寸差异很大，固定段分配容易产生碎片；可扩展段让分配器在虚拟地址空间里合并相邻段，降低“显存还剩很多但分配不出大块”的概率。

任务队列：TASK_QUEUE_ENABLE=0

1
os.environ["TASK_QUEUE_ENABLE"] = "0"

Ascend 的 task queue 是一种把算子下发批处理以降低 host 开销的机制。关掉它通常意味着更直接的下发路径，代价是 host 侧开销可能上升。引擎选择关闭，说明在这个场景里确定性/时序可控比批量下发更重要——尤其是与 CUDA Graph 捕获、ZOS 步编排配合时，隐式的队列缓冲会让“这一步到底下发了哪些任务”变得难以界定。

HCCL：缓冲、超时与端口范围

1
os.environ["HCCL_BUFFSIZE"] = "16"
2
os.environ.setdefault("HCCL_CONNECT_TIMEOUT", "240")
3
os.environ.setdefault("HCCL_EXEC_TIMEOUT", "200")
4
os.environ.setdefault("HCCL_HOST_SOCKET_PORT_RANGE", "auto")
5
os.environ.setdefault("HCCL_NPU_SOCKET_PORT_RANGE", "auto")

HCCL_BUFFSIZE=16 把通信缓冲设为 16 MB。这个值直接影响 all-to-all / all-reduce 的分块策略：太小会增加同步次数，太大会挤占 KV 池空间。CONNECT_TIMEOUT=240 与 EXEC_TIMEOUT=200 是秒级的连接与执行超时，比默认值更宽松，目的是容忍大规模集群里慢节点与初始化抖动。端口范围设为 auto 则让 HCCL 自己选可用端口，避免多实例部署时端口冲突。

对比 GPU 侧的 runtime_config.py：

1
os.environ["NCCL_NO_CACHE"] = "NCCL_WIN_STRIDE,NCCL_GIN_TYPE,NCCL_GIN_GDAKI_MAX_QP_RD_ATOMIC"
2
os.environ["NCCL_PXN_DISABLE"] = "1"
3
os.environ["NCCL_CUMEM_ENABLE"] = "1"
4
os.environ["NCCL_GIN_ERROR_QUERY_SEC"] = "432000"

这些变量针对 NCCL 的窗口缓存、PXN 路径、cuMem 与 GIN 错误查询。NCCL_PXN_DISABLE=1 的注释是“为专家重排禁用 NCCL peer access”；NCCL_CUMEM_ENABLE=1 则与可扩展内存/对称内存有关。两边没有一一对应关系，因为它们解决的是不同运行时的问题。

运行时行为开关

1
torch_npu.npu.config.allow_internal_format = True
2
torch_npu.npu.set_compile_mode(jit_compile=False)

allow_internal_format 允许 NPU 使用内部私有格式（例如 NZ 布局），这对卷积与矩阵乘的性能很关键；jit_compile=False 关闭即时编译模式，避免运行时编译带来的不确定延迟。

还有一个临时 hack：

1
# TODO(NPU): Waiting for torch_npu bug fix, hack it temporarily.
2
os.environ.setdefault("TORCH_NPU_DEVICE_CAPABILITY", "9.0")

把设备能力谎报成 9.0，是为了绕过 torch_npu 某个版本对能力的检查。这类 hack 在平台层很常见，也最应该被记录：它不影响正确性，但会在升级 torch_npu 后变成隐藏的技术债。

确定性与 batch invariant

runtime_config.py 里有一段很硬的约束：

1
if PLATFORM.is_npu:
2
    # TODO(NPU): 待完整支持 batch invariant 后移除
3
    assert not envs.DEEPINFER_BATCH_INVARIANT
4

5
if PLATFORM.is_gpu:
6
    deep_gemm.set_pdl(True)
7
    tile_kernels.set_pdl(True)
8

9
if envs.DEEPINFER_BATCH_INVARIANT or PLATFORM.is_npu:
10
    deep_gemm.use_deterministic_algorithms(True)
11
    tile_kernels.use_deterministic_algorithms(True)

三件事值得展开：

1. batch invariant 在 NPU 上被断言禁用。第六篇讲过 batch invariant 的目标：同样的输入 token，不管同批还有谁，输出必须完全一致。它依赖 kernel 内部的归约顺序固定。NPU 上这套保证还没做全，所以引擎选择直接拒绝，而不是给出一个“看起来能用但不保证”的开关。
2. PDL 只在 GPU 上开启。PDL（Programmatic Dependent Launch）是 Hopper 之后的特性，允许下一个 kernel 在前一个 kernel 还没完全结束时启动。NPU 没有对应机制。
3. NPU 无条件启用确定性算法。即使没有开 DEEPINFER_BATCH_INVARIANT，NPU 也会调用 use_deterministic_algorithms(True)。注释写的是“防止 CI 数值不稳定”，说明这是为了测试可复现性，而不是线上性能优先的选择。

Engram 的存储位置

environ.py 里有一行很能说明平台差异的默认值：

1
DEEPINFER_ENGRAM_FORCE_GPU_STORAGE = EnvVar(
2
    lambda v: _bool(v) or envs.DEEPINFER_ACCELERATOR == "ascend"
3
)

在 Ascend 上，这个开关被强制为真。它影响 deepseek.py 里 Engram 参数的物化位置：

1
engram_device = "cuda" if envs.DEEPINFER_ENGRAM_FORCE_GPU_STORAGE else "cpu"

以及 userbuffer.py 里的 RDMA 存储选择：

1
use_cpu_rdma_storage = not envs.DEEPINFER_ENGRAM_FORCE_GPU_STORAGE

换句话说，GPU 上可以把 Engram 表放在主机内存并通过 RDMA 访问；Ascend 上则强制放在设备侧（"cuda" 在别名层下就是 NPU）。这是一个典型的“平台能力决定数据放置”的例子：不是算法变了，而是可用路径变了。

scripts/build.sh：<accelerator>-<arch> 标签

1
PY_BUILDTAG=$(python -c "import sys, sysconfig; print(f'{sysconfig.get_platform()}-{sys.implementation.cache_tag}')")
2
PLATFORM_TAG="${DEEPINFER_ACCELERATOR:-nvidia}-$(uname -m)"
3
BUILD_BASE="build/$PLATFORM_TAG"
4
SO_DST="deepinfer/_so/$PLATFORM_TAG"

构建时先按平台标签建立独立的 build 目录，再把产物软链到 deepinfer/_so/<tag>/：

1
python setup.py build --build-base "$BUILD_BASE"
2
mkdir -p "$SO_DST"
3
for so in "${BUILD_BASE}/lib.${PY_BUILDTAG}/deepinfer/"*.so; do
4
  ln -fs "$(realpath --relative-to="$SO_DST" "$so")" "$SO_DST/"
5
done

用软链而不是复制，有两个好处：一是同一份构建产物不会被复制成两份占空间；二是重新构建时链接自动指向新文件，不需要额外同步。脚本最后还会检查 $SO_DST 非空，否则直接失败——这是一个很实用的防呆：如果 setup.py 因为平台判断错误没有产出 .so，构建会在这里停住，而不是等到运行时 import 失败。

deepinfer/__init__.py：把平台目录加入 __path__

1
_so_dir = os.path.join(os.path.dirname(__file__), "_so", _so_platform_tag())
2
if os.path.isdir(_so_dir):
3
    __path__.append(_so_dir)

deepinfer 是一个常规包，__path__ 决定 import deepinfer._C 时去哪里找子模块。把 _so/<tag> 追加进去后，Python 会在原包目录之后继续搜索这个目录。由于两份 .so 的文件名都是 canonical soabi 名（例如 _C.cpython-312-x86_64-linux-gnu.so），import 机制不需要知道平台标签，只需要在搜索路径里找到第一个匹配项。

这就是“同 arch 的 nvidia / ascend 可以共存于同一棵代码树”的实现方式。标签里的 <arch> 很重要：nvidia-x86_64 与 ascend-aarch64 不会互相干扰；如果是 nvidia-aarch64 与 ascend-aarch64，它们仍然各自在独立目录里，靠 _so_platform_tag() 选择。

setup.py：两套 Extension

setup.py 在导入阶段就判断平台：

1
IS_NPU = accelerator_kind() == "ascend"

然后分两条路构造扩展：

1
if IS_NPU:
2
    from torch_npu.utils.cpp_extension import NpuExtension
3

4
    deepinfer_extension = NpuExtension(
5
        name="deepinfer._C",
6
        sources=[
7
            *COMMON_SOURCES,
8
            "csrc/asc_utils.cpp",
9
            "csrc/graph_timer/graph_timer_asc.cpp",
10
        ],
11
        extra_compile_args={
12
            "cxx": [*CXX_FLAGS, "-fopenmp", "-fvisibility=hidden", "-DWITH_NPU"],
13
        },
14
        libraries=COMMON_LIBRARIES,
15
        include_dirs=INCLUDE_DIRS,
16
    )
17
else:
18
    deepinfer_extension = CUDAExtension(
19
        name="deepinfer._C",
20
        sources=[
21
            *COMMON_SOURCES,
22
            "csrc/cuda_utils.cpp",
23
            "csrc/graph_timer/graph_timer_cuda.cpp",
24
        ],
25
        extra_compile_args={
26
            "cxx": [*CXX_FLAGS, "-fvisibility=hidden"],
27
        },
28
        libraries=[*COMMON_LIBRARIES, "cuda"],
29
        library_dirs=[os.path.join(CUDA_HOME, "lib64", "stubs")],
30
        include_dirs=INCLUDE_DIRS,
31
    )

共同部分包括 csrc/pybind.cpp、调度器、ZOS、输入元数据、采样准备与图计时器的 pybind 层；差异只在工具链与两个平台源文件：

• asc_utils.cpp 用 ACL 的 aclrtMalloc、aclrtMemcpyAsync、aclrtIpcMemGetExportKey、aclrtIpcMemImportByKey 实现内存分配、异步拷贝与 IPC；
• graph_timer_asc.cpp 则是前面看到的“全部抛异常”的占位实现；
• 编译宏 WITH_NPU 控制头文件里的条件编译。

cuda_utils.hpp 里能看到这个宏的具体作用：

1
inline torch::Device get_default_device() {
2
#ifdef WITH_NPU
3
    return torch::kPrivateUse1;
4
#else
5
    return torch::kCUDA;
6
#endif
7
}
8

9
template <typename Deleter>
10
inline torch::Tensor make_from_blob(void *ptr, at::IntArrayRef sizes, Deleter deleter,
11
                                    const at::TensorOptions &options) {
12
#ifdef WITH_NPU
13
    return at_npu::native::from_blob(ptr, sizes, std::move(deleter), options);
14
#else
15
    return torch::from_blob(ptr, sizes, std::move(deleter), options);
16
#endif
17
}

注释解释了为什么 NPU 不能用 torch::from_blob：at::from_blob 在 PrivateUse1 设备上推断 dtype 会失效，必须走 at_npu::native::from_blob。这是一个非常底层的差异——同样是“把裸指针包成 tensor”，两套运行时的元数据路径不同。

IPC：Ascend 的 2 MB 对齐与 PID 校验

asc_utils.cpp 里还有两条硬约束，值得单独说：

1
// Ascend 上 kernel 需要直接跨卡访问 IPC 导入的内存（CP>1 fan-out），
2
// 必须用 ENABLE_PEER_ACCESS，否则导入在跨进程/跨卡时直接失败（507899）。
3
aclrtIpcMemImportByKey(&ptr, info.key, ACL_RT_IPC_MEM_IMPORT_FLAG_ENABLE_PEER_ACCESS);
4

5
// ACL 要求被导出的地址 2MB 对齐，否则 aclrtIpcMemGetExportKey 返回
6
// ACL_ERROR_RT_DRV_INTERNAL_ERROR(507899)，且不区分大小写仅看地址。
7
TORCH_CHECK(ptr % (2 * 1024 * 1024) == 0, "NPU IPC export requires a 2MB-aligned address");

这两条直接决定了 CP（上下文并行）共享内存的分配方式：导出的地址必须 2 MB 对齐，导入时必须显式允许 peer access。GPU 侧对应的是 cudaIpcGetMemHandle / cudaIpcOpenMemHandle，对齐要求通常是 4 KB 或更小。平台层把这些差异收在 C++ 实现里，Python 调用方只看到 export_ipc_tensor / import_ipc_tensor 两个同签名函数。

vendor 构建：换的是包，不是接口

scripts/build_vendor.sh 用同一个脚本构建四个 vendor 包，但根据平台换源码集：

1
case "${DEEPINFER_ACCELERATOR:-nvidia}" in
2
  nvidia)
3
    BUILD_TARGET=CUDA
4
    PACKAGES=(DeepSelect FlashMLA DeepEP DeepGEMM)
5
    INSTALL=(deep_select flash_mla deep_ep deep_gemm)
6
    ;;
7
  ascend)
8
    BUILD_TARGET=ASCEND
9
    PACKAGES=(DeepSelect FlashMLA DeepEP-Ascend DeepGEMM-Ascend)
10
    INSTALL=(deep_select flash_mla deep_ep deep_gemm)
11
    ;;
12
esac

注意 INSTALL 的名字在两套平台上是相同的：deep_ep、deep_gemm。也就是说，Python 侧的 import deep_ep 不需要平台分支；差异被封装在 wheel 里。DeepSelect 与 FlashMLA 通过 DEEP_SELECT_BUILD_TARGET_PLATFORM、FLASH_MLA_BUILD_TARGET_PLATFORM 环境变量选择目标；DeepGEMM 则强制 DG_FORCE_BUILD=1 本地编译。

requirements 也分平台：

1
-r requirements-dev.txt
2
--find-links https://download.pytorch.org/whl/cpu/torch/
3
torch==2.13.0+cpu
4
torch-npu==2.13.0rc1

Ascend 上安装的是 CPU 版 torch 加 torch-npu，不是 CUDA 版 torch。这一点很关键：torch-npu 是作为 torch 的设备后端插件存在的，transfer_to_npu 才能把 torch.cuda 映射过去。如果误装 CUDA 版 torch，不仅浪费体积，还可能让 torch.cuda 指向真实 CUDA 而不是 NPU。

Makefile：把平台选择显式化

1
build-nvidia:
2
  @echo "🛠️ Building Python extension (nvidia) ..."
3
  DEEPINFER_ACCELERATOR=nvidia bash ./scripts/build.sh build
4

5
build-ascend:
6
  @echo "🛠️ Building Python extension (ascend) ..."
7
  DEEPINFER_ACCELERATOR=ascend bash ./scripts/build.sh build

build.sh 本身也要求必须显式给出 DEEPINFER_ACCELERATOR：

1
case "${DEEPINFER_ACCELERATOR:-}" in
2
    nvidia|ascend) ;;
3
    *)
4
        echo "Usage: DEEPINFER_ACCELERATOR=nvidia|ascend $0" >&2
5
        exit 1
6
        ;;
7
esac

这是有意的：构建平台不能靠“探测当前机器”猜。交叉编译、容器构建、CI 缓存都可能让探测结果与目标不符；显式变量把选择权交给构建者。

差异一：store_kv_cache 的缓存视图

GPU 侧的分页 KV 缓存是“值 + scale”两段式布局。store_kv_cache 会先把 kv_cache 拆成 cache_value 与 cache_sf，再按具体格式选择 kernel：

1
block_size, cache = split_kv_cache(kv_cache, layout)
2
if layout.quantized_dtype is None:
3
    cache_value = cache_sf = cache.bf16_data
4
else:
5
    cache_value = cache.quantized_data
6
    cache_sf = cache.sf_data
7

8
store_kernel = ops.store_kv_cache_tl(layout, block_size, **kernel_kwargs)
9
store_kernel(kv, cache_value, cache_sf, slot_mapping)

NPU 侧则直接走一条更简单的路径：

1
if PLATFORM.is_npu:
2
    cache = kv_cache.view(*kv_cache.shape[:2], layout.num_elements) \
3
        if layout.quantized_dtype is None else kv_cache[:, :, 0, :]
4
    ops.store_kv_cache_tl(layout, kv_cache.shape[1])(kv, cache, slot_mapping)
5
    return

差别有三点：

1. 不拆 value/scale。NPU 的缓存被当作一整块 uint8 字节存储，layout 描述里的字节偏移由 kernel 自己解析；GPU 侧则把两段作为独立张量传入。
2. block size 直接取 kv_cache.shape[1]。GPU 侧通过 split_kv_cache 推断，NPU 侧信任张量形状。
3. 没有 fp4_cache_kernel_kwargs。FP4 的值/scale stride 参数只在 GPU 侧需要。

这不是“NPU 实现更粗糙”，而是两边对缓存布局的约定不同。GPU 侧的 kernel 工厂需要知道 value 与 scale 的 stride；NPU 侧的 kernel 在一个连续的字节缓冲里按 layout 计算偏移，天然不需要外部再传 stride。

同一文件里还有一个 Ascend 专用的 workaround：

1
# TODO(A5): 待 torch_npu 修复或有替代 kernel 后移除
2
def _slot_scatter(cache, block, offset, values):
3
    """``cache[block, offset] = values``，在 Ascend 上绕开 IndexPutV2 的 32 位偏移回绕。"""
4
    if not PLATFORM.is_npu:
5
        cache[block, offset] = values
6
        return
7
    flat = torch.as_strided(cache, (cache.untyped_storage().nbytes() // cache.element_size(),), (1,), 0)
8
    dst = cache.storage_offset() + block * cache.stride(0) + offset * cache.stride(1)
9
    ...
10
    flat.index_put_((dst.reshape(-1),), values.reshape(-1))

IndexPutV2 在 Ascend 上用 32 位偏移，大缓存会回绕。代码把多维索引摊平成一维偏移，绕过这个限制。这是平台层最典型的一类补丁：不是算法问题，是运行时实现的位宽问题。

另外，NPU 上写 FP8 缓存前会把张量视图转成 uint8：

1
if PLATFORM.is_npu and cache.dtype == torch.float8_e4m3fn:
2
    cache = cache.view(torch.uint8)
3
    values = values.view(torch.uint8)

原因很可能是 Ascend 的 index_put_ 对 FP8 类型的支持不完整；转成字节视图后，写入的是同样的位模式，语义不变。

差异二：prepare_v4_decode_metadata 的切分方式

GPU 侧的 prepare_decode_metadata.py 用 TileLang 的 CUDA 后端，关闭 TMA 与 warp specialization：

1
pass_configs={
2
    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
3
    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
4
    tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
5
}

Ascend 侧的 prepare_v4_decode_metadata.py 则使用 tilelang.ascend.language：

1
import tilelang
2
from tilelang.ascend import language as T
3

4
@tilelang.jit(
5
    out_idx=None,
6
    pass_configs={
7
        tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True,
8
        tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
9
    },
10
)
11
def copy_block_tables_tl(max_seqlen_q: int, max_num_blocks: int) -> Any:
12
    ...
13
    num_cores = tile_kernels.get_num_vec_cores()
14
    ...
15
    with T.Kernel(num_cores) as pid:
16
        bt_ub = T.alloc_shared((bt_len_aligned,), "int32")
17
        ...

两边的并行模型完全不同：CUDA 侧用线程块/线程，Ascend 侧用向量核（vector core）。tile_kernels.get_num_vec_cores() 在 asc/ 目录里出现 27 次、共 40 处调用，几乎所有 Ascend kernel 都用它决定 T.Kernel 的网格大小。与之配套的还有 get_max_ub_per_vector_core(use_simt=...)，用来计算 Unified Buffer（UB）能放多少数据。

（图源：华为昇腾官网 Ascend 950 产品页“混合编程”特性图）

这解释了一个常见疑问：为什么不能把 CUDA TileLang kernel 直接编译到 Ascend？因为两边不只是指令集不同，并行抽象也不同。CUDA 的 T.Kernel 映射到 block，Ascend 的 T.Kernel(num_cores) 映射到向量核；UB 与 shared memory 的容量、对齐、DMA 方式都不一样。DeepInfer 选择为每个 kernel 写两份实现，而不是追求一份源码跨平台。

差异三：fused_transform 的 RoPE 与 cast

第三十二篇拆过 fused_transform 的五段流水。这里只看平台分支：

1
rope_offset = head_dim - rope_args.rotary_dim if rope_args.offset is None else rope_args.offset
2
half = rope_args.rotary_dim // 2
3
cos = rope_args.cos_sin_cache[:, :half]
4
sin = rope_args.cos_sin_cache[:, half:]
5
if PLATFORM.is_npu:
6
    cos, sin = rope_args.cos_sin_cache, None

GPU 侧把 cos/sin cache 拆成两半分别传入；NPU 侧传整个 cache，sin 传 None。这说明 Ascend kernel 内部自己完成了拆分，或者使用了不同的旋转公式实现。

写缓存时的实参也不同：

1
cache_args = (cache_out.kv_cache[:, :, 0, :] if cache_out is not None else None,) \
2
    if PLATFORM.is_npu else (cache_quant, cache_sf)

GPU 传 value 与 scale 两段；NPU 传整块缓存的一个切片。这与 store_kv_cache 的差异同源。

cast 路径还有一条限制：

1
elif PLATFORM.is_npu and cast_args.num_per_channels != 32:
2
    # Ascend per_token_cast only supports g32. Other generic
3
    # group sizes use the library's PyTorch reference.
4
    encoded, sf = torch_cast(cast_input, cast_args.fmt,
5
                             block_size=(1, cast_args.num_per_channels), ...)
6
else:
7
    encoded, sf = quant.per_token_cast(...)

Ascend 的 per_token_cast 只支持 group size 32；其它 group size 退回 PyTorch 参考实现。这是一个“性能路径有约束，正确性路径仍在”的设计：不会因为不支持的粒度直接报错，但会明显变慢。

差异四：手写 SIMD 原语

asc/kv_quant.py 的开头写得很清楚：

1
"""Register-only BF16 KV quantization primitives for Ascend SIMD.
2

3
Scale exponent/inverse math and sticky-bit FP4 packing are adapted locally from
4
``tile_kernels/quant/per_token_cast_asc.py`` ...
5
"""

它用 tilelang.ascend.language.simd 里的 vdup、vsel、vcmps、vdiv、vsqrt 等原语，在寄存器里完成倒数表构建、scale 编码与 FP4 打包。asc/rstd_1ulp.py 则实现“补偿式 rsqrt”：

1
"""Compensated RSTD, matching hai_llm.tilelang_kernels.ascend.helper.rstd_1ulp.
2

3
Shared by norm and fused_transform; keep the FP32/FMA operation order identical
4
so both inference paths follow the hai-llm numerical standard.
5
"""

这两个文件在 CUDA 侧没有对应物，因为 CUDA 侧可以直接调用库函数或使用硬件指令。Ascend 侧为了保证与训练侧（hai-llm）的数值标准一致，必须把运算顺序、FMA 次序、甚至 1 ULP 的误差都复刻出来。

差异五：阈值与形状约束

平台差异不只在实现里，还渗进了常量。下面这张表列的是源码里能直接读到的约束：

这些常量不是“调参”，而是接口契约的一部分。它们说明：即使函数签名相同，可接受的输入域也可能不同。调用方如果只按 CUDA 的上限准备数据，迁到 Ascend 上就会撞到断言。

（图为依据 tilelang_ops/ops.py、两套后端目录与 store_kv_cache 实现整理）

GPU 侧：两行调用

1
def pin_host_memory(self, tensor: torch.Tensor) -> None:
2
    import torch
3
    torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), tensor.nbytes, 0)
4

5
def unpin_host_memory(self, ptr: int) -> None:
6
    import torch
7
    torch.cuda.cudart().cudaHostUnregister(ptr)

第三个参数 0 表示不传 flags。cudaHostRegister 默认会把整页锁定，失败时返回错误码；这里没有检查返回值，也没有告警。

NPU 侧：ctypes + ACL 错误码

NPU 侧复杂得多，原因有三：torch_npu 没有暴露 host register 的 Python API；ACL 的调用需要当前线程的 device context；错误码有“已注册/未注册”两种良性情况。

第一步是在 __init__ 里解析函数符号：

1
@classmethod
2
def _resolve_acl_host_register_api(cls) -> None:
3
    """Resolve register/unregister once at init; raise RuntimeError if
4
    unavailable (check the driver version)."""
5
    try:
6
        lib = ctypes.CDLL("libascendcl.so")
7
    except OSError as e:
8
        raise RuntimeError(f"failed to load libascendcl.so: {e}; is an unsupported driver version in use?") from e
9
    register = getattr(lib, "aclrtHostRegisterV2", None)
10
    unregister = getattr(lib, "aclrtHostUnregister", None)
11
    if register is None or unregister is None:
12
        raise RuntimeError(
13
            "aclrtHostRegisterV2/aclrtHostUnregister not found in libascendcl.so; ..."
14
        )
15
    register.argtypes = (ctypes.c_void_p, ctypes.c_uint64, ctypes.c_uint32)
16
    register.restype = ctypes.c_int
17
    unregister.argtypes = (ctypes.c_void_p,)
18
    unregister.restype = ctypes.c_int
19
    cls._acl_host_register_v2 = register
20
    cls._acl_host_unregister = unregister

注释说 libascendcl.so 已经在 LD_LIBRARY_PATH 上（torch_npu 依赖它），所以按默认搜索路径加载即可，运行时状态与 torch_npu 共享。argtypes 与 restype 是必须的：在 64 位地址上，如果不声明 c_void_p，ctypes 会把指针截成 32 位。

第二步是线程本地 context：

1
@staticmethod
2
def _bind_current_device() -> None:
3
    """Establish the ACL device context for the current thread (the context
4
    is thread-local, otherwise a bare host register/unregister call fails
5
    with 107002); set_device is idempotent and lazily initializes."""
6
    import torch
7
    torch.npu.set_device(torch.npu.current_device())

ACL 的 device context 是线程本地的。如果共享内存池是在一个线程里创建、在另一个线程里销毁，直接调 aclrtHostRegisterV2 会因为“没有当前 context”返回 107002。set_device 是幂等的，调用它可以把当前线程绑定到同一张设备。

第三步是注册与错误码处理：

1
def pin_host_memory(self, tensor: torch.Tensor) -> None:
2
    self._bind_current_device()
3
    err = self._acl_host_register_v2(tensor.data_ptr(), tensor.nbytes, self._ACL_HOST_REG_PINNED)
4
    if err != self._ACL_SUCCESS and err != self._ACL_ERROR_HOST_MEMORY_ALREADY_REGISTERED:
5
        warnings.warn(
6
            f"aclrtHostRegisterV2(0x{tensor.data_ptr():x}, {tensor.nbytes}, PINNED) failed "
7
            f"(aclError={err}): H2D/D2H from shm will use the pageable slow path.",
8
            stacklevel=2,
9
        )

三个常量：

1
_ACL_SUCCESS = 0
2
_ACL_HOST_REG_PINNED = 0x10000000
3
_ACL_ERROR_HOST_MEMORY_ALREADY_REGISTERED = 507910
4
_ACL_ERROR_HOST_MEMORY_NOT_REGISTERED = 507911

0x10000000 是 acl_rt.h 里的 ACL_HOST_REG_PINNED 宏值，注释说它与 npu.c 的编译期宏一致。507910 表示这块内存已经注册过，视为成功；507911 表示未注册，unpin 时视为成功。其它错误码则告警。

告警的措辞很讲究：“H2D/D2H from shm will use the pageable slow path”。这不是失败退出，而是降级：pin 失败后，共享内存照样能用，只是拷贝走慢路径。为什么不直接抛异常？因为 pin 失败往往与驱动版本、权限或节点状态有关，而这些问题的严重程度因部署而异；让整个进程起不来可能比慢一点更糟。告警作为一种“降级金丝雀”，让运维在日志里发现，而不是在性能上盲猜。

unpin 的注释说明了为什么不能静默：

1
def unpin_host_memory(self, ptr: int) -> None:
2
    """Unpin via aclrtHostUnregister ... Not-registered (507911) is silenced;
3
    other failures always warn (a leaked pin must stay visible)."""

pin 泄漏是宿主内存被长期锁定的问题：进程不退出就无法回收，久了会把主机内存耗尽。所以除了“未注册”这一良性情况，其它失败一律告警。

（图为依据 deepinfer/platform.py 与 deepinfer/shared_memory.py 核心路径整理）

这张图想说明的三件事

1. 同一个契约，两种调用约定：GPU 侧失败不检查、NPU 侧按错误码分级处理。这不是 NVIDIA 不会报错，而是 CUDA 的 cudaHostRegister 在这条路径上失败的概率与后果在引擎看来可以接受；ACL 的良性错误码需要显式区分。
2. context 是线程本地的：GPU 的 CUDA context 在进程级隐式存在，NPU 的 ACL context 需要每个线程显式绑定。这个差异决定了 _bind_current_device 必须出现在每次 pin/unpin 之前。
3. 失败是降级而不是崩溃：pin 只是性能路径，不是正确性前提。

什么是 UE8M0

FP8（E4M3/E5M2）与 FP4（E2M1）都是“块量化”格式：每 32 个元素（或 16 个、128 个，取决于格式）共享一个缩放因子。这个因子通常写成 float32 或 UE8M0——一种只有指数位、没有尾数位与符号位的 8 位格式。UE8M0 的取值是：

也就是“2 的整数次幂”。DeepGEMM 的 UE8M0 路径要求 scale 是 power-of-two，这样解码时只需要指数加减，不需要乘法。一个 1 字节的 UE8M0 完全可以单独存成 uint8，但很多 kernel 希望一次读多个 scale，减少访存指令与索引计算。于是就有了“打包”。

GPU：4 个 scale 打进一个 int32

int32 是 4 字节，所以 CUDA 上的打包宽度是 4。看 cache_common.py 的分配逻辑：

1
if layout.sf_format == "ue8m0":
2
    sf_packs = layout.sf_bytes // PLATFORM.fp8_sf_type.itemsize
3
    assert sf_packs * PLATFORM.fp8_sf_type.itemsize == layout.sf_bytes
4
    sf_shape = leading_shape if sf_packs == 1 else (*leading_shape, sf_packs)
5
    sf_data = torch.empty(sf_shape, dtype=PLATFORM.fp8_sf_type, device=device)

假设一个 token 的 head 维度是 128，elements_per_sf=32，那么 sf_bytes = 4。在 GPU 上，sf_packs = 4 // 4 = 1，sf_shape 就是 leading_shape，一个 int32 装 4 个 scale。在 Ascend 上，sf_packs = 4 // 2 = 2，sf_shape 多出一维 2，每个 int16 装 2 个 scale。

而 kernel 侧永远拿到的是一维字节视图：

1
def sf_data_as_int32(sf_data: torch.Tensor, leading_shape: tuple[int, ...]) -> torch.Tensor:
2
    """View packed UE8M0 scale bytes in the int32 form used by load kernels."""
3
    sf_bytes = sf_data.numel() * sf_data.element_size() // torch.Size(leading_shape).numel()
4
    assert sf_bytes % torch.int32.itemsize == 0
5
    return sf_data.reshape(*leading_shape, -1).view(torch.uint8).view(torch.int32)

注意它把 Ascend 的 int16 张量先 view(torch.uint8) 再 view(torch.int32)。一个 int16 占 2 字节，两个 int16 拼接后正好是 4 字节，可以塞进一个 int32。sf_bytes 必须是 4 的倍数，否则断言失败。

这张“存储宽度”与“kernel 读取宽度”分离的设计，让存储层可以按平台选择容器，而加载 kernel 只需要一种 int32 视图。代价是：如果某个 layout 的 sf_bytes 不是 4 的倍数（例如 single scale 的情况），sf_data_as_int32 会断言失败，必须走另一条路径。

norm 与 fused_transform：显式断言 4 的倍数

_alloc_fp8_sf 的注释写得更细：

1
"""分配 fp8 cast 的 sf，返回 (对外视图, kernel 实参)。布局与 tile_kernels.quant 一致：
2

3
- row-major: float32 (B, M, d_scale) 连续
4
- row-major packed ue8m0（mega-moe x_sf 布局）: PLATFORM.fp8_sf_type
5
  (B, M, d_scale/pack_factor) 连续，每个元素打包 k 方向连续 pack_factor 个 ue8m0；
6
  kernel 以 uint8 (B, M, d_scale) 视图逐字节写入
7
...
8
"""
9
pack_dtype = PLATFORM.fp8_sf_type
10
pack_factor = pack_dtype.itemsize
11
if not use_tma_aligned_col_major_sf:
12
    if use_packed_ue8m0:
13
        assert d_scale % pack_factor == 0
14
        sf = torch.empty(batch, seqlen, d_scale // pack_factor, dtype=pack_dtype, device=device)
15
        return sf, sf.view(torch.uint8)

这里有一个容易忽略的推论：GPU 上 pack_factor=4，d_scale 必须是 4 的倍数；Ascend 上 pack_factor=2，d_scale 必须是 2 的倍数。如果 d_scale 是 2 但不被 4 整除（例如 head_dim=64、group=32 时 d_scale=2），GPU 的 fast path 会断言失败，而 Ascend 可以走。这也是为什么 cache_common.py 在 NPU 上干脆把 FP8_FLOAT32_SF_INDEX_KV_FORMAT 从支持列表里去掉：

1
SUPPORTED_KV_CACHE_FORMATS = [
2
    V41_FORMAT,
3
    V41_FP4_FORMAT,
4
    FP4_INDEX_KV_FORMAT,
5
    *([] if PLATFORM.is_npu else [FP8_FLOAT32_SF_INDEX_KV_FORMAT]),
6
    MXFP8_INDEX_KV_FORMAT,
7
]

GPU 上多一个“FP8 值 + FP32 scale”的索引缓存格式。FP32 scale 不需要打包，但对应的 NPU kernel 没有实现，所以直接从支持列表里删掉。这个选择的用意是让不支持在构造期就暴露：如果有代码试图用该格式，resolve_kv_cache_layout 会直接拒绝，而不是在 kernel 里产生错误数据。

fused_sample 与 sampler 的阈值

fp8_sf_type 不是唯一会“按平台变”的常量。fused_sample 的测试代码里：

1
vocab_size = 64 if PLATFORM.is_npu else 32  # ASC 要求至少一个完整的 SIMD 向量。

Ascend 的一条 SIMD 指令处理 64 个元素（以 float32 计，具体取决于向量宽度）。如果测试词表只有 32，无法填满一个向量，kernel 的边界路径不会被覆盖。这个常量只出现在测试里，但它暴露了一个事实：Ascend kernel 的最小处理粒度更大，很多形状约束来自硬件而不是算法。

sampler.py 里：

1
assert force_topk <= (4096 if PLATFORM.is_npu else 16384)

GPU 上允许 force_topk 到 16384，Ascend 上只到 4096。原因是 Ascend sampler 的 aligned_logprobs、draft_local_ids、identity_rows 三个输出都要按 force_topk 分配，并做一次 remap kernel；4096 是 UB/寄存器能稳定放下的上限。

什么时候该用 fp8_sf_type，什么时候该用分支

从代码分布看，规则很清楚：

• 只是“容器宽度不同”：用 PLATFORM.fp8_sf_type，不要写分支。cache_common、norm、fused_transform 的 _alloc、linear 的类型断言、containers.py 的 sf_dtype 都属于这一类。
• “格式/能力不同”：用 is_npu / is_gpu 分支。cache_common 删除格式、fused_transform 的 cast 粒度、load_paged_kv_cache 的 index_batch_size 都属于这一类。

这条规则的价值在于：容器宽度是一个标量，分支是一个行为。把标量变化做成字段，可以让 25 处调用点共享同一套形状推导；把行为变化做成分支，则必须在每个调用点重新论证。

问题形态

GPU 侧有两条 import：

1
# sampler.py / sampler_utils.py / index_attn.py
2
from deep_select.interface import topk
3

4
from flash_attn.cute.interface import _flash_attn_fwd

deep_select 的 topk 在 NPU 上没有可用 kernel；flash_attn.cute.interface._flash_attn_fwd 在 NPU 上也没有对应实现（ViT 的 varlen 注意力）。直接修改每个调用点意味着把平台分支散落到 sampler、indexer、ViT 等多处，而且第三方库升级后还要重新检查。

方案：拦截 from import 语句

patch_imports.py 用替换 builtins.__import__ 的方式，在首次 import 时把目标符号换成另一个模块里的同名符号：

1
_patch_map: dict[str, dict[str, str]] = {}
2
_orig_import = builtins.__import__
3

4
def _hooked_import(name, globals=None, locals=None, fromlist=(), level=0):
5
    # 只处理 `from name import a, b`（非 *）；fromlist 可能为 None
6
    if level != 0 or not fromlist:
7
        return _orig_import(name, globals, locals, fromlist, level)
8

9
    names = tuple(n for n in fromlist if n != "*")
10
    target_map = _patch_map.get(name)
11
    if not names or not target_map or target_map.keys().isdisjoint(names):
12
        return _orig_import(name, globals, locals, fromlist, level)
13

14
    caller = (globals or {}).get("__name__", "")
15
    m = ModuleType(name)
16
    real: ModuleType | None = None
17

18
    def real_module() -> ModuleType:
19
        nonlocal real
20
        if real is None:
21
            # 拿原始 target module，其内部 import 仍会被 patch
22
            real = _orig_import(name, globals, locals, names, level)
23
        return real
24

25
    for n in names:
26
        src = target_map.get(n)
27
        if src is None:
28
            continue
29
        src_mod = src.rsplit(".", 1)[0]
30
        if caller == src_mod or caller.startswith(src_mod + "."):
31
            # src 自身 import 回对应 target 时取原始对象
32
            setattr(m, n, getattr(real_module(), n))
33
        else:
34
            setattr(m, n, Lazy(f"{name}.{n}", src))
35
    m.__getattr__ = lambda attr: getattr(real_module(), attr)
36
    return m

这个实现有几个值得注意的点：

1. 只拦截 from X import a，不拦截 import X 与 import X as Y。后者拿到的是原始模块对象，如果使用方调用 X.a(...)，补丁不生效。DeepInfer 的调用点恰好都是 from ... import ...，所以这个限制没有造成问题；但这意味着新增调用点必须遵守同样的写法。
2. 构造一个假的 ModuleType 返回给 import 语句，模块名仍然是 deep_select.interface，但其中被替换的符号是 Lazy 对象。
3. Lazy 首次调用才加载：Lazy.__call__、Lazy.__getattr__ 都会触发 _load()，从源模块 import 真实对象。这样启动时不需要导入被替换的实现，避免循环依赖与初始化顺序问题。
4. real_module() 兜底：未被替换的符号通过 m.__getattr__ 从原始模块取。也就是说，补丁只替换指定符号，其余行为不变。
5. 自引用保护：如果源模块自己 import 回目标符号，取原始对象，避免把 deep_select.interface.topk 换成它自己。
6. 只调用一次：patch_imports._called 防止重复安装 hook，否则 hook 会层层嵌套。

两个补丁

1
from deepinfer.npu_ops.patch_imports import patch_imports
2

3
patch_imports(
4
    {
5
        "deepinfer.npu_ops.deep_select": {
6
            "deep_select.interface": "topk",
7
        },
8
        "deepinfer.npu_ops.vit_attn": {
9
            "flash_attn.cute.interface": "_flash_attn_fwd",
10
        },
11
    }
12
)

格式是 {source_module: {target_module: symbol}}：把 deep_select.interface.topk 换成 deepinfer.npu_ops.deep_select.topk，把 flash_attn.cute.interface._flash_attn_fwd 换成 deepinfer.npu_ops.vit_attn._flash_attn_fwd。

npu_ops/__init__.py 在 PlatformNPU.__init__ 的末尾被导入：

1
# TODO(NPU): Intended to be removed.
2
import deepinfer.npu_ops  # noqa: F401

注释直接写了“打算移除”。也就是说，这是一个过渡方案；等 deep_select 的 sampler topk 与 ViT 注意力在 NPU 上补齐，整个 npu_ops 目录可以删掉。

deep_select.py：部分路径回退

1
def topk(*args: Any, **kwargs: Any) -> tuple[torch.Tensor | None, torch.Tensor]:
2
    if kwargs.get("return_value", True):
3
        # Sampler is not supported yet, fallback to torch.
4
        return _torch_topk(*args, **kwargs)
5
    # Indexer use real topk kernel.
6
    return _topk(*args, **kwargs)

同一个 topk 接口有两个用途：sampler 需要返回值（return_value=True），indexer（稀疏注意力）只需要索引。NPU 上 indexer 用的 topk kernel 已经实现，sampler 用的还没有，所以按参数分流：sampler 走 torch.topk，indexer 走真实 kernel。

_torch_topk 还处理了索引类型：

1
def _torch_topk(input, topk, sorted=False, indices_type=torch.int64, **kwargs):
2
    data, index = torch.topk(input, topk, dim=-1, sorted=sorted)
3
    if index.dtype is not indices_type:
4
        index = index.to(indices_type)
5
    return data, index

这和第十一篇里 indexer 的 index_attn.py 对得上：indexer 期望的索引类型与 sampler 不同。

vit_attn.py：用 NPU Fusion Attention 重写 varlen 注意力

1
def _flash_attn_fwd(q, k, v, cu_seqlens_q=None, **kwargs) -> tuple[torch.Tensor, None]:
2
    """ViT 用的 varlen attention：非 causal，每段内做全量注意力。
3

4
    cu_seqlens_q: host list（InputMetadataVL.vit_attn_cu_seqlens，每段累计末端）。
5
    """
6
    result = torch_npu.npu_fusion_attention(
7
        q, k, v,
8
        q.shape[-2],
9
        "TND",
10
        scale=1.0 / math.sqrt(q.shape[-1]),
11
        actual_seq_qlen=cu_seqlens_q,
12
        actual_seq_kvlen=cu_seqlens_q,
13
        sparse_mode=0,
14
    )[0]
15
    return (result, None)