evox.triton_kernels.kernels.virtual_noise#

RNG: fast centered-uniform approximation#

This is a forward-only performance demo. The noise is generated with a cheap centered-uniform RNG (uniform [0, 1) shifted by -0.5, i.e. values in [-0.5, 0.5) with mean ~0 and variance 1/12) rather than a standard normal. This deliberately avoids the transcendental functions (sqrt / log / cos of Box-Muller, or the Philox path behind tl.randn) which are the expensive part of the noise generation: a centered uniform is just tl.rand(seed, offsets) - 0.5 and fuses trivially with no overhead. RNG distributional quality is irrelevant for this demo — only determinism (the same (seed, element_index) must always yield the same value, so the forward and gradient paths regenerate identical noise) matters, and that is preserved.

Noise indexing scheme#

For a parameter block at flat element offset off:

  • Weight element (j, k) of shape (out, in): noise = PRNG(seed_i, off + j * in + k)

  • Bias element j of shape (out,) (placed immediately after the weight block, so the caller passes offset = off + out * in): noise = PRNG(seed_i, offset + j)

offset is the cumulative element count across all preceding parameter blocks (see :func:compute_offsets).

Virtual (never-materialized) noise fused into linear / matmul kernels.

This module implements virtual perturbation of neural-network weight matrices: instead of generating a full (pop_size, out_features, in_features) noise tensor (which would be prohibitively large), the noise is generated on-the-fly inside the matmul kernel and added to the weight tile in registers.

The population-based zeroth-order / evolution-strategies model is:

Y[i] = X[i] @ (W + sigma * N_i)^T + (b + sigma * nb_i)

where N_i is a full (out_features, in_features) noise matrix that is unique per individual i (derived from seed_i). The corresponding gradient estimate w.r.t. W is:

grad_W[j, k] = sum_i fitness_i * N_i[j, k] / (pop_size * sigma)

CRITICAL requirement: the forward pass (:func:virtual_perturbed_linear) and the gradient estimate (:func:virtual_weight_gradient / :func:virtual_bias_gradient) must regenerate the exact same noise N_i for a given (seed_i, offset, element_index) triple. Each path (Triton / PyTorch) is internally self-consistent; the two paths do not need to match each other.

Module Contents#

Functions#

_cpu_logical_rshift

Logical (zero-filling) right shift on (possibly negative) int64 tensors.

_splitmix64_step

One round of the splitmix64 mixing function (operates in-place logically).

_cpu_normal_noise

Generate (pop_size, n_elements) cheap centered-uniform noise deterministically.

compute_offsets

Compute cumulative flat-element offsets for a list of parameter blocks.

_virtual_perturbed_linear_fake

_virtual_weight_gradient_fake

_virtual_bias_gradient_fake

_virtual_reduce_metric_fake

Abstract evaluation (torch.compile tracing) for :func:virtual_reduce_metric.

_choose_num_stages

Pick a software-pipelining depth that fits the device’s shared-memory budget.

_triton_virtual_perturbed_linear

Launch the fused Triton virtual-noise perturbed linear kernel.

_triton_virtual_weight_gradient

Launch the fused Triton virtual weight-gradient kernel.

_triton_virtual_bias_gradient

Launch the fused Triton virtual bias-gradient kernel.

_triton_virtual_reduce_metric

Launch the fused Triton virtual-noise mean-abs reduction kernel.

virtual_perturbed_linear

Virtual-noise perturbed linear transformation.

virtual_weight_gradient

Population-based virtual weight gradient estimate.

virtual_bias_gradient

Population-based virtual bias gradient estimate.

virtual_reduce_metric

Virtual-noise mean-absolute-value reduction metric.

Data#

_SPLITMIX64_GAMMA

_SPLITMIX64_GAMMA_I64

_MASK64

API#

evox.triton_kernels.kernels.virtual_noise._SPLITMIX64_GAMMA#

11400714819323198485

evox.triton_kernels.kernels.virtual_noise._SPLITMIX64_GAMMA_I64#

None

evox.triton_kernels.kernels.virtual_noise._MASK64#

None

evox.triton_kernels.kernels.virtual_noise._cpu_logical_rshift(x: torch.Tensor, shift: int) → torch.Tensor[source]#

Logical (zero-filling) right shift on (possibly negative) int64 tensors.

PyTorch’s >> on signed int64 is arithmetic (sign-extending). For values whose high bit is set (i.e. that represent large unsigned 64-bit integers), arithmetic shift corrupts the bits. We reconstruct the logical shift by masking off the sign-extended bits.

Parameters:

  • x – int64 tensor.

  • shift – Number of bits to shift right (0 <= shift < 64).

Returns:

Logical right-shifted int64 tensor.

evox.triton_kernels.kernels.virtual_noise._splitmix64_step(z: torch.Tensor) → torch.Tensor[source]#

One round of the splitmix64 mixing function (operates in-place logically).

Given a 64-bit state z, returns z' such that the full 64-bit result matches the canonical splitmix64 finalizer:

z' = (z ^ (z >> 30)) * GAMMA
z' = (z' ^ (z' >> 27)) * GAMMA
z' = z' ^ (z' >> 31)

All arithmetic wraps mod 2**64. PyTorch int64 multiply/add/xor wrap correctly; logical right shifts use :func:_cpu_logical_rshift.

Parameters:

z – int64 tensor of states.

Returns:

int64 tensor of mixed (64-bit) values.

evox.triton_kernels.kernels.virtual_noise._cpu_normal_noise(seeds: torch.Tensor, n_elements: int, offset: int) → torch.Tensor[source]#

Generate (pop_size, n_elements) cheap centered-uniform noise deterministically.

This is NOT a standard-normal distribution. It is a deliberately cheap centered-uniform approximation: a single uniform value in [0, 1) (derived from the low 32 bits of a splitmix64 hash) shifted by -0.5 so the output lies in [-0.5, 0.5) with mean ~0 and variance 1/12.

This is a forward-only performance demo: RNG distributional quality is irrelevant here — only determinism matters (the forward and gradient paths must regenerate identical noise for the same (seed, element_index)), and that is fully preserved. Avoiding the transcendental functions of Box-Muller (sqrt / log / cos) keeps generation as cheap as possible. The output is trivially finite (bounded in [-0.5, 0.5)), so it can never produce inf / nan.

The element index is offset + flat_index where flat_index ranges over [0, n_elements). For a weight block of shape (out, in) the flat index is j * in + k (row-major); for a bias block (out,) it is just j. Forward and gradient MUST call this helper with the same offset and the same flat layout for the noise to match exactly.

Parameters:

  • seeds – 1-D int64/int32 tensor of per-individual seeds, shape (pop_size,).

  • n_elements – Number of output values per individual.

  • offset – Flat element offset for this block’s noise.

Returns:

(pop_size, n_elements) float32 centered-uniform noise tensor with values in [-0.5, 0.5).

evox.triton_kernels.kernels.virtual_noise.compute_offsets(param_shapes: list[tuple]) → list[int][source]#

Compute cumulative flat-element offsets for a list of parameter blocks.

For each block the number of elements is prod(shape). The returned list holds the starting offset of each block (the cumulative element count of all preceding blocks).

Example::

>>> compute_offsets([(256, 784), (256,), (10, 256), (10,)])
[0, 200704, 200960, 203520]
Parameters:

param_shapes – List of parameter block shapes (each a tuple of ints).

Returns:

List of starting offsets (one per block), length len(param_shapes).

evox.triton_kernels.kernels.virtual_noise._virtual_perturbed_linear_fake(x: torch.Tensor, weight: torch.Tensor, bias, seeds: torch.Tensor, sigma: float, offset: int) → torch.Tensor[source]#
evox.triton_kernels.kernels.virtual_noise._virtual_weight_gradient_fake(fitness: torch.Tensor, seeds: torch.Tensor, weight_shape: list[int], sigma: float, pop_size: int, offset: int) → torch.Tensor[source]#
evox.triton_kernels.kernels.virtual_noise._virtual_bias_gradient_fake(fitness: torch.Tensor, seeds: torch.Tensor, bias_shape: list[int], sigma: float, pop_size: int, offset: int) → torch.Tensor[source]#
evox.triton_kernels.kernels.virtual_noise._virtual_reduce_metric_fake(center: torch.Tensor, seeds: torch.Tensor, sigma: float, n_params: int, offset: int) → torch.Tensor[source]#

Abstract evaluation (torch.compile tracing) for :func:virtual_reduce_metric.

Returns:

An empty (pop_size,) float32 tensor on center’s device.

evox.triton_kernels.kernels.virtual_noise._choose_num_stages(device: torch.device, per_stage_bytes: int, max_stages: int = 4) → int[source]#

Pick a software-pipelining depth that fits the device’s shared-memory budget.

Triton’s default num_stages (>= 3) can over-allocate shared memory for large tl.dot tiles (e.g. the worst-case (BLOCK_BATCH=128, BLOCK_IN=512, BLOCK_OUT=64) forward tile), triggering an OutOfResources crash on shared-memory-limited GPUs (e.g. sm_86 with a ~99 KB per-block opt-in limit).

This defensively queries the device’s per-block shared-memory budget and clamps the pipeline depth to what fits. The footprint estimate (per_stage_bytes) is provided by the caller and is intentionally pessimistic vs Triton’s real (MMA-tiled) usage, so the resulting num_stages is conservative — which is safe (it only avoids OOM).

Parameters:

  • device – The device the kernel will run on.

  • per_stage_bytes – Estimated shared-memory bytes consumed per pipeline stage for the kernel’s largest tile (pessimistic).

  • max_stages – Upper bound on the returned pipeline depth.

Returns:

An int in [1, max_stages].

evox.triton_kernels.kernels.virtual_noise._triton_virtual_perturbed_linear(x: torch.Tensor, weight: torch.Tensor, bias, seeds: torch.Tensor, sigma: float, offset: int) → torch.Tensor[source]#

Launch the fused Triton virtual-noise perturbed linear kernel.

evox.triton_kernels.kernels.virtual_noise._triton_virtual_weight_gradient(fitness: torch.Tensor, seeds: torch.Tensor, weight_shape: list[int], sigma: float, pop_size: int, offset: int) → torch.Tensor[source]#

Launch the fused Triton virtual weight-gradient kernel.

evox.triton_kernels.kernels.virtual_noise._triton_virtual_bias_gradient(fitness: torch.Tensor, seeds: torch.Tensor, bias_shape: list[int], sigma: float, pop_size: int, offset: int) → torch.Tensor[source]#

Launch the fused Triton virtual bias-gradient kernel.

evox.triton_kernels.kernels.virtual_noise._triton_virtual_reduce_metric(center: torch.Tensor, seeds: torch.Tensor, sigma: float, n_params: int, offset: int) → torch.Tensor[source]#

Launch the fused Triton virtual-noise mean-abs reduction kernel.

evox.triton_kernels.kernels.virtual_noise.virtual_perturbed_linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], seeds: torch.Tensor, sigma: float, offset: int) → torch.Tensor#

