@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e93f4cc9e3

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Add the support for the qwen3 next model (a hybrid attention model). (#24526)

Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Tao He committed Sep 11, 2025 at 15:32 UTC e93f4cc9e37484009f74e15d3111a1f335c532a5

29 files changed +2476 -61

.yapfignore

@@ -1 +1,2 @@

1 collect_env.py

2 +vllm/model_executor/layers/fla/ops/*.py

docs/models/supported_models.md

@@ -403,6 +403,7 @@ th {

403 | `Qwen2MoeForCausalLM` | Qwen2MoE | `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |

404 | `Qwen3ForCausalLM` | Qwen3 | `Qwen/Qwen3-8B`, etc. | ✅︎ | ✅︎ | ✅︎ |

405 | `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ | ✅︎ |

406 +| `Qwen3NextForCausalLM` | Qwen3.5MoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |

407 | `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |

409 | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ | ✅︎ |

pyproject.toml

@@ -228,6 +228,7 @@ fo = "fo"

228 ba = "ba"

229

230 [tool.typos.type.py.extend-words]

231 +ba = "ba"

232

233 [tool.typos.type.cpp]

234 extend-glob = ["*.cu"]

tests/models/registry.py

+5 -1

@@ -326,6 +326,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

326 "Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),

327 "Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),

328 "Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),

329 + "Qwen3NextForCausalLM": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct",

330 + min_transformers_version="4.56.2"),

331 "RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),

332 "SeedOssForCausalLM": _HfExamplesInfo("ByteDance-Seed/Seed-OSS-36B-Instruct", # noqa: E501

333 trust_remote_code=True,

@@ -640,7 +642,9 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

642 is_available_online=False),

643 "MiMoMTPModel": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",

644 trust_remote_code=True,

643 - speculative_model="XiaomiMiMo/MiMo-7B-RL")

645 + speculative_model="XiaomiMiMo/MiMo-7B-RL"),

646 + "Qwen3NextMTP": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct",

647 + min_transformers_version="4.56.2"),

648 }

649

650 _TRANSFORMERS_BACKEND_MODELS = {

vllm/config/__init__.py

+46 -10

@@ -1508,7 +1508,8 @@ class ModelConfig:

1508 if (self.hf_text_config.model_type == "deepseek_mtp"

1509 or self.hf_config.model_type == "mimo_mtp"

1510 or self.hf_config.model_type == "glm4_moe_mtp"

1511 - or self.hf_config.model_type == "ernie_mtp"):

1511 + or self.hf_config.model_type == "ernie_mtp"

1512 + or self.hf_config.model_type == "qwen3_next_mtp"):

1513 total_num_hidden_layers = getattr(self.hf_text_config,

1514 "num_nextn_predict_layers", 0)

1515 else:

@@ -1571,15 +1572,28 @@ class ModelConfig:

1572 if attn_type_list:

1573 return sum(t == 1 for t in attn_type_list[start:end])

1574

1574 - if layers_block_type_value is None and attn_type_list is None:

1575 + # Hybrid model Qwen3Next

1576 + layer_types_value = getattr(self.hf_config, "layer_types", None)

1577 + if layer_types_value is not None:

1578 + if getattr(block_type, "value", block_type) == "attention":

1579 + return sum(t == "full_attention"

1580 + for t in layer_types_value[start:end])

1581 + elif getattr(block_type, "value",

1582 + block_type) == "linear_attention":

1583 + return sum(t == "linear_attention"

1584 + for t in layer_types_value[start:end])

1585 + else:

1586 + return sum(t == getattr(block_type, "value", block_type)

1587 + for t in layer_types_value[start:end])

1588 +

1589 + if (layers_block_type_value is None and attn_type_list is None

1590 + and layer_types_value is None):

1591 raise ValueError(

1592 "The model is an hybrid without a"

1577 - "layers_block_type or an attn_type_list in the hf_config,"

1578 - "cannot determine the num of "

1593 + "layers_block_type or an attn_type_list, or a layer_types "

1594 + "in the hf_config, cannot determine the num of "

1595 f"{block_type.value} layers")

1596

1581 - return sum(t == 1 for t in attn_type_list[start:end])

1582 -

1597 def get_mamba_chunk_size(self) -> Optional[int]:

1598 """

1599 Returns the mamba chunk size if it exists

@@ -1866,7 +1880,7 @@ class DeviceConfig:

1880

1881 SpeculativeMethod = Literal["ngram", "eagle", "eagle3", "medusa",

1882 "mlp_speculator", "draft_model", "deepseek_mtp",

1869 - "ernie_mtp"]

1883 + "ernie_mtp", "qwen3_next_mtp"]

1884

1885

1886 @config

@@ -2007,7 +2021,15 @@ class SpeculativeConfig:

2021 "n_predict": n_predict,

2022 "architectures": ["ErnieMTPModel"]

2023 })

2010 - return hf_config

2024 +

2025 + if hf_config.model_type == "qwen3_next":

2026 + hf_config.model_type = "qwen3_next_mtp"

2027 + if hf_config.model_type == "qwen3_next_mtp":

2028 + n_predict = getattr(hf_config, "num_nextn_predict_layers", None)

2029 + hf_config.update({

2030 + "n_predict": n_predict,

2031 + "architectures": ["Qwen3NextMTP"]

2032 + })

2033

2034 return hf_config

2035

@@ -2028,9 +2050,13 @@ class SpeculativeConfig:

2050 (self.target_model_config.hf_text_config.model_type \

2051 == "deepseek_v3" or

2052 self.target_model_config.hf_text_config.model_type in

2031 - ("mimo","ernie4_5_moe")):

2053 + ("mimo","ernie4_5_moe", "qwen3_next")):

2054 # use the draft model from the same model:

2055 self.model = self.target_model_config.model

2056 + # Align the quantization of draft model for cases such as

2057 + # --quantization fp8 with a bf16 checkpoint.

2058 + if not self.quantization:

2059 + self.quantization = self.target_model_config.quantization

2060 elif self.method in ("ngram", "[ngram]"):

2061 self.model = "ngram"

2062 else:

@@ -2140,6 +2166,15 @@ class SpeculativeConfig:

2166 "one layer. Might need some code changes " \

2167 "to support multiple layers."

2168 )

2169 + elif (self.draft_model_config.hf_config.model_type ==

2170 + "qwen3_next_mtp"):

2171 + self.method = "qwen3_next_mtp"

2172 + if self.num_speculative_tokens > 1:

2173 + logger.warning(

2174 + "All Qwen3Next MTP models only have " \

2175 + "one layer. Might need some code changes " \

2176 + "to support multiple layers."

2177 + )

2178 else:

2179 self.method = "draft_model"

2180 raise NotImplementedError(

@@ -2355,7 +2390,8 @@ class SpeculativeConfig:

2390 return self.num_speculative_tokens

2391

2392 def use_eagle(self) -> bool:

2358 - return self.method in ("eagle", "eagle3", "deepseek_mtp", "ernie_mtp")

2393 + return self.method in ("eagle", "eagle3", "deepseek_mtp", "ernie_mtp",

2394 + "qwen3_next_mtp")

2395

2396 def __repr__(self) -> str:

2397 method = self.method

vllm/config/compilation.py

@@ -341,6 +341,7 @@ class CompilationConfig:

341 "vllm.short_conv",

342 "vllm.linear_attention",

343 "vllm.plamo2_mamba_mixer",

344 + "vllm.gdn_attention",

345 ]

346

347 def compute_hash(self) -> str:

vllm/model_executor/layers/fla/ops/chunk_delta_h.py

+3 -2

@@ -14,7 +14,7 @@ import torch

14 from vllm.triton_utils import tl, triton

16 from .index import prepare_chunk_indices, prepare_chunk_offsets

17 -from .op import exp, safe_exp

17 +from .op import exp

18 from .utils import is_nvidia_hopper, use_cuda_graph

20 NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16]

@@ -175,12 +175,13 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(

175 boundary_check=(0, 1))

176

177 if USE_G:

178 + m_t = (i_t * BT + tl.arange(0, BT)) < T

179 last_idx = min((i_t + 1) * BT, T) - 1

180 b_g_last = tl.load(g + bos * H + last_idx * H + i_h)

181 p_g = tl.make_block_ptr(g + bos * H + i_h, (T, ), (H, ),

182 (i_t * BT, ), (BT, ), (0, ))

183 b_g = tl.load(p_g, boundary_check=(0, ))

183 - b_v_new = b_v_new * safe_exp(b_g_last - b_g)[:, None]

184 + b_v_new = b_v_new * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None]

185 b_g_last = exp(b_g_last)

186 b_h1 = b_h1 * b_g_last

187 if K > 64:

vllm/model_executor/layers/fla/ops/chunk_o.py

+5 -4

@@ -16,7 +16,7 @@ import torch

16 from vllm.triton_utils import tl, triton

18 from .index import prepare_chunk_indices

19 -from .op import exp, safe_exp

19 +from .op import exp

20 from .utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper

22 BKV_LIST = [64, 128] if check_shared_mem() else [32, 64]

@@ -112,10 +112,11 @@ def chunk_fwd_kernel_o(

112 p_g = tl.make_block_ptr(g, (T, ), (H, ), (i_t * BT, ), (BT, ), (0, ))

113 b_g = tl.load(p_g, boundary_check=(0, ))

114 b_o = b_o * exp(b_g)[:, None]

115 - b_A = b_A * safe_exp(b_g[:, None] - b_g[None, :])

115 + b_A = b_A * exp(b_g[:, None] - b_g[None, :])

116

117 - o_i = tl.arange(0, BT)

118 - m_A = o_i[:, None] >= o_i[None, :]

117 + o_t = i_t * BT + tl.arange(0, BT)

118 + m_t = o_t < T

119 + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)

120 b_A = tl.where(m_A, b_A, 0)

121

122 p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV),

vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py

+6 -4

@@ -14,7 +14,7 @@ import torch

14 from vllm.triton_utils import tl, triton

16 from .index import prepare_chunk_indices

17 -from .op import safe_exp

17 +from .op import exp

20 @triton.heuristics({

@@ -56,7 +56,8 @@ def chunk_scaled_dot_kkt_fwd_kernel(

56 T = eos - bos

57 else:

58 bos, eos = i_b * T, i_b * T + T

59 - o_t = tl.arange(0, BT)

59 + o_t = i_t * BT + tl.arange(0, BT)

60 + m_t = o_t < T

62 p_beta = tl.make_block_ptr(beta + bos * H + i_h, (T, ), (H, ),

63 (i_t * BT, ), (BT, ), (0, ))

@@ -76,9 +77,10 @@ def chunk_scaled_dot_kkt_fwd_kernel(

77 (i_t * BT, ), (BT, ), (0, ))

78 b_g = tl.load(p_g, boundary_check=(0, ))

79 b_g_diff = b_g[:, None] - b_g[None, :]

79 - b_A = b_A * safe_exp(b_g_diff)

80 + b_A = b_A * exp(b_g_diff)

81 - b_A = tl.where(o_t[:, None] > o_t[None, :], b_A, 0)

82 + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)

83 + b_A = tl.where(m_A, b_A, 0)

84 p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1),

85 (i_t * BT, 0), (BT, BT), (1, 0))

86 tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))

vllm/model_executor/layers/fla/ops/fused_recurrent.py

+2 -2

@@ -116,8 +116,8 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(

116 b_g = tl.load(p_g).to(tl.float32)

117

118 if USE_QK_L2NORM_IN_KERNEL:

119 - b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6)

120 - b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6)

119 + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)

120 + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)

121 b_q = b_q * scale

122 # [BK, BV]

123 b_h *= exp(b_g)

vllm/model_executor/layers/fla/ops/l2norm.py

+1 -1

@@ -78,7 +78,7 @@ def l2norm_fwd_kernel2(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr):

78 row_idx = xoffset + tl.arange(0, MBLOCK)[:, None]

79 xmask = row_idx < M

80 rindex = tl.arange(0, N)[None, :]

81 - xs = tl.load(X + (rindex + N * row_idx), None).to(tl.float32)

81 + xs = tl.load(X + (rindex + N * row_idx), xmask).to(tl.float32)

82 square = tl.broadcast_to(xs * xs, [MBLOCK, N])

83 square_sum = tl.sum(tl.where(xmask, square, 0), 1)[:, None]

84 rsqrt = tl.rsqrt(square_sum + eps)

vllm/model_executor/layers/fla/ops/op.py

-5

@@ -28,11 +28,6 @@ else:

28 log2 = tl.log2

31 -@triton.jit

32 -def safe_exp(x):

33 - return exp(tl.where(x <= 0, x, float('-inf')))

34 -

35 -

31 if not hasattr(tl, 'gather'):

33 @triton.jit

vllm/model_executor/layers/mamba/mamba_utils.py

+37

@@ -70,6 +70,15 @@ class MambaStateDtypeCalculator:

70 model_dtype)

71 return (conv_state_dtype, )

73 + @classmethod

74 + def gated_delta_net_state_dtype(

75 + cls,

76 + model_dtype: Union[ModelDType, torch.dtype],

77 + mamba_cache_dtype: MambaDType,

78 + ) -> tuple[torch.dtype, torch.dtype]:

79 + state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)

80 + return (state_dtype, state_dtype)

81 +

83 class MambaStateShapeCalculator:

@@ -163,3 +172,31 @@ class MambaStateShapeCalculator:

172

173 # for n_groups == 1, this is exactly tp_size - n_groups

174 return tp_size - ngroups

175 +

176 + @classmethod

177 + def gated_delta_net_state_shape(

178 + cls,

179 + tp_world_size: int,

180 + num_k_heads: int,

181 + num_v_heads: int,

182 + head_k_dim: int,

183 + head_v_dim: int,

184 + conv_kernel_size: int,

185 + num_spec: int = 0,

186 + use_v1: bool = True,

187 + ):

188 + conv_dim = (head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads)

189 + conv_state_shape = (

190 + divide(conv_dim, tp_world_size),

191 + conv_kernel_size - 1 + num_spec,

192 + )

193 +

194 + # In V0, the conv_state shape was swapped during allocation in

195 + # MambaCacheManager, but in V1 it needs to be determined here at the

196 + # calculation level

197 + if use_v1:

198 + conv_state_shape = conv_state_shape[1], conv_state_shape[0]

199 +

200 + temporal_state_shape = (divide(num_v_heads,

201 + tp_world_size), head_k_dim, head_v_dim)

202 + return conv_state_shape, temporal_state_shape

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

+41 -7

@@ -464,7 +464,9 @@ def causal_conv1d_fn(

464 # 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx]

465 # 4. computation can be skipped if cache_indices[idx] == pad_slot_id

466 num_cache_lines = conv_states.size(0)

467 - assert (num_cache_lines, dim, width - 1) == conv_states.shape

467 + assert (num_cache_lines == conv_states.shape[0]

468 + and dim == conv_states.shape[1]

469 + and width - 1 <= conv_states.shape[2])

470 stride_istate_seq = conv_states.stride(0)

471 stride_istate_dim = conv_states.stride(1)

472 stride_istate_token = conv_states.stride(2)

@@ -623,6 +625,7 @@ def _causal_conv1d_update_kernel(

625 conv_state_ptr,

626 cache_seqlens_ptr, # circular buffer

627 conv_state_indices_ptr,

628 + num_accepted_tokens_ptr,

629 o_ptr, # (batch, dim, seqlen)

630 # Matrix dimensions

631 batch: int,

@@ -639,6 +642,7 @@ def _causal_conv1d_update_kernel(

642 stride_conv_state_seq: tl.constexpr,

643 stride_conv_state_dim: tl.constexpr,

644 stride_conv_state_tok: tl.constexpr,

645 + stride_state_indices: tl.constexpr,

646 stride_o_seq: tl.constexpr,

647 stride_o_dim: tl.constexpr,

648 stride_o_token: tl.constexpr,

@@ -649,6 +653,7 @@ def _causal_conv1d_update_kernel(

653 KERNEL_WIDTH: tl.constexpr,

654 SILU_ACTIVATION: tl.constexpr,

655 IS_CONTINUOUS_BATCHING: tl.constexpr,

656 + IS_SPEC_DECODING: tl.constexpr,

657 NP2_STATELEN: tl.constexpr,

658 USE_PAD_SLOT: tl.constexpr,

659 BLOCK_N: tl.constexpr,

@@ -663,8 +668,9 @@ def _causal_conv1d_update_kernel(

668

669 if IS_CONTINUOUS_BATCHING:

670 # mask = idx_seq < batch

666 - conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq).to(

667 - tl.int64)

671 + conv_state_batch_coord = tl.load(conv_state_indices_ptr +

672 + idx_seq * stride_state_indices).to(

673 + tl.int64)

674 else:

675 conv_state_batch_coord = idx_seq

676 if USE_PAD_SLOT: # noqa

@@ -672,13 +678,32 @@ def _causal_conv1d_update_kernel(

678 # not processing as this is not the actual sequence

679 return

680

681 + if IS_SPEC_DECODING:

682 + # The rolling of conv state:

683 + #

684 + # Before forward, the conv_state is:

685 + # [history1, history2, ..., historyM].

686 + #

687 + # After forward, the conv_state becomes:

688 + # [history2, ..., historyM, draft1, draft2, ..., draftN].

689 + #

690 + # After acceptance, it becomes:

691 + #

692 + # - accept 1 tokens: [history2, ..., historyM, draft1]

693 + # - accept 2 tokens: [history3, ..., historyM, draft1, draft2]

694 + # - and so on.

695 + conv_state_token_offset = (tl.load(num_accepted_tokens_ptr + idx_seq) -

696 + 1)

697 + else:

698 + conv_state_token_offset = 0

699 +

700 # STEP 1: READ init_state data

701 conv_states_base = (conv_state_ptr +

702 (conv_state_batch_coord * stride_conv_state_seq) +

703 (idx_feats * stride_conv_state_dim))

704 mask_w = idx_feats < dim

705

681 - prior_tokens = conv_states_base

706 + prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok

707 if KERNEL_WIDTH >= 2:

708 conv_states_ptrs = prior_tokens # [BLOCK_N]

709 col0 = tl.load(conv_states_ptrs, mask_w, 0.0)

@@ -695,10 +720,14 @@ def _causal_conv1d_update_kernel(

720 # STEP 2: assume state_len > seqlen

721 idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M]

722

723 + # The conv_state updates works in a sliding window manner,

724 + # at each forward pass, the tokens are shift by 1, so we

725 + # load since idx_tokens + 1.

726 conv_state_ptrs_source = (

727 conv_state_ptr + (conv_state_batch_coord * stride_conv_state_seq) +

728 + conv_state_token_offset * stride_conv_state_tok +

729 (idx_feats * stride_conv_state_dim)[None, :] +

701 - ((idx_tokens + seqlen) * stride_conv_state_tok)[:, None]

730 + ((idx_tokens + 1) * stride_conv_state_tok)[:, None]

731 ) # [BLOCK_M, BLOCK_N]

732 mask = ((conv_state_batch_coord < num_cache_lines)

733 & ((idx_tokens + seqlen) < state_len)[:, None]

@@ -820,6 +849,7 @@ def causal_conv1d_update(

849 activation: Union[bool, str, None] = None,

850 cache_seqlens: Optional[torch.Tensor] = None,

851 conv_state_indices: Optional[torch.Tensor] = None,

852 + num_accepted_tokens: Optional[torch.Tensor] = None,

853 pad_slot_id: int = PAD_SLOT_ID,

854 metadata=None,

855 validate_data=False,

@@ -890,10 +920,11 @@ def causal_conv1d_update(

920 ) # X (batch, dim, seqlen)

921

922 stride_o_seq, stride_o_dim, stride_o_token = out.stride()

893 -

923 stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride(

924 )

896 - state_len = width - 1

925 + stride_state_indices = conv_state_indices.stride(

926 + 0) if conv_state_indices is not None else 0

927 + state_len = width - 1 + (seqlen - 1) # effective state_len needed

928 np2_statelen = triton.next_power_of_2(state_len)

929

930 def grid(META):

@@ -910,6 +941,7 @@ def causal_conv1d_update(

941 conv_state,

942 cache_seqlens,

943 conv_state_indices,

944 + num_accepted_tokens,

945 out,

946 # Matrix dimensions

947 batch,

@@ -926,6 +958,7 @@ def causal_conv1d_update(

958 stride_istate_seq,

959 stride_istate_dim,

960 stride_istate_token,

961 + stride_state_indices,

962 stride_o_seq,

963 stride_o_dim,

964 stride_o_token,

@@ -936,6 +969,7 @@ def causal_conv1d_update(

969 KERNEL_WIDTH=width,

970 SILU_ACTIVATION=activation in ["silu", "swish"],

971 IS_CONTINUOUS_BATCHING=conv_state_indices is not None,

972 + IS_SPEC_DECODING=num_accepted_tokens is not None,

973 NP2_STATELEN=np2_statelen,

974 USE_PAD_SLOT=pad_slot_id is not None,

975 BLOCK_N=256,

vllm/model_executor/models/config.py

+2 -1

@@ -312,7 +312,8 @@ class MambaModelConfig(VerifyAndUpdateConfig):

312

313 # TODO(tdoublep): remove as full cuda graph support is added

314 FCG_NOT_SUPPORTED_MODELS = [

315 - "Lfm2ForCausalLM", "MiniMaxText01ForCausalLM"

315 + "Lfm2ForCausalLM",

316 + "MiniMaxText01ForCausalLM",

317 ]

318

319 if (model_config.architecture not in FCG_NOT_SUPPORTED_MODELS

vllm/model_executor/models/qwen3_next.py new

+1294

@@ -0,0 +1,1294 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Inference-only Qwen3Next model."""

4 +from collections.abc import Iterable

5 +from typing import Optional

6 +

7 +import torch

8 +import torch.nn.functional as F

9 +from einops import rearrange

10 +from torch import nn

11 +from transformers.activations import ACT2FN

12 +

13 +from vllm import envs

14 +from vllm.attention import Attention, AttentionBackend, AttentionMetadata

15 +from vllm.compilation.decorators import support_torch_compile

16 +from vllm.config import (CacheConfig, ModelConfig, SpeculativeConfig,

17 + VllmConfig, get_current_vllm_config)

18 +from vllm.distributed import (divide, get_ep_group, get_pp_group,

19 + get_tensor_model_parallel_rank,

20 + get_tensor_model_parallel_world_size)

21 +from vllm.forward_context import ForwardContext, get_forward_context

22 +from vllm.logger import init_logger

23 +from vllm.model_executor.layers.fla.ops import (

24 + RMSNormGated, chunk_gated_delta_rule, fused_recurrent_gated_delta_rule)

25 +from vllm.model_executor.layers.fused_moe import FusedMoE

26 +# yapf conflicts with isort for this block

27 +# yapf: disable

28 +from vllm.model_executor.layers.layernorm import (

29 + GemmaRMSNorm as Qwen3NextRMSNorm)

30 +# yapf: enable

31 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

32 + MergedColumnParallelLinear,

33 + QKVParallelLinear,

34 + ReplicatedLinear,

35 + RowParallelLinear)

36 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

37 +from vllm.model_executor.layers.mamba.abstract import MambaBase

38 +from vllm.model_executor.layers.mamba.mamba_mixer2 import (

39 + mamba_v2_sharded_weight_loader)

40 +from vllm.model_executor.layers.mamba.mamba_utils import (

41 + MambaStateDtypeCalculator, MambaStateShapeCalculator)

42 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (

43 + causal_conv1d_fn, causal_conv1d_update)

44 +from vllm.model_executor.layers.quantization import QuantizationConfig

45 +from vllm.model_executor.layers.quantization.gptq import GPTQConfig

46 +from vllm.model_executor.layers.quantization.gptq_marlin import (

47 + GPTQMarlinConfig)

48 +from vllm.model_executor.layers.rotary_embedding import get_rope

49 +from vllm.model_executor.layers.vocab_parallel_embedding import (

50 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

51 +from vllm.model_executor.model_loader.weight_utils import (

52 + default_weight_loader, sharded_weight_loader)

53 +from vllm.model_executor.models.mamba_cache import MambaCacheParams

54 +from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP

55 +from vllm.model_executor.sampling_metadata import SamplingMetadata

56 +from vllm.model_executor.utils import set_weight_attrs

57 +from vllm.platforms import current_platform

58 +from vllm.sequence import IntermediateTensors

59 +from vllm.transformers_utils.configs import Qwen3NextConfig

60 +from vllm.triton_utils import tl, triton

61 +from vllm.utils import direct_register_custom_op

62 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

63 +

64 +from .interfaces import (HasInnerState, IsHybrid, MixtureOfExperts,

65 + SupportsLoRA, SupportsPP)

66 +from .utils import (AutoWeightsLoader, PPMissingLayer, extract_layer_index,

67 + is_pp_missing_parameter,

68 + make_empty_intermediate_tensors_factory, make_layers,

69 + maybe_prefix)

70 +

71 +logger = init_logger(__name__)

72 +

73 +KVCache = tuple[torch.Tensor, torch.Tensor]

74 +

75 +

76 +class Qwen3NextSparseMoeBlock(nn.Module):

77 +

78 + def __init__(

79 + self,

80 + config: Qwen3NextConfig,

81 + quant_config: Optional[QuantizationConfig] = None,

82 + prefix: str = "",

83 + enable_eplb: bool = False,

84 + ):

85 + super().__init__()

86 + self.tp_size = get_tensor_model_parallel_world_size()

87 +

88 + self.ep_group = get_ep_group().device_group

89 + self.ep_rank = self.ep_group.rank()

90 + self.ep_size = self.ep_group.size()

91 + self.n_routed_experts = config.num_experts

92 +

93 + if self.tp_size > config.num_experts:

94 + raise ValueError(

95 + f"Tensor parallel size {self.tp_size} is greater than "

96 + f"the number of experts {config.num_experts}.")

97 +

98 + # Load balancing settings.

99 + vllm_config = get_current_vllm_config()

100 + eplb_config = vllm_config.parallel_config.eplb_config

101 + self.enable_eplb = enable_eplb

102 +

103 + self.n_logical_experts = self.n_routed_experts

104 + self.n_redundant_experts = eplb_config.num_redundant_experts

105 + self.n_physical_experts = (self.n_logical_experts +

106 + self.n_redundant_experts)

107 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size

108 +

109 + self.physical_expert_start = (self.ep_rank *

110 + self.n_local_physical_experts)

111 + self.physical_expert_end = (self.physical_expert_start +

112 + self.n_local_physical_experts)

113 +

114 + self.experts = FusedMoE(num_experts=self.n_routed_experts,

115 + top_k=config.num_experts_per_tok,

116 + hidden_size=config.hidden_size,

117 + intermediate_size=config.moe_intermediate_size,

118 + reduce_results=False,

119 + renormalize=config.norm_topk_prob,

120 + quant_config=quant_config,

121 + prefix=f"{prefix}.experts",

122 + enable_eplb=self.enable_eplb,

123 + num_redundant_experts=self.n_redundant_experts)

124 +

125 + self.gate = ReplicatedLinear(

126 + config.hidden_size,

127 + config.num_experts,

128 + bias=False,

129 + quant_config=self._maybe_ignore_quant_config(quant_config),

130 + prefix=f"{prefix}.gate")

131 +

132 + if config.shared_expert_intermediate_size > 0:

133 + self.shared_expert = Qwen3NextMLP(

134 + hidden_size=config.hidden_size,

135 + intermediate_size=config.shared_expert_intermediate_size,

136 + hidden_act=config.hidden_act,

137 + quant_config=quant_config,

138 + reduce_results=self.experts.must_reduce_shared_expert_outputs(

139 + ),

140 + )

141 + else:

142 + self.shared_expert = None

143 + self.shared_expert_gate = torch.nn.Linear(config.hidden_size,

144 + 1,

145 + bias=False)

146 +

147 + def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):

148 + # GPTQ configs do not have a list of ignored modules, however AutoGPTQ

149 + # seems to avoid gate quantization.

150 + # See: https://huggingface.co/Qwen/Qwen3-30B-A3B-GPTQ-Int4

151 + if isinstance(quant_config, (GPTQConfig, GPTQMarlinConfig)):

152 + return None

153 + return quant_config

154 +

155 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

156 + # NOTE: hidden_states can have either 1D or 2D shape.

157 + orig_shape = hidden_states.shape

158 + hidden_dim = hidden_states.shape[-1]

159 + hidden_states = hidden_states.view(-1, hidden_dim)

160 +

161 + shared_output = None

162 + if self.shared_expert is not None:

163 + shared_output = self.shared_expert(hidden_states)

164 + if self.shared_expert_gate is not None:

165 + shared_output = F.sigmoid(

166 + self.shared_expert_gate(hidden_states)) * shared_output

167 +

168 + # router_logits: (num_tokens, n_experts)

169 + router_logits, _ = self.gate(hidden_states)

170 + final_hidden_states = self.experts(hidden_states=hidden_states,

171 + router_logits=router_logits)

172 +

173 + if shared_output is not None:

174 + final_hidden_states = final_hidden_states + shared_output

175 + if self.tp_size > 1:

176 + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501

177 + final_hidden_states)

178 +

179 + return final_hidden_states.view(orig_shape)

180 +

181 +

182 +class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

183 +

184 + @property

185 + def mamba_type(self) -> str:

186 + return "linear_attention"

187 +

188 + def get_attn_backend(self) -> type["AttentionBackend"]:

189 + from vllm.v1.attention.backends.gdn_attn import GDNAttentionBackend

190 + return GDNAttentionBackend

191 +

192 + def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:

193 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(

194 + self.model_config.dtype, self.cache_config.mamba_cache_dtype)

195 +

196 + def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:

197 + return MambaStateShapeCalculator.gated_delta_net_state_shape(

198 + self.tp_size,

199 + self.num_k_heads,

200 + self.num_v_heads,

201 + self.head_k_dim,

202 + self.head_v_dim,

203 + self.conv_kernel_size,

204 + self.num_spec,

205 + use_v1=True)

206 +

207 + def __init__(

208 + self,

209 + config: Qwen3NextConfig,

210 + model_config: Optional[ModelConfig] = None,

211 + cache_config: Optional[CacheConfig] = None,

212 + quant_config: Optional[QuantizationConfig] = None,

213 + speculative_config: Optional[SpeculativeConfig] = None,

214 + prefix: str = "",

215 + ) -> None:

216 + super().__init__()

217 + self.tp_size = get_tensor_model_parallel_world_size()

218 + self.tp_rank = get_tensor_model_parallel_rank()

219 + self.hidden_size = config.hidden_size

220 + self.num_v_heads = config.linear_num_value_heads

221 + self.num_k_heads = config.linear_num_key_heads

222 + self.head_k_dim = config.linear_key_head_dim

223 + self.head_v_dim = config.linear_value_head_dim

224 + self.key_dim = self.head_k_dim * self.num_k_heads

225 + self.value_dim = self.head_v_dim * self.num_v_heads

226 +

227 + self.conv_kernel_size = config.linear_conv_kernel_dim

228 + self.layer_idx = extract_layer_index(prefix)

229 + self.activation = config.hidden_act

230 + self.act = ACT2FN[config.hidden_act]

231 + self.layer_norm_epsilon = config.rms_norm_eps

232 + self.prefix = prefix

233 +

234 + self.config = config

235 + self.model_config = model_config

236 + self.cache_config = cache_config

237 + self.quant_config = quant_config

238 + self.speculative_config = speculative_config

239 + self.num_spec = (self.speculative_config.num_speculative_tokens

240 + if self.speculative_config else 0)

241 +

242 + # QKV

243 + self.conv_dim = self.key_dim * 2 + self.value_dim

244 + self.conv1d = ColumnParallelLinear(

245 + input_size=self.conv_kernel_size,

246 + output_size=self.conv_dim,

247 + bias=False,

248 + prefix=f"{prefix}.conv1d",

249 + )

250 + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

251 +

252 + # projection of the input hidden states

253 + self.projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2

254 + self.projection_size_ba = self.num_v_heads * 2

255 + self.in_proj = MergedColumnParallelLinear(

256 + input_size=self.hidden_size,

257 + output_sizes=[self.projection_size_qkvz, self.projection_size_ba],

258 + bias=False,

259 + quant_config=quant_config,

260 + prefix=f"{prefix}.in_proj",

261 + )

262 +

263 + query_key_settings = (self.key_dim, 0, False)

264 + value_settings = (self.value_dim, 0, False)

265 +

266 + delattr(self.conv1d.weight, "weight_loader")

267 + set_weight_attrs(

268 + self.conv1d.weight, {

269 + "weight_loader":

270 + mamba_v2_sharded_weight_loader([

271 + query_key_settings,

272 + query_key_settings,

273 + value_settings,

274 + ], self.tp_size, self.tp_rank)

275 + })

276 +

277 + # selective projection used to make dt, B and C input dependant

278 +

279 + # time step projection (discretization)

280 + # instantiate once and copy inv_dt in init_weights of PretrainedModel

281 + self.dt_bias = nn.Parameter(

282 + torch.ones(self.num_v_heads // self.tp_size), )

283 + self.A_log = nn.Parameter(

284 + torch.empty(

285 + divide(self.num_v_heads, self.tp_size),

286 + dtype=torch.float32,

287 + ))

288 +

289 + set_weight_attrs(self.A_log,

290 + {"weight_loader": sharded_weight_loader(0)})

291 + set_weight_attrs(self.dt_bias,

292 + {"weight_loader": sharded_weight_loader(0)})

293 +

294 + self.norm = RMSNormGated(

295 + self.head_v_dim,

296 + eps=self.layer_norm_epsilon,

297 + group_size=None,

298 + norm_before_gate=True,

299 + device=torch.cuda.current_device(),

300 + dtype=config.torch_dtype,

301 + )

302 +

303 + self.out_proj = RowParallelLinear(self.value_dim,

304 + self.hidden_size,

305 + bias=False,

306 + input_is_parallel=True,

307 + quant_config=quant_config,

308 + prefix=f"{prefix}.out_proj")

309 +

310 + compilation_config = get_current_vllm_config().compilation_config

311 + if prefix in compilation_config.static_forward_context:

312 + raise ValueError(f"Duplicate layer name: {prefix}")

313 + compilation_config.static_forward_context[prefix] = self

314 +

315 + def fix_query_key_value_ordering(

316 + self,

317 + mixed_qkvz,

318 + mixed_ba,

319 + ):

320 + """

321 + Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.

322 + """

323 + new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (

324 + self.num_k_heads // self.tp_size,

325 + (self.head_k_dim + self.head_k_dim +

326 + (self.head_v_dim + self.head_v_dim) * self.num_v_heads //

327 + self.num_k_heads),

328 + )

329 + new_tensor_shape_ba = mixed_qkvz.size()[:-1] + (

330 + self.num_k_heads // self.tp_size,

331 + 2 * self.num_v_heads // self.num_k_heads,

332 + )

333 +

334 + mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)

335 + mixed_ba = mixed_ba.view(*new_tensor_shape_ba)

336 +

337 + split_arg_list_qkvz = [

338 + self.head_k_dim,

339 + self.head_k_dim,

340 + (self.num_v_heads // self.num_k_heads * self.head_v_dim),

341 + (self.num_v_heads // self.num_k_heads * self.head_v_dim),

342 + ]

343 + split_arg_list_ba = [

344 + self.num_v_heads // self.num_k_heads,

345 + self.num_v_heads // self.num_k_heads

346 + ]

347 +

348 + # [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)]

349 + # --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn],

350 + # [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng]

351 + (query, key, value, z) = torch.split(mixed_qkvz,

352 + split_arg_list_qkvz,

353 + dim=2)

354 + (b, a) = torch.split(mixed_ba, split_arg_list_ba, dim=2)

355 +

356 + # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn]

357 + value = value.reshape(value.size(0), -1, self.head_v_dim)

358 + z = z.reshape(z.size(0), -1, self.head_v_dim)

359 + b = b.reshape(b.size(0), self.num_v_heads // self.tp_size)

360 + a = a.reshape(a.size(0), self.num_v_heads // self.tp_size)

361 +

362 + return query, key, value, z, b, a

363 +

364 + def rearrange_mixed_qkv(self, mixed_qkv):

365 + if mixed_qkv is None:

366 + return None, None, None

367 + query, key, value = torch.split(

368 + mixed_qkv,

369 + [

370 + self.key_dim // self.tp_size,

371 + self.key_dim // self.tp_size,

372 + self.value_dim // self.tp_size,

373 + ],

374 + dim=-1,

375 + )

376 + query, key = map(

377 + lambda x: rearrange(x, 'l (h d) -> 1 l h d', d=self.head_k_dim),

378 + (query, key))

379 + value = rearrange(value, 'l (h d) -> 1 l h d', d=self.head_v_dim)

380 + return query, key, value

381 +

382 + def forward(

383 + self,

384 + hidden_states: torch.Tensor,

385 + output: torch.Tensor,

386 + cache_params: Optional[MambaCacheParams] = None,

387 + ):

388 + return torch.ops.vllm.gdn_attention(

389 + hidden_states,

390 + output,

391 + self.prefix,

392 + )

393 +

394 + def _forward(

395 + self,

396 + hidden_states: torch.Tensor,

397 + output: torch.Tensor,

398 + ):

399 + forward_context = get_forward_context()

400 + attn_metadata: AttentionMetadata = forward_context.attn_metadata

401 +

402 + if attn_metadata is None:

403 + # V1 profile run

404 + return

405 +

406 + assert isinstance(attn_metadata, dict)

407 + attn_metadata = attn_metadata[self.prefix]

408 + assert isinstance(attn_metadata, GDNAttentionMetadata)

409 + has_initial_state = attn_metadata.has_initial_state

410 + spec_query_start_loc = attn_metadata.spec_query_start_loc

411 + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc

412 + spec_sequence_masks = attn_metadata.spec_sequence_masks

413 + spec_token_masks = attn_metadata.spec_token_masks

414 + spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501

415 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501

416 + self_kv_cache = self.kv_cache[forward_context.virtual_engine]

417 + conv_state = self_kv_cache[0].transpose(-1, -2)

418 + ssm_state = self_kv_cache[1]

419 + num_actual_tokens = (attn_metadata.num_prefill_tokens +

420 + attn_metadata.num_decode_tokens +

421 + attn_metadata.num_spec_decode_tokens)

422 + num_accepted_tokens = attn_metadata.num_accepted_tokens

423 +

424 + # 1. Set up dimensions for reshapes later

425 + projected_states, _ = self.in_proj(hidden_states[:num_actual_tokens])

426 + if spec_token_masks is not None:

427 + spec_token_masks = spec_token_masks[:num_actual_tokens]

428 + projected_states_qkvz, projected_states_ba = torch.split(

429 + projected_states,

430 + [

431 + self.projection_size_qkvz // self.tp_size,

432 + self.projection_size_ba // self.tp_size

433 + ],

434 + dim=-1,

435 + )

436 + query, key, value, z, b, a = self.fix_query_key_value_ordering(

437 + projected_states_qkvz, projected_states_ba)

438 + query, key, value = map(lambda x: rearrange(x, 'l p d -> l (p d)'),

439 + (query, key, value))

440 + mixed_qkv = torch.cat((query, key, value), dim=-1)

441 +

442 + # 2. Convolution sequence transformation

443 + conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),

444 + self.conv1d.weight.size(2))

445 +

446 + if spec_sequence_masks is not None:

447 + if (attn_metadata.num_prefills == 0

448 + and attn_metadata.num_decodes == 0):

449 + mixed_qkv_spec = mixed_qkv

450 + mixed_qkv_non_spec = None

451 + else:

452 + mixed_qkv_spec = mixed_qkv[spec_token_masks]

453 + mixed_qkv_non_spec = mixed_qkv[~spec_token_masks]

454 + else:

455 + mixed_qkv_spec = None

456 + mixed_qkv_non_spec = mixed_qkv

457 +

458 + # 2.1: process the mutli-query part

459 + if spec_sequence_masks is not None:

460 + mixed_qkv_spec = mixed_qkv_spec.view(

461 + attn_metadata.num_spec_decodes, -1, mixed_qkv_spec.size(-1))

462 + mixed_qkv_spec = rearrange(mixed_qkv_spec, 'b l d -> b d l')

463 + mixed_qkv_spec = causal_conv1d_update(

464 + mixed_qkv_spec,

465 + conv_state,

466 + conv_weights,

467 + self.conv1d.bias,

468 + self.activation,

469 + conv_state_indices=spec_state_indices_tensor[:, 0]

470 + [:attn_metadata.num_spec_decodes],

471 + num_accepted_tokens=num_accepted_tokens,

472 + validate_data=False,

473 + )

474 + mixed_qkv_spec = rearrange(mixed_qkv_spec, 'b d l -> (b l) d')

475 +

476 + # 2.2: process the remaining part

477 + if attn_metadata.num_prefills > 0:

478 + # - "cache_indices" updates the conv_state cache in positions

479 + # pointed to by "mamba_cache_params.state_indices_tensor"

480 + mixed_qkv_non_spec = causal_conv1d_fn(

481 + mixed_qkv_non_spec.transpose(0, 1),

482 + conv_weights,

483 + self.conv1d.bias,

484 + activation=self.activation,

485 + conv_states=conv_state,

486 + has_initial_state=has_initial_state,

487 + cache_indices=non_spec_state_indices_tensor,

488 + query_start_loc=non_spec_query_start_loc,

489 + ).transpose(0, 1)

490 + elif attn_metadata.num_decodes > 0:

491 + mixed_qkv_non_spec = causal_conv1d_update(

492 + mixed_qkv_non_spec,

493 + conv_state,

494 + conv_weights,

495 + self.conv1d.bias,

496 + self.activation,

497 + conv_state_indices=non_spec_state_indices_tensor[:attn_metadata

498 + .num_decodes],

499 + validate_data=True,

500 + )

501 + else:

502 + mixed_qkv_non_spec = None

503 +

504 + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(

505 + mixed_qkv_spec)

506 + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(

507 + mixed_qkv_non_spec)

508 +

509 + beta = b.sigmoid()

510 + # g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)

511 + g = fused_gdn_gating(self.A_log, a, self.dt_bias)

512 + g, beta = map(lambda x: rearrange(x, 'l d -> 1 l d'), (g, beta))

513 +

514 + if spec_sequence_masks is not None:

515 + if (attn_metadata.num_prefills == 0

516 + and attn_metadata.num_decodes == 0):

517 + g_spec = g

518 + beta_spec = beta

519 + g_non_spec = None

520 + beta_non_spec = None

521 + else:

522 + g_spec = g[:, spec_token_masks]

523 + beta_spec = beta[:, spec_token_masks]

524 + g_non_spec = g[:, ~spec_token_masks]

525 + beta_non_spec = beta[:, ~spec_token_masks]

526 + else:

527 + g_spec = None

528 + beta_spec = None

529 + g_non_spec = g

530 + beta_non_spec = beta

531 +

532 + # 3. Recurrent attention

533 +

534 + # 3.1: process the mutlti-query part

535 + if spec_sequence_masks is not None:

536 + core_attn_out_spec, last_recurrent_state = (

537 + fused_recurrent_gated_delta_rule(

538 + q=query_spec,

539 + k=key_spec,

540 + v=value_spec,

541 + g=g_spec,

542 + beta=beta_spec,

543 + initial_state=ssm_state,

544 + inplace_final_state=True,

545 + cu_seqlens=spec_query_start_loc[:attn_metadata.

546 + num_spec_decodes + 1],

547 + ssm_state_indices=spec_state_indices_tensor,

548 + num_accepted_tokens=num_accepted_tokens,

549 + use_qk_l2norm_in_kernel=True,

550 + ))

551 + else:

552 + core_attn_out_spec, last_recurrent_state = None, None

553 +

554 + # 3.2: process the remaining part

555 + if attn_metadata.num_prefills > 0:

556 + initial_state = ssm_state[

557 + non_spec_state_indices_tensor].contiguous()

558 + initial_state[~has_initial_state, ...] = 0

559 + (

560 + core_attn_out_non_spec,

561 + last_recurrent_state,

562 + ) = chunk_gated_delta_rule(

563 + q=query_non_spec,

564 + k=key_non_spec,

565 + v=value_non_spec,

566 + g=g_non_spec,

567 + beta=beta_non_spec,

568 + initial_state=initial_state,

569 + output_final_state=True,

570 + cu_seqlens=non_spec_query_start_loc,

571 + head_first=False,

572 + use_qk_l2norm_in_kernel=True,

573 + )

574 + # Init cache

575 + ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(

576 + ssm_state.dtype)

577 + elif attn_metadata.num_decodes > 0:

578 + core_attn_out_non_spec, last_recurrent_state = (

579 + fused_recurrent_gated_delta_rule(

580 + q=query_non_spec,

581 + k=key_non_spec,

582 + v=value_non_spec,

583 + g=g_non_spec,

584 + beta=beta_non_spec,

585 + initial_state=ssm_state,

586 + inplace_final_state=True,

587 + cu_seqlens=non_spec_query_start_loc[:attn_metadata.

588 + num_decodes + 1],

589 + ssm_state_indices=non_spec_state_indices_tensor,

590 + use_qk_l2norm_in_kernel=True,

591 + ))

592 + else:

593 + core_attn_out_non_spec, last_recurrent_state = None, None

594 +

595 + # Merge core attention output

596 + if (spec_sequence_masks is not None

597 + and core_attn_out_non_spec is not None):

598 + core_attn_out = torch.empty(

599 + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),

600 + dtype=core_attn_out_non_spec.dtype,

601 + device=core_attn_out_non_spec.device,

602 + )

603 + core_attn_out[:, spec_token_masks] = core_attn_out_spec

604 + core_attn_out[:, ~spec_token_masks] = core_attn_out_non_spec

605 + elif spec_sequence_masks is not None:

606 + core_attn_out = core_attn_out_spec

607 + else:

608 + core_attn_out = core_attn_out_non_spec

609 +

610 + z_shape_og = z.shape

611 + # reshape input data into 2D tensor

612 + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])

613 + z = z.reshape(-1, z.shape[-1])

614 + core_attn_out = self.norm(core_attn_out, z)

615 + core_attn_out = core_attn_out.reshape(z_shape_og)

616 + core_attn_out = rearrange(core_attn_out, '... h d -> ... (h d)')

617 +

618 + output[:num_actual_tokens], _ = self.out_proj(core_attn_out)

619 +

620 +

621 +class Qwen3NextAttention(nn.Module):

622 +

623 + def __init__(

624 + self,

625 + config: Qwen3NextConfig,

626 + model_config: Optional[ModelConfig] = None,

627 + cache_config: Optional[CacheConfig] = None,

628 + quant_config: Optional[QuantizationConfig] = None,

629 + prefix: str = "",

630 + ) -> None:

631 + super().__init__()

632 + self.config = config

633 + self.hidden_size = config.hidden_size

634 + tp_size = get_tensor_model_parallel_world_size()

635 + self.total_num_heads = config.num_attention_heads

636 + assert self.total_num_heads % tp_size == 0

637 + self.num_heads = self.total_num_heads // tp_size

638 + self.total_num_kv_heads = config.num_key_value_heads

639 + if self.total_num_kv_heads >= tp_size:

640 + # Number of KV heads is greater than TP size, so we partition

641 + # the KV heads across multiple tensor parallel GPUs.

642 + assert self.total_num_kv_heads % tp_size == 0

643 + else:

644 + # Number of KV heads is less than TP size, so we replicate

645 + # the KV heads across multiple tensor parallel GPUs.

646 + assert tp_size % self.total_num_kv_heads == 0

647 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

648 + self.head_dim = config.head_dim or (self.hidden_size // self.num_heads)

649 + self.q_size = self.num_heads * self.head_dim

650 + self.kv_size = self.num_kv_heads * self.head_dim

651 + self.scaling = self.head_dim**-0.5

652 + self.dual_chunk_attention_config = getattr(

653 + config, "dual_chunk_attention_config", None)

654 + self.attn_output_gate = getattr(config, "attn_output_gate", True)

655 +

656 + self.qkv_proj = QKVParallelLinear(

657 + config.hidden_size,

658 + self.head_dim,

659 + self.total_num_heads * (1 + self.attn_output_gate),

660 + self.total_num_kv_heads,

661 + bias=getattr(config, "qkv_bias", False),

662 + quant_config=quant_config,

663 + prefix=f"{prefix}.qkv_proj",

664 + )

665 +

666 + self.o_proj = RowParallelLinear(

667 + self.total_num_heads * self.head_dim,

668 + config.hidden_size,

669 + bias=False,

670 + quant_config=quant_config,

671 + prefix=f"{prefix}.o_proj",

672 + )

673 +

674 + self.rotary_emb = get_rope(

675 + head_size=self.head_dim,

676 + rotary_dim=self.head_dim,

677 + max_position=config.max_position_embeddings,

678 + base=config.rope_theta,

679 + rope_scaling=config.rope_scaling,

680 + partial_rotary_factor=config.partial_rotary_factor,

681 + dual_chunk_attention_config=self.dual_chunk_attention_config,

682 + )

683 +

684 + self.attn = Attention(

685 + self.num_heads,

686 + self.head_dim,

687 + self.scaling,

688 + num_kv_heads=self.num_kv_heads,

689 + cache_config=cache_config,

690 + quant_config=quant_config,

691 + prefix=f"{prefix}.attn",

692 + **{

693 + "layer_idx": extract_layer_index(prefix),

694 + "dual_chunk_attention_config":

695 + self.dual_chunk_attention_config,

696 + } if self.dual_chunk_attention_config else {},

697 + )

698 +

699 + self.q_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps)

700 + self.k_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps)

701 +

702 + def forward(

703 + self,

704 + positions: torch.Tensor,

705 + output: torch.Tensor,

706 + hidden_states: torch.Tensor,

707 + ):

708 + qkv, _ = self.qkv_proj(hidden_states)

709 +

710 + if self.attn_output_gate:

711 + q_gate, k, v = qkv.split(

712 + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1)

713 + orig_shape = q_gate.shape[:-1]

714 + q_gate = q_gate.view(*orig_shape, self.num_heads, -1)

715 + q, gate = torch.chunk(q_gate, 2, dim=-1)

716 + q = q.reshape(*orig_shape, -1)

717 + gate = gate.reshape(*orig_shape, -1)

718 + else:

719 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size],

720 + dim=-1)

721 +

722 + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(

723 + -1, self.num_heads * self.head_dim)

724 + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(

725 + -1, self.num_kv_heads * self.head_dim)

726 +

727 + q, k = self.rotary_emb(positions, q, k)

728 +

729 + attn_output = self.attn(q, k, v)

730 +

731 + if self.attn_output_gate:

732 + gate = torch.sigmoid(gate)

733 + attn_output = attn_output * gate

734 +

735 + output[:], _ = self.o_proj(attn_output)

736 +

737 +

738 +class Qwen3NextDecoderLayer(nn.Module):

739 +

740 + def __init__(

741 + self,

742 + config: Qwen3NextConfig,

743 + layer_type: str,

744 + model_config: Optional[ModelConfig] = None,

745 + cache_config: Optional[CacheConfig] = None,

746 + quant_config: Optional[QuantizationConfig] = None,

747 + speculative_config: Optional[SpeculativeConfig] = None,

748 + prefix: str = "",

749 + enable_eplb: bool = False,

750 + ) -> None:

751 + super().__init__()

752 + self.config = config

753 +

754 + self.layer_type = layer_type

755 + self.layer_idx = extract_layer_index(prefix)

756 +

757 + if self.layer_type == "linear_attention":

758 + self.linear_attn = Qwen3NextGatedDeltaNet(

759 + config,

760 + model_config=model_config,

761 + cache_config=cache_config,

762 + quant_config=quant_config,

763 + speculative_config=speculative_config,

764 + prefix=f'{prefix}.linear_attn')

765 + elif self.layer_type == "full_attention":

766 + self.self_attn = Qwen3NextAttention(

767 + config,

768 + model_config=model_config,

769 + cache_config=cache_config,

770 + quant_config=quant_config,

771 + prefix=f'{prefix}.self_attn',

772 + )

773 + else:

774 + raise ValueError(f"Invalid layer_type {self.layer_type}")

775 +

776 + mlp_only_layers = ([] if not hasattr(config, "mlp_only_layers") else

777 + config.mlp_only_layers)

778 + if (self.layer_idx not in mlp_only_layers) and (

779 + config.num_experts > 0 and

780 + (self.layer_idx + 1) % config.decoder_sparse_step == 0):

781 + self.mlp = Qwen3NextSparseMoeBlock(

782 + config=config,

783 + quant_config=quant_config,

784 + prefix=f"{prefix}.mlp",

785 + enable_eplb=enable_eplb,

786 + )

787 + else:

788 + self.mlp = Qwen3NextMLP(

789 + hidden_size=config.hidden_size,

790 + intermediate_size=config.intermediate_size,

791 + hidden_act=config.hidden_act,

792 + quant_config=quant_config,

793 + )

794 +

795 + self.input_layernorm = Qwen3NextRMSNorm(config.hidden_size,

796 + eps=config.rms_norm_eps)

797 + self.post_attention_layernorm = Qwen3NextRMSNorm(

798 + config.hidden_size, eps=config.rms_norm_eps)

799 +

800 + self.layer_scale = getattr(config, "layer_scale", False)

801 + if self.layer_scale:

802 + self.attn_layer_scale = torch.nn.Parameter(

803 + torch.zeros(

804 + 1,

805 + 1,

806 + self.config.hidden_size,

807 + dtype=config.torch_dtype,

808 + ), )

809 + self.ffn_layer_scale = torch.nn.Parameter(

810 + torch.zeros(

811 + 1,

812 + 1,

813 + self.config.hidden_size,

814 + dtype=config.torch_dtype,

815 + ), )

816 +

817 + def forward(

818 + self,

819 + hidden_states: torch.Tensor,

820 + residual: Optional[torch.Tensor],

821 + positions: torch.Tensor = None,

822 + **kwargs: object,

823 + ):

824 + if residual is None:

825 + residual = hidden_states

826 + hidden_states = self.input_layernorm(hidden_states)

827 + else:

828 + hidden_states, residual = self.input_layernorm(

829 + hidden_states, residual)

830 +

831 + self_attention_output = torch.empty_like(hidden_states)

832 + if self.layer_type == "linear_attention":

833 + self.linear_attn(

834 + hidden_states=hidden_states,

835 + output=self_attention_output,

836 + )

837 + elif self.layer_type == "full_attention":

838 + self.self_attn(

839 + hidden_states=hidden_states,

840 + output=self_attention_output,

841 + positions=positions,

842 + )

843 + else:

844 + raise ValueError("Invalid layer_type")

845 + hidden_states = self_attention_output

846 +

847 + if self.layer_scale:

848 + if len(hidden_states.shape) == 2:

849 + hidden_states = hidden_states * (

850 + self.attn_layer_scale.to(hidden_states.dtype)[0] + 1)

851 + else:

852 + hidden_states = hidden_states * (

853 + self.attn_layer_scale.to(hidden_states.dtype) + 1)

854 +

855 + # Fully Connected

856 + hidden_states, residual = self.post_attention_layernorm(

857 + hidden_states, residual)

858 + hidden_states = self.mlp(hidden_states)

859 +

860 + if self.layer_scale:

861 + if len(hidden_states.shape) == 2:

862 + hidden_states = hidden_states * (

863 + self.ffn_layer_scale.to(hidden_states.dtype)[0] + 1)

864 + else:

865 + assert len(hidden_states.shape) == len(

866 + self.ffn_layer_scale.shape

867 + ), f'shape must be the same {len(hidden_states.shape)}, {len(self.ffn_layer_scale.shape)}' # noqa: E501

868 + hidden_states = hidden_states * (

869 + self.ffn_layer_scale.to(hidden_states.dtype) + 1)

870 +

871 + return hidden_states, residual

872 +

873 +

874 +@support_torch_compile

875 +class Qwen3NextModel(nn.Module):

876 +

877 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

878 + super().__init__()

879 +

880 + config: Qwen3NextConfig = vllm_config.model_config.hf_config

881 + model_config = vllm_config.model_config

882 + cache_config = vllm_config.cache_config

883 + quant_config = vllm_config.quant_config

884 + parallel_config = vllm_config.parallel_config

885 + lora_config = vllm_config.lora_config

886 + speculative_config = vllm_config.speculative_config

887 + enable_eplb = parallel_config.enable_eplb

888 + eplb_config = parallel_config.eplb_config

889 + self.num_redundant_experts = eplb_config.num_redundant_experts

890 +

891 + self.config = config

892 + lora_vocab = ((lora_config.lora_extra_vocab_size *

893 + (lora_config.max_loras or 1)) if lora_config else 0)

894 + self.vocab_size = config.vocab_size + lora_vocab

895 +

896 + self.embed_tokens = VocabParallelEmbedding(

897 + self.vocab_size,

898 + config.hidden_size,

899 + org_num_embeddings=config.vocab_size,

900 + )

901 +

902 + def get_layer(prefix: str):

903 + return Qwen3NextDecoderLayer(

904 + config,

905 + layer_type=config.layer_types[extract_layer_index(prefix)],

906 + model_config=model_config,

907 + cache_config=cache_config,

908 + quant_config=quant_config,

909 + speculative_config=speculative_config,

910 + prefix=prefix,

911 + enable_eplb=enable_eplb,

912 + )

913 +

914 + self.start_layer, self.end_layer, self.layers = make_layers(

915 + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers")

916 + self.make_empty_intermediate_tensors = (

917 + make_empty_intermediate_tensors_factory(

918 + ["hidden_states", "residual"], config.hidden_size))

919 +

920 + self.norm = Qwen3NextRMSNorm(config.hidden_size,

921 + eps=config.rms_norm_eps)

922 +

923 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

924 + return self.embed_tokens(input_ids)

925 +

926 + def forward(

927 + self,

928 + input_ids: torch.Tensor,

929 + positions: torch.Tensor,

930 + intermediate_tensors: Optional[IntermediateTensors] = None,

931 + inputs_embeds: Optional[torch.Tensor] = None,

932 + ) -> torch.Tensor:

933 + if get_pp_group().is_first_rank:

934 + if inputs_embeds is not None:

935 + hidden_states = inputs_embeds

936 + else:

937 + hidden_states = self.get_input_embeddings(input_ids)

938 + residual = None

939 + else:

940 + assert intermediate_tensors is not None

941 + hidden_states = intermediate_tensors["hidden_states"]

942 + residual = intermediate_tensors["residual"]

943 +

944 + for layer in self.layers:

945 + hidden_states, residual = layer(

946 + positions=positions,

947 + hidden_states=hidden_states,

948 + residual=residual,

949 + )

950 +

951 + if not get_pp_group().is_last_rank:

952 + return IntermediateTensors({

953 + "hidden_states": hidden_states,

954 + "residual": residual

955 + })

956 + hidden_states, _ = self.norm(hidden_states, residual)

957 + return hidden_states

958 +

959 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

960 + # Params for weights, fp8 weight scales, fp8 activation scales

961 + # (param_name, weight_name, expert_id, shard_id)

962 + return FusedMoE.make_expert_params_mapping(

963 + ckpt_gate_proj_name="gate_proj",

964 + ckpt_down_proj_name="down_proj",

965 + ckpt_up_proj_name="up_proj",

966 + num_experts=self.config.num_experts,

967 + num_redundant_experts=self.num_redundant_experts)

968 +

969 + def load_weights(self, weights: Iterable[tuple[str,

970 + torch.Tensor]]) -> set[str]:

971 + stacked_params_mapping = [

972 + # (param_name, shard_name, shard_id)

973 + ("qkv_proj", "q_proj", "q"),

974 + ("qkv_proj", "k_proj", "k"),

975 + ("qkv_proj", "v_proj", "v"),

976 + ("gate_up_proj", "gate_proj", 0),

977 + ("gate_up_proj", "up_proj", 1),

978 + ("in_proj", "in_proj_qkvz", 0),

979 + ("in_proj", "in_proj_ba", 1),

980 + ]

981 +

982 + params_dict = dict(self.named_parameters())

983 + loaded_params: set[str] = set()

984 + expert_params_mapping = self.get_expert_mapping()

985 + for name, loaded_weight in weights:

986 + if "rotary_emb.inv_freq" in name:

987 + continue

988 +

989 + if name.startswith("mtp."):

990 + continue

991 +

992 + for param_name, weight_name, shard_id in stacked_params_mapping:

993 + if weight_name not in name:

994 + continue

995 +

996 + if "mlp.experts" in name:

997 + continue

998 +

999 + name = name.replace(weight_name, param_name)

1000 + # Skip loading extra bias for GPTQ models.

1001 + if name.endswith(".bias") and name not in params_dict:

1002 + continue

1003 + # Skip layers on other devices.

1004 + if is_pp_missing_parameter(name, self):

1005 + continue

1006 + # name = apply_attn_prefix(name, params_dict)

1007 + if name not in params_dict:

1008 + continue

1009 + param = params_dict[name]

1010 + weight_loader = param.weight_loader

1011 + weight_loader(param, loaded_weight, shard_id)

1012 + break

1013 + else:

1014 + for mapping in expert_params_mapping:

1015 + param_name, weight_name, expert_id, shard_id = mapping

1016 + if weight_name not in name:

1017 + continue

1018 + name = name.replace(weight_name, param_name)

1019 + # Skip layers on other devices.

1020 + if is_pp_missing_parameter(name, self):

1021 + continue

1022 + # Skip loading extra bias for GPTQ models.

1023 + if ((name.endswith(".bias") or name.endswith("_bias"))

1024 + and name not in params_dict):

1025 + continue

1026 + param = params_dict[name]

1027 + weight_loader = param.weight_loader

1028 + weight_loader(param,

1029 + loaded_weight,

1030 + name,

1031 + shard_id=shard_id,

1032 + expert_id=expert_id)

1033 + break

1034 + else:

1035 + # Skip loading extra bias for GPTQ models.

1036 + if name.endswith(".bias") and name not in params_dict:

1037 + continue

1038 + if is_pp_missing_parameter(name, self):

1039 + continue

1040 + param = params_dict[name]

1041 + weight_loader = getattr(param, "weight_loader",

1042 + default_weight_loader)

1043 + weight_loader(param, loaded_weight)

1044 + loaded_params.add(name)

1045 + return loaded_params

1046 +

1047 +

1048 +class Qwen3NextForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

1049 + MixtureOfExperts, IsHybrid):

1050 + packed_modules_mapping = {

1051 + "qkv_proj": [

1052 + "q_proj",

1053 + "k_proj",

1054 + "v_proj",

1055 + ],

1056 + "gate_up_proj": ["up_proj", "down_proj"],

1057 + "in_proj": ["in_proj_qkvz", "in_proj_ba"],

1058 + }

1059 +

1060 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1061 + config = vllm_config.model_config.hf_config

1062 + self.vllm_config = vllm_config

1063 + self.model_config = vllm_config.model_config

1064 + cache_config = vllm_config.cache_config

1065 + lora_config = vllm_config.lora_config

1066 + scheduler_config = vllm_config.scheduler_config

1067 + assert not cache_config.enable_prefix_caching, \

1068 + "Qwen3Next currently does not support prefix caching"

1069 + assert envs.VLLM_USE_V1, "Qwen3Next requires VLLM_USE_V1"

1070 + self.quant_config = vllm_config.quant_config

1071 +

1072 + super().__init__()

1073 + self.config = config

1074 + self.scheduler_config = scheduler_config

1075 + self.model = Qwen3NextModel(vllm_config=vllm_config,

1076 + prefix=maybe_prefix(prefix, "model"))

1077 + self.unpadded_vocab_size = config.vocab_size

1078 + if lora_config:

1079 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

1080 + self.lm_head = ParallelLMHead(

1081 + self.unpadded_vocab_size,

1082 + config.hidden_size,

1083 + org_num_embeddings=config.vocab_size,

1084 + padding_size=DEFAULT_VOCAB_PADDING_SIZE

1085 + # We need bigger padding if using lora for kernel

1086 + # compatibility

1087 + if not lora_config else lora_config.lora_vocab_padding_size,

1088 + )

1089 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

1090 + config.vocab_size)

1091 + self.make_empty_intermediate_tensors = (

1092 + self.model.make_empty_intermediate_tensors)

1093 +

1094 + # Set MoE hyperparameters

1095 + self.expert_weights = []

1096 +

1097 + self.moe_layers: list[FusedMoE] = []

1098 + example_layer = None

1099 + for layer in self.model.layers:

1100 + if isinstance(layer, PPMissingLayer):

1101 + continue

1102 +

1103 + assert isinstance(layer, Qwen3NextDecoderLayer)

1104 + if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):

1105 + example_layer = layer.mlp

1106 + self.moe_layers.append(layer.mlp.experts)

1107 +

1108 + if example_layer is None:

1109 + raise RuntimeError("No Qwen3Next layer found in the model.layers.")

1110 +

1111 + self.num_moe_layers = len(self.moe_layers)

1112 + self.num_expert_groups = 1

1113 + self.num_shared_experts = 0

1114 + self.num_logical_experts = example_layer.n_logical_experts

1115 + self.num_physical_experts = example_layer.n_physical_experts

1116 + self.num_local_physical_experts = example_layer.n_local_physical_experts

1117 + self.num_routed_experts = example_layer.n_routed_experts

1118 + self.num_redundant_experts = example_layer.n_redundant_experts

1119 +

1120 + def set_eplb_state(

1121 + self,

1122 + expert_load_view: torch.Tensor,

1123 + logical_to_physical_map: torch.Tensor,

1124 + logical_replica_count: torch.Tensor,

1125 + ) -> None:

1126 + for layer_idx, layer in enumerate(self.moe_layers):

1127 + # Register the expert weights.

1128 + self.expert_weights.append(layer.get_expert_weights())

1129 + layer.set_eplb_state(

1130 + moe_layer_idx=layer_idx,

1131 + expert_load_view=expert_load_view,

1132 + logical_to_physical_map=logical_to_physical_map,

1133 + logical_replica_count=logical_replica_count,

1134 + )

1135 +

1136 + def update_physical_experts_metadata(

1137 + self,

1138 + num_physical_experts: int,

1139 + num_local_physical_experts: int,

1140 + ) -> None:

1141 + assert self.num_local_physical_experts == num_local_physical_experts

1142 + self.num_physical_experts = num_physical_experts

1143 + self.num_local_physical_experts = num_local_physical_experts

1144 + self.num_redundant_experts = (num_physical_experts -

1145 + self.num_logical_experts)

1146 + for layer in self.model.layers:

1147 + if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):

1148 + moe = layer.mlp

1149 + moe.n_local_physical_experts = num_local_physical_experts

1150 + moe.n_physical_experts = num_physical_experts

1151 + moe.n_redundant_experts = self.num_redundant_experts

1152 + moe.experts.update_expert_map()

1153 +

1154 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

1155 + return self.model.get_input_embeddings(input_ids)

1156 +

1157 + def forward(

1158 + self,

1159 + input_ids: torch.Tensor,

1160 + positions: torch.Tensor,

1161 + intermediate_tensors: Optional[IntermediateTensors] = None,

1162 + inputs_embeds: Optional[torch.Tensor] = None,

1163 + **kwargs: object,

1164 + ):

1165 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

1166 + inputs_embeds)

1167 +

1168 + return hidden_states

1169 +

1170 + @classmethod

1171 + def get_mamba_state_dtype_from_config(

1172 + cls,

1173 + vllm_config: "VllmConfig",

1174 + ) -> tuple[torch.dtype, torch.dtype]:

1175 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(

1176 + vllm_config.model_config.dtype,

1177 + vllm_config.cache_config.mamba_cache_dtype)

1178 +

1179 + @classmethod

1180 + def get_mamba_state_shape_from_config(

1181 + cls, vllm_config: "VllmConfig"

1182 + ) -> tuple[tuple[int, int], tuple[int, int]]:

1183 + parallel_config = vllm_config.parallel_config

1184 + hf_config = vllm_config.model_config.hf_config

1185 + tp_size = parallel_config.tensor_parallel_size

1186 + num_spec = (vllm_config.speculative_config.num_speculative_tokens

1187 + if vllm_config.speculative_config else 0)

1188 + return MambaStateShapeCalculator.gated_delta_net_state_shape(

1189 + tp_size,

1190 + hf_config.linear_num_key_heads,

1191 + hf_config.linear_num_value_heads,

1192 + hf_config.linear_key_head_dim,

1193 + hf_config.linear_value_head_dim,

1194 + hf_config.linear_conv_kernel_dim,

1195 + num_spec,

1196 + use_v1=True)

1197 +

1198 + def compute_logits(

1199 + self,

1200 + hidden_states: torch.Tensor,

1201 + sampling_metadata: SamplingMetadata,

1202 + ) -> Optional[torch.Tensor]:

1203 + return self.logits_processor(self.lm_head, hidden_states,

1204 + sampling_metadata)

1205 +

1206 + def load_weights(self, weights: Iterable[tuple[str,

1207 + torch.Tensor]]) -> set[str]:

1208 + loader = AutoWeightsLoader(

1209 + self,

1210 + skip_prefixes=["mtp."],

1211 + )

1212 + return loader.load_weights(weights)

1213 +

1214 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

1215 + return self.model.get_expert_mapping()

1216 +

1217 +

1218 +def gdn_attention(

1219 + hidden_states: torch.Tensor,

1220 + output: torch.Tensor,

1221 + layer_name: str,

1222 +) -> None:

1223 + forward_context: ForwardContext = get_forward_context()

1224 + self = forward_context.no_compile_layers[layer_name]

1225 + self._forward(hidden_states=hidden_states, output=output)

1226 +

1227 +

1228 +def gdn_attention_fake(

1229 + hidden_states: torch.Tensor,

1230 + output: torch.Tensor,

1231 + layer_name: str,

1232 +) -> None:

1233 + return

1234 +

1235 +

1236 +direct_register_custom_op(

1237 + op_name="gdn_attention",

1238 + op_func=gdn_attention,

1239 + mutates_args=["output"],

1240 + fake_impl=gdn_attention_fake,

1241 + dispatch_key=current_platform.dispatch_key,

1242 +)

1243 +

1244 +

1245 +# g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)

1246 +@triton.jit

1247 +def fused_gdn_gating_kernel(

1248 + g,

1249 + A_log,

1250 + a,

1251 + dt_bias,

1252 + seq_len,

1253 + NUM_HEADS: tl.constexpr,

1254 + beta: tl.constexpr,

1255 + threshold: tl.constexpr,

1256 + BLK_HEADS: tl.constexpr,

1257 +):

1258 + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)

1259 + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)

1260 + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off

1261 + mask = head_off < NUM_HEADS

1262 + blk_A_log = tl.load(A_log + head_off, mask=mask)

1263 + blk_a = tl.load(a + off, mask=mask)

1264 + blk_bias = tl.load(dt_bias + head_off, mask=mask)

1265 + # If the model is loaded in fp16, without the .float() here, A might be -inf

1266 + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)

1267 + softplus_x = tl.where(beta * x <= threshold,

1268 + (1 / beta) * tl.log(1 + tl.exp(beta * x)), x)

1269 + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x

1270 + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)

1271 +

1272 +

1273 +def fused_gdn_gating(

1274 + A_log: torch.Tensor,

1275 + a: torch.Tensor,

1276 + dt_bias: torch.Tensor,

1277 + beta: float = 1.0,

1278 + threshold: float = 20.0,

1279 +) -> torch.Tensor:

1280 + batch, num_heads = a.shape

1281 + seq_len = 1

1282 + grid = (batch, seq_len, triton.cdiv(num_heads, 8))

1283 + g = torch.empty_like(a, dtype=torch.float32)

1284 + fused_gdn_gating_kernel[grid](g,

1285 + A_log,

1286 + a,

1287 + dt_bias,

1288 + seq_len,

1289 + num_heads,

1290 + beta,

1291 + threshold,

1292 + 8,

1293 + num_warps=1)

1294 + return g

vllm/model_executor/models/qwen3_next_mtp.py new

+285

@@ -0,0 +1,285 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Inference-only Qwen3Next MTP model."""

4 +from collections.abc import Iterable

5 +from typing import Optional

6 +

7 +import torch

8 +from torch import nn

9 +

10 +from vllm.compilation.decorators import support_torch_compile

11 +from vllm.config import VllmConfig

12 +from vllm.distributed.parallel_state import get_pp_group

13 +from vllm.logger import init_logger

14 +from vllm.model_executor.layers.fused_moe import FusedMoE

15 +from vllm.model_executor.layers.linear import ColumnParallelLinear

16 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

17 +from vllm.model_executor.layers.vocab_parallel_embedding import (

18 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

19 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

20 +from vllm.model_executor.models.qwen3_next import (Qwen3NextDecoderLayer,

21 + Qwen3NextRMSNorm)

22 +from vllm.model_executor.sampling_metadata import SamplingMetadata

23 +from vllm.sequence import IntermediateTensors

24 +from vllm.transformers_utils.configs import Qwen3NextConfig

25 +

26 +from .interfaces import SupportsPP

27 +from .utils import (AutoWeightsLoader, is_pp_missing_parameter,

28 + make_empty_intermediate_tensors_factory, maybe_prefix)

29 +

30 +logger = init_logger(__name__)

31 +

32 +KVCache = tuple[torch.Tensor, torch.Tensor]

33 +

34 +

35 +@support_torch_compile

36 +class Qwen3NextMultiTokenPredictor(nn.Module):

37 +

38 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

39 + super().__init__()

40 +

41 + model_config = vllm_config.model_config

42 + cache_config = vllm_config.cache_config

43 + quant_config = vllm_config.quant_config

44 + lora_config = vllm_config.lora_config

45 + config: Qwen3NextConfig = model_config.hf_config

46 +

47 + self.config = config

48 + lora_vocab = ((lora_config.lora_extra_vocab_size *

49 + (lora_config.max_loras or 1)) if lora_config else 0)

50 + self.vocab_size = config.vocab_size + lora_vocab

51 + self.org_vocab_size = config.vocab_size

52 +

53 + self.mtp_start_layer_idx = config.num_hidden_layers

54 + self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1)

55 +

56 + self.embed_tokens = VocabParallelEmbedding(

57 + self.vocab_size,

58 + config.hidden_size,

59 + org_num_embeddings=config.vocab_size,

60 + )

61 +

62 + self.fc = ColumnParallelLinear(self.config.hidden_size * 2,

63 + self.config.hidden_size,

64 + gather_output=True,

65 + bias=False,

66 + return_bias=False)

67 +

68 + self.layers = torch.nn.ModuleList(

69 + Qwen3NextDecoderLayer(

70 + config,

71 + layer_type="full_attention",

72 + model_config=model_config,

73 + cache_config=cache_config,

74 + quant_config=quant_config,

75 + prefix=f'{prefix}.layers.{self.mtp_start_layer_idx + idx}',

76 + ) for idx in range(self.num_mtp_layers))

77 +

78 + self.make_empty_intermediate_tensors = (

79 + make_empty_intermediate_tensors_factory(

80 + ["hidden_states", "residual"], config.hidden_size))

81 +

82 + self.norm = Qwen3NextRMSNorm(config.hidden_size,

83 + eps=config.rms_norm_eps)

84 + self.pre_fc_norm_hidden = Qwen3NextRMSNorm(config.hidden_size,

85 + eps=config.rms_norm_eps)

86 + self.pre_fc_norm_embedding = Qwen3NextRMSNorm(config.hidden_size,

87 + eps=config.rms_norm_eps)

88 +

89 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

90 + return self.embed_tokens(input_ids)

91 +

92 + def forward(

93 + self,

94 + input_ids: torch.Tensor,

95 + positions: torch.Tensor,

96 + hidden_states: torch.Tensor,

97 + intermediate_tensors: Optional[IntermediateTensors] = None,

98 + inputs_embeds: Optional[torch.Tensor] = None,

99 + spec_step_idx: int = 0,

100 + ) -> torch.Tensor:

101 + if get_pp_group().is_first_rank:

102 + if inputs_embeds is None:

103 + inputs_embeds = self.get_input_embeddings(input_ids)

104 + assert hidden_states.shape[-1] == inputs_embeds.shape[-1]

105 + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)

106 + hidden_states = self.pre_fc_norm_hidden(hidden_states)

107 + hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1)

108 + hidden_states = self.fc(hidden_states)

109 + residual = None

110 + else:

111 + assert intermediate_tensors is not None

112 + hidden_states = intermediate_tensors["hidden_states"]

113 + residual = intermediate_tensors["residual"]

114 +

115 + current_step_idx = (spec_step_idx % self.num_mtp_layers)

116 + hidden_states, residual = self.layers[current_step_idx](

117 + positions=positions,

118 + hidden_states=hidden_states,

119 + residual=residual,

120 + )

121 +

122 + if not get_pp_group().is_last_rank:

123 + return IntermediateTensors({

124 + "hidden_states": hidden_states,

125 + "residual": residual

126 + })

127 +

128 + hidden_states, _ = self.norm(hidden_states, residual)

129 + return hidden_states

130 +

131 + def load_weights(self, weights: Iterable[tuple[str,

132 + torch.Tensor]]) -> set[str]:

133 + stacked_params_mapping = [

134 + # (param_name, shard_name, shard_id)

135 + ("qkv_proj", "q_proj", "q"),

136 + ("qkv_proj", "k_proj", "k"),

137 + ("qkv_proj", "v_proj", "v"),

138 + ("gate_up_proj", "gate_proj", 0),

139 + ("gate_up_proj", "up_proj", 1),

140 + ]

141 +

142 + # Params for weights, fp8 weight scales, fp8 activation scales

143 + # (param_name, weight_name, expert_id, shard_id)

144 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

145 + ckpt_gate_proj_name="gate_proj",

146 + ckpt_down_proj_name="down_proj",

147 + ckpt_up_proj_name="up_proj",

148 + num_experts=self.config.num_experts)

149 +

150 + params_dict = dict(self.named_parameters())

151 + loaded_params: set[str] = set()

152 + for name, loaded_weight in weights:

153 + if "rotary_emb.inv_freq" in name:

154 + continue

155 +

156 + for param_name, weight_name, shard_id in stacked_params_mapping:

157 + if weight_name not in name:

158 + continue

159 +

160 + if "mlp.experts" in name:

161 + continue

162 +

163 + name = name.replace(weight_name, param_name)

164 + # Skip loading extra bias for GPTQ models.

165 + if name.endswith(".bias") and name not in params_dict:

166 + continue

167 + # Skip layers on other devices.

168 + if is_pp_missing_parameter(name, self):

169 + continue

170 + if name not in params_dict:

171 + continue

172 + param = params_dict[name]

173 + weight_loader = param.weight_loader

174 + weight_loader(param, loaded_weight, shard_id)

175 + break

176 + else:

177 + for mapping in expert_params_mapping:

178 + param_name, weight_name, expert_id, shard_id = mapping

179 + if weight_name not in name:

180 + continue

181 + name = name.replace(weight_name, param_name)

182 + # Skip layers on other devices.

183 + if is_pp_missing_parameter(name, self):

184 + continue

185 + # Skip loading extra bias for GPTQ models.

186 + if ((name.endswith(".bias") or name.endswith("_bias"))

187 + and name not in params_dict):

188 + continue

189 + param = params_dict[name]

190 + weight_loader = param.weight_loader

191 + weight_loader(param,

192 + loaded_weight,

193 + name,

194 + shard_id=shard_id,

195 + expert_id=expert_id)

196 + break

197 + else:

198 + # Skip loading extra bias for GPTQ models.

199 + if name.endswith(".bias") and name not in params_dict:

200 + continue

201 + if is_pp_missing_parameter(name, self):

202 + continue

203 +

204 + param = params_dict[name]

205 + weight_loader = getattr(param, "weight_loader",

206 + default_weight_loader)

207 + weight_loader(param, loaded_weight)

208 + loaded_params.add(name)

209 + return loaded_params

210 +

211 +

212 +@support_torch_compile

213 +class Qwen3NextMTP(nn.Module, SupportsPP):

214 + packed_modules_mapping = {

215 + "qkv_proj": [

216 + "q_proj",

217 + "k_proj",

218 + "v_proj",

219 + ],

220 + "gate_up_proj": ["up_proj", "down_proj"]

221 + }

222 +

223 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

224 + config = vllm_config.model_config.hf_config

225 + self.vllm_config = vllm_config

226 + cache_config = vllm_config.cache_config

227 + assert not cache_config.enable_prefix_caching, \

228 + "Qwen3NextMTP currently does not support prefix caching"

229 +

230 + self.quant_config = vllm_config.quant_config

231 +

232 + super().__init__()

233 + self.config = config

234 + self.model = Qwen3NextMultiTokenPredictor(vllm_config=vllm_config,

235 + prefix=maybe_prefix(

236 + prefix, "model"))

237 + self.unpadded_vocab_size = config.vocab_size

238 + self.lm_head = ParallelLMHead(self.unpadded_vocab_size,

239 + config.hidden_size,

240 + org_num_embeddings=config.vocab_size,

241 + padding_size=DEFAULT_VOCAB_PADDING_SIZE)

242 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

243 + config.vocab_size)

244 + self.make_empty_intermediate_tensors = (

245 + self.model.make_empty_intermediate_tensors)

246 +

247 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

248 + return self.model.get_input_embeddings(input_ids)

249 +

250 + def forward(

251 + self,

252 + input_ids: torch.Tensor,

253 + positions: torch.Tensor,

254 + hidden_states: torch.Tensor,

255 + intermediate_tensors: Optional[IntermediateTensors] = None,

256 + inputs_embeds: Optional[torch.Tensor] = None,

257 + **kwargs: object,

258 + ):

259 + hidden_states = self.model(input_ids, positions, hidden_states,

260 + intermediate_tensors, inputs_embeds)

261 + return hidden_states

262 +

263 + def compute_logits(

264 + self,

265 + hidden_states: torch.Tensor,

266 + sampling_metadata: SamplingMetadata,

267 + spec_step_idx: int = 0,

268 + ) -> Optional[torch.Tensor]:

269 + return self.logits_processor(self.lm_head, hidden_states,

270 + sampling_metadata)

271 +

272 + def load_weights(self, weights: Iterable[tuple[str,

273 + torch.Tensor]]) -> set[str]:

274 + shared_weight_names = ["embed_tokens", "lm_head"]

275 +

276 + def remap_weight_names(weights):

277 + for name, weight in weights:

278 + if name.startswith("mtp."):

279 + name = name.replace("mtp.", "model.")

280 + elif not any(key in name for key in shared_weight_names):

281 + continue

282 + yield name, weight

283 +

284 + loader = AutoWeightsLoader(self)

285 + return loader.load_weights(remap_weight_names(weights))

vllm/model_executor/models/registry.py

@@ -74,6 +74,7 @@ _TEXT_GENERATION_MODELS = {

74 "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),

75 "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),

76 "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),

77 + "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),

78 "GlmForCausalLM": ("glm", "GlmForCausalLM"),

79 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),

80 "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),

@@ -285,6 +286,7 @@ _SPECULATIVE_DECODING_MODELS = {

286 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),

287 "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),

288 "MedusaModel": ("medusa", "Medusa"),

289 + "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),

290 # Temporarily disabled.

291 # # TODO(woosuk): Re-enable this once the MLP Speculator is supported in V1.

292 # "MLPSpeculatorPreTrainedModel": ("mlp_speculator", "MLPSpeculator"),

vllm/transformers_utils/config.py

+1 -1

@@ -79,7 +79,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

79 ultravox="UltravoxConfig",

80 step3_vl="Step3VLConfig",

81 step3_text="Step3TextConfig",

82 -)

82 + qwen3_next="Qwen3NextConfig")

84 _CONFIG_ATTRS_MAPPING: dict[str, str] = {

85 "llm_config": "text_config",

vllm/transformers_utils/configs/__init__.py

@@ -24,6 +24,7 @@ from vllm.transformers_utils.configs.nemotron import NemotronConfig

24 from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig

25 from vllm.transformers_utils.configs.nemotron_vl import Nemotron_Nano_VL_Config

26 from vllm.transformers_utils.configs.ovis import OvisConfig

27 +from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig

28 from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig

29 from vllm.transformers_utils.configs.step3_vl import (Step3TextConfig,

30 Step3VisionEncoderConfig,

@@ -50,4 +51,5 @@ __all__ = [

51 "Step3VLConfig",

52 "Step3VisionEncoderConfig",

53 "Step3TextConfig",

54 + "Qwen3NextConfig",

55 ]

vllm/transformers_utils/configs/qwen3_next.py new

+275

@@ -0,0 +1,275 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

5 +#

6 +# Licensed under the Apache License, Version 2.0 (the "License");

7 +# you may not use this file except in compliance with the License.

8 +# You may obtain a copy of the License at

9 +#

10 +# http://www.apache.org/licenses/LICENSE-2.0

11 +#

12 +# Unless required by applicable law or agreed to in writing, software

13 +# distributed under the License is distributed on an "AS IS" BASIS,

14 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

15 +# See the License for the specific language governing permissions and

16 +# limitations under the License.

17 +"""Qwen3-Next model configuration"""

18 +

19 +from transformers.configuration_utils import (PretrainedConfig,

20 + layer_type_validation)

21 +from transformers.modeling_rope_utils import rope_config_validation

22 +from transformers.utils import logging

23 +

24 +logger = logging.get_logger(__name__)

25 +

26 +

27 +class Qwen3NextConfig(PretrainedConfig):

28 + r"""

29 + This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a

30 + Qwen3-Next model according to the specified arguments, defining the model architecture.

31 + Instantiating a configuration with the defaults will yield a similar configuration to that of

32 + Qwen3-Next-80B-A3B-Instruct [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct).

33 +

34 + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the

35 + documentation from [`PretrainedConfig`] for more information.

36 +

37 +

38 + Args:

39 + vocab_size (`int`, *optional*, defaults to 151936):

40 + Vocabulary size of the model. Defines the number of different tokens that can be represented by the

41 + `inputs_ids`.

42 + hidden_size (`int`, *optional*, defaults to 2048):

43 + Dimension of the hidden representations.

44 + intermediate_size (`int`, *optional*, defaults to 5632):

45 + Dimension of the MLP representations.

46 + num_hidden_layers (`int`, *optional*, defaults to 48):

47 + Number of hidden layers in the Transformer encoder.

48 + num_attention_heads (`int`, *optional*, defaults to 16):

49 + Number of attention heads for each attention layer in the Transformer encoder.

50 + num_key_value_heads (`int`, *optional*, defaults to 2):

51 + This is the number of key_value heads that should be used to implement Grouped Query Attention. If

52 + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if

53 + `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When

54 + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed

55 + by meanpooling all the original heads within that group. For more details checkout [this

56 + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.

57 + hidden_act (`str`, *optional*, defaults to `"silu"`):

58 + The non-linear activation function in the decoder.

59 + max_position_embeddings (`int`, *optional*, defaults to 32768):

60 + The maximum sequence length that this model might ever be used with.

61 + initializer_range (`float`, *optional*, defaults to 0.02):

62 + The standard deviation of the truncated_normal_initializer for initializing all weight matrices.

63 + rms_norm_eps (`float`, *optional*, defaults to 1e-06):

64 + The epsilon used by the rms normalization layers.

65 + use_cache (`bool`, *optional*, defaults to `True`):

66 + Whether or not the model should return the last key/values attentions (not used by all models). Only

67 + relevant if `config.is_decoder=True`.

68 + tie_word_embeddings (`bool`, *optional*, defaults to `False`):

69 + Whether the model's input and output word embeddings should be tied.

70 + rope_theta (`float`, *optional*, defaults to 10000.0):

71 + The base period of the RoPE embeddings.

72 + rope_scaling (`Dict`, *optional*):

73 + Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type

74 + and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value

75 + accordingly.

76 + Expected contents:

77 + `rope_type` (`str`):

78 + The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',

79 + 'llama3'], with 'default' being the original RoPE implementation.

80 + `factor` (`float`, *optional*):

81 + Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In

82 + most scaling types, a `factor` of x will enable the model to handle sequences of length x *

83 + original maximum pre-trained length.

84 + `original_max_position_embeddings` (`int`, *optional*):

85 + Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during

86 + pretraining.

87 + `attention_factor` (`float`, *optional*):

88 + Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention

89 + computation. If unspecified, it defaults to value recommended by the implementation, using the

90 + `factor` field to infer the suggested value.

91 + `beta_fast` (`float`, *optional*):

92 + Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear

93 + ramp function. If unspecified, it defaults to 32.

94 + `beta_slow` (`float`, *optional*):

95 + Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear

96 + ramp function. If unspecified, it defaults to 1.

97 + `short_factor` (`List[float]`, *optional*):

98 + Only used with 'longrope'. The scaling factor to be applied to short contexts (<

99 + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden

100 + size divided by the number of attention heads divided by 2

101 + `long_factor` (`List[float]`, *optional*):

102 + Only used with 'longrope'. The scaling factor to be applied to long contexts (<

103 + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden

104 + size divided by the number of attention heads divided by 2

105 + `low_freq_factor` (`float`, *optional*):

106 + Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE

107 + `high_freq_factor` (`float`, *optional*):

108 + Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE

109 + partial_rotary_factor (`float`, *optional*, defaults to 0.25):

110 + Percentage of the query and keys which will have rotary embedding.

111 + attention_bias (`bool`, *optional*, defaults to `False`):

112 + Whether to use a bias in the query, key, value and output projection layers during self-attention.

113 + attention_dropout (`float`, *optional*, defaults to 0.0):

114 + The dropout ratio for the attention probabilities.

115 + head_dim (`int`, *optional*, defaults to 256):

116 + Projection weights dimension in multi-head attention.

117 + linear_conv_kernel_dim (`int`, *optional*, defaults to 4):

118 + Kernel size of the convolution used in linear attention layers.

119 + linear_key_head_dim (`int`, *optional*, defaults to 128):

120 + Dimension of each key head in linear attention.

121 + linear_value_head_dim (`int`, *optional*, defaults to 128):

122 + Dimension of each value head in linear attention.

123 + linear_num_key_heads (`int`, *optional*, defaults to 16):

124 + Number of key heads used in linear attention layers.

125 + linear_num_value_heads (`int`, *optional*, defaults to 32):

126 + Number of value heads used in linear attention layers.

127 + decoder_sparse_step (`int`, *optional*, defaults to 1):

128 + The frequency of the MoE layer.

129 + moe_intermediate_size (`int`, *optional*, defaults to 512):

130 + Intermediate size of the routed expert.

131 + shared_expert_intermediate_size (`int`, *optional*, defaults to 512):

132 + Intermediate size of the shared expert.

133 + num_experts_per_tok (`int`, *optional*, defaults to 10):

134 + Number of selected experts.

135 + num_experts (`int`, *optional*, defaults to 512):

136 + Number of routed experts.

137 + norm_topk_prob (`bool`, *optional*, defaults to `True`):

138 + Whether to normalize the topk probabilities.

139 + output_router_logits (`bool`, *optional*, defaults to `False`):

140 + Whether or not the router logits should be returned by the model. Enabling this will also

141 + allow the model to output the auxiliary loss, including load balancing loss and router z-loss.

142 + router_aux_loss_coef (`float`, *optional*, defaults to 0.001):

143 + The aux loss factor for the total loss.

144 + mlp_only_layers (`list[int]`, *optional*, defaults to `[]`):

145 + Indicate which layers use Qwen3NextMLP rather than Qwen3NextSparseMoeBlock

146 + The list contains layer index, from 0 to num_layers-1 if we have num_layers layers

147 + If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.

148 + layer_types (`list[str]`, *optional*):

149 + Types of each layer (attention or linear).

150 +

151 + ```python

152 + >>> from transformers import Qwen3NextModel, Qwen3NextConfig

153 +

154 + >>> # Initializing a Qwen3Next style configuration

155 + >>> configuration = Qwen3NextConfig()

156 +

157 + >>> # Initializing a model from the Qwen3-Next-80B-A3B style configuration

158 + >>> model = Qwen3NextModel(configuration)

159 +

160 + >>> # Accessing the model configuration

161 + >>> configuration = model.config

162 + ```

163 + """ # noqa: E501

164 +

165 + model_type = "qwen3_next"

166 + keys_to_ignore_at_inference = ["past_key_values"]

167 +

168 + base_model_tp_plan = {

169 + "layers.*.self_attn.q_proj": "colwise",

170 + "layers.*.self_attn.k_proj": "colwise",

171 + "layers.*.self_attn.v_proj": "colwise",

172 + "layers.*.self_attn.o_proj": "rowwise",

173 + "layers.*.mlp.experts.*.gate_proj": "colwise",

174 + "layers.*.mlp.experts.*.up_proj": "colwise",

175 + "layers.*.mlp.experts.*.down_proj": "rowwise",

176 + "layers.*.mlp.shared_experts.gate_proj": "colwise",

177 + "layers.*.mlp.shared_experts.up_proj": "colwise",

178 + "layers.*.mlp.shared_experts.down_proj": "rowwise",

179 + "layers.*.mlp.gate_proj": "colwise",

180 + "layers.*.mlp.up_proj": "colwise",

181 + "layers.*.mlp.down_proj": "rowwise",

182 + }

183 + base_model_pp_plan = {

184 + "embed_tokens": (["input_ids"], ["inputs_embeds"]),

185 + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),

186 + "norm": (["hidden_states"], ["hidden_states"]),

187 + }

188 +

189 + def __init__(

190 + self,

191 + vocab_size=151936,

192 + hidden_size=2048,

193 + intermediate_size=5632,

194 + num_hidden_layers=48,

195 + num_attention_heads=16,

196 + num_key_value_heads=2,

197 + hidden_act="silu",

198 + max_position_embeddings=32768,

199 + initializer_range=0.02,

200 + rms_norm_eps=1e-6,

201 + use_cache=True,

202 + tie_word_embeddings=False,

203 + rope_theta=10000.0,

204 + rope_scaling=None,

205 + partial_rotary_factor=0.25,

206 + attention_bias=False,

207 + attention_dropout=0.0,

208 + head_dim=256,

209 + linear_conv_kernel_dim=4,

210 + linear_key_head_dim=128,

211 + linear_value_head_dim=128,

212 + linear_num_key_heads=16,

213 + linear_num_value_heads=32,

214 + decoder_sparse_step=1,

215 + moe_intermediate_size=512,

216 + shared_expert_intermediate_size=512,

217 + num_experts_per_tok=10,

218 + num_experts=512,

219 + norm_topk_prob=True,

220 + output_router_logits=False,

221 + router_aux_loss_coef=0.001,

222 + mlp_only_layers=None,

223 + layer_types=None,

224 + **kwargs,

225 + ):

226 + if mlp_only_layers is None:

227 + mlp_only_layers = []

228 + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)

229 + self.vocab_size = vocab_size

230 + self.max_position_embeddings = max_position_embeddings

231 + self.hidden_size = hidden_size

232 + self.intermediate_size = intermediate_size

233 + self.num_hidden_layers = num_hidden_layers

234 + self.num_attention_heads = num_attention_heads

235 + self.num_key_value_heads = num_key_value_heads

236 + self.hidden_act = hidden_act

237 + self.initializer_range = initializer_range

238 + self.rms_norm_eps = rms_norm_eps

239 + self.use_cache = use_cache

240 + self.rope_theta = rope_theta

241 + self.rope_scaling = rope_scaling

242 + self.partial_rotary_factor = partial_rotary_factor

243 + self.attention_bias = attention_bias

244 + self.attention_dropout = attention_dropout

245 + self.head_dim = head_dim

246 + rope_config_validation(self)

247 +

248 + self.layer_types = layer_types

249 + if self.layer_types is None:

250 + self.layer_types = [

251 + "linear_attention" if bool((i + 1) % 4) else "full_attention"

252 + for i in range(self.num_hidden_layers)

253 + ]

254 + layer_type_validation(self.layer_types)

255 +

256 + # linear attention part

257 + self.linear_conv_kernel_dim = linear_conv_kernel_dim

258 + self.linear_key_head_dim = linear_key_head_dim

259 + self.linear_value_head_dim = linear_value_head_dim

260 + self.linear_num_key_heads = linear_num_key_heads

261 + self.linear_num_value_heads = linear_num_value_heads

262 +

263 + # MoE arguments

264 + self.decoder_sparse_step = decoder_sparse_step

265 + self.moe_intermediate_size = moe_intermediate_size

266 + self.shared_expert_intermediate_size = shared_expert_intermediate_size

267 + self.num_experts_per_tok = num_experts_per_tok

268 + self.num_experts = num_experts

269 + self.norm_topk_prob = norm_topk_prob

270 + self.output_router_logits = output_router_logits

271 + self.router_aux_loss_coef = router_aux_loss_coef

272 + self.mlp_only_layers = mlp_only_layers

273 +

274 +

275 +__all__ = ["Qwen3NextConfig"]

vllm/v1/attention/backends/gdn_attn.py new

+319

@@ -0,0 +1,319 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Backend for GatedDeltaNet attention."""

4 +from dataclasses import dataclass

5 +from typing import ClassVar, Optional

6 +

7 +import torch

8 +

9 +from vllm.attention.backends.abstract import AttentionBackend

10 +from vllm.attention.backends.utils import PAD_SLOT_ID

11 +from vllm.config import VllmConfig

12 +from vllm.v1.attention.backends.utils import (AttentionCGSupport,

13 + AttentionMetadataBuilder,

14 + CommonAttentionMetadata,

15 + split_decodes_and_prefills)

16 +from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec

17 +

18 +

19 +class GDNAttentionBackend(AttentionBackend):

20 +

21 + @staticmethod

22 + def get_builder_cls() -> type["GDNAttentionMetadataBuilder"]:

23 + return GDNAttentionMetadataBuilder

24 +

25 +

26 +@dataclass

27 +class GDNAttentionMetadata:

28 + num_prefills: int

29 + num_prefill_tokens: int

30 + num_decodes: int

31 + num_decode_tokens: int

32 + num_spec_decodes: int

33 + num_spec_decode_tokens: int

34 +

35 + has_initial_state: Optional[torch.Tensor] = None

36 +

37 + spec_query_start_loc: Optional[

38 + torch.Tensor] = None # shape: [num_spec_decodes + 1,]

39 + non_spec_query_start_loc: Optional[

40 + torch.Tensor] = None # shape: [batch - num_spec_decodes + 1,]

41 +

42 + spec_state_indices_tensor: Optional[

43 + torch.Tensor] = None # shape: [batch, num_spec]

44 + non_spec_state_indices_tensor: Optional[

45 + torch.Tensor] = None # shape: [batch - num_spec_decodes,]

46 + spec_sequence_masks: Optional[torch.Tensor] = None # shape: [batch,]

47 + spec_token_masks: Optional[

48 + torch.

49 + Tensor] = None # shape: [num_prefill_tokens + num_decode_tokens,]

50 + num_accepted_tokens: Optional[torch.Tensor] = None # shape: [batch,]

51 +

52 +

53 +class GDNAttentionMetadataBuilder(

54 + AttentionMetadataBuilder[GDNAttentionMetadata]):

55 +

56 + cudagraph_support = AttentionCGSupport.UNIFORM_BATCH

57 +

58 + reorder_batch_threshold: ClassVar[int] = 1

59 +

60 + def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],

61 + vllm_config: VllmConfig, device: torch.device):

62 + assert isinstance(kv_cache_spec, MambaSpec)

63 + self.vllm_config = vllm_config

64 + self.compilation_config = vllm_config.compilation_config

65 + self.speculative_config = vllm_config.speculative_config

66 + self.kv_cache_spec = kv_cache_spec

67 + if self.speculative_config:

68 + self.num_spec = self.speculative_config.num_speculative_tokens # noqa: E501

69 + else:

70 + self.num_spec = 0

71 + self.use_spec_decode = self.num_spec > 0

72 + self.reorder_batch_threshold = self.num_spec + 1 # type: ignore[misc]

73 +

74 + self.use_full_cuda_graph = \

75 + self.compilation_config.cudagraph_mode.has_full_cudagraphs()

76 + self.decode_cudagraph_max_bs = min(

77 + self.vllm_config.scheduler_config.max_num_seqs,

78 + self.compilation_config.max_capture_size)

79 +

80 + self.spec_state_indices_tensor = torch.empty(

81 + (self.decode_cudagraph_max_bs, self.num_spec + 1),

82 + dtype=torch.int32,

83 + device=device,

84 + )

85 + self.non_spec_state_indices_tensor = torch.empty(

86 + (self.decode_cudagraph_max_bs, ),

87 + dtype=torch.int32,

88 + device=device,

89 + )

90 + self.spec_sequence_masks = torch.empty(

91 + (self.decode_cudagraph_max_bs, ),

92 + dtype=torch.bool,

93 + device=device,

94 + )

95 + self.spec_token_masks = torch.empty(

96 + (self.decode_cudagraph_max_bs * (self.num_spec + 1), ),

97 + dtype=torch.bool,

98 + device=device,

99 + )

100 + self.spec_query_start_loc = torch.empty(

101 + (self.decode_cudagraph_max_bs + 1, ),

102 + dtype=torch.int32,

103 + device=device,

104 + )

105 + self.non_spec_query_start_loc = torch.empty(

106 + (self.decode_cudagraph_max_bs + 1, ),

107 + dtype=torch.int32,

108 + device=device,

109 + )

110 + self.num_accepted_tokens = torch.empty(

111 + (self.decode_cudagraph_max_bs, ),

112 + dtype=torch.int32,

113 + device=device,

114 + )

115 +

116 + def build( # type: ignore[override]

117 + self,

118 + common_prefix_len: int,

119 + common_attn_metadata: CommonAttentionMetadata,

120 + num_accepted_tokens: Optional[torch.Tensor] = None,

121 + num_draft_tokens: Optional[torch.Tensor] = None,

122 + fast_build: bool = False,

123 + ) -> GDNAttentionMetadata:

124 + m = common_attn_metadata

125 +

126 + query_start_loc = m.query_start_loc

127 + context_lens = m.num_computed_tokens_cpu

128 + context_lens_tensor = context_lens.to(query_start_loc.device)

129 + seq_lens_tensor = m.seq_lens

130 +

131 + if (not self.use_spec_decode or num_draft_tokens is None

132 + or num_draft_tokens.sum().item() == 0):

133 + spec_sequence_masks = None

134 + else:

135 + spec_sequence_masks = (num_draft_tokens > 0) & (

136 + context_lens_tensor +

137 + (num_draft_tokens + 1) == seq_lens_tensor)

138 + if spec_sequence_masks.sum().item() == 0:

139 + spec_sequence_masks = None

140 +

141 + if spec_sequence_masks is None:

142 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

143 + split_decodes_and_prefills(m, decode_threshold=1))

144 + num_spec_decodes = 0

145 + num_spec_decode_tokens = 0

146 + spec_token_masks = None

147 + spec_state_indices_tensor = None

148 + non_spec_state_indices_tensor = m.block_table_tensor[:, 0]

149 + spec_query_start_loc = None

150 + non_spec_query_start_loc = query_start_loc

151 + num_accepted_tokens = None

152 + else:

153 + num_spec_decodes = spec_sequence_masks.sum().item()

154 + query_lens = query_start_loc[1:] - query_start_loc[:-1]

155 +

156 + non_spec_query_lens = query_lens[~spec_sequence_masks]

157 + num_decodes = (non_spec_query_lens == 1).sum().item()

158 + num_prefills = non_spec_query_lens.size(0) - num_decodes

159 + num_decode_tokens = num_decodes

160 + num_prefill_tokens = non_spec_query_lens.sum().item(

161 + ) - num_decode_tokens

162 +

163 + if num_prefills == 0 and num_decodes == 0:

164 + spec_token_masks = torch.ones(

165 + (min(num_spec_decodes *

166 + (self.num_spec + 1), query_start_loc[-1].item())),

167 + dtype=torch.bool,

168 + device=query_start_loc.device)

169 + spec_state_indices_tensor = m.block_table_tensor[:, :self.

170 + num_spec + 1]

171 + non_spec_state_indices_tensor = None

172 + spec_query_start_loc = query_start_loc

173 + non_spec_query_start_loc = None

174 + else:

175 + spec_token_masks = torch.repeat_interleave(

176 + spec_sequence_masks, query_lens)

177 + spec_state_indices_tensor = m.block_table_tensor[

178 + spec_sequence_masks, :self.num_spec + 1]

179 + non_spec_state_indices_tensor = \

180 + m.block_table_tensor[~spec_sequence_masks, 0]

181 +

182 + spec_query_start_loc = torch.zeros(

183 + num_spec_decodes + 1,

184 + dtype=torch.int32,

185 + device=query_start_loc.device)

186 + torch.cumsum(query_lens[spec_sequence_masks],

187 + dim=0,

188 + out=spec_query_start_loc[1:])

189 + non_spec_query_start_loc = torch.zeros(

190 + query_lens.size(0) - num_spec_decodes + 1,

191 + dtype=torch.int32,

192 + device=query_start_loc.device)

193 + torch.cumsum(query_lens[~spec_sequence_masks],

194 + dim=0,

195 + out=non_spec_query_start_loc[1:])

196 +

197 + num_spec_decode_tokens = min(

198 + num_spec_decodes * (self.num_spec + 1),

199 + spec_token_masks.size(0))

200 + assert num_accepted_tokens is not None

201 + num_accepted_tokens = num_accepted_tokens[spec_sequence_masks]

202 +

203 + if num_prefills > 0:

204 + has_initial_state = context_lens_tensor > 0

205 + if spec_sequence_masks is not None:

206 + has_initial_state = has_initial_state[~spec_sequence_masks]

207 + else:

208 + has_initial_state = None

209 +

210 + # prepare tensors for cudagraph

211 + if (self.use_full_cuda_graph and num_prefills == 0 and num_decodes == 0

212 + and num_spec_decodes <= self.decode_cudagraph_max_bs):

213 + num_total_tokens = self.vllm_config.pad_for_cudagraph(

214 + m.num_actual_tokens)

215 + batch_size = num_total_tokens // (self.num_spec + 1)

216 +

217 + self.spec_state_indices_tensor[:num_spec_decodes].copy_(

218 + spec_state_indices_tensor, non_blocking=True)

219 + spec_state_indices_tensor = self.spec_state_indices_tensor[:

220 + batch_size]

221 + spec_state_indices_tensor[num_spec_decodes:].fill_(PAD_SLOT_ID)

222 +

223 + self.spec_sequence_masks[:num_spec_decodes].copy_(

224 + spec_sequence_masks, non_blocking=True)

225 + spec_sequence_masks = self.spec_sequence_masks[:batch_size]

226 + spec_sequence_masks[num_spec_decodes:].fill_(False)

227 +

228 + assert spec_token_masks is not None

229 + self.spec_token_masks[:spec_token_masks.size(0)].copy_(

230 + spec_token_masks, non_blocking=True)

231 + spec_token_masks = self.spec_token_masks[:m.num_actual_tokens]

232 + spec_token_masks[spec_token_masks.size(0):].fill_(False)

233 +

234 + self.spec_query_start_loc[:num_spec_decodes + 1].copy_(

235 + spec_query_start_loc, non_blocking=True)

236 + spec_num_query_tokens = spec_query_start_loc[

237 + -1] # type: ignore[index]

238 + spec_query_start_loc = self.spec_query_start_loc[:batch_size + 1]

239 + spec_query_start_loc[num_spec_decodes +

240 + 1:].fill_(spec_num_query_tokens)

241 +

242 + self.num_accepted_tokens[:num_spec_decodes].copy_(

243 + num_accepted_tokens, non_blocking=True)

244 + num_accepted_tokens = self.num_accepted_tokens[:batch_size]

245 + num_accepted_tokens[num_spec_decodes:].fill_(1)

246 +

247 + if (self.use_full_cuda_graph and num_prefills == 0

248 + and num_spec_decodes == 0

249 + and num_decodes <= self.decode_cudagraph_max_bs):

250 + num_total_tokens = self.vllm_config.pad_for_cudagraph(

251 + m.num_actual_tokens)

252 + batch_size = num_total_tokens

253 +

254 + self.non_spec_state_indices_tensor[:num_decodes].copy_(

255 + non_spec_state_indices_tensor, non_blocking=True)

256 + non_spec_state_indices_tensor = \

257 + self.non_spec_state_indices_tensor[:batch_size]

258 + non_spec_state_indices_tensor[num_decodes:].fill_(PAD_SLOT_ID)

259 +

260 + self.non_spec_query_start_loc[:num_decodes + 1].copy_(

261 + non_spec_query_start_loc, non_blocking=True)

262 + non_spec_num_query_tokens = non_spec_query_start_loc[

263 + -1] # type: ignore[index]

264 + non_spec_query_start_loc = \

265 + self.non_spec_query_start_loc[:batch_size + 1]

266 + non_spec_query_start_loc[num_decodes +

267 + 1:].fill_(non_spec_num_query_tokens)

268 +

269 + attn_metadata = GDNAttentionMetadata(

270 + num_prefills=num_prefills,

271 + num_prefill_tokens=num_prefill_tokens,

272 + num_decodes=num_decodes,

273 + num_decode_tokens=num_decode_tokens,

274 + num_spec_decodes=num_spec_decodes,

275 + num_spec_decode_tokens=num_spec_decode_tokens,

276 + has_initial_state=has_initial_state,

277 + spec_query_start_loc=spec_query_start_loc,

278 + non_spec_query_start_loc=non_spec_query_start_loc,

279 + spec_state_indices_tensor=spec_state_indices_tensor,

280 + non_spec_state_indices_tensor=non_spec_state_indices_tensor,

281 + spec_sequence_masks=spec_sequence_masks,

282 + spec_token_masks=spec_token_masks,

283 + num_accepted_tokens=num_accepted_tokens,

284 + )

285 + return attn_metadata

286 +

287 + def build_for_cudagraph_capture(

288 + self, common_attn_metadata: CommonAttentionMetadata):

289 + """

290 + This method builds the metadata for full cudagraph capture.

291 + Currently, only decode is supported for full cudagraphs with Mamba.

292 + """

293 + m = common_attn_metadata

294 +

295 + assert (m.num_reqs * (self.num_spec + 1) <= m.num_actual_tokens

296 + and ((m.num_reqs + 1) * (self.num_spec + 1)

297 + >= m.num_actual_tokens)), \

298 + "GDN only supports decode-only full CUDAGraph capture. " \

299 + "Make sure all cudagraph capture sizes <= max_num_seq."

300 +

301 + num_accepted_tokens = torch.full((m.num_reqs, ),

302 + m.max_query_len,

303 + dtype=torch.int32,

304 + device=m.query_start_loc.device)

305 + num_drafted_tokens = torch.full((m.num_reqs, ),

306 + self.num_spec,

307 + dtype=torch.int32,

308 + device=m.query_start_loc.device)

309 +

310 + # Fixes query-start loc for spec-sequence-indices.

311 + m.query_start_loc = torch.arange(0,

312 + m.num_actual_tokens + 1,

313 + step=m.max_query_len,

314 + device=m.query_start_loc.device,

315 + dtype=torch.int32)

316 + m.num_computed_tokens_cpu = (m.seq_lens_cpu - torch.full(

317 + (m.num_reqs, ), m.max_query_len, dtype=torch.int32, device='cpu'))

318 +

319 + return self.build(0, m, num_accepted_tokens, num_drafted_tokens)

vllm/v1/core/single_type_kv_cache_manager.py

+40 -4

@@ -559,12 +559,48 @@ class MambaManager(SingleTypeKVCacheManager):

559 num_running_requests: int) -> int:

560 return 0

561

562 + def get_num_blocks_to_allocate(

563 + self, request_id: str, num_tokens: int,

564 + new_computed_blocks: list[KVCacheBlock]) -> int:

565 + """

566 + Get the number of blocks needed to be allocated for the request.

567 +

568 + Args:

569 + request_id: The request ID.

570 + num_tokens: The total number of tokens that need a slot (including

571 + tokens that are already allocated).

572 + new_computed_blocks: The new computed blocks just hitting the

573 + prefix caching.

574 +

575 + Returns:

576 + The number of blocks

577 + """

578 +

579 + assert isinstance(self.kv_cache_spec, MambaSpec)

580 + if self.kv_cache_spec.num_speculative_blocks > 0:

581 + num_tokens += (self.kv_cache_spec.block_size *

582 + self.kv_cache_spec.num_speculative_blocks)

583 + num_required_blocks = cdiv(num_tokens, self.block_size)

584 + num_new_blocks = (num_required_blocks - len(new_computed_blocks) -

585 + len(self.req_to_blocks[request_id]))

586 + # If a computed block of a request is an eviction candidate (in the

587 + # free queue and ref_cnt == 0), it will be changed from a free block

588 + # to a computed block when the request is allocated, so we also count

589 + # it as needed to be allocated.

590 + num_evictable_computed_blocks = sum(

591 + blk.ref_cnt == 0 and not blk.is_null

592 + for blk in new_computed_blocks)

593 + return num_new_blocks + num_evictable_computed_blocks

594 +

595 def allocate_new_blocks(self, request_id: str,

596 num_tokens: int) -> list[KVCacheBlock]:

564 - new_blocks = super().allocate_new_blocks(request_id, num_tokens)

565 - assert len(self.req_to_blocks[request_id]) == 1, (

566 - "MambaManager should only allocate 1 block for each request.")

567 - return new_blocks

597 + # Allocate extra `num_speculative_blocks` blocks for

598 + # speculative decoding (MTP/EAGLE) with linear attention.

599 + assert isinstance(self.kv_cache_spec, MambaSpec)

600 + if self.kv_cache_spec.num_speculative_blocks > 0:

601 + num_tokens += (self.kv_cache_spec.block_size *

602 + self.kv_cache_spec.num_speculative_blocks)

603 + return super().allocate_new_blocks(request_id, num_tokens)

604

605

606 class CrossAttentionManager(SingleTypeKVCacheManager):

vllm/v1/kv_cache_interface.py

@@ -194,6 +194,7 @@ class MambaSpec(KVCacheSpec):

194 dtypes: tuple[torch.dtype]

195 page_size_padded: Optional[int] = None

196 mamba_type: str = "mamba2"

197 + num_speculative_blocks: int = 0

198

199 @property

200 def page_size_bytes(self) -> int:

vllm/v1/spec_decode/eagle.py

+8 -2

@@ -218,7 +218,7 @@ class EagleProposer:

218 hidden_states=self.hidden_states[:num_input_tokens],

219 inputs_embeds=inputs_embeds,

220 )

221 - if self.method in ("deepseek_mtp", "ernie_mtp"):

221 + if self.method in ("deepseek_mtp", "ernie_mtp", "qwen3_next_mtp"):

222 last_hidden_states = ret_hidden_states

223 hidden_states = last_hidden_states

224 else:

@@ -322,12 +322,18 @@ class EagleProposer:

322 with set_forward_context(per_layer_attn_metadata,

323 self.vllm_config,

324 num_tokens=input_batch_size):

325 - last_hidden_states, hidden_states = self.model(

325 + ret_hidden_states = self.model(

326 input_ids=input_ids,

327 positions=self.positions[:input_batch_size],

328 hidden_states=self.hidden_states[:input_batch_size],

329 inputs_embeds=inputs_embeds,

330 )

331 + if self.method in ("deepseek_mtp", "ernie_mtp",

332 + "qwen3_next_mtp"):

333 + last_hidden_states = ret_hidden_states

334 + hidden_states = ret_hidden_states

335 + else:

336 + last_hidden_states, hidden_states = ret_hidden_states

337 hidden_states = hidden_states[:batch_size]

338 logits = self.model.compute_logits(last_hidden_states[:batch_size],

339 None)

vllm/v1/worker/block_table.py

+13 -7

@@ -156,9 +156,14 @@ class BlockTable:

156 class MultiGroupBlockTable:

157 """The BlockTables for each KV cache group."""

158

159 - def __init__(self, max_num_reqs: int, max_model_len: int,

160 - max_num_batched_tokens: int, pin_memory: bool,

161 - device: torch.device, block_sizes: list[int]) -> None:

159 + def __init__(self,

160 + max_num_reqs: int,

161 + max_model_len: int,

162 + max_num_batched_tokens: int,

163 + pin_memory: bool,

164 + device: torch.device,

165 + block_sizes: list[int],

166 + num_speculative_tokens: int = 0) -> None:

167 # Note(hc): each dcp rank only store

168 # (max_model_len//dcp_world_size) tokens in kvcache,

169 # so the block_size which used for calc max_num_blocks_per_req

@@ -170,10 +175,11 @@ class MultiGroupBlockTable:

175 dcp_world_size = 1

176

177 self.block_tables = [

173 - BlockTable(block_size, max_num_reqs,

174 - cdiv(max_model_len, block_size * dcp_world_size),

175 - max_num_batched_tokens, pin_memory, device)

176 - for block_size in block_sizes

178 + BlockTable(

179 + block_size, max_num_reqs,

180 + max(cdiv(max_model_len, block_size * dcp_world_size),

181 + 1 + num_speculative_tokens), max_num_batched_tokens,

182 + pin_memory, device) for block_size in block_sizes

183 ]

184

185 def append_row(self, block_ids: tuple[list[int], ...],

vllm/v1/worker/gpu_input_batch.py

+17

@@ -83,6 +83,7 @@ class InputBatch:

83 logitsprocs: Optional[LogitsProcessors] = None,

84 is_spec_decode: bool = False,

85 is_pooling_model: bool = False,

86 + num_speculative_tokens: int = 0,

87 ):

88 self.is_pooling_model = is_pooling_model

89 self.is_spec_decode = is_spec_decode

@@ -127,6 +128,7 @@ class InputBatch:

128 pin_memory=pin_memory,

129 device=device,

130 block_sizes=block_sizes,

131 + num_speculative_tokens=num_speculative_tokens,

132 )

133

134 # Sampling-related.

@@ -202,6 +204,14 @@ class InputBatch:

204 self.repetition_penalties_cpu_tensor.numpy()

205 self.repetition_penalties_reqs: set[str] = set()

206

207 + # Speculative decoding

208 + self.num_accepted_tokens_cpu_tensor = torch.ones((max_num_reqs, ),

209 + dtype=torch.int64,

210 + device="cpu",

211 + pin_memory=pin_memory)

212 + self.num_accepted_tokens_cpu = \

213 + self.num_accepted_tokens_cpu_tensor.numpy()

214 +

215 # lora related

216 self.request_lora_mapping = np.zeros((self.max_num_reqs, ),

217 dtype=np.int32)

@@ -394,6 +404,9 @@ class InputBatch:

404 else:

405 raise NotImplementedError("Unrecognized request type")

406

407 + # Speculative decoding: by default 1 token is generated.

408 + self.num_accepted_tokens_cpu[req_index] = 1

409 +

410 # Add request lora ID

411 if request.lora_request:

412 lora_id = request.lora_request.lora_int_id

@@ -515,6 +528,8 @@ class InputBatch:

528 self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]

529 self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = \

530 self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]

531 + self.num_accepted_tokens_cpu[i1], self.num_accepted_tokens_cpu[i2] =\

532 + self.num_accepted_tokens_cpu[i2], self.num_accepted_tokens_cpu[i1]

533

534 swap_dict_values(self.generators, i1, i2)

535 swap_dict_values(self.bad_words_token_ids, i1, i2)

@@ -609,6 +624,8 @@ class InputBatch:

624 empty_index] = self.presence_penalties_cpu[last_req_index]

625 self.repetition_penalties_cpu[

626 empty_index] = self.repetition_penalties_cpu[last_req_index]

627 + self.num_accepted_tokens_cpu[

628 + empty_index] = self.num_accepted_tokens_cpu[last_req_index]

629 generator = self.generators.pop(last_req_index, None)

630 if generator is not None:

631 self.generators[empty_index] = generator

vllm/v1/worker/gpu_model_runner.py

+65 -9

@@ -53,9 +53,9 @@ from vllm.sampling_params import SamplingType

53 from vllm.sequence import IntermediateTensors, PoolerOutput

54 from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

55 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,

56 - GiB_bytes, LazyLoader, cdiv, check_use_alibi,

57 - get_dtype_size, is_pin_memory_available, round_up,

58 - supports_dynamo)

56 + GiB_bytes, LazyLoader, check_use_alibi, get_dtype_size,

57 + is_pin_memory_available, round_up, supports_dynamo)

58 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder

59 from vllm.v1.attention.backends.utils import (

60 AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,

61 create_fast_prefill_custom_backend,

@@ -324,6 +324,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

324 self.hidden_size,

325 dtype=self.dtype,

326 numpy=False)

327 + self.num_draft_tokens = self._make_buffer(self.max_num_reqs,

328 + dtype=torch.int32)

329 + self.num_accepted_tokens = self._make_buffer(self.max_num_reqs,

330 + dtype=torch.int64)

331

332 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)

333 if self.uses_mrope:

@@ -663,6 +667,31 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

667 # Refresh batch metadata with any pending updates.

668 self.input_batch.refresh_metadata()

669

670 + def _update_states_after_model_execute(

671 + self, output_token_ids: torch.Tensor) -> None:

672 + """Update the cached states after model execution.

673 +

674 + This is used for MTP/EAGLE for hybrid models, as in linear attention,

675 + only the last token's state is kept. In MTP/EAGLE, for draft tokens

676 + the state are kept util we decide how many tokens are accepted for

677 + each sequence, and a shifting is done during the next iteration

678 + based on the number of accepted tokens.

679 + """

680 + if not self.model_config.is_hybrid or not self.speculative_config:

681 + return

682 +

683 + # Find the number of accepted tokens for each sequence.

684 + num_accepted_tokens = (torch.cat(

685 + [

686 + output_token_ids,

687 + torch.full((output_token_ids.size(0), 1),

688 + -1,

689 + device=output_token_ids.device),

690 + ],

691 + dim=1) == -1).int().argmax(-1).cpu().numpy()

692 + for i, num_tokens in enumerate(num_accepted_tokens):

693 + self.input_batch.num_accepted_tokens_cpu[i] = num_tokens

694 +

695 def _init_mrope_positions(self, req_state: CachedRequestState):

696 image_grid_thw = []

697 video_grid_thw = []

@@ -936,6 +965,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

965 # We will ignore the sampled tokens from the partial requests.

966 # TODO: Support prompt logprobs.

967 logits_indices = query_start_loc[1:] - 1

968 + num_draft_tokens = None

969 spec_decode_metadata = None

970 else:

971 # Get the number of draft tokens for each request.

@@ -950,6 +980,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

980 spec_decode_metadata = self._calc_spec_decode_metadata(

981 num_draft_tokens, cu_num_tokens)

982 logits_indices = spec_decode_metadata.logits_indices

983 + self.num_draft_tokens.np[:num_reqs] = num_draft_tokens

984 + self.num_draft_tokens.np[num_reqs:].fill(0)

985 + self.num_draft_tokens.copy_to_gpu()

986

987 logits_indices_padded = None

988 if self.cache_config.kv_sharing_fast_prefill:

@@ -964,6 +997,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

997 num_computed_tokens_cpu = (

998 self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs])

999 spec_decode_common_attn_metadata = None

1000 + if use_spec_decode:

1001 + self.num_accepted_tokens.np[:num_reqs] = (

1002 + self.input_batch.num_accepted_tokens_cpu[:num_reqs])

1003 + self.num_accepted_tokens.np[num_reqs:].fill(1)

1004 + self.num_accepted_tokens.copy_to_gpu()

1005

1006 # Prepare the attention metadata for each KV cache group and make layers

1007 # in the same group share the same metadata.

@@ -1034,10 +1072,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1072 builder,

1073 )

1074

1075 + extra_attn_metadata_args = {}

1076 + if use_spec_decode and isinstance(builder,

1077 + GDNAttentionMetadataBuilder):

1078 + extra_attn_metadata_args = dict(

1079 + num_accepted_tokens=self.num_accepted_tokens.

1080 + gpu[:num_reqs],

1081 + num_draft_tokens=self.num_draft_tokens.gpu[:num_reqs],

1082 + )

1083 +

1084 attn_metadata_i = builder.build(

1085 common_prefix_len=common_prefix_len,

1086 common_attn_metadata=common_attn_metadata,

1040 - )

1087 + **extra_attn_metadata_args)

1088

1089 for layer_name in attn_group.layer_names:

1090 attn_metadata[layer_name] = attn_metadata_i

@@ -1814,6 +1861,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1861 sampling_metadata,

1862 )

1863 sampler_output.sampled_token_ids = output_token_ids

1864 + self._update_states_after_model_execute(output_token_ids)

1865

1866 return sampler_output

1867

@@ -2644,13 +2692,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2692 # Note: Overriding max_query_len to be the prefill tokens

2693 max_query_len = num_prefill_tokens

2694 elif uniform_decode:

2647 - assert not create_mixed_batch

2648 - num_reqs = cdiv(num_tokens, max_query_len)

2695 + num_reqs = num_tokens // max_query_len

2696 assert num_reqs <= max_num_reqs, \

2697 "Do not capture num_reqs > max_num_reqs for uniform batch"

2698 num_scheduled_tokens_list = [max_query_len] * num_reqs

2699 if num_tokens % max_query_len != 0:

2653 - num_scheduled_tokens_list[-1] = num_tokens % max_query_len

2700 + num_scheduled_tokens_list[-1] += num_tokens % max_query_len

2701 else:

2702 num_reqs = min(num_tokens, max_num_reqs)

2703 min_tokens_per_req = num_tokens // num_reqs

@@ -3297,6 +3344,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3344 is_spec_decode=bool(self.vllm_config.speculative_config),

3345 logitsprocs=self.input_batch.logitsprocs,

3346 is_pooling_model=self.is_pooling_model,

3347 + num_speculative_tokens=(

3348 + self.vllm_config.speculative_config.num_speculative_tokens

3349 + if self.vllm_config.speculative_config else 0),

3350 )

3351

3352 def _allocate_kv_cache_tensors(

@@ -3647,7 +3697,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3697

3698 mamba_layers = get_layers_from_vllm_config(self.vllm_config, MambaBase)

3699 if len(mamba_layers) > 0:

3650 - if self.vllm_config.speculative_config is not None:

3700 + if (self.vllm_config.speculative_config is not None

3701 + and self.vllm_config.model_config.hf_config.model_type

3702 + not in ["qwen3_next"]):

3703 raise NotImplementedError(

3704 "Mamba with speculative decoding is not supported yet.")

3705 if self.vllm_config.cache_config.enable_prefix_caching:

@@ -3666,7 +3718,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3718 dtypes=mamba_module.get_state_dtype(),

3719 block_size=max_model_len,

3720 page_size_padded=page_size_padded,

3669 - mamba_type=mamba_module.mamba_type)

3721 + mamba_type=mamba_module.mamba_type,

3722 + num_speculative_blocks=(

3723 + self.speculative_config.num_speculative_tokens

3724 + if self.speculative_config else 0),

3725 + )

3726

3727 return kv_cache_spec

3728

vllm/worker/worker.py

+2 -1

@@ -78,7 +78,8 @@ class Worker(LocalOrDistributedWorkerBase):

78 "deepseek_mtp",

79 "glm4_moe_mtp",

80 "mimo_mtp",

81 - "ernie_mtp")) \

81 + "ernie_mtp",

82 + "qwen3_next_mtp")) \

83 else {"return_hidden_states": True}

85 ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner