Class QwenMtp

java.lang.Object
smile.deep.layer.LayerBlock
smile.llm.model.qwen.QwenMtp
All Implemented Interfaces:
Function<Tensor,Tensor>, Layer
public class QwenMtp extends LayerBlock
Qwen3.5/3.8 native multi-token prediction (MTP) draft head.

Fuses the previous backbone hidden state with the embedding of the last sampled token (pre_fc_norm_* + fc), runs one or more full-attention transformer blocks, applies norm, then scores with the shared lm_head. Multi-step drafts chain this module autoregressively (NEXTN-style).

  • Constructor Details

  • Method Details

    • bindShared

      public void bindShared(EmbeddingLayer embeddings, LinearLayer lmHead, PartialRotaryEncoding.CosSin rope)
      Binds shared backbone embeddings, LM head, and RoPE tables.
      Parameters:
      embeddings - token embeddings.
      lmHead - vocabulary projection.
      rope - partial RoPE cos/sin tables.

    • setKvCachePool

      public void setKvCachePool(KvCachePool pool, boolean closePrevious)
      Installs an MTP-local KV pool and wires full-attention layers.
      Parameters:
      pool - MTP KV pool.
      closePrevious - whether to close a previous pool.

    • kvCacheLayout

      public KvCacheLayout kvCacheLayout()
      Returns the MTP KV layout for a short draft window.
      Returns:
      MTP KV layout for a short draft window.

    • kvCachePool

      public KvCachePool kvCachePool()
      Returns the MTP KV pool.
      Returns:
      MTP KV pool, or null before install.

    • numLayers

      public int numLayers()
      Returns the number of MTP transformer layers.
      Returns:
      number of MTP transformer layers.

    • draftStep

      public Tensor draftStep(Tensor tokenIds, Tensor prevHidden, int position, int draftStep)
      One MTP draft step (single logical request; broadcasts position to every row of a batch — only correct when every row is truly at the same absolute position). Prefer draftStep(Tensor, Tensor, int[], int) for the request-pooled batched path.
      Parameters:
      tokenIds - token ids [B] (last sampled / previous draft).
      prevHidden - previous hidden [B, D] (backbone or prior MTP).
      position - absolute RoPE / cache write position for this step.
      draftStep - zero-based index within the current draft window (MTP KV).
      Returns:
      logits [B, V] in float32 (caller owns).

    • draftStep

      public Tensor draftStep(Tensor tokenIds, Tensor prevHidden, int[] positions, int draftStep)
      One MTP draft step: fuse embedding(token) with previous hidden, run layers, return float logits [B, V] and update lastDraftHidden with the pre-lm-head hidden for the next draft step.
      Parameters:
      tokenIds - token ids [B] (last sampled / previous draft).
      prevHidden - previous hidden [B, D] (backbone or prior MTP).
      positions - absolute RoPE / cache write position per batch row — concurrent requests are generally at different absolute positions, unlike the relative in-round draftStep index below.
      draftStep - zero-based index within the current draft window (MTP KV); uniform across the batch — every row in a round is always at the same relative draft-step.
      Returns:
      logits [B, V] in float32 (caller owns).

    • absorb

      public Tensor absorb(Tensor tokens, Tensor hidden, int startPos, boolean wantLogits)
      History-aware MTP step: fuses m consecutive positions [startPos, startPos+m) of one request, writes their K/V into the (persistent, request-bound) MTP KV pool, and attends over the request's entire MTP history [0, startPos+m) — unlike draftStep(Tensor, Tensor, int, int), whose per-round pool only ever contains the current round's own draft positions (so the head never sees the prompt).

      Position i's input is (token t_{i+1}, backbone hidden h_i) and is stored at RoPE/KV position i, matching vLLM's Qwen3.5 MTP proposer. The caller must have activated the MTP pool for the request. Always refreshes lastDraftHidden() with the last row's post-norm hidden [1, 1, D] (the next chained draft step's input).

      Parameters:
      tokens - next-token ids [1, m] (int64, on this device).
      hidden - backbone/MTP hidden [1, m, D].
      startPos - absolute position of the first row.
      wantLogits - whether to score the last row with the shared LM head.
      Returns:
      float logits [1, 1, V] for the last row, or null.

    • lastDraftHidden

      public Tensor lastDraftHidden()
      Returns the last draft hidden from the most recent draftStep(Tensor, Tensor, int, int).
      Returns:
      last draft hidden [B, 1, D] or [B, D], or null.

    • beginRound

      public void beginRound(int draftWindow)
      Clears MTP KV for a new single-request speculative round (bind capacity for the draft window). Prefer beginRound(int, int) for the request-pooled batched path.
      Parameters:
      draftWindow - number of draft steps to reserve.

    • beginRound

      public void beginRound(int batchSize, int draftWindow)
      Clears MTP KV for a new speculative round sized for batchSize concurrent requests, each reserving draftWindow + 1 slots.
      Parameters:
      batchSize - number of concurrent requests in this round's cohort.
      draftWindow - number of draft steps to reserve per request.

    • absorbBatch

      public Tensor absorbBatch(Tensor tokens, Tensor hidden, int[] startPositions, int[] lastRealRow)
      Batched, history-aware MTP step for a cohort of concurrent requests, each at its own absolute position. Writes K/V for every column of every row into the persistent request-bound pool and attends over each row's whole history, like absorb(Tensor, Tensor, int, boolean).

      Rows are padded to a common width S: row i has lastRealRow[i]+1 real columns (positions startPositions[i]..) followed by dummy columns. Dummy columns come after the real ones, so causal attention keeps real columns unaffected; their K/V lands in the row's speculative region (past its committed history) and is overwritten before it can be read. The logits and chained hidden are taken at each row's last real column.

      Parameters:
      tokens - next-token ids [B, S] (int64, on this device).
      hidden - backbone/MTP hidden [B, S, D].
      startPositions - absolute KV/RoPE position of column 0, per row.
      lastRealRow - index of each row's last real column, in [0, S).
      Returns:
      float logits [B, 1, V] for each row's last real column. The chained hidden [B, 1, D] is available from lastDraftHidden().

    • clearDraftHidden

      public void clearDraftHidden()
      Drops the chained draft hidden (history-aware path keeps the KV pool bound).

    • endRound

      public void endRound()
      Releases MTP draft KV after a speculative round.

    • to

      public QwenMtp to(Device device)
      Description copied from interface: Layer
      Moves the layer block to a device.
      Specified by:
      to in interface Layer
      Overrides:
      to in class LayerBlock
      Parameters:
      device - the compute device.
      Returns:
      this layer.

    • to

      public QwenMtp to(Device device, ScalarType dtype)
      Description copied from interface: Layer
      Moves the layer block to a device.
      Specified by:
      to in interface Layer
      Overrides:
      to in class LayerBlock
      Parameters:
      device - the compute device.
      dtype - the data type.
      Returns:
      this layer.

    • forward

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