Class PartialRotaryEncoding

java.lang.Object
smile.llm.model.qwen.PartialRotaryEncoding
public final class PartialRotaryEncoding extends Object
Partial rotary positional encoding for Qwen3.5 gated attention.

Matches HuggingFace apply_rotary_pos_emb / rotate_half (NeoX-style half-rotation), not Meta interleaved complex RoPE. Only the leading rotaryDim features of each head are rotated; the remainder is passed through unchanged.

  • Method Details

    • apply

      public static Tuple2<Tensor,Tensor> apply(Tensor xq, Tensor xk, Tensor cos, Tensor sin, int rotaryDim)
      Applies HF rotate_half RoPE to query and key tensors.
      Parameters:
      xq - query [B, S, H, D].
      xk - key [B, S, Hkv, D].
      cos - cosines for the current positions [S, rotaryDim], [B, S, rotaryDim] (per-row decode), or broadcastable.
      sin - sines with the same layout as cos.
      rotaryDim - number of leading head dims to rotate (even).
      Returns:
      rotated (query, key).

    • gather

      public static Tensor gather(Tensor table, int[] positions)
      Gathers cos/sin rows for per-request decode positions.
      Parameters:
      table - [maxPos, rotaryDim] table.
      positions - absolute positions, one per batch row.
      Returns:
      [B, 1, rotaryDim] (caller owns).

    • gatherInto

      public static void gatherInto(Tensor table, int[] positions, Tensor out)
      Gathers RoPE rows into a pre-allocated [B, 1, R] buffer (CUDA graph).
      Parameters:
      table - [maxPos, rotaryDim] table.
      positions - absolute positions, one per batch row.
      out - destination [B, 1, rotaryDim] (caller-owned, stable address).

    • gatherWindowInto

      public static void gatherWindowInto(Tensor table, int startPos, int windowLen, Tensor out)
      Gathers a contiguous window of RoPE rows into a pre-allocated [S, R] buffer (CUDA graph), for a verify step where every batch row shares the same absolute start position (broadcasts across batch like an ordinary [S, R] table slice — see broadcastCosSin). Uses an index_select gather rather than a table slice-view so a captured graph replay reads the buffer's stable address, not a capture-time storage offset baked into a view.
      Parameters:
      table - [maxPos, rotaryDim] table.
      startPos - absolute position of the window's first token.
      windowLen - number of contiguous positions to gather.
      out - destination [windowLen, rotaryDim] (caller-owned, stable address).

    • gatherWindowInto

      public static void gatherWindowInto(Tensor table, int[] startPositions, int windowLen, Tensor out)
      Per-row form of gatherWindowInto(Tensor, int, int, Tensor) for a ragged batched verify graph: writes each row's window into a stable [B, S, R] buffer in place.
      Parameters:
      table - [maxPos, rotaryDim] table.
      startPositions - absolute position of each row's first window token.
      windowLen - positions per row.
      out - destination [B, windowLen, rotaryDim] (stable address).

    • gatherWindow

      public static Tensor gatherWindow(Tensor table, int[] startPositions, int windowLen)
      Gathers a per-row contiguous window of RoPE rows, for a batched verify step where concurrent requests are at different absolute positions (unlike gatherWindowInto(Tensor, int, int, Tensor), which broadcasts one shared window across the whole batch — only correct when every row shares it).
      Parameters:
      table - [maxPos, rotaryDim] table.
      startPositions - absolute position of the window's first token per batch row.
      windowLen - number of contiguous positions to gather per row.
      Returns:
      owned [B, windowLen, rotaryDim] tensor.

    • computeCosSin

      public static PartialRotaryEncoding.CosSin computeCosSin(int rotaryDim, int end, double theta)
      Precomputes HF-style cos/sin tables for the rotary subspace.

      For each position t and frequency index i: θ = t / base^(2i/dim), then cos = cat(cos θ, cos θ), sin = cat(sin θ, sin θ) along the last dim (matches Transformers emb = cat(freqs, freqs)).

      Parameters:
      rotaryDim - rotary feature count (even).
      end - table length (typically 2 * maxSeqLen).
      theta - RoPE base theta.
      Returns:
      cos/sin tables of shape [end, rotaryDim] in float32.

    • computeFreqCis

      @Deprecated public static Tensor computeFreqCis(int rotaryDim, int end, double theta)
      Deprecated.
      use computeCosSin(int, int, double); kept only if callers still expect a single complex table.
      Legacy complex frequency table (unsupported).
      Parameters:
      rotaryDim - partial RoPE dimension.
      end - exclusive position end.
      theta - RoPE base theta.
      Returns:
      never returns.