TriangleHyena#

class TriangleHyena(
hidden_dim,
global_conv_cfg,
short_conv_kernel_size=3,
short_conv_backend='torch',
start_node=True,
right_padding=False,
)#

Bases: Module

Starting- or ending-node Hyena mixer for pair tensors [B, L, L, C].

Inner mixer for QKVSequenceMixer. It sets takes_fused_qkv so the mixer hands it the un-split [B, L, L, 3C] projection; forward(q, k, v, ...) with three tensors works as well.

Parameters:

  • hidden_dim (int)

  • global_conv_cfg (LazyConfig)

  • short_conv_kernel_size (int)

  • short_conv_backend (Literal['torch', 'subq_ops'])

  • start_node (bool)

  • right_padding (bool)

global_conv#

The 1D CKConvND long convolution (SIREN filter, window, FFT / Toeplitz backend).

Type:

torch.nn.Module

short_conv_weight#

Depthwise causal taps [K, 3C] for the query, key and value streams, in that channel order.

Type:

torch.nn.Parameter

start_node#

True for the starting-node update (rows are the sequences), False for the ending-node update (columns).

Type:

bool

takes_fused_qkv = True#

Opts in to the un-split [B, L, L, 3C] projection from QKVSequenceMixer.

__init__(
hidden_dim,
global_conv_cfg,
short_conv_kernel_size=3,
short_conv_backend='torch',
start_node=True,
right_padding=False,
)#

Build the short-conv taps and instantiate the long convolution.

Parameters:

  • hidden_dim (int) – Channel width C of the pair tensor and of each projection.

  • global_conv_cfg (LazyConfig) – LazyConfig for a 1D non-causal CKConvND with hidden_dim channels (see the module docstring).

  • short_conv_kernel_size (int) – Taps of the depthwise causal short conv.

  • short_conv_backend (Literal['torch', 'subq_ops']) – "torch" runs the short conv with F.conv1d; "subq_ops" with the channel-last subquadratic_ops_torch.causal_conv1d CUDA kernel. Under torch.compile both use the fused stencil form.

  • start_node (bool) – True mixes along rows (starting node), False along columns (ending node).

  • right_padding (bool) – Skip zeroing the projections before the short conv. True only if every padding mask marks a suffix of each sample; other layouts then leak padded values.

Raises:

ValueError – If short_conv_backend is not "torch" or "subq_ops", or global_conv_cfg is not a 1D non-causal zero-padded double-grid CKConvND with hidden_dim channels, or its kernel is FiLM-conditioned.

extra_repr()#

Return a short summary string.

Return type:

str

flop_count(spatial_dims, inference=False)#

Count FLOPs for one forward pass over an L x L pair tensor.

Short conv (three depthwise streams), the two gates, and the long conv once per row; the SIREN filter and its FFT run once per forward, so their cost is counted once.

Parameters:

  • spatial_dims (tuple[int, int]) – (L, L).

  • inference (bool) – Forwarded to global_conv.flop_count.

Returns:

Total FLOP count; a multiply-add counts as 2 FLOPs, as in QKVSequenceMixer.flop_count().

Return type:

int

forward(
query,
key=None,
value=None,
cp_group=None,
padding_mask=None,
**mixer_kwargs,
)#

Apply the row- or column-wise Hyena update.

Parameters:

  • query (Tensor) – [B, L, L, C] gate-input projection (reference x1), or, when key and value are None, the fused [B, L, L, 3C] projection [query | key | value] as QKVSequenceMixer passes it (takes_fused_qkv).

  • key (Tensor | None) – [B, L, L, C] projection whose columns feed the long conv (reference v).

  • value (Tensor | None) – [B, L, L, C] output-gate projection (reference x2).

  • cp_group (ProcessGroup) – Context-parallel group; only None or a group of size 1 is supported.

  • padding_mask (Tensor | None) – Optional [B, L] boolean tensor, True at padded tokens. The projections are zeroed at every pair position in a padded row or column before the short conv (skipped with right_padding), and the long-conv input x1 * v and the output gate after it. Padded inputs must be finite for training: valid outputs ignore them, but non-finite inputs make upstream weight gradients (e.g. the input projection) non-finite.

  • **mixer_kwargs – Only conditioning is accepted, and it is ignored (no FiLM kernel conditioning).

Returns:

[B, L, L, C] update of the pair tensor.

Raises:

  • ValueError – If the pair tensor is not square, the channel width does not match hidden_dim, only one of key / value is given, or cp_group has size > 1.

  • TypeError – If mixer_kwargs holds any key other than conditioning.

Return type:

Tensor