Perf moe input generation

[feixi139]

改动一：向量化 rowmap 输入生成

涉及文件：

• tester/input_generation/generation_rules.py
• generate_rowmap_input_value

zipped_expertwise_rowmap 的构造原先在 (token, expert) 上开双层 Python 循环逐格赋值，对每个格子单独 nonzero 判断该 token 是否路由到该 expert，再手工维护 num_experts 个计数器发号。迭代次数为 seqlen × num_experts，本文件最大 case 达 3000 万次。

改为两步向量化：

• present 矩阵按 topk 列 scatter 标记 (token, expert) 是否有效，循环次数为 topk，与 seqlen 无关；
• expert 内的流水号等于 present 的列向 exclusive cumsum（cumsum(present, axis=0) - 1），等价于原来“按 row 升序、每遇到一个有效格子自增一次”的语义。

索引用 rule.ops.nonzero(column >= 0)[0] 而不是布尔掩码，避免依赖 torch/paddle backend 的布尔高级索引；全部走 rule.ops.*，三个 backend 通用。

最大 Case（seqlen=3759104, num_experts=8, topk=4）该步骤	改前	改后
NumPy Backend	92.51 s	1.17 s
Paddle GPU Backend	2095 us/行，外推 2.2 h	0.093 s

GPU 那一栏是这个改动的关键：

旧实现每次内层迭代都有一次 nonzero + 一次取值判空导致的 device sync + 一次 setitem，实测 2095 us/行。

这是 --use_gpu_mode=True 之前对 1M 级 moe_unpermute case 不可用的直接原因（主机侧构造阶段跑不完 / 被 OOM kill）。

2.2 h 是按前 300 行线性外推，未实跑全量；NumPy 那一栏两个数字均为全量实测。

一处需要 Reviewer 注意的语义变化

expert_counts 从：

count_nonzero(routemap == expert)

改为：

present.sum(axis=0)

即由“计入一行内重复出现的同一 expert”改为“去重”。

当前 generate_expert_routemap_input_value 填充的值为：

(row + column) % num_experts

且各 column 互不相同，因此一行内不会出现重复 expert，两者对所有现有配置结果相同（已在 topk > num_experts 的边界上验证）。

去重后也与“每个 (row, expert) 只发一个号”的实际写入行为更自洽——原实现在假想的重复场景下算出的 offset 会与实际发号数量不匹配。

---

改动二：大 Tensor symmetric 输入改为分块生成

涉及文件：

• tester/input_generation/value_generators.py
• generate_symmetric_input_value

原表达式：

rng.cast(
    (rng.random(spec.shape) - 0.5) * (2 * max_abs),
    dtype,
)

在大 Tensor 上存在四层内存放大。

以 bfloat16 [1966080, 1024]（20.1 亿元素，目标 3.8 GiB）为例：

步骤	dtype / 大小	耗时
rng.random(shape)，未传 dtype → storage_dtype 为 None → 跳过 astype	float64 15.0 GiB	14.23 s
- 0.5 eager 临时	float64 15.0 GiB	16.59 s
* (2 * max_abs) eager 临时	float64 15.0 GiB	13.54 s
astype（bfloat16 经 _INPUT_INTERMEDIATE_DTYPES 映射为 float32）	float32 7.5 GiB	4.92 s
合计	共写入约 52 GiB 主机内存	49.28 s

全程单线程（MT19937 串行 + NumPy elementwise 无并行）。

backend.py:_direct_float_dtype 的规避快路径只在：

storage_dtype == "float16"

时生效。

由于 bfloat16 映射到 float32，因此永远不会命中该优化路径。

当元素数超过：

1 << 24

时，改走新增的：

_chunked_symmetric_numpy

实现。

具体策略：

• 按首维切块生成；
• 原地执行数值变换；
• 逐块写入最终输出缓冲。

这样可将 float64 临时峰值从：

3 × 15 GiB

降低到：

32 MiB

同时消除 1M 级 case 在非 GPU mode 下的主机 OOM 风险。

随机流保持不变

本次改动不改变随机流。

NumPy backend 使用的是 RandomState（MT19937）顺序流。

random_sample(size)：

• 按 C-order 连续填充；
• 后续调用从同一随机流继续向后推进。

因此：

random((N, 1024))

与：

random((k, 1024))

调用 N/k 次后再拼接，

得到的随机序列完全一致（bit-exact）。

这是有意保持的行为：

derive_input_seed
+ InputConfigRandomState

保证同一条 config 永远生成 bit-identical 输入。

虽然改成：

Generator(PCG64)

并直接生成 float32（实测约 10.1 s）会更快，但会导致：

• --retest 无法复现历史失败；
• accuracy_manual_threshold_config 已调好的阈值失效；
• paddle_bitwise_knows 基线失效；
• 历史问题无法继续 bisect。

本质上等价于一次全局 baseline reset，因此未采用。

Backend 兼容性

启用条件：

getattr(rng, "name", "numpy") == "numpy"

其中：

• InputNumPyRandomState 无 name 属性；
• InputConfigRandomState 无 name 属性；

因此默认按 NumPy 路径处理。

而：

torch backend  -> "torch"
paddle backend -> "paddle"

因此：

torch / paddle backend 行为完全不变。

由于：

value_generators -> backend
backend -> value_generators

会形成循环依赖，

因此这里采用 duck-typing，而非 isinstance 判断。

[cangtianhuang]

LGTM