FfnWorkerBatching TTK 测试描述文档

1. 文件总览

ffn/ffn_worker_batching/tests/
├── assets/
│   ├── spec.py                          # TestSpec 适配器（精度容差 + 自定义比较逻辑）
│   └── impl/
│       ├── inputs.py                    # 输入数据生成（customize_inputs）
│       └── golden.py                    # CPU 参考实现（golden）
└── ut/op_api/
    ├── test_aclnn_ffn_worker_batching.csv      # 基础用例集（25 个）
    └── test_aclnn_ffn_worker_batching_full.csv # 泛化用例集（200 个）

---

2. spec.py — TestSpec 适配器

文件路径：tests/assets/spec.py

作用：将 golden、customize_inputs、精度容差和自定义比较逻辑统一注册到一个 TestSpec 类中，供 TTK 框架在 ACLNN 模式下加载使用。

关键设计：

2.1 模块加载与缓存共享

golden_module = _load_impl("golden")
inputs_module = _load_impl("inputs")
golden_module._DATA_CACHE = inputs_module._DATA_CACHE

inputs.py 在执行 customize_inputs 时会将生成的 CPU 测试数据缓存到 _DATA_CACHE 字典（key=testcase_name）。golden.py 的 golden 函数需要这些 CPU 数据来计算参考结果（因为 schedule_context tensor 中存的是 device 偏移量，CPU 侧无法解析）。通过 golden_module._DATA_CACHE = inputs_module._DATA_CACHE 将两个模块的缓存打通。

2.2 精度容差定义

tolerance = {
    "float16":  {"standard": "binary_equal"},   # FP16: 二进制精确匹配
    "bfloat16": {"standard": "stat_rel_err"},   # BF16: 统计相对误差（位模式可能有差异）
    "int8":     {"standard": "binary_equal"},   # INT8: 二进制精确匹配
    "float32":  {"standard": "binary_equal"},   # float32: 二进制精确匹配
    "int32":    {"standard": "binary_equal"},   # int32: 二进制精确匹配
    "int64":    {"standard": "binary_equal"},   # int64: 二进制精确匹配
}

该算子是整数索引排序 + 原样字节搬运，无浮点运算，因此除 BF16 外均使用二进制精确匹配。BF16 使用统计相对误差是因为搬运过程中可能存在位模式差异。

2.3 自定义 compare 函数

kernel 只写入前 actual_token_num 个有效 token 的输出，其余位置保留 TTK 的初始化值（1）。golden 用 0 填充未使用位置。因此需要自定义比较逻辑，只比较有效区域：

输出索引	名称	比较范围
0	y	前 actual_token_num 行（每行 H 个元素）
1	group_list	第一个 [0, 0] 终止符之前的行
2	session_ids	前 actual_token_num 个元素
3	micro_batch_ids	前 actual_token_num 个元素
4	token_ids	前 actual_token_num 个元素
5	expert_offsets	前 actual_token_num 个元素
6	dynamic_scale	全量比较（非 INT8 时 golden 填 1，与 TTK 初始化值一致）
7	actual_token_num	全量比较（1 个 int64 元素）

BF16 tensor 的特殊处理：to_np 辅助函数将 BF16 tensor 先转 float32 再转 numpy，避免 np.asarray 不支持 BF16 的 TypeError。

2.4 注册

__spec__ = {"aclnnFfnWorkerBatching": "AclnnFfnWorkerBatchingTestSpec"}

key 必须与 CSV 中的 api_name 完全一致。

---

3. inputs.py — 输入数据生成

文件路径：tests/assets/impl/inputs.py

作用：为每个测试用例生成确定性输入数据，构建 schedule_context 结构体，将附属数据追加到 schedule_context storage 后面（相对偏移模式），并缓存 CPU 数据供 golden 使用。

3.1 确定性随机数据生成

def _seed_from_name(name):
    h = 0
    for ch in name:
        h = (h * 131 + ord(ch)) & 0x7FFFFFFF
    return h

从用例名生成种子，确保同一用例每次运行数据一致。数据包括：

