| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | from typing import Any |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | from torch import nn |
| 11 | |
| 12 | from vllm.config import ( |
| 13 | CompilationConfig, |
| 14 | CompilationMode, |
| 15 | DeviceConfig, |
| 16 | LoadConfig, |
| 17 | VllmConfig, |
| 18 | set_current_vllm_config, |
| 19 | ) |
| 20 | from vllm.distributed import parallel_state |
| 21 | from vllm.model_executor.layers import logits_processor |
| 22 | from vllm.model_executor.layers.linear import LinearBase |
| 23 | from vllm.model_executor.models.lilicorr import LiLiCorrHead |
| 24 | from vllm.transformers_utils.configs.eagle import EAGLEConfig |
| 25 | from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import CandidateSampler |
| 26 | from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import LiLiCorrSpeculator |
| 27 | |
| 28 | |
| 29 | @pytest.fixture(autouse=True) |
| 30 | def head_tp_metadata(monkeypatch): |
| 31 | # Single-rank unit tests only need TP metadata. |
| 32 | monkeypatch.setattr( |
| 33 | parallel_state, "_TP", SimpleNamespace(rank_in_group=0, world_size=1) |
| 34 | ) |
| 35 | |
| 36 | |
| 37 | def _config(**overrides): |
| 38 | config = dict( |
| 39 | candidate_topk=4, |
| 40 | hidden_size=8, |
| 41 | num_layers=2, |
| 42 | num_heads=2, |
| 43 | mlp_ratio=2.0, |
| 44 | factor_dim=4, |
| 45 | vector_eps=1e-6, |
| 46 | logit_scale=3.0, |
| 47 | ) |
| 48 | config.update(overrides) |
| 49 | return {f"lilicorr_{key}": value for key, value in config.items()} |
| 50 | |
| 51 | |
| 52 | def _reference_scores(head, embeddings, log_probs, hidden, anchor, valid): |
| 53 | """Unfused exported head: MHA attention, concatenated factors, separate edges.""" |
| 54 | batch, slots, top_k = log_probs.shape |
| 55 | rank = torch.arange(top_k, device=log_probs.device) / max(top_k - 1, 1) |
| 56 | top1 = (torch.arange(top_k, device=log_probs.device) == 0).float() |
| 57 | features = torch.stack( |
| 58 | ( |
| 59 | log_probs, |
| 60 | log_probs.exp(), |
| 61 | log_probs - log_probs.amax(-1, keepdim=True), |
| 62 | rank.expand_as(log_probs), |
| 63 | top1.expand_as(log_probs), |
| 64 | ), |
| 65 | -1, |
| 66 | ) |
| 67 | x = head.token_proj(embeddings) + head.pass_hidden_proj(hidden).unsqueeze(-2) |
| 68 | features = head.feature_norm(features.to(x.dtype)) |
| 69 | x = x + head.feature_mlp.down_proj(F.silu(head.feature_mlp.up_proj(features)[0]))[0] |
| 70 | x = x + head.slot_embedding[:, 0, :slots] + head.rank_embedding[:, 0] |
| 71 | x = x.flatten(1, 2) |
| 72 | positions = torch.arange(slots, device=x.device).repeat_interleave(top_k) |
| 73 | relative = positions[:, None] - positions[None, :] |
| 74 | bias = head.relative_slot_bias[:, relative + head.block_size - 1] |
| 75 | bias = bias + (relative == 0) * head.same_slot_bias[:, None, None] |
| 76 | bias = bias.repeat(batch, 1, 1) |
| 77 | for layer in head.layers: |
| 78 | attention = nn.MultiheadAttention( |
| 79 | head.hidden_size, |
| 80 | head.num_heads, |
| 81 | batch_first=True, |
| 82 | dtype=x.dtype, |
| 83 | device=x.device, |
| 84 | ) |
| 85 | attention.load_state_dict(layer.attn.state_dict()) |
| 86 | normalized = layer.attn_norm(x) |
| 87 | x = ( |
| 88 | x |
| 89 | + attention( |
| 90 | normalized, normalized, normalized, attn_mask=bias, need_weights=False |
| 91 | )[0] |
| 92 | ) |
| 93 | x = x + layer.mlp.down_proj(F.silu(layer.mlp.up_proj(layer.mlp_norm(x))[0]))[0] |
| 94 | x = head.output_norm(x).view(batch, slots, top_k, -1) |
| 95 | a = head.anchor_norm(head.context_proj(anchor) * valid[:, None]) |
| 96 | expanded = a[:, None, None].expand_as(x) |
| 97 | factors = F.silu(head.factor_input_proj(torch.cat((x, expanded, x * expanded), -1))) |
| 98 | incoming = F.normalize(head.in_head(factors), dim=-1, eps=head.vector_eps) |
| 99 | outgoing = F.normalize(head.out_head(factors), dim=-1, eps=head.vector_eps) |
| 100 | a = F.normalize(head.anchor_out_head(a), dim=-1, eps=head.vector_eps) |
| 101 | start = (a[:, None] * incoming[:, 0]).sum(-1) |
| 102 | pairs = outgoing[:, :-1] @ incoming[:, 1:].transpose(-1, -2) |
| 103 | return ( |
| 104 | head.logit_scale |
| 105 | * torch.cat((start[:, None, None].expand(-1, 1, top_k, -1), pairs), 1).float() |
| 106 | ) |
| 107 | |
| 108 | |
| 109 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm") |
| 110 | @pytest.mark.parametrize( |
| 111 | "head_width,slots,dtype", |
| 112 | [ |
| 113 | pytest.param(8, 3, torch.float32, id="projected-multi-slot"), |
| 114 | pytest.param(16, 1, torch.float32, id="identity-shortened-single-slot"), |
| 115 | pytest.param(8, 2, torch.bfloat16, id="bf16-shortened-multi-slot"), |
| 116 | ], |
| 117 | ) |
| 118 | def test_lilicorr_matches_exported_head(head_width, slots, dtype): |
| 119 | device = torch.device("cuda") |
| 120 | torch.manual_seed(7) |
| 121 | with ( |
| 122 | torch.device(device), |
| 123 | set_current_vllm_config( |
| 124 | VllmConfig( |
| 125 | device_config=DeviceConfig("cuda"), |
| 126 | compilation_config=CompilationConfig(mode=CompilationMode.NONE), |
| 127 | ) |
| 128 | ), |
| 129 | ): |
| 130 | head = LiLiCorrHead( |
| 131 | model_hidden_size=16, |
| 132 | block_size=4, |
| 133 | rms_norm_eps=1e-6, |
| 134 | config=_config(hidden_size=head_width), |
| 135 | ).to(dtype) |
| 136 | # vLLM linear parameters are initialized by the checkpoint loader. |
| 137 | with torch.no_grad(), torch.device(device): |
| 138 | for module in head.modules(): |
| 139 | if isinstance(module, LinearBase): |
| 140 | for parameter in module.parameters(): |
| 141 | parameter.weight_loader( |
| 142 | parameter, torch.randn_like(parameter) * 0.1 |
| 143 | ) |
| 144 | for layer in head.layers: |
| 145 | layer.attn.in_proj_weight.normal_(std=0.1) |
| 146 | layer.attn.in_proj_bias.normal_(std=0.1) |
| 147 | head.relative_slot_bias.normal_(std=0.1) |
| 148 | head.same_slot_bias.normal_(std=0.1) |
| 149 | head.slot_embedding.normal_(std=0.1) |
| 150 | head.rank_embedding.normal_(std=0.1) |
| 151 | head.materialize_inference_buffers(device, dtype) |
| 152 | inputs = ( |
| 153 | torch.randn(2, slots, 4, 16, dtype=dtype), |
| 154 | torch.randn(2, slots, 23).log_softmax(-1).topk(4).values, |
| 155 | torch.randn(2, slots, 16, dtype=dtype), |
| 156 | torch.randn(2, 16, dtype=dtype), |
| 157 | torch.tensor([True, False]), |
| 158 | ) |
| 159 | actual = head(*inputs) |
| 160 | expected = _reference_scores(head, *inputs) |
| 161 | tolerance = 0.04 if dtype == torch.bfloat16 else 2e-6 |
| 162 | torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) |
| 163 | |
| 164 | |
| 165 | def test_runtime_length_keeps_trained_checkpoint_geometry(monkeypatch): |
| 166 | from vllm.model_executor.models.lilicorr import LiLiCorr |
| 167 | from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model |
| 168 | |
| 169 | def init_backbone(self, **kwargs): |
| 170 | nn.Module.__init__(self) |
| 171 | self.quant_config = None |
| 172 | |
| 173 | monkeypatch.setattr(DFlashQwen3Model, "__init__", init_backbone) |
| 174 | spec = SimpleNamespace( |
| 175 | num_speculative_tokens=1, |
| 176 | draft_model_config=SimpleNamespace( |
| 177 | hf_config=SimpleNamespace( |
| 178 | hidden_size=16, |
| 179 | rms_norm_eps=1e-6, |
| 180 | dflash_config=_config() | {"block_size": 4}, |
| 181 | ) |
| 182 | ), |
| 183 | ) |
| 184 | with set_current_vllm_config( |
| 185 | VllmConfig( |
| 186 | device_config=DeviceConfig("cpu"), |
| 187 | compilation_config=CompilationConfig(mode=CompilationMode.NONE), |
| 188 | ) |
| 189 | ): |
| 190 | for slots in (1, 3): |
| 191 | spec.num_speculative_tokens = slots |
| 192 | head = LiLiCorr( |
| 193 | vllm_config=SimpleNamespace(speculative_config=spec) |
| 194 | ).lilicorr |
| 195 | assert head.slot_embedding.shape == (1, 1, 3, 1, 8) |
| 196 | assert head.relative_slot_bias.shape == (2, 7) |
| 197 | spec.draft_model_config.hf_config.dflash_config["lilicorr_enabled"] = False |
| 198 | with pytest.raises(ValueError, match="lilicorr_enabled"): |
| 199 | LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec)) |
| 200 | spec.draft_model_config.hf_config.dflash_config["lilicorr_enabled"] = True |
| 201 | spec.num_speculative_tokens = 4 |
| 202 | with pytest.raises(ValueError, match="trained block_size"): |
| 203 | LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec)) |
| 204 | |
| 205 | |
| 206 | @pytest.mark.parametrize("tp_size", [1, 2]) |
| 207 | def test_candidates_use_global_partition_and_exclude_padding(monkeypatch, tp_size): |
| 208 | """Candidate features must include mass outside top-k and on other TP ranks.""" |
| 209 | |
| 210 | class Gather: |
| 211 | def __init__(self, calls): |
| 212 | self.calls = calls |
| 213 | self.index = 0 |
| 214 | |
| 215 | def __call__(self, value, dim=-1): |
| 216 | expected_input, result = self.calls[self.index] |
| 217 | assert dim == -1 |
| 218 | torch.testing.assert_close(value, expected_input) |
| 219 | self.index += 1 |
| 220 | return result |
| 221 | |
| 222 | shards = [ |
| 223 | torch.tensor([[1.0, 3.0, 2.0, 100.0], [4.0, 1.0, 2.0, 100.0]]), |
| 224 | torch.tensor([[6.0, 5.0, 0.0, 100.0], [3.0, 6.0, 5.0, 100.0]]), |
| 225 | ][:tp_size] |
| 226 | for rank in range(tp_size): |
| 227 | fake = SimpleNamespace( |
| 228 | scale=1.0, |
| 229 | soft_cap=None, |
| 230 | _apply_head=lambda *args, rank=rank: shards[rank].clone(), |
| 231 | ) |
| 232 | lm_head = SimpleNamespace( |
| 233 | tp_size=tp_size, |
| 234 | shard_indices=SimpleNamespace( |
| 235 | num_org_vocab_padding=1, org_vocab_start_index=rank * 3 |
| 236 | ), |
| 237 | ) |
| 238 | valid_shards = [shard[:, :3] for shard in shards] |
| 239 | local_topk = [shard.topk(2) for shard in valid_shards] |
| 240 | gather_calls = [ |
| 241 | ( |
| 242 | valid_shards[rank].logsumexp(-1, keepdim=True), |
| 243 | torch.cat( |
| 244 | [shard.logsumexp(-1, keepdim=True) for shard in valid_shards], |
| 245 | -1, |
| 246 | ), |
| 247 | ), |
| 248 | ( |
| 249 | local_topk[rank].values, |
| 250 | torch.cat([topk.values for topk in local_topk], -1), |
| 251 | ), |
| 252 | ( |
| 253 | local_topk[rank].indices + rank * 3, |
| 254 | torch.cat( |
| 255 | [ |
| 256 | topk.indices + shard_rank * 3 |
| 257 | for shard_rank, topk in enumerate(local_topk) |
| 258 | ], |
| 259 | -1, |
| 260 | ), |
| 261 | ), |
| 262 | ] |
| 263 | all_gather = Gather(gather_calls) |
| 264 | |
| 265 | monkeypatch.setattr( |
| 266 | logits_processor, |
| 267 | "tensor_model_parallel_all_gather", |
| 268 | all_gather, |
| 269 | ) |
| 270 | monkeypatch.setattr(logits_processor, "_topk", lambda x, k: x.topk(k, dim=-1)) |
| 271 | ids, values = logits_processor.LogitsProcessor.get_top_k_tokens( |
| 272 | fake, |
| 273 | lm_head, |
| 274 | torch.empty(2, 1), |
| 275 | 2, |
| 276 | return_log_probs=True, |
| 277 | ) |
| 278 | expected = torch.cat(valid_shards, -1).log_softmax(-1).topk(2) |
| 279 | torch.testing.assert_close(ids, expected.indices) |
| 280 | torch.testing.assert_close(values, expected.values) |
| 281 | assert all_gather.index == (3 if tp_size > 1 else 0) |
| 282 | |
| 283 | |
| 284 | def test_context_anchor_is_last_committed_normalized_feature(): |
| 285 | hidden = torch.arange(1, 25, dtype=torch.float32).view(6, 4) |
| 286 | norm = nn.RMSNorm(4) |
| 287 | spec = SimpleNamespace( |
| 288 | hidden_states=hidden, |
| 289 | model=SimpleNamespace(model=SimpleNamespace(hidden_norm=norm)), |
| 290 | anchor_hidden=torch.full((4, 4), 123.0), |
| 291 | anchor_valid=torch.ones(4, dtype=torch.bool), |
| 292 | ) |
| 293 | batch = SimpleNamespace(num_reqs=2, query_start_loc=torch.tensor([0, 2, 6])) |
| 294 | LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([0, 2])) |
| 295 | torch.testing.assert_close(spec.anchor_hidden[:2], norm(hidden[[1, 3]])) |
| 296 | assert spec.anchor_valid.tolist() == [True, True, False, False] |
| 297 | assert not spec.anchor_hidden[2:].any() |
| 298 | # A smaller reordered batch must not retain an old request's anchor. |
| 299 | batch = SimpleNamespace(num_reqs=1, query_start_loc=torch.tensor([0, 3])) |
| 300 | LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([1])) |
| 301 | torch.testing.assert_close(spec.anchor_hidden[0], norm(hidden[1])) |
| 302 | assert not spec.anchor_hidden[1:].any() |
| 303 | |
| 304 | |
| 305 | @pytest.mark.parametrize("lilicorr", [False, True]) |
| 306 | def test_candidate_generation_routes_scores_and_adaptive_inputs(monkeypatch, lilicorr): |
| 307 | """Both heads share the walk, including a padded request and adaptive scores.""" |
| 308 | from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator |
| 309 | from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator |
| 310 | |
| 311 | key = "lilicorr_candidate_topk" if lilicorr else "selector_top_k" |
| 312 | |
| 313 | def init_base(self, config, device): |
| 314 | self.draft_model_config = SimpleNamespace( |
| 315 | hf_config=SimpleNamespace(hidden_size=4, dflash_config={key: 4}) |
| 316 | ) |
| 317 | self.max_num_reqs = 3 |
| 318 | self.num_speculative_steps = 2 |
| 319 | self.num_query_per_req = 3 |
| 320 | self.dtype = torch.float32 |
| 321 | |
| 322 | monkeypatch.setattr(DFlashSpeculator, "__init__", init_base) |
| 323 | cls = LiLiCorrSpeculator if lilicorr else DFlash2Speculator |
| 324 | spec = cls(None, torch.device("cpu")) |
| 325 | hidden = torch.arange(36).float().view(9, 4) |
| 326 | spec._run_model = lambda *args: hidden |
| 327 | spec.sample_indices = torch.tensor([1, 2, 4, 5, 7, 8]) |
| 328 | spec.sample_pos = torch.arange(6) |
| 329 | spec.sample_idx_mapping = torch.tensor([1, 1, 0, 0, -1, -1]) |
| 330 | spec.sample_col = torch.tensor([0, 1, 0, 1, 0, 1]) |
| 331 | spec.temperature = torch.ones(3) |
| 332 | spec.seeds = torch.arange(3) |
| 333 | spec.draft_tokens = torch.empty(3, 2, dtype=torch.long) |
| 334 | spec.draft_logits = None |
| 335 | spec.use_fp64_gumbel = False |
| 336 | spec.enable_adaptive_verification = True |
| 337 | spec.input_buffers = SimpleNamespace(input_ids=torch.arange(9)) |
| 338 | candidates = torch.arange(24).view(6, 4) |
| 339 | log_probs = candidates.float().log_softmax(-1) |
| 340 | scores = torch.randn(3, 2, 4, 4) |
| 341 | scoring_inputs: list[torch.Tensor] = [] |
| 342 | sampling_inputs: list[Any] = [] |
| 343 | adaptive_inputs: list[torch.Tensor] = [] |
| 344 | |
| 345 | def score(*args): |
| 346 | scoring_inputs.extend(args) |
| 347 | return scores |
| 348 | |
| 349 | def sample(*args): |
| 350 | sampling_inputs.extend(args) |
| 351 | spec.candidate_sampler.scores.copy_(scores[:, :, 0]) |
| 352 | |
| 353 | spec.model = SimpleNamespace( |
| 354 | compute_candidates=lambda h: (candidates, log_probs), |
| 355 | model=SimpleNamespace(lilicorr=score, candidate_selector=score), |
| 356 | ) |
| 357 | if lilicorr: |
| 358 | assert isinstance(spec, LiLiCorrSpeculator) |
| 359 | monkeypatch.setattr( |
| 360 | spec, |
| 361 | "target_embeddings", |
| 362 | lambda ids: ids.float().unsqueeze(-1).expand(*ids.shape, 4), |
| 363 | raising=False, |
| 364 | ) |
| 365 | spec.anchor_hidden.copy_(torch.arange(12).view(3, 4)) |
| 366 | spec.anchor_valid[:2] = True |
| 367 | monkeypatch.setattr(spec.candidate_sampler, "sample", sample) |
| 368 | monkeypatch.setattr( |
| 369 | spec, "_maybe_predict_acceptance", lambda *args: adaptive_inputs.extend(args) |
| 370 | ) |
| 371 | spec._generate_draft(3, 9, None, None, None) |
| 372 | |
| 373 | expected_ids = candidates.view(3, 2, 4) |
| 374 | expected_first = ( |
| 375 | spec.target_embeddings(expected_ids) |
| 376 | if isinstance(spec, LiLiCorrSpeculator) |
| 377 | else expected_ids |
| 378 | ) |
| 379 | torch.testing.assert_close(scoring_inputs[0], expected_first) |
| 380 | torch.testing.assert_close(scoring_inputs[1], log_probs.view(3, 2, 4)) |
| 381 | torch.testing.assert_close( |
| 382 | scoring_inputs[2], hidden[spec.sample_indices].view(3, 2, 4) |
| 383 | ) |
| 384 | if lilicorr: |
| 385 | assert isinstance(spec, LiLiCorrSpeculator) |
| 386 | torch.testing.assert_close(scoring_inputs[3], spec.anchor_hidden) |
| 387 | torch.testing.assert_close(scoring_inputs[4], spec.anchor_valid) |
| 388 | else: |
| 389 | torch.testing.assert_close(scoring_inputs[3], torch.tensor([0, 3, 6])) |
| 390 | torch.testing.assert_close(sampling_inputs[0], expected_ids) |
| 391 | assert sampling_inputs[1] is scores |
| 392 | assert sampling_inputs[4] is spec.sample_idx_mapping |
| 393 | torch.testing.assert_close(adaptive_inputs[0], scores[:, :, 0].flatten(0, 1)) |
| 394 | torch.testing.assert_close(adaptive_inputs[1], spec.sample_idx_mapping) |
| 395 | torch.testing.assert_close(adaptive_inputs[2], spec.sample_col) |
| 396 | |
| 397 | |
| 398 | def test_lilicorr_config_preserves_architecture(): |
| 399 | from transformers import Qwen3Config |
| 400 | |
| 401 | config = Qwen3Config(architectures=["LiLiCorrDraftModel"]) |
| 402 | wrapped = EAGLEConfig(config, method="dflash") |
| 403 | assert wrapped.architectures == ["LiLiCorrDraftModel"] |
| 404 | |
| 405 | |
| 406 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 407 | @pytest.mark.parametrize("probabilistic", [False, True]) |
| 408 | def test_candidate_walk_preserves_conditional_scores_and_padding(probabilistic): |
| 409 | """Check the shared sampler's actual q, stale support, and graph padding.""" |
| 410 | device = torch.device("cuda") |
| 411 | batch, steps, top_k, vocab = 3, 3, 4, 32 |
| 412 | sampler = CandidateSampler(batch, steps, top_k, device) |
| 413 | candidates = ( |
| 414 | torch.arange(steps * top_k, device=device) |
| 415 | .view(1, steps, top_k) |
| 416 | .repeat(batch, 1, 1) |
| 417 | ) |
| 418 | scores = torch.randn(batch, steps, top_k, top_k, device=device) |
| 419 | tokens = torch.full((batch, steps), -1, dtype=torch.long, device=device) |
| 420 | cache = ( |
| 421 | torch.full((batch, steps, vocab), -float("inf"), device=device) |
| 422 | if probabilistic |
| 423 | else None |
| 424 | ) |
| 425 | positions = torch.arange(batch * steps * 2, device=device)[::2] |
| 426 | mappings = torch.tensor( |
| 427 | [2] * steps + [0] * steps + [-1] * steps, device=device |
| 428 | ).repeat_interleave(2)[::2] |
| 429 | temperature = torch.tensor([0.0, 1.0, 0.8], device=device) |
| 430 | seeds = torch.tensor([11, 22, 33], device=device) |
| 431 | |
| 432 | def run(): |
| 433 | sampler.sample( |
| 434 | candidates, |
| 435 | scores, |
| 436 | batch, |
| 437 | positions, |
| 438 | mappings, |
| 439 | temperature, |
| 440 | seeds, |
| 441 | tokens, |
| 442 | cache, |
| 443 | False, |
| 444 | ) |
| 445 | |
| 446 | run() |
| 447 | graph = torch.cuda.CUDAGraph() |
| 448 | with torch.cuda.graph(graph): |
| 449 | run() |
| 450 | old_candidates = candidates.clone() |
| 451 | candidates.add_(12) |
| 452 | graph.replay() |
| 453 | for row, state in ((0, 2), (1, 0)): |
| 454 | previous = 0 |
| 455 | for step in range(steps): |
| 456 | realized = scores[row, step, previous] |
| 457 | torch.testing.assert_close(sampler.scores[row, step], realized) |
| 458 | index = (candidates[row, step] == tokens[row, step]).nonzero().item() |
| 459 | if not probabilistic or temperature[state] == 0: |
| 460 | assert index == realized.argmax().item() |
| 461 | if cache is not None: |
| 462 | torch.testing.assert_close( |
| 463 | cache[state, step, candidates[row, step]], realized |
| 464 | ) |
| 465 | assert cache[state, step, old_candidates[row, step]].isneginf().all() |
| 466 | previous = index |
| 467 | assert (tokens[2] == -1).all() |
| 468 | if cache is not None: |
| 469 | assert cache[1].isneginf().all() |
| 470 | |
| 471 | |
| 472 | @pytest.mark.parametrize( |
| 473 | ("convolution", "mismatch", "quantized"), |
| 474 | [ |
| 475 | pytest.param(False, None, False, id="valid-plain"), |
| 476 | pytest.param(True, None, False, id="valid-convolution"), |
| 477 | pytest.param(False, "missing_head", False, id="missing-head"), |
| 478 | pytest.param(False, "extra_head", False, id="extra-head"), |
| 479 | pytest.param(True, "missing_conv", False, id="missing-convolution"), |
| 480 | pytest.param(False, "extra_conv", False, id="unexpected-convolution"), |
| 481 | pytest.param(False, None, True, id="optional-w4a16-input-scale"), |
| 482 | pytest.param(False, "missing_weight", True, id="missing-quantized-weight"), |
| 483 | pytest.param(False, "missing_bias", True, id="missing-quantized-bias"), |
| 484 | pytest.param(False, "missing_weight_scale", True, id="missing-block-scale"), |
| 485 | pytest.param(False, "missing_weight_scale_2", True, id="missing-global-scale"), |
| 486 | ], |
| 487 | ) |
| 488 | def test_checkpoint_coverage_rejects_incomplete_or_wrong_heads( |
| 489 | monkeypatch, convolution, mismatch, quantized |
| 490 | ): |
| 491 | from vllm.model_executor.models.lilicorr import LiLiCorr, LiLiCorrForCausalLM |
| 492 | |
| 493 | wrapper = LiLiCorrForCausalLM.__new__(LiLiCorrForCausalLM) # type: ignore[type-abstract] |
| 494 | nn.Module.__init__(wrapper) |
| 495 | wrapper.model = LiLiCorr.__new__(LiLiCorr) |
| 496 | nn.Module.__init__(wrapper.model) |
| 497 | wrapper.has_own_lm_head = False |
| 498 | with set_current_vllm_config( |
| 499 | VllmConfig( |
| 500 | device_config=DeviceConfig("cpu"), |
| 501 | compilation_config=CompilationConfig(mode=CompilationMode.NONE), |
| 502 | ) |
| 503 | ): |
| 504 | wrapper.model.lilicorr = LiLiCorrHead( |
| 505 | model_hidden_size=16, block_size=4, rms_norm_eps=1e-6, config=_config() |
| 506 | ) |
| 507 | layer = nn.Module() |
| 508 | if convolution: |
| 509 | layer.attention_conv = nn.Linear(2, 2, bias=False) |
| 510 | wrapper.model.layers = nn.ModuleList([layer]) |
| 511 | if quantized: |
| 512 | from vllm.config.quantization import QuantSpec |
| 513 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 514 | CkptCtx, |
| 515 | ModelOptLinearMethod, |
| 516 | ) |
| 517 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 518 | kNvfp4Static, |
| 519 | ) |
| 520 | |
| 521 | # Match the W4A16 NVFP4 checkpoint contract without GPU kernel setup. |
| 522 | projection = wrapper.model.lilicorr.pass_hidden_proj |
| 523 | projection.quant_method = ModelOptLinearMethod( |
| 524 | QuantSpec(weight=kNvfp4Static, activation=None), CkptCtx(group_size=16) |
| 525 | ) |
| 526 | projection.weight = nn.Parameter( |
| 527 | torch.zeros(8, 8, dtype=torch.uint8), requires_grad=False |
| 528 | ) |
| 529 | for name, value in ( |
| 530 | ("weight_scale", torch.ones(8, 1).to(torch.float8_e4m3fn)), |
| 531 | ("weight_scale_2", torch.ones(1)), |
| 532 | ("input_scale", torch.full((1,), torch.nan)), |
| 533 | ): |
| 534 | projection.register_parameter( |
| 535 | name, nn.Parameter(value, requires_grad=False) |
| 536 | ) |
| 537 | weights = dict(wrapper.model.state_dict()) |
| 538 | if quantized: |
| 539 | del weights["lilicorr.pass_hidden_proj.input_scale"] |
| 540 | if mismatch == "missing_head": |
| 541 | del weights["lilicorr.slot_embedding"] |
| 542 | elif mismatch in ( |
| 543 | "missing_weight", |
| 544 | "missing_bias", |
| 545 | "missing_weight_scale", |
| 546 | "missing_weight_scale_2", |
| 547 | ): |
| 548 | del weights[f"lilicorr.pass_hidden_proj.{mismatch.removeprefix('missing_')}"] |
| 549 | elif mismatch == "extra_head": |
| 550 | weights["lilicorr.untrained.weight"] = torch.zeros(1) |
| 551 | elif mismatch == "missing_conv": |
| 552 | del weights["layers.0.attention_conv.weight"] |
| 553 | elif mismatch == "extra_conv": |
| 554 | weights["layers.0.mlp_conv.untrained"] = torch.zeros(1) |
| 555 | |
| 556 | # Keep production loading and mapping, bypass only backbone GPU KV setup. |
| 557 | wrapper.model.use_aux_hidden_state = True |
| 558 | wrapper.model.has_separate_mask_embedding = False |
| 559 | monkeypatch.setattr(wrapper, "_read_mask_embedding", lambda: None) |
| 560 | monkeypatch.setattr( |
| 561 | wrapper.model, "_build_fused_kv_buffers", lambda: None, raising=False |
| 562 | ) |
| 563 | supplied = [("model." + name, value) for name, value in weights.items()] |
| 564 | if mismatch: |
| 565 | message = ( |
| 566 | "no module or parameter" |
| 567 | if mismatch.startswith("extra") |
| 568 | else "coverage mismatch" |
| 569 | ) |
| 570 | with pytest.raises(ValueError, match=message): |
| 571 | wrapper.load_weights(supplied) |
| 572 | else: |
| 573 | wrapper.load_weights(supplied) |
| 574 | assert wrapper.model.lilicorr._attn_bias is not None |
| 575 | |
| 576 | |
| 577 | @pytest.fixture |
| 578 | def owned_head_model(monkeypatch): |
| 579 | from vllm.model_executor.layers.quantization import modelopt |
| 580 | from vllm.model_executor.models import qwen3_dflash |
| 581 | from vllm.model_executor.models.lilicorr import LiLiCorrForCausalLM |
| 582 | |
| 583 | quant_config = modelopt.ModelOptNvFp4Config( |
| 584 | quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True |
| 585 | ) |
| 586 | monkeypatch.setattr( |
| 587 | modelopt, |
| 588 | "select_linear_kernel", |
| 589 | lambda *a, **kw: SimpleNamespace(input_quant_key=lambda: None), |
| 590 | ) |
| 591 | monkeypatch.setattr(qwen3_dflash, "get_draft_quant_config", lambda _: quant_config) |
| 592 | |
| 593 | class Backbone(nn.Module): |
| 594 | def __init__(self, **kwargs): |
| 595 | super().__init__() |
| 596 | self.quant_config = quant_config |
| 597 | self.use_aux_hidden_state = True |
| 598 | self.has_separate_mask_embedding = False |
| 599 | self.lilicorr = LiLiCorrHead( |
| 600 | model_hidden_size=16, block_size=4, rms_norm_eps=1e-6, config=_config() |
| 601 | ) |
| 602 | |
| 603 | def _build_fused_kv_buffers(self): |
| 604 | pass |
| 605 | |
| 606 | monkeypatch.setattr(LiLiCorrForCausalLM, "model_cls", Backbone) |
| 607 | monkeypatch.setattr(LiLiCorrForCausalLM, "_read_mask_embedding", lambda _: None) |
| 608 | |
| 609 | def build(owned=True): |
| 610 | config = SimpleNamespace( |
| 611 | vocab_size=64, |
| 612 | draft_vocab_size=64, |
| 613 | hidden_size=16, |
| 614 | has_own_lm_head=owned, |
| 615 | num_hidden_layers=0, |
| 616 | dflash_config={}, |
| 617 | ) |
| 618 | vllm_config = SimpleNamespace( |
| 619 | speculative_config=SimpleNamespace( |
| 620 | draft_model_config=SimpleNamespace(hf_config=config), |
| 621 | draft_load_config=None, |
| 622 | attention_backend=None, |
| 623 | kv_cache_dtype=None, |
| 624 | ), |
| 625 | model_config=SimpleNamespace( |
| 626 | get_total_num_hidden_layers=lambda: 0, |
| 627 | get_vocab_size=lambda: 64, |
| 628 | ), |
| 629 | attention_config=SimpleNamespace(), |
| 630 | cache_config=SimpleNamespace(), |
| 631 | load_config=LoadConfig(), |
| 632 | ) |
| 633 | with set_current_vllm_config( |
| 634 | VllmConfig( |
| 635 | device_config=DeviceConfig("cpu"), |
| 636 | compilation_config=CompilationConfig(mode=CompilationMode.NONE), |
| 637 | ) |
| 638 | ): |
| 639 | model = LiLiCorrForCausalLM(vllm_config=vllm_config) |
| 640 | return model, vllm_config |
| 641 | |
| 642 | return build |
| 643 | |
| 644 | |
| 645 | @pytest.mark.parametrize("missing", [None, "weight", "weight_scale", "weight_scale_2"]) |
| 646 | def test_owned_nvfp4_lm_head_loads_required_checkpoint_tensors( |
| 647 | owned_head_model, missing |
| 648 | ): |
| 649 | model, _ = owned_head_model() |
| 650 | assert model.lm_head.weight.dtype == torch.uint8 |
| 651 | assert model.lm_head.weight.shape == (64, 8) |
| 652 | weights = { |
| 653 | name: torch.ones_like(value) for name, value in model.state_dict().items() |
| 654 | } |
| 655 | del weights["lm_head.input_scale"] # Deprecated and unused for W4A16. |
| 656 | if missing: |
| 657 | del weights[f"lm_head.{missing}"] |
| 658 | with pytest.raises(ValueError, match=f"lm_head.{missing}"): |
| 659 | model.load_weights(weights.items()) |
| 660 | else: |
| 661 | model.load_weights(weights.items()) |
| 662 | assert model.has_own_lm_head |
| 663 | for name in ("weight", "weight_scale", "weight_scale_2"): |
| 664 | torch.testing.assert_close( |
| 665 | getattr(model.lm_head, name), weights[f"lm_head.{name}"] |
| 666 | ) |
| 667 | |
| 668 | |
| 669 | @pytest.mark.parametrize("owned", [False, True]) |
| 670 | def test_selected_lm_head_survives_dflash_and_lilicorr_sharing( |
| 671 | monkeypatch, owned_head_model, owned |
| 672 | ): |
| 673 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 674 | VocabParallelEmbedding, |
| 675 | ) |
| 676 | from vllm.v1.worker.gpu.spec_decode.dflash import utils as dflash_utils |
| 677 | from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils |
| 678 | |
| 679 | model, config = owned_head_model(owned) |
| 680 | own_head = model.lm_head |
| 681 | target_head = owned_head_model(owned)[0].lm_head |
| 682 | # Identical packed bytes do not imply identical quantized heads. |
| 683 | with torch.no_grad(): |
| 684 | own_head.weight.fill_(1) |
| 685 | target_head.weight.fill_(1) |
| 686 | if owned: |
| 687 | own_head.weight_scale_2.fill_(1) |
| 688 | target_head.weight_scale_2.fill_(2) |
| 689 | embedding = VocabParallelEmbedding(64, 16) |
| 690 | target = SimpleNamespace( |
| 691 | model=SimpleNamespace(embed_tokens=embedding), lm_head=target_head |
| 692 | ) |
| 693 | monkeypatch.setattr( |
| 694 | eagle_utils, "get_pp_group", lambda: SimpleNamespace(world_size=1) |
| 695 | ) |
| 696 | monkeypatch.setattr( |
| 697 | dflash_utils, "replace", lambda obj, **kw: SimpleNamespace(**(vars(obj) | kw)) |
| 698 | ) |
| 699 | monkeypatch.setattr(dflash_utils, "get_pp_safe_draft_load_config", lambda c: c) |
| 700 | monkeypatch.setattr(dflash_utils, "get_model", lambda **kw: model) |
| 701 | if owned: |
| 702 | |
| 703 | def no_weight_comparison(*args): |
| 704 | raise AssertionError("Explicit ownership must bypass raw-weight comparison") |
| 705 | |
| 706 | monkeypatch.setattr(dflash_utils, "_should_share", no_weight_comparison) |
| 707 | spec = LiLiCorrSpeculator.__new__(LiLiCorrSpeculator) |
| 708 | spec.vllm_config = config |
| 709 | assert spec.load_draft_model(target, set()) is model |
| 710 | assert model.lm_head is (own_head if owned else target_head) |
| 711 | assert spec.target_embeddings is embedding |
| 712 | if owned: |
| 713 | target.lm_head = None |
| 714 | assert spec.load_draft_model(target, set()).lm_head is own_head |
| 715 | |
| 716 | |
| 717 | def test_quantized_head_calls_methods_without_reading_packed_weights(monkeypatch): |
| 718 | from vllm.model_executor.layers import linear |
| 719 | from vllm.model_executor.layers.linear import UnquantizedLinearMethod |
| 720 | |
| 721 | calls = [] |
| 722 | configured = [] |
| 723 | quant_config = object() |
| 724 | |
| 725 | class PackedMethod(UnquantizedLinearMethod): |
| 726 | def apply(self, layer, x, bias=None): |
| 727 | calls.append(layer.prefix) |
| 728 | # Packed storage has a different name/dtype from the activations. |
| 729 | assert x.dtype == torch.float32 |
| 730 | return F.linear(x, layer.packed_weight.float() * layer.weight_scale, bias) |
| 731 | |
| 732 | def resolve(config, layer, prefix): |
| 733 | assert config is quant_config |
| 734 | configured.append(prefix) |
| 735 | return PackedMethod() |
| 736 | |
| 737 | monkeypatch.setattr(linear, "resolve_quant_method", resolve) |
| 738 | with set_current_vllm_config( |
| 739 | VllmConfig( |
| 740 | device_config=DeviceConfig("cpu"), |
| 741 | compilation_config=CompilationConfig(mode=CompilationMode.NONE), |
| 742 | ) |
| 743 | ): |
| 744 | head = LiLiCorrHead( |
| 745 | model_hidden_size=16, |
| 746 | block_size=4, |
| 747 | rms_norm_eps=1e-6, |
| 748 | config=_config(hidden_size=8), |
| 749 | quant_config=quant_config, |
| 750 | prefix="model.lilicorr", |
| 751 | ) |
| 752 | with torch.no_grad(): |
| 753 | for module in head.modules(): |
| 754 | if isinstance(module, LinearBase): |
| 755 | if module.quant_config is None: |
| 756 | module.weight.normal_(std=0.1) |
| 757 | module.bias.zero_() |
| 758 | continue |
| 759 | module.register_parameter( |
| 760 | "packed_weight", |
| 761 | nn.Parameter( |
| 762 | torch.randint(-4, 5, module.weight.shape, dtype=torch.int8), |
| 763 | requires_grad=False, |
| 764 | ), |
| 765 | ) |
| 766 | module.register_parameter( |
| 767 | "weight_scale", |
| 768 | nn.Parameter( |
| 769 | torch.tensor(0.05), |
| 770 | requires_grad=False, |
| 771 | ), |
| 772 | ) |
| 773 | del module.weight |
| 774 | module.bias.zero_() |
| 775 | for layer in head.layers: |
| 776 | layer.attn.in_proj_weight.normal_(std=0.1) |
| 777 | layer.attn.in_proj_bias.zero_() |
| 778 | head.materialize_inference_buffers(torch.device("cpu"), torch.float32) |
| 779 | result = head( |
| 780 | torch.randn(2, 3, 4, 16), |
| 781 | torch.randn(2, 3, 4).log_softmax(-1), |
| 782 | torch.randn(2, 3, 16), |
| 783 | torch.randn(2, 16), |
| 784 | torch.tensor([True, False]), |
| 785 | ) |
| 786 | assert result.shape == (2, 3, 4, 4) |
| 787 | assert torch.isfinite(result).all() |
| 788 | for name in ("factor_input_proj", "out_head", "in_head"): |
| 789 | assert getattr(head, name).quant_config is None |
| 790 | assert f"model.lilicorr.{name}" not in configured |
| 791 | assert set(calls) == set(configured) |
| 792 | assert len(configured) == len(set(configured)) |
| 793 | for name, module in head.named_modules(): |
| 794 | if isinstance(module, LinearBase): |
| 795 | assert module.prefix == f"model.lilicorr.{name}" |