evox.triton_kernels.kernels.philox

Philox4x32-10 counter-based pseudo-random number generator (PRNG).

This module provides deterministic, seedable random number generation using the Philox4x32-10 algorithm (the same family used by cuRAND / NumPy's Philox bit generator). Two public ops are registered via

func:

~evox.triton_kernels.register_triton_op:

  • func:

    philox_uniform — uniform floats in [0, 1) (24 mantissa bits).

  • func:

    philox_normal — standard normal floats via Box-Muller transform.

On CUDA devices with Triton installed the ops dispatch to hand-written Triton kernels; on every other device (CPU, MPS, …) they fall back to the pure PyTorch implementations. The two paths are designed to produce identical output for the same inputs.

The algorithm follows the canonical Philox4x32-10 specification:

  • Multiplication keys: M0 = 0xD2511F53, M1 = 0xCD9E8D57

  • Key bump constants: W0 = 0x9E3779B9, W1 = 0xBB67AE85

  • 10 rounds of the Feistel-mixing single_round.

  • Each Philox call consumes one (4×uint32) counter and produces 4 uint32 outputs. The low counter word c0 is incremented per call so successive groups of 4 outputs are independent.

模块内容

函数

_philox_uniform_fallback

Pure-PyTorch Philox4x32-10 uniform generator (the fallback path).

_philox_normal_fallback

Pure-PyTorch standard normal generator via Box-Muller (the fallback path).

_philox_uniform_fake

_philox_normal_fake

_triton_philox_uniform

Launch the Triton Philox uniform kernel.

_triton_philox_normal

Launch the Triton Philox normal kernel via Box-Muller.

philox_uniform

Generate deterministic uniform random floats in [0, 1) via Philox4x32-10.

philox_normal

Generate deterministic standard normal floats via Philox4x32-10 + Box-Muller.

数据

PHILOX_M4x32_0

PHILOX_M4x32_1

PHILOX_W32_0

PHILOX_W32_1

MASK

_ROUNDS

API

evox.triton_kernels.kernels.philox.PHILOX_M4x32_0

3528531795

evox.triton_kernels.kernels.philox.PHILOX_M4x32_1

3449720151

evox.triton_kernels.kernels.philox.PHILOX_W32_0

2654435769

evox.triton_kernels.kernels.philox.PHILOX_W32_1

3144134277

evox.triton_kernels.kernels.philox.MASK

4294967295

evox.triton_kernels.kernels.philox._ROUNDS

10

evox.triton_kernels.kernels.philox._philox_uniform_fallback(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor[源代码]

Pure-PyTorch Philox4x32-10 uniform generator (the fallback path).

Produces n uniform float32 values in [0, 1) per seed using 24 mantissa bits: (uint32 >> 8) * (1 / 2**24).

参数:

  • seeds -- 1-D int64 tensor of per-individual seeds.

  • n -- Number of output values per seed.

  • counter -- Starting 64-bit counter value (jump-ahead offset).

返回:

(pop_size, n) float32 tensor of uniform values in [0, 1).

evox.triton_kernels.kernels.philox._philox_normal_fallback(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor[源代码]

Pure-PyTorch standard normal generator via Box-Muller (the fallback path).

Generates ceil(n/2) uniform pairs (u1, u2) and applies the Box-Muller transform. If n is odd, the last (extra) value is discarded.

参数:

  • seeds -- 1-D int64 tensor of per-individual seeds.

  • n -- Number of output values per seed.

  • counter -- Starting 64-bit counter value (jump-ahead offset).

返回:

(pop_size, n) float32 tensor of standard normal values.

evox.triton_kernels.kernels.philox._philox_uniform_fake(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor[源代码]
evox.triton_kernels.kernels.philox._philox_normal_fake(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor[源代码]
evox.triton_kernels.kernels.philox._triton_philox_uniform(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor[源代码]

Launch the Triton Philox uniform kernel.

Produces the same values as :func:_philox_uniform_fallback for identical inputs.

evox.triton_kernels.kernels.philox._triton_philox_normal(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor[源代码]

Launch the Triton Philox normal kernel via Box-Muller.

Generates 2*ceil(n/2) uniforms using the uniform kernel, then applies Box-Muller. This guarantees bit-identical results with

Func:

_philox_normal_fallback because it reuses the exact same uniform stream and the same Box-Muller arithmetic.

evox.triton_kernels.kernels.philox.philox_uniform(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor

Generate deterministic uniform random floats in [0, 1) via Philox4x32-10.

Each seed in seeds produces an independent stream of n values. The optional counter provides jump-ahead so multiple sub-streams (e.g. for different parameter blocks) can be drawn from the same seed without overlap.

On CUDA devices with Triton, a fused Triton kernel is used; otherwise the pure-PyTorch fallback runs on all devices (CPU, MPS, …). Both paths produce identical output for the same inputs.

参数:

  • seeds -- 1-D int64 tensor of per-individual seeds.

  • n -- Number of output values per seed.

  • counter -- Starting 64-bit counter value (jump-ahead offset). Default 0.

返回:

(pop_size, n) float32 tensor of uniform values in [0, 1).

evox.triton_kernels.kernels.philox.philox_normal(seeds: torch.Tensor, n: int, counter: int = 0) → torch.Tensor

Generate deterministic standard normal floats via Philox4x32-10 + Box-Muller.

Each seed in seeds produces an independent stream of n standard normal (mean 0, std 1) values using the Box-Muller transform applied to the Philox uniform stream.

On CUDA devices with Triton, the uniform stream is generated by the Triton kernel and Box-Muller is applied in PyTorch; otherwise the pure-PyTorch fallback runs on all devices. Both paths produce identical output.

参数:

  • seeds -- 1-D int64 tensor of per-individual seeds.

  • n -- Number of output values per seed.

  • counter -- Starting 64-bit counter value (jump-ahead offset). Default 0.

返回:

(pop_size, n) float32 tensor of standard normal values.