@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 73a78e6f1f

[SpecDecode] Add LiLiCorr drafter (#57934)

Signed-off-by: Andrii Skliar <askliar@nvidia.com> Signed-off-by: Andrii Skliar <andreyws96@gmail.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>

Andrii Skliar committed Sep 25, 2026 at 16:08 UTC 73a78e6f1f38e280986b81e0f2a9aa5e1ee6fe47

21 files changed +1933 -99

docs/features/speculative_decoding/README.md

@@ -9,6 +9,7 @@ To train your own draft models for optimized speculative decoding, see [vllm-pro

9 vLLM supports a variety of methods of speculative decoding. Model-based methods such as EAGLE, MTP, draft models, PARD and MLP provide the best latency reduction, while simpler methods such as n-gram and suffix decoding provide modest speedups without increasing workload during peak traffic.

11 - [EAGLE](eagle.md)

12 +- [LiLiCorr](lilicorr.md)

13 - [Multi-Token Prediction (MTP)](mtp.md)

14 - [Draft Model](draft_model.md)

15 - [Parallel Draft Model (PARD)](parallel_draft_model.md)

docs/features/speculative_decoding/lilicorr.md new

+77

@@ -0,0 +1,77 @@

1 +# LiLiCorr

2 +

3 +LiLiCorr uses a DFlash backbone to produce per-position candidates, then scores

4 +candidate combinations with a learned correlator. It requires a compatible

5 +trained checkpoint and Model Runner V2, which is selected automatically.

6 +Both plain and grouped-convolution checkpoints are supported.

7 +

8 +## Usage

9 +

10 +For a checkpoint trained with block size 16:

11 +

12 +```bash

13 +vllm serve /path/to/target --dtype bfloat16 \

14 + --speculative-config '{"method":"dflash","model":"/path/to/lilicorr","num_speculative_tokens":15,"draft_sample_method":"probabilistic"}'

15 +```

16 +

17 +We strongly recommend explicitly setting `num_speculative_tokens` to the

18 +checkpoint's trained `block_size - 1` (for example, `15` for `block_size=16`).

19 +The trained block includes one anchor position. Other supported lengths are

20 +available, but may degrade acceptance length unpredictably; benchmark them with

21 +your checkpoint and workload before using them.

22 +

23 +Shorter drafts, down to one token, use the corresponding prefix of the learned

24 +slot embeddings and attention bias without changing checkpoint parameter shapes.

25 +Lengths above `block_size - 1` are unsupported because the checkpoint has no

26 +learned slot embeddings for those positions.

27 +

28 +Fixed-length drafting with shorter verification is a separate option provided by

29 +adaptive verification (`enable_adaptive_verification`). LiLiCorr shares the

30 +candidate-sampling and acceptance-estimation path, but this combination has not

31 +been validated end-to-end. Adaptive verification also requires compatible target

32 +attention backends and variable-length CUDA graph support; do not assume it works

33 +with Mamba2 targets.

34 +

35 +## Proposal sampling

36 +

37 +- `"draft_sample_method":"greedy"` selects the highest-scoring candidate at each

38 + step of the correlator's conditional walk.

39 +- `"draft_sample_method":"probabilistic"` samples from those conditional scores

40 + using the request's temperature and supplies the realized proposal

41 + distributions to the rejection sampler.

42 +

43 +These settings control proposals; they do not replace the target model's

44 +sampling settings. With request temperature zero, probabilistic proposals also

45 +use argmax. For nonzero temperatures, benchmark both modes with your checkpoint

46 +and workload: compare acceptance length and end-to-end throughput or latency.

47 +Neither mode is universally faster.

48 +

49 +The default rejection method is `"rejection_sample_method":"standard"`.

50 +The shared Model Runner V2 rejection sampler also exposes `"block"`, but

51 +LiLiCorr-specific GPU correctness and performance validation for that combination

52 +is still pending. Do not assume that changing the rejection method improves

53 +throughput; evaluate it separately from the proposal sampling mode.

54 +

55 +## Checkpoint requirements

56 +

57 +The checkpoint must declare `LiLiCorrDraftModel` and include all `lilicorr_*`

58 +geometry fields and trained head weights. Geometry is read from `dflash_config`.

59 +Convolution tensors must match the configured `conv_kernel_size` and

60 +`conv_group_size`. Target input embeddings must be available on the draft rank.

61 +By default, candidates use the target LM head, which must also be available there.

62 +A checkpoint with top-level `"has_own_lm_head": true` instead uses its own

63 +`lm_head` with the draft quantization configuration and exclusions. Its weights

64 +must be present; for ModelOpt NVFP4 this includes `lm_head.weight`,

65 +`lm_head.weight_scale`, and `lm_head.weight_scale_2`. The owned head is preserved

66 +and is not replaced by the target head. Candidate token embeddings still come

67 +from the target.

68 +

69 +Correlator linear layers and convolution kernel projections use the draft

70 +quantization configuration and its module exclusions. The following parameters

71 +must remain in floating-point model dtype:

72 +

73 +- Lattice QKV `in_proj_weight` and `in_proj_bias`, which retain the exported

74 + parameter layout.

75 +- `lilicorr.factor_input_proj`, `lilicorr.out_head`, and `lilicorr.in_head`, which

76 + are used by the split factor and fused edge projections. Exclude these modules

77 + when exporting a quantized checkpoint.

tests/models/registry.py

@@ -1502,6 +1502,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

1502 max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env

1503 max_num_seqs=32,

1504 ),

1505 + # LiLiCorr checkpoints are not published yet.

1506 + "LiLiCorrDraftModel": _HfExamplesInfo(

1507 + "Qwen/Qwen3-8B",

1508 + speculative_model="LiLiCorrDraftModel",

1509 + is_available_online=False,

1510 + use_original_num_layers=True,

1511 + ),

1512 "DFlash2DraftModel": _HfExamplesInfo(

1513 "Qwen/Qwen3.8-27B",

1514 speculative_model="z-lab/Qwen3.8-27B-DFlash2",

tests/test_config.py

+14 -9

@@ -582,7 +582,7 @@ def test_rocm_mrv1_default_yields_to_v1_unsupported_config(monkeypatch):

582 speculative_config=None,

583 )

584 config._dflash_needs_multi_kv_group = lambda: False

585 - config._is_dflash2_draft = lambda: False

585 + config._is_dflash_candidate_draft = lambda: False

586 config._get_v2_model_runner_unsupported_features = lambda: []

587 # The real predicate, so the test also pins where dspark lands in it.

588 config._get_v1_model_runner_unsupported_features = lambda: (

@@ -813,7 +813,8 @@ def test_v2_model_runner_supports_custom_logits_processors():

813 assert config._get_v2_model_runner_unsupported_features() == []

814

815

816 -def test_dflash2_draft_forces_v2_model_runner():

816 +@pytest.mark.parametrize("architecture", ["DFlash2DraftModel", "LiLiCorrDraftModel"])

817 +def test_dflash_candidate_draft_forces_v2_model_runner(architecture):

818 """A DFlash2 draft must reach the V2 speculator, the only one that runs its

819 candidate selector; on V1 it would draft as DFlash1 without raising."""

820

@@ -825,11 +826,15 @@ def test_dflash2_draft_forces_v2_model_runner():

826 )

827 )

828

828 - assert VllmConfig._is_dflash2_draft(config("dflash", ["DFlash2DraftModel"]))

829 - assert not VllmConfig._is_dflash2_draft(config("dflash", ["DFlashDraftModel"]))

830 - assert not VllmConfig._is_dflash2_draft(config("eagle", ["DFlash2DraftModel"]))

831 - assert not VllmConfig._is_dflash2_draft(SimpleNamespace(speculative_config=None))

832 - assert not VllmConfig._is_dflash2_draft(

829 + assert VllmConfig._is_dflash_candidate_draft(config("dflash", [architecture]))

830 + assert not VllmConfig._is_dflash_candidate_draft(

831 + config("dflash", ["DFlashDraftModel"])

832 + )

833 + assert not VllmConfig._is_dflash_candidate_draft(config("eagle", [architecture]))

834 + assert not VllmConfig._is_dflash_candidate_draft(

835 + SimpleNamespace(speculative_config=None)

836 + )

837 + assert not VllmConfig._is_dflash_candidate_draft(

838 SimpleNamespace(

839 speculative_config=SimpleNamespace(method="dflash", draft_model_config=None)

840 )

@@ -1243,7 +1248,7 @@ def test_v1_model_runner_rejects_v2_only_features():

1248 model_config=None,

1249 )

1250 config._dflash_needs_multi_kv_group = lambda: False

1246 - config._is_dflash2_draft = lambda: False

1251 + config._is_dflash_candidate_draft = lambda: False

1252 config._get_v1_model_runner_unsupported_features = lambda: (

1253 VllmConfig._get_v1_model_runner_unsupported_features(config)

1254 )

@@ -1430,7 +1435,7 @@ def test_v1_model_runner_rejects_pipeline_parallelism_with_async_scheduling():

1435 model_config=None,

1436 )

1437 config._dflash_needs_multi_kv_group = lambda: False

1433 - config._is_dflash2_draft = lambda: False

1438 + config._is_dflash_candidate_draft = lambda: False

1439

1440 assert VllmConfig._get_v1_model_runner_unsupported_features(config) == []

1441

tests/v1/spec_decode/test_dflash2.py

+180 -5

@@ -6,6 +6,9 @@ from types import SimpleNamespace

6 import pytest

7 import torch

9 +from vllm.model_executor.models.qwen3_dflash import (

10 + _add_global_draft_layer_exclusions,

11 +)

12 from vllm.model_executor.models.qwen3_dflash2 import _grouped_conv, _score_edges

13 from vllm.platforms import current_platform

14 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator

@@ -85,6 +88,24 @@ def test_grouped_conv_triton_matches_reference(

88 )

91 +def test_draft_quant_exclusions_include_global_layer_indices():

92 + quant_config = SimpleNamespace(

93 + exclude_modules=[

94 + "layers.0.mlp_conv*",

95 + "*layers.4.self_attn.q_proj",

96 + "layers.88.already_global",

97 + "lilicorr.layers.0.mlp.0",

98 + ]

99 + )

100 +

101 + _add_global_draft_layer_exclusions(quant_config, 88, 5)

