pymllm.layers.attention.gdn_chunkwise

Chunkwise parallel GDN extend (prefill) using WY representation.

Implements the WY decomposition from the Gated DeltaNet paper to parallelize the GDN recurrent scan over token chunks. Within each chunk the key operations (output and state computation) are expressed as batch matrix multiplications, enabling efficient GPU utilisation.

The inter-chunk state propagation remains sequential, but reduces the number of iterations from T to ceil(T / C) where C is the chunk size.

Algorithm (per chunk of C tokens with initial state S_0)

  1. WY construction (sequential, O(C²K) per head):

    w_i = β_i k_i - K_{<i}(W_{<i}^T (β_i k_i)) ũ_i = (β_i / γ_i) v_i - Ũ_{<i}(K_{<i}^T (β_i k_i))

    where γ_i = Π_{j=0}^{i} α_j is the cumulative decay.

  2. Output computation (parallel, batch matmul):

    O = γ · (S_0 Q − (S_0 W) · triu(K^T Q) + Ũ · triu(K^T Q))

    triu is the causal mask (upper-triangular: j ≤ r).

  3. State update (parallel, batch matmul):

    S_new = γ_C · (S_0 − (S_0 W) K^T + Ũ^T K)

Usage:

from pymllm.layers.attention.gdn_chunkwise import gdn_extend_chunkwise

output = gdn_extend_chunkwise(
    q, k, v, a, b, A_log, dt_bias,
    state_pool, cache_indices, cu_seqlens,
    chunk_size=64,
)

Functions

gdn_extend_chunkwise_torch(q, k, v, a, b, A_log, ...)

Fully-batched chunkwise GDN extend (prefill).

gdn_extend_chunkwise(q, k, v, a, b, A_log, dt_bias, ...)

Chunkwise parallel GDN extend with WY representation.

Module Contents

pymllm.layers.attention.gdn_chunkwise.gdn_extend_chunkwise_torch(q, k, v, a, b, A_log, dt_bias, state_pool, cache_indices, cu_seqlens, chunk_size=128)

Fully-batched chunkwise GDN extend (prefill).

Drop-in replacement for gdn_extend_chunkwise() / the CUDA cuda_chunkwise kernel with identical inputs/outputs. The recurrent state in state_pool is updated in-place.

Parameters match gdn_extend_chunkwise().

Parameters:

  • q (torch.Tensor)

  • k (torch.Tensor)

  • v (torch.Tensor)

  • a (torch.Tensor)

  • b (torch.Tensor)

  • A_log (torch.Tensor)

  • dt_bias (torch.Tensor)

  • state_pool (torch.Tensor)

  • cache_indices (torch.Tensor)

  • cu_seqlens (torch.Tensor)

  • chunk_size (int)

Return type:

torch.Tensor

pymllm.layers.attention.gdn_chunkwise.gdn_extend_chunkwise(q, k, v, a, b, A_log, dt_bias, state_pool, cache_indices, cu_seqlens, chunk_size=64)

Chunkwise parallel GDN extend with WY representation.

Parameters:

  • q (Tensor [total_tokens, H, K]) – Query tensor (bf16/fp16). NOT pre-normalised.

  • k (Tensor [total_tokens, H, K]) – Key tensor (bf16/fp16).

  • v (Tensor [total_tokens, HV, V]) – Value tensor (bf16/fp16).

  • a (Tensor [total_tokens, HV]) – Raw decay-gate input (before softplus/exp).

  • b (Tensor [total_tokens, HV]) – Raw update-gate input (before sigmoid).

  • A_log (Tensor [HV]) – Log-space decay parameter, float32.

  • dt_bias (Tensor [HV]) – Bias for decay gate, float32.

  • state_pool (Tensor [pool_size, HV, V, K]) – Pooled recurrent state, float32. Modified in-place.

  • cache_indices (Tensor [batch_size]) – Pool index per request, int64.

  • cu_seqlens (Tensor [batch_size + 1]) – Cumulative sequence lengths, int64.

  • chunk_size (int) – Number of tokens per chunk (default 64).

Returns:

Output tensor, same dtype as v.

Return type:

Tensor [total_tokens, HV, V]