• expert_ids：NORM 模式 shape=(A, BS, K)，RECV 模式 shape=(A, M, BS, K)
• token_data：shape=(A, M, BS, K, H)，dtype 取决于 tokenDtype
• token_scales：仅 INT8 量化时存在，shape=(A, M, BS, K)，float32
• session_ids：shape=(A,)，值为 0..A-1
• micro_batch_ids：shape=(A,)，全 0

3.2 特殊用例数据变异

• masked 用例：用例名含 masked 时，~15% 的 expert_id 被设为 1000000（_MASK_VALUE），验证 mask 过滤逻辑
• single_expert 用例：用例名含 single_expert 时，所有 expert_id 设为 0，验证退化场景

3.3 数据布局与 HBM 可见性方案

问题：TTK 框架在 customize_inputs 阶段调用 aclrtMalloc 分配的 HBM，在后续 create_context 后对 kernel 不可见（返回 107000 错误）。

方案：把 token_data / expert_ids / session_ids / micro_batch_ids 等附属数据追加到 schedule_context tensor 的 storage 后面（保持 view shape (1024,) 不变，只扩展 untyped_storage）。TTK 的 copy_torch_tensor_to_hbm 用 storage().nbytes() 拷贝完整 storage 到 device。

schedule_context 中存相对偏移量（从 schedule_context 起始到数据起始的字节偏移），offset 640 写 TEST_MAGIC = 0x54455354 标志位。kernel 读取此标志后，将 buffer 指针字段解释为 schedule_context 地址 + 偏移。

3.4 BF16 数据转换

_build_token_data_bytes 对 tokenDtype=1（BF16）先把 float32 转 torch.bfloat16 再取字节，确保 HBM 中是 BF16 位模式而非 float32。

3.5 INT8 量化 scale 打包

对 tokenDtype=2，把每个 token 的 int8 数据（H 字节）和 float32 scale（4 字节）连续排布为 [A, M, BS, K, H+4] 的 uint8 数组。

3.6 token_info 构建（RECV 模式）

_build_token_info_buf 生成 [A, M, F] 的 int32 数组，F = 2 + BSK。每个 [a, m] 的前两个元素是 flag=1（数据就绪）和 layer_id=0，后面是 BSK 个 expert_id。

3.7 storage 扩展

用 scheduleContext.set_(new_storage, offset, shape, stride) 替换底层 storage，然后用 untyped_storage()[:total_size].copy_(...) 拷贝数据。

---

4. golden.py — CPU 参考实现

文件路径：tests/assets/impl/golden.py

作用：模拟 kernel 的排序 + gather + group_list 逻辑，作为精度比对的参考实现。

4.1 数据获取

通过 _get_cached_data(testcase_name) 从 _DATA_CACHE 获取 inputs.py 缓存的 CPU 数据。因为 schedule_context 中的偏移量在 CPU 侧无意义。

4.2 排序逻辑 (_sort_expert_ids)

1. 把 >= EXPERT_MASK_VALUE (1000000) 的 expert_id 映射为 int32.max
2. 稳定排序（np.argsort(kind='stable')），masked token 排到最后
3. 返回排序顺序、有效 token 数、排序后的有效 expert_id

4.3 group_list 生成 (_generate_group_list)

1. 遍历排序后的 expert_id，检测跳变点生成 [expert_id, token_count] 行
2. 最后一个有效专家后写一行 [0, 0] 作为终止符
3. 未使用行初始化为 [0, 0]

4.4 gather 逻辑

按排序顺序遍历每个有效 token：

1. 从 sorted_order[i] 获取原始索引 gidx
2. 分解为 a_idx = gidx // (BS*K), bs_idx = (gidx % (BS*K)) // K, k_idx = gidx % K
3. NORM 模式：从 session_ids_buf[a_idx] 和 micro_batch_ids_buf[a_idx] 查表
4. RECV 模式：session = a_idx, micro_batch = cur_micro_batch_id
5. 从 token_data[a_idx, mb_idx, bs_idx, k_idx, :] 提取 hidden states
6. INT8 量化时额外输出 token_scales[a_idx, mb_idx, bs_idx, k_idx]

4.5 dynamic_scale 初始化

用 np.ones 而非 np.zeros。因为 tokenDtype≠2 时 kernel 不写 dynamic_scale 输出，TTK 把纯输出 tensor 初始化为 1，golden 需要匹配这个初始值。