102 +

103 + assert "layers.88.mlp_conv*" in quant_config.exclude_modules

104 + assert "*layers.92.self_attn.q_proj" in quant_config.exclude_modules

105 + assert quant_config.exclude_modules.count("layers.88.already_global") == 1

106 + assert "lilicorr.layers.88.mlp.0" not in quant_config.exclude_modules

107 +

108 +

109 def test_selector_edges_match_sequential_reference():

110 torch.manual_seed(1)

111 batch, steps, top_k, rank = 2, 4, 3, 5

@@ -169,10 +190,13 @@ def test_selector_asks_for_fp32_proposal_logits():

190

191

192 @pytest.mark.skip_global_cleanup

172 -def test_dflash2_model_decoder_layer_cls(monkeypatch):

193 +@pytest.mark.parametrize("variant", ["dflash2", "lilicorr", "lilicorr_plain"])

194 +def test_candidate_model_decoder_layer_cls(monkeypatch, variant):

195 from types import SimpleNamespace

196

197 from vllm.config import set_current_vllm_config

198 + from vllm.model_executor.models.lilicorr import LiLiCorr

199 + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3DecoderLayer

200 from vllm.model_executor.models.qwen3_dflash2 import (

201 DFlash2Qwen3DecoderLayer,

202 DFlash2Qwen3Model,

@@ -249,8 +273,17 @@ def test_dflash2_model_decoder_layer_cls(monkeypatch):

273 dflash_config={

274 "selector_rank": 4,

275 "selector_top_k": 3,

252 - "conv_kernel_size": 3,

253 - "conv_group_size": 2,

276 + "conv_kernel_size": 0 if variant == "lilicorr_plain" else 3,

277 + "conv_group_size": 0 if variant == "lilicorr_plain" else 2,

278 + "block_size": 5,

279 + "lilicorr_candidate_topk": 4,

280 + "lilicorr_hidden_size": 8,

281 + "lilicorr_num_layers": 2,

282 + "lilicorr_num_heads": 2,

283 + "lilicorr_mlp_ratio": 2.0,

284 + "lilicorr_factor_dim": 4,

285 + "lilicorr_vector_eps": 1e-6,

286 + "lilicorr_logit_scale": 3.0,

287 "use_aux_hidden_state": False,

288 },

289 )

@@ -277,8 +310,150 @@ def test_dflash2_model_decoder_layer_cls(monkeypatch):

310

311 # 3. Instantiate the model under meta device to avoid parameter allocation issues

312 with set_current_vllm_config(mock_current_vllm_config), torch.device("meta"):

280 - model = DFlash2Qwen3Model(vllm_config=vllm_config)

313 + model_cls = DFlash2Qwen3Model if variant == "dflash2" else LiLiCorr

314 + model = model_cls(vllm_config=vllm_config)

315

316 # 4. Assert that the layers are DFlash2Qwen3DecoderLayer (the subclass)

317 assert len(model.layers) == 2

284 - assert isinstance(model.layers[0], DFlash2Qwen3DecoderLayer)

318 + expected = (

319 + DFlashQwen3DecoderLayer

320 + if variant == "lilicorr_plain"

321 + else DFlash2Qwen3DecoderLayer

322 + )

323 + assert type(model.layers[0]) is expected

324 +

325 +

326 +def test_conv_projections_use_draft_quant_config(monkeypatch):

327 + from torch import nn

328 +

329 + from vllm.distributed import parallel_state

330 + from vllm.model_executor.layers.quantization import modelopt

331 + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3DecoderLayer

332 + from vllm.model_executor.models.qwen3_dflash2 import DFlash2Qwen3DecoderLayer

333 + from vllm.model_executor.models.utils import AutoWeightsLoader

334 +

335 + monkeypatch.setattr(

336 + parallel_state, "_TP", SimpleNamespace(rank_in_group=0, world_size=1)

337 + )

338 + monkeypatch.setattr(

339 + DFlashQwen3DecoderLayer,

340 + "__init__",

341 + lambda self, *a, **kw: nn.Module.__init__(self),

342 + )

343 + # Exercise the real quantized parameter allocation without selecting a GPU kernel.

344 + monkeypatch.setattr(

345 + modelopt,

346 + "select_linear_kernel",

347 + lambda *a, **kw: SimpleNamespace(input_quant_key=lambda: None),

348 + )

349 + quant_config = modelopt.ModelOptNvFp4Config(

350 + quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True

351 + )

352 + layer = DFlash2Qwen3DecoderLayer(

353 + SimpleNamespace(

354 + speculative_config=SimpleNamespace(num_speculative_tokens=7),

355 + model_config=SimpleNamespace(dtype=torch.bfloat16),

356 + ),

357 + config=SimpleNamespace(

358 + hidden_size=16,

359 + dflash_config={"conv_kernel_size": 2, "conv_group_size": 2},

360 + ),

361 + layer_idx=0,

362 + prefix="model.layers.0",

363 + quant_config=quant_config,

364 + )

365 + for name in ("attention_conv", "mlp_conv"):

366 + module = getattr(layer, name)

367 + projection = module.kernel_projection

368 + assert projection.weight.dtype == torch.uint8

369 + assert projection.weight.shape == (32, 8)

370 + weight = torch.ones_like(projection.weight)

371 + AutoWeightsLoader(module).load_weights([("kernel_projection.weight", weight)])

372 + torch.testing.assert_close(projection.weight, weight)

373 + assert module.base_kernel.dtype == torch.bfloat16

374 +

375 +

376 +def test_context_kv_uses_quantized_projection_fallback(monkeypatch):

377 + from torch import nn

378 +

379 + from vllm.model_executor.models import qwen3_dflash

380 +

381 + class Projection(nn.Module):

382 + def __init__(self, packed_weight):

383 + super().__init__()

384 + self.register_buffer("packed_weight", packed_weight)

385 + self.quant_method = object()

386 + self.calls = 0

387 +

388 + def forward(self, hidden_states):

389 + self.calls += 1

390 + return torch.nn.functional.linear(hidden_states, self.packed_weight), None

391 +

392 + monkeypatch.setattr(

393 + qwen3_dflash.ops,

394 + "rms_norm",

395 + lambda output, hidden_states, weight, eps: output.copy_(hidden_states),

396 + )

397 + context_states = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

398 + projections = [

399 + Projection(

400 + torch.tensor(

401 + [

402 + [0.0, 0.0],

403 + [0.0, 0.0],

404 + [1.0, 0.0],

405 + [0.0, 1.0],

406 + [1.0, 1.0],

407 + [1.0, -1.0],

408 + ]

409 + )

410 + ),

411 + Projection(

412 + torch.tensor(

413 + [

414 + [0.0, 0.0],

415 + [0.0, 0.0],

416 + [2.0, 0.0],

417 + [0.0, 2.0],

418 + [-1.0, 0.0],

419 + [0.0, -1.0],

420 + ]

421 + )

422 + ),

423 + ]

424 + model = SimpleNamespace(

425 + hidden_norm=SimpleNamespace(weight=nn.Parameter(torch.ones(2))),

426 + _rms_norm_eps=1e-6,

427 + )

428 + layers_attn = [

429 + SimpleNamespace(

430 + qkv_proj=projection,

431 + q_size=2,

432 + k_norm=SimpleNamespace(weight=nn.Parameter(torch.ones(2))),

433 + )

434 + for projection in projections

435 + ]

436 + qwen3_dflash.DFlashQwen3Model._build_context_kv_buffers(

437 + model, layers_attn, has_bias=False

438 + )

439 + assert model._fused_kv_weight is None

440 + assert all(not hasattr(projection, "weight") for projection in projections)

441 +

442 + actual_k, actual_v = qwen3_dflash.DFlashQwen3Model._project_context_kv(

443 + model,

444 + context_states,

445 + num_ctx=2,

446 + num_layers=2,

447 + num_kv_heads=1,

448 + head_dim=2,

449 + )

450 +

451 + expected_k = torch.tensor(

452 + [[[1.0, 2.0], [3.0, 4.0]], [[2.0, 4.0], [6.0, 8.0]]]

453 + ).unsqueeze(2)

454 + expected_v = torch.tensor(

455 + [[[3.0, -1.0], [7.0, -1.0]], [[-1.0, -2.0], [-3.0, -4.0]]]

456 + ).unsqueeze(2)

457 + torch.testing.assert_close(actual_k, expected_k)

458 + torch.testing.assert_close(actual_v, expected_v)

459 + assert [projection.calls for projection in projections] == [1, 1]

tests/v1/spec_decode/test_lilicorr.py new

+777

@@ -0,0 +1,777 @@

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 + DeviceConfig,

15 + LoadConfig,

16 + VllmConfig,

17 + set_current_vllm_config,

18 +)

19 +from vllm.distributed import parallel_state

20 +from vllm.model_executor.layers import logits_processor

21 +from vllm.model_executor.layers.linear import LinearBase

22 +from vllm.model_executor.models.lilicorr import LiLiCorrHead

23 +from vllm.transformers_utils.configs.eagle import EAGLEConfig

24 +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import CandidateSampler

25 +from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import LiLiCorrSpeculator

26 +

27 +

28 +@pytest.fixture(autouse=True)

29 +def head_tp_metadata(monkeypatch):

30 + # Single-rank unit tests only need TP metadata.

31 + monkeypatch.setattr(

32 + parallel_state, "_TP", SimpleNamespace(rank_in_group=0, world_size=1)

33 + )

34 +

35 +

36 +def _config(**overrides):

37 + config = dict(

38 + candidate_topk=4,

39 + hidden_size=8,

40 + num_layers=2,

41 + num_heads=2,

42 + mlp_ratio=2.0,

43 + factor_dim=4,

44 + vector_eps=1e-6,

45 + logit_scale=3.0,

46 + )

47 + config.update(overrides)

48 + return {f"lilicorr_{key}": value for key, value in config.items()}

49 +

50 +

51 +def _reference_scores(head, embeddings, log_probs, hidden, anchor, valid):

52 + """Unfused exported head: MHA attention, concatenated factors, separate edges."""

53 + batch, slots, top_k = log_probs.shape

54 + rank = torch.arange(top_k, device=log_probs.device) / max(top_k - 1, 1)

55 + top1 = (torch.arange(top_k, device=log_probs.device) == 0).float()

56 + features = torch.stack(

57 + (

58 + log_probs,

59 + log_probs.exp(),

60 + log_probs - log_probs.amax(-1, keepdim=True),

61 + rank.expand_as(log_probs),

62 + top1.expand_as(log_probs),

63 + ),

64 + -1,

65 + )

66 + x = head.token_proj(embeddings) + head.pass_hidden_proj(hidden).unsqueeze(-2)

67 + features = head.feature_norm(features.to(x.dtype))

68 + x = x + head.feature_mlp.down_proj(F.silu(head.feature_mlp.up_proj(features)[0]))[0]

69 + x = x + head.slot_embedding[:, 0, :slots] + head.rank_embedding[:, 0]

70 + x = x.flatten(1, 2)

71 + positions = torch.arange(slots, device=x.device).repeat_interleave(top_k)

72 + relative = positions[:, None] - positions[None, :]

73 + bias = head.relative_slot_bias[:, relative + head.block_size - 1]

74 + bias = bias + (relative == 0) * head.same_slot_bias[:, None, None]

75 + bias = bias.repeat(batch, 1, 1)

76 + for layer in head.layers:

77 + attention = nn.MultiheadAttention(

78 + head.hidden_size,

79 + head.num_heads,

80 + batch_first=True,

81 + dtype=x.dtype,

82 + device=x.device,

83 + )

84 + attention.load_state_dict(layer.attn.state_dict())

85 + normalized = layer.attn_norm(x)

86 + x = (

87 + x

88 + + attention(

89 + normalized, normalized, normalized, attn_mask=bias, need_weights=False

90 + )[0]

91 + )

92 + x = x + layer.mlp.down_proj(F.silu(layer.mlp.up_proj(layer.mlp_norm(x))[0]))[0]

93 + x = head.output_norm(x).view(batch, slots, top_k, -1)

94 + a = head.anchor_norm(head.context_proj(anchor) * valid[:, None])

95 + expanded = a[:, None, None].expand_as(x)

96 + factors = F.silu(head.factor_input_proj(torch.cat((x, expanded, x * expanded), -1)))

97 + incoming = F.normalize(head.in_head(factors), dim=-1, eps=head.vector_eps)

98 + outgoing = F.normalize(head.out_head(factors), dim=-1, eps=head.vector_eps)

99 + a = F.normalize(head.anchor_out_head(a), dim=-1, eps=head.vector_eps)

100 + start = (a[:, None] * incoming[:, 0]).sum(-1)

101 + pairs = outgoing[:, :-1] @ incoming[:, 1:].transpose(-1, -2)

102 + return (

103 + head.logit_scale

104 + * torch.cat((start[:, None, None].expand(-1, 1, top_k, -1), pairs), 1).float()

105 + )

106 +

107 +

108 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm")

109 +@pytest.mark.parametrize(

110 + "head_width,slots,dtype",

111 + [

112 + pytest.param(8, 3, torch.float32, id="projected-multi-slot"),

113 + pytest.param(16, 1, torch.float32, id="identity-shortened-single-slot"),

114 + pytest.param(8, 2, torch.bfloat16, id="bf16-shortened-multi-slot"),

115 + ],

116 +)

117 +def test_lilicorr_matches_exported_head(head_width, slots, dtype):

118 + device = torch.device("cuda")

119 + torch.manual_seed(7)

120 + with (

121 + torch.device(device),

122 + set_current_vllm_config(

123 + VllmConfig(

124 + device_config=DeviceConfig("cuda"),

125 + compilation_config=CompilationConfig(mode=0),

126 + )

127 + ),

128 + ):

129 + head = LiLiCorrHead(

130 + model_hidden_size=16,

131 + block_size=4,

132 + rms_norm_eps=1e-6,

133 + config=_config(hidden_size=head_width),

134 + ).to(dtype)

135 + # vLLM linear parameters are initialized by the checkpoint loader.

136 + with torch.no_grad(), torch.device(device):

137 + for module in head.modules():

138 + if isinstance(module, LinearBase):

139 + for parameter in module.parameters():

140 + parameter.weight_loader(

141 + parameter, torch.randn_like(parameter) * 0.1

142 + )

143 + for layer in head.layers:

144 + layer.attn.in_proj_weight.normal_(std=0.1)

145 + layer.attn.in_proj_bias.normal_(std=0.1)

146 + head.relative_slot_bias.normal_(std=0.1)

147 + head.same_slot_bias.normal_(std=0.1)

148 + head.slot_embedding.normal_(std=0.1)

149 + head.rank_embedding.normal_(std=0.1)

150 + head.materialize_inference_buffers(device, dtype)

151 + inputs = (

152 + torch.randn(2, slots, 4, 16, dtype=dtype),

153 + torch.randn(2, slots, 23).log_softmax(-1).topk(4).values,

154 + torch.randn(2, slots, 16, dtype=dtype),

155 + torch.randn(2, 16, dtype=dtype),

156 + torch.tensor([True, False]),

157 + )

158 + actual = head(*inputs)

159 + expected = _reference_scores(head, *inputs)

160 + tolerance = 0.04 if dtype == torch.bfloat16 else 2e-6

161 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance)

162 +

163 +

164 +def test_runtime_length_keeps_trained_checkpoint_geometry(monkeypatch):

165 + from vllm.model_executor.models.lilicorr import LiLiCorr

166 + from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model

167 +

168 + def init_backbone(self, **kwargs):

169 + nn.Module.__init__(self)

170 + self.quant_config = None

171 +

172 + monkeypatch.setattr(DFlashQwen3Model, "__init__", init_backbone)

173 + spec = SimpleNamespace(

174 + num_speculative_tokens=1,

175 + draft_model_config=SimpleNamespace(

176 + hf_config=SimpleNamespace(

177 + hidden_size=16,

178 + rms_norm_eps=1e-6,

179 + dflash_config=_config() | {"block_size": 4},

180 + )

181 + ),

182 + )

183 + with set_current_vllm_config(

184 + VllmConfig(

185 + device_config=DeviceConfig("cpu"),

186 + compilation_config=CompilationConfig(mode=0),

187 + )

188 + ):

189 + for slots in (1, 3):

190 + spec.num_speculative_tokens = slots

191 + head = LiLiCorr(

192 + vllm_config=SimpleNamespace(speculative_config=spec)

193 + ).lilicorr

194 + assert head.slot_embedding.shape == (1, 1, 3, 1, 8)

195 + assert head.relative_slot_bias.shape == (2, 7)

196 + spec.draft_model_config.hf_config.dflash_config["lilicorr_enabled"] = False

197 + with pytest.raises(ValueError, match="lilicorr_enabled"):

198 + LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec))

199 + spec.draft_model_config.hf_config.dflash_config["lilicorr_enabled"] = True

200 + spec.num_speculative_tokens = 4

201 + with pytest.raises(ValueError, match="trained block_size"):

202 + LiLiCorr(vllm_config=SimpleNamespace(speculative_config=spec))

203 +

204 +

205 +@pytest.mark.parametrize("tp_size", [1, 2])

206 +def test_candidates_use_global_partition_and_exclude_padding(monkeypatch, tp_size):

207 + """Candidate features must include mass outside top-k and on other TP ranks."""

208 +

209 + class Gather:

210 + def __init__(self, calls):

211 + self.calls = calls

212 + self.index = 0

213 +

214 + def __call__(self, value, dim=-1):

215 + expected_input, result = self.calls[self.index]

216 + assert dim == -1

217 + torch.testing.assert_close(value, expected_input)

218 + self.index += 1

219 + return result

220 +

221 + shards = [

222 + torch.tensor([[1.0, 3.0, 2.0, 100.0], [4.0, 1.0, 2.0, 100.0]]),

223 + torch.tensor([[6.0, 5.0, 0.0, 100.0], [3.0, 6.0, 5.0, 100.0]]),

224 + ][:tp_size]

225 + for rank in range(tp_size):

226 + fake = SimpleNamespace(

227 + scale=1.0,

228 + soft_cap=None,

229 + _apply_head=lambda *args, rank=rank: shards[rank].clone(),

230 + )

231 + lm_head = SimpleNamespace(

232 + tp_size=tp_size,

233 + shard_indices=SimpleNamespace(

234 + num_org_vocab_padding=1, org_vocab_start_index=rank * 3

235 + ),

236 + )

237 + valid_shards = [shard[:, :3] for shard in shards]

238 + local_topk = [shard.topk(2) for shard in valid_shards]

239 + gather_calls = [

240 + (

241 + valid_shards[rank].logsumexp(-1, keepdim=True),

242 + torch.cat(

243 + [shard.logsumexp(-1, keepdim=True) for shard in valid_shards],

244 + -1,

245 + ),

246 + ),

247 + (

248 + local_topk[rank].values,

249 + torch.cat([topk.values for topk in local_topk], -1),

250 + ),

251 + (

252 + local_topk[rank].indices + rank * 3,

253 + torch.cat(

254 + [

255 + topk.indices + shard_rank * 3

256 + for shard_rank, topk in enumerate(local_topk)

257 + ],

258 + -1,

259 + ),

260 + ),

261 + ]

262 + all_gather = Gather(gather_calls)

263 +

264 + monkeypatch.setattr(

265 + logits_processor,

266 + "tensor_model_parallel_all_gather",

267 + all_gather,

268 + )

269 + monkeypatch.setattr(logits_processor, "_topk", lambda x, k: x.topk(k, dim=-1))

270 + ids, values = logits_processor.LogitsProcessor.get_top_k_tokens(

271 + fake, lm_head, torch.empty(2, 1), 2, return_log_probs=True

272 + )

273 + expected = torch.cat(valid_shards, -1).log_softmax(-1).topk(2)

274 + torch.testing.assert_close(ids, expected.indices)

275 + torch.testing.assert_close(values, expected.values)

276 + assert all_gather.index == (3 if tp_size > 1 else 0)

277 +

278 +

279 +def test_context_anchor_is_last_committed_normalized_feature():

280 + hidden = torch.arange(1, 25, dtype=torch.float32).view(6, 4)

281 + norm = nn.RMSNorm(4)

282 + spec = SimpleNamespace(

283 + hidden_states=hidden,

284 + model=SimpleNamespace(model=SimpleNamespace(hidden_norm=norm)),

285 + anchor_hidden=torch.full((4, 4), 123.0),

286 + anchor_valid=torch.ones(4, dtype=torch.bool),

287 + )

288 + batch = SimpleNamespace(num_reqs=2, query_start_loc=torch.tensor([0, 2, 6]))

289 + LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([0, 2]))

290 + torch.testing.assert_close(spec.anchor_hidden[:2], norm(hidden[[1, 3]]))

291 + assert spec.anchor_valid.tolist() == [True, True, False, False]

292 + assert not spec.anchor_hidden[2:].any()

293 + # A smaller reordered batch must not retain an old request's anchor.

294 + batch = SimpleNamespace(num_reqs=1, query_start_loc=torch.tensor([0, 3]))

295 + LiLiCorrSpeculator.prepare_context_anchor(spec, batch, torch.tensor([1]))

296 + torch.testing.assert_close(spec.anchor_hidden[0], norm(hidden[1]))

297 + assert not spec.anchor_hidden[1:].any()

298 +

299 +

300 +@pytest.mark.parametrize("lilicorr", [False, True])

301 +def test_candidate_generation_routes_scores_and_adaptive_inputs(monkeypatch, lilicorr):

302 + """Both heads share the walk, including a padded request and adaptive scores."""

303 + from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator

304 + from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator

305 +

306 + key = "lilicorr_candidate_topk" if lilicorr else "selector_top_k"

307 +

308 + def init_base(self, config, device):

309 + self.draft_model_config = SimpleNamespace(

310 + hf_config=SimpleNamespace(hidden_size=4, dflash_config={key: 4})

311 + )

312 + self.max_num_reqs = 3

313 + self.num_speculative_steps = 2

314 + self.num_query_per_req = 3

315 + self.dtype = torch.float32

316 +

317 + monkeypatch.setattr(DFlashSpeculator, "__init__", init_base)

318 + cls = LiLiCorrSpeculator if lilicorr else DFlash2Speculator

319 + spec = cls(None, torch.device("cpu"))

320 + hidden = torch.arange(36).float().view(9, 4)

321 + spec._run_model = lambda *args: hidden

322 + spec.sample_indices = torch.tensor([1, 2, 4, 5, 7, 8])

323 + spec.sample_pos = torch.arange(6)

324 + spec.sample_idx_mapping = torch.tensor([1, 1, 0, 0, -1, -1])

325 + spec.sample_col = torch.tensor([0, 1, 0, 1, 0, 1])

326 + spec.temperature = torch.ones(3)

327 + spec.seeds = torch.arange(3)

328 + spec.draft_tokens = torch.empty(3, 2, dtype=torch.long)

329 + spec.draft_logits = None

330 + spec.use_fp64_gumbel = False

331 + spec.enable_adaptive_verification = True

332 + spec.input_buffers = SimpleNamespace(input_ids=torch.arange(9))

333 + candidates = torch.arange(24).view(6, 4)

334 + log_probs = candidates.float().log_softmax(-1)

335 + scores = torch.randn(3, 2, 4, 4)

336 + scoring_inputs: list[torch.Tensor] = []

337 + sampling_inputs: list[Any] = []

338 + adaptive_inputs: list[torch.Tensor] = []

339 +

340 + def score(*args):

341 + scoring_inputs.extend(args)

342 + return scores

343 +

344 + def sample(*args):

345 + sampling_inputs.extend(args)

346 + spec.candidate_sampler.scores.copy_(scores[:, :, 0])

347 +

348 + spec.model = SimpleNamespace(

349 + compute_candidates=lambda h: (candidates, log_probs),

350 + model=SimpleNamespace(lilicorr=score, candidate_selector=score),

351 + )

352 + spec.target_embeddings = lambda ids: ids.float().unsqueeze(-1).expand(*ids.shape, 4)

353 + if lilicorr:

354 + spec.anchor_hidden.copy_(torch.arange(12).view(3, 4))

355 + spec.anchor_valid[:2] = True

356 + spec.candidate_sampler.sample = sample

357 + spec._maybe_predict_acceptance = lambda *args: adaptive_inputs.extend(args)

358 + spec._generate_draft(3, 9, None, None, None)

359 +

360 + expected_ids = candidates.view(3, 2, 4)

361 + expected_first = spec.target_embeddings(expected_ids) if lilicorr else expected_ids

362 + torch.testing.assert_close(scoring_inputs[0], expected_first)

363 + torch.testing.assert_close(scoring_inputs[1], log_probs.view(3, 2, 4))

364 + torch.testing.assert_close(

365 + scoring_inputs[2], hidden[spec.sample_indices].view(3, 2, 4)

366 + )

367 + if lilicorr:

368 + torch.testing.assert_close(scoring_inputs[3], spec.anchor_hidden)

369 + torch.testing.assert_close(scoring_inputs[4], spec.anchor_valid)

370 + else:

371 + torch.testing.assert_close(scoring_inputs[3], torch.tensor([0, 3, 6]))

372 + torch.testing.assert_close(sampling_inputs[0], expected_ids)

373 + assert sampling_inputs[1] is scores

374 + assert sampling_inputs[4] is spec.sample_idx_mapping

375 + torch.testing.assert_close(adaptive_inputs[0], scores[:, :, 0].flatten(0, 1))

376 + torch.testing.assert_close(adaptive_inputs[1], spec.sample_idx_mapping)

377 + torch.testing.assert_close(adaptive_inputs[2], spec.sample_col)

378 +

379 +

380 +def test_lilicorr_config_preserves_architecture():

381 + from transformers import Qwen3Config

382 +

383 + config = Qwen3Config(architectures=["LiLiCorrDraftModel"])

384 + wrapped = EAGLEConfig(config, method="dflash")

385 + assert wrapped.architectures == ["LiLiCorrDraftModel"]

386 +

387 +

388 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")

389 +@pytest.mark.parametrize("probabilistic", [False, True])

390 +def test_candidate_walk_preserves_conditional_scores_and_padding(probabilistic):

391 + """Check the shared sampler's actual q, stale support, and graph padding."""

392 + device = torch.device("cuda")

393 + batch, steps, top_k, vocab = 3, 3, 4, 32

394 + sampler = CandidateSampler(batch, steps, top_k, device)

395 + candidates = (

396 + torch.arange(steps * top_k, device=device)

397 + .view(1, steps, top_k)

398 + .repeat(batch, 1, 1)

399 + )

400 + scores = torch.randn(batch, steps, top_k, top_k, device=device)

401 + tokens = torch.full((batch, steps), -1, dtype=torch.long, device=device)

402 + cache = (

403 + torch.full((batch, steps, vocab), -float("inf"), device=device)

404 + if probabilistic

405 + else None

406 + )

407 + positions = torch.arange(batch * steps * 2, device=device)[::2]

408 + mappings = torch.tensor(

409 + [2] * steps + [0] * steps + [-1] * steps, device=device

410 + ).repeat_interleave(2)[::2]

411 + temperature = torch.tensor([0.0, 1.0, 0.8], device=device)

412 + seeds = torch.tensor([11, 22, 33], device=device)

413 +

414 + def run():

415 + sampler.sample(

416 + candidates,

417 + scores,

418 + batch,

419 + positions,

420 + mappings,

421 + temperature,

422 + seeds,

423 + tokens,

424 + cache,

425 + False,

426 + )

427 +

428 + run()

429 + graph = torch.cuda.CUDAGraph()

430 + with torch.cuda.graph(graph):

431 + run()

432 + old_candidates = candidates.clone()

433 + candidates.add_(12)

434 + graph.replay()

435 + for row, state in ((0, 2), (1, 0)):

436 + previous = 0

437 + for step in range(steps):

438 + realized = scores[row, step, previous]

439 + torch.testing.assert_close(sampler.scores[row, step], realized)

440 + index = (candidates[row, step] == tokens[row, step]).nonzero().item()

441 + if not probabilistic or temperature[state] == 0:

442 + assert index == realized.argmax().item()

443 + if cache is not None:

444 + torch.testing.assert_close(

445 + cache[state, step, candidates[row, step]], realized

446 + )

447 + assert cache[state, step, old_candidates[row, step]].isneginf().all()

448 + previous = index

449 + assert (tokens[2] == -1).all()

450 + if cache is not None:

451 + assert cache[1].isneginf().all()

452 +

453 +

454 +@pytest.mark.parametrize(

455 + ("convolution", "mismatch", "quantized"),

456 + [

457 + pytest.param(False, None, False, id="valid-plain"),

458 + pytest.param(True, None, False, id="valid-convolution"),

459 + pytest.param(False, "missing_head", False, id="missing-head"),

460 + pytest.param(False, "extra_head", False, id="extra-head"),

461 + pytest.param(True, "missing_conv", False, id="missing-convolution"),

462 + pytest.param(False, "extra_conv", False, id="unexpected-convolution"),

463 + pytest.param(False, None, True, id="optional-w4a16-input-scale"),

464 + pytest.param(False, "missing_weight", True, id="missing-quantized-weight"),

465 + pytest.param(False, "missing_bias", True, id="missing-quantized-bias"),

466 + pytest.param(False, "missing_weight_scale", True, id="missing-block-scale"),

467 + pytest.param(False, "missing_weight_scale_2", True, id="missing-global-scale"),

468 + ],

469 +)

470 +def test_checkpoint_coverage_rejects_incomplete_or_wrong_heads(

471 + monkeypatch, convolution, mismatch, quantized

472 +):

473 + from vllm.model_executor.models.lilicorr import LiLiCorr, LiLiCorrForCausalLM

474 +

475 + wrapper = LiLiCorrForCausalLM.__new__(LiLiCorrForCausalLM)

476 + nn.Module.__init__(wrapper)

477 + wrapper.model = LiLiCorr.__new__(LiLiCorr)

478 + nn.Module.__init__(wrapper.model)

479 + wrapper.has_own_lm_head = False

480 + with set_current_vllm_config(

481 + VllmConfig(

482 + device_config=DeviceConfig("cpu"),

483 + compilation_config=CompilationConfig(mode=0),

484 + )

485 + ):

486 + wrapper.model.lilicorr = LiLiCorrHead(

487 + model_hidden_size=16, block_size=4, rms_norm_eps=1e-6, config=_config()

488 + )

489 + layer = nn.Module()

490 + if convolution:

