| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # LiLiCorr head adapted from sgl-project/sglang PR #37462 (Apache-2.0). |
| 4 | """DFlash backbone with the LiLiCorr candidate-lattice correlator.""" |
| 5 | |
| 6 | from collections.abc import Iterable |
| 7 | from typing import Any |
| 8 | |
| 9 | import torch |
| 10 | import torch.nn.functional as F |
| 11 | from torch import nn |
| 12 | |
| 13 | from vllm.compilation.backends import set_model_tag |
| 14 | from vllm.compilation.decorators import support_torch_compile |
| 15 | from vllm.config import VllmConfig |
| 16 | from vllm.model_executor.layers.linear import ( |
| 17 | LinearBase, |
| 18 | ReplicatedLinear, |
| 19 | UnquantizedLinearMethod, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 22 | from vllm.model_executor.layers.quantization.base_config import QuantizationConfig |
| 23 | from vllm.model_executor.layers.quantization.modelopt import ModelOptLinearMethod |
| 24 | from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Static |
| 25 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 26 | |
| 27 | from .nemotron_h import NemotronHMLP |
| 28 | from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model |
| 29 | from .qwen3_dflash2 import DFlash2Qwen3DecoderLayer |
| 30 | from .utils import AutoWeightsLoader, maybe_prefix |
| 31 | |
| 32 | |
| 33 | class LiLiCorrLatticeAttention(nn.Module): |
| 34 | def __init__( |
| 35 | self, |
| 36 | hidden_size: int, |
| 37 | num_heads: int, |
| 38 | quant_config: QuantizationConfig | None = None, |
| 39 | prefix: str = "", |
| 40 | ) -> None: |
| 41 | super().__init__() |
| 42 | if num_heads <= 0 or hidden_size % num_heads != 0: |
| 43 | raise ValueError( |
| 44 | f"LiLiCorr hidden_size={hidden_size} must be divisible " |
| 45 | f"by num_heads={num_heads}." |
| 46 | ) |
| 47 | self.hidden_size = hidden_size |
| 48 | self.num_heads = num_heads |
| 49 | self.head_dim = self.hidden_size // self.num_heads |
| 50 | # Exported lattice QKV parameters stay in the model's floating dtype. |
| 51 | self.in_proj_weight = nn.Parameter( |
| 52 | torch.empty(3 * hidden_size, hidden_size), requires_grad=False |
| 53 | ) |
| 54 | self.in_proj_bias = nn.Parameter( |
| 55 | torch.empty(3 * hidden_size), requires_grad=False |
| 56 | ) |
| 57 | self.out_proj = ReplicatedLinear( |
| 58 | hidden_size, |
| 59 | hidden_size, |
| 60 | return_bias=False, |
| 61 | quant_config=quant_config, |
| 62 | prefix=maybe_prefix(prefix, "out_proj"), |
| 63 | ) |
| 64 | |
| 65 | def forward( |
| 66 | self, hidden_states: torch.Tensor, attention_bias: torch.Tensor |
| 67 | ) -> torch.Tensor: |
| 68 | bsz, seq_len, _ = hidden_states.shape |
| 69 | qkv = F.linear(hidden_states, self.in_proj_weight, self.in_proj_bias) |
| 70 | q, k, v = qkv.chunk(3, dim=-1) |
| 71 | shape = (bsz, seq_len, self.num_heads, self.head_dim) |
| 72 | q = q.view(shape).transpose(1, 2) |
| 73 | k = k.view(shape).transpose(1, 2) |
| 74 | v = v.view(shape).transpose(1, 2) |
| 75 | # Attention and MMEncoderAttention cannot accept this learned per-head |
| 76 | # additive bias. The full candidate lattice has no KV-cache state. |
| 77 | out = F.scaled_dot_product_attention( |
| 78 | q, |
| 79 | k, |
| 80 | v, |
| 81 | attn_mask=attention_bias, |
| 82 | ) |
| 83 | return self.out_proj( |
| 84 | out.transpose(1, 2).reshape(bsz, seq_len, self.hidden_size) |
| 85 | ) |
| 86 | |
| 87 | |
| 88 | class LiLiCorrLayer(nn.Module): |
| 89 | def __init__( |
| 90 | self, |
| 91 | hidden_size: int, |
| 92 | num_heads: int, |
| 93 | mlp_ratio: float, |
| 94 | rms_norm_eps: float, |
| 95 | quant_config: QuantizationConfig | None = None, |
| 96 | prefix: str = "", |
| 97 | ) -> None: |
| 98 | super().__init__() |
| 99 | self.attn_norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps) |
| 100 | self.attn = LiLiCorrLatticeAttention( |
| 101 | hidden_size, num_heads, quant_config, maybe_prefix(prefix, "attn") |
| 102 | ) |
| 103 | self.mlp_norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps) |
| 104 | mlp_hidden_size = int(hidden_size * mlp_ratio) |
| 105 | self.mlp = NemotronHMLP( |
| 106 | config=None, |
| 107 | hidden_size=hidden_size, |
| 108 | intermediate_size=mlp_hidden_size, |
| 109 | hidden_act="silu", |
| 110 | bias=True, |
| 111 | quant_config=quant_config, |
| 112 | prefix=maybe_prefix(prefix, "mlp"), |
| 113 | ) |
| 114 | |
| 115 | def forward( |
| 116 | self, hidden_states: torch.Tensor, attention_bias: torch.Tensor |
| 117 | ) -> torch.Tensor: |
| 118 | hidden_states = hidden_states + self.attn( |
| 119 | self.attn_norm(hidden_states), attention_bias |
| 120 | ) |
| 121 | return hidden_states + self.mlp(self.mlp_norm(hidden_states)) |
| 122 | |
| 123 | |
| 124 | @support_torch_compile |
| 125 | class LiLiCorrHead(nn.Module): |
| 126 | num_candidate_features = 5 |
| 127 | |
| 128 | def __init__( |
| 129 | self, |
| 130 | *, |
| 131 | model_hidden_size: int, |
| 132 | block_size: int, |
| 133 | rms_norm_eps: float, |
| 134 | config: dict[str, Any], |
| 135 | quant_config: QuantizationConfig | None = None, |
| 136 | prefix: str = "", |
| 137 | ) -> None: |
| 138 | super().__init__() |
| 139 | hidden_size = config["lilicorr_hidden_size"] or model_hidden_size |
| 140 | self.block_size = block_size |
| 141 | self.num_candidate_slots = block_size - 1 |
| 142 | self.candidate_topk = config["lilicorr_candidate_topk"] |
| 143 | if ( |
| 144 | self.candidate_topk <= 0 |
| 145 | or self.candidate_topk > 16 |
| 146 | or self.candidate_topk & (self.candidate_topk - 1) |
| 147 | ): |
| 148 | raise ValueError("LiLiCorr candidate_topk must be a power of two <= 16.") |
| 149 | self.hidden_size = hidden_size |
| 150 | self.num_heads = config["lilicorr_num_heads"] |
| 151 | self.mlp_ratio = config["lilicorr_mlp_ratio"] |
| 152 | self.factor_dim = config["lilicorr_factor_dim"] |
| 153 | self.vector_eps = config["lilicorr_vector_eps"] |
| 154 | self.logit_scale = config["lilicorr_logit_scale"] |
| 155 | # Candidate IDs and embeddings are replicated across TP ranks, so each |
| 156 | # rank scores the same lattice; the MLPs use standard tensor parallelism. |
| 157 | self.token_proj = ( |
| 158 | nn.Identity() |
| 159 | if model_hidden_size == hidden_size |
| 160 | else ReplicatedLinear( |
| 161 | model_hidden_size, |
| 162 | hidden_size, |
| 163 | return_bias=False, |
| 164 | quant_config=quant_config, |
| 165 | prefix=maybe_prefix(prefix, "token_proj"), |
| 166 | ) |
| 167 | ) |
| 168 | self.pass_hidden_proj = ReplicatedLinear( |
| 169 | model_hidden_size, |
| 170 | hidden_size, |
| 171 | return_bias=False, |
| 172 | quant_config=quant_config, |
| 173 | prefix=maybe_prefix(prefix, "pass_hidden_proj"), |
| 174 | ) |
| 175 | self.feature_norm = nn.LayerNorm(self.num_candidate_features) |
| 176 | self.feature_mlp = NemotronHMLP( |
| 177 | config=None, |
| 178 | hidden_size=self.num_candidate_features, |
| 179 | intermediate_size=hidden_size, |
| 180 | hidden_act="silu", |
| 181 | output_size=hidden_size, |
| 182 | bias=True, |
| 183 | quant_config=quant_config, |
| 184 | prefix=maybe_prefix(prefix, "feature_mlp"), |
| 185 | ) |
| 186 | self.slot_embedding = nn.Parameter( |
| 187 | torch.zeros(1, 1, self.num_candidate_slots, 1, hidden_size) |
| 188 | ) |
| 189 | self.rank_embedding = nn.Parameter( |
| 190 | torch.zeros(1, 1, 1, self.candidate_topk, hidden_size) |
| 191 | ) |
| 192 | self.relative_slot_bias = nn.Parameter( |
| 193 | torch.zeros(self.num_heads, 2 * self.block_size - 1) |
| 194 | ) |
| 195 | self.same_slot_bias = nn.Parameter(torch.zeros(self.num_heads)) |
| 196 | self.context_proj = ReplicatedLinear( |
| 197 | model_hidden_size, |
| 198 | hidden_size, |
| 199 | return_bias=False, |
| 200 | quant_config=quant_config, |
| 201 | prefix=maybe_prefix(prefix, "context_proj"), |
| 202 | ) |
| 203 | self.layers = nn.ModuleList( |
| 204 | [ |
| 205 | LiLiCorrLayer( |
| 206 | hidden_size=hidden_size, |
| 207 | num_heads=self.num_heads, |
| 208 | mlp_ratio=self.mlp_ratio, |
| 209 | rms_norm_eps=rms_norm_eps, |
| 210 | quant_config=quant_config, |
| 211 | prefix=maybe_prefix(prefix, f"layers.{i}"), |
| 212 | ) |
| 213 | for i in range(config["lilicorr_num_layers"]) |
| 214 | ] |
| 215 | ) |
| 216 | self.output_norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps) |
| 217 | self.anchor_norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps) |
| 218 | self.factor_input_proj = ReplicatedLinear( |
| 219 | hidden_size * 3, |
| 220 | hidden_size, |
| 221 | return_bias=False, |
| 222 | quant_config=None, |
| 223 | prefix=maybe_prefix(prefix, "factor_input_proj"), |
| 224 | ) |
| 225 | # Each candidate gets outgoing and incoming vectors. After normalization, |
| 226 | # their dot product scores a transition from slot s to slot s+1: |
| 227 | # pair_score[s, i, j] = dot(out[s, i], incoming[s+1, j]) |
| 228 | # The anchor supplies the outgoing vector for the first slot. |
| 229 | self.out_head = ReplicatedLinear( |
| 230 | hidden_size, |
| 231 | self.factor_dim, |
| 232 | return_bias=False, |
| 233 | quant_config=None, |
| 234 | prefix=maybe_prefix(prefix, "out_head"), |
| 235 | ) |
| 236 | self.in_head = ReplicatedLinear( |
| 237 | hidden_size, |
| 238 | self.factor_dim, |
| 239 | return_bias=False, |
| 240 | quant_config=None, |
| 241 | prefix=maybe_prefix(prefix, "in_head"), |
| 242 | ) |
| 243 | self.anchor_out_head = ReplicatedLinear( |
| 244 | hidden_size, |
| 245 | self.factor_dim, |
| 246 | return_bias=False, |
| 247 | quant_config=quant_config, |
| 248 | prefix=maybe_prefix(prefix, "anchor_out_head"), |
| 249 | ) |
| 250 | self._fused_edge_weight: torch.Tensor | None = None |
| 251 | self._fused_edge_bias: torch.Tensor | None = None |
| 252 | self._factor_input_splits: tuple[torch.Tensor, ...] |
| 253 | self._attn_bias: torch.Tensor | None = None |
| 254 | self._rank_frac_col: torch.Tensor |
| 255 | self._is_top1_col: torch.Tensor |
| 256 | |
| 257 | @torch.no_grad() |
| 258 | def materialize_inference_buffers( |
| 259 | self, device: torch.device, dtype: torch.dtype |
| 260 | ) -> None: |
| 261 | """Build inference-only views/copies after checkpoint weights are loaded.""" |
| 262 | topk = self.candidate_topk |
| 263 | self._attn_bias = self._build_attention_bias(device=device, dtype=dtype) |
| 264 | # Both projections consume the same candidate features f: |
| 265 | # out = W_out @ f + b_out |
| 266 | # incoming = W_in @ f + b_in |
| 267 | # Stack weights and biases to compute both in one linear call. score() |
| 268 | # splits the result and normalizes the two vectors separately. |
| 269 | self._fused_edge_weight = ( |
| 270 | torch.cat([self.out_head.weight, self.in_head.weight], dim=0) |
| 271 | .to(device=device, dtype=dtype) |
| 272 | .contiguous() |
| 273 | ) |
| 274 | self._fused_edge_bias = ( |
| 275 | torch.cat([self.out_head.bias, self.in_head.bias], dim=0) |
| 276 | .to(device=device, dtype=dtype) |
| 277 | .contiguous() |
| 278 | ) |
| 279 | # For candidate state x and request anchor a, the trained projection is: |
| 280 | # W @ concat(x, a, x*a) + b |
| 281 | # Split W into input blocks to evaluate the same expression as: |
| 282 | # W_x @ x + W_a @ a + W_cross @ (x*a) + b |
| 283 | # Here x*a is elementwise. This avoids concatenating inputs and reuses |
| 284 | # W_a @ a across candidates, adding the original bias only once. |
| 285 | weight = self.factor_input_proj.weight |
| 286 | hdim = self.hidden_size |
| 287 | self._factor_input_splits = ( |
| 288 | weight[:, :hdim].contiguous(), |
| 289 | weight[:, hdim : 2 * hdim].contiguous(), |
| 290 | weight[:, 2 * hdim :].contiguous(), |
| 291 | ) |
| 292 | if topk > 1: |
| 293 | rank_frac = torch.arange(topk, device=device, dtype=torch.float32).view( |
| 294 | 1, 1, topk |
| 295 | ) / float(topk - 1) |
| 296 | else: |
| 297 | rank_frac = torch.zeros(1, 1, topk, device=device, dtype=torch.float32) |
| 298 | is_top1 = torch.zeros(1, 1, topk, device=device, dtype=torch.float32) |
| 299 | is_top1[..., 0] = 1.0 |
| 300 | self._rank_frac_col = rank_frac.contiguous() |
| 301 | self._is_top1_col = is_top1.contiguous() |
| 302 | |
| 303 | def _build_attention_bias( |
| 304 | self, *, device: torch.device, dtype: torch.dtype |
| 305 | ) -> torch.Tensor: |
| 306 | topk = self.candidate_topk |
| 307 | slot_ids = torch.arange( |
| 308 | self.num_candidate_slots, device=device, dtype=torch.long |
| 309 | ).repeat_interleave(topk) |
| 310 | # Attention flattens [slot, candidate] into one sequence. With topk=2, |
| 311 | # the candidate slot IDs are [0, 0, 1, 1, ...]. Each head learns: |
| 312 | # bias(q, k) = relative_bias[slot(q) - slot(k) + offset] |
| 313 | # + same_slot_bias * (slot(q) == slot(k)) |
| 314 | # The same-slot term includes distinct candidates at the same position. |
| 315 | # This is an additive preference, not a causal mask. Keep the trained |
| 316 | # offset = block_size-1 when score() selects a shorter draft's prefix, |
| 317 | # so each distance continues to use the same learned weight. |
| 318 | rel = slot_ids.view(-1, 1) - slot_ids.view(1, -1) |
| 319 | bias = self.relative_slot_bias[:, rel + self.block_size - 1] |
| 320 | same_slot = slot_ids.view(-1, 1) == slot_ids.view(1, -1) |
| 321 | bias = bias + same_slot.unsqueeze(0).to( |
| 322 | dtype=bias.dtype |
| 323 | ) * self.same_slot_bias.view(-1, 1, 1) |
| 324 | return bias.to(device=device, dtype=dtype).contiguous() |
| 325 | |
| 326 | def score( |
| 327 | self, |
| 328 | *, |
| 329 | token_embeddings: torch.Tensor, |
| 330 | candidate_log_probs: torch.Tensor, |
| 331 | pass_hidden: torch.Tensor, |
| 332 | anchor_hidden: torch.Tensor, |
| 333 | anchor_valid: torch.Tensor, |
| 334 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 335 | if self._attn_bias is None: |
| 336 | raise RuntimeError( |
| 337 | "Call materialize_inference_buffers() after loading the LiLiCorr head." |
| 338 | ) |
| 339 | bsz, n_slots, topk = candidate_log_probs.shape |
| 340 | if not 1 <= n_slots <= self.num_candidate_slots: |
| 341 | raise ValueError( |
| 342 | f"LiLiCorr supports 1..{self.num_candidate_slots} candidate slots, " |
| 343 | f"got {n_slots}." |
| 344 | ) |
| 345 | if topk != self.candidate_topk: |
| 346 | raise ValueError( |
| 347 | f"LiLiCorr expects candidate_topk={self.candidate_topk}, got {topk}." |
| 348 | ) |
| 349 | proj_dtype = self.slot_embedding.dtype |
| 350 | if token_embeddings.dtype != proj_dtype: |
| 351 | token_embeddings = token_embeddings.to(proj_dtype) |
| 352 | if pass_hidden.dtype != proj_dtype: |
| 353 | pass_hidden = pass_hidden.to(proj_dtype) |
| 354 | token_states = self.token_proj(token_embeddings) |
| 355 | pass_states = self.pass_hidden_proj(pass_hidden).unsqueeze(-2) |
| 356 | log_probs = candidate_log_probs.float() |
| 357 | features = torch.stack( |
| 358 | [ |
| 359 | log_probs, |
| 360 | log_probs.exp(), |
| 361 | log_probs - log_probs.max(dim=-1, keepdim=True).values, |
| 362 | self._rank_frac_col.expand_as(log_probs), |
| 363 | self._is_top1_col.expand_as(log_probs), |
| 364 | ], |
| 365 | dim=-1, |
| 366 | ) |
| 367 | hidden_states = token_states + pass_states |
| 368 | hidden_states = hidden_states + self.feature_mlp( |
| 369 | self.feature_norm(features.to(dtype=token_states.dtype)) |
| 370 | ) |
| 371 | hidden_states = hidden_states + self.slot_embedding[:, 0, :n_slots] |
| 372 | hidden_states = hidden_states + self.rank_embedding[:, 0] |
| 373 | hidden_states = hidden_states.reshape(bsz, n_slots * topk, self.hidden_size) |
| 374 | anchor_state = self.context_proj(anchor_hidden) |
| 375 | anchor_state = anchor_state * anchor_valid[:, None].to(anchor_state.dtype) |
| 376 | # Shorter drafts use the learned prefix; keep checkpoint parameter shapes. |
| 377 | lattice = n_slots * topk |
| 378 | attention_bias = self._attn_bias[None, :, :lattice, :lattice] |
| 379 | for layer in self.layers: |
| 380 | hidden_states = layer(hidden_states, attention_bias) |
| 381 | hidden_states = self.output_norm(hidden_states).reshape( |
| 382 | bsz, n_slots, topk, self.hidden_size |
| 383 | ) |
| 384 | anchor_state = self.anchor_norm(anchor_state) |
| 385 | w_self, w_anchor, w_cross = self._factor_input_splits |
| 386 | anchor_row = anchor_state[:, None, None, :] |
| 387 | pre = F.linear(hidden_states, w_self, self.factor_input_proj.bias) |
| 388 | pre = pre + F.linear(anchor_row, w_anchor) |
| 389 | pre = pre + F.linear(hidden_states * anchor_row, w_cross) |
| 390 | factor_hidden = F.silu(pre) |
| 391 | edges = F.linear(factor_hidden, self._fused_edge_weight, self._fused_edge_bias) |
| 392 | out_vec, in_vec = F.normalize( |
| 393 | edges.unflatten(-1, (2, self.factor_dim)), dim=-1, eps=self.vector_eps |
| 394 | ).unbind(-2) |
| 395 | anchor_out = F.normalize( |
| 396 | self.anchor_out_head(anchor_state), dim=-1, eps=self.vector_eps |
| 397 | ) |
| 398 | # start_scores[b, j] scores anchor -> candidate j in slot zero. |
| 399 | # pair_scores[b, s, i, j] scores candidate i in slot s -> candidate j |
| 400 | # in slot s+1. Normalized factor dot products encode learned compatibility; |
| 401 | # forward() scales them into logits for the candidate-path selector. |
| 402 | start_scores = (anchor_out[:, None, :] * in_vec[:, 0, :, :]).sum(dim=-1) |
| 403 | pair_scores = torch.matmul(out_vec[:, :-1], in_vec[:, 1:].transpose(-1, -2)) |
| 404 | return (start_scores, pair_scores) |
| 405 | |
| 406 | def forward( |
| 407 | self, |
| 408 | token_embeddings: torch.Tensor, |
| 409 | candidate_log_probs: torch.Tensor, |
| 410 | pass_hidden: torch.Tensor, |
| 411 | anchor_hidden: torch.Tensor, |
| 412 | anchor_valid: torch.Tensor, |
| 413 | ) -> torch.Tensor: |
| 414 | start, pairs = self.score( |
| 415 | token_embeddings=token_embeddings, |
| 416 | candidate_log_probs=candidate_log_probs, |
| 417 | pass_hidden=pass_hidden, |
| 418 | anchor_hidden=anchor_hidden, |
| 419 | anchor_valid=anchor_valid, |
| 420 | ) |
| 421 | start = self.logit_scale * start.float() |
| 422 | pairs = self.logit_scale * pairs.float() |
| 423 | # The shared walk starts at predecessor index zero; all first rows agree. |
| 424 | first = start[:, None, None, :].expand(-1, 1, self.candidate_topk, -1) |
| 425 | return torch.cat((first, pairs), dim=1) |
| 426 | |
| 427 | |
| 428 | class LiLiCorr(DFlashQwen3Model): |
| 429 | def __init__( |
| 430 | self, *, vllm_config: VllmConfig, start_layer_id: int = 0, prefix: str = "" |
| 431 | ): |
| 432 | spec = vllm_config.speculative_config |
| 433 | assert spec is not None |
| 434 | config = spec.draft_model_config.hf_config |
| 435 | draft_config = config.dflash_config |
| 436 | if not draft_config.get("lilicorr_enabled", True): |
| 437 | raise ValueError("LiLiCorr requires lilicorr_enabled=True.") |
| 438 | block_size = draft_config.get("block_size", getattr(config, "block_size", 0)) |
| 439 | if not 1 <= spec.num_speculative_tokens < block_size: |
| 440 | raise ValueError( |
| 441 | "LiLiCorr requires 1 <= num_speculative_tokens < trained block_size." |
| 442 | ) |
| 443 | taps = draft_config.get("conv_kernel_size", 0) |
| 444 | groups = draft_config.get("conv_group_size", 0) |
| 445 | if taps < 0 or groups < 0 or bool(taps) != bool(groups): |
| 446 | raise ValueError( |
| 447 | "LiLiCorr convolution geometry must enable both taps and group size." |
| 448 | ) |
| 449 | if taps: |
| 450 | self.decoder_layer_cls = DFlash2Qwen3DecoderLayer |
| 451 | super().__init__( |
| 452 | vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix |
| 453 | ) |
| 454 | with set_model_tag("lilicorr_head"): |
| 455 | self.lilicorr = LiLiCorrHead( |
| 456 | model_hidden_size=config.hidden_size, |
| 457 | block_size=block_size, |
| 458 | rms_norm_eps=config.rms_norm_eps, |
| 459 | config=draft_config, |
| 460 | quant_config=self.quant_config, |
| 461 | prefix=maybe_prefix(prefix, "lilicorr"), |
| 462 | ) |
| 463 | |
| 464 | |
| 465 | class LiLiCorrForCausalLM(DFlashQwen3ForCausalLM): |
| 466 | model_cls = LiLiCorr |
| 467 | |
| 468 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 469 | spec = vllm_config.speculative_config |
| 470 | assert spec is not None |
| 471 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 472 | if self.draft_id_to_target_id is not None: |
| 473 | raise ValueError("LiLiCorr candidates require the full target vocabulary.") |
| 474 | self.has_own_lm_head = bool(getattr(self.config, "has_own_lm_head", False)) |
| 475 | # The correlator consumes raw candidate-head log-softmax features. |
| 476 | self.candidate_logits_processor = LogitsProcessor(self.config.vocab_size) |
| 477 | |
| 478 | def compute_candidates( |
| 479 | self, hidden_states: torch.Tensor |
| 480 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 481 | return self.candidate_logits_processor.get_top_k_tokens( |
| 482 | self.lm_head, |
| 483 | hidden_states, |
| 484 | self.model.lilicorr.candidate_topk, |
| 485 | return_log_probs=True, |
| 486 | ) |
| 487 | |
| 488 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): |
| 489 | quantized_metadata: set[str] = set() |
| 490 | for name, module in self.named_modules(): |
| 491 | name = name.removeprefix("model.") |
| 492 | if not isinstance(module, (LinearBase, ParallelLMHead)) or isinstance( |
| 493 | module.quant_method, UnquantizedLinearMethod |
| 494 | ): |
| 495 | continue |
| 496 | required = {"weight", "bias"} |
| 497 | method = module.quant_method |
| 498 | if ( |
| 499 | isinstance(method, ModelOptLinearMethod) |
| 500 | and method.spec.weight == kNvfp4Static |
| 501 | ): |
| 502 | # NVFP4 scales are trained checkpoint data. Only W4A16's |
| 503 | # deprecated, unused input_scale may be absent. |
| 504 | required.update(("weight_scale", "weight_scale_2")) |
| 505 | if method.spec.activation is not None: |
| 506 | required.add("input_scale") |
| 507 | quantized_metadata.update( |
| 508 | f"{name}.{parameter}" |
| 509 | for parameter, _ in module.named_parameters(recurse=False) |
| 510 | if parameter not in required |
| 511 | ) |
| 512 | expected = { |
| 513 | name.removeprefix("model.") |
| 514 | for name, _ in self.named_parameters() |
| 515 | if ( |
| 516 | name.startswith("model.lilicorr.") |
| 517 | or ".attention_conv." in name |
| 518 | or ".mlp_conv." in name |
| 519 | or (self.has_own_lm_head and name.startswith("lm_head.")) |
| 520 | ) |
| 521 | and name.removeprefix("model.") not in quantized_metadata |
| 522 | } |
| 523 | seen: set[str] = set() |
| 524 | head_weights: list[tuple[str, torch.Tensor]] = [] |
| 525 | |
| 526 | def normalized_weights(): |
| 527 | for name, value in weights: |
| 528 | name = name.removeprefix("model.") |
| 529 | seen.add(name) |
| 530 | # The backbone stacks up_proj into gate_up_proj. Load the |
| 531 | # nongated correlator separately to avoid that transformation. |
| 532 | if name.startswith("lilicorr."): |
| 533 | head_weights.append((name.removeprefix("lilicorr."), value)) |
| 534 | else: |
| 535 | yield name, value |
| 536 | |
| 537 | super().load_weights(normalized_weights()) |
| 538 | AutoWeightsLoader(self.model.lilicorr).load_weights(head_weights) |
| 539 | missing = expected - seen |
| 540 | if missing: |
| 541 | raise ValueError( |
| 542 | f"LiLiCorr checkpoint coverage mismatch: missing={sorted(missing)}" |
| 543 | ) |
| 544 | head = self.model.lilicorr |
| 545 | parameter = head.slot_embedding |
| 546 | head.materialize_inference_buffers(parameter.device, parameter.dtype) |