pymllm.layers.attention.gdn_chunkwise
=====================================

.. py:module:: pymllm.layers.attention.gdn_chunkwise

.. autoapi-nested-parse::

   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
---------

.. autoapisummary::

   pymllm.layers.attention.gdn_chunkwise.gdn_extend_chunkwise_torch
   pymllm.layers.attention.gdn_chunkwise.gdn_extend_chunkwise


Module Contents
---------------

.. py:function:: 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 :func:`gdn_extend_chunkwise` / the CUDA
   ``cuda_chunkwise`` kernel with identical inputs/outputs.  The recurrent
   state in ``state_pool`` is updated in-place.

   Parameters match :func:`gdn_extend_chunkwise`.


.. py:function:: 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.

   :param q: Query tensor (bf16/fp16).  NOT pre-normalised.
   :type q: Tensor [total_tokens, H, K]
   :param k: Key tensor (bf16/fp16).
   :type k: Tensor [total_tokens, H, K]
   :param v: Value tensor (bf16/fp16).
   :type v: Tensor [total_tokens, HV, V]
   :param a: Raw decay-gate input (before softplus/exp).
   :type a: Tensor [total_tokens, HV]
   :param b: Raw update-gate input (before sigmoid).
   :type b: Tensor [total_tokens, HV]
   :param A_log: Log-space decay parameter, float32.
   :type A_log: Tensor [HV]
   :param dt_bias: Bias for decay gate, float32.
   :type dt_bias: Tensor [HV]
   :param state_pool: Pooled recurrent state, float32.  Modified **in-place**.
   :type state_pool: Tensor [pool_size, HV, V, K]
   :param cache_indices: Pool index per request, int64.
   :type cache_indices: Tensor [batch_size]
   :param cu_seqlens: Cumulative sequence lengths, int64.
   :type cu_seqlens: Tensor [batch_size + 1]
   :param chunk_size: Number of tokens per chunk (default 64).
   :type chunk_size: int

   :returns: Output tensor, same dtype as *v*.
   :rtype: Tensor [total_tokens, HV, V]