`evox.triton_kernels.kernels.virtual_noise`#

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) Gaussian 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) Gaussian 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) Gaussian 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)` standard normal 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`
`_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.
`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.

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) standard normal noise deterministically.

Uses a splitmix64-based integer hash of (seed, element_index) (the full 64-bit hash state mixed from seed + (offset + flat_index)), then a Box-Muller transform to produce standard normals.

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 standard-normal tensor.

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._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.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#

Noise indexing scheme#

Module Contents#

Functions#

Data#

API#

`evox.triton_kernels.kernels.virtual_noise`#