491 + layer.attention_conv = nn.Linear(2, 2, bias=False)

492 + wrapper.model.layers = nn.ModuleList([layer])

493 + if quantized:

494 + from vllm.config.quantization import QuantSpec

495 + from vllm.model_executor.layers.quantization.modelopt import (

496 + CkptCtx,

497 + ModelOptLinearMethod,

498 + )

499 + from vllm.model_executor.layers.quantization.utils.quant_utils import (

500 + kNvfp4Static,

501 + )

502 +

503 + # Match the W4A16 NVFP4 checkpoint contract without GPU kernel setup.

504 + projection = wrapper.model.lilicorr.pass_hidden_proj

505 + projection.quant_method = ModelOptLinearMethod(

506 + QuantSpec(weight=kNvfp4Static, activation=None), CkptCtx(group_size=16)

507 + )

508 + projection.weight = nn.Parameter(

509 + torch.zeros(8, 8, dtype=torch.uint8), requires_grad=False

510 + )

511 + for name, value in (

512 + ("weight_scale", torch.ones(8, 1).to(torch.float8_e4m3fn)),

513 + ("weight_scale_2", torch.ones(1)),

514 + ("input_scale", torch.full((1,), torch.nan)),

515 + ):

516 + projection.register_parameter(

517 + name, nn.Parameter(value, requires_grad=False)

518 + )

519 + weights = dict(wrapper.model.state_dict())

520 + if quantized:

521 + del weights["lilicorr.pass_hidden_proj.input_scale"]

522 + if mismatch == "missing_head":

523 + del weights["lilicorr.slot_embedding"]

524 + elif mismatch in (

525 + "missing_weight",

526 + "missing_bias",

527 + "missing_weight_scale",

528 + "missing_weight_scale_2",

529 + ):

530 + del weights[f"lilicorr.pass_hidden_proj.{mismatch.removeprefix('missing_')}"]

531 + elif mismatch == "extra_head":

532 + weights["lilicorr.untrained.weight"] = torch.zeros(1)

533 + elif mismatch == "missing_conv":

534 + del weights["layers.0.attention_conv.weight"]

535 + elif mismatch == "extra_conv":

536 + weights["layers.0.mlp_conv.untrained"] = torch.zeros(1)

537 +

538 + # Keep production loading and mapping, bypass only backbone GPU KV setup.

539 + wrapper.model.use_aux_hidden_state = True

540 + wrapper.model.has_separate_mask_embedding = False

541 + monkeypatch.setattr(wrapper, "_read_mask_embedding", lambda: None)

542 + monkeypatch.setattr(

543 + wrapper.model, "_build_fused_kv_buffers", lambda: None, raising=False

544 + )

545 + supplied = [("model." + name, value) for name, value in weights.items()]

546 + if mismatch:

547 + message = (

548 + "no module or parameter"

549 + if mismatch.startswith("extra")

550 + else "coverage mismatch"

551 + )

552 + with pytest.raises(ValueError, match=message):

553 + wrapper.load_weights(supplied)

554 + else:

555 + wrapper.load_weights(supplied)

556 + assert wrapper.model.lilicorr._attn_bias is not None

557 +

558 +

559 +@pytest.fixture

560 +def owned_head_model(monkeypatch):

561 + from vllm.model_executor.layers.quantization import modelopt

562 + from vllm.model_executor.models import qwen3_dflash

563 + from vllm.model_executor.models.lilicorr import LiLiCorrForCausalLM

564 +

565 + quant_config = modelopt.ModelOptNvFp4Config(

566 + quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True

567 + )

568 + monkeypatch.setattr(

569 + modelopt,

570 + "select_linear_kernel",

571 + lambda *a, **kw: SimpleNamespace(input_quant_key=lambda: None),

572 + )

573 + monkeypatch.setattr(qwen3_dflash, "get_draft_quant_config", lambda _: quant_config)

574 +

575 + class Backbone(nn.Module):

576 + def __init__(self, **kwargs):

577 + super().__init__()

578 + self.quant_config = quant_config

579 + self.use_aux_hidden_state = True

580 + self.has_separate_mask_embedding = False

581 + self.lilicorr = LiLiCorrHead(

582 + model_hidden_size=16, block_size=4, rms_norm_eps=1e-6, config=_config()

583 + )

584 +

585 + def _build_fused_kv_buffers(self):

586 + pass

587 +

588 + monkeypatch.setattr(LiLiCorrForCausalLM, "model_cls", Backbone)

589 + monkeypatch.setattr(LiLiCorrForCausalLM, "_read_mask_embedding", lambda _: None)

590 +

591 + def build(owned=True):

592 + config = SimpleNamespace(

593 + vocab_size=64,

594 + draft_vocab_size=64,

595 + hidden_size=16,

596 + has_own_lm_head=owned,

597 + num_hidden_layers=0,

598 + dflash_config={},

599 + )

600 + vllm_config = SimpleNamespace(

601 + speculative_config=SimpleNamespace(

602 + draft_model_config=SimpleNamespace(hf_config=config),

603 + draft_load_config=None,

604 + attention_backend=None,

605 + kv_cache_dtype=None,

606 + ),

607 + model_config=SimpleNamespace(

608 + get_total_num_hidden_layers=lambda: 0,

609 + get_vocab_size=lambda: 64,

610 + ),

611 + attention_config=SimpleNamespace(),

612 + cache_config=SimpleNamespace(),

613 + load_config=LoadConfig(),

614 + )

615 + with set_current_vllm_config(

616 + VllmConfig(

617 + device_config=DeviceConfig("cpu"),

618 + compilation_config=CompilationConfig(mode=0),

619 + )

620 + ):

621 + model = LiLiCorrForCausalLM(vllm_config=vllm_config)

622 + return model, vllm_config

623 +

624 + return build

625 +

626 +

627 +@pytest.mark.parametrize("missing", [None, "weight", "weight_scale", "weight_scale_2"])

628 +def test_owned_nvfp4_lm_head_loads_required_checkpoint_tensors(

629 + owned_head_model, missing

630 +):

631 + model, _ = owned_head_model()

632 + assert model.lm_head.weight.dtype == torch.uint8

633 + assert model.lm_head.weight.shape == (64, 8)

634 + weights = {

635 + name: torch.ones_like(value) for name, value in model.state_dict().items()

636 + }

637 + del weights["lm_head.input_scale"] # Deprecated and unused for W4A16.

638 + if missing:

639 + del weights[f"lm_head.{missing}"]

640 + with pytest.raises(ValueError, match=f"lm_head.{missing}"):

641 + model.load_weights(weights.items())

642 + else:

643 + model.load_weights(weights.items())

644 + assert model.has_own_lm_head

645 + for name in ("weight", "weight_scale", "weight_scale_2"):

646 + torch.testing.assert_close(

647 + getattr(model.lm_head, name), weights[f"lm_head.{name}"]

648 + )

649 +

650 +

651 +@pytest.mark.parametrize("owned", [False, True])

652 +def test_selected_lm_head_survives_dflash_and_lilicorr_sharing(

653 + monkeypatch, owned_head_model, owned

654 +):

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

656 + VocabParallelEmbedding,

657 + )

658 + from vllm.v1.worker.gpu.spec_decode.dflash import utils as dflash_utils

659 + from vllm.v1.worker.gpu.spec_decode.eagle import utils as eagle_utils

660 +

661 + model, config = owned_head_model(owned)

662 + own_head = model.lm_head

663 + target_head = owned_head_model(owned)[0].lm_head

664 + # Identical packed bytes do not imply identical quantized heads.

665 + with torch.no_grad():

666 + own_head.weight.fill_(1)

667 + target_head.weight.fill_(1)

668 + if owned:

669 + own_head.weight_scale_2.fill_(1)

670 + target_head.weight_scale_2.fill_(2)

671 + embedding = VocabParallelEmbedding(64, 16)

672 + target = SimpleNamespace(

673 + model=SimpleNamespace(embed_tokens=embedding), lm_head=target_head

674 + )

675 + monkeypatch.setattr(

676 + eagle_utils, "get_pp_group", lambda: SimpleNamespace(world_size=1)

677 + )

678 + monkeypatch.setattr(

679 + dflash_utils, "replace", lambda obj, **kw: SimpleNamespace(**(vars(obj) | kw))

680 + )

681 + monkeypatch.setattr(dflash_utils, "get_pp_safe_draft_load_config", lambda c: c)

682 + monkeypatch.setattr(dflash_utils, "get_model", lambda **kw: model)

683 + if owned:

684 +

685 + def no_weight_comparison(*args):

686 + raise AssertionError("Explicit ownership must bypass raw-weight comparison")

687 +

688 + monkeypatch.setattr(dflash_utils, "_should_share", no_weight_comparison)

689 + spec = LiLiCorrSpeculator.__new__(LiLiCorrSpeculator)

690 + spec.vllm_config = config

691 + assert spec.load_draft_model(target, set()) is model

692 + assert model.lm_head is (own_head if owned else target_head)

693 + assert spec.target_embeddings is embedding

694 + if owned:

695 + target.lm_head = None

696 + assert spec.load_draft_model(target, set()).lm_head is own_head

697 +

698 +

699 +def test_quantized_head_calls_methods_without_reading_packed_weights(monkeypatch):

700 + from vllm.model_executor.layers import linear

701 + from vllm.model_executor.layers.linear import UnquantizedLinearMethod

702 +

703 + calls = []

704 + configured = []

705 + quant_config = object()

706 +

707 + class PackedMethod(UnquantizedLinearMethod):

708 + def apply(self, layer, x, bias=None):

709 + calls.append(layer.prefix)

710 + # Packed storage has a different name/dtype from the activations.

711 + assert x.dtype == torch.float32

712 + return F.linear(x, layer.packed_weight.float() * layer.weight_scale, bias)

713 +

714 + def resolve(config, layer, prefix):

715 + assert config is quant_config

716 + configured.append(prefix)

717 + return PackedMethod()

718 +

719 + monkeypatch.setattr(linear, "resolve_quant_method", resolve)

720 + with set_current_vllm_config(

721 + VllmConfig(

722 + device_config=DeviceConfig("cpu"),

723 + compilation_config=CompilationConfig(mode=0),

724 + )

725 + ):

726 + head = LiLiCorrHead(

727 + model_hidden_size=16,

728 + block_size=4,

729 + rms_norm_eps=1e-6,

730 + config=_config(hidden_size=8),

731 + quant_config=quant_config,

732 + prefix="model.lilicorr",

733 + )

734 + with torch.no_grad():

735 + for module in head.modules():

736 + if isinstance(module, LinearBase):

737 + if module.quant_config is None:

738 + module.weight.normal_(std=0.1)

739 + module.bias.zero_()

740 + continue

741 + module.register_parameter(

742 + "packed_weight",

743 + nn.Parameter(

744 + torch.randint(-4, 5, module.weight.shape, dtype=torch.int8),

745 + requires_grad=False,

746 + ),

747 + )

748 + module.register_parameter(

749 + "weight_scale",

750 + nn.Parameter(

751 + torch.tensor(0.05),

752 + requires_grad=False,

753 + ),

754 + )

755 + del module.weight

756 + module.bias.zero_()

757 + for layer in head.layers:

758 + layer.attn.in_proj_weight.normal_(std=0.1)

759 + layer.attn.in_proj_bias.zero_()

760 + head.materialize_inference_buffers(torch.device("cpu"), torch.float32)

761 + result = head(

762 + torch.randn(2, 3, 4, 16),

763 + torch.randn(2, 3, 4).log_softmax(-1),

764 + torch.randn(2, 3, 16),

765 + torch.randn(2, 16),

766 + torch.tensor([True, False]),

767 + )

768 + assert result.shape == (2, 3, 4, 4)

769 + assert torch.isfinite(result).all()

770 + for name in ("factor_input_proj", "out_head", "in_head"):

771 + assert getattr(head, name).quant_config is None

772 + assert f"model.lilicorr.{name}" not in configured

773 + assert set(calls) == set(configured)

774 + assert len(configured) == len(set(configured))

775 + for name, module in head.named_modules():

776 + if isinstance(module, LinearBase):

777 + assert module.prefix == f"model.lilicorr.{name}"

vllm/config/vllm.py

+10 -6

@@ -760,8 +760,8 @@ class VllmConfig:

760

761 return True

762

763 - def _is_dflash2_draft(self) -> bool:

764 - """Whether the DFlash draft is a DFlash2 one, by the architecture the

763 + def _is_dflash_candidate_draft(self) -> bool:

764 + """Whether the DFlash draft has a candidate head, by the architecture the

765 speculator selects on (v1/worker/gpu/spec_decode/__init__.py)."""

766 spec = self.speculative_config

767 if spec is None or spec.method != "dflash":

@@ -769,7 +769,11 @@ class VllmConfig:

769 draft_config = getattr(spec, "draft_model_config", None)

770 if draft_config is None:

771 return False

772 - return "DFlash2DraftModel" in (draft_config.architectures or [])

772 + return bool(

773 + {"DFlash2DraftModel", "LiLiCorrDraftModel"}.intersection(

774 + draft_config.architectures or []

775 + )

776 + )

777

778 def _dflash_needs_multi_kv_group(self) -> bool:

779 """Whether a DFlash draft mixes sliding-window and full attention."""

@@ -3097,11 +3101,11 @@ class VllmConfig:

3101 if self._dflash_needs_multi_kv_group():

3102 unsupported.append("mixed sliding/full dflash drafts")

3103

3100 - # The DFlash2 candidate selector exists only in the V2 speculator. On

3104 + # DFlash candidate heads exist only in the V2 speculator. On

3105 # V1 the same checkpoint drafts through DFlashProposer, which never

3106 # calls it, so the draft would degrade to DFlash1 silently.

3103 - if self._is_dflash2_draft():

3104 - unsupported.append("dflash2 drafts")

3107 + if self._is_dflash_candidate_draft():

3108 + unsupported.append("DFlash candidate-head drafts")

3109

3110 if self.model_config is not None and self.model_config.is_diffusion:

3111 unsupported.append("diffusion models")

vllm/model_executor/layers/logits_processor.py

+21 -1

@@ -256,6 +256,8 @@ class LogitsProcessor(PluggableLayer):

256 hidden_states: torch.Tensor,

257 k: int,

258 embedding_bias: torch.Tensor | None = None,

259 + *,

260 + return_log_probs: bool = False,

261 ) -> tuple[torch.Tensor, torch.Tensor]:

262 """Vocab-parallel top-k without all-gathering full logits.

263

@@ -263,7 +265,9 @@ class LogitsProcessor(PluggableLayer):

265 the values as well as the global ids. Communication is

266 O(batch * 2k * tp_size) rather than O(batch * vocab_size).

267

266 - Scale and soft cap are applied to the k selected values rather than

268 + With return_log_probs, values use the full-vocabulary log-partition.

269 +

270 + Otherwise scale and soft cap are applied to the k selected values rather than

271 the whole vocabulary; both are monotonic, so the selection is the same

272 and only k entries are touched.

273 """

@@ -274,12 +278,26 @@ class LogitsProcessor(PluggableLayer):

278 )

279

280 logits = self._apply_head(lm_head, hidden_states, embedding_bias)

281 + if return_log_probs:

282 + if self.soft_cap is not None:

283 + logits = torch.tanh(logits / self.soft_cap) * self.soft_cap

284 + if self.scale != 1.0:

285 + logits = logits * self.scale

286

287 # Mask out padding entries beyond org_vocab_size on this shard.

288 num_pad = lm_head.shard_indices.num_org_vocab_padding

289 if num_pad > 0:

290 logits[..., -num_pad:] = -float("inf")

291

292 + if return_log_probs:

293 + log_partition = torch.logsumexp(logits.float(), dim=-1, keepdim=True)

294 + if lm_head.tp_size > 1:

295 + log_partition = torch.logsumexp(

296 + tensor_model_parallel_all_gather(log_partition, dim=-1),

297 + dim=-1,

298 + keepdim=True,

299 + )

300 +

301 values, ids = _topk(logits, k)

302 # Convert shard-local indices to global vocab indices.

303 ids = ids.to(torch.int64) + lm_head.shard_indices.org_vocab_start_index

@@ -291,6 +309,8 @@ class LogitsProcessor(PluggableLayer):

309 ids = ids.gather(-1, selected)

310

311 values = values.float()

312 + if return_log_probs:

313 + return ids, values - log_partition

314 if self.scale != 1.0:

315 values = values * self.scale

316 if self.soft_cap is not None:

vllm/model_executor/models/lilicorr.py new

+546

@@ -0,0 +1,546 @@

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)

vllm/model_executor/models/nano_nemotron_vl.py

@@ -34,6 +34,7 @@ from vllm.model_executor.models.interfaces import (

34 HasInnerState,

35 IsHybrid,

36 MultiModalEmbeddings,

37 + SupportsEagle3,

38 SupportsLoRA,

39 SupportsMultiModal,

40 SupportsMultiModalPruning,

@@ -928,6 +929,7 @@ class NemotronH_Nano_VL_V2(

929 nn.Module,

930 HasInnerState,

931 IsHybrid,

932 + SupportsEagle3,

933 SupportsMultiModal,

934 SupportsMultiModalPruning,

935 SupportsLoRA,

vllm/model_executor/models/nemotron_h.py

+7 -4

@@ -34,7 +34,7 @@ from vllm.config.parallel import ParallelConfig

34 from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size

35 from vllm.distributed.communication_op import tensor_model_parallel_all_gather

36 from vllm.distributed.parallel_state import get_pp_group

37 -from vllm.model_executor.layers.activation import ReLUSquaredActivation

37 +from vllm.model_executor.layers.activation import get_act_fn

38 from vllm.model_executor.layers.attention import Attention

39 from vllm.model_executor.layers.fused_moe import (

40 FusedMoEFactory,

@@ -93,7 +93,7 @@ from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig

93 class NemotronHMLP(nn.Module):

94 def __init__(

95 self,

96 - config: NemotronHConfig,

96 + config: NemotronHConfig | None,

97 hidden_size: int,

98 intermediate_size: int,

99 quant_config: QuantizationConfig | None = None,

@@ -101,6 +101,9 @@ class NemotronHMLP(nn.Module):

101 reduce_results: bool = True,

102 is_sequence_parallel: bool = False,

103 prefix: str = "",

104 + *,

105 + hidden_act: str = "relu2",

106 + output_size: int | None = None,

107 ) -> None:

108 super().__init__()

109

@@ -114,14 +117,14 @@ class NemotronHMLP(nn.Module):

117 )

118 self.down_proj = RowParallelLinear(

119 input_size=intermediate_size,

117 - output_size=hidden_size,

120 + output_size=hidden_size if output_size is None else output_size,

121 bias=bias,

122 quant_config=quant_config,

123 reduce_results=reduce_results,

124 disable_tp=is_sequence_parallel,

125 prefix=f"{prefix}.down_proj",

126 )

124 - self.act_fn = ReLUSquaredActivation()

127 + self.act_fn = get_act_fn(hidden_act)

128

129 def forward(self, x: torch.Tensor):

130 x, _ = self.up_proj(x)

vllm/model_executor/models/qwen3_dflash.py

+70 -15

@@ -4,6 +4,7 @@

4 import io

5 from collections.abc import Iterable

7 +import regex as re

8 import torch

9 import torch.nn.functional as F

10 from torch import nn

@@ -23,6 +24,7 @@ from vllm.model_executor.layers.linear import (

24 QKVParallelLinear,

25 ReplicatedLinear,

26 RowParallelLinear,

27 + UnquantizedLinearMethod,

28 )

29 from vllm.model_executor.layers.logits_processor import LogitsProcessor

30 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

@@ -54,6 +56,31 @@ logger = init_logger(__name__)

58 _SLIDING_ATTENTION = "sliding_attention"

59 +_DRAFT_LAYER_PATTERN = re.compile(r"(?<![A-Za-z0-9_.])layers\.(\d+)")

60 +

61 +

62 +def _add_global_draft_layer_exclusions(

63 + quant_config: QuantizationConfig | None,

64 + start_layer_id: int,

65 + num_hidden_layers: int,

66 +) -> None:

67 + """Add runtime layer aliases for checkpoint-local quant exclusions."""

68 + if quant_config is None or start_layer_id == 0:

69 + return

70 + exclusions = getattr(quant_config, "exclude_modules", None)

71 + if not isinstance(exclusions, list):

72 + return

73 +

74 + def offset_local_layer(match: re.Match[str]) -> str:

75 + layer_idx = int(match.group(1))

76 + if layer_idx >= num_hidden_layers:

77 + return match.group(0)

78 + return f"layers.{layer_idx + start_layer_id}"

79 +

80 + for exclusion in tuple(exclusions):

81 + global_exclusion = _DRAFT_LAYER_PATTERN.sub(offset_local_layer, exclusion)

82 + if global_exclusion != exclusion and global_exclusion not in exclusions:

83 + exclusions.append(global_exclusion)

86 def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool:

@@ -394,6 +421,9 @@ class DFlashQwen3Model(nn.Module):

421 self.config = vllm_config.speculative_config.draft_model_config.hf_config

422 self.vocab_size = self.config.vocab_size

423 self.quant_config = get_draft_quant_config(vllm_config)

424 + _add_global_draft_layer_exclusions(

425 + self.quant_config, start_layer_id, self.config.num_hidden_layers

426 + )

427

428 drafter_config = getattr(self.config, "eagle_config", {})

429 drafter_config.update(getattr(self.config, "dflash_config", {}))

@@ -473,14 +503,25 @@ class DFlashQwen3Model(nn.Module):

503 has_bias: bool,

504 ) -> None:

505 self._hidden_norm_weight = self.hidden_norm.weight.data

506 + self._context_qkv_projs = [a.qkv_proj for a in layers_attn]

507 + self._context_q_sizes = [a.q_size for a in layers_attn]

508

477 - # KV projection weights: [num_layers * 2 * kv_size, hidden_size]

478 - kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]

479 - self._fused_kv_weight = torch.cat(kv_weights, dim=0)

480 - if has_bias:

481 - kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn]

482 - self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0)

509 + if all(

510 + isinstance(proj.quant_method, UnquantizedLinearMethod)

511 + for proj in self._context_qkv_projs

512 + ):

513 + # KV projection weights: [num_layers * 2 * kv_size, hidden_size]

514 + kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]

515 + self._fused_kv_weight: torch.Tensor | None = torch.cat(kv_weights, dim=0)

516 + if has_bias:

517 + kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn]

518 + self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0)

519 + else:

520 + self._fused_kv_bias = None

521 else:

522 + # Quantized linear weights may use packed storage that cannot be

523 + # consumed by F.linear. Run each projection through its quant method.

524 + self._fused_kv_weight = None

525 self._fused_kv_bias = None

526

527 # K-norm weights stacked into one contiguous [num_layers, head_dim]

@@ -548,17 +589,26 @@ class DFlashQwen3Model(nn.Module):

589 self._hidden_norm_weight,

590 self._rms_norm_eps,

591 )

551 - all_kv_flat = F.linear(

552 - normed_context_states, self._fused_kv_weight, self._fused_kv_bias

553 - )

592 + if self._fused_kv_weight is not None:

593 + all_kv_flat = F.linear(

594 + normed_context_states, self._fused_kv_weight, self._fused_kv_bias

595 + )

596 + all_kv = all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim)

597 + else:

598 + layer_kv = []

599 + for proj, q_size in zip(

600 + self._context_qkv_projs, self._context_q_sizes, strict=True

601 + ):

602 + qkv, _ = proj(normed_context_states)

603 + layer_kv.append(

604 + qkv[:, q_size:].view(num_ctx, 2, num_kv_heads, head_dim)

605 + )

606 + all_kv = torch.stack(layer_kv, dim=1)

607 +

608 # Single contiguous copy that separates K/V and transposes to

555 - # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous.

609 + # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous.

610 # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.

557 - all_kv = (

558 - all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim)

559 - .permute(2, 1, 0, 3, 4)

560 - .contiguous()

561 - )

611 + all_kv = all_kv.permute(2, 1, 0, 3, 4).contiguous()

612 all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous

613 all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous

614 return all_k, all_v

@@ -724,6 +774,11 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):

774 self.lm_head = ParallelLMHead(

775 self.config.draft_vocab_size,

776 self.config.hidden_size,

777 + quant_config=(

778 + get_draft_quant_config(vllm_config)

779 + if getattr(self.config, "has_own_lm_head", False)

780 + else None

781 + ),

782 prefix=maybe_prefix(prefix, "lm_head"),

783 )

784 self.logits_processor = LogitsProcessor(

vllm/model_executor/models/qwen3_dflash2.py

+3 -1

@@ -176,6 +176,7 @@ class DFlashGroupedConv(nn.Module):

176 block_size: int,

177 params_dtype: torch.dtype,

178 prefix: str,

179 + quant_config: QuantizationConfig | None = None,

180 ) -> None:

181 super().__init__()

182 if hidden_size % group_size:

@@ -195,7 +196,7 @@ class DFlashGroupedConv(nn.Module):

196 2 * taps * self.num_groups,

197 bias=False,

198 params_dtype=params_dtype,

198 - quant_config=None,

199 + quant_config=quant_config,

200 prefix=maybe_prefix(prefix, "kernel_projection"),

201 return_bias=False,

202 )

@@ -254,6 +255,7 @@ class DFlash2Qwen3DecoderLayer(DFlashQwen3DecoderLayer):

255 # Query tokens per request: the bonus token plus the mask tokens.

256 block_size=1 + speculative_config.num_speculative_tokens,

257 params_dtype=vllm_config.model_config.dtype,

258 + quant_config=quant_config,

259 )

260 self.attention_conv = DFlashGroupedConv(

261 **conv_args, prefix=maybe_prefix(prefix, "attention_conv")

vllm/model_executor/models/registry.py

@@ -633,6 +633,7 @@ _SPECULATIVE_DECODING_MODELS = {

633 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),

634 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),

635 "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),

636 + "LiLiCorrDraftModel": ("lilicorr", "LiLiCorrForCausalLM"),

637 "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"),

638 # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash

639 # implementation. EAGLEConfig rewrites a dflash draft's architecture to

vllm/transformers_utils/configs/eagle.py

+3 -1

@@ -68,7 +68,9 @@ class EAGLEConfig(PreTrainedConfig):

68 )

69 kwargs["architectures"] = [

70 arch

71 - if arch.startswith("DFlash") or arch.endswith("DFlash")

71 + if arch.startswith("DFlash")

72 + or arch.endswith("DFlash")

73 + or arch == "LiLiCorrDraftModel"

74 else f"DFlash{arch}"

75 for arch in self.model.architectures

76 ]

vllm/v1/worker/gpu/spec_decode/__init__.py

@@ -15,6 +15,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):

16 return ExtractHiddenStatesSpeculator(vllm_config, device)

17 elif speculative_config.method == "dflash":

18 + if "LiLiCorrDraftModel" in speculative_config.draft_model_config.architectures:

19 + from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import (

20 + LiLiCorrSpeculator,

21 + )

22 +

23 + return LiLiCorrSpeculator(vllm_config, device)

24 if "DFlash2DraftModel" in speculative_config.draft_model_config.architectures:

25 from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import (

26 DFlash2Speculator,

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -322,6 +322,11 @@ class DFlashSpeculator(DraftModelSpeculator):

322 context_slots,

323 )

324

325 + def prepare_context_anchor(

326 + self, input_batch: InputBatch, num_rejected: torch.Tensor

327 + ) -> None:

328 + """Publish context features required by a draft's candidate head."""

329 +

330 @torch.inference_mode()

331 def propose(

332 self,

@@ -369,6 +374,7 @@ class DFlashSpeculator(DraftModelSpeculator):

374 else:

375 hidden_states = last_hidden_states

376 self.hidden_states[:num_target_tokens].copy_(hidden_states[:num_target_tokens])

377 + self.prepare_context_anchor(input_batch, num_rejected)

378

379 if dummy_run and skip_attn_for_dummy_run:

380 # Memory profiling path: block_tables / kv_cache_config are not initialized.

vllm/v1/worker/gpu/spec_decode/dflash/utils.py

+6 -2

@@ -61,8 +61,12 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo

62 target_lm_head = get_target_lm_head(target_model, target_language_model)

63 draft_lm_head = getattr(dflash_model, "lm_head", None)

64 - if target_lm_head is not None and _should_share(

65 - dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head

64 + if (

65 + target_lm_head is not None

66 + and not getattr(dflash_model.config, "has_own_lm_head", False)

67 + and _should_share(

68 + dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head

69 + )

70 ):

71 if draft_lm_head is not None:

72 del dflash_model.lm_head

vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py

+97 -55

@@ -108,76 +108,101 @@ def _cache_draft_logits_kernel(

108 tl.store(cached_candidate_ptr + cache_base + offsets, token_ids, mask=mask)

109

110

111 -class DFlash2Speculator(DFlashSpeculator):

112 - _speculator_name = "DFlash2"

111 +class CandidateSampler:

112 + """The shared DFlash2/LiLiCorr candidate walk and realized proposal cache."""

113

114 - def __init__(self, vllm_config: VllmConfig, device: torch.device):

115 - super().__init__(vllm_config, device)

116 - draft_config = self.draft_model_config.hf_config.dflash_config

117 - self.selector_top_k = int(draft_config["selector_top_k"])

118 - self._anchor_indices = (

119 - torch.arange(self.max_num_reqs, dtype=torch.int64, device=device)

120 - * self.num_query_per_req

121 - )

122 - self._selector_scores = torch.empty(

123 - self.max_num_reqs,

124 - self.num_speculative_steps,

125 - self.selector_top_k,

126 - dtype=torch.float32,

127 - device=device,

114 + def __init__(

115 + self, max_num_reqs: int, num_steps: int, top_k: int, device: torch.device

116 + ):

117 + self.num_steps = num_steps

118 + self.top_k = top_k

119 + self.scores = torch.empty(

120 + max_num_reqs, num_steps, top_k, dtype=torch.float32, device=device

121 )

129 - self._cached_candidate_ids = torch.zeros(

130 - self._selector_scores.shape, dtype=torch.int64, device=device

122 + self.cached_candidate_ids = torch.zeros(

123 + self.scores.shape, dtype=torch.int64, device=device

124 )

125

133 - def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]:

134 - # fp32 so the walk and the rejection that checks it read the same

135 - # distribution; -inf because the cache kernel writes only the K

136 - # candidates.

137 - return torch.float32, -float("inf")

138 -

139 - def _sample_path(

126 + def sample(

127 self,

128 candidate_ids: torch.Tensor,

129 scores: torch.Tensor,

130 num_reqs: int,

131 + sample_pos: torch.Tensor,

132 + idx_mapping: torch.Tensor,

133 + temperature: torch.Tensor,

134 + seeds: torch.Tensor,

135 + draft_tokens: torch.Tensor,

136 + draft_logits: torch.Tensor | None,

137 + use_fp64: bool,

138 ) -> None:

145 - block_k = triton.next_power_of_2(self.selector_top_k)

139 + candidate_ids = candidate_ids.contiguous()

140 + sample_pos = sample_pos.contiguous()

141 + idx_mapping = idx_mapping.contiguous()

142 + block_k = triton.next_power_of_2(self.top_k)

143 _selector_walk_kernel[(num_reqs,)](

144 scores.contiguous(),

145 candidate_ids.contiguous(),

149 - self.sample_pos,

150 - self.sample_idx_mapping,

151 - self.temperature,

152 - self.seeds,

153 - self.draft_tokens,

154 - self._selector_scores,

155 - num_steps=self.num_speculative_steps,

156 - top_k=self.selector_top_k,

146 + sample_pos,

147 + idx_mapping,

148 + temperature,

149 + seeds,

150 + draft_tokens,

151 + self.scores,

152 + num_steps=self.num_steps,

153 + top_k=self.top_k,

154 BLOCK_K=block_k,

158 - SAMPLE_PROBABILISTIC=self.draft_logits is not None,

159 - USE_FP64=self.use_fp64_gumbel,

155 + SAMPLE_PROBABILISTIC=draft_logits is not None,

156 + USE_FP64=use_fp64,

157 num_warps=1,

158 )

159

163 - def _cache_draft_logits(self, candidate_ids: torch.Tensor, num_sample: int) -> None:

164 - draft_logits = self.draft_logits

165 - assert draft_logits is not None

166 - block_k = triton.next_power_of_2(self.selector_top_k)

160 + if draft_logits is not None:

161 + self._cache_draft_logits(

162 + candidate_ids, num_reqs * self.num_steps, idx_mapping, draft_logits

163 + )

164 +

165 + def _cache_draft_logits(

166 + self,

167 + candidate_ids: torch.Tensor,

168 + num_sample: int,

169 + idx_mapping: torch.Tensor,

170 + draft_logits: torch.Tensor,

171 + ) -> None:

172 + block_k = triton.next_power_of_2(self.top_k)

173 _cache_draft_logits_kernel[(num_sample,)](

174 draft_logits,

169 - self._cached_candidate_ids,

175 + self.cached_candidate_ids,

176 candidate_ids,

171 - self._selector_scores,

172 - self.sample_idx_mapping,

177 + self.scores,

178 + idx_mapping,

179 draft_logits.stride(0),

180 draft_logits.stride(1),

175 - num_steps=self.num_speculative_steps,

176 - top_k=self.selector_top_k,

181 + num_steps=self.num_steps,

182 + top_k=self.top_k,

183 BLOCK_K=block_k,

184 num_warps=1,

185 )

186

187 +

188 +class DFlash2Speculator(DFlashSpeculator):

189 + _speculator_name = "DFlash2"

190 + _candidate_top_k_key = "selector_top_k"

191 +

192 + def __init__(self, vllm_config: VllmConfig, device: torch.device):

193 + super().__init__(vllm_config, device)

194 + draft_config = self.draft_model_config.hf_config.dflash_config

195 + self.top_k = int(draft_config[self._candidate_top_k_key])

196 + self.candidate_sampler = CandidateSampler(

197 + self.max_num_reqs, self.num_speculative_steps, self.top_k, device

198 + )

199 +

200 + def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]:

201 + # fp32 so the walk and the rejection that checks it read the same

202 + # distribution; -inf because the cache kernel writes only the K

203 + # candidates.

204 + return torch.float32, -float("inf")

205 +

206 def _generate_draft(

207 self,

208 num_reqs: int,

@@ -202,22 +227,39 @@ class DFlash2Speculator(DFlashSpeculator):

227 hidden_states.flatten(0, 1)

228 )

229 candidate_ids = candidate_ids.view(

205 - num_reqs, self.num_speculative_steps, self.selector_top_k

230 + num_reqs, self.num_speculative_steps, self.top_k

231 )

232 unary_logits = unary_logits.view_as(candidate_ids)

208 - anchor_token_ids = self.input_buffers.input_ids[self._anchor_indices[:num_reqs]]

209 - scores = self.model.model.candidate_selector(

233 + scores = self._score_candidates(candidate_ids, unary_logits, hidden_states)

234 + self.candidate_sampler.sample(

235 candidate_ids,

211 - unary_logits,

212 - hidden_states,

213 - anchor_token_ids,

236 + scores,

237 + num_reqs,

238 + self.sample_pos,

239 + self.sample_idx_mapping,

240 + self.temperature,

241 + self.seeds,

242 + self.draft_tokens,

243 + self.draft_logits,

244 + self.use_fp64_gumbel,

245 )

215 - self._sample_path(candidate_ids, scores, num_reqs)

246 if self.enable_adaptive_verification:

247 self._maybe_predict_acceptance(

218 - self._selector_scores[:num_reqs].flatten(0, 1),

248 + self.candidate_sampler.scores[:num_reqs].flatten(0, 1),

249 self.sample_idx_mapping[:num_sample],

250 self.sample_col[:num_sample],

251 )

222 - if self.draft_logits is not None:

223 - self._cache_draft_logits(candidate_ids, num_sample)

252 +

253 + def _score_candidates(

254 + self,

255 + candidate_ids: torch.Tensor,

256 + unary_logits: torch.Tensor,

257 + hidden_states: torch.Tensor,

258 + ) -> torch.Tensor:

259 + num_reqs = candidate_ids.shape[0]

260 + anchor_token_ids = self.input_buffers.input_ids[

261 + : num_reqs * self.num_query_per_req : self.num_query_per_req

262 + ]

263 + return self.model.model.candidate_selector(

264 + candidate_ids, unary_logits, hidden_states, anchor_token_ids

265 + )

vllm/v1/worker/gpu/spec_decode/lilicorr/__init__.py new

@@ -0,0 +1,2 @@

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

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

vllm/v1/worker/gpu/spec_decode/lilicorr/speculator.py new

+97

@@ -0,0 +1,97 @@

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

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

3 +

4 +import torch

5 +from torch import nn

6 +

7 +from vllm.config import VllmConfig

8 +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

9 +from vllm.v1.worker.gpu.input_batch import InputBatch

10 +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator

11 +from vllm.v1.worker.gpu.spec_decode.eagle.utils import get_target_lm_head

12 +

13 +

14 +class LiLiCorrSpeculator(DFlash2Speculator):

15 + _speculator_name = "LiLiCorr"

16 + _candidate_top_k_key = "lilicorr_candidate_topk"

17 +

18 + def __init__(self, vllm_config: VllmConfig, device: torch.device):

19 + super().__init__(vllm_config, device)

20 + config = self.draft_model_config.hf_config

21 + self.anchor_hidden = torch.zeros(

22 + self.max_num_reqs, config.hidden_size, dtype=self.dtype, device=device

23 + )

24 + self.anchor_valid = torch.zeros(

25 + self.max_num_reqs, dtype=torch.bool, device=device

26 + )

27 +

28 + def load_draft_model(

29 + self, target_model: nn.Module, target_attn_layer_names: set[str]

30 + ) -> nn.Module:

31 + model = super().load_draft_model(target_model, target_attn_layer_names)

32 + language_model = (

33 + target_model.get_language_model()

34 + if hasattr(target_model, "get_language_model")

35 + else target_model

36 + )

37 + inner = getattr(language_model, "model", language_model)

38 + embedding = getattr(inner, "embed_tokens", None) or getattr(

39 + inner, "embedding", None

40 + )

41 + lm_head = (

42 + model.lm_head

43 + if model.has_own_lm_head

44 + else get_target_lm_head(target_model, language_model)

45 + )

46 + if not isinstance(embedding, VocabParallelEmbedding) or not isinstance(

47 + lm_head, VocabParallelEmbedding

48 + ):

49 + raise ValueError(

50 + "LiLiCorr requires the target input embedding and its selected "

51 + "LM head on the draft rank."

52 + )

53 + if embedding.embedding_dim != model.config.hidden_size:

54 + raise ValueError(

55 + "LiLiCorr target embedding width must match its trained "

56 + "token projection input."

57 + )

58 + if (

59 + embedding.shard_indices.num_added_elements

60 + or lm_head.shard_indices.num_added_elements

61 + ):

62 + raise ValueError("LiLiCorr does not support added vocabulary entries.")

63 + # Candidate embeddings always come from the target; the checkpoint may

64 + # supply its own candidate LM head.

65 + self.target_embeddings = embedding

66 + model.lm_head = lm_head

67 + return model

68 +

69 + def prepare_context_anchor(

70 + self, input_batch: InputBatch, num_rejected: torch.Tensor

71 + ) -> None:

72 + num_reqs = input_batch.num_reqs

73 + starts = input_batch.query_start_loc[:num_reqs]

74 + ends = input_batch.query_start_loc[1 : num_reqs + 1] - num_rejected[:num_reqs]

75 + valid = ends > starts

76 + indices = (ends - 1).clamp_min(0).long()

77 + self.anchor_hidden.zero_()

78 + self.anchor_valid.zero_()

79 + self.anchor_hidden[:num_reqs].copy_(

80 + self.model.model.hidden_norm(self.hidden_states[indices]) * valid[:, None]

81 + )

82 + self.anchor_valid[:num_reqs].copy_(valid)

83 +

84 + def _score_candidates(

85 + self,

86 + candidate_ids: torch.Tensor,

87 + unary_logits: torch.Tensor,

88 + hidden_states: torch.Tensor,

89 + ) -> torch.Tensor:

90 + num_reqs = candidate_ids.shape[0]

91 + return self.model.model.lilicorr(

92 + self.target_embeddings(candidate_ids),

93 + unary_logits,

94 + hidden_states,

95 + self.anchor_hidden[:num_reqs],

96 + self.anchor_valid[:num_reqs],

97 + )