Virtual-noise perturbed linear transformation.

For each individual i computes::

Y[i] = X[i] @ (W + sigma * N_i)^T + (b + sigma * nb_i)

where N_i is a full (out_features, in_features) Gaussian noise matrix generated on-the-fly from seed_i and the block offset. The noise is never materialized as a full tensor in the Triton path; on the CPU fallback it is generated per-individual (performance is not critical on CPU).

Weight element (j, k) uses noise element index offset + j * in + k; bias element j uses offset + out * in + j.

Parameters:

  • x – Input tensor, either (batch, in_features) (shared across all individuals) or (pop_size, batch, in_features) (per-individual).

  • weight – Weight tensor of shape (out_features, in_features).

  • bias – Bias tensor of shape (out_features,) or None.

  • seeds – 1-D int tensor of per-individual seeds, shape (pop_size,).

  • sigma – Perturbation scale (Python float).

  • offset – Flat element offset for this block’s noise.

Returns:

Output tensor of shape (pop_size, batch, out_features).

evox.triton_kernels.kernels.virtual_noise.virtual_weight_gradient(fitness: torch.Tensor, seeds: torch.Tensor, weight_shape: list[int], sigma: float, pop_size: int, offset: int) → torch.Tensor#

Population-based virtual weight gradient estimate.

Computes::

grad[j, k] = sum_i fitness_i * N_i[j, k] / (pop_size * sigma)

regenerating the same noise N_i as :func:virtual_perturbed_linear for weight element (j, k) (element index offset + j * in + k).

Parameters:

  • fitness – 1-D float tensor of per-individual fitness, (pop_size,).

  • seeds – 1-D int tensor of per-individual seeds, (pop_size,).

  • weight_shape – Target weight shape (out_features, in_features).

  • sigma – Perturbation scale used in the forward pass.

  • pop_size – Population size (number of individuals).

  • offset – Flat element offset for this block’s noise.

Returns:

Gradient tensor of shape weight_shape.

evox.triton_kernels.kernels.virtual_noise.virtual_bias_gradient(fitness: torch.Tensor, seeds: torch.Tensor, bias_shape: list[int], sigma: float, pop_size: int, offset: int) → torch.Tensor#

Population-based virtual bias gradient estimate.

Computes::

grad[j] = sum_i fitness_i * nb_i[j] / (pop_size * sigma)

regenerating the same bias noise nb_i as the bias contribution of

Func:

virtual_perturbed_linear for bias element j (element index offset + j).

Parameters:

  • fitness – 1-D float tensor of per-individual fitness, (pop_size,).

  • seeds – 1-D int tensor of per-individual seeds, (pop_size,).

  • bias_shape – Target bias shape (out_features,).

  • sigma – Perturbation scale used in the forward pass.

  • pop_size – Population size (number of individuals).

  • offset – Flat element offset for this block’s bias noise (the caller should pass the offset pointing at the bias region, i.e. weight_offset + out_features * in_features).

Returns:

Gradient tensor of shape bias_shape.

evox.triton_kernels.kernels.virtual_noise.virtual_reduce_metric(center: torch.Tensor, seeds: torch.Tensor, sigma: float, n_params: int, offset: int) → torch.Tensor#

Virtual-noise mean-absolute-value reduction metric.

Computes, per individual i::

fitness[i] = mean_k( | center[k] + sigma * noise[i, k] | )

where noise[i, k] is generated on-the-fly from seeds[i] using element index offset + k with the SAME fast centered-uniform RNG as the other kernels in this module.

On the Triton (CUDA) path the (pop_size, n_params) noise tensor is NEVER materialized: noise is generated tile-by-tile in registers, fused with the center + sigma * noise perturbation, abs()-ed and reduced to a single scalar per individual (one program per individual). Output memory is O(pop_size).

This CPU fallback (the function body) MAY materialize the noise — performance is not critical on CPU — but uses the same :func:_cpu_normal_noise so the noise is internally consistent.

Parameters:

  • center – 1-D float tensor of length n_params (the center/mean of the perturbation, e.g. a flattened parameter vector).

  • seeds – 1-D int tensor of per-individual seeds, shape (pop_size,).

  • sigma – Perturbation scale (Python float).

  • n_params – Number of parameters (length of center and the reduction dimension). Pass explicitly so it is a trace-time constant on the Triton path; it should equal center.numel().

  • offset – Flat element offset for this block’s noise (default 0).

Returns:

(pop_size,) float32 tensor of per-individual mean-abs values.