@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 267eee54bf

[Qwen4Exp][ROCm] PLE n-gram table CPU offload (#57497)

Signed-off-by: Mathew Odden <modden@redhat.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: opencode+deepseek-v4-flash+vllm <opencode+deepseek-v4-flash+vllm@example.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Mathew Odden committed Sep 24, 2026 at 23:28 UTC 267eee54bf232e4008a7e0b510f19a62e23aaaca
7 files changed +887 -529
.buildkite/test_areas/models_basic.yaml
+2
@@ -169,6 +169,8 @@ steps:
169 - image-build-amd
170 commands:
171 - pytest -v -s models/qwen4_exp/test_qsa_amd.py
172 + # PLE and QSA import conflicting HC registrations in a shared process.
173 + - pytest -v -s models/qwen4_exp/test_ple.py
174
175 - label: ":computer: (CPU) Qwen4 Exp"
176 key: qwen4-exp-unit-tests-cpu
tests/models/qwen4_exp/test_ple.py
+214 -1
@@ -13,7 +13,7 @@ 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
16 +import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module
17 from vllm.config.quantization import QuantizationConfigArgs
18 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
19 from vllm.model_executor.layers.quantization.modelopt import (
@@ -21,11 +21,16 @@ from vllm.model_executor.layers.quantization.modelopt import (
21 ModelOptNvFp4Config,
22 )
23 from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
24 +from vllm.models.qwen4_exp.amd import ple_layer as amd_ple_layer
25 +from vllm.models.qwen4_exp.amd.ple_layer import (
26 + Qwen4ExpPLELayer as Qwen4ExpPLELayerAMD,
27 +)
28 from vllm.models.qwen4_exp.common.ple import (
29 PLEShardOverlap,
30 compute_ple_shard_overlap,
31 copy_ple_embedding_shard_,
32 )
33 +from vllm.models.qwen4_exp.config import Qwen4ExpTextConfig
34 from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
35 Qwen4ExpNGramEmbedding,
36 Qwen4ExpPLEDeviceEmbedding,
@@ -1750,3 +1755,211 @@ def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None:
1755 expected_normed = grouped_norm(expected_gated, norm_conv)
1756 assert torch.equal(gated, expected_gated)
1757 torch.testing.assert_close(normed, expected_normed, atol=1e-2, rtol=1e-2)
1758 +
1759 +
1760 +def _build_amd_ngram_embedding(
1761 + monkeypatch: pytest.MonkeyPatch,
1762 + *,
1763 + device: str,
1764 + cpu_offload: bool,
1765 + fp8_checkpoint: bool,
1766 +) -> tuple[amd_ple_layer.Qwen4ExpNGramEmbedding, torch.Tensor]:
1767 + """Construct and load a small AMD embedding through its checkpoint loader."""
1768 + _mock_etp_group(monkeypatch)
1769 + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
1770 + monkeypatch.setattr(
1771 + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1
1772 + )
1773 + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0)
1774 + monkeypatch.setattr(
1775 + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1
1776 + )
1777 +
1778 + monkeypatch.setattr(
1779 + amd_ple_layer,
1780 + "get_current_vllm_config",
1781 + lambda: SimpleNamespace(engram_config=SimpleNamespace(cpu_offload=cpu_offload)),
1782 + )
1783 + config = Qwen4ExpTextConfig(
1784 + ngram_size=3,
1785 + heads_per_ngram=1,
1786 + ngram_vocab_size_base=5,
1787 + make_ngram_vocab_size_divisible_by=8,
1788 + split_ngram_parts=2,
1789 + ple_embed_dim=6,
1790 + eos_token_id=0,
1791 + vocab_size=64,
1792 + )
1793 + with torch.device(device):
1794 + module = amd_ple_layer.Qwen4ExpNGramEmbedding(
1795 + config,
1796 + embedding_dim=config.ple_embed_dim,
1797 + ple_dense_layer_id=0,
1798 + max_total_tokens=4,
1799 + max_num_reqs=2,
1800 + prefix="test.ple_embedding",
1801 + layer_name="test.ple",
1802 + quant_config=(
1803 + Fp8Config(is_checkpoint_fp8_serialized=True) if fp8_checkpoint else None
1804 + ),
1805 + params_dtype=torch.bfloat16,
1806 + )
1807 + # Only the outer owner is a stub; the embedding constructor and loader run.
1808 + layer = Qwen4ExpPLELayerAMD.__new__(Qwen4ExpPLELayerAMD)
1809 + nn.Module.__init__(layer)
1810 + layer.ple_embedding = module
1811 +
1812 + monkeypatch.setattr(
1813 + amd_ple_layer,
1814 + "get_forward_context",
1815 + lambda: SimpleNamespace(no_compile_layers={module.layer_name: layer}),
1816 + )
1817 +
1818 + embedding = module.ngram_embedding
1819 + loaded_weight = (
1820 + torch.arange(embedding.org_vocab_size * embedding.embedding_dim)
1821 + .reshape(embedding.org_vocab_size, embedding.embedding_dim)
1822 + .to(embedding.weight.dtype)
1823 + )
1824 + shard_size = (embedding.org_vocab_size + 1) // 2
1825 + weights = [
1826 + (f"ngram_embedding.shard_{i}.weight", shard)
1827 + for i, shard in enumerate(loaded_weight.split(shard_size))
1828 + ]
1829 + expected_loaded = {"ngram_embedding.weight"}
1830 + if fp8_checkpoint:
1831 + weights.append(("ngram_embedding.weight_scale", torch.tensor([0.25])))
1832 + expected_loaded.add("ngram_embedding.weight_scale")
1833 + assert module.load_weights(weights) == expected_loaded
1834 + embedding.quant_method.process_weights_after_loading(embedding)
1835 + return module, loaded_weight
1836 +
1837 +
1838 +@pytest.mark.parametrize(
1839 + "device,cpu_offload",
1840 + [
1841 + pytest.param("cpu", False, id="cpu-resident"),
1842 + pytest.param(
1843 + "cuda:0",
1844 + False,
1845 + id="gpu-resident",
1846 + marks=pytest.mark.skipif(
1847 + not torch.cuda.is_available(), reason="requires CUDA or ROCm"
1848 + ),
1849 + ),
1850 + pytest.param(
1851 + "cuda:0",
1852 + True,
1853 + id="gpu-offloaded",
1854 + marks=pytest.mark.skipif(
1855 + not torch.cuda.is_available(), reason="requires CUDA or ROCm"
1856 + ),
1857 + ),
1858 + ],
1859 +)
1860 +def test_amd_fp8_embedding_loads_checkpoint_shards_and_global_scale(
1861 + monkeypatch: pytest.MonkeyPatch,
1862 + device: str,
1863 + cpu_offload: bool,
1864 +) -> None:
1865 + """AMD checkpoint loading must accept and apply the global FP8 scale."""
1866 + module, loaded_weight = _build_amd_ngram_embedding(
1867 + monkeypatch, device=device, cpu_offload=cpu_offload, fp8_checkpoint=True
1868 + )
1869 + embedding = module.ngram_embedding
1870 + assert embedding.weight.dtype == torch.float8_e4m3fn
1871 + assert embedding.weight.device == torch.device("cpu" if cpu_offload else device)
1872 + if cpu_offload:
1873 + assert embedding.weight.is_pinned()
1874 + assert embedding.weight_scale.device == torch.device(device)
1875 + torch.testing.assert_close(embedding.weight.float().cpu(), loaded_weight.float())
1876 + dequantized = embedding.dequantize(embedding.weight.to(device), torch.bfloat16)
1877 + torch.testing.assert_close(
1878 + dequantized.cpu(), loaded_weight.to(torch.bfloat16) * 0.25, rtol=0, atol=0
1879 + )
1880 +
1881 +
1882 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm")
1883 +@pytest.mark.parametrize("fp8_checkpoint", [False, True], ids=["bf16", "fp8"])
1884 +def test_amd_pinned_embedding_output_written_under_compile(
1885 + monkeypatch: pytest.MonkeyPatch,
1886 + fp8_checkpoint: bool,
1887 +) -> None:
1888 + """The AMD pinned lookup op must survive aot_autograd/Inductor.
1889 +
1890 + The op writes via the declared-mutated ``output``. Without
1891 + ``mutates_args=["output"]`` the call is dead-code-eliminated, the UVA lookup
1892 + never runs, and the returned ``output`` stays untouched.
1893 + """
1894 + module, loaded_weight = _build_amd_ngram_embedding(
1895 + monkeypatch, device="cuda:0", cpu_offload=True, fp8_checkpoint=fp8_checkpoint
1896 + )
1897 +
1898 + ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")
1899 + output = torch.zeros(
1900 + ngram_ids.shape[0],
1901 + module.embedding_dim,
1902 + dtype=module.ngram_embedding.weight.dtype,
1903 + device="cuda:0",
1904 + )
1905 +
1906 + def lookup(ids: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
1907 + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned(
1908 + ids, out, module.layer_name
1909 + )
1910 + return out
1911 +
1912 + torch.compile(lookup, fullgraph=True)(ngram_ids, output)
1913 + torch.cuda.current_stream().synchronize()
1914 +
1915 + expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2)
1916 + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
1917 +
1918 +
1919 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm")
1920 +@pytest.mark.parametrize("fp8_checkpoint", [False, True], ids=["bf16", "fp8"])
1921 +def test_amd_pinned_embedding_output_written_under_cudagraph_capture(
1922 + monkeypatch: pytest.MonkeyPatch,
1923 + fp8_checkpoint: bool,
1924 +) -> None:
1925 + """The AMD pinned lookup op must be capture-safe.
1926 +
1927 + It launches a single kernel on the current stream, so a cudagraph capture
1928 + reproduces the lookup into ``output`` without cross-stream work.
1929 + """
1930 + module, loaded_weight = _build_amd_ngram_embedding(
1931 + monkeypatch, device="cuda:0", cpu_offload=True, fp8_checkpoint=fp8_checkpoint
1932 + )
1933 +
1934 + ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")
1935 + output = torch.zeros(
1936 + ngram_ids.shape[0],
1937 + module.embedding_dim,
1938 + dtype=module.ngram_embedding.weight.dtype,
1939 + device="cuda:0",
1940 + )
1941 +
1942 + def lookup(ids: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
1943 + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned(
1944 + ids, out, module.layer_name
1945 + )
1946 + return out
1947 +
1948 + graph = torch.cuda.CUDAGraph()
1949 + warmup_stream = torch.cuda.Stream(device="cuda:0")
1950 + warmup_stream.wait_stream(torch.cuda.current_stream())
1951 + with torch.cuda.stream(warmup_stream):
1952 + for _ in range(3):
1953 + lookup(ngram_ids, output)
1954 + torch.cuda.current_stream().wait_stream(warmup_stream)
1955 +
1956 + with torch.cuda.graph(graph):
1957 + lookup(ngram_ids, output)
1958 + for _ in range(2):
1959 + ngram_ids.copy_(ngram_ids.flip(-1) + 3)
1960 + output.fill_(float("nan"))
1961 + graph.replay()
1962 + torch.cuda.current_stream().synchronize()
1963 +
1964 + expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2)
1965 + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
tests/test_config.py
+20 -14
@@ -1530,26 +1530,32 @@ def test_engram_dp_shared_memory_config_validation(
1530
1531
1532 @pytest.mark.parametrize(
1533 - "architecture, ple_layers, accelerator, supported",
1533 + "architecture, ple_layers, platform, supported",
1534 [
1535 - ("DeepseekV41ForCausalLM", [1], True, True),
1536 - ("DeepseekV41ForCausalLM", [], True, False),
1537 - ("DeepseekV41ForCausalLM", [1], False, False),
1538 - ("Qwen4ExpForCausalLM", [1], True, True),
1539 - ("Qwen4ExpForConditionalGeneration", [1], True, True),
1540 - ("Qwen4ExpForCausalLM", [], True, False),
1541 - ("Qwen4ExpForCausalLM", None, True, False),
1542 - ("Qwen4ExpForCausalLM", [1], False, False),
1543 - ("LlamaForCausalLM", [1], True, False),
1544 - ("Qwen4ExpMTP", [], True, False),
1545 - (None, None, True, False),
1535 + ("DeepseekV41ForCausalLM", [1], "cuda", True),
1536 + ("DeepseekV41ForCausalLM", [1], "rocm", True),
1537 + ("DeepseekV41ForCausalLM", [], "cuda", False),
1538 + ("DeepseekV41ForCausalLM", [1], "cpu", False),
1539 + ("Qwen4ExpForCausalLM", [1], "cuda", True),
1540 + ("Qwen4ExpForCausalLM", [1], "rocm", True),
1541 + ("Qwen4ExpForConditionalGeneration", [1], "cuda", True),
1542 + ("Qwen4ExpForConditionalGeneration", [1], "rocm", True),
1543 + ("Qwen4ExpForCausalLM", [], "cuda", False),
1544 + ("Qwen4ExpForCausalLM", None, "cuda", False),
1545 + ("Qwen4ExpForCausalLM", [1], "cpu", False),
1546 + ("LlamaForCausalLM", [1], "cuda", False),
1547 + ("Qwen4ExpMTP", [], "cuda", False),
1548 + (None, None, "cuda", False),
1549 ],
1550 )
1551 def test_engram_model_support(
1549 - monkeypatch, architecture, ple_layers, accelerator, supported
1552 + monkeypatch, architecture, ple_layers, platform, supported
1553 ):
1554 """A similarly named HF field must not enable unsupported implementations."""
1552 - monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: accelerator)
1555 + monkeypatch.setattr(current_platform, "is_cuda", lambda: platform == "cuda")
1556 + monkeypatch.setattr(
1557 + current_platform, "is_cuda_alike", lambda: platform in ("cuda", "rocm")
1558 + )
1559 model = (
1560 cast(
1561 ModelConfig,
vllm/config/engram.py
+2 -1
@@ -82,7 +82,8 @@ class EngramConfig:
82 ):
83 raise ValueError(
84 "EngramConfig requires a model with supported Engram "
85 - "embeddings, non-empty n-gram layer ids, and CUDA."
85 + "embeddings, non-empty n-gram layer ids, and a CUDA-alike "
86 + "device (CUDA or ROCm)."
87 )
88
89 def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None:
vllm/models/qwen4_exp/amd/ple_layer.py
+90 -4
@@ -11,6 +11,7 @@ from torch import nn
11
12 from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
13 from vllm.forward_context import get_forward_context
14 +from vllm.logger import init_logger
15 from vllm.model_executor.layers.linear import ReplicatedLinear
16 from vllm.model_executor.layers.mamba.abstract import MambaBase
17 from vllm.model_executor.layers.mamba.mamba_utils import (
@@ -18,6 +19,9 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
19 MambaStateShapeCalculator,
20 is_conv_state_dim_first,
21 )
22 +from vllm.model_executor.layers.quantization.base_config import (
23 + QuantizationConfig,
24 +)
25 from vllm.model_executor.models.utils import AutoWeightsLoader
26 from vllm.transformers_utils.configs.qwen4_exp import (
27 Qwen4ExpTextConfig,
@@ -30,7 +34,13 @@ from vllm.v1.attention.backends.short_conv_attn import (
34 )
35 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
36
33 -from ..common.ple import PLEVocabParallelEmbedding
37 +from ..common.ngram_embedding import (
38 + Qwen4ExpPLEDeviceEmbedding,
39 + Qwen4ExpPLEEmbeddingMethod,
40 + Qwen4ExpPLEPinnedHostEmbedding,
41 +)
42 +
43 +logger = init_logger(__name__)
44
45
46 class Qwen4ExpPLEGroupedNorm(nn.Module):
@@ -172,6 +182,10 @@ class Qwen4ExpNGramEmbedding(nn.Module):
182 max_num_reqs: int,
183 prefix: str,
184 layer_name: str,
185 + *,
186 + data_parallel_rank: int = 0,
187 + quant_config: QuantizationConfig | None = None,
188 + params_dtype: torch.dtype | None = None,
189 ) -> None:
190 super().__init__()
191 self.embedding_dim = embedding_dim
@@ -224,11 +238,39 @@ class Qwen4ExpNGramEmbedding(nn.Module):
238 )
239 divisor = int(config.make_ngram_vocab_size_divisible_by)
240 padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
227 - self.ngram_embedding = PLEVocabParallelEmbedding(
241 + if params_dtype is None:
242 + params_dtype = torch.get_default_dtype()
243 + embedding_prefix = f"{prefix}.ngram_embedding"
244 + embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config(
245 + quant_config,
246 + embedding_prefix,
247 + getattr(config, "ple_embedding_dtype", None),
248 + )
249 + engram_config = get_current_vllm_config().engram_config
250 + embedding_cls = (
251 + Qwen4ExpPLEPinnedHostEmbedding
252 + if engram_config is not None and engram_config.cpu_offload
253 + else Qwen4ExpPLEDeviceEmbedding
254 + )
255 + self.ngram_embedding = embedding_cls(
256 padded_vocab_size,
257 self.head_dim,
258 + params_dtype=params_dtype,
259 padding_size=divisor,
231 - prefix=f"{prefix}.ngram_embedding",
260 + prefix=embedding_prefix,
261 + embedding_method=embedding_quant_method,
262 + num_ngram_heads=self.ngram_heads,
263 + max_total_tokens=max_total_tokens,
264 + data_parallel_rank=data_parallel_rank,
265 + )
266 + logger.info(
267 + "Initialized AMD PLE embedding %s: quantization_method=%s, "
268 + "weight_dtype=%s, weight_device=%s, pinned=%s",
269 + embedding_prefix,
270 + type(embedding_quant_method).__name__,
271 + self.ngram_embedding.weight.dtype,
272 + self.ngram_embedding.weight.device,
273 + self.ngram_embedding.weight.is_pinned(),
274 )
275 self.register_buffer(
276 "positions_buffer",
@@ -344,9 +386,21 @@ class Qwen4ExpNGramEmbedding(nn.Module):
386 ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
387 id_blocks.append(ids[request_indices, adjusted_columns])
388 ngram_ids = torch.cat(id_blocks, dim=-1)
389 + embedding = self.ngram_embedding
390 + if embedding.supports_prefetch:
391 + output = ngram_ids.new_empty(
392 + (ngram_ids.shape[0], self.embedding_dim),
393 + dtype=embedding.weight.dtype,
394 + )
395 + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned(
396 + ngram_ids,
397 + output,
398 + self.layer_name,
399 + )
400 + return output
401 output = ngram_ids.new_empty(
402 (ngram_ids.shape[0], self.embedding_dim),
349 - dtype=self.ngram_embedding.params_dtype,
403 + dtype=embedding.params_dtype,
404 )
405 torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding(
406 ngram_ids,
@@ -459,6 +513,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
513 vllm_config.scheduler_config.max_num_seqs,
514 f"{prefix}.ple_embedding",
515 prefix,
516 + data_parallel_rank=vllm_config.parallel_config.data_parallel_rank,
517 + quant_config=quant_config,
518 + params_dtype=model_config.dtype,
519 )
520 self.key_proj = ReplicatedLinear(
521 int(config.ple_embed_dim),
@@ -1044,6 +1101,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
1101 f"{hidden_states.shape[0]}"
1102 )
1103 embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
1104 + embeddings = self.ple_embedding.ngram_embedding.dequantize(
1105 + embeddings,
1106 + self.ple_embedding.ngram_embedding.params_dtype,
1107 + )
1108 key, _ = self.key_proj(embeddings)
1109 value, _ = self.value_proj(embeddings)
1110 token_count = hidden_states.shape[0]
@@ -1081,6 +1142,24 @@ def qwen4_exp_amd_ple_ngram_embedding(
1142 output.copy_(result)
1143
1144
1145 +def qwen4_exp_amd_ple_ngram_embedding_pinned(
1146 + ngram_ids: torch.Tensor,
1147 + output: torch.Tensor,
1148 + layer_name: str,
1149 +) -> None:
1150 + """Run the pinned PLE UVA lookup outside Inductor's FX graph.
1151 +
1152 + Same rationale as the device-path escape: keeping the large embedding weight
1153 + out of the graph prevents AOT compile-time autotuning from materializing a
1154 + synthetic copy of the weight.
1155 + """
1156 + layer = get_forward_context().no_compile_layers[layer_name]
1157 + if not isinstance(layer, Qwen4ExpPLELayer):
1158 + raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner")
1159 + result = layer.ple_embedding.ngram_embedding.sync_lookup(ngram_ids).flatten(-2)
1160 + output.copy_(result)
1161 +
1162 +
1163 def qwen4_exp_ple_short_conv(
1164 inputs: torch.Tensor,
1165 output: torch.Tensor,
@@ -1098,6 +1177,13 @@ direct_register_custom_op(
1177 )
1178
1179
1180 +direct_register_custom_op(
1181 + op_name="qwen4_exp_amd_ple_ngram_embedding_pinned",
1182 + op_func=qwen4_exp_amd_ple_ngram_embedding_pinned,
1183 + mutates_args=["output"],
1184 +)
1185 +
1186 +
1187 direct_register_custom_op(
1188 op_name="qwen4_exp_ple_short_conv",
1189 op_func=qwen4_exp_ple_short_conv,
vllm/models/qwen4_exp/common/ngram_embedding.py new
+542
@@ -0,0 +1,542 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared Qwen4Exp n-gram embedding storage with device and pinned-host backends.
4 +
5 +Both the NVIDIA and AMD Qwen4Exp implementations use these classes so the (large)
6 +n-gram embedding table can be kept in pinned host memory and looked up through
7 +Unified Virtual Addressing on any CUDA-alike platform.
8 +"""
9 +
10 +from abc import ABC, abstractmethod
11 +from typing import ClassVar
12 +
13 +import torch
14 +import torch.nn.functional as F
15 +from torch import nn
16 +
17 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
18 +from vllm.distributed import get_dp_group, get_etp_group, get_tp_group
19 +from vllm.forward_context import DPMetadata, get_forward_context
20 +from vllm.logger import init_logger
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 + ModelOptQuantConfigBase,
29 +)
30 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
31 + create_fp8_scale_parameter,
32 +)
33 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
34 + is_layer_skipped,
35 +)
36 +from vllm.model_executor.parameter import (
37 + ModelWeightParameter,
38 + PerTensorScaleParameter,
39 +)
40 +from vllm.model_executor.utils import set_weight_attrs
41 +from vllm.triton_utils import tl, triton
42 +from vllm.utils.platform_utils import is_uva_available
43 +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
44 +
45 +from .ple import PLEVocabParallelEmbedding
46 +
47 +logger = init_logger(__name__)
48 +
49 +
50 +class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC):
51 + """ETP-sharded PLE table shared by device and pinned-host backends."""
52 +
53 + supports_prefetch: ClassVar[bool] = False
54 +
55 + def __init__(
56 + self,
57 + num_embeddings: int,
58 + embedding_dim: int,
59 + *,
60 + params_dtype: torch.dtype,
61 + padding_size: int,
62 + prefix: str,
63 + embedding_method: "Qwen4ExpPLEEmbeddingMethod",
64 + num_ngram_heads: int = 1,
65 + max_total_tokens: int = 0,
66 + data_parallel_rank: int = 0,
67 + ) -> None:
68 + del num_ngram_heads, max_total_tokens
69 + super().__init__(
70 + num_embeddings,
71 + embedding_dim,
72 + params_dtype=params_dtype,
73 + padding_size=padding_size,
74 + prefix=prefix,
75 + quant_method=embedding_method,
76 + parallel_group=get_etp_group(),
77 + )
78 + self.embedding_method = embedding_method
79 + self.data_parallel_rank = data_parallel_rank
80 + tp_size = get_tp_group().world_size
81 + if self.tp_size % tp_size:
82 + raise ValueError(
83 + "ETP size must be divisible by TP size, but got "
84 + f"ETP={self.tp_size} and TP={tp_size}"
85 + )
86 + self.etp_data_parallel_size = self.tp_size // tp_size
87 +
88 + @abstractmethod
89 + def allocate_embedding_weight(
90 + self,
91 + num_embeddings: int,
92 + embedding_dim: int,
93 + dtype: torch.dtype,
94 + ) -> torch.Tensor:
95 + """Allocate storage for the complete embedding weight."""
96 + raise NotImplementedError
97 +
98 + def dequantize(
99 + self,
100 + embeddings: torch.Tensor,
101 + output_dtype: torch.dtype,
102 + ) -> torch.Tensor:
103 + """Delegate storage-format conversion to the embedding method."""
104 + return self.embedding_method.dequantize(self, embeddings, output_dtype)
105 +
106 + def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]:
107 + """Return the per-DP slot size and this rank's slot offset."""
108 + if self.etp_data_parallel_size == 1:
109 + return local_num_tokens, 0
110 + dp_metadata: DPMetadata | None = get_forward_context().dp_metadata
111 + if dp_metadata is None:
112 + raise RuntimeError("ETP spanning DP requires DP token metadata")
113 + group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * (
114 + self.etp_data_parallel_size
115 + )
116 + group_end = group_start + self.etp_data_parallel_size
117 + token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist()
118 + group_counts = token_counts[group_start:group_end]
119 + slot_size = max(group_counts)
120 + dp_rank = get_dp_group().rank_in_group
121 + return slot_size, dp_rank * slot_size
122 +
123 + def _gather_dp_ids(
124 + self,
125 + ngram_ids: torch.Tensor,
126 + slot_size: int,
127 + ) -> torch.Tensor:
128 + """Gather DP-local IDs that share one ETP-sharded PLE table."""
129 + if self.etp_data_parallel_size == 1:
130 + return ngram_ids
131 + if ngram_ids.shape[0] < slot_size:
132 + padding = ngram_ids.new_zeros(
133 + slot_size - ngram_ids.shape[0], ngram_ids.shape[1]
134 + )
135 + ngram_ids = torch.cat((ngram_ids, padding), dim=0)
136 + return get_dp_group().all_gather(ngram_ids, dim=0)
137 +
138 + def _select_embeddings(
139 + self,
140 + embeddings: torch.Tensor,
141 + local_num_tokens: int,
142 + slot_offset: int,
143 + ) -> torch.Tensor:
144 + """Select this DP rank's rows from the ETP-reduced embeddings."""
145 + if self.etp_data_parallel_size == 1:
146 + return embeddings
147 + return embeddings.narrow(0, slot_offset, local_num_tokens)
148 +
149 + @abstractmethod
150 + def start_prefetch(
151 + self,
152 + hidden_states: torch.Tensor,
153 + ngram_ids: torch.Tensor,
154 + ) -> None:
155 + """Start an asynchronous lookup when supported."""
156 + raise NotImplementedError
157 +
158 +
159 +class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
160 + """Quantization interface shared by resident and pinned PLE tables."""
161 +
162 + # PLE post-load processing only validates scales in their current storage.
163 + requires_device_loading: bool = False
164 +
165 + @staticmethod
166 + def from_quant_config(
167 + quant_config: QuantizationConfig | None,
168 + prefix: str,
169 + embedding_dtype: str | None = None,
170 + ) -> "Qwen4ExpPLEEmbeddingMethod":
171 + """Select the concrete PLE embedding format for a layer."""
172 + if embedding_dtype == "float8_e4m3fn":
173 + return Qwen4ExpPLEFp8EmbeddingMethod()
174 + if quant_config is None:
175 + return Qwen4ExpPLEUnquantizedEmbeddingMethod()
176 + if isinstance(quant_config, ModelOptMixedPrecisionConfig):
177 + if quant_config._resolve_quant_algo(prefix) == "FP8":
178 + return Qwen4ExpPLEFp8EmbeddingMethod()
179 + return Qwen4ExpPLEUnquantizedEmbeddingMethod()
180 + if isinstance(
181 + quant_config, ModelOptQuantConfigBase
182 + ) and quant_config.is_layer_excluded(prefix):
183 + return Qwen4ExpPLEUnquantizedEmbeddingMethod()
184 + if not isinstance(quant_config, Fp8Config):
185 + raise NotImplementedError(
186 + "Qwen4Exp PLE embedding does not support quantization config "
187 + f"{type(quant_config).__name__}"
188 + )
189 +
190 + ignored_layers = quant_config.ignored_layers
191 + if is_layer_skipped(
192 + prefix,
193 + ignored_layers,
194 + quant_config.packed_modules_mapping,
195 + match_mode=quant_config.ignored_layers_match_mode,
196 + ):
197 + return Qwen4ExpPLEUnquantizedEmbeddingMethod()
198 + # PLE checkpoint shards form one runtime embedding parameter.
199 + shard_prefix = f"{prefix}.shard_"
200 + if any(name.startswith(shard_prefix) for name in ignored_layers):
201 + return Qwen4ExpPLEUnquantizedEmbeddingMethod()
202 + if not quant_config.is_checkpoint_fp8_serialized:
203 + raise NotImplementedError(
204 + "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
205 + )
206 + return Qwen4ExpPLEFp8EmbeddingMethod()
207 +
208 + def apply(
209 + self,
210 + layer: nn.Module,
211 + x: torch.Tensor,
212 + bias: torch.Tensor | None = None,
213 + ) -> torch.Tensor:
214 + raise NotImplementedError("PLE weights only support embedding lookup")
215 +
216 + def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
217 + return F.embedding(input_, layer.weight)
218 +
219 + @abstractmethod
220 + def dequantize(
221 + self,
222 + layer: nn.Module,
223 + embeddings: torch.Tensor,
224 + output_dtype: torch.dtype,
225 + ) -> torch.Tensor:
226 + """Convert looked-up PLE rows to the activation dtype."""
227 + raise NotImplementedError
228 +
229 +
230 +class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
231 + """Unquantized PLE embedding storage and lookup semantics."""
232 +
233 + def create_weights(
234 + self,
235 + layer: Qwen4ExpPLEEmbedding,
236 + input_size_per_partition: int,
237 + output_partition_sizes: list[int],
238 + input_size: int,
239 + output_size: int,
240 + params_dtype: torch.dtype,
241 + **extra_weight_attrs,
242 + ) -> None:
243 + del input_size, output_size
244 + weight = nn.Parameter(
245 + layer.allocate_embedding_weight(
246 + sum(output_partition_sizes),
247 + input_size_per_partition,
248 + params_dtype,
249 + ),
250 + requires_grad=False,
251 + )
252 + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
253 + set_weight_attrs(weight, extra_weight_attrs)
254 + layer.register_parameter("weight", weight)
255 +
256 + def dequantize(
257 + self,
258 + layer: nn.Module,
259 + embeddings: torch.Tensor,
260 + output_dtype: torch.dtype,
261 + ) -> torch.Tensor:
262 + del layer, output_dtype
263 + return embeddings
264 +
265 +
266 +class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
267 + """FP8 PLE embedding with one global checkpoint scale."""
268 +
269 + def create_weights(
270 + self,
271 + layer: Qwen4ExpPLEEmbedding,
272 + input_size_per_partition: int,
273 + output_partition_sizes: list[int],
274 + input_size: int,
275 + output_size: int,
276 + params_dtype: torch.dtype,
277 + **extra_weight_attrs,
278 + ) -> None:
279 + del input_size, output_size, params_dtype
280 + weight_loader = extra_weight_attrs.get("weight_loader")
281 + weight = ModelWeightParameter(
282 + data=layer.allocate_embedding_weight(
283 + sum(output_partition_sizes),
284 + input_size_per_partition,
285 + torch.float8_e4m3fn,
286 + ),
287 + input_dim=1,
288 + output_dim=0,
289 + weight_loader=weight_loader,
290 + )
291 + layer.register_parameter("weight", weight)
292 +
293 + weight_scale = create_fp8_scale_parameter(
294 + PerTensorScaleParameter,
295 + output_partition_sizes,
296 + input_size_per_partition,
297 + None,
298 + weight_loader,
299 + scale_dtype=torch.float32,
300 + )
301 + layer.register_parameter("weight_scale", weight_scale)
302 +
303 + def process_weights_after_loading(self, layer: nn.Module) -> None:
304 + """Reject FP8 PLE checkpoints without a global scale."""
305 + sentinel = torch.finfo(torch.float32).min
306 + if torch.any(layer.weight_scale == sentinel):
307 + raise ValueError("FP8 PLE checkpoint is missing its global scale")
308 +
309 + def dequantize(
310 + self,
311 + layer: nn.Module,
312 + embeddings: torch.Tensor,
313 + output_dtype: torch.dtype,
314 + ) -> torch.Tensor:
315 + weight_scale = getattr(layer, "weight_scale", None)
316 + if weight_scale is None:
317 + raise RuntimeError("FP8 PLE embedding is missing its global scale")
318 + if weight_scale.device != embeddings.device:
319 + raise RuntimeError("FP8 PLE embedding scale must be on the output device")
320 + return embeddings.to(output_dtype) * weight_scale.to(output_dtype)
321 +
322 +
323 +class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding):
324 + """PLE table allocated on the active model device."""
325 +
326 + def allocate_embedding_weight(
327 + self,
328 + num_embeddings: int,
329 + embedding_dim: int,
330 + dtype: torch.dtype,
331 + ) -> torch.Tensor:
332 + """Allocate the complete PLE weight on the active device."""
333 + return torch.empty(num_embeddings, embedding_dim, dtype=dtype)
334 +
335 + def start_prefetch(
336 + self,
337 + hidden_states: torch.Tensor,
338 + ngram_ids: torch.Tensor,
339 + ) -> None:
340 + """Resident embedding prefetch is a no-op."""
341 + return None
342 +
343 + def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor:
344 + """Gather ETP inputs, look up embeddings, and select local rows."""
345 + slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
346 + gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
347 + embeddings = super().forward(gathered_ids)
348 + return self._select_embeddings(
349 + embeddings,
350 + ngram_ids.shape[0],
351 + slot_offset,
352 + )
353 +
354 +
355 +@triton.jit
356 +def _lookup_ple_embedding_from_pinned_kernel(
357 + weight_ptr,
358 + ids_ptr,
359 + output_ptr,
360 + embedding_dim,
361 + tp_vocab_start,
362 + tp_vocab_end,
363 + BLOCK_D: tl.constexpr,
364 +):
365 + """Look up TP-owned PLE rows through a CUDA view of pinned host memory."""
366 + row_id = tl.program_id(0)
367 + global_idx = tl.load(ids_ptr + row_id)
368 + in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end)
369 + local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0)
370 + offsets = tl.arange(0, BLOCK_D)
371 + store_mask = offsets < embedding_dim
372 + load_mask = store_mask & in_range
373 + values = tl.load(
374 + weight_ptr + local_idx * embedding_dim + offsets,
375 + mask=load_mask,
376 + other=0.0,
377 + )
378 + tl.store(
379 + output_ptr + row_id * embedding_dim + offsets,
380 + values,
381 + mask=store_mask,
382 + )
383 +
384 +
385 +class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):
386 + """PLE table loaded into pinned CPU memory and looked up through UVA."""
387 +
388 + supports_prefetch: ClassVar[bool] = True
389 +
390 + def __init__(
391 + self,
392 + num_embeddings: int,
393 + embedding_dim: int,
394 + *,
395 + params_dtype: torch.dtype,
396 + padding_size: int,
397 + prefix: str,
398 + embedding_method: Qwen4ExpPLEEmbeddingMethod,
399 + num_ngram_heads: int = 1,
400 + max_total_tokens: int = 0,
401 + data_parallel_rank: int = 0,
402 + ) -> None:
403 + if not is_uva_available():
404 + raise RuntimeError("Engram CPU offload requires UVA support")
405 + super().__init__(
406 + num_embeddings,
407 + embedding_dim,
408 + params_dtype=params_dtype,
409 + padding_size=padding_size,
410 + prefix=prefix,
411 + embedding_method=embedding_method,
412 + num_ngram_heads=num_ngram_heads,
413 + max_total_tokens=max_total_tokens,
414 + data_parallel_rank=data_parallel_rank,
415 + )
416 + self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight)
417 + self._block_d = triton.next_power_of_2(self.embedding_dim)
418 + self._prefetch_stream = torch.cuda.Stream(device=self._uva_weight.device)
419 + self._prefetch_buffer = torch.empty(
420 + max_total_tokens * self.etp_data_parallel_size,
421 + num_ngram_heads,
422 + self.embedding_dim,
423 + dtype=self.weight.dtype,
424 + device=self._uva_weight.device,
425 + )
426 + self._output_dim = num_ngram_heads * self.embedding_dim
427 +
428 + def allocate_embedding_weight(
429 + self,
430 + num_embeddings: int,
431 + embedding_dim: int,
432 + dtype: torch.dtype,
433 + ) -> torch.Tensor:
434 + """Allocate the complete PLE weight directly in pinned CPU memory."""
435 + return torch.empty(
436 + num_embeddings,
437 + embedding_dim,
438 + dtype=dtype,
439 + device="cpu",
440 + pin_memory=True,
441 + )
442 +
443 + def _lookup(
444 + self,
445 + input_ids: torch.Tensor,
446 + output: torch.Tensor | None = None,
447 + ) -> torch.Tensor:
448 + """Look up local ETP rows while preserving the weight storage dtype."""
449 + expected_shape = (*input_ids.shape, self.embedding_dim)
450 + if output is None:
451 + output = torch.empty(
452 + expected_shape,
453 + dtype=self.weight.dtype,
454 + device=input_ids.device,
455 + )
456 + elif (
457 + tuple(output.shape) != expected_shape
458 + or output.dtype != self.weight.dtype
459 + or output.device != input_ids.device
460 + ):
461 + raise ValueError(
462 + "PLE prefetch output must match the input shape, weight dtype, "
463 + "and input device"
464 + )
465 +
466 + flat_ids = input_ids.reshape(-1).long()
467 + if flat_ids.numel():
468 + _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
469 + self._uva_weight,
470 + flat_ids,
471 + output,
472 + self.embedding_dim,
473 + self.shard_indices.org_vocab_start_index,
474 + self.shard_indices.org_vocab_end_index,
475 + BLOCK_D=self._block_d,
476 + )
477 + return output
478 +
479 + def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor:
480 + """Synchronous UVA lookup for platforms without prefetch wiring."""
481 + slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
482 + gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
483 + embeddings = self._lookup(gathered_ids)
484 + embeddings = self._reduce_etp_embeddings(embeddings)
485 + return self._select_embeddings(
486 + embeddings,
487 + ngram_ids.shape[0],
488 + slot_offset,
489 + )
490 +
491 + def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
492 + """Combine pinned lookup results owned by different ETP ranks."""
493 + if self.tp_size == 1:
494 + return embeddings
495 + assert self.parallel_group is not None
496 + if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
497 + # Each vocabulary row has one owner, so reduce the raw FP8 bytes.
498 + reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
499 + return reduced.view(embeddings.dtype)
500 + return self.parallel_group.all_reduce(embeddings)
501 +
502 + @eager_break_during_capture
503 + def start_prefetch(
504 + self,
505 + hidden_states: torch.Tensor,
506 + ngram_ids: torch.Tensor,
507 + ) -> None:
508 + """Gather ETP IDs and launch their UVA lookup on the side stream."""
509 + slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0])
510 + gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
511 + active_output = self._prefetch_buffer[: gathered_ids.shape[0]]
512 + prefetch_stream = self._prefetch_stream
513 + prefetch_stream.wait_stream(torch.cuda.current_stream())
514 + gathered_ids.record_stream(prefetch_stream)
515 + with torch.cuda.stream(prefetch_stream):
516 + self._lookup(gathered_ids, output=active_output)
517 +
518 + @eager_break_during_capture
519 + def _finalize_prefetch(
520 + self,
521 + prefetch_output: torch.Tensor,
522 + output: torch.Tensor,
523 + ) -> None:
524 + """Join the side stream, reduce ETP shards, and select local rows."""
525 + torch.cuda.current_stream().wait_stream(self._prefetch_stream)
526 + slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0])
527 + active_output = prefetch_output[: slot_size * self.etp_data_parallel_size]
528 + embeddings = self._reduce_etp_embeddings(active_output)
529 + embeddings = self._select_embeddings(
530 + embeddings,
531 + output.shape[0],
532 + slot_offset,
533 + )
534 + output.copy_(embeddings.flatten(-2))
535 +
536 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
537 + """Finish the pinned lookup into graph-owned output storage."""
538 + output = self._prefetch_buffer.new_empty(
539 + (hidden_states.shape[0], self._output_dim)
540 + )
541 + self._finalize_prefetch(self._prefetch_buffer, output)
542 + return output
vllm/models/qwen4_exp/nvidia/ngram_embedding.py
+17 -509
@@ -2,534 +2,42 @@
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
5 from collections.abc import Iterable
7 -from typing import ClassVar
6
7 import torch
10 -import torch.nn.functional as F
8 from torch import nn
9
13 -from vllm.compilation.breakable_cudagraph import eager_break_during_capture
10 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
11 from vllm.logger import init_logger
12 from vllm.model_executor.layers.quantization.base_config import (
13 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,
14 )
15 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
16 from vllm.transformers_utils.configs.qwen4_exp import (
17 Qwen4ExpTextConfig,
18 )
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
19
46 -from ..common.ple import PLEVocabParallelEmbedding
20 +from ..common.ngram_embedding import (
21 + Qwen4ExpPLEDeviceEmbedding,
22 + Qwen4ExpPLEEmbedding,
23 + Qwen4ExpPLEEmbeddingMethod,
24 + Qwen4ExpPLEFp8EmbeddingMethod,
25 + Qwen4ExpPLEPinnedHostEmbedding,
26 + Qwen4ExpPLEUnquantizedEmbeddingMethod,
27 +)
28 from .ops.ple import ple_ngram_ids
29
30 logger = init_logger(__name__)
31
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
32 +__all__ = [
33 + "Qwen4ExpPLEDeviceEmbedding",
34 + "Qwen4ExpPLEEmbedding",
35 + "Qwen4ExpPLEEmbeddingMethod",
36 + "Qwen4ExpPLEFp8EmbeddingMethod",
37 + "Qwen4ExpPLEPinnedHostEmbedding",
38 + "Qwen4ExpPLEUnquantizedEmbeddingMethod",
39 + "Qwen4ExpNGramEmbedding",
40 +]
41
42
43 class Qwen4ExpNGramEmbedding(nn.Module):