Class QwenModel

java.lang.Object
smile.deep.layer.LayerBlock
smile.llm.model.qwen.QwenModel
All Implemented Interfaces:
Function<Tensor,Tensor>, Layer
public class QwenModel extends LayerBlock
Qwen3.5 hybrid text model: embeddings, hybrid blocks, final norm, LM head.

When constructed with a TensorShardSpec, attention / FFN / DeltaNet projections are locally sized for that TP rank. Embeddings and the LM head remain replicated (full vocab) on each rank in phase 1.

Construct on CPU, load weights, then call to(Device) to place parameters and the RoPE frequency table. Install the KV pool afterward via setKvCachePool(KvCachePool, boolean) when full-attention layers are present.

  • Constructor Details

    • QwenModel

      public QwenModel(QwenModelArgs args, DeltaNetStatePool statePool)
      Constructs the module graph on CPU. Call to(Device) after weight load; then setKvCachePool(KvCachePool, boolean) when full-attention layers exist.
      Parameters:
      args - hyperparameters.
      statePool - DeltaNet state pool (may be null when no linear layers).

    • QwenModel

      public QwenModel(QwenModelArgs args, DeltaNetStatePool statePool, TensorShardSpec shard, TensorParallelGroup tpGroup)
      Tensor-parallel shard constructor (CPU). Call to(Device) after load.
      Parameters:
      args - hyperparameters.
      statePool - DeltaNet state pool (may be null when no linear layers).
      shard - local head / FFN shard description, or null for full width.
      tpGroup - tensor-parallel group, or null for single-device.

    • QwenModel

      public QwenModel(QwenModelArgs args, DeltaNetStatePool statePool, TensorShardSpec shard, TensorParallelGroup tpGroup, QwenVisionArgs visionArgs)
      Multimodal constructor with optional vision tower.
      Parameters:
      args - text hyperparameters.
      statePool - DeltaNet state pool.
      shard - TP shard, or null.
      tpGroup - TP group, or null.
      visionArgs - vision hyperparameters, or null for text-only.

  • Method Details

    • to

      public QwenModel to(Device device)
      Moves parameters and the RoPE cos/sin tables to device.
      Specified by:
      to in interface Layer
      Overrides:
      to in class LayerBlock
      Parameters:
      device - the compute device.
      Returns:
      this layer.

    • to

      public QwenModel to(Device device, ScalarType dtype)
      Moves parameters and the RoPE cos/sin tables to device / dtype. RoPE tables stay float32 (device move only).
      Specified by:
      to in interface Layer
      Overrides:
      to in class LayerBlock
      Parameters:
      device - the compute device.
      dtype - the data type.
      Returns:
      this layer.

    • params

      public QwenModelArgs params()
      Returns model hyperparameters.
      Returns:
      model args.

    • numLayers

      public int numLayers()
      Returns the number of hybrid decoder blocks.
      Returns:
      number of hybrid decoder blocks.

    • layers

      public List<QwenBlock> layers()
      Returns the decoder blocks.
      Returns:
      decoder blocks (owned by this model).

    • visual

      public QwenVisionTower visual()
      Returns the vision tower.
      Returns:
      vision tower, or null when text-only.

    • visionArgs

      public QwenVisionArgs visionArgs()
      Returns the vision args.
      Returns:
      vision args, or null when text-only.

    • embedTokens

      public Tensor embedTokens(Tensor tokens)
      Token embedding lookup (for multimodal splice).
      Parameters:
      tokens - token ids.
      Returns:
      embeddings.

    • kvCachePool

      public KvCachePool kvCachePool()
      Returns the KV cache pool for full-attention layers, if installed.
      Returns:
      KV pool, or null if unset.

    • deltaNetStatePool

      public DeltaNetStatePool deltaNetStatePool()
      Returns the DeltaNet recurrent/conv state pool.
      Returns:
      DeltaNet state pool, or null when unused.

    • mtpAnchorPool

      public MtpAnchorPool mtpAnchorPool()
      Returns the per-request MTP anchor store.
      Returns:
      anchor pool, or null when mtp() is null.

    • mtp

      public QwenMtp mtp()
      Returns the native MTP draft head.
      Returns:
      native MTP draft head, or null when not configured.

    • lastPreNormHidden

      public Tensor lastPreNormHidden()
      Returns the last captured post-final-norm hidden for MTP.
      Returns:
      last pre-norm hidden, or null.

    • setMtpAnchorAtWindowPosition

      public void setMtpAnchorAtWindowPosition(int r)
      On a partial MTP accept at window position r (checkpoint-replay path), restores the MTP anchor from the per-position hidden retained in verifyWindowNormalizedBuf instead of a second full-window forward.
      Parameters:
      r - accepted window position (0-indexed).

    • setMtpAnchorAtWindowPositions

      public void setMtpAnchorAtWindowPositions(int[] positions)
      Batched counterpart of setMtpAnchorAtWindowPosition(int): restores each row i's MTP anchor from its own retained window position positions[i] — needed when concurrent requests verified together in one round accept different numbers of draft tokens.
      Parameters:
      positions - accepted window position per row (0-indexed), length B.

    • shard

      public TensorShardSpec shard()
      Returns the tensor-parallel shard description for this rank.
      Returns:
      shard spec, or null for unsharded models.

    • tpRank

      public int tpRank()
      Returns this rank's tensor-parallel index.
      Returns:
      TP rank (0 when unsharded).

    • setKvCachePool

      public void setKvCachePool(KvCachePool pool, boolean closePrevious)
      Replaces the KV cache pool on every full-attention layer.
      Parameters:
      pool - new KV pool (must not be null).
      closePrevious - true to close the previous pool when replaced.

    • forward

      public Tensor forward(Tensor tokens, int startPos)
      Forward pass.
      Parameters:
      tokens - token ids [B, S].
      startPos - cache start position.
      Returns:
      logits [B, S, V] in float32 (or [B, 1, V] when only the last position is scored — see forward(Tensor, int, boolean)).

    • forward

      public Tensor forward(Tensor tokens, int startPos, boolean allTokenLogits)
      Forward pass.
      Parameters:
      tokens - token ids [B, S].
      startPos - cache start position.
      allTokenLogits - when false and S > 1, run lm_head only on the last hidden state (sampling / decode). When true, score every position (logprobs).
      Returns:
      logits in float32.

    • forwardEmbeds

      public Tensor forwardEmbeds(Tensor inputsEmbeds, int startPos, Tensor cos, Tensor sin, boolean allTokenLogits)
      Prefill from precomputed embeddings (multimodal splice) with optional interleaved mRoPE cos/sin. When cos/sin are null, uses the standard 1D RoPE table slice for [startPos, startPos+S).
      Parameters:
      inputsEmbeds - [B, S, D] hidden states.
      startPos - cache start position.
      cos - optional [S, rotaryDim] (or null).
      sin - optional [S, rotaryDim] (or null).
      allTokenLogits - whether to score every position.
      Returns:
      logits in float32.

    • spliceVisionEmbeds

      public static Tensor spliceVisionEmbeds(Tensor embeds, int[] inputIds, Tensor visionEmbeds, int imageTokenId, int videoTokenId)
      Replaces image/video pad rows in text embeddings with vision features.
      Parameters:
      embeds - [1, S, D] text embeddings (mutated copy returned).
      inputIds - length S token ids.
      visionEmbeds - [N, D] vision tokens in pad order.
      imageTokenId - image pad id.
      videoTokenId - video pad id.
      Returns:
      spliced embeddings [1, S, D] (caller owns).

    • forward

      public Tensor forward(Tensor tokens, int[] positions)
      Decode forward with per-row absolute positions (seqLen must be 1).
      Parameters:
      tokens - token ids [B, 1].
      positions - write position per batch row.
      Returns:
      logits [B, 1, V].

    • forward

      public Tensor forward(Tensor tokens, int[] positions, boolean allTokenLogits)
      Decode forward with per-row absolute positions.
      Parameters:
      tokens - token ids [B, 1].
      positions - write position per batch row.
      allTokenLogits - unused for S == 1 (kept for API symmetry).
      Returns:
      logits in float32.

    • forward

      public Tensor forward(Tensor tokens, int[] cachePositions, int[] ropePositions, boolean allTokenLogits)
      Decode forward with separate KV write positions and RoPE gather positions (needed for multimodal rope_delta).
      Parameters:
      tokens - token ids [B, 1].
      cachePositions - KV write positions.
      ropePositions - RoPE table gather positions.
      allTokenLogits - unused for S == 1.
      Returns:
      logits in float32.

    • forwardDecodeGraph

      public Tensor forwardDecodeGraph(Tensor tokens, int[] cachePositions, int[] ropePositions)
      Uniform decode forward with optional CUDA graph capture / replay.
      Parameters:
      tokens - token ids [B, 1].
      cachePositions - KV write positions (uniform length across rows).
      ropePositions - RoPE table gather positions.
      Returns:
      logits in float32 [B, 1, V] (graph path returns persistent buffer).

    • closeDecodeGraph

      public void closeDecodeGraph()
      Releases CUDA graph resources for this rank.

    • invalidateDecodeCudaGraphs

      public void invalidateDecodeCudaGraphs()
      Drops captured decode / prefetch CUDA graphs while keeping durable token/RoPE/logits buffers. Required after MTP window verify or KvCachePool.truncateTo(int, int), which rebuild FlashInfer CSR tensors that a captured graph may still reference (replay → illegal memory access).

    • forwardVerifyGraph

      public Tensor forwardVerifyGraph(Tensor tokens, int startPos)
      Uniform MTP window-verify forward (query length S = numDrafts + 1, every batch row sharing the same start position) with optional CUDA graph capture / replay.

      Stage 4 shadow-run: whenever VerifyCudaGraph.enabled(), every call runs eager through the new stable-buffer / graph-capturable kernel path (see AttentionContext dispatch in GatedAttention.forwardUniform), and bucket/warmup bookkeeping in verifyGraphSession runs for real against live traffic shapes — but VerifyCudaGraph.captureEnabled() is false until Stage 5, so beginCapture is never reached. This lets the existing eager forward(Tensor, int, boolean) path serve as the numeric go/no-go reference for the new kernel/buffer plumbing before Stage 5 adds the (separate) capture/replay risk.

      Parameters:
      tokens - token ids [B, S].
      startPos - cache start position (uniform across the batch).
      Returns:
      logits in float32 [B, S, V] (graph path returns persistent buffer).

    • forwardVerifyGraph

      public Tensor forwardVerifyGraph(Tensor tokens, int[] startPositions)
      Ragged form of forwardVerifyGraph(Tensor, int): the B rows (concurrent requests) may be at different absolute positions. One graph is captured per (batch, windowLen); per-row positions, RoPE rows, KV write slots and the page-table CSR are rewritten in place before each replay. Falls back to the eager path when graphs are unavailable.
      Parameters:
      tokens - token ids [B, S].
      startPositions - KV write position of each row's first window token.
      Returns:
      logits in float32 [B, S, V] (the graph path returns a persistent buffer).

    • forwardVerifyGraphCodeEager

      public Tensor forwardVerifyGraphCodeEager(Tensor tokens, int startPos)
      Diagnostic only: runs the exact code a verify CUDA graph captures (forwardVerifyGraphCore(Tensor, int[], Tensor, Tensor, Tensor) with the graph-mode stable buffers and FlashInfer verify kernel) eagerly, so replay output can be compared against identical kernels with no graph involved. Returns an owned copy of the float logits [B, S, V].
      Parameters:
      tokens - window token ids [B, S].
      startPos - cache start position.
      Returns:
      owned logits.

    • forwardVerifyGraphCodeEager

      public Tensor forwardVerifyGraphCodeEager(Tensor tokens, int[] startPositions)
      Ragged form of forwardVerifyGraphCodeEager(Tensor, int).
      Parameters:
      tokens - window token ids [B, S].
      startPositions - per-row cache start position, length B.
      Returns:
      owned logits.

    • closeVerifyGraph

      public void closeVerifyGraph()
      Releases verify CUDA graph resources for this rank.

    • invalidateVerifyCudaGraphs

      public void invalidateVerifyCudaGraphs()
      Drops a captured verify CUDA graph while keeping durable token/RoPE/ logits buffers. Required after KvCachePool.truncateTo(int, int) rebuilds (not bumps-in-place) the FlashInfer CSR tensors a captured graph may still reference (replay → illegal memory access).

    • forwardBatchedVerify

      public Tensor forwardBatchedVerify(Tensor tokens, int[] startPositions, int[] cacheLengths)
      Batched eager verify forward for a cohort of concurrent requests that may be at different absolute positions — the multi-request counterpart of forwardVerifyGraph(Tensor, int), which requires a single shared start position across the batch so it can be CUDA-graph-captured (see VerifyCudaGraph#canGraphVerify). This always runs eager; heterogeneous-position batched cohorts do not graph-capture (a deliberately deferred fast-follow — see the batching plan).
      Parameters:
      tokens - token ids [B, windowLen].
      startPositions - KV write position of the window's first token per row.
      cacheLengths - inclusive cache length per row (prompt + already-sealed generation).
      Returns:
      logits [B, windowLen, V] in float32.

    • forward

      public Tensor forward(Tensor tokens)
      Description copied from interface: Layer
      Forward propagation (or forward pass) through the layer.
      Parameters:
      tokens - the input tensor.
      Returns:
      the output tensor.