4.6 输出

返回 8 个 torch tensor，顺序与 CSV 的 output_tensor_indexes 一致：y, group_list, session_ids, micro_batch_ids, token_ids, expert_offsets, dynamic_scale, actual_token_num。

---

5. test_aclnn_ffn_worker_batching.csv — 基础用例集

文件路径：tests/ut/op_api/test_aclnn_ffn_worker_batching.csv

作用：25 个基础测试用例，覆盖算子的核心功能和主要参数组合。

CSV 字段说明

字段	含义	示例
testcase_name	用例名（全局唯一，决定随机种子）	aclnnFfnWorkerBatching_fp16_basic
api_name	ACLNN API 名（与 spec.py 的 __spec__ key 一致）	aclnnFfnWorkerBatching
tensor_view_shapes	9 个 tensor 的 view shape	((1024,),(72,128),(8,2),(72,),...)
tensor_dtypes	9 个 tensor 的 dtype	('int8','float16','int64','int32',...)
attributes	算子属性 JSON	{'expertNum': 8, 'maxOutShape': [1,8,9,128], ...}
output_tensor_indexes	哪些 tensor 是输出	(1,2,3,4,5,6,7,8)

9 个 tensor 的索引

索引	名称	shape	dtype	说明
0	schedule_context	(1024,)	int8	输入：调度上下文
1	y	(Y, H)	float16/bfloat16/int8	输出：gather 后的 token 数据
2	group_list	(E, 2)	int64	输出：专家→token数映射
3	session_ids	(Y,)	int32	输出：per-token session ID
4	micro_batch_ids	(Y,)	int32	输出：per-token micro-batch ID
5	token_ids	(Y,)	int32	输出：per-token batch 内索引
6	expert_offsets	(Y,)	int32	输出：per-token expert 内偏移
7	dynamic_scale	(Y,)	float32	输出：per-token scale（INT8 量化）
8	actual_token_num	(1,)	int64	输出：有效 token 数

其中 Y = A × BS × K，E = expertNum。

25 个用例分类

索引	用例名	测试维度
0	fp16_basic	FP16 基础场景：A=1, BS=8, K=9, H=128, E=8
1	fp16_h256	H 维度：H=256
2	fp16_h512	H 维度 + BS 维度：H=512, BS=16
3	fp16_h4096	超大 H：H=4096
4	bf16_basic	BF16 dtype
5	bf16_h512	BF16 + 大 shape
6	int8_quant_basic	INT8 动态量化
7	int8_quant_h256	INT8 + H=256
8	multi_session_a2	多 session：A=2
9	multi_session_a4	多 session：A=4
10	large_bs	大 batch：BS=32
11	expert16	16 专家：K=17
12	expert32	32 专家 + 大 shape
13	expert64	64 专家
14	fp16_masked	~15% token 被 mask
15	fp16_single_expert	所有 token 归同一专家
16	recv_basic	RECV 模式基础
17	recv_a2	RECV 多 session
18	large_combined	大组合：A=4, BS=16, K=9, H=256, E=16
19	bf16_large	BF16 大场景
20	fp16_h4096_large_bs	H=4096 + BS=32
21	fp16_h4096_a4_e32	H=4096 + A=4 + E=32
22	bf16_h4096_large	BF16 + H=4096
23	int8_h4096_large	INT8 + H=4096
24	recv_h4096	RECV + H=4096

---

6. test_aclnn_ffn_worker_batching_full.csv — 泛化用例集

文件路径：tests/ut/op_api/test_aclnn_ffn_worker_batching_full.csv

作用：200 个泛化测试用例，在基础用例集之上大幅扩展参数覆盖范围，用于全面验证算子在各 shape 组合下的正确性。

参数覆盖范围

参数	取值范围	数量
A (session 数)	1, 2, 4, 8, 16, 32, 64, 128, 256	9
BS (batch size)	1, 2, 4, 8, 16, 32, 64, 128	8
K (selected+1)	2, 3, 5, 9, 17, 33, 64	7
H (hidden size)	1, 64, 128, 256, 512, 1024, 2048, 4096, 8192	9
expertNum	4, 8, 16, 32, 64, 128, 256	7
token_dtype	0 (FP16), 1 (BF16), 2 (INT8)	3
need_schedule	0 (NORM), 1 (RECV)	2

用例分布

分类	数量
按 dtype	FP16: 109, BF16: 44, INT8: 47
按模式	NORM: 156, RECV: 44

用例组成（8 个部分）

部分	覆盖内容	用例数
1. FP16 NORM H×A 扫描	H ∈ {64..8192} × A ∈ {1,2,4,8} + BS/K/expertNum 扫描	~51
2. BF16 NORM H×A 扫描	H ∈ {64..8192} × A ∈ {1,2,4} + BS/K 扫描	~32
3. INT8 NORM H×A 扫描	H ∈ {64..8192} × A ∈ {1,2,4} + BS/K/expertNum 扫描	~35
4. FP16 RECV H×A 扫描	H ∈ {64..8192} × A ∈ {1,2,4}	~24
5. BF16+INT8 RECV	H ∈ {128,512,4096} × A ∈ {1,2}	~12
6. 大组合 stress	多维度大 shape 组合	~20
7. 边界用例	最小 shape、masked、single_expert、layerNum 变化	~26
8. 补充用例	补齐至 200 个的额外 shape 组合	动态

边界用例说明

用例名	场景
fp16_min_a1_bs1_k2_h1_e4_norm	最小合法 shape：A=1,BS=1,K=2,H=1,E=4
fp16_min_a1_bs1_k2_h1_e4_recv	最小 shape + RECV 模式
fp16_masked_h128	FP16 + ~15% mask + H=128
fp16_masked_h4096	FP16 + mask + 超大 H
bf16_masked_h256	BF16 + mask
int8_masked_h128	INT8 + mask
fp16_masked_a4_h512	mask + 多 session
fp16_single_expert_h128	所有 expert_id=0（退化场景）
fp16_single_expert_h4096	退化 + 超大 H
bf16_single_expert_h256	BF16 退化
int8_single_expert_h128	INT8 退化
fp16_layer0_h128	layerNum=0
fp16_layer2_h128	layerNum=2
fp16_layer0_h4096	layerNum=0 + 超大 H
fp16_a128_bs1_k9_h128_norm	高 A=128（多 session 极限）
fp16_a256_bs1_k2_h128_norm	A=256 + K=2
fp16_a1_bs8_k2_h4096_norm	K=2（topK=1 最小）
int8_k2_h512_norm	INT8 + K=2

---

7. TTK 执行流程

TTK 框架读取 CSV
  ↓
对每行用例：
  1. 加载 spec.py → AclnnFfnWorkerBatchingTestSpec
  2. 调用 customize_inputs(inputs.py) → 修改 schedule_context tensor + 缓存 CPU 数据
  3. TTK 将 schedule_context 拷贝到 HBM
  4. 调用 aclnnFfnWorkerBatchingGetWorkspaceSize → 执行 tiling
  5. 调用 aclnnFfnWorkerBatching → 执行 kernel
  6. 调用 golden(golden.py) → 生成 CPU 参考结果
  7. 调用 compare(spec.py) → 自定义精度比较
  8. 输出 precision_status: PASS/FAIL

执行命令

# 基础用例（25 个）
cd /workspace/ops-test-kit
python3 -m ttk aclnn \
  -i .../test_aclnn_ffn_worker_batching.csv \
  -o /tmp/result.csv \
  --plugin .../spec.py \
  --pc 1 --ti=0-25

# 泛化用例（200 个）
python3 -m ttk aclnn \
  -i .../test_aclnn_ffn_worker_batching_full.csv \
  -o /tmp/result_full.csv \
  --plugin .../spec.py \
  --pc 1 --ti=0-199

结果检查

python3 -c "
import csv
p=0;f=0
with open('/tmp/result.csv') as fh:
    for row in csv.DictReader(fh):
        if row.get('precision_status')=='PASS': p+=1
        else: f+=1
print(f'Total:{p+f} PASS:{p} FAIL:{f}')
"

---

8. 测试验证结果

用例集	用例数	PASS	FAIL
基础用例集	25	25	0
泛化用例集	200	200	0