| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 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 |
| 15 | from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator |
| 16 | |
| 17 | |
| 18 | @pytest.mark.parametrize("block_size", [5, 8]) |
| 19 | def test_grouped_conv_matches_reference(block_size: int): |
| 20 | torch.manual_seed(0) |
| 21 | batch, taps, num_groups, group_size = 3, 3, 4, 2 |
| 22 | hidden = torch.randn(batch * block_size, num_groups * group_size) |
| 23 | delta = torch.randn(batch * block_size, taps, num_groups) |
| 24 | base = torch.randn(taps, num_groups * group_size) |
| 25 | |
| 26 | actual = _grouped_conv( |
| 27 | hidden, delta, base, block_size, num_groups, group_size, taps |
| 28 | ) |
| 29 | hidden_blocks = hidden.view(batch, block_size, num_groups, group_size) |
| 30 | expected = torch.zeros_like(hidden_blocks) |
| 31 | base = base.view(taps, num_groups, group_size) |
| 32 | delta = delta.view(batch, block_size, taps, num_groups) |
| 33 | for position in range(block_size): |
| 34 | for tap in range(min(taps, position + 1)): |
| 35 | expected[:, position] += ( |
| 36 | base[tap] + delta[:, position, tap, :, None] |
| 37 | ) * hidden_blocks[:, position - tap] |
| 38 | |
| 39 | torch.testing.assert_close(actual, expected.flatten(0, 1).flatten(-2)) |
| 40 | |
| 41 | |
| 42 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") |
| 43 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 44 | @pytest.mark.parametrize( |
| 45 | "batch,num_groups,group_size,block_size,taps", |
| 46 | [ |
| 47 | (7, 13, 11, 5, 1), |
| 48 | (7, 13, 11, 5, 3), |
| 49 | (7, 13, 11, 8, 2), |
| 50 | (16, 64, 16, 8, 2), |
| 51 | (16, 160, 16, 8, 2), |
| 52 | ], |
| 53 | ) |
| 54 | def test_grouped_conv_triton_matches_reference( |
| 55 | dtype: torch.dtype, |
| 56 | batch: int, |
| 57 | num_groups: int, |
| 58 | group_size: int, |
| 59 | block_size: int, |
| 60 | taps: int, |
| 61 | ): |
| 62 | torch.manual_seed(0) |
| 63 | rows = batch * block_size |
| 64 | hidden = torch.randn(rows, num_groups * group_size, device="cuda", dtype=dtype) |
| 65 | base = torch.randn(taps, num_groups * group_size, device="cuda", dtype=dtype) |
| 66 | projected = torch.randn(rows, 2, taps, num_groups, device="cuda", dtype=dtype) |
| 67 | delta = projected[:, 1] |
| 68 | |
| 69 | actual = _grouped_conv( |
| 70 | hidden, delta, base, block_size, num_groups, group_size, taps |
| 71 | ) |
| 72 | |
| 73 | hidden_blocks = hidden.float().view(batch, block_size, num_groups, group_size) |
| 74 | expected = torch.zeros_like(hidden_blocks) |
| 75 | base_blocks = base.float().view(taps, num_groups, group_size) |
| 76 | delta_blocks = delta.float().view(batch, block_size, taps, num_groups) |
| 77 | for position in range(block_size): |
| 78 | for tap in range(min(taps, position + 1)): |
| 79 | expected[:, position] += ( |
| 80 | base_blocks[tap] + delta_blocks[:, position, tap, :, None] |
| 81 | ) * hidden_blocks[:, position - tap] |
| 82 | |
| 83 | torch.testing.assert_close( |
| 84 | actual, |
| 85 | expected.flatten(0, 1).flatten(-2).to(dtype), |
| 86 | rtol=1e-2 if dtype is torch.bfloat16 else 1e-5, |
| 87 | atol=1e-2 if dtype is torch.bfloat16 else 1e-5, |
| 88 | ) |
| 89 | |
| 90 | |
| 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 |
| 112 | vocab = 17 |
| 113 | predecessors = torch.randn(vocab, rank) |
| 114 | successors = torch.randn(vocab, rank) |
| 115 | candidate_ids = torch.randint(vocab, (batch, steps, top_k)) |
| 116 | unary = torch.randn(batch, steps, top_k) |
| 117 | hidden = torch.randn(batch, steps, rank) |
| 118 | anchors = torch.randint(vocab, (batch,)) |
| 119 | |
| 120 | actual = _score_edges( |
| 121 | predecessors, |
| 122 | successors, |
| 123 | candidate_ids, |
| 124 | unary, |
| 125 | hidden, |
| 126 | anchors, |
| 127 | top_k, |
| 128 | ) |
| 129 | expected = torch.empty_like(actual) |
| 130 | for step in range(steps): |
| 131 | pred = ( |
| 132 | anchors[:, None].expand(-1, top_k) |
| 133 | if step == 0 |
| 134 | else candidate_ids[:, step - 1] |
| 135 | ) |
| 136 | expected[:, step] = unary[:, step, None] + torch.einsum( |
| 137 | "bpr,bcr->bpc", |
| 138 | predecessors[pred] * hidden[:, step, None], |
| 139 | successors[candidate_ids[:, step]], |
| 140 | ) |
| 141 | |
| 142 | torch.testing.assert_close(actual, expected) |
| 143 | |
| 144 | |
| 145 | def _stub_base(monkeypatch, draft_logits): |
| 146 | """A DFlashSpeculator.__init__ that allocates only what the base class would. |
| 147 | |
| 148 | The real base class fills draft_logits from draft_logits_spec, so callers |
| 149 | pass a tensor already in that state. |
| 150 | """ |
| 151 | |
| 152 | def init_base(self, _vllm_config, device): |
| 153 | self.draft_model_config = SimpleNamespace( |
| 154 | hf_config=SimpleNamespace(dflash_config={"selector_top_k": 3}) |
| 155 | ) |
| 156 | self.max_num_reqs = 2 |
| 157 | self.num_query_per_req = 5 |
| 158 | self.num_speculative_steps = 4 |
| 159 | self.vocab_size = 17 |
| 160 | self.draft_tokens = torch.empty((2, 4), dtype=torch.int64, device=device) |
| 161 | self.draft_logits = draft_logits |
| 162 | |
| 163 | monkeypatch.setattr(DFlashSpeculator, "__init__", init_base) |
| 164 | |
| 165 | |
| 166 | def test_selector_leaves_greedy_drafting_without_proposal_logits(monkeypatch): |
| 167 | """Greedy is the default, and it caches no proposal distribution. |
| 168 | |
| 169 | The base class allocates draft_logits only for "probabilistic"; verification |
| 170 | reads `draft_logits is None` to decide whether a distribution is on offer, so |
| 171 | allocating one here would claim a proposal the walk never sampled from. |
| 172 | """ |
| 173 | _stub_base(monkeypatch, None) |
| 174 | speculator = DFlash2Speculator(None, torch.device("cpu")) |
| 175 | |
| 176 | assert speculator.draft_logits is None |
| 177 | |
| 178 | |
| 179 | def test_selector_asks_for_fp32_proposal_logits(): |
| 180 | """The spec the base class allocates from: fp32, filled -inf. |
| 181 | |
| 182 | Not the head dtype -- rounding selector scores to bf16 moves the argmax of a |
| 183 | candidate row often enough that the walk and the rejection sampler checking it |
| 184 | would no longer read the same distribution. |
| 185 | """ |
| 186 | dtype, fill = DFlash2Speculator.draft_logits_spec(None, None) |
| 187 | |
| 188 | assert dtype is torch.float32 |
| 189 | assert fill == float("-inf") |
| 190 | |
| 191 | |
| 192 | @pytest.mark.skip_global_cleanup |
| 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, |
| 203 | ) |
| 204 | |
| 205 | # 1. Mock get_current_vllm_config and TP groups |
| 206 | mock_current_vllm_config = SimpleNamespace( |
| 207 | cache_config=SimpleNamespace( |
| 208 | block_size=16, |
| 209 | user_specified_block_size=False, |
| 210 | kv_cache_dtype_skip_layers=[], |
| 211 | cache_dtype="auto", |
| 212 | sliding_window=None, |
| 213 | enable_prefix_caching=False, |
| 214 | ), |
| 215 | kv_transfer_config=None, |
| 216 | speculative_config=None, |
| 217 | attention_config=SimpleNamespace( |
| 218 | use_non_causal=False, |
| 219 | backend=None, |
| 220 | backend_per_kind={}, |
| 221 | ), |
| 222 | parallel_config=SimpleNamespace( |
| 223 | prefill_context_parallel_size=1, |
| 224 | decode_context_parallel_size=1, |
| 225 | ), |
| 226 | compilation_config=SimpleNamespace( |
| 227 | compile_custom_ops=False, |
| 228 | custom_ops="all", |
| 229 | enabled_custom_ops=set(), |
| 230 | static_forward_context={}, |
| 231 | mode=0, # CompilationMode.NONE is 0 |
| 232 | ), |
| 233 | model_config=SimpleNamespace( |
| 234 | dtype=torch.float32, |
| 235 | is_mm_prefix_lm=False, |
| 236 | rswa_window=None, |
| 237 | ), |
| 238 | kernel_config=SimpleNamespace( |
| 239 | linear_backend="auto", |
| 240 | ), |
| 241 | ) |
| 242 | from vllm.platforms import current_platform |
| 243 | |
| 244 | monkeypatch.setattr( |
| 245 | current_platform, |
| 246 | "get_attn_backend_cls", |
| 247 | lambda *args, **kwargs: ( |
| 248 | "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend" |
| 249 | ), |
| 250 | ) |
| 251 | |
| 252 | class MockGroup: |
| 253 | rank_in_group = 0 |
| 254 | world_size = 1 |
| 255 | |
| 256 | monkeypatch.setattr( |
| 257 | "vllm.distributed.parallel_state._TP", |
| 258 | MockGroup(), |
| 259 | ) |
| 260 | |
| 261 | # 2. Mock vllm_config |
| 262 | hf_config = SimpleNamespace( |
| 263 | vocab_size=1000, |
| 264 | hidden_size=256, |
| 265 | num_hidden_layers=2, |
| 266 | num_attention_heads=8, |
| 267 | num_key_value_heads=2, |
| 268 | max_position_embeddings=2048, |
| 269 | rms_norm_eps=1e-6, |
| 270 | rope_parameters={}, |
| 271 | intermediate_size=512, |
| 272 | hidden_act="silu", |
| 273 | dflash_config={ |
| 274 | "selector_rank": 4, |
| 275 | "selector_top_k": 3, |
| 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 | ) |
| 290 | vllm_config = SimpleNamespace( |
| 291 | speculative_config=SimpleNamespace( |
| 292 | draft_model_config=SimpleNamespace( |
| 293 | hf_config=hf_config, |
| 294 | quantization=None, |
| 295 | ), |
| 296 | num_speculative_tokens=4, |
| 297 | enable_adaptive_verification=False, |
| 298 | ), |
| 299 | model_config=SimpleNamespace( |
| 300 | dtype=torch.float32, |
| 301 | is_mm_prefix_lm=False, |
| 302 | ), |
| 303 | load_config=SimpleNamespace( |
| 304 | quantization=None, |
| 305 | quantization_param_path=None, |
| 306 | ), |
| 307 | ) |
| 308 | mock_current_vllm_config.speculative_config = vllm_config.speculative_config |
| 309 | vllm_config.compilation_config = mock_current_vllm_config.compilation_config |
| 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"): |
| 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 |
| 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, |
| 438 | layers_attn, |
| 439 | has_bias=False, |
| 440 | ) |
| 441 | assert model._fused_kv_weight is None |
| 442 | assert all(not hasattr(projection, "weight") for projection in projections) |
| 443 | |
| 444 | actual_k, actual_v = qwen3_dflash.DFlashQwen3Model._project_context_kv( |
| 445 | model, |
| 446 | context_states, |
| 447 | num_ctx=2, |
| 448 | num_layers=2, |
| 449 | num_kv_heads=1, |
| 450 | head_dim=2, |
| 451 | ) |
| 452 | |
| 453 | expected_k = torch.tensor( |
| 454 | [[[1.0, 2.0], [3.0, 4.0]], [[2.0, 4.0], [6.0, 8.0]]] |
| 455 | ).unsqueeze(2) |
| 456 | expected_v = torch.tensor( |
| 457 | [[[3.0, -1.0], [7.0, -1.0]], [[-1.0, -2.0], [-3.0, -4.0]]] |
| 458 | ).unsqueeze(2) |
| 459 | torch.testing.assert_close(actual_k, expected_k) |
| 460 | torch.testing.assert_close(actual_v, expected_v) |
| 461 | assert [projection.calls for projection in projections] == [1, 1] |