Matmul Layout 与数据格式开发指南

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills

本文覆盖 Matmul 算子所有 Layout 与数据格式相关的开发指导，包括 ND/DN/NZ/ZN 格式定义、 数据生成流程、LayoutPtn 选型、kernel 适配点和排障。

API 参考：LayoutPattern 谱系、MakeFrameLayout 签名、Routing 表等 API 级内容详见tensor-api-reference.md。

1. 格式定义

NZ 是一种分形存储格式。对原始 tensor(dim0, dim1)ND 排布，NZ 格式的物理排列为(dim1/C0, dim0/16, 16, C0)， 其中C0 = 32 / sizeof(dtype)（fp16/bf16: C0=16, fp8/int8: C0=32, fp4: C0=64）。 非对齐时先补 0 到 16 对齐。

记矩阵计算的三个维度分别为 M、K、N（即 A[M,K] × B[K,N] = C[M,N]），则 A、B 矩阵的各种格式情况对应的 GM 数据排布及 LayoutPtn 如下。

A 矩阵

B 矩阵

2. 数据生成流程

随机数据生成固定以数学描述 A(M,K)、B(K,N) 生成，golden 计算固定为 A @ B。 数据生成后分两路：一路计算 golden 保存为 bin，另一路按 kernel 入参要求对 A、B 各自独立做转换后保存为 bin。

转换函数

def to_nz_format(data, c0): """ND (dim0, dim1) → NZ 分形 (dim1/C0, dim0/16, 16, C0) c0 必须按 dtype 显式传入：fp16/bf16=16, int8/fp8=32, fp4=64 """ dim0, dim1 = data.shape dim0_pad = ((dim0 + 15) // 16) * 16 dim1_pad = ((dim1 + c0 - 1) // c0) * c0 padded = torch.zeros((dim0_pad, dim1_pad), dtype=data.dtype) padded[:dim0, :dim1] = data b_4d = padded.reshape(dim0_pad // 16, 16, dim1_pad // c0, c0) return b_4d.permute(2, 0, 1, 3).contiguous()

关键：permute(2, 0, 1, 3)产生物理排列(dim1/C0, dim0/16, 16, C0)。 常见错误是写成permute(0, 2, 1, 3)产出(dim0/16, dim1/C0, 16, C0)， 与 tensor_api NZ layout 的 stride 不匹配，导致全 FAIL。

注意：c0参数必须按 dtype 显式传入。不同 dtype 的 C0 不同（fp16/bf16=16, int8/fp8=32）。

gen_data 范例

def gen_data(m, k, n, dtype=torch.int8, trans_a=False, trans_b=False, a_format="nd", b_format="nd"): C0 = 32 // torch.tensor([], dtype=dtype).element_size() # 动态计算 C0 # 生成原始数据（固定 shape） if dtype.is_floating_point: A = torch.randn(m, k, dtype=dtype) B = torch.randn(k, n, dtype=dtype) else: A = torch.randint(-128, 128, (m, k), dtype=dtype) B = torch.randint(-128, 128, (k, n), dtype=dtype) # golden 计算（固定 A @ B） golden = (A.float() @ B.float()).to(torch.bfloat16) # 按 kernel 需求转换 A（先 transpose，再 to_nz_format） if trans_a: A_for_kernel = A.T.contiguous() else: A_for_kernel = A if a_format == "nz": A_bin = to_nz_format(A_for_kernel, C0) else: A_bin = A_for_kernel # 按 kernel 需求转换 B（先 transpose，再 to_nz_format） if trans_b: B_for_kernel = B.T.contiguous() else: B_for_kernel = B if b_format == "nz": B_bin = to_nz_format(B_for_kernel, C0) else: B_bin = B_for_kernel

3. LayoutPtn 选择

Launcher 直接传 tensor_api pattern 作为模板参数：

TagToTrans<Pattern>在layout_utils.h派生 transA/transB：

常见错误：新增 transA=true 但 launcher 里 layoutA 仍硬编码 NDExtLayoutPtn → 编译过但 ≈100% mismatch（K 维和 M 维错位）。

4. L1 layout 自动选择

使用L1LayoutHelper<LayoutPtn, Type, TransVal>统一处理：

• ND/DN 输入：L1 按 trans 标志选 NZ/ZN（走硬件 ND→NZ 格式转换）
• NZ/ZN 输入：L1 与 GM pattern 一致（走 NZ→NZ / ZN→ZN 块拷贝，省掉格式转换带宽）

using MakeLayoutAL1 = typename L1LayoutHelper<LayoutA, AType, transA>::type; using MakeLayoutBL1 = typename L1LayoutHelper<LayoutB, BType, transB>::type;

L1→L0A/L0B 的 routing 由 tensor_api 根据 DstPattern（L0A=NZ, L0B=ZN）和 SrcPattern（L1 的 pattern）自动派发NORMAL 或 TRANS 模式，无需手动添加LoadDataTrait{transposed=true}。

5. GM 端 layout 构造

GM 端 NZ layout 构造必须使用正确的 C0：

// ⚠️ 关键：FrameLayoutFormat 默认 C0=16（基于 uint16_t），int8/fp8 需要 C0=32 static constexpr uint64_t A_C0 = 32 / sizeof(AType); static constexpr uint64_t B_C0 = 32 / sizeof(BType); using MakeLayoutA = AscendC::Te::FrameLayoutFormat<LayoutA, AscendC::Std::Int<A_C0>>; using MakeLayoutB = AscendC::Te::FrameLayoutFormat<LayoutB, AscendC::Std::Int<B_C0>>;

⚠️ 致命陷阱：FrameLayoutFormat<NZLayoutPtn>默认使用LayoutTraitDefault<>，其 C0 = 32/sizeof(uint16_t) =16。 当 AType/BType 为 int8 时，正确的 C0 应为 32/sizeof(int8_t) =32。 C0 错误会导致 GM 端 NZ layout 的 Shape[Column][1]（列块数）和 Stride[Column][1]（列块间距）全部错误，CopyGmToCbufAlignV2NZ读取错误的 blockCount/blockLen/srcStride，多 tile 场景全 FAIL。 L1 端的L1LayoutHelper已正确按 dtype 计算 C0，此问题仅影响 GM 端 layout 构造。

6. Host 侧 size 计算

ND 格式：size = dim0 * dim1 * sizeof(dtype)

NZ 格式：buffer size 按物理维度计算：

// CalcNzSize(dim0, dim1, c0): 分形物理排列的字节数 uint64_t dim0Blocks = (dim0 + 15) / 16; uint64_t dim1Blocks = (dim1 + c0 - 1) / c0; return dim1Blocks * dim0Blocks * 16 * c0 * sizeof(dtype);

按转换后的物理数据 shape 传入参数：

A 矩阵（golden 固定有 A[M,K]）

B 矩阵（golden 固定有 B[K,N]）

7. Kernel 端适配点（5 个）

适配点 1 —layout_utils.h：新增 NZ 的TagToTrans特化（NZLayoutPtn和ZNLayoutPtn各一个）+IsNzOrZn检测 +L1LayoutHelper辅助模板。

适配点 2 —matmul_kernel.h/matmul_kernel_fused.h：GM 端 layout 构造必须使用正确的 C0（详见 §5）。

适配点 3 —matmul_block_mmad.h：MakeLayoutAL1/MakeLayoutBL1改用L1LayoutHelper（详见 §4）。

适配点 4 —matmul_kernel.h：Slice 顺序——kernel 层做 M/N-slice（保留 fullK stride），block 层只做 K-slice：

// Kernel 层：N-slice（保留 fullK stride，NZ column stride 依赖 fullK） auto gmBlockB = gmB.Slice(Coord(0, nPos), Shape(fullK, tileN)); // Block 层：K-slice（gmBlockB 已经是 N-tile 大小） auto gmTileB = gmB.Slice(Coord(kL1Offset, 0), Shape(curKL1, curN));

关键：NZ 的 column stride =C0 × ceil_align(K, 16)，依赖 fullK。 Slice 保留父 tensor 的 stride 不重算。如果 block 层同时切 K+N，CopyGmToCbufAlignV2NZ用 sliced K 计算smallFractalSize，但 stride 基于 fullK → 地址错位。

适配点 5 — Launcher：按 transA/transB + format 分发 LayoutPtn：

8. 排障速查

9. Tiling 对齐约束（NZ 场景）

NZ 格式的物理排列为(dim1/C0, dim0/16, 16, C0)，其中：

• 外轴 dim0：需要 16 元素对齐
• 内轴 dim1：需要 C0 元素对齐（C0 = 32 / sizeof(dtype)）

A 矩阵（NZ 格式）

B 矩阵（NZ 格式）

各轴约束

Tiling 引擎代码：

constexpr uint64_t C0 = 32 / sizeof(AType); runInfo_.baseM = Align(runInfo_.baseM, args_.isANz && args_.isATrans ? C0 : static_cast<uint64_t>(16)); runInfo_.baseN = Align(runInfo_.baseN, args_.isBNz && !args_.isBTrans ? C0 : static_cast<uint64_t>(16)); runInfo_.baseK = Align(runInfo_.baseK, (args_.isANz && !args_.isATrans) || (args_.isBNz && args_.isBTrans) ? C0 : static_cast<uint64_t>(16));

常见错误：用条件表达式（如if (isBTrans || !isATrans)）控制 baseN 的 C0 对齐，在 transA=true+transB=false 时条件为 false，baseN 未对齐。此时 B 为 NZ of (K,N)，N 是内轴需要 C0 对齐，未对齐导致CopyGmToCbufAlignV2NZblock copy 失败。对齐 shape（如 N=256）不受影响因为 N 本身已是 C0 的倍数。

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills