vLLM-Ascend 的 GroupedMatmulSituQuant：Ascend 950 上 MXFP8×MXFP4 分组矩阵乘、SiTU 激活与动态 MX 量化的单次融合算子

• 人工智能
• 大模型
• 模型推理服务
• Ascend
• CANN

GroupedMatmulSituQuant 是 vLLM-Ascend 为Ascend 950（A5，arch35）平台引入的融合自定义算子：它将npu_grouped_matmul + situ_mx_quant这一串行算子链合并为一次 kernel 启动（single launch），在 MoE（Mixture of Experts）推理场景中把「分组矩阵乘法（MXFP8 激活 × MXFP4 专家权重）、SiTU 激活、动态 MX 量化」三段计算一体完成。本文以 csrc/moe/grouped_matmul_situ_quant/README.md 为主线，结合算子定义、Host Tiling、设备 Kernel、Torch 适配层与端到端测试源码，完整讲解该算子的接口契约、四形态权重分发、图捕获执行机制与数值精度实现。读完本文，你将掌握该算子在 A5 上的调用方式、参数语义、shape 约束以及其背后的融合设计思路。

一、为什么需要单次启动的融合算子

在常规的 MoE 前向链路中，Kimi 等 W4A8（4-bit 权重、8-bit 激活）模型的一层 FFN 通常被拆成多段执行：

1. npu_grouped_matmul：把当前 batch 的 token 按路由结果分组，分别与各专家的低比特权重做矩阵乘，得到 BF16 中间结果；
2. situ_mx_quant：对中间结果施加 SiTU 激活（beta * tanh(gate/beta) * sigmoid(gate) * up），再做动态 MX 量化，输出 FP8-E4M3 数据与 E8M0 的 MX scale。

两次 kernel 启动意味着两次设备访存、两段调度开销，以及中间 BF16 结果的额外写读。GroupedMatmulSituQuant 的目标就是把这条拆分链（split chain）熔成一个 launch：设备侧直接消费 MXFP8 激活与 MXFP4 专家权重，完成分组矩阵乘后，在 epilogue 阶段立即计算 SiTU 激活并动态量化，最终只产出output（FP8）与outputScale（E8M0）两个输出张量。从源码看，该算子是从 PR #15871 引入，且仅面向 Ascend 950（arch35）设计，是 A5 专属的 MX 数据通路实现。

二、算子接口：输入、输出与核心约束

2.1 输入输出一览

官方接口文档 docs/aclnnGroupedMatmulSituQuant.md 定义了如下张量契约：

2.2 官方约束条件

• groupListType仅支持0（cumulative，累加式）与1（count，计数式）；
• linearBeta必须为正数才能走融合路径；
• 只实现 MX A8W4 组合；bias与smoothScale有意不支持；
• K以及输出宽度的一半（N/2）必须是 64 的倍数；
• 唯一支持平台：Ascend 950。

2.3 算子定义层的形状与数据类型契约

在算子定义 op_host/grouped_matmul_situ_quant_def.cpp 中可以确认上述约束的落地细节：

• x、x_scale、group_list为必填（REQUIRED）ND 格式输入；
• weight与weight_scale为动态（DYNAMIC）输入，其中权重格式被固定为FORMAT_FRACTAL_NZ_C0_16（对应 K3 A5 MXFP4 的 C0=16 布局），并通过IgnoreContiguous()接受非连续张量——这是因为权重以 NZ 存储的转置视图到达（loader 先做npu_format_cast(...).transpose(1, 2)），NZ 张量无法强制 contiguous，而 kernel 本身通过格式元数据读取它们；
• 两个属性group_list_type（默认 0）与beta、linear_beta（默认 1.0）均为可选；
• AICore 配置注册在ascend950平台，开启动态编译静态化（DynamicCompileStaticFlag）与动态形状支持（DynamicShapeSupportFlag）。

InferShape 实现 op_host/grouped_matmul_situ_quant_infershape.cpp 给出了输出形状的精确规则：设n = packedWeightElements / ((isTensorList ? 1 : E) * (K/2))，n2 = n/2，则输出为y = (M, n2)、yScale = (M, ceil(n2/64), 2)。这里yScale的第三维固定为 2，正是 gate/up 两半输出的 scale 拼接结构。

三、Host Tiling：形状校验、块切分与参数解码

Host 侧 Tiling 实现在 op_host/grouped_matmul_situ_quant_tiling.cpp，它负责把运行时形状与属性编码进 64 字节的 tiling 头（结构体定义见 op_kernel/grouped_matmul_situ_quant_tiling.h）。

3.1 形状校验链

Tiling 入口按以下顺序做硬性校验，任何一项失败都返回GRAPH_FAILED：

• x必须为二维[M, K]，groupList必须为一维，权重维数至少为 1；
• K > 0且K % 64 == 0（MX scale 配对要求）；
• 从权重打包元素数反推N：N = packedWeightElements / (E * (K/2))（TensorList 形式时按groupList长度取 E）；
• N > 0、N % 2 == 0且n2 = N/2必须是MAIN_BLOCK_N2 = 64的倍数；
• 属性校验：groupListType ∈ {0, 1}，beta != 0，linearBeta != 0。

3.2 Tiling 头关键字段

struct SituTilingHeader { uint32_t coreNum; // 使用 AIC 核心总数（GetCoreNumAic） uint32_t activeCount; // 专家数 E（设备入口语义） uint32_t kSize; // K uint32_t nSize; // N uint32_t baseM; // 固定 128（命中 kbL1Size=512 快路径） uint32_t mainBlockSize; // 主块 N2 宽度：64 或 128（见下文） uint32_t firstTailBlockSize; // 预留（当前为 0） uint32_t reserved; // 低 1 bit = groupListType，bit1 = TensorList 标志 uint64_t mainBlockCount; uint64_t firstTailBlockCount; float beta, invBeta, linearBeta, invLinearBeta; };

其中beta与linearBeta在 Tiling 阶段即预计算好各自的倒数，设备端 epilogue 直接用乘法代替除法。

3.3 N/K 联合 tile（nk_tile）策略

Tiling 代码中值得注意的优化：当n2 % 128 == 0时，mainBlockSize取 128（宽块），设备侧nL1Size = mainBlockSize * 2 = 256，配合kbL1Size自动落到 256；B 单槽nL1Align × kbL1 = 256×256的 FP8 数据恰为 64KB，与基线128×512的容量相同。若n2不被 128 整除，则回退到基线块宽 64，保证合法 shape 全覆盖（详见 tiling 源码中关于 nk_tile 的注释）。

四、Torch 适配层：V2 对齐的双入口与四种权重形态

该算子的对外调用不是直接暴露 aclnn 接口，而是通过 grouped_matmul_situ_quant_torch_adpt.h 注册到torch.ops._C_ascend（注册入口位于 csrc/torch_binding.cpp，图/编译期的 Meta 实现在 csrc/torch_binding_meta.cpp）。

4.1 镜像官方 V2 API 的四个入口

该适配层刻意对齐官方aclnnGroupedMatmulSwigluQuantWeightNzV2的调用习惯，提供**两个名字 × 两种重载（stacked /.list）**共四个入口，共享同一份 Host Tiling 与同一个设备 kernel：

torch.ops._C_ascend.grouped_matmul_situ_quant(x, weight(ND), weightScale, weightAssistMatrix?, bias?, xScale, smoothScale?, groupList, dequantMode, dequantDtype, quantMode, groupListType, tuningConfigOptional?, beta, linearBeta) -> (output, outputScale) torch.ops._C_ascend.grouped_matmul_situ_quant_weight_nz(...) # NZ 权重 # 两者各自的 .list 重载接受 per-expert TensorList

4.2 权重分发四形态（源码 §3.2）

• ND stacked：(E, N, K/2)FP4x2 打包连续张量。入口层通过生产级 aclnn 格式转换（npu_format_cast，目标 id 29，见CastNdToNz）转成 NZ 字节流后喂给 kernel；
• NZ stacked：format-29 存储（FRACTAL_NZ_C0_16或FRACTAL_NZ标签），可以是(E,N,K/2)、转置视图(E,K/2,N)或规范 5D[E,K/32,N/16,16,32]。kernel 按字节直读、零转换（这是融合 kernel 的原生形态）；
• ND list：每次调用将 per-expert 张量拼接（CatTensorList，单次at::cat设备算子），再统一转 NZ；
• NZ list：作为 ACL 动态输入传入，kernel 通过地址表直接解引用各专家张量，无 cast、无 cat。

4.3 参数行为矩阵（源码 §3.3）

• bias/smoothScale：传None或空张量即跳过；传入真实值会触发TORCH_CHECK显式报错——宁可报错也绝不静默给出错误结果；
• weightAssistMatrix：接收但忽略（仅打印一次提示），因为 vendoredvf_nz寻址路径不消费 NZ assist matrix；
• tuningConfigOptional：接收但忽略，算子使用自己的 Tiling 策略（baseM=128）；
• 模式枚举只接受一组组合：dequantMode=1（MX 联合 data+scale 反量化）、dequantDtype=0（BF16 中间 GEMM 结果）、quantMode=1（动态 MX 量化，FP8-E4M3 输出 + E8M0 scale），其他取值一律报错；
• groupList：设备侧 int64(E,)张量，type0 为累加和、type1 为计数，由 kernel 内部 preamble 解码（图安全）。

4.4 输出形状与空 token 保护

输出张量由入口层直接分配：y = (M_cap, N/2)FP8-E4M3，yScale = (M_cap, ceil(N/2/64), 2)E8M0，与 golden 拆分链形状完全一致。另外，代码显式处理了EP 秩可以合法地收到零路由 token的情况：当M_cap == 0时直接返回空输出，不启动 kernel，但元数据校验仍然执行（RunV2Core中的空分支）。

五、设备 Kernel：图捕获友好的静态网格与内核内剪枝

设备侧入口 op_kernel/grouped_matmul_situ_quant.cpp 非常简洁：它只负责注册 tiling 结构（REGISTER_TILING_DEFAULT），真正的逻辑全部在GmmSituController（op_kernel/gmsq_vcv_controller.h）中。

5.1 设备 group_list 与内核内 preamble

与 eager 入口（Host 在 tiling blob 中预剪枝零 token 组的 X14 活跃表）不同，图捕获入口要求group_list 保持为设备张量，由 kernel 内部的 preamble 段（per-core 对 E≤32 个条目做 running cumsum）构建等价的分组表——刻意不做独立的 cumsum 设备 kernel，从而守住「融合 launch」的契约红线。每条专家按形状推导静态 stride（devPerGroupW_、devKScaleRow_等），不依赖 group_list 内容。

5.2 静态网格与重放语义

• 网格 = 全量 AIC 核心数：启动配置对捕获图的每次重放完全一致，不依赖 Host 端活跃专家表；
• 无基本块的 core 在 preamble 之后直接退出（约 +1%~3% 的低活跃度成本桶）；
• 重放语义：kernel 每次 launch 都重新读取 group_list 设备缓冲，因此路由器输出被就地更新后，图的每次重放都能感知新路由，不存在冻结的 Host 副本。

这正好解释了端到端测试的用例设计：测试用NPUGraph捕获一次融合调用后，反复修改groups张量的内容并graph.replay()，验证跨小 M 边界与双角色边界的路由变化（详见第六节）。

5.3 SplitNByMultiCore 与块切分

ProcessGroup对每个专家组按baseM=128切分 M 方向，再经SplitNByMultiCore把 gate/up 两半输出按核心轮转分配；其中nL1Size = basicBlockSize * 2（覆盖两半），nOffset语义为前 N/2（gate）输出列偏移。宽块（nL1Size > 128且mL1Size > 64）时按每子块 64 个 gate 列拆分为连续 128 宽子块，kbL1Size依据生产动态规则（小 M 窄 N 走 512 深度，否则 256）逐子块重算。

5.4 双角色（dual-role）模式的条件化启用

控制器按「每 L1 块 M 行数」maxL1M条件化启用供数专核 + epilogue 专核的双角色拆分：阈值 34 依据 M 分布间隙（maxL1M ∈ [33,35]观测中点）选定，且额外要求maxL1M × maxNL1Size ≤ 16384（64KB relay 窗 / 4B），否则整体回退到交错路径。该判断由 tiling 头与组表内容在三个核（AIC/AIV sub0/AIV sub1）上独立推导，结果一致、全程恒定，无需额外同步。

六、SiTU 激活与 MX 量化 Epilogue 的数值实现

epilogue 实现在 op_kernel/situ_epilogue.h，其文件头明确标注：这是对生产级situ_mx_quantarch35 regbase 实现的逐字（verbatim）改编（源自 csrc/moe/situ_mx_quant），要求与生产 kernel位级一致（bit-exact），因此禁止改动指令顺序。

6.1 SiTU 激活的数学形态

核心公式（ComputeVfSitu）：

situ(gate, up) = beta * tanh(gate / beta) * sigmoid(gate) * up （可选：up 再施加 linear_beta * tanh(up / linear_beta)）

实现细节值得注意：

• tanh采用双路径：|x| < 0.6走 9 次多项式 FMA Horner（系数c1..c4与参考 tanh.h 完全一致），|x| >= 0.6走 sigmoid 分解2/(1+exp(-2x)) - 1，符号天然保持；
• 所有运算在 FP32 域完成，最终Cast<BF16>输出，保证与生产 kernel 的舍入行为一致；
• 有符号性边界处理：MAX_EXP_FOR_BF16、FP8_E4M3FN_MAX/MIN_VALUE = ±448.0等常量被用于量化前 clamp（有序比较不改变上游 NaN）。

6.2 动态 MX 量化三步走

量化辅助函数改编自swiglu_mx_quant_common.h的 BF16-only 路径：

1. ComputeVfMaxExpVfLast：按 32 元素块提取 BF16 指数并求块内最大值；
2. ComputeScaleLast：按 OCP 算法由最大指数计算 E8M0 scale 及其倒数（含 NaN/零/特殊指数掩码处理）；
3. ComputeDataF8Last：数据乘以倒数 scale，clamp 到 E4M3FN 有限域后Cast<FP8>，通过寄存器布局交错（CAST_32_TO_80/81/82/83）与字节加法完成打包输出。

七、构建、调用与 MoE 量化方法的自动融合

7.1 构建与打包

kernel 随vllm_ascend_C一起构建，vLLM-Ascend 在SOC_VERSION=ascend950*环境下正常安装即可打包进算子库（无需额外编译步骤）。平台注册见 op_host/grouped_matmul_situ_quant_def.cpp 中AICore().AddConfig("ascend950", ...)。

7.2 调用入口

A5 上现有的DeviceOperator.npu_grouped_matmul_situ_quant入口调用torch.ops._C_ascend.grouped_matmul_situ_quant_weight_nz；TensorList 分发使用.list后缀，ND 条目保留参考名。一个最小 NZ-list 调用示意（与测试用例同形）：

torch.ops._C_ascend.grouped_matmul_situ_quant_weight_nz.list( x, # (M_cap, K) fp8_e4m3fn weights, # per-expert FRACTAL_NZ FP4 权重列表 scales, # per-expert E8M0 weightScale 列表 None, # weightAssistMatrix（忽略） None, # bias（不支持） x_scale, # (M_cap, K/32) fp8_e8m0fnu None, # smoothScale（不支持） groups, # (E,) int64 设备张量 dequant_mode=1, # MX 反量化 dequant_dtype=0, # BF16 中间结果 quant_mode=1, # 动态 MX 量化 group_list_type=0, # 0=cumsum / 1=counts tuning_config=None, beta=4.0, linear_beta=25.0, )

7.3 W4A8 MoE 方法的自动选择

W4A8 MXFP MoE 量化方法在满足以下条件时会自动选择融合路径：SiTU 激活、group size 32、group-list 类型 0（cumulative）或 1（counts）、linear_beta为正。此时 W13 以原生 FP4 元数据加载，scale 使用 E8M0 元数据与 N-major 视图，不复制存储；W2 与既有 GMM2 路径保持不变。

八、端到端验证：小专家数 + 图捕获下的动态路由

Ascend 950 的端到端用例位于 tests/e2e/nightly/single_node/ops/singlecard_ops/test_gmm_situ_small_experts.py，它同时验证了两件关键能力：

1. 与 golden 拆分链的位级一致性：测试构造参考实现npu_grouped_matmul(...)+situ_mx_quant(...)，用torch.testing.assert_close(..., rtol=0, atol=0)断言融合结果与其逐字节相等；
2. 图捕获下的动态路由正确性：在torch.npu.NPUGraph()中捕获一次融合调用后，用 6 组路由模式（均衡、仅 8 个专家活跃、单专家独享、全量路由到单专家等）反复set_groups(...)并graph.replay()，每次重放后仍与 reference 严格一致——这直接验证了「设备 group_list + 内核内 preamble + 静态网格」机制对路由变化的响应能力。

测试还参数化覆盖了experts ∈ {14, 15}与group_list_type ∈ {0, 1}，并在非 Ascend950 设备上自动 skip，进一步印证了该算子「仅 A5 支持」的约束。

九、设计要点速查

• 融合边界：一次 launch 完成「MXFP8×MXFP4 分组矩阵乘 → SiTU → 动态 MX 量化」，消除中间 BF16 结果的写读；
• 图捕获友好：group_list 设备常驻、静态满核网格、内核内 preamble 建表，重放时自动感知路由变化；
• 数值位级一致：epilogue 逐字复用生产situ_mx_quant实现，禁止优化指令顺序；
• 显式拒绝而非静默错误：不支持的bias/smoothScale/模式枚举直接TORCH_CHECK报错；
• shape 契约：K % 64 == 0、N % 2 == 0、N/2 % 64 == 0；输出y=(M, N/2)、yScale=(M, ceil(N/2/64), 2)。

十、继续深入阅读

• 算子总览：csrc/moe/grouped_matmul_situ_quant/README.md
• aclnn 接口规格：csrc/moe/grouped_matmul_situ_quant/docs/aclnnGroupedMatmulSituQuant.md
• 算子定义与 InferShape：op_host/grouped_matmul_situ_quant_def.cpp、op_host/grouped_matmul_situ_quant_infershape.cpp
• Host Tiling：op_host/grouped_matmul_situ_quant_tiling.cpp
• 设备 Kernel：op_kernel/grouped_matmul_situ_quant.cpp、op_kernel/gmsq_vcv_controller.h
• SiTU/MXQuant epilogue：op_kernel/situ_epilogue.h
• Torch 适配与注册：grouped_matmul_situ_quant_torch_adpt.h、csrc/torch_binding.cpp、csrc/torch_binding_meta.cpp
• 端到端测试：tests/e2e/nightly/single_node/ops/singlecard_ops/test_gmm_situ_small_experts.py

• 人工智能
• 大模型
• 模型推理服务
• Ascend
• CANN