dew.nn.kernels

Pallas kernels for the operations XLA does not fuse.

Each module here owns one operation, exports the predicate that says which shapes and backends its kernel covers, and leaves the XLA form that it replaces in the module the operation belongs to, as the oracle its tests compare against and as the path every other shape and backend takes.

Name	Summary
bf16_dot_runs	Whether the default device multiplies bf16 operands into an fp32 sum as one dot algorithm, BF16_BF16_F32: every TPU and CPU, and a GPU from BF16_GPU on.
device_generation	The default device’s hardware generation: ‘sm89’ for a GPU of compute capability 8.9, ‘v6e’ for a TPU v6e, and the platform’s name otherwise.
grouped_projection	dew.nn.moe.expert_projection on the kernels, first-order reverse mode.
ragged_dot_runs	Whether the kernels compute the product a caller asked for.
ssd_chunk_scan
ssd_kernel_platform	The backend to build this scan’s kernel for, or None for the XLA path.
ssd_kernel_runs	Whether the SSD kernel is chosen for this geometry: a tpu backend, a chunk long enough to pay for a program, three widths that are powers of two so that Mosaic’s tiling throws no lanes away, and a tile inside the per-program budget.
triton_runs	The one eligibility rule for Dew’s Pallas GPU (Triton) kernels: a GPU of compute capability 8.0 or later, the bound JAX’s own Pallas lowerings apply (_backend_supports_triton).

bf16_dot_runs

function source

def bf16_dot_runs() -> bool

Whether the default device multiplies bf16 operands into an fp32 sum as one dot algorithm, BF16_BF16_F32: every TPU and CPU, and a GPU from BF16_GPU on.

device_generation

function source

def device_generation() -> str

The default device’s hardware generation: ‘sm89’ for a GPU of compute capability 8.9, ‘v6e’ for a TPU v6e, and the platform’s name otherwise.

grouped_projection

function source

def grouped_projection(
    x: jax.Array,
    kernel: jax.Array,
    group_sizes: jax.Array,
    compute: Dtype,
    interpret_on_cpu: bool,
) -> jax.Array

dew.nn.moe.expert_projection on the kernels, first-order reverse mode.

The operands are cast to compute, the product accumulates in fp32 and rounds once to compute. The input gradient takes x’s dtype and the kernel gradient the kernel’s, each summed in fp32 from the compute-dtype cotangent and rounded operands: with 16-bit compute, exact products. Only a CUDA lowering runs the kernels (interpret_on_cpu adds the CPU’s interpreter); every other lowering runs jax.lax.ragged_dot under the same contract.

ragged_dot_runs

function source

def ragged_dot_runs(
    compute: Dtype,
    operands: tuple[Dtype, ...],
    precision: PrecisionLike,
) -> bool

Whether the kernels compute the product a caller asked for.

They multiply in compute, accumulate in fp32 and ignore precision. With 16-bit compute that is exact products summed in fp32, which any precision asks for; with fp32 it is TF32, which only the default precision asks for (explicitly or through jax_default_matmul_precision). An operand or master wider than fp32 needs its gradient summed wider than the kernels do, and x64 widens their int32 group offsets, so both are refused.

ssd_chunk_scan

attribute source

ssd_chunk_scan = jax.custom_vjp(_ssd_scan, nondiff_argnums=(5,))

ssd_kernel_platform

function source

def ssd_kernel_platform(chunk_size: int, head_dim: int, state_size: int) -> str | None

The backend to build this scan’s kernel for, or None for the XLA path.

ssd_kernel_runs

function source

def ssd_kernel_runs(
    chunk_size: int,
    head_dim: int,
    state_size: int,
    backend: str,
) -> bool

Whether the SSD kernel is chosen for this geometry: a tpu backend, a chunk long enough to pay for a program, three widths that are powers of two so that Mosaic’s tiling throws no lanes away, and a tile inside the per-program budget.

A GPU takes the XLA path. A Triton port of this kernel ran 6x to 12x slower than XLA on an RTX 4080 wherever it compiled, and every chunk of 128 or 256 overflowed shared memory (docs/performance.md); it was removed.

chunk_ssd asks this at trace time and takes the XLA path when it says no, the way attention’s ‘auto’ asks cudnn_runs.

triton_runs

function source

def triton_runs() -> bool

The one eligibility rule for Dew’s Pallas GPU (Triton) kernels: a GPU of compute capability 8.0 or later, the bound JAX’s own Pallas lowerings apply (_backend_supports_triton). A T4 fails to compile them (“Triton support is only enabled for cc>=8.0”).