diff --git a/tester/input_generation/backend.py b/tester/input_generation/backend.py index 01466bd6..f6f24c12 100644 --- a/tester/input_generation/backend.py +++ b/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 diff --git a/tester/input_generation/backend_runtime.py b/tester/input_generation/backend_runtime.py index 29d7cb7b..99ce7c77 100644 --- a/tester/input_generation/backend_runtime.py +++ b/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] diff --git a/tester/input_generation/generation_rules.py b/tester/input_generation/generation_rules.py index effa3ab9..cfb48978 100644 --- a/tester/input_generation/generation_rules.py +++ b/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): diff --git a/tester/input_generation/value_generators.py b/tester/input_generation/value_generators.py index 0b25f2eb..ba710ea2 100644 --- a/tester/input_generation/value_generators.py +++ b/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,7 +191,7 @@ def generate_normal_input_value( def generate_default_input_value( spec: InputTensorSpec, - rng=INPUT_NUMPY_RANDOM_STATE, + backend, *, max_abs=0.6, ) -> object: @@ -199,11 +199,11 @@ def generate_default_input_value( 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: