core.ssm.ops.gdp.common

Shared helpers for the Gated Delta Product kernels.

The kernel modules in this package reference a handful of names that do not belong to any one of them: the hardware-capability probes that select autotune configurations, the chunk-descriptor builders, exp / exp2, and the L2 normalization applied to the queries and keys. Keeping them here makes the package self-contained, with no fla import at run time.

Take care when changing the probes: they choose which autotune configurations exist, so editing one changes which kernel variants get benchmarked and picked. Every GDP autotune config list must pass through autotune_configs: timing-based selection can otherwise choose numerically different reduction tilings between deterministic-mode processes.

Module Contents

Functions

_is_nvidia
check_shared_mem	Whether the current device has at least arch’s shared memory budget.
exp	Exponentiate in fp32 regardless of the input dtype.
exp2	Base-2 exponentiate in fp32 regardless of the input dtype.
_segmented_arange	Expand per-segment counts into flat per-slot index tensors.
prepare_chunk_indices	Flattened (sequence, chunk-within-sequence) pairs, one per chunk.
prepare_chunk_offsets	Per-sequence prefix sum of chunk counts, with a leading zero.
l2norm_fwd_kernel1	Row-per-program L2 normalization, used when D > 512.
l2norm_fwd_kernel	Block-of-rows L2 normalization, used when D <= 512.
l2norm_fwd	Row-wise L2 normalization over the last dimension, computed in fp32.

Data

CHUNK_SIZE
RCP_LN2
IS_NVIDIA
IS_NVIDIA_HOPPER
IS_NVIDIA_BLACKWELL
IS_TMA_SUPPORTED
_SHARED_MEM_BY_ARCH
_BT_LIST

API

core.ssm.ops.gdp.common.CHUNK_SIZE

64

core.ssm.ops.gdp.common.RCP_LN2

1.4426950216

core.ssm.ops.gdp.common._is_nvidia() → bool

core.ssm.ops.gdp.common.IS_NVIDIA

‘_is_nvidia(…)’

core.ssm.ops.gdp.common.IS_NVIDIA_HOPPER

None

core.ssm.ops.gdp.common.IS_NVIDIA_BLACKWELL

None

core.ssm.ops.gdp.common.IS_TMA_SUPPORTED

None

core.ssm.ops.gdp.common._SHARED_MEM_BY_ARCH

None

core.ssm.ops.gdp.common.check_shared_mem(arch: str = 'none') → bool

Whether the current device has at least arch’s shared memory budget.

core.ssm.ops.gdp.common.exp(x)

Exponentiate in fp32 regardless of the input dtype.

core.ssm.ops.gdp.common.exp2(x)

Base-2 exponentiate in fp32 regardless of the input dtype.

core.ssm.ops.gdp.common._segmented_arange(

counts: torch.Tensor,

) → tuple[torch.Tensor, torch.Tensor]

Expand per-segment counts into flat per-slot index tensors.

Given segment sizes counts = [c0, c1, ...], return two 1-D tensors of length counts.sum() labelling every slot with its segment and its position within that segment. For counts = [2, 3]::

seg_id    = [0, 0, 1, 1, 1]
intra_idx = [0, 1, 0, 1, 2]

core.ssm.ops.gdp.common.prepare_chunk_indices(

cu_seqlens: torch.Tensor,

chunk_size: int,

cu_seqlens_cpu: torch.Tensor | None = None,

) → torch.Tensor

Flattened (sequence, chunk-within-sequence) pairs, one per chunk.

Note that this synchronizes on the device when the counts live on the GPU: repeat_interleave reads counts.sum() on the host to size its output. Passing cu_seqlens_cpu avoids the sync.

core.ssm.ops.gdp.common.prepare_chunk_offsets(

cu_seqlens: torch.Tensor,

chunk_size: int,

) → torch.Tensor

Per-sequence prefix sum of chunk counts, with a leading zero.

core.ssm.ops.gdp.common._BT_LIST

[8, 16, 32, 64, 128]

core.ssm.ops.gdp.common.l2norm_fwd_kernel1(x, y, rstd, eps, D, BD: triton.language.constexpr)

Row-per-program L2 normalization, used when D > 512.

core.ssm.ops.gdp.common.l2norm_fwd_kernel(

x,

y,

rstd,

eps,

T,

D: triton.language.constexpr,

BD: triton.language.constexpr,

NB: triton.language.constexpr,

BT: triton.language.constexpr,

)

Block-of-rows L2 normalization, used when D <= 512.

core.ssm.ops.gdp.common.l2norm_fwd(x: torch.Tensor, eps: float = 1e-06) → torch.Tensor

Row-wise L2 normalization over the last dimension, computed in fp32.

The kernel also writes rstd, which only a backward pass would consume, so it is not returned.