[Qwen4Exp] Support UVA PLE-offload and Engram tensor parallelism (#54371)
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Harry Huang committed
Sep 9, 2026 at 22:32 UTC
3116c5d06bfe76501b3dd6b5434bfc7f3274f5e7
16 files changed
+1712
-495
tests/engine/test_arg_utils.py
+33
@@ -46,6 +46,39 @@ def test_optional_type():
46
assert optional_type_func("42") == 42
47
48
49
+@pytest.mark.parametrize(
50
+ "options",
51
+ [
52
+ ["--engram-config", '{"cpu_offload": false, "embedding_across_dp": true}'],
53
+ [
54
+ "--engram-config.cpu_offload",
55
+ "false",
56
+ "--engram-config.embedding_across_dp",
57
+ "true",
58
+ ],
59
+ ],
60
+)
61
+def test_engram_config_cli(options, monkeypatch):
62
+ """CLI settings take precedence over the legacy offload environment."""
63
+ monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
64
+ parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
65
+ args = EngineArgs.from_cli_args(parser.parse_args(options))
66
+ assert args.engram_config is not None
67
+ assert args.engram_config.cpu_offload is False
68
+ assert args.engram_config.embedding_across_dp is True
69
+
70
+
71
+@pytest.mark.parametrize(
72
+ "options, provided",
73
+ [([], False), (["--engram-config", "{}"], True)],
74
+)
75
+def test_engram_config_cli_optional(options, provided):
76
+ """An explicit empty config must remain distinct from an omitted config."""
77
+ parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
78
+ args = EngineArgs.from_cli_args(parser.parse_args(options))
79
+ assert (args.engram_config is not None) == provided
80
+
81
+
82
@pytest.mark.parametrize(
83
("type_hint", "type", "expected"),
84
[
tests/models/qwen4_exp/test_ple.py
+327
-21
@@ -13,27 +13,50 @@ from torch.nn import functional as F
13
14
import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
15
import vllm.model_executor.parameter as parameter_module
16
+import vllm.models.qwen4_exp.nvidia.ngram_embedding as ngram_embedding_module
17
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
18
from vllm.model_executor.layers.quantization.modelopt import (
19
ModelOptMixedPrecisionConfig,
20
+ ModelOptNvFp4Config,
21
)
22
from vllm.models.qwen4_exp.common.ple import (
23
PLEShardOverlap,
24
compute_ple_shard_overlap,
25
copy_ple_embedding_shard_,
26
)
25
-from vllm.models.qwen4_exp.nvidia.ple_layer import (
27
+from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
28
Qwen4ExpNGramEmbedding,
29
+ Qwen4ExpPLEDeviceEmbedding,
30
+ Qwen4ExpPLEEmbeddingMethod,
31
Qwen4ExpPLEFp8EmbeddingMethod,
28
- Qwen4ExpPLELayer,
29
- _get_ple_embedding_quant_method,
32
+ Qwen4ExpPLEPinnedHostEmbedding,
33
+ Qwen4ExpPLEUnquantizedEmbeddingMethod,
34
)
35
+from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
36
from vllm.v1.attention.backends.short_conv_attn import (
37
PleShortConvAttentionMetadata,
38
)
39
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
40
41
42
+def _mock_etp_group(
43
+ monkeypatch: pytest.MonkeyPatch,
44
+ world_size: int = 1,
45
+ all_reduce=lambda tensor: tensor,
46
+) -> None:
47
+ group = SimpleNamespace(
48
+ rank_in_group=0,
49
+ world_size=world_size,
50
+ all_reduce=all_reduce,
51
+ )
52
+ monkeypatch.setattr(ngram_embedding_module, "get_etp_group", lambda: group)
53
+ monkeypatch.setattr(
54
+ ngram_embedding_module,
55
+ "get_tp_group",
56
+ lambda: SimpleNamespace(world_size=world_size),
57
+ )
58
+
59
+
60
def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
61
module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
62
nn.Module.__init__(module)
@@ -171,11 +194,57 @@ def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None:
194
)
195
196
197
+@pytest.mark.parametrize(("dp_rank", "local_tokens"), [(0, 2), (1, 3)])
198
+def test_etp_lookup_gathers_and_returns_dp_local_rows(
199
+ monkeypatch: pytest.MonkeyPatch,
200
+ dp_rank: int,
201
+ local_tokens: int,
202
+) -> None:
203
+ embedding = Qwen4ExpPLEDeviceEmbedding.__new__(Qwen4ExpPLEDeviceEmbedding)
204
+ nn.Module.__init__(embedding)
205
+ embedding.tp_size = 2
206
+ embedding.etp_data_parallel_size = 2
207
+ embedding.data_parallel_rank = dp_rank
208
+ gathered_ids = torch.tensor([[10], [11], [0], [20], [21], [22]])
209
+ group = SimpleNamespace(
210
+ rank_in_group=dp_rank,
211
+ all_gather=lambda input_ids, dim: gathered_ids,
212
+ all_reduce=lambda embeddings: embeddings,
213
+ )
214
+ embedding.parallel_group = group
215
+ forward_context = SimpleNamespace(
216
+ dp_metadata=SimpleNamespace(
217
+ num_tokens_across_dp_cpu=torch.tensor([2, 3]),
218
+ )
219
+ )
220
+ monkeypatch.setattr(ngram_embedding_module, "get_dp_group", lambda: group)
221
+ monkeypatch.setattr(
222
+ ngram_embedding_module, "get_forward_context", lambda: forward_context
223
+ )
224
+ local_ids = gathered_ids[dp_rank * 3 : dp_rank * 3 + local_tokens]
225
+ embeddings = torch.arange(12).reshape(6, 2)
226
+
227
+ def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
228
+ assert torch.equal(input_ids, gathered_ids)
229
+ return embeddings
230
+
231
+ monkeypatch.setattr(
232
+ ngram_embedding_module.PLEVocabParallelEmbedding,
233
+ "forward",
234
+ forward,
235
+ )
236
+ output = embedding(local_ids)
237
+
238
+ expected = embeddings[dp_rank * 3 : dp_rank * 3 + local_tokens]
239
+ assert torch.equal(output, expected)
240
+
241
+
242
def _make_fp8_embedding_layer(
243
monkeypatch: pytest.MonkeyPatch,
244
*,
245
load_scale: bool = True,
178
-) -> embedding_module.VocabParallelEmbedding:
246
+) -> Qwen4ExpPLEDeviceEmbedding:
247
+ _mock_etp_group(monkeypatch)
248
monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
249
monkeypatch.setattr(
250
embedding_module, "get_tensor_model_parallel_world_size", lambda: 1
@@ -185,12 +254,13 @@ def _make_fp8_embedding_layer(
254
parameter_module, "get_tensor_model_parallel_world_size", lambda: 1
255
)
256
method = Qwen4ExpPLEFp8EmbeddingMethod()
188
- layer = embedding_module.VocabParallelEmbedding(
257
+ layer = Qwen4ExpPLEDeviceEmbedding(
258
3,
259
2,
260
params_dtype=torch.bfloat16,
261
padding_size=1,
193
- quant_method=method,
262
+ prefix="test.ple_embedding",
263
+ embedding_method=method,
264
)
265
weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]])
266
layer.weight.data.copy_(weight.to(torch.float8_e4m3fn))
@@ -228,7 +298,58 @@ def test_ple_fp8_embedding_rejects_missing_global_scale(monkeypatch) -> None:
298
layer.quant_method.process_weights_after_loading(layer)
299
300
231
-def test_ple_fp8_embedding_uses_int8_for_tp_reduce(monkeypatch) -> None:
301
+@pytest.mark.parametrize("requires_device_loading", [False, True])
302
+@pytest.mark.parametrize("load_scale", [False, True])
303
+def test_pinned_ple_post_load_validates_scale_without_staging(
304
+ monkeypatch: pytest.MonkeyPatch, requires_device_loading: bool, load_scale: bool
305
+) -> None:
306
+ """Validate scales through the loader while honoring the staging policy."""
307
+ from contextlib import nullcontext
308
+
309
+ from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
310
+ from vllm.model_executor.model_loader import utils as loader_utils
311
+
312
+ source = _make_fp8_embedding_layer(monkeypatch, load_scale=load_scale)
313
+ embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
314
+ nn.Module.__init__(embedding)
315
+ embedding.weight = source.weight
316
+ embedding.weight_scale = source.weight_scale
317
+ embedding.embedding_method = source.embedding_method
318
+ embedding.quant_method = source.embedding_method
319
+ embedding.tp_rank = source.tp_rank
320
+ embedding.tp_size = source.tp_size
321
+ assert QuantizeMethodBase.requires_device_loading
322
+ assert not embedding.quant_method.requires_device_loading
323
+ if requires_device_loading:
324
+ embedding.quant_method.requires_device_loading = True
325
+ model = nn.Module()
326
+ model.embedding = embedding
327
+ weight_ptr = embedding.weight.data_ptr()
328
+ staged_modules = []
329
+
330
+ def track_staging(module, target_device):
331
+ staged_modules.append(module)
332
+ return nullcontext()
333
+
334
+ monkeypatch.setattr(loader_utils, "device_loading_context", track_staging)
335
+ monkeypatch.setattr(loader_utils, "maybe_retie_word_embeddings", lambda *args: None)
336
+ monkeypatch.setattr(
337
+ loader_utils, "release_device_memory_under_pressure", lambda *args: None
338
+ )
339
+ config = SimpleNamespace(quantization="fp8")
340
+ if load_scale:
341
+ loader_utils.process_weights_after_loading(model, config, torch.device("cuda"))
342
+ else:
343
+ with pytest.raises(ValueError, match="missing its global scale"):
344
+ loader_utils.process_weights_after_loading(
345
+ model, config, torch.device("cuda")
346
+ )
347
+ assert embedding.weight.device.type == "cpu"
348
+ assert embedding.weight.data_ptr() == weight_ptr
349
+ assert staged_modules == ([embedding] if requires_device_loading else [])
350
+
351
+
352
+def test_ple_fp8_embedding_uses_int8_for_parallel_reduce(monkeypatch) -> None:
353
monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
354
monkeypatch.setattr(
355
embedding_module, "get_tensor_model_parallel_world_size", lambda: 2
@@ -251,17 +372,14 @@ def test_ple_fp8_embedding_uses_int8_for_tp_reduce(monkeypatch) -> None:
372
reduced_dtypes.append(tensor.dtype)
373
return tensor.clone()
374
254
- monkeypatch.setattr(
255
- embedding_module,
256
- "tensor_model_parallel_all_reduce",
257
- all_reduce,
258
- )
259
- layer = embedding_module.VocabParallelEmbedding(
375
+ _mock_etp_group(monkeypatch, world_size=2, all_reduce=all_reduce)
376
+ layer = Qwen4ExpPLEDeviceEmbedding(
377
4,
378
2,
379
params_dtype=torch.bfloat16,
380
padding_size=1,
264
- quant_method=Qwen4ExpPLEFp8EmbeddingMethod(),
381
+ prefix="test.ple_embedding",
382
+ embedding_method=Qwen4ExpPLEFp8EmbeddingMethod(),
383
)
384
layer.weight.data.copy_(
385
torch.tensor([[1.0, 2.0], [4.0, 8.0]]).to(torch.float8_e4m3fn)
@@ -283,12 +401,200 @@ def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None:
401
weight_block_size=[128, 128],
402
)
403
assert isinstance(
286
- _get_ple_embedding_quant_method(quant_config, prefix),
404
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
405
Qwen4ExpPLEFp8EmbeddingMethod,
406
)
407
408
quant_config.ignored_layers = [f"{prefix}.shard_0"]
291
- assert _get_ple_embedding_quant_method(quant_config, prefix) is None
409
+ assert isinstance(
410
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
411
+ Qwen4ExpPLEUnquantizedEmbeddingMethod,
412
+ )
413
+ assert isinstance(
414
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(None, prefix),
415
+ Qwen4ExpPLEUnquantizedEmbeddingMethod,
416
+ )
417
+
418
+
419
+def test_ple_embedding_rejects_unsupported_quantization_configs() -> None:
420
+ prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
421
+ nvfp4_config = ModelOptNvFp4Config(exclude_modules=[])
422
+ with pytest.raises(NotImplementedError, match="ModelOptNvFp4Config"):
423
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(nvfp4_config, prefix)
424
+
425
+ dynamic_fp8_config = Fp8Config(
426
+ is_checkpoint_fp8_serialized=False,
427
+ ignored_layers=[],
428
+ )
429
+ with pytest.raises(NotImplementedError, match="serialized FP8"):
430
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(dynamic_fp8_config, prefix)
431
+
432
+
433
+def test_ple_embedding_respects_modelopt_exclusion() -> None:
434
+ prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
435
+ quant_config = ModelOptNvFp4Config(exclude_modules=[prefix])
436
+
437
+ assert isinstance(
438
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
439
+ Qwen4ExpPLEUnquantizedEmbeddingMethod,
440
+ )
441
+
442
+
443
+def test_ple_embedding_dtype_overrides_modelopt_exclusion() -> None:
444
+ prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
445
+ quant_config = ModelOptNvFp4Config(exclude_modules=[prefix])
446
+
447
+ assert isinstance(
448
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(
449
+ quant_config,
450
+ prefix,
451
+ "float8_e4m3fn",
452
+ ),
453
+ Qwen4ExpPLEFp8EmbeddingMethod,
454
+ )
455
+
456
+
457
+def test_pinned_embedding_forward_finalizes_prefetched_output(
458
+ monkeypatch: pytest.MonkeyPatch,
459
+) -> None:
460
+ embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
461
+ nn.Module.__init__(embedding)
462
+ embedding._prefetch_buffer = torch.empty(4, 2, 3, dtype=torch.float8_e4m3fn)
463
+ embedding._output_dim = 6
464
+ hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16)
465
+ expected = torch.arange(12).reshape(2, 6).to(torch.float8_e4m3fn)
466
+
467
+ def finalize_prefetched(
468
+ prefetch_output: torch.Tensor,
469
+ output: torch.Tensor,
470
+ ) -> None:
471
+ assert prefetch_output is embedding._prefetch_buffer
472
+ output.copy_(expected)
473
+
474
+ monkeypatch.setattr(
475
+ embedding,
476
+ "_finalize_prefetch",
477
+ finalize_prefetched,
478
+ )
479
+
480
+ output = embedding(hidden_states)
481
+
482
+ assert output.dtype == embedding._prefetch_buffer.dtype
483
+ assert torch.equal(output, expected)
484
+
485
+
486
+def test_pinned_fp8_embedding_uses_int8_for_parallel_reduce() -> None:
487
+ embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
488
+ nn.Module.__init__(embedding)
489
+ embedding.tp_size = 2
490
+ reduced_dtypes = []
491
+
492
+ def all_reduce(tensor: torch.Tensor) -> torch.Tensor:
493
+ reduced_dtypes.append(tensor.dtype)
494
+ return tensor.clone()
495
+
496
+ embedding.parallel_group = SimpleNamespace(all_reduce=all_reduce)
497
+ embeddings = torch.arange(8).reshape(2, 4).to(torch.float8_e4m3fn)
498
+
499
+ output = embedding._reduce_etp_embeddings(embeddings)
500
+
501
+ assert reduced_dtypes == [torch.int8]
502
+ assert output.dtype == embeddings.dtype
503
+ torch.testing.assert_close(output.float(), embeddings.float())
504
+
505
+
506
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
507
+def test_ple_device_embedding_allocates_on_active_device(
508
+ monkeypatch: pytest.MonkeyPatch,
509
+) -> None:
510
+ _mock_etp_group(monkeypatch)
511
+ monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
512
+ monkeypatch.setattr(
513
+ embedding_module, "get_tensor_model_parallel_world_size", lambda: 1
514
+ )
515
+ monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0)
516
+ monkeypatch.setattr(
517
+ parameter_module, "get_tensor_model_parallel_world_size", lambda: 1
518
+ )
519
+ embedding_method = Qwen4ExpPLEUnquantizedEmbeddingMethod()
520
+
521
+ with torch.device("cuda:0"):
522
+ embedding = Qwen4ExpPLEDeviceEmbedding(
523
+ 4,
524
+ 3,
525
+ params_dtype=torch.bfloat16,
526
+ padding_size=1,
527
+ prefix="test.ple_embedding",
528
+ embedding_method=embedding_method,
529
+ )
530
+
531
+ assert embedding.weight.device == torch.device("cuda:0")
532
+
533
+
534
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
535
+@pytest.mark.parametrize("fp8_checkpoint", [False, True])
536
+def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva(
537
+ monkeypatch: pytest.MonkeyPatch,
538
+ fp8_checkpoint: bool,
539
+) -> None:
540
+ _mock_etp_group(monkeypatch)
541
+ monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
542
+ monkeypatch.setattr(
543
+ embedding_module, "get_tensor_model_parallel_world_size", lambda: 1
544
+ )
545
+ monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0)
546
+ monkeypatch.setattr(
547
+ parameter_module, "get_tensor_model_parallel_world_size", lambda: 1
548
+ )
549
+
550
+ with torch.device("cuda:0"):
551
+ embedding_method = (
552
+ Qwen4ExpPLEFp8EmbeddingMethod()
553
+ if fp8_checkpoint
554
+ else Qwen4ExpPLEUnquantizedEmbeddingMethod()
555
+ )
556
+ embedding = Qwen4ExpPLEPinnedHostEmbedding(
557
+ 4,
558
+ 3,
559
+ params_dtype=torch.bfloat16,
560
+ padding_size=1,
561
+ prefix="test.ple_embedding",
562
+ embedding_method=embedding_method,
563
+ )
564
+
565
+ storage_dtype = torch.float8_e4m3fn if fp8_checkpoint else torch.bfloat16
566
+ loaded_weight = (
567
+ torch.arange(12, dtype=torch.float32).reshape(4, 3).to(storage_dtype)
568
+ )
569
+ copied = copy_ple_embedding_shard_(
570
+ embedding.weight,
571
+ loaded_weight,
572
+ checkpoint_start=0,
573
+ tp_start=0,
574
+ tp_end=4,
575
+ )
576
+ input_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")
577
+ output = embedding._lookup(input_ids)
578
+ if fp8_checkpoint:
579
+ embedding.weight_scale.data.fill_(0.25)
580
+ weight_ptr = embedding.weight.data_ptr()
581
+ embedding.quant_method.process_weights_after_loading(embedding)
582
+ assert embedding.weight.data_ptr() == weight_ptr
583
+ dequantized = embedding.dequantize(output, torch.bfloat16)
584
+
585
+ assert copied == 4
586
+ assert embedding.weight.device.type == "cpu"
587
+ assert embedding.weight.is_pinned()
588
+ assert output.dtype == storage_dtype
589
+ expected = loaded_weight[input_ids.cpu()].to(device="cuda:0")
590
+ torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
591
+ expected_scale = 0.25 if fp8_checkpoint else 1.0
592
+ torch.testing.assert_close(
593
+ dequantized,
594
+ expected.to(torch.bfloat16) * expected_scale,
595
+ rtol=0,
596
+ atol=0,
597
+ )
598
599
600
def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
@@ -310,15 +616,15 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:
616
)
617
618
assert isinstance(
313
- _get_ple_embedding_quant_method(quant_config, prefix),
619
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
620
Qwen4ExpPLEFp8EmbeddingMethod,
621
)
316
- assert (
317
- _get_ple_embedding_quant_method(
622
+ assert isinstance(
623
+ Qwen4ExpPLEEmbeddingMethod.from_quant_config(
624
quant_config,
625
"model.language_model.layers.2.moe.gate_proj",
320
- )
321
- is None
626
+ ),
627
+ Qwen4ExpPLEUnquantizedEmbeddingMethod,
628
)
629
630
tests/test_config.py
+154
@@ -22,6 +22,7 @@ from vllm.config import (
22
CacheConfig,
23
CompilationConfig,
24
DeviceConfig,
25
+ EngramConfig,
26
KernelConfig,
27
KVTransferConfig,
28
ModelConfig,
@@ -913,6 +914,159 @@ def test_all2all_backend_has_portable_default():
914
assert ParallelConfig().all2all_backend == "allgather_reducescatter"
915
916
917
+@pytest.mark.parametrize(
918
+ "dp_size, across_dp, expected",
919
+ [(1, False, 4), (1, True, 4), (2, False, 4), (2, True, 8)],
920
+)
921
+def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: int):
922
+ parallel = ParallelConfig(
923
+ tensor_parallel_size=4,
924
+ data_parallel_size=dp_size,
925
+ distributed_executor_backend="mp",
926
+ )
927
+ config = EngramConfig(embedding_across_dp=across_dp)
928
+ assert config.get_parallel_size(parallel) == expected
929
+
930
+
931
+def test_engram_rejects_elastic_cross_dp():
932
+ parallel = ParallelConfig(
933
+ tensor_parallel_size=4,
934
+ data_parallel_size=2,
935
+ distributed_executor_backend="mp",
936
+ )
937
+ parallel.enable_elastic_ep = True
938
+ with pytest.raises(ValueError, match="embedding_across_dp.*elastic EP"):
939
+ EngramConfig(embedding_across_dp=True).verify_parallel_config(parallel)
940
+
941
+
942
+@pytest.mark.parametrize("legacy", [None, "0", "1"])
943
+def test_engram_cpu_offload_environment_fallback(monkeypatch, legacy):
944
+ """Explicit settings must override the legacy environment fallback."""
945
+ monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
946
+ if legacy is not None:
947
+ monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
948
+ assert EngramConfig().cpu_offload == (legacy == "1")
949
+ assert EngramConfig(cpu_offload=False).cpu_offload is False
950
+ assert EngramConfig(cpu_offload=True).cpu_offload is True
951
+
952
+
953
+@pytest.mark.parametrize(
954
+ "architecture, ple_layers, cuda, supported",
955
+ [
956
+ ("Qwen4ExpForCausalLM", [1], True, True),
957
+ ("Qwen4ExpForConditionalGeneration", [1], True, True),
958
+ ("Qwen4ExpForCausalLM", [], True, False),
959
+ ("Qwen4ExpForCausalLM", None, True, False),
960
+ ("Qwen4ExpForCausalLM", [1], False, False),
961
+ ("LlamaForCausalLM", [1], True, False),
962
+ ("Qwen4ExpMTP", [], True, False),
963
+ (None, None, True, False),
964
+ ],
965
+)
966
+def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, supported):
967
+ """A similarly named HF field must not enable unsupported implementations."""
968
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
969
+ model = (
970
+ cast(
971
+ ModelConfig,
972
+ SimpleNamespace(
973
+ architecture=architecture,
974
+ hf_text_config=SimpleNamespace(ple_layer_ids=ple_layers),
975
+ ),
976
+ )
977
+ if architecture is not None
978
+ else None
979
+ )
980
+ config = EngramConfig(cpu_offload=False, embedding_across_dp=False)
981
+ if supported:
982
+ config.verify_model_config(model)
983
+ else:
984
+ with pytest.raises(ValueError, match="requires a model with supported Engram"):
985
+ config.verify_model_config(model)
986
+
987
+
988
+def test_engram_config_defaults_to_none(monkeypatch):
989
+ monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False)
990
+ config = VllmConfig()
991
+ assert config.engram_config is None
992
+ assert config.compute_hash()
993
+
994
+
995
+@pytest.mark.parametrize("legacy", ["0", "1"])
996
+def test_engram_none_resolves_legacy_offload(monkeypatch, legacy):
997
+ """Legacy enablement materializes a config before model validation."""
998
+ monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
999
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
1000
+ config = cast(
1001
+ VllmConfig,
1002
+ SimpleNamespace(
1003
+ model_config=SimpleNamespace(
1004
+ architecture="Qwen4ExpForCausalLM",
1005
+ hf_text_config=SimpleNamespace(ple_layer_ids=[1]),
1006
+ ),
1007
+ speculative_config=None,
1008
+ engram_config=None,
1009
+ parallel_config=ParallelConfig(),
1010
+ ),
1011
+ )
1012
+ VllmConfig._resolve_and_verify_engram_config(config)
1013
+ if legacy == "1":
1014
+ assert config.engram_config is not None
1015
+ assert config.engram_config.cpu_offload is True
1016
+ else:
1017
+ assert config.engram_config is None
1018
+
1019
+
1020
+@pytest.mark.parametrize("legacy", ["0", "1"])
1021
+def test_engram_explicit_config_requires_supported_model(monkeypatch, legacy):
1022
+ """Explicit all-false settings still opt into model validation."""
1023
+ monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy)
1024
+ with pytest.raises(ValueError, match="requires a model with supported Engram"):
1025
+ VllmConfig(engram_config=EngramConfig(cpu_offload=False))
1026
+
1027
+
1028
+def test_engram_legacy_offload_requires_supported_model(monkeypatch):
1029
+ monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1")
1030
+ with pytest.raises(ValueError, match="requires a model with supported Engram"):
1031
+ VllmConfig()
1032
+
1033
+
1034
+@pytest.mark.parametrize("target_has_ple", [False, True])
1035
+def test_engram_draft_config_validates_target(monkeypatch, target_has_ple):
1036
+ """MTP may inherit cross-DP sharding without having its own PLE layers."""
1037
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
1038
+ target = SimpleNamespace(
1039
+ architecture="Qwen4ExpForCausalLM",
1040
+ hf_text_config=SimpleNamespace(ple_layer_ids=[1] if target_has_ple else []),
1041
+ )
1042
+ draft = SimpleNamespace(architecture="Qwen4ExpMTP")
1043
+ config = cast(
1044
+ VllmConfig,
1045
+ SimpleNamespace(
1046
+ model_config=draft,
1047
+ speculative_config=SimpleNamespace(
1048
+ draft_model_config=draft, target_model_config=target
1049
+ ),
1050
+ engram_config=EngramConfig(embedding_across_dp=True),
1051
+ parallel_config=ParallelConfig(),
1052
+ ),
1053
+ )
1054
+ if target_has_ple:
1055
+ VllmConfig._resolve_and_verify_engram_config(config)
1056
+ else:
1057
+ with pytest.raises(ValueError, match="requires a model with supported Engram"):
1058
+ VllmConfig._resolve_and_verify_engram_config(config)
1059
+
1060
+
1061
+def test_engram_hash_tracks_storage_and_sharding():
1062
+ configs = [
1063
+ EngramConfig(cpu_offload=offload, embedding_across_dp=across_dp)
1064
+ for offload in (False, True)
1065
+ for across_dp in (False, True)
1066
+ ]
1067
+ assert len({config.compute_hash() for config in configs}) == 4
1068
+
1069
+
1070
@pytest.mark.parametrize("port", [1, 29550, 65535])
1071
def test_data_parallel_rpc_port_accepts_valid_ports(port: int):
1072
assert ParallelConfig(data_parallel_rpc_port=port).data_parallel_rpc_port == port
vllm/config/__init__.py
+3
@@ -13,6 +13,7 @@ from vllm.config.device import DeviceConfig
13
from vllm.config.diffusion import DiffusionConfig
14
from vllm.config.ec_manager_config import EncoderCacheManagerConfig
15
from vllm.config.ec_transfer import ECTransferConfig
16
+from vllm.config.engram import EngramConfig
17
from vllm.config.fault_tolerance import FaultToleranceConfig
18
from vllm.config.kernel import KernelConfig
19
from vllm.config.kv_events import KVEventsConfig
@@ -81,6 +82,8 @@ __all__ = [
82
"EncoderCacheManagerConfig",
83
# From vllm.config.ec_transfer
84
"ECTransferConfig",
85
+ # From vllm.config.engram
86
+ "EngramConfig",
87
# From vllm.config.kernel
88
"KernelConfig",
89
# From vllm.config.kv_events
vllm/config/engram.py
new
+82
@@ -0,0 +1,82 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import TYPE_CHECKING
5
+
6
+from pydantic import Field
7
+
8
+import vllm.envs as envs
9
+from vllm.config.utils import config, get_hash_factors, hash_factors
10
+from vllm.logger import init_logger
11
+
12
+if TYPE_CHECKING:
13
+ from vllm.config.model import ModelConfig
14
+ from vllm.config.parallel import ParallelConfig
15
+
16
+logger = init_logger(__name__)
17
+
18
+
19
+def _default_cpu_offload() -> bool:
20
+ """Honor the legacy environment variable only when the field is omitted."""
21
+ if envs.is_set("VLLM_PLE_CPU_OFFLOAD"):
22
+ logger.warning_once(
23
+ "VLLM_PLE_CPU_OFFLOAD is a legacy setting and may be removed in a "
24
+ "future release. Use --engram-config.cpu_offload instead."
25
+ )
26
+ return envs.VLLM_PLE_CPU_OFFLOAD
27
+
28
+
29
+@config
30
+class EngramConfig:
31
+ """Configuration for Engram embedding storage and sharding."""
32
+
33
+ cpu_offload: bool = Field(default_factory=_default_cpu_offload)
34
+ """Store embedding weights in pinned CPU memory for UVA lookup.
35
+ Defaults to False, or VLLM_PLE_CPU_OFFLOAD when set for compatibility.
36
+ An explicit value takes precedence over the legacy environment variable."""
37
+
38
+ embedding_across_dp: bool = False
39
+ """Shard embeddings across TP and all DP ranks when enabled.
40
+ Otherwise, each DP rank has a separate TP-sharded embedding replica."""
41
+
42
+ def verify_model_config(self, model_config: "ModelConfig | None") -> None:
43
+ """Reject Engram configuration for models without supported embeddings."""
44
+ from vllm.platforms import current_platform
45
+
46
+ supported_architectures = {
47
+ "Qwen4ExpForCausalLM",
48
+ "Qwen4ExpForConditionalGeneration",
49
+ }
50
+ if (
51
+ model_config is None
52
+ or model_config.architecture not in supported_architectures
53
+ or not current_platform.is_cuda()
54
+ or not getattr(model_config.hf_text_config, "ple_layer_ids", None)
55
+ ):
56
+ raise ValueError(
57
+ "EngramConfig requires a model with supported Engram "
58
+ "embeddings. Currently only the CUDA Qwen4Exp implementation "
59
+ "with non-empty ple_layer_ids is supported."
60
+ )
61
+
62
+ def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None:
63
+ """Reject unsupported embedding parallel topologies."""
64
+ if (
65
+ self.embedding_across_dp
66
+ and parallel_config.data_parallel_size > 1
67
+ and parallel_config.enable_elastic_ep
68
+ ):
69
+ raise ValueError(
70
+ "Engram embedding_across_dp is not supported with elastic EP yet."
71
+ )
72
+
73
+ def get_parallel_size(self, parallel_config: "ParallelConfig") -> int:
74
+ """Derive the embedding group size from the parallel configuration."""
75
+ size = parallel_config.tensor_parallel_size
76
+ if self.embedding_across_dp and parallel_config.data_parallel_size > 1:
77
+ size *= parallel_config.data_parallel_size
78
+ return size
79
+
80
+ def compute_hash(self) -> str:
81
+ """Hash settings that affect embedding execution and graph structure."""
82
+ return hash_factors(get_hash_factors(self, set()))
vllm/config/vllm.py
+28
@@ -34,6 +34,7 @@ from .device import DeviceConfig
34
from .diffusion import DiffusionConfig
35
from .ec_manager_config import EncoderCacheManagerConfig
36
from .ec_transfer import ECTransferConfig
37
+from .engram import EngramConfig
38
from .kernel import KernelConfig
39
from .kv_events import KVEventsConfig
40
from .kv_transfer import KVTransferConfig
@@ -370,6 +371,8 @@ class VllmConfig:
371
"""Model weight offloading configuration."""
372
attention_config: AttentionConfig = Field(default_factory=AttentionConfig)
373
"""Attention configuration."""
374
+ engram_config: EngramConfig | None = None
375
+ """Optional Engram configuration, only valid for supported PLE models."""
376
mamba_config: MambaConfig = Field(default_factory=MambaConfig)
377
"""Mamba configuration."""
378
kernel_config: KernelConfig = Field(default_factory=KernelConfig)
@@ -506,6 +509,11 @@ class VllmConfig:
509
vllm_factors.append(self.attention_config.compute_hash())
510
else:
511
vllm_factors.append("None")
512
+ vllm_factors.append(
513
+ self.engram_config.compute_hash()
514
+ if self.engram_config is not None
515
+ else "None"
516
+ )
517
if self.lora_config:
518
vllm_factors.append(self.lora_config.compute_hash())
519
else:
@@ -1083,6 +1091,25 @@ class VllmConfig:
1091
if not self.use_v2_model_runner:
1092
raise ValueError("trace replay requires Model Runner V2")
1093
1094
+ def _resolve_and_verify_engram_config(self) -> None:
1095
+ """Resolve legacy offload settings and validate model and parallel configs."""
1096
+ if self.engram_config is None:
1097
+ if not envs.VLLM_PLE_CPU_OFFLOAD:
1098
+ return
1099
+ self.engram_config = EngramConfig()
1100
+ model_config = self.model_config
1101
+ speculative_config = self.speculative_config
1102
+ # Draft configs inherit the target's communication groups and settings.
1103
+ # Qwen4Exp MTP itself disables PLE, so validate its target instead.
1104
+ if (
1105
+ speculative_config is not None
1106
+ and model_config is speculative_config.draft_model_config
1107
+ ):
1108
+ model_config = speculative_config.target_model_config
1109
+ self.engram_config.verify_model_config(model_config)
1110
+ self.engram_config.verify_parallel_config(self.parallel_config)
1111
+ logger.info_once("Resolved Engram configuration: %s", str(self.engram_config))
1112
+
1113
def __post_init__(self):
1114
"""Verify configs are valid & consistent with each other."""
1115
@@ -1095,6 +1122,7 @@ class VllmConfig:
1122
logger.info_once("Performance mode set to '%s'.", self.performance_mode)
1123
1124
self.try_verify_and_update_config()
1125
+ self._resolve_and_verify_engram_config()
1126
1127
# Models may have supplied their own DCP defaults above; anything still
1128
# unset falls back to the stock ones.
vllm/distributed/device_communicators/cuda_communicator.py
+2
-1
@@ -47,7 +47,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
47
global_world_size,
48
use_all2all=use_all2all,
49
)
50
- if "tp" not in unique_name:
50
+ # Match the group name exactly so ETP does not enable TP-only backends.
51
+ if unique_name.split(":")[0] != "tp":
52
# custom allreduce or torch symm mem can be used only by tp
53
use_custom_allreduce = False
54
use_torch_symm_mem = False
vllm/distributed/parallel_state.py
+38
-2
@@ -1551,6 +1551,15 @@ def get_tp_group() -> GroupCoordinator:
1551
return _TP
1552
1553
1554
+_ETP: GroupCoordinator | None = None
1555
+
1556
+
1557
+def get_etp_group() -> GroupCoordinator:
1558
+ """Return the Engram Tensor Parallel (ETP) group that shards one embedding table."""
1559
+ assert _ETP is not None, "Engram tensor-parallel group is not initialized"
1560
+ return _ETP
1561
+
1562
+
1563
_DCP: GroupCoordinator | None = None
1564
1565
@@ -2008,6 +2017,28 @@ def initialize_model_parallel(
2017
group_name="tp",
2018
)
2019
2020
+ global _ETP
2021
+ assert _ETP is None, "Engram tensor-parallel group is already initialized"
2022
+ engram_tensor_parallel_size = (
2023
+ config.engram_config.get_parallel_size(parallel_config)
2024
+ if config.engram_config is not None
2025
+ else tensor_model_parallel_size
2026
+ )
2027
+ if engram_tensor_parallel_size == tensor_model_parallel_size:
2028
+ _ETP = _TP
2029
+ else:
2030
+ group_ranks = (
2031
+ all_ranks.permute(0, 2, 3, 1, 4)
2032
+ .reshape(-1, engram_tensor_parallel_size)
2033
+ .unbind(0)
2034
+ )
2035
+ _ETP = init_model_parallel_group(
2036
+ [ranks.tolist() for ranks in group_ranks],
2037
+ get_world_group().local_rank,
2038
+ backend,
2039
+ group_name="etp",
2040
+ )
2041
+
2042
# Build the DCP model-parallel groups.
2043
global _DCP
2044
assert _DCP is None, "decode context model parallel group is already initialized"
@@ -2142,13 +2173,14 @@ def initialize_model_parallel(
2173
logger.info_once(
2174
"rank %s in world size %s is assigned as "
2175
"DP rank %s, PP rank %s, PCP rank %s, "
2145
- "TP rank %s, EP rank %s, EPLB rank %s",
2176
+ "TP rank %s, ETP rank %s, EP rank %s, EPLB rank %s",
2177
rank,
2178
world_size,
2179
_DP.rank_in_group,
2180
_PP.rank_in_group,
2181
_PCP.rank_in_group,
2182
_TP.rank_in_group,
2183
+ _ETP.rank_in_group,
2184
_EP.rank_in_group if _EP is not None else "N/A",
2185
_EPLB.rank_in_group if _EPLB is not None else "N/A",
2186
)
@@ -2243,7 +2275,11 @@ def get_node_count() -> int:
2275
2276
def destroy_model_parallel():
2277
"""Set the groups to none and destroy them."""
2246
- global _TP
2278
+ global _TP, _ETP
2279
+
2280
+ if _ETP and _ETP is not _TP:
2281
+ _ETP.destroy()
2282
+ _ETP = None
2283
2284
if _TP:
2285
_TP.destroy()
vllm/engine/arg_utils.py
+6
@@ -42,6 +42,7 @@ from vllm.config import (
42
DiffusionConfig,
43
ECTransferConfig,
44
EncoderCacheManagerConfig,
45
+ EngramConfig,
46
EPLBConfig,
47
FaultToleranceConfig,
48
KernelConfig,
@@ -688,6 +689,7 @@ class EngineArgs:
689
pooler_config: PoolerConfig | None = ModelConfig.pooler_config
690
compilation_config: CompilationConfig = get_field(VllmConfig, "compilation_config")
691
attention_config: AttentionConfig = get_field(VllmConfig, "attention_config")
692
+ engram_config: EngramConfig | None = VllmConfig.engram_config
693
mamba_config: MambaConfig = get_field(VllmConfig, "mamba_config")
694
kernel_config: KernelConfig = get_field(VllmConfig, "kernel_config")
695
enable_flashinfer_autotune: bool = get_field(
@@ -788,6 +790,8 @@ class EngineArgs:
790
self.compilation_config = CompilationConfig(**self.compilation_config)
791
if isinstance(self.attention_config, dict):
792
self.attention_config = AttentionConfig(**self.attention_config)
793
+ if isinstance(self.engram_config, dict):
794
+ self.engram_config = EngramConfig(**self.engram_config)
795
if isinstance(self.mamba_config, dict):
796
self.mamba_config = MambaConfig(**self.mamba_config)
797
if isinstance(self.kernel_config, dict):
@@ -1701,6 +1705,7 @@ class EngineArgs:
1705
vllm_group.add_argument(
1706
"--attention-config", "-ac", **vllm_kwargs["attention_config"]
1707
)
1708
+ vllm_group.add_argument("--engram-config", **vllm_kwargs["engram_config"])
1709
vllm_group.add_argument("--reasoning-config", **vllm_kwargs["reasoning_config"])
1710
vllm_group.add_argument("--kernel-config", **vllm_kwargs["kernel_config"])
1711
vllm_group.add_argument(
@@ -2599,6 +2604,7 @@ class EngineArgs:
2604
load_config=load_config,
2605
offload_config=offload_config,
2606
attention_config=attention_config,
2607
+ engram_config=self.engram_config,
2608
mamba_config=mamba_config,
2609
kernel_config=kernel_config,
2610
lora_config=lora_config,
vllm/envs.py
+6
@@ -155,6 +155,7 @@ if TYPE_CHECKING:
155
VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False
156
VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
157
VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
158
+ VLLM_PLE_CPU_OFFLOAD: bool = False
159
VLLM_DISABLE_COMPILE_CACHE: bool = False
160
VLLM_REPLICATE_EMBED: bool = False
161
VLLM_USE_LAYERNAME: bool = True
@@ -2077,6 +2078,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
2078
"VLLM_LOG_MODEL_INSPECTION": lambda: bool(
2079
int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0"))
2080
),
2081
+ # Keep Qwen4Exp PLE embedding tables in pinned CPU memory and gather their
2082
+ # rows through UVA on a dedicated CUDA stream.
2083
+ # Legacy fallback for EngramConfig.cpu_offload, which takes precedence.
2084
+ # This environment variable may be removed in a future release.
2085
+ "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))),
2086
# Debug logging for --enable-mfu-metrics
2087
"VLLM_DEBUG_MFU_METRICS": lambda: bool(
2088
int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0"))
vllm/model_executor/layers/quantization/base_config.py
+3
@@ -23,6 +23,9 @@ else:
23
class QuantizeMethodBase(ABC):
24
"""Base class for different quantized methods."""
25
26
+ requires_device_loading: bool = True
27
+ """Whether post-load processing requires parameters on the target device."""
28
+
29
uses_meta_device: bool = False
30
"""
31
Whether this method creates weights on meta device for online quantization.
vllm/model_executor/layers/vocab_parallel_embedding.py
+14
-1
@@ -11,6 +11,7 @@ from torch.nn.parameter import Parameter
11
import vllm.envs as envs
12
from vllm import _custom_ops as ops
13
from vllm.distributed import (
14
+ GroupCoordinator,
15
divide,
16
get_tensor_model_parallel_rank,
17
get_tensor_model_parallel_world_size,
@@ -240,6 +241,7 @@ class VocabParallelEmbedding(PluggableLayer):
241
prefix: full name of the layer in the state dict
242
disable_tp: If true, tensor parallelism will be disabled for this layer.
243
quant_method: Preselected quantization method for model-specific layers.
244
+ parallel_group: Process group used to shard and reduce the embedding.
245
""" # noqa: E501
246
247
# --8<-- [end:vocab_parallel_embedding]
@@ -256,13 +258,18 @@ class VocabParallelEmbedding(PluggableLayer):
258
*,
259
disable_tp: bool = False,
260
quant_method: QuantizeMethodBase | None = None,
261
+ parallel_group: GroupCoordinator | None = None,
262
):
263
super().__init__()
264
265
# Keep the input dimensions.
266
self.disable_tp = disable_tp
267
+ self.parallel_group = parallel_group
268
if disable_tp:
269
tp_rank, self.tp_size = 0, 1
270
+ elif parallel_group is not None:
271
+ tp_rank = parallel_group.rank_in_group
272
+ self.tp_size = parallel_group.world_size
273
else:
274
tp_rank = get_tensor_model_parallel_rank()
275
self.tp_size = get_tensor_model_parallel_world_size()
@@ -535,10 +542,16 @@ class VocabParallelEmbedding(PluggableLayer):
542
# Each vocab token has one owner, so FP8 bytes can use int8 SUM.
543
comm_output = output_parallel.view(torch.int8)
544
comm_output.masked_fill_(input_mask.unsqueeze(-1), 0)
538
- output = tensor_model_parallel_all_reduce(comm_output)
545
+ output = (
546
+ self.parallel_group.all_reduce(comm_output)
547
+ if self.parallel_group is not None
548
+ else tensor_model_parallel_all_reduce(comm_output)
549
+ )
550
return output.view(output_parallel.dtype)
551
output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
552
# Reduce across all the model parallel GPUs.
553
+ if self.parallel_group is not None:
554
+ return self.parallel_group.all_reduce(output_parallel)
555
return tensor_model_parallel_all_reduce(output_parallel)
556
557
def extra_repr(self) -> str:
vllm/model_executor/model_loader/utils.py
+10
-3
@@ -4,7 +4,7 @@
4
5
import inspect
6
import warnings
7
-from contextlib import contextmanager
7
+from contextlib import contextmanager, nullcontext
8
from typing import Any
9
10
import torch
@@ -120,11 +120,18 @@ def process_weights_after_loading(
120
"does not support pre-processed weights"
121
)
122
# When quant methods need to process weights after loading
123
- # (for repacking, quantizing, etc), they expect parameters
123
+ # (for repacking, quantizing, etc), they typically expect parameters
124
# to be on the global target device. This scope is for the
125
# case where cpu offloading is used, where we will move the
126
# parameters onto device for processing and back off after.
127
- with device_loading_context(module, target_device):
127
+ # Methods that can process weights in place (e.g. PLE scale
128
+ # validation) set requires_device_loading=False to skip this move.
129
+ loading_context = (
130
+ device_loading_context(module, target_device)
131
+ if quant_method.requires_device_loading
132
+ else nullcontext()
133
+ )
134
+ with loading_context:
135
quant_method.process_weights_after_loading(module)
136
# process_weights_after_loading may swap in freshly-created
137
# Parameters (e.g. FP8 requantization), which are stamped with the
vllm/models/qwen4_exp/nvidia/model.py
+37
@@ -461,6 +461,27 @@ class Qwen4ExpModel(nn.Module):
461
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
462
return self.embed_tokens(input_ids)
463
464
+ @staticmethod
465
+ def _start_layer_ple_prefetch(
466
+ layer: nn.Module,
467
+ hidden_states: torch.Tensor,
468
+ input_ids: torch.Tensor | None,
469
+ query_start_loc: torch.Tensor | None,
470
+ ngram_context: torch.Tensor | None,
471
+ ) -> None:
472
+ """Start a layer's PLE prefetch when the required inputs exist."""
473
+ ple: Qwen4ExpPLELayer | None = getattr(layer, "ple", None)
474
+ if ple is None:
475
+ return
476
+ if input_ids is None or query_start_loc is None or ngram_context is None:
477
+ raise RuntimeError("PLE inputs were not prepared")
478
+ ple.start_prefetch(
479
+ hidden_states,
480
+ input_ids,
481
+ query_start_loc,
482
+ ngram_context,
483
+ )
484
+
485
def forward(
486
self,
487
input_ids: torch.Tensor | None,
@@ -487,10 +508,26 @@ class Qwen4ExpModel(nn.Module):
508
block_output = None
509
injection = None
510
last_layer = None
511
+ if self.start_layer < self.end_layer:
512
+ self._start_layer_ple_prefetch(
513
+ self.layers[self.start_layer],
514
+ hidden_states,
515
+ input_ids,
516
+ query_start_loc,
517
+ ngram_context,
518
+ )
519
for layer_idx, layer in islice(
520
enumerate(self.layers), self.start_layer, self.end_layer
521
):
522
last_layer = layer
523
+ if layer_idx + 1 < self.end_layer:
524
+ self._start_layer_ple_prefetch(
525
+ self.layers[layer_idx + 1],
526
+ hidden_states,
527
+ input_ids,
528
+ query_start_loc,
529
+ ngram_context,
530
+ )
531
hidden_states, block_output, injection = layer(
532
hidden_states=hidden_states,
533
prev_block_output=block_output,
vllm/models/qwen4_exp/nvidia/ngram_embedding.py
new
+939
@@ -0,0 +1,939 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Qwen4Exp n-gram embeddings with device and pinned-host storage."""
4
+
5
+from abc import ABC, abstractmethod
6
+from collections.abc import Iterable
7
+from typing import ClassVar
8
+
9
+import torch
10
+import torch.nn.functional as F
11
+from torch import nn
12
+
13
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
14
+from vllm.config import get_current_vllm_config
15
+from vllm.distributed import get_dp_group, get_etp_group, get_tp_group
16
+from vllm.forward_context import DPMetadata, get_forward_context
17
+from vllm.logger import init_logger
18
+from vllm.model_executor.layers.quantization.base_config import (
19
+ QuantizationConfig,
20
+ QuantizeMethodBase,
21
+)
22
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
23
+from vllm.model_executor.layers.quantization.modelopt import (
24
+ ModelOptMixedPrecisionConfig,
25
+ ModelOptQuantConfigBase,
26
+)
27
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
28
+ create_fp8_scale_parameter,
29
+)
30
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
+ is_layer_skipped,
32
+)
33
+from vllm.model_executor.models.utils import AutoWeightsLoader
34
+from vllm.model_executor.parameter import (
35
+ ModelWeightParameter,
36
+ PerTensorScaleParameter,
37
+)
38
+from vllm.model_executor.utils import set_weight_attrs
39
+from vllm.transformers_utils.configs.qwen4_exp import (
40
+ Qwen4ExpTextConfig,
41
+)
42
+from vllm.triton_utils import tl, triton
43
+from vllm.utils.platform_utils import is_uva_available
44
+from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
45
+
46
+from ..common.ple import PLEVocabParallelEmbedding
47
+from .ops.ple import ple_ngram_ids
48
+
49
+logger = init_logger(__name__)
50
+
51
+
52
+class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC):
53
+ """ETP-sharded PLE table shared by device and pinned-host backends."""
54
+
55
+ supports_prefetch: ClassVar[bool] = False
56
+
57
+ def __init__(
58
+ self,
59
+ num_embeddings: int,
60
+ embedding_dim: int,
61
+ *,
62
+ params_dtype: torch.dtype,
63
+ padding_size: int,
64
+ prefix: str,
65
+ embedding_method: "Qwen4ExpPLEEmbeddingMethod",
66
+ num_ngram_heads: int = 1,
67
+ max_total_tokens: int = 0,
68
+ data_parallel_rank: int = 0,
69
+ ) -> None:
70
+ del num_ngram_heads, max_total_tokens
71
+ super().__init__(
72
+ num_embeddings,
73
+ embedding_dim,
74
+ params_dtype=params_dtype,
75
+ padding_size=padding_size,
76
+ prefix=prefix,
77
+ quant_method=embedding_method,
78
+ parallel_group=get_etp_group(),
79
+ )
80
+ self.embedding_method = embedding_method
81
+ self.data_parallel_rank = data_parallel_rank
82
+ tp_size = get_tp_group().world_size
83
+ if self.tp_size % tp_size:
84
+ raise ValueError(
85
+ "ETP size must be divisible by TP size, but got "
86
+ f"ETP={self.tp_size} and TP={tp_size}"
87
+ )
88
+ self.etp_data_parallel_size = self.tp_size // tp_size
89
+
90
+ @abstractmethod
91
+ def allocate_embedding_weight(
92
+ self,
93
+ num_embeddings: int,
94
+ embedding_dim: int,
95
+ dtype: torch.dtype,
96
+ ) -> torch.Tensor:
97
+ """Allocate storage for the complete embedding weight."""
98
+ raise NotImplementedError
99
+
100
+ def dequantize(
101
+ self,
102
+ embeddings: torch.Tensor,
103
+ output_dtype: torch.dtype,
104
+ ) -> torch.Tensor:
105
+ """Delegate storage-format conversion to the embedding method."""
106
+ return self.embedding_method.dequantize(self, embeddings, output_dtype)
107
+
108
+ def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]:
109
+ """Return the per-DP slot size and this rank's slot offset."""
110
+ if self.etp_data_parallel_size == 1:
111
+ return local_num_tokens, 0
112
+ dp_metadata: DPMetadata | None = get_forward_context().dp_metadata
113
+ if dp_metadata is None:
114
+ raise RuntimeError("ETP spanning DP requires DP token metadata")
115
+ group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * (
116
+ self.etp_data_parallel_size
117
+ )
118
+ group_end = group_start + self.etp_data_parallel_size
119
+ token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist()
120
+ group_counts = token_counts[group_start:group_end]
121
+ slot_size = max(group_counts)
122
+ dp_rank = get_dp_group().rank_in_group
123
+ return slot_size, dp_rank * slot_size
124
+
125
+ def _gather_dp_ids(
126
+ self,
127
+ ngram_ids: torch.Tensor,
128
+ slot_size: int,
129
+ ) -> torch.Tensor:
130
+ """Gather DP-local IDs that share one ETP-sharded PLE table."""
131
+ if self.etp_data_parallel_size == 1:
132
+ return ngram_ids
133
+ if ngram_ids.shape[0] < slot_size:
134
+ padding = ngram_ids.new_zeros(
135
+ slot_size - ngram_ids.shape[0], ngram_ids.shape[1]
136
+ )
137
+ ngram_ids = torch.cat((ngram_ids, padding), dim=0)
138
+ return get_dp_group().all_gather(ngram_ids, dim=0)
139
+
140
+ def _select_embeddings(
141
+ self,
142
+ embeddings: torch.Tensor,
143
+ local_num_tokens: int,
144
+ slot_offset: int,
145
+ ) -> torch.Tensor:
146
+ """Select this DP rank's rows from the ETP-reduced embeddings."""
147
+ if self.etp_data_parallel_size == 1:
148
+ return embeddings
149
+ return embeddings.narrow(0, slot_offset, local_num_tokens)
150
+
151
+ @abstractmethod
152
+ def start_prefetch(
153
+ self,
154
+ hidden_states: torch.Tensor,
155
+ ngram_ids: torch.Tensor,
156
+ ) -> None:
157
+ """Start an asynchronous lookup when supported."""
158
+ raise NotImplementedError
159
+
160
+
161
+class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
162
+ """Quantization interface shared by resident and pinned PLE tables."""
163
+
164
+ # PLE post-load processing only validates scales in their current storage.
165
+ requires_device_loading: bool = False
166
+
167
+ @staticmethod
168
+ def from_quant_config(
169
+ quant_config: QuantizationConfig | None,
170
+ prefix: str,
171
+ embedding_dtype: str | None = None,
172
+ ) -> "Qwen4ExpPLEEmbeddingMethod":
173
+ """Select the concrete PLE embedding format for a layer."""
174
+ if embedding_dtype == "float8_e4m3fn":
175
+ return Qwen4ExpPLEFp8EmbeddingMethod()
176
+ if quant_config is None:
177
+ return Qwen4ExpPLEUnquantizedEmbeddingMethod()
178
+ if isinstance(quant_config, ModelOptMixedPrecisionConfig):
179
+ if quant_config._resolve_quant_algo(prefix) == "FP8":
180
+ return Qwen4ExpPLEFp8EmbeddingMethod()
181
+ return Qwen4ExpPLEUnquantizedEmbeddingMethod()
182
+ if isinstance(
183
+ quant_config, ModelOptQuantConfigBase
184
+ ) and quant_config.is_layer_excluded(prefix):
185
+ return Qwen4ExpPLEUnquantizedEmbeddingMethod()
186
+ if not isinstance(quant_config, Fp8Config):
187
+ raise NotImplementedError(
188
+ "Qwen4Exp PLE embedding does not support quantization config "
189
+ f"{type(quant_config).__name__}"
190
+ )
191
+
192
+ ignored_layers = quant_config.ignored_layers
193
+ if is_layer_skipped(
194
+ prefix,
195
+ ignored_layers,
196
+ quant_config.packed_modules_mapping,
197
+ match_mode=quant_config.ignored_layers_match_mode,
198
+ ):
199
+ return Qwen4ExpPLEUnquantizedEmbeddingMethod()
200
+ # PLE checkpoint shards form one runtime embedding parameter.
201
+ shard_prefix = f"{prefix}.shard_"
202
+ if any(name.startswith(shard_prefix) for name in ignored_layers):
203
+ return Qwen4ExpPLEUnquantizedEmbeddingMethod()
204
+ if not quant_config.is_checkpoint_fp8_serialized:
205
+ raise NotImplementedError(
206
+ "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
207
+ )
208
+ return Qwen4ExpPLEFp8EmbeddingMethod()
209
+
210
+ def apply(
211
+ self,
212
+ layer: nn.Module,
213
+ x: torch.Tensor,
214
+ bias: torch.Tensor | None = None,
215
+ ) -> torch.Tensor:
216
+ raise NotImplementedError("PLE weights only support embedding lookup")
217
+
218
+ def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
219
+ return F.embedding(input_, layer.weight)
220
+
221
+ @abstractmethod
222
+ def dequantize(
223
+ self,
224
+ layer: nn.Module,
225
+ embeddings: torch.Tensor,
226
+ output_dtype: torch.dtype,
227
+ ) -> torch.Tensor:
228
+ """Convert looked-up PLE rows to the activation dtype."""
229
+ raise NotImplementedError
230
+
231
+
232
+class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
233
+ """Unquantized PLE embedding storage and lookup semantics."""
234
+
235
+ def create_weights(
236
+ self,
237
+ layer: Qwen4ExpPLEEmbedding,
238
+ input_size_per_partition: int,
239
+ output_partition_sizes: list[int],
240
+ input_size: int,
241
+ output_size: int,
242
+ params_dtype: torch.dtype,
243
+ **extra_weight_attrs,
244
+ ) -> None:
245
+ del input_size, output_size
246
+ weight = nn.Parameter(
247
+ layer.allocate_embedding_weight(
248
+ sum(output_partition_sizes),
249
+ input_size_per_partition,
250
+ params_dtype,
251
+ ),
252
+ requires_grad=False,
253
+ )
254
+ set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
255
+ set_weight_attrs(weight, extra_weight_attrs)
256
+ layer.register_parameter("weight", weight)
257
+
258
+ def dequantize(
259
+ self,
260
+ layer: nn.Module,
261
+ embeddings: torch.Tensor,
262
+ output_dtype: torch.dtype,
263
+ ) -> torch.Tensor:
264
+ del layer, output_dtype
265
+ return embeddings
266
+
267
+
268
+class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
269
+ """FP8 PLE embedding with one global checkpoint scale."""
270
+
271
+ def create_weights(
272
+ self,
273
+ layer: Qwen4ExpPLEEmbedding,
274
+ input_size_per_partition: int,
275
+ output_partition_sizes: list[int],
276
+ input_size: int,
277
+ output_size: int,
278
+ params_dtype: torch.dtype,
279
+ **extra_weight_attrs,
280
+ ) -> None:
281
+ del input_size, output_size, params_dtype
282
+ weight_loader = extra_weight_attrs.get("weight_loader")
283
+ weight = ModelWeightParameter(
284
+ data=layer.allocate_embedding_weight(
285
+ sum(output_partition_sizes),
286
+ input_size_per_partition,
287
+ torch.float8_e4m3fn,
288
+ ),
289
+ input_dim=1,
290
+ output_dim=0,
291
+ weight_loader=weight_loader,
292
+ )
293
+ layer.register_parameter("weight", weight)
294
+
295
+ weight_scale = create_fp8_scale_parameter(
296
+ PerTensorScaleParameter,
297
+ output_partition_sizes,
298
+ input_size_per_partition,
299
+ None,
300
+ weight_loader,
301
+ scale_dtype=torch.float32,
302
+ )
303
+ layer.register_parameter("weight_scale", weight_scale)
304
+
305
+ def process_weights_after_loading(self, layer: nn.Module) -> None:
306
+ """Reject FP8 PLE checkpoints without a global scale."""
307
+ sentinel = torch.finfo(torch.float32).min
308
+ if torch.any(layer.weight_scale == sentinel):
309
+ raise ValueError("FP8 PLE checkpoint is missing its global scale")
310
+
311
+ def dequantize(
312
+ self,
313
+ layer: nn.Module,
314
+ embeddings: torch.Tensor,
315
+ output_dtype: torch.dtype,
316
+ ) -> torch.Tensor:
317
+ weight_scale = getattr(layer, "weight_scale", None)
318
+ if weight_scale is None:
319
+ raise RuntimeError("FP8 PLE embedding is missing its global scale")
320
+ if weight_scale.device != embeddings.device:
321
+ raise RuntimeError("FP8 PLE embedding scale must be on the output device")
322
+ return embeddings.to(output_dtype) * weight_scale.to(output_dtype)
323
+
324
+
325
+class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding):
326
+ """PLE table allocated on the active model device."""
327
+
328
+ def allocate_embedding_weight(
329
+ self,
330
+ num_embeddings: int,
331
+ embedding_dim: int,
332
+ dtype: torch.dtype,
333
+ ) -> torch.Tensor:
334
+ """Allocate the complete PLE weight on the active device."""
335
+ return torch.empty(num_embeddings, embedding_dim, dtype=dtype)
336
+
337
+ def start_prefetch(
338
+ self,
339
+ hidden_states: torch.Tensor,
340
+ ngram_ids: torch.Tensor,
341
+ ) -> None:
342
+ """Resident embedding prefetch is a no-op."""
343
+ return None
344
+
345
+ def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor:
346
+ """Gather ETP inputs, look up embeddings, and select local rows."""
347
+ slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
348
+ gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
349
+ embeddings = super().forward(gathered_ids)
350
+ return self._select_embeddings(
351
+ embeddings,
352
+ ngram_ids.shape[0],
353
+ slot_offset,
354
+ )
355
+
356
+
357
+@triton.jit
358
+def _lookup_ple_embedding_from_pinned_kernel(
359
+ weight_ptr,
360
+ ids_ptr,
361
+ output_ptr,
362
+ embedding_dim,
363
+ tp_vocab_start,
364
+ tp_vocab_end,
365
+ BLOCK_D: tl.constexpr,
366
+):
367
+ """Look up TP-owned PLE rows through a CUDA view of pinned host memory."""
368
+ row_id = tl.program_id(0)
369
+ global_idx = tl.load(ids_ptr + row_id)
370
+ in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end)
371
+ local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0)
372
+ offsets = tl.arange(0, BLOCK_D)
373
+ store_mask = offsets < embedding_dim
374
+ load_mask = store_mask & in_range
375
+ values = tl.load(
376
+ weight_ptr + local_idx * embedding_dim + offsets,
377
+ mask=load_mask,
378
+ other=0.0,
379
+ )
380
+ tl.store(
381
+ output_ptr + row_id * embedding_dim + offsets,
382
+ values,
383
+ mask=store_mask,
384
+ )
385
+
386
+
387
+class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):
388
+ """PLE table loaded into pinned CPU memory and looked up through UVA."""
389
+
390
+ supports_prefetch: ClassVar[bool] = True
391
+
392
+ def __init__(
393
+ self,
394
+ num_embeddings: int,
395
+ embedding_dim: int,
396
+ *,
397
+ params_dtype: torch.dtype,
398
+ padding_size: int,
399
+ prefix: str,
400
+ embedding_method: Qwen4ExpPLEEmbeddingMethod,
401
+ num_ngram_heads: int = 1,
402
+ max_total_tokens: int = 0,
403
+ data_parallel_rank: int = 0,
404
+ ) -> None:
405
+ if not is_uva_available():
406
+ raise RuntimeError("Engram CPU offload requires UVA support")
407
+ super().__init__(
408
+ num_embeddings,
409
+ embedding_dim,
410
+ params_dtype=params_dtype,
411
+ padding_size=padding_size,
412
+ prefix=prefix,
413
+ embedding_method=embedding_method,
414
+ num_ngram_heads=num_ngram_heads,
415
+ max_total_tokens=max_total_tokens,
416
+ data_parallel_rank=data_parallel_rank,
417
+ )
418
+ self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight)
419
+ self._block_d = triton.next_power_of_2(self.embedding_dim)
420
+ self._prefetch_stream = torch.cuda.Stream(device=self._uva_weight.device)
421
+ self._prefetch_buffer = torch.empty(
422
+ max_total_tokens * self.etp_data_parallel_size,
423
+ num_ngram_heads,
424
+ self.embedding_dim,
425
+ dtype=self.weight.dtype,
426
+ device=self._uva_weight.device,
427
+ )
428
+ self._output_dim = num_ngram_heads * self.embedding_dim
429
+
430
+ def allocate_embedding_weight(
431
+ self,
432
+ num_embeddings: int,
433
+ embedding_dim: int,
434
+ dtype: torch.dtype,
435
+ ) -> torch.Tensor:
436
+ """Allocate the complete PLE weight directly in pinned CPU memory."""
437
+ return torch.empty(
438
+ num_embeddings,
439
+ embedding_dim,
440
+ dtype=dtype,
441
+ device="cpu",
442
+ pin_memory=True,
443
+ )
444
+
445
+ def _lookup(
446
+ self,
447
+ input_ids: torch.Tensor,
448
+ output: torch.Tensor | None = None,
449
+ ) -> torch.Tensor:
450
+ """Look up local ETP rows while preserving the weight storage dtype."""
451
+ expected_shape = (*input_ids.shape, self.embedding_dim)
452
+ if output is None:
453
+ output = torch.empty(
454
+ expected_shape,
455
+ dtype=self.weight.dtype,
456
+ device=input_ids.device,
457
+ )
458
+ elif (
459
+ tuple(output.shape) != expected_shape
460
+ or output.dtype != self.weight.dtype
461
+ or output.device != input_ids.device
462
+ ):
463
+ raise ValueError(
464
+ "PLE prefetch output must match the input shape, weight dtype, "
465
+ "and input device"
466
+ )
467
+
468
+ flat_ids = input_ids.reshape(-1).long()
469
+ if flat_ids.numel():
470
+ _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
471
+ self._uva_weight,
472
+ flat_ids,
473
+ output,
474
+ self.embedding_dim,
475
+ self.shard_indices.org_vocab_start_index,
476
+ self.shard_indices.org_vocab_end_index,
477
+ BLOCK_D=self._block_d,
478
+ )
479
+ return output
480
+
481
+ def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
482
+ """Combine pinned lookup results owned by different ETP ranks."""
483
+ if self.tp_size == 1:
484
+ return embeddings
485
+ assert self.parallel_group is not None
486
+ if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
487
+ # Each vocabulary row has one owner, so reduce the raw FP8 bytes.
488
+ reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
489
+ return reduced.view(embeddings.dtype)
490
+ return self.parallel_group.all_reduce(embeddings)
491
+
492
+ @eager_break_during_capture
493
+ def start_prefetch(
494
+ self,
495
+ hidden_states: torch.Tensor,
496
+ ngram_ids: torch.Tensor,
497
+ ) -> None:
498
+ """Gather ETP IDs and launch their UVA lookup on the side stream."""
499
+ slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0])
500
+ gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
501
+ active_output = self._prefetch_buffer[: gathered_ids.shape[0]]
502
+ prefetch_stream = self._prefetch_stream
503
+ prefetch_stream.wait_stream(torch.cuda.current_stream())
504
+ gathered_ids.record_stream(prefetch_stream)
505
+ with torch.cuda.stream(prefetch_stream):
506
+ self._lookup(gathered_ids, output=active_output)
507
+
508
+ @eager_break_during_capture
509
+ def _finalize_prefetch(
510
+ self,
511
+ prefetch_output: torch.Tensor,
512
+ output: torch.Tensor,
513
+ ) -> None:
514
+ """Join the side stream, reduce ETP shards, and select local rows."""
515
+ torch.cuda.current_stream().wait_stream(self._prefetch_stream)
516
+ slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0])
517
+ active_output = prefetch_output[: slot_size * self.etp_data_parallel_size]
518
+ embeddings = self._reduce_etp_embeddings(active_output)
519
+ embeddings = self._select_embeddings(
520
+ embeddings,
521
+ output.shape[0],
522
+ slot_offset,
523
+ )
524
+ output.copy_(embeddings.flatten(-2))
525
+
526
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
527
+ """Finish the pinned lookup into graph-owned output storage."""
528
+ output = self._prefetch_buffer.new_empty(
529
+ (hidden_states.shape[0], self._output_dim)
530
+ )
531
+ self._finalize_prefetch(self._prefetch_buffer, output)
532
+ return output
533
+
534
+
535
+class Qwen4ExpNGramEmbedding(nn.Module):
536
+ _MASK64 = (1 << 64) - 1
537
+ _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
538
+ _SPLITMIX_M1 = 0xBF58476D1CE4E5B9
539
+ _SPLITMIX_M2 = 0x94D049BB133111EB
540
+ _PLE_LAYER_PRIME = 10007
541
+
542
+ @classmethod
543
+ def _splitmix64(cls, value: int) -> int:
544
+ """Mix an integer into a deterministic unsigned 64-bit value."""
545
+ value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64
546
+ value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64
547
+ value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64
548
+ return (value ^ (value >> 31)) & cls._MASK64
549
+
550
+ @staticmethod
551
+ def _is_prime_64(value: int) -> bool:
552
+ """Return whether a 64-bit integer is prime."""
553
+ if value < 2:
554
+ return False
555
+ for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
556
+ if value % prime == 0:
557
+ return value == prime
558
+ exponent = value - 1
559
+ shifts = 0
560
+ while exponent % 2 == 0:
561
+ exponent //= 2
562
+ shifts += 1
563
+ for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022):
564
+ if base % value == 0:
565
+ continue
566
+ witness = pow(base, exponent, value)
567
+ if witness in (1, value - 1):
568
+ continue
569
+ for _ in range(shifts - 1):
570
+ witness = pow(witness, 2, value)
571
+ if witness == value - 1:
572
+ break
573
+ else:
574
+ return False
575
+ return True
576
+
577
+ @classmethod
578
+ def _nth_prime_after(cls, start: int, count: int) -> int:
579
+ """Return the ``count``-th prime strictly greater than ``start``."""
580
+ prime = int(start)
581
+ for _ in range(count):
582
+ candidate = prime + 1
583
+ if candidate <= 2:
584
+ prime = 2
585
+ continue
586
+ if candidate % 2 == 0:
587
+ candidate += 1
588
+ while not cls._is_prime_64(candidate):
589
+ candidate += 2
590
+ prime = candidate
591
+ return prime
592
+
593
+ @classmethod
594
+ def _make_layer_multipliers(
595
+ cls,
596
+ *,
597
+ ngram_size: int,
598
+ unigram_vocab_size: int,
599
+ seed: int,
600
+ ple_dense_layer_id: int,
601
+ ) -> list[int]:
602
+ """Build deterministic hash multipliers for one PLE layer."""
603
+ max_multiplier = ((1 << 63) - 1) // unigram_vocab_size
604
+ half_bound = max(1, max_multiplier // 2)
605
+ base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id
606
+ multipliers = []
607
+ for index in range(ngram_size):
608
+ value = base_seed + cls._SPLITMIX_GAMMA * (index + 1)
609
+ multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1)
610
+ return multipliers
611
+
612
+ @classmethod
613
+ def _make_vocab_layout(
614
+ cls,
615
+ *,
616
+ ngram_vocab_size_base: int,
617
+ ngram_heads: int,
618
+ ple_dense_layer_id: int,
619
+ ) -> tuple[list[int], list[int], int]:
620
+ """Build per-head vocabulary sizes, offsets, and total row count."""
621
+ sizes: list[int] = []
622
+ offsets: list[int] = []
623
+ offset = 0
624
+ for local_head in range(ngram_heads):
625
+ global_head = ple_dense_layer_id * ngram_heads + local_head
626
+ size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1)
627
+ sizes.append(size)
628
+ offsets.append(offset)
629
+ offset += size
630
+ return sizes, offsets, offset
631
+
632
+ def __init__(
633
+ self,
634
+ config: Qwen4ExpTextConfig,
635
+ embedding_dim: int,
636
+ ple_dense_layer_id: int,
637
+ max_total_tokens: int,
638
+ *,
639
+ data_parallel_rank: int,
640
+ prefix: str,
641
+ quant_config: QuantizationConfig | None = None,
642
+ params_dtype: torch.dtype | None = None,
643
+ ) -> None:
644
+ super().__init__()
645
+ self.embedding_dim = embedding_dim
646
+ self.ngram_size = int(config.ngram_size)
647
+ self.heads_per_ngram = int(config.heads_per_ngram)
648
+ self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
649
+ if self.ngram_size < 2:
650
+ raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}")
651
+ if self.heads_per_ngram <= 0:
652
+ raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}")
653
+ if embedding_dim % self.ngram_heads:
654
+ raise ValueError(
655
+ "ple_embed_dim must be divisible by total ngram heads: "
656
+ f"{embedding_dim} % {self.ngram_heads} != 0"
657
+ )
658
+ self.head_dim = embedding_dim // self.ngram_heads
659
+ self.eos_token_id = int(config.eos_token_id)
660
+ self.unigram_vocab_size = int(config.vocab_size)
661
+ self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512))
662
+ if self.split_ngram_parts <= 0:
663
+ raise ValueError("split_ngram_parts must be positive")
664
+
665
+ multipliers = self._make_layer_multipliers(
666
+ ngram_size=self.ngram_size,
667
+ unigram_vocab_size=self.unigram_vocab_size,
668
+ seed=int(getattr(config, "seed", 1234)),
669
+ ple_dense_layer_id=ple_dense_layer_id,
670
+ )
671
+ self.register_buffer(
672
+ "layer_multipliers",
673
+ torch.tensor(multipliers, dtype=torch.long),
674
+ persistent=True,
675
+ )
676
+
677
+ sizes, offsets, total_vocab_size = self._make_vocab_layout(
678
+ ngram_vocab_size_base=int(config.ngram_vocab_size_base),
679
+ ngram_heads=self.ngram_heads,
680
+ ple_dense_layer_id=ple_dense_layer_id,
681
+ )
682
+ self.register_buffer(
683
+ "ngram_heads_vocab_sizes",
684
+ torch.tensor(sizes, dtype=torch.long),
685
+ persistent=True,
686
+ )
687
+ self.register_buffer(
688
+ "ngram_heads_offsets",
689
+ torch.tensor(offsets, dtype=torch.long),
690
+ persistent=True,
691
+ )
692
+ divisor = int(config.make_ngram_vocab_size_divisible_by)
693
+ padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
694
+ embedding_prefix = f"{prefix}.ngram_embedding"
695
+ embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config(
696
+ quant_config,
697
+ embedding_prefix,
698
+ getattr(config, "ple_embedding_dtype", None),
699
+ )
700
+ if params_dtype is None:
701
+ params_dtype = torch.get_default_dtype()
702
+ engram_config = get_current_vllm_config().engram_config
703
+ embedding_cls = (
704
+ Qwen4ExpPLEPinnedHostEmbedding
705
+ if engram_config is not None and engram_config.cpu_offload
706
+ else Qwen4ExpPLEDeviceEmbedding
707
+ )
708
+ self.ngram_embedding = embedding_cls(
709
+ padded_vocab_size,
710
+ self.head_dim,
711
+ params_dtype=params_dtype,
712
+ padding_size=divisor,
713
+ prefix=embedding_prefix,
714
+ embedding_method=embedding_quant_method,
715
+ num_ngram_heads=self.ngram_heads,
716
+ max_total_tokens=max_total_tokens,
717
+ data_parallel_rank=data_parallel_rank,
718
+ )
719
+ weight = self.ngram_embedding.weight
720
+ logger.info(
721
+ "Initialized PLE embedding %s: quantization_method=%s, "
722
+ "weight_dtype=%s, weight_device=%s, pinned=%s",
723
+ embedding_prefix,
724
+ type(embedding_quant_method).__name__,
725
+ weight.dtype,
726
+ weight.device,
727
+ weight.is_pinned(),
728
+ )
729
+
730
+ @staticmethod
731
+ def _shift_precompute(
732
+ tokens: torch.Tensor, eos_token_id: int
733
+ ) -> tuple[torch.Tensor, torch.Tensor]:
734
+ if tokens.dim() != 2:
735
+ raise ValueError("tokens must be a 2D tensor")
736
+ batch_size, seq_len = tokens.shape
737
+ positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64)
738
+ eos_positions = torch.where(tokens == eos_token_id, positions, -1)
739
+ previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
740
+ previous_eos = torch.cat(
741
+ [
742
+ eos_positions.new_full((batch_size, 1), -1),
743
+ previous_eos_inclusive[:, :-1],
744
+ ],
745
+ dim=1,
746
+ )
747
+ return positions, positions.unsqueeze(0) - previous_eos - 1
748
+
749
+ @staticmethod
750
+ def _shift_apply(
751
+ tokens: torch.Tensor,
752
+ positions: torch.Tensor,
753
+ position_in_segment: torch.Tensor,
754
+ shift: int,
755
+ eos_token_id: int,
756
+ ) -> torch.Tensor:
757
+ if shift == 0:
758
+ return tokens
759
+ source = positions - shift
760
+ gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1)
761
+ shifted = tokens.gather(1, gather_indices)
762
+ valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift)
763
+ return torch.where(valid, shifted, tokens.new_full((), eos_token_id))
764
+
765
+ def compute_ngram_ids(
766
+ self,
767
+ input_ids: torch.Tensor,
768
+ query_start_loc: torch.Tensor,
769
+ ngram_context: torch.Tensor,
770
+ output: torch.Tensor | None = None,
771
+ ) -> torch.Tensor:
772
+ """Compute n-gram embedding indices for the current request layout."""
773
+ input_ids = input_ids.reshape(-1)
774
+ num_reqs = query_start_loc.numel() - 1
775
+ num_tokens = input_ids.shape[0]
776
+
777
+ if input_ids.is_cuda:
778
+ return ple_ngram_ids(
779
+ input_ids=input_ids,
780
+ query_start_loc=query_start_loc,
781
+ ngram_context=ngram_context,
782
+ layer_multipliers=self.layer_multipliers,
783
+ ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes,
784
+ ngram_heads_offsets=self.ngram_heads_offsets,
785
+ eos_token_id=self.eos_token_id,
786
+ heads_per_ngram=self.heads_per_ngram,
787
+ output=output,
788
+ )
789
+ input_ids = input_ids.long()
790
+ query_start_loc = query_start_loc.long()
791
+ positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64)
792
+ packed = torch.full(
793
+ (num_reqs, num_tokens),
794
+ self.eos_token_id,
795
+ device=input_ids.device,
796
+ dtype=torch.int64,
797
+ )
798
+ request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1
799
+ request_indices.clamp_(max=num_reqs - 1)
800
+ columns = (positions - query_start_loc[request_indices]).clamp(
801
+ 0, packed.shape[1] - 1
802
+ )
803
+ packed[request_indices, columns] = input_ids
804
+ ngram_context = ngram_context[:num_reqs].to(
805
+ device=input_ids.device, dtype=torch.long
806
+ )
807
+
808
+ context = torch.cat([ngram_context, packed], dim=-1)
809
+ positions_2d, position_in_segment = self._shift_precompute(
810
+ context, self.eos_token_id
811
+ )
812
+ shifted = [context]
813
+ for shift in range(1, self.ngram_size):
814
+ shifted.append(
815
+ self._shift_apply(
816
+ context,
817
+ positions_2d,
818
+ position_in_segment,
819
+ shift,
820
+ self.eos_token_id,
821
+ )
822
+ )
823
+ adjusted_columns = columns + self.ngram_size - 1
824
+ id_blocks = []
825
+ for ngram in range(2, self.ngram_size + 1):
826
+ start = (ngram - 2) * self.heads_per_ngram
827
+ end = start + self.heads_per_ngram
828
+ mixed = shifted[0] * self.layer_multipliers[0]
829
+ for index in range(1, ngram):
830
+ mixed = torch.bitwise_xor(
831
+ mixed, shifted[index] * self.layer_multipliers[index]
832
+ )
833
+ sizes = self.ngram_heads_vocab_sizes[start:end]
834
+ offsets = self.ngram_heads_offsets[start:end]
835
+ ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
836
+ id_blocks.append(ids[request_indices, adjusted_columns])
837
+ return torch.cat(id_blocks, dim=-1)
838
+
839
+ def forward(
840
+ self,
841
+ hidden_states: torch.Tensor,
842
+ input_ids: torch.Tensor,
843
+ query_start_loc: torch.Tensor,
844
+ ngram_context: torch.Tensor,
845
+ ) -> torch.Tensor:
846
+ embedding = self.ngram_embedding
847
+ if embedding.supports_prefetch:
848
+ return embedding(hidden_states)
849
+ ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
850
+ return self.ngram_embedding(ngram_ids).flatten(-2)
851
+
852
+ def start_prefetch(
853
+ self,
854
+ hidden_states: torch.Tensor,
855
+ input_ids: torch.Tensor,
856
+ query_start_loc: torch.Tensor,
857
+ ngram_context: torch.Tensor,
858
+ ) -> None:
859
+ """Start the pinned lookup while the preceding decoder layer runs."""
860
+ embedding = self.ngram_embedding
861
+ if not embedding.supports_prefetch:
862
+ return
863
+ ngram_ids = self.compute_ngram_ids(
864
+ input_ids,
865
+ query_start_loc,
866
+ ngram_context,
867
+ )
868
+ embedding.start_prefetch(hidden_states, ngram_ids)
869
+
870
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
871
+ """Load hash buffers and checkpoint-split embedding rows."""
872
+
873
+ persistent_buffers = {
874
+ "layer_multipliers": self.layer_multipliers,
875
+ "ngram_heads_offsets": self.ngram_heads_offsets,
876
+ "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
877
+ }
878
+ loaded: set[str] = set()
879
+ regular_weights: list[tuple[str, torch.Tensor]] = []
880
+ shard_prefix = "ngram_embedding.shard_"
881
+
882
+ for name, loaded_weight in weights:
883
+ leaf_name = name.rsplit(".", 1)[-1]
884
+ if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
885
+ continue
886
+ if name in persistent_buffers:
887
+ buffer = persistent_buffers[name]
888
+ if buffer.shape != loaded_weight.shape:
889
+ raise ValueError(
890
+ f"Shape mismatch for {name}: expected "
891
+ f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}"
892
+ )
893
+ buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype))
894
+ loaded.add(name)
895
+ continue
896
+ if name.startswith(shard_prefix) and name.endswith(".weight"):
897
+ shard_text = name[len(shard_prefix) : -len(".weight")]
898
+ if not shard_text.isdigit():
899
+ regular_weights.append((name, loaded_weight))
900
+ continue
901
+ shard_index = int(shard_text)
902
+ if shard_index >= self.split_ngram_parts:
903
+ raise ValueError(
904
+ f"PLE embedding shard index {shard_index} exceeds "
905
+ f"split_ngram_parts={self.split_ngram_parts}"
906
+ )
907
+ embedding = self.ngram_embedding
908
+ shard_size = (
909
+ embedding.org_vocab_size + self.split_ngram_parts - 1
910
+ ) // self.split_ngram_parts
911
+ checkpoint_start = shard_index * shard_size
912
+ expected_rows = max(
913
+ 0,
914
+ min(shard_size, embedding.org_vocab_size - checkpoint_start),
915
+ )
916
+ expected_shape = (expected_rows, embedding.embedding_dim)
917
+ if tuple(loaded_weight.shape) != expected_shape:
918
+ raise ValueError(
919
+ f"Shape mismatch for PLE embedding shard {shard_index}: "
920
+ f"expected {expected_shape}, got "
921
+ f"{tuple(loaded_weight.shape)}"
922
+ )
923
+ embedding.weight.weight_loader(
924
+ embedding.weight,
925
+ loaded_weight,
926
+ checkpoint_start=checkpoint_start,
927
+ )
928
+ loaded.add("ngram_embedding.weight")
929
+ continue
930
+ regular_weights.append((name, loaded_weight))
931
+
932
+ if regular_weights:
933
+ loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
934
+ return loaded
935
+
936
+
937
+__all__ = [
938
+ "Qwen4ExpNGramEmbedding",
939
+]
vllm/models/qwen4_exp/nvidia/ple_layer.py
+30
-467
@@ -2,10 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""GPU-resident Qwen4Exp position-learning enhancement layers."""
4
5
-from collections.abc import Iterable, Sequence
5
+from collections.abc import Sequence
6
7
import torch
8
-import torch.nn.functional as F
8
from torch import nn
9
10
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
@@ -18,24 +17,6 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
17
MambaStateShapeCalculator,
18
is_conv_state_dim_first,
19
)
21
-from vllm.model_executor.layers.quantization.base_config import (
22
- QuantizationConfig,
23
- QuantizeMethodBase,
24
-)
25
-from vllm.model_executor.layers.quantization.fp8 import Fp8Config
26
-from vllm.model_executor.layers.quantization.modelopt import (
27
- ModelOptMixedPrecisionConfig,
28
-)
29
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
30
- create_fp8_scale_parameter,
31
- create_fp8_weight_parameter,
32
- is_fp8,
33
-)
34
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
35
- is_layer_skipped,
36
-)
37
-from vllm.model_executor.models.utils import AutoWeightsLoader
38
-from vllm.model_executor.parameter import PerTensorScaleParameter
20
from vllm.transformers_utils.configs.qwen4_exp import (
21
Qwen4ExpTextConfig,
22
)
@@ -45,8 +26,8 @@ from vllm.v1.attention.backends.short_conv_attn import (
26
PleShortConvAttentionMetadata,
27
)
28
48
-from ..common.ple import PLEVocabParallelEmbedding
49
-from .ops.ple import ple_conv, ple_gate, ple_ngram_ids
29
+from .ngram_embedding import Qwen4ExpNGramEmbedding
30
+from .ops.ple import ple_conv, ple_gate
31
32
33
class Qwen4ExpPLEGroupedNorm(nn.Module):
@@ -82,437 +63,6 @@ class Qwen4ExpPLEGroupedNorm(nn.Module):
63
return (normalized * (1.0 + self.weight.float())).to(input_dtype)
64
65
85
-class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase):
86
- """FP8 PLE embedding with one global checkpoint scale."""
87
-
88
- def create_weights(
89
- self,
90
- layer: nn.Module,
91
- input_size_per_partition: int,
92
- output_partition_sizes: list[int],
93
- input_size: int,
94
- output_size: int,
95
- params_dtype: torch.dtype,
96
- **extra_weight_attrs,
97
- ) -> None:
98
- del input_size, output_size, params_dtype
99
- weight_loader = extra_weight_attrs.get("weight_loader")
100
- weight = create_fp8_weight_parameter(
101
- sum(output_partition_sizes), input_size_per_partition, weight_loader
102
- )
103
- layer.register_parameter("weight", weight)
104
-
105
- weight_scale = create_fp8_scale_parameter(
106
- PerTensorScaleParameter,
107
- output_partition_sizes,
108
- input_size_per_partition,
109
- None,
110
- weight_loader,
111
- scale_dtype=torch.float32,
112
- )
113
- layer.register_parameter("weight_scale", weight_scale)
114
-
115
- def process_weights_after_loading(self, layer: nn.Module) -> None:
116
- """Reject FP8 PLE checkpoints without a global scale."""
117
- sentinel = torch.finfo(torch.float32).min
118
- if torch.any(layer.weight_scale == sentinel):
119
- raise ValueError("FP8 PLE checkpoint is missing its global scale")
120
-
121
- def apply(
122
- self,
123
- layer: nn.Module,
124
- x: torch.Tensor,
125
- bias: torch.Tensor | None = None,
126
- ) -> torch.Tensor:
127
- raise NotImplementedError("PLE FP8 weights only support embedding lookup")
128
-
129
- def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
130
- return F.embedding(input_, layer.weight)
131
-
132
-
133
-def _get_ple_embedding_quant_method(
134
- quant_config: QuantizationConfig | None,
135
- prefix: str,
136
-) -> QuantizeMethodBase | None:
137
- """Select global-scale FP8 only for quantized PLE checkpoint shards."""
138
-
139
- if isinstance(quant_config, ModelOptMixedPrecisionConfig):
140
- if quant_config._resolve_quant_algo(prefix) == "FP8":
141
- return Qwen4ExpPLEFp8EmbeddingMethod()
142
- return None
143
-
144
- if not isinstance(quant_config, Fp8Config):
145
- return None
146
- if not quant_config.is_checkpoint_fp8_serialized:
147
- return None
148
-
149
- ignored_layers = quant_config.ignored_layers
150
- if is_layer_skipped(
151
- prefix,
152
- ignored_layers,
153
- quant_config.packed_modules_mapping,
154
- match_mode=quant_config.ignored_layers_match_mode,
155
- ):
156
- return None
157
- # PLE checkpoint shards form one runtime embedding parameter.
158
- shard_prefix = f"{prefix}.shard_"
159
- if any(name.startswith(shard_prefix) for name in ignored_layers):
160
- return None
161
- return Qwen4ExpPLEFp8EmbeddingMethod()
162
-
163
-
164
-class Qwen4ExpNGramEmbedding(nn.Module):
165
- _MASK64 = (1 << 64) - 1
166
- _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
167
- _SPLITMIX_M1 = 0xBF58476D1CE4E5B9
168
- _SPLITMIX_M2 = 0x94D049BB133111EB
169
- _PLE_LAYER_PRIME = 10007
170
-
171
- @classmethod
172
- def _splitmix64(cls, value: int) -> int:
173
- """Mix an integer into a deterministic unsigned 64-bit value."""
174
- value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64
175
- value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64
176
- value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64
177
- return (value ^ (value >> 31)) & cls._MASK64
178
-
179
- @staticmethod
180
- def _is_prime_64(value: int) -> bool:
181
- """Return whether a 64-bit integer is prime."""
182
- if value < 2:
183
- return False
184
- for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
185
- if value % prime == 0:
186
- return value == prime
187
- exponent = value - 1
188
- shifts = 0
189
- while exponent % 2 == 0:
190
- exponent //= 2
191
- shifts += 1
192
- for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022):
193
- if base % value == 0:
194
- continue
195
- witness = pow(base, exponent, value)
196
- if witness in (1, value - 1):
197
- continue
198
- for _ in range(shifts - 1):
199
- witness = pow(witness, 2, value)
200
- if witness == value - 1:
201
- break
202
- else:
203
- return False
204
- return True
205
-
206
- @classmethod
207
- def _nth_prime_after(cls, start: int, count: int) -> int:
208
- """Return the ``count``-th prime strictly greater than ``start``."""
209
- prime = int(start)
210
- for _ in range(count):
211
- candidate = prime + 1
212
- if candidate <= 2:
213
- prime = 2
214
- continue
215
- if candidate % 2 == 0:
216
- candidate += 1
217
- while not cls._is_prime_64(candidate):
218
- candidate += 2
219
- prime = candidate
220
- return prime
221
-
222
- @classmethod
223
- def _make_layer_multipliers(
224
- cls,
225
- *,
226
- ngram_size: int,
227
- unigram_vocab_size: int,
228
- seed: int,
229
- ple_dense_layer_id: int,
230
- ) -> list[int]:
231
- """Build deterministic hash multipliers for one PLE layer."""
232
- max_multiplier = ((1 << 63) - 1) // unigram_vocab_size
233
- half_bound = max(1, max_multiplier // 2)
234
- base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id
235
- multipliers = []
236
- for index in range(ngram_size):
237
- value = base_seed + cls._SPLITMIX_GAMMA * (index + 1)
238
- multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1)
239
- return multipliers
240
-
241
- @classmethod
242
- def _make_vocab_layout(
243
- cls,
244
- *,
245
- ngram_vocab_size_base: int,
246
- ngram_heads: int,
247
- ple_dense_layer_id: int,
248
- ) -> tuple[list[int], list[int], int]:
249
- """Build per-head vocabulary sizes, offsets, and total row count."""
250
- sizes: list[int] = []
251
- offsets: list[int] = []
252
- offset = 0
253
- for local_head in range(ngram_heads):
254
- global_head = ple_dense_layer_id * ngram_heads + local_head
255
- size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1)
256
- sizes.append(size)
257
- offsets.append(offset)
258
- offset += size
259
- return sizes, offsets, offset
260
-
261
- def __init__(
262
- self,
263
- config: Qwen4ExpTextConfig,
264
- embedding_dim: int,
265
- ple_dense_layer_id: int,
266
- prefix: str,
267
- quant_config: QuantizationConfig | None = None,
268
- params_dtype: torch.dtype | None = None,
269
- ) -> None:
270
- super().__init__()
271
- self.embedding_dim = embedding_dim
272
- self.ngram_size = int(config.ngram_size)
273
- self.heads_per_ngram = int(config.heads_per_ngram)
274
- self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
275
- if self.ngram_size < 2:
276
- raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}")
277
- if self.heads_per_ngram <= 0:
278
- raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}")
279
- if embedding_dim % self.ngram_heads:
280
- raise ValueError(
281
- "ple_embed_dim must be divisible by total ngram heads: "
282
- f"{embedding_dim} % {self.ngram_heads} != 0"
283
- )
284
- self.head_dim = embedding_dim // self.ngram_heads
285
- self.eos_token_id = int(config.eos_token_id)
286
- self.unigram_vocab_size = int(config.vocab_size)
287
- self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512))
288
- if self.split_ngram_parts <= 0:
289
- raise ValueError("split_ngram_parts must be positive")
290
-
291
- multipliers = self._make_layer_multipliers(
292
- ngram_size=self.ngram_size,
293
- unigram_vocab_size=self.unigram_vocab_size,
294
- seed=int(getattr(config, "seed", 1234)),
295
- ple_dense_layer_id=ple_dense_layer_id,
296
- )
297
- self.register_buffer(
298
- "layer_multipliers",
299
- torch.tensor(multipliers, dtype=torch.long),
300
- persistent=True,
301
- )
302
-
303
- sizes, offsets, total_vocab_size = self._make_vocab_layout(
304
- ngram_vocab_size_base=int(config.ngram_vocab_size_base),
305
- ngram_heads=self.ngram_heads,
306
- ple_dense_layer_id=ple_dense_layer_id,
307
- )
308
- self.register_buffer(
309
- "ngram_heads_vocab_sizes",
310
- torch.tensor(sizes, dtype=torch.long),
311
- persistent=True,
312
- )
313
- self.register_buffer(
314
- "ngram_heads_offsets",
315
- torch.tensor(offsets, dtype=torch.long),
316
- persistent=True,
317
- )
318
- divisor = int(config.make_ngram_vocab_size_divisible_by)
319
- padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
320
- self.ngram_embedding = PLEVocabParallelEmbedding(
321
- padded_vocab_size,
322
- self.head_dim,
323
- params_dtype=params_dtype,
324
- padding_size=divisor,
325
- prefix=f"{prefix}.ngram_embedding",
326
- quant_method=_get_ple_embedding_quant_method(
327
- quant_config, f"{prefix}.ngram_embedding"
328
- ),
329
- )
330
-
331
- @staticmethod
332
- def _shift_precompute(
333
- tokens: torch.Tensor, eos_token_id: int
334
- ) -> tuple[torch.Tensor, torch.Tensor]:
335
- if tokens.dim() != 2:
336
- raise ValueError("tokens must be a 2D tensor")
337
- batch_size, seq_len = tokens.shape
338
- positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64)
339
- eos_positions = torch.where(tokens == eos_token_id, positions, -1)
340
- previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
341
- previous_eos = torch.cat(
342
- [
343
- eos_positions.new_full((batch_size, 1), -1),
344
- previous_eos_inclusive[:, :-1],
345
- ],
346
- dim=1,
347
- )
348
- return positions, positions.unsqueeze(0) - previous_eos - 1
349
-
350
- @staticmethod
351
- def _shift_apply(
352
- tokens: torch.Tensor,
353
- positions: torch.Tensor,
354
- position_in_segment: torch.Tensor,
355
- shift: int,
356
- eos_token_id: int,
357
- ) -> torch.Tensor:
358
- if shift == 0:
359
- return tokens
360
- source = positions - shift
361
- gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1)
362
- shifted = tokens.gather(1, gather_indices)
363
- valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift)
364
- return torch.where(valid, shifted, tokens.new_full((), eos_token_id))
365
-
366
- def compute_ngram_ids(
367
- self,
368
- input_ids: torch.Tensor,
369
- query_start_loc: torch.Tensor,
370
- ngram_context: torch.Tensor,
371
- output: torch.Tensor | None = None,
372
- ) -> torch.Tensor:
373
- """Compute n-gram embedding indices for the current request layout."""
374
- input_ids = input_ids.reshape(-1)
375
- num_reqs = query_start_loc.numel() - 1
376
- num_tokens = input_ids.shape[0]
377
-
378
- if input_ids.is_cuda:
379
- return ple_ngram_ids(
380
- input_ids=input_ids,
381
- query_start_loc=query_start_loc,
382
- ngram_context=ngram_context,
383
- layer_multipliers=self.layer_multipliers,
384
- ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes,
385
- ngram_heads_offsets=self.ngram_heads_offsets,
386
- eos_token_id=self.eos_token_id,
387
- heads_per_ngram=self.heads_per_ngram,
388
- output=output,
389
- )
390
- input_ids = input_ids.long()
391
- query_start_loc = query_start_loc.long()
392
- positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64)
393
- packed = torch.full(
394
- (num_reqs, num_tokens),
395
- self.eos_token_id,
396
- device=input_ids.device,
397
- dtype=torch.int64,
398
- )
399
- request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1
400
- request_indices.clamp_(max=num_reqs - 1)
401
- columns = (positions - query_start_loc[request_indices]).clamp(
402
- 0, packed.shape[1] - 1
403
- )
404
- packed[request_indices, columns] = input_ids
405
- ngram_context = ngram_context[:num_reqs].to(
406
- device=input_ids.device, dtype=torch.long
407
- )
408
-
409
- context = torch.cat([ngram_context, packed], dim=-1)
410
- positions_2d, position_in_segment = self._shift_precompute(
411
- context, self.eos_token_id
412
- )
413
- shifted = [context]
414
- for shift in range(1, self.ngram_size):
415
- shifted.append(
416
- self._shift_apply(
417
- context,
418
- positions_2d,
419
- position_in_segment,
420
- shift,
421
- self.eos_token_id,
422
- )
423
- )
424
- adjusted_columns = columns + self.ngram_size - 1
425
- id_blocks = []
426
- for ngram in range(2, self.ngram_size + 1):
427
- start = (ngram - 2) * self.heads_per_ngram
428
- end = start + self.heads_per_ngram
429
- mixed = shifted[0] * self.layer_multipliers[0]
430
- for index in range(1, ngram):
431
- mixed = torch.bitwise_xor(
432
- mixed, shifted[index] * self.layer_multipliers[index]
433
- )
434
- sizes = self.ngram_heads_vocab_sizes[start:end]
435
- offsets = self.ngram_heads_offsets[start:end]
436
- ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
437
- id_blocks.append(ids[request_indices, adjusted_columns])
438
- return torch.cat(id_blocks, dim=-1)
439
-
440
- def forward(
441
- self,
442
- input_ids: torch.Tensor,
443
- query_start_loc: torch.Tensor,
444
- ngram_context: torch.Tensor,
445
- ) -> torch.Tensor:
446
- ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
447
- return self.ngram_embedding(ngram_ids).flatten(-2)
448
-
449
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
450
- """Load hash buffers and checkpoint-split embedding rows."""
451
-
452
- persistent_buffers = {
453
- "layer_multipliers": self.layer_multipliers,
454
- "ngram_heads_offsets": self.ngram_heads_offsets,
455
- "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
456
- }
457
- loaded: set[str] = set()
458
- regular_weights: list[tuple[str, torch.Tensor]] = []
459
- shard_prefix = "ngram_embedding.shard_"
460
-
461
- for name, loaded_weight in weights:
462
- leaf_name = name.rsplit(".", 1)[-1]
463
- if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
464
- continue
465
- if name in persistent_buffers:
466
- buffer = persistent_buffers[name]
467
- if buffer.shape != loaded_weight.shape:
468
- raise ValueError(
469
- f"Shape mismatch for {name}: expected "
470
- f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}"
471
- )
472
- buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype))
473
- loaded.add(name)
474
- continue
475
- if name.startswith(shard_prefix) and name.endswith(".weight"):
476
- shard_text = name[len(shard_prefix) : -len(".weight")]
477
- if not shard_text.isdigit():
478
- regular_weights.append((name, loaded_weight))
479
- continue
480
- shard_index = int(shard_text)
481
- if shard_index >= self.split_ngram_parts:
482
- raise ValueError(
483
- f"PLE embedding shard index {shard_index} exceeds "
484
- f"split_ngram_parts={self.split_ngram_parts}"
485
- )
486
- embedding = self.ngram_embedding
487
- shard_size = (
488
- embedding.org_vocab_size + self.split_ngram_parts - 1
489
- ) // self.split_ngram_parts
490
- checkpoint_start = shard_index * shard_size
491
- expected_rows = max(
492
- 0,
493
- min(shard_size, embedding.org_vocab_size - checkpoint_start),
494
- )
495
- expected_shape = (expected_rows, embedding.embedding_dim)
496
- if tuple(loaded_weight.shape) != expected_shape:
497
- raise ValueError(
498
- f"Shape mismatch for PLE embedding shard {shard_index}: "
499
- f"expected {expected_shape}, got "
500
- f"{tuple(loaded_weight.shape)}"
501
- )
502
- embedding.weight.weight_loader(
503
- embedding.weight,
504
- loaded_weight,
505
- checkpoint_start=checkpoint_start,
506
- )
507
- loaded.add("ngram_embedding.weight")
508
- continue
509
- regular_weights.append((name, loaded_weight))
510
-
511
- if regular_weights:
512
- loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
513
- return loaded
514
-
515
-
66
class Qwen4ExpPLELayer(nn.Module, MambaBase):
67
def __init__(
68
self,
@@ -547,6 +97,8 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
97
config,
98
int(config.ple_embed_dim),
99
self.ple_dense_layer_id,
100
+ vllm_config.scheduler_config.max_num_batched_tokens,
101
+ data_parallel_rank=vllm_config.parallel_config.data_parallel_rank,
102
prefix=f"{prefix}.ple_embedding",
103
quant_config=quant_config,
104
params_dtype=model_config.dtype,
@@ -588,10 +140,6 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
140
raise ValueError(f"Duplicate layer name: {prefix}")
141
compilation_config.static_forward_context[prefix] = self
142
591
- def _get_embedding_weight_scale(self) -> torch.Tensor | None:
592
- embedding = getattr(self.ple_embedding, "ngram_embedding", None)
593
- return getattr(embedding, "weight_scale", None)
594
-
143
def _dequantize_embeddings(
144
self,
145
embeddings: torch.Tensor,
@@ -599,14 +147,25 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
147
) -> torch.Tensor:
148
"""Dequantize PLE lookup output."""
149
602
- if not is_fp8(embeddings):
603
- return embeddings
604
- weight_scale = self._get_embedding_weight_scale()
605
- if weight_scale is None:
606
- raise RuntimeError("FP8 PLE embedding is missing its global scale")
607
- if weight_scale.device != embeddings.device:
608
- raise RuntimeError("FP8 PLE embedding scale must be on the output device")
609
- return embeddings.to(output_dtype) * weight_scale.to(output_dtype)
150
+ return self.ple_embedding.ngram_embedding.dequantize(
151
+ embeddings,
152
+ output_dtype,
153
+ )
154
+
155
+ def start_prefetch(
156
+ self,
157
+ hidden_states: torch.Tensor,
158
+ input_ids: torch.Tensor,
159
+ query_start_loc: torch.Tensor,
160
+ ngram_context: torch.Tensor,
161
+ ) -> None:
162
+ """Start the pinned PLE lookup while the preceding decoder layer runs."""
163
+ self.ple_embedding.start_prefetch(
164
+ hidden_states,
165
+ input_ids,
166
+ query_start_loc,
167
+ ngram_context,
168
+ )
169
170
@property
171
def mamba_type(self) -> MambaAttentionBackendEnum:
@@ -827,7 +386,12 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
386
f"token length, got {input_ids.shape[0]} and "
387
f"{hidden_states.shape[0]}"
388
)
830
- embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
389
+ embeddings = self.ple_embedding(
390
+ hidden_states,
391
+ input_ids,
392
+ query_start_loc,
393
+ ngram_context,
394
+ )
395
embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype)
396
kv, _ = self.kv_proj(embeddings)
397
key, value = kv.split(self.kv_proj.output_sizes, dim=-1)
@@ -845,7 +409,6 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
409
410
411
__all__ = [
848
- "Qwen4ExpNGramEmbedding",
412
"Qwen4ExpPLEGroupedNorm",
413
"Qwen4ExpPLELayer",
414
]