Perf moe input generation

tester/input_generation/backend.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import math
 import numbers
 from dataclasses import dataclass, field
 from typing import Protocol, runtime_checkable
@@ -17,6 +18,9 @@
 
 # 大 Tensor 才切换目标 storage dtype，避免改变小配置既有随机序列。
 DIRECT_DTYPE_NUMEL_THRESHOLD = 1 << 20
+# NumPy 对称随机值按元素数限制 float64 中间块大小；原生 backend 保持整块生成。
+_SYMMETRIC_CHUNK_MIN_ELEMENTS = 1 << 24
+_SYMMETRIC_CHUNK_ELEMENTS = 1 << 22
 
 
 def _normalize_shape(shape, *, scalar_empty):
@@ -70,6 +74,8 @@ def random(self, shape=None, dtype=None): ...
 
     def uniform(self, low=0.0, high=1.0, shape=None, dtype=None): ...
 
+    def symmetric(self, shape, dtype, max_abs): ...
+
     def randint(self, low, high=None, shape=None, dtype=None): ...
 
     def randn(self, *shape, dtype=None): ...
@@ -187,6 +193,27 @@ def random(self, shape=None, dtype=None):
         storage_dtype = self._storage_dtype(dtype)
         return numpy.asarray(value).astype(storage_dtype) if storage_dtype is not None else value
 
+    def symmetric(self, shape, dtype, max_abs):
+        normalized_shape = _normalize_shape(shape, scalar_empty=False)
+        storage_dtype = self._storage_dtype(dtype)
+        if normalized_shape and math.prod(normalized_shape) >= _SYMMETRIC_CHUNK_MIN_ELEMENTS:
+            value = numpy.empty(normalized_shape, dtype=storage_dtype)
+            flat_value = value.reshape(-1)
+            total_elements = int(flat_value.size)
+            # 一维顺序分块保持与整块 random(shape) 相同的 RNG 消费顺序。
+            for start in range(0, total_elements, _SYMMETRIC_CHUNK_ELEMENTS):
+                end = min(total_elements, start + _SYMMETRIC_CHUNK_ELEMENTS)
+                block = numpy.asarray(self.input_random_state.random(end - start))
+                # block 是本次调用新分配的数组，原地缩放不会改变其他生成结果。
+                block -= 0.5
+                block *= 2 * max_abs
+                flat_value[start:end] = block
+            return value
+        return self.cast(
+            (self.input_random_state.random(normalized_shape) - 0.5) * (2 * max_abs),
+            dtype,
+        )
+
     def uniform(self, low=0.0, high=1.0, shape=None, dtype=None):
         value = self.input_random_state.uniform(low=low, high=high, shape=shape)
         storage_dtype = self._storage_dtype(dtype)
@@ -421,6 +448,10 @@ def uniform(self, low=0.0, high=1.0, shape=None, dtype=None):
         )
         return self.cast(value, dtype) if dtype is not None and not direct_dtype else value
 
+    def symmetric(self, shape, dtype, max_abs):
+        # Torch 保持原生 Tensor 整块生成，避免 Python 分块增加 kernel launch。
+        return self.cast((self.random(shape) - 0.5) * (2 * max_abs), dtype)
+
     def randint(self, low, high=None, shape=None, dtype=None):
         torch = self._torch()
         torch_shape = self._torch_shape(shape)
@@ -779,6 +810,10 @@ def uniform(self, low=0.0, high=1.0, shape=None, dtype=None):
         )
         return self.cast(value, dtype) if dtype is not None and not direct_dtype else value
 
+    def symmetric(self, shape, dtype, max_abs):
+        # Paddle 保持原生 Tensor 整块生成，随机值不会回落到 NumPy 主存。
+        return self.cast((self.random(shape) - 0.5) * (2 * max_abs), dtype)
+
     def randint(self, low, high=None, shape=None, dtype=None):
         if high is None:
             low, high = 0, low

tester/input_generation/backend_runtime.py

@@ -337,7 +337,7 @@ def cached_numpy_output_grad(
             else:
                 # cached complex 复用公共路径，保证实部和虚部独立采样。
                 spec = InputTensorSpec(shape, dtype, None, False, None)
-                value = generate_symmetric_input_value(spec, max_abs, rng)
+                value = generate_symmetric_input_value(spec, max_abs, NumPyInputBackend(rng))
             self._cached_numpy_output_grads[key] = value
         return self._cached_numpy_output_grads[key]
 

tester/input_generation/generation_rules.py

@@ -1191,24 +1191,24 @@ def generate_rowmap_input_value(input_binding):
         rowmap = rule.ops.full(input_binding.shape, -1, dtype="int32")
         if rule.is_tensor_config(routemap_config) and routemap_binding is not None:
             routemap = rule.value(routemap_binding)
-            expert_counts = rule.ops.asarray(
-                [rule.ops.count_nonzero(routemap == expert) for expert in range(num_experts)],
-                dtype="int64",
-            )
+            # present[row, expert] 标记该 token 是否被路由到该 expert；只按 topk 迭代，
+            # 避免 seqlen x num_experts 级别的 python 双循环。
+            present = rule.ops.zeros(input_binding.shape, dtype="int64")
+            for topk_index in range(routemap.shape[1]):
+                column = rule.ops.cast(routemap[:, topk_index], "int64")
+                selected = rule.ops.nonzero(column >= 0)[0]
+                # Paddle 不接受空高级索引，空列保持 present 全 0。
+                if selected.shape[0] == 0:
+                    continue
+                present[selected, column[selected]] = 1
+            expert_counts = rule.ops.cast(rule.ops.sum(present, axis=0), "int64")
             if int(rule.ops.sum(expert_counts)) > unzipped_seqlen:
                 raise ValueError("routemap assignments exceed hidden_states_unzipped capacity")
             expert_offsets = rule.ops.zeros(num_experts, dtype="int64")
             expert_offsets[1:] = rule.ops.cumsum(expert_counts[:-1])
-            expert_counters = rule.ops.zeros(num_experts, dtype="int64")
-            for row_index in range(seqlen):
-                for expert in range(num_experts):
-                    positions = rule.ops.nonzero(routemap[row_index] == expert)[0]
-                    if rule.ops.prod(positions.shape) == 0:
-                        continue
-                    rowmap[row_index, expert] = rule.ops.cast(
-                        expert_offsets[expert] + expert_counters[expert], "int32"
-                    )
-                    expert_counters[expert] += 1
+            # 每个 expert 内按 row 升序自增编号，等价于列向 present 的 exclusive cumsum。
+            ranks = rule.ops.cumsum(present, axis=0) - 1
+            rowmap = rule.ops.cast(rule.ops.where(present > 0, expert_offsets + ranks, -1), "int32")
         return rowmap
 
     def generate_token_prob_input_value(input_binding):

tester/input_generation/value_generators.py

@@ -8,8 +8,8 @@
 
 from .values import InputTensorSpec
 
-# 单值生成器只消费 InputTensorSpec 和 RNG，不读取 API 名称或修改 TensorConfig。
-# `spec` 与 `rng` 是本模块内部的数值计算惯例，完整输入标识由函数名和类型提供。
+# 单值生成器只消费 InputTensorSpec 和 backend/RNG，不读取 API 名称或修改 TensorConfig。
+# `spec` 与 backend/RNG 是本模块内部的数值计算惯例，完整输入标识由函数名和类型提供。
 # 这些中间 dtype 转换要保持固定，才能保证输出字节稳定。
 _INPUT_INTERMEDIATE_DTYPES = {
     "bfloat16": "float32",
@@ -147,27 +147,27 @@ def _complex_value(dtype, shape, rng, **kwargs):
 def generate_symmetric_input_value(
     spec: InputTensorSpec,
     max_abs,
-    rng=INPUT_NUMPY_RANDOM_STATE,
+    backend,
 ) -> object:
     """生成实部和虚部分量均位于对称区间的数值。"""
     dtype = resolve_input_dtype(spec.dtype)
     if dtype.startswith("complex"):
         # max_abs 约束每个分量，复数模长允许达到 sqrt(2) 倍上界。
-        return _complex_value(dtype, spec.shape, rng, offset=-max_abs, scale=2 * max_abs)
-    return rng.cast((rng.random(spec.shape) - 0.5) * (2 * max_abs), dtype)
+        return _complex_value(dtype, spec.shape, backend, offset=-max_abs, scale=2 * max_abs)
+    return backend.symmetric(spec.shape, dtype, max_abs)
 
 
 def generate_nonzero_symmetric_input_value(
     spec: InputTensorSpec,
     max_abs,
-    rng=INPUT_NUMPY_RANDOM_STATE,
+    backend,
 ) -> object:
     """生成可配置对称范围并替换量化后产生的零值。"""
     dtype = resolve_input_dtype(spec.dtype)
-    value = generate_symmetric_input_value(spec, max_abs, rng)
+    value = generate_symmetric_input_value(spec, max_abs, backend)
     # replacement 保持相同 dtype，避免低精度 Tensor 被 Python 标量提升。
-    replacement = rng.asarray(max_abs, dtype=dtype)
-    return rng.where(value == 0, replacement, value)
+    replacement = backend.asarray(max_abs, dtype=dtype)
+    return backend.where(value == 0, replacement, value)
 
 
 def generate_normal_input_value(
@@ -191,19 +191,19 @@ def generate_normal_input_value(
 
 def generate_default_input_value(
     spec: InputTensorSpec,
-    rng=INPUT_NUMPY_RANDOM_STATE,
+    backend,
     *,
     max_abs=0.6,
 ) -> object:
     """生成默认值。"""
     dtype = resolve_input_dtype(spec.dtype)
     if dtype == "bool":
         # 连续随机数 cast 到 bool 几乎恒为 True，必须直接采样二值空间。
-        return rng.cast(rng.randint(0, 2, shape=spec.shape), dtype)
+        return backend.cast(backend.randint(0, 2, shape=spec.shape), dtype)
     if "int" in dtype:
         # 运行级浮点范围不能收窄已有的整数压力测试范围。
-        return rng.cast(rng.randint(-65535, 65535, shape=spec.shape), dtype)
-    return generate_symmetric_input_value(spec, max_abs, rng)
+        return backend.cast(backend.randint(-65535, 65535, shape=spec.shape), dtype)
+    return generate_symmetric_input_value(spec, max_abs, backend)
 
 
 def generate_nonzero_input_value(spec: InputTensorSpec, rng=INPUT_NUMPY_RANDOM_STATE) -> object: