ejkernel.kernels._xla.gla._interface

Contents

ejkernel.kernels._xla.gla._interface#

Gated Linear Attention interface using recurrent formulation.

This module provides the public API for GLA implemented via the recurrent attention mechanism. Wraps the core recurrent implementation with GLA-specific gating support for linear-time sequence processing.

ejkernel.kernels._xla.gla._interface.recurrent_gla(query: Float[jaxlib._jax.Array, 'batch seq_len num_heads qk_head_dim'], key: Float[jaxlib._jax.Array, 'batch seq_len num_kv_heads qk_head_dim'], value: Float[jaxlib._jax.Array, 'batch seq_len num_kv_heads v_head_dim'], g: jaxtyping.Float[jaxlib._jax.Array, 'batch seq_len num_heads qk_head_dim'] | None = None, g_gamma: jaxtyping.Float[jaxlib._jax.Array, '... num_heads'] | None = None, softmax_scale: float | None = None, initial_state: jaxtyping.Float[jaxlib._jax.Array, '... num_heads qk_head_dim v_head_dim'] | None = None, reverse: bool = False, cu_seqlens: jaxtyping.Int[jaxlib._jax.Array, 'num_seqs_plus_one'] | None = None) → tuple[jaxtyping.Float[jaxlib._jax.Array, 'batch seq_len num_heads v_head_dim'], jaxtyping.Float[jaxlib._jax.Array, '... num_heads qk_head_dim v_head_dim']][source]#

Computes Gated Linear Attention (GLA) in a recurrent, linear-time manner using JAX/XLA.

This function provides a convenient wrapper around the core recurrent implementation, tailored for GLA. It processes sequences step-by-step, making it highly efficient for very long sequences and suitable for autoregressive decoding.

It supports both standard batch processing and variable-length sequence processing using cumulative sequence lengths (cu_seqlens).

Parameters

  • query – The query tensor. Expected shape is (batch, seq_len, num_heads, head_dim) or (total_tokens, num_heads, head_dim) if cu_seqlens is used.

  • key – The key tensor. Must have the same shape as q.

  • value – The value tensor. Must have the same shape as q.

  • g – The gate tensor, specific to Gated Linear Attention. If provided, it should have the same shape as q.

  • g_gamma – The gate decay factor.

  • softmax_scale – A scaling factor applied to the query before the recurrent computation. If None, it defaults to 1 / sqrt(head_dim).

  • initial_state – The initial hidden state for the recurrence. Useful for chunked processing of long sequences.

  • reverse – If True, the sequence is processed in reverse order.

  • cu_seqlens – Cumulative sequence lengths for variable-length inputs. This is a 1D tensor like [0, len_seq1, len_seq1+len_seq2, …]. If provided, the input tensors query, key, value, g are expected to be “packed” with a shape of (total_tokens, …).

Returns

  • o (jax.Array): The output tensor, with the same shape as q.

  • final_state (jax.Array): The final hidden state of the recurrence, which can be used as initial_state for a subsequent segment.

Return type

A tuple containing

Raises

  • ValueError – If cu_seqlens is provided and the batch size of q is not 1.

  • ValueError – If cu_seqlens is provided and the number of initial states does not match the number of sequences.

Examples

>>>
>>> q = jnp.ones((2, 100, 8, 64))
>>> k = jnp.ones((2, 100, 8, 64))
>>> v = jnp.ones((2, 100, 8, 64))
>>> g = jnp.ones((2, 100, 8, 64))
>>> output, final_state = recurrent_gla(query, key, value, g=g)
>>> output.shape
(2, 100, 8, 64)

>>>
>>> q = jnp.ones((150, 8, 64))
>>> k = jnp.ones((150, 8, 64))
>>> v = jnp.ones((150, 8, 64))
>>> g = jnp.ones((150, 8, 64))
>>> cu_seqlens = jnp.array([0, 50, 100, 150])
>>> output, final_state = recurrent_gla(query, key, value, g=g, cu_seqlens=cu_seqlens)
>>> output.shape
(150, 8, 64)
Contents