| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """GPU-resident Qwen4Exp position-learning enhancement layers.""" |
| 4 | |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | from transformers import Qwen4ExpTextConfig |
| 10 | |
| 11 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 12 | from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config |
| 13 | from vllm.forward_context import get_forward_context |
| 14 | from vllm.model_executor.layers.linear import MergedColumnParallelLinear |
| 15 | from vllm.model_executor.layers.mamba.abstract import MambaBase |
| 16 | from vllm.model_executor.layers.mamba.mamba_utils import ( |
| 17 | MambaStateDtypeCalculator, |
| 18 | MambaStateShapeCalculator, |
| 19 | is_conv_state_dim_first, |
| 20 | ) |
| 21 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 22 | from vllm.v1.attention.backends.short_conv_attn import ( |
| 23 | PleShortConvAttentionBackend, |
| 24 | PleShortConvAttentionMetadata, |
| 25 | ) |
| 26 | |
| 27 | from .ngram_embedding import Qwen4ExpNGramEmbedding |
| 28 | from .ops.ple import ple_conv, ple_gate |
| 29 | |
| 30 | |
| 31 | class Qwen4ExpPLEGroupedNorm(nn.Module): |
| 32 | def __init__( |
| 33 | self, |
| 34 | hidden_size: int, |
| 35 | eps: float, |
| 36 | group_size: int | None, |
| 37 | dtype: torch.dtype | None, |
| 38 | ) -> None: |
| 39 | super().__init__() |
| 40 | if group_size is not None and hidden_size % group_size: |
| 41 | raise ValueError( |
| 42 | f"hidden_size ({hidden_size}) must be divisible by " |
| 43 | f"group_size ({group_size})" |
| 44 | ) |
| 45 | self.eps = eps |
| 46 | self.group_size = group_size |
| 47 | self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) |
| 48 | |
| 49 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 50 | input_dtype = hidden_states.dtype |
| 51 | hidden_states = hidden_states.float() |
| 52 | if self.group_size is None: |
| 53 | variance = hidden_states.square().mean(dim=-1, keepdim=True) |
| 54 | normalized = hidden_states * torch.rsqrt(variance + self.eps) |
| 55 | else: |
| 56 | grouped = hidden_states.unflatten( |
| 57 | -1, (hidden_states.shape[-1] // self.group_size, self.group_size) |
| 58 | ) |
| 59 | variance = grouped.square().mean(dim=-1, keepdim=True) |
| 60 | normalized = (grouped * torch.rsqrt(variance + self.eps)).flatten(-2) |
| 61 | return (normalized * (1.0 + self.weight.float())).to(input_dtype) |
| 62 | |
| 63 | |
| 64 | class Qwen4ExpPLELayer(nn.Module, MambaBase): |
| 65 | def __init__( |
| 66 | self, |
| 67 | config: Qwen4ExpTextConfig, |
| 68 | vllm_config: VllmConfig, |
| 69 | layer_idx: int = 0, |
| 70 | ple_dense_layer_id: int | None = None, |
| 71 | prefix: str = "", |
| 72 | ) -> None: |
| 73 | super().__init__() |
| 74 | model_config = vllm_config.model_config |
| 75 | cache_config = vllm_config.cache_config |
| 76 | quant_config = vllm_config.quant_config |
| 77 | self.model_config: ModelConfig = model_config |
| 78 | self.cache_config: CacheConfig = cache_config |
| 79 | self.layer_idx = layer_idx |
| 80 | self.ple_dense_layer_id = ( |
| 81 | int(ple_dense_layer_id) |
| 82 | if ple_dense_layer_id is not None |
| 83 | else int(layer_idx) |
| 84 | ) |
| 85 | self.prefix = prefix |
| 86 | self.hidden_size = int(config.hidden_size) |
| 87 | self.hc_count = config.hc_count |
| 88 | self.hc_hidden_size = self.hidden_size * self.hc_count |
| 89 | self.conv_kernel_size = int(config.ple_conv_kernel_size) |
| 90 | self.short_conv_dilation = int(config.ngram_size) |
| 91 | self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation |
| 92 | self.num_spec_tokens = vllm_config.num_speculative_tokens |
| 93 | self.activation = "silu" |
| 94 | self.ple_embedding = Qwen4ExpNGramEmbedding( |
| 95 | config, |
| 96 | int(config.ple_embed_dim), |
| 97 | self.ple_dense_layer_id, |
| 98 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 99 | data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, |
| 100 | prefix=f"{prefix}.ple_embedding", |
| 101 | quant_config=quant_config, |
| 102 | params_dtype=model_config.dtype, |
| 103 | ) |
| 104 | # The PLE cache is TP-replicated, so this merged projection is too. |
| 105 | self.kv_proj = MergedColumnParallelLinear( |
| 106 | int(config.ple_embed_dim), |
| 107 | [self.hc_hidden_size, self.hidden_size], |
| 108 | bias=False, |
| 109 | params_dtype=model_config.dtype, |
| 110 | quant_config=quant_config, |
| 111 | prefix=f"{prefix}.kv_proj", |
| 112 | disable_tp=True, |
| 113 | ) |
| 114 | norm_args = ( |
| 115 | self.hc_hidden_size, |
| 116 | config.rms_norm_eps, |
| 117 | self.hidden_size, |
| 118 | model_config.dtype, |
| 119 | ) |
| 120 | self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args) |
| 121 | self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args) |
| 122 | self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args) |
| 123 | self.conv1d = nn.Conv1d( |
| 124 | self.hc_hidden_size, |
| 125 | self.hc_hidden_size, |
| 126 | self.conv_kernel_size, |
| 127 | groups=self.hc_hidden_size, |
| 128 | padding=self.conv_state_len, |
| 129 | dilation=self.short_conv_dilation, |
| 130 | bias=False, |
| 131 | dtype=model_config.dtype, |
| 132 | ) |
| 133 | nn.init.zeros_(self.conv1d.weight) |
| 134 | self.conv1d.weight._no_reinit = True |
| 135 | self.kv_cache = (torch.tensor([]),) |
| 136 | compilation_config = get_current_vllm_config().compilation_config |
| 137 | if prefix in compilation_config.static_forward_context: |
| 138 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 139 | compilation_config.static_forward_context[prefix] = self |
| 140 | |
| 141 | def _dequantize_embeddings( |
| 142 | self, |
| 143 | embeddings: torch.Tensor, |
| 144 | output_dtype: torch.dtype, |
| 145 | ) -> torch.Tensor: |
| 146 | """Dequantize PLE lookup output.""" |
| 147 | return self.ple_embedding.ngram_embedding.dequantize( |
| 148 | embeddings, |
| 149 | output_dtype, |
| 150 | ) |
| 151 | |
| 152 | def start_prefetch( |
| 153 | self, |
| 154 | hidden_states: torch.Tensor, |
| 155 | input_ids: torch.Tensor, |
| 156 | query_start_loc: torch.Tensor, |
| 157 | ngram_context: torch.Tensor, |
| 158 | ) -> None: |
| 159 | """Start the pinned PLE lookup while the preceding decoder layer runs.""" |
| 160 | self.ple_embedding.start_prefetch( |
| 161 | hidden_states, |
| 162 | input_ids, |
| 163 | query_start_loc, |
| 164 | ngram_context, |
| 165 | ) |
| 166 | |
| 167 | @property |
| 168 | def mamba_type(self) -> MambaAttentionBackendEnum: |
| 169 | return MambaAttentionBackendEnum.SHORT_CONV |
| 170 | |
| 171 | @property |
| 172 | def is_kv_cache_tp_replicated(self) -> bool: |
| 173 | return True |
| 174 | |
| 175 | def get_attn_backend(self) -> type[PleShortConvAttentionBackend]: |
| 176 | return PleShortConvAttentionBackend |
| 177 | |
| 178 | def get_state_dtype(self) -> tuple[torch.dtype, ...]: |
| 179 | return MambaStateDtypeCalculator.short_conv_state_dtype( |
| 180 | self.model_config.dtype, self.cache_config.mamba_cache_dtype |
| 181 | ) |
| 182 | |
| 183 | def get_state_shape(self) -> Sequence[tuple[int, ...]]: |
| 184 | return MambaStateShapeCalculator.short_conv_state_shape( |
| 185 | tp_world_size=1, |
| 186 | intermediate_size=self.hc_hidden_size, |
| 187 | conv_kernel=self.conv_state_len + 1, |
| 188 | num_spec=self.num_spec_tokens, |
| 189 | ) |
| 190 | |
| 191 | def _short_conv_dilated_dispatch( |
| 192 | self, |
| 193 | inputs: torch.Tensor, |
| 194 | residual: torch.Tensor, |
| 195 | outer_residual: torch.Tensor, |
| 196 | metadata: PleShortConvAttentionMetadata, |
| 197 | conv_state: torch.Tensor, |
| 198 | conv_weights: torch.Tensor, |
| 199 | ) -> None: |
| 200 | num_prefills = metadata.num_prefills |
| 201 | num_decodes = metadata.num_decodes |
| 202 | num_decode_tokens = metadata.num_decode_tokens |
| 203 | num_prefill_tokens = metadata.num_prefill_tokens |
| 204 | has_prefill = num_prefills > 0 |
| 205 | has_decode = num_decodes > 0 |
| 206 | has_spec = metadata.spec_sequence_masks is not None |
| 207 | has_non_spec = has_prefill or has_decode |
| 208 | inputs = inputs[: metadata.num_actual_tokens] |
| 209 | residual = residual[: metadata.num_actual_tokens] |
| 210 | outer_residual = outer_residual[: metadata.num_actual_tokens] |
| 211 | |
| 212 | spec_token_indices = None |
| 213 | non_spec_token_indices = None |
| 214 | if has_spec and has_non_spec: |
| 215 | assert metadata.spec_token_indx is not None |
| 216 | assert metadata.non_spec_token_indx is not None |
| 217 | spec_token_indices = metadata.spec_token_indx |
| 218 | non_spec_token_indices = metadata.non_spec_token_indx |
| 219 | |
| 220 | if has_spec: |
| 221 | assert metadata.spec_state_indices_tensor is not None |
| 222 | query_start_loc = metadata.spec_query_start_loc |
| 223 | num_accepted_tokens = metadata.num_accepted_tokens |
| 224 | assert query_start_loc is not None |
| 225 | assert num_accepted_tokens is not None |
| 226 | spec_state_indices = metadata.spec_state_indices_tensor[ |
| 227 | : metadata.num_spec_decodes |
| 228 | ] |
| 229 | # Mixed batches stay in their original row order; the kernels map |
| 230 | # logical spec/non-spec rows instead of materializing both groups. |
| 231 | ple_conv( |
| 232 | inputs=inputs, |
| 233 | residual=residual, |
| 234 | conv_state=conv_state, |
| 235 | conv_weights=conv_weights, |
| 236 | state_indices=spec_state_indices, |
| 237 | outer_residual=outer_residual, |
| 238 | mode="spec", |
| 239 | dilation=self.short_conv_dilation, |
| 240 | query_start_loc=query_start_loc, |
| 241 | num_accepted_tokens=num_accepted_tokens, |
| 242 | spec_query_len=metadata.spec_query_len, |
| 243 | token_indices=spec_token_indices, |
| 244 | ) |
| 245 | |
| 246 | if not has_non_spec: |
| 247 | return |
| 248 | |
| 249 | state_indices = metadata.state_indices_tensor |
| 250 | assert state_indices is not None |
| 251 | if has_prefill: |
| 252 | state_indices_d, state_indices_p = torch.split( |
| 253 | state_indices, [num_decodes, num_prefills], dim=0 |
| 254 | ) |
| 255 | if non_spec_token_indices is None: |
| 256 | inputs_d, inputs_p = torch.split( |
| 257 | inputs, [num_decode_tokens, num_prefill_tokens], dim=0 |
| 258 | ) |
| 259 | residual_d, residual_p = torch.split( |
| 260 | residual, [num_decode_tokens, num_prefill_tokens], dim=0 |
| 261 | ) |
| 262 | outer_residual_d, outer_residual_p = torch.split( |
| 263 | outer_residual, |
| 264 | [num_decode_tokens, num_prefill_tokens], |
| 265 | dim=0, |
| 266 | ) |
| 267 | token_indices_d = None |
| 268 | token_indices_p = None |
| 269 | else: |
| 270 | inputs_d = inputs_p = inputs |
| 271 | residual_d = residual_p = residual |
| 272 | outer_residual_d = outer_residual_p = outer_residual |
| 273 | token_indices_d, token_indices_p = torch.split( |
| 274 | non_spec_token_indices, |
| 275 | [num_decode_tokens, num_prefill_tokens], |
| 276 | dim=0, |
| 277 | ) |
| 278 | |
| 279 | if has_decode: |
| 280 | ple_conv( |
| 281 | inputs=inputs_d, |
| 282 | residual=residual_d, |
| 283 | conv_state=conv_state, |
| 284 | conv_weights=conv_weights, |
| 285 | state_indices=state_indices_d, |
| 286 | outer_residual=outer_residual_d, |
| 287 | mode="decode", |
| 288 | dilation=self.short_conv_dilation, |
| 289 | has_initial_states=metadata.has_initial_states_d, |
| 290 | token_indices=token_indices_d, |
| 291 | ) |
| 292 | |
| 293 | query_start_loc = metadata.query_start_loc_p |
| 294 | if query_start_loc is None: |
| 295 | raise ValueError("query_start_loc is required for prefill short-conv") |
| 296 | has_initial_states = metadata.has_initial_states_p |
| 297 | if has_initial_states is None: |
| 298 | raise ValueError( |
| 299 | "has_initial_states_p is required for prefill short-conv" |
| 300 | ) |
| 301 | ple_conv( |
| 302 | inputs=inputs_p, |
| 303 | residual=residual_p, |
| 304 | conv_state=conv_state, |
| 305 | conv_weights=conv_weights, |
| 306 | state_indices=state_indices_p, |
| 307 | outer_residual=outer_residual_p, |
| 308 | mode="prefill", |
| 309 | dilation=self.short_conv_dilation, |
| 310 | query_start_loc=query_start_loc, |
| 311 | has_initial_states=has_initial_states, |
| 312 | token_indices=token_indices_p, |
| 313 | ) |
| 314 | else: |
| 315 | num_decode_rows = ( |
| 316 | non_spec_token_indices.numel() |
| 317 | if non_spec_token_indices is not None |
| 318 | else inputs.size(0) |
| 319 | ) |
| 320 | ple_conv( |
| 321 | inputs=inputs, |
| 322 | residual=residual, |
| 323 | conv_state=conv_state, |
| 324 | conv_weights=conv_weights, |
| 325 | state_indices=state_indices[:num_decode_rows], |
| 326 | outer_residual=outer_residual, |
| 327 | mode="decode", |
| 328 | dilation=self.short_conv_dilation, |
| 329 | has_initial_states=metadata.has_initial_states_d, |
| 330 | token_indices=non_spec_token_indices, |
| 331 | ) |
| 332 | |
| 333 | # State routing consumes the current request metadata on every replay. |
| 334 | @eager_break_during_capture |
| 335 | def _short_conv( |
| 336 | self, |
| 337 | inputs: torch.Tensor, |
| 338 | residual: torch.Tensor, |
| 339 | outer_residual: torch.Tensor, |
| 340 | ) -> None: |
| 341 | forward_context = get_forward_context() |
| 342 | attn_metadata = forward_context.attn_metadata |
| 343 | # Profiling omits all metadata or this Mamba entry. Short convolution |
| 344 | # is a no-op there, but preserve the outer residual addition. |
| 345 | if attn_metadata is None: |
| 346 | residual.add_(outer_residual) |
| 347 | return |
| 348 | |
| 349 | if not isinstance(attn_metadata, dict): |
| 350 | raise RuntimeError( |
| 351 | "PLE short-conv expects per-layer attention metadata dict " |
| 352 | f"during inference, got {type(attn_metadata).__name__}." |
| 353 | ) |
| 354 | |
| 355 | layer_attn_metadata = attn_metadata.get(self.prefix) |
| 356 | if layer_attn_metadata is None: |
| 357 | residual.add_(outer_residual) |
| 358 | return |
| 359 | if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): |
| 360 | raise TypeError( |
| 361 | "Expected PleShortConvAttentionMetadata for layer " |
| 362 | f"'{self.prefix}', got " |
| 363 | f"{type(layer_attn_metadata).__name__}." |
| 364 | ) |
| 365 | |
| 366 | conv_state = self.kv_cache[0] |
| 367 | # Canonicalize both backend cache layouts to [slot, channel, window]. |
| 368 | if not is_conv_state_dim_first(): |
| 369 | conv_state = conv_state.transpose(-1, -2) |
| 370 | conv_weights = self.conv1d.weight.squeeze(1) |
| 371 | |
| 372 | state_capacity = self.conv_state_len + self.num_spec_tokens |
| 373 | if state_capacity > 0: |
| 374 | state_size = conv_state.size(-1) |
| 375 | if state_size < state_capacity: |
| 376 | raise RuntimeError( |
| 377 | "PLE short-conv cache is smaller than expected for " |
| 378 | f"dilated convolution: got {state_size}, " |
| 379 | f"expect at least {state_capacity}." |
| 380 | ) |
| 381 | conv_state = conv_state[..., -state_capacity:] |
| 382 | self._short_conv_dilated_dispatch( |
| 383 | inputs=inputs, |
| 384 | residual=residual, |
| 385 | outer_residual=outer_residual, |
| 386 | metadata=layer_attn_metadata, |
| 387 | conv_state=conv_state, |
| 388 | conv_weights=conv_weights.to(dtype=inputs.dtype), |
| 389 | ) |
| 390 | |
| 391 | def forward( |
| 392 | self, |
| 393 | hidden_states: torch.Tensor, |
| 394 | input_ids: torch.Tensor, |
| 395 | query_start_loc: torch.Tensor, |
| 396 | ngram_context: torch.Tensor, |
| 397 | ) -> torch.Tensor: |
| 398 | input_ids = input_ids.reshape(-1) |
| 399 | if input_ids.shape[0] != hidden_states.shape[0]: |
| 400 | raise ValueError( |
| 401 | "PLE expects input_ids and hidden_states to have the same " |
| 402 | f"token length, got {input_ids.shape[0]} and " |
| 403 | f"{hidden_states.shape[0]}" |
| 404 | ) |
| 405 | embeddings = self.ple_embedding( |
| 406 | hidden_states, |
| 407 | input_ids, |
| 408 | query_start_loc, |
| 409 | ngram_context, |
| 410 | ) |
| 411 | embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) |
| 412 | kv, _ = self.kv_proj(embeddings) |
| 413 | key, value = kv.split(self.kv_proj.output_sizes, dim=-1) |
| 414 | gated_output, conv_input = ple_gate( |
| 415 | key, |
| 416 | value, |
| 417 | hidden_states, |
| 418 | self.norm_key.weight, |
| 419 | self.norm_query.weight, |
| 420 | self.norm_conv.weight, |
| 421 | self.norm_key.eps, |
| 422 | ) |
| 423 | # State routing depends on runtime request metadata and remains outside |
| 424 | # the graph; short convolution accumulates into gated_output. |
| 425 | self._short_conv(conv_input, gated_output, hidden_states) |
| 426 | return gated_output |
| 427 | |
| 428 | |
| 429 | __all__ = [ |
| 430 | "Qwen4ExpPLEGroupedNorm", |
| 431 | "Qwen4ExpPLELayer", |
| 432 | ] |