Class GatedAttention

java.lang.Object
smile.llm.model.qwen.GatedAttention
All Implemented Interfaces:
Attention
public class GatedAttention extends Object implements Attention
Qwen3.5 gated full attention (GQA + partial RoPE + sigmoid output gate).

The query projection emits twice the head channels; the second half is a per-token sigmoid gate applied to the attention output before o_proj.

Under tensor parallelism, Q/K/V projections are column-sharded by heads and o_proj is row-parallel (all-reduce after the projection).

  • Constructor Details

    • GatedAttention

      public GatedAttention(int dim, int numHeads, int numKvHeads, int headDim, int rotaryDim, double normEps, int kvLayerId)
      Constructor. Install a KvCachePool via setCachePool(KvCachePool) before the first forward(Tensor, int, Tensor, Tensor).
      Parameters:
      dim - hidden size.
      numHeads - query head count (local under TP).
      numKvHeads - key/value head count (local under TP).
      headDim - per-head dimension.
      rotaryDim - partial RoPE dimension.
      normEps - RMSNorm epsilon.
      kvLayerId - layer index inside the KV pool.

    • GatedAttention

      public GatedAttention(int dim, int numHeads, int numKvHeads, int headDim, int rotaryDim, double normEps, int kvLayerId, TensorParallelGroup tpGroup, int tpRank)
      Tensor-parallel constructor.
      Parameters:
      dim - hidden size.
      numHeads - query head count (local under TP).
      numKvHeads - key/value head count (local under TP).
      headDim - per-head dimension.
      rotaryDim - partial RoPE dimension.
      normEps - RMSNorm epsilon.
      kvLayerId - layer index inside the KV pool.
      tpGroup - tensor-parallel group, or null for single-device.
      tpRank - this rank's TP index.

    • GatedAttention

      public GatedAttention(int dim, int numHeads, int numKvHeads, int headDim, int rotaryDim, double normEps, KvCacheLayout layout)
      Test helper that allocates a private KV pool.
      Parameters:
      dim - hidden size.
      numHeads - query head count.
      numKvHeads - key/value head count.
      headDim - per-head dimension.
      rotaryDim - partial RoPE dimension.
      normEps - RMSNorm epsilon.
      layout - KV cache layout used to size the private pool.

  • Method Details

    • replaceProjections

      public void replaceProjections(LinearOp q, LinearOp k, LinearOp v, LinearOp o)
      Replaces Q/K/V/O with quantized ops (already sharded/packed).
      Parameters:
      q - query projection op.
      k - key projection op.
      v - value projection op.
      o - output projection op.

    • forShard

      public static GatedAttention forShard(int dim, int headDim, int rotaryDim, double normEps, int kvLayerId, TensorShardSpec shard, TensorParallelGroup tpGroup)
      Builds attention from a shard spec (local head counts).
      Parameters:
      dim - hidden size.
      headDim - per-head dimension.
      rotaryDim - partial RoPE dimension.
      normEps - RMSNorm epsilon.
      kvLayerId - layer index inside the KV pool.
      shard - local head / rank shard description.
      tpGroup - tensor-parallel group.
      Returns:
      gated attention for the shard.

    • module

      public MemorySegment module()
      Description copied from interface: Attention
      Returns pytorch module.
      Specified by:
      module in interface Attention
      Returns:
      pytorch module.

    • forward

      public Tensor forward(Tensor x, int startPos, Tensor cis, Tensor mask)
      Description copied from interface: Attention
      Forward pass through the attention module.
      Specified by:
      forward in interface Attention
      Parameters:
      x - the input tensor.
      startPos - the starting position for attention caching.
      cis - the precomputed frequency tensor.
      mask - the attention mask tensor.
      Returns:
      the output tensor.

    • forward

      public Tensor forward(Tensor x, int startPos, Tensor cos, Tensor sin, Tensor mask)
      Forward with HuggingFace-style partial RoPE cos/sin tables.
      Parameters:
      x - hidden states [B, S, D].
      startPos - KV cache write position.
      cos - cosines for this window [S, rotaryDim].
      sin - sines for this window [S, rotaryDim].
      mask - causal attention mask, or null when S == 1.
      Returns:
      attention output [B, S, D].

    • forward

      public Tensor forward(Tensor x, int[] positions, Tensor cos, Tensor sin, Tensor mask)
      Decode with per-row cache write positions (seqLen == 1), uniform prefill/verify when every row shares the same positions[0], or a genuinely ragged multi-token verify window (concurrent requests at different absolute positions, seqLen > 1) via forwardVerifyRagged(Tensor, int[], Tensor, Tensor).
      Parameters:
      x - hidden states [B, S, D].
      positions - absolute write position per batch row.
      cos - cosines [S, R] or per-row [B, S, R].
      sin - sines with the same layout as cos.
      mask - causal mask; ignored by the ragged verify path (that kernel's causal masking is unconditional), used otherwise, or null for decode.
      Returns:
      attention output.