diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 6f570b3459..cb9e25c0f9 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -169,6 +169,8 @@ steps: - image-build-amd commands: - pytest -v -s models/qwen4_exp/test_qsa_amd.py + # PLE and QSA import conflicting HC registrations in a shared process. + - pytest -v -s models/qwen4_exp/test_ple.py - label: ":computer: (CPU) Qwen4 Exp" key: qwen4-exp-unit-tests-cpu diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 54af705b95..85928d293b 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -13,7 +13,7 @@ from torch.nn import functional as F import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module import vllm.model_executor.parameter as parameter_module -import vllm.models.qwen4_exp.nvidia.ngram_embedding as ngram_embedding_module +import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.modelopt import ( @@ -21,11 +21,16 @@ from vllm.model_executor.layers.quantization.modelopt import ( ModelOptNvFp4Config, ) from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig +from vllm.models.qwen4_exp.amd import ple_layer as amd_ple_layer +from vllm.models.qwen4_exp.amd.ple_layer import ( + Qwen4ExpPLELayer as Qwen4ExpPLELayerAMD, +) from vllm.models.qwen4_exp.common.ple import ( PLEShardOverlap, compute_ple_shard_overlap, copy_ple_embedding_shard_, ) +from vllm.models.qwen4_exp.config import Qwen4ExpTextConfig from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( Qwen4ExpNGramEmbedding, Qwen4ExpPLEDeviceEmbedding, @@ -1750,3 +1755,211 @@ def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None: expected_normed = grouped_norm(expected_gated, norm_conv) assert torch.equal(gated, expected_gated) torch.testing.assert_close(normed, expected_normed, atol=1e-2, rtol=1e-2) + + +def _build_amd_ngram_embedding( + monkeypatch: pytest.MonkeyPatch, + *, + device: str, + cpu_offload: bool, + fp8_checkpoint: bool, +) -> tuple[amd_ple_layer.Qwen4ExpNGramEmbedding, torch.Tensor]: + """Construct and load a small AMD embedding through its checkpoint loader.""" + _mock_etp_group(monkeypatch) + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + + monkeypatch.setattr( + amd_ple_layer, + "get_current_vllm_config", + lambda: SimpleNamespace(engram_config=SimpleNamespace(cpu_offload=cpu_offload)), + ) + config = Qwen4ExpTextConfig( + ngram_size=3, + heads_per_ngram=1, + ngram_vocab_size_base=5, + make_ngram_vocab_size_divisible_by=8, + split_ngram_parts=2, + ple_embed_dim=6, + eos_token_id=0, + vocab_size=64, + ) + with torch.device(device): + module = amd_ple_layer.Qwen4ExpNGramEmbedding( + config, + embedding_dim=config.ple_embed_dim, + ple_dense_layer_id=0, + max_total_tokens=4, + max_num_reqs=2, + prefix="test.ple_embedding", + layer_name="test.ple", + quant_config=( + Fp8Config(is_checkpoint_fp8_serialized=True) if fp8_checkpoint else None + ), + params_dtype=torch.bfloat16, + ) + # Only the outer owner is a stub; the embedding constructor and loader run. + layer = Qwen4ExpPLELayerAMD.__new__(Qwen4ExpPLELayerAMD) + nn.Module.__init__(layer) + layer.ple_embedding = module + + monkeypatch.setattr( + amd_ple_layer, + "get_forward_context", + lambda: SimpleNamespace(no_compile_layers={module.layer_name: layer}), + ) + + embedding = module.ngram_embedding + loaded_weight = ( + torch.arange(embedding.org_vocab_size * embedding.embedding_dim) + .reshape(embedding.org_vocab_size, embedding.embedding_dim) + .to(embedding.weight.dtype) + ) + shard_size = (embedding.org_vocab_size + 1) // 2 + weights = [ + (f"ngram_embedding.shard_{i}.weight", shard) + for i, shard in enumerate(loaded_weight.split(shard_size)) + ] + expected_loaded = {"ngram_embedding.weight"} + if fp8_checkpoint: + weights.append(("ngram_embedding.weight_scale", torch.tensor([0.25]))) + expected_loaded.add("ngram_embedding.weight_scale") + assert module.load_weights(weights) == expected_loaded + embedding.quant_method.process_weights_after_loading(embedding) + return module, loaded_weight + + +@pytest.mark.parametrize( + "device,cpu_offload", + [ + pytest.param("cpu", False, id="cpu-resident"), + pytest.param( + "cuda:0", + False, + id="gpu-resident", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA or ROCm" + ), + ), + pytest.param( + "cuda:0", + True, + id="gpu-offloaded", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA or ROCm" + ), + ), + ], +) +def test_amd_fp8_embedding_loads_checkpoint_shards_and_global_scale( + monkeypatch: pytest.MonkeyPatch, + device: str, + cpu_offload: bool, +) -> None: + """AMD checkpoint loading must accept and apply the global FP8 scale.""" + module, loaded_weight = _build_amd_ngram_embedding( + monkeypatch, device=device, cpu_offload=cpu_offload, fp8_checkpoint=True + ) + embedding = module.ngram_embedding + assert embedding.weight.dtype == torch.float8_e4m3fn + assert embedding.weight.device == torch.device("cpu" if cpu_offload else device) + if cpu_offload: + assert embedding.weight.is_pinned() + assert embedding.weight_scale.device == torch.device(device) + torch.testing.assert_close(embedding.weight.float().cpu(), loaded_weight.float()) + dequantized = embedding.dequantize(embedding.weight.to(device), torch.bfloat16) + torch.testing.assert_close( + dequantized.cpu(), loaded_weight.to(torch.bfloat16) * 0.25, rtol=0, atol=0 + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm") +@pytest.mark.parametrize("fp8_checkpoint", [False, True], ids=["bf16", "fp8"]) +def test_amd_pinned_embedding_output_written_under_compile( + monkeypatch: pytest.MonkeyPatch, + fp8_checkpoint: bool, +) -> None: + """The AMD pinned lookup op must survive aot_autograd/Inductor. + + The op writes via the declared-mutated ``output``. Without + ``mutates_args=["output"]`` the call is dead-code-eliminated, the UVA lookup + never runs, and the returned ``output`` stays untouched. + """ + module, loaded_weight = _build_amd_ngram_embedding( + monkeypatch, device="cuda:0", cpu_offload=True, fp8_checkpoint=fp8_checkpoint + ) + + ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") + output = torch.zeros( + ngram_ids.shape[0], + module.embedding_dim, + dtype=module.ngram_embedding.weight.dtype, + device="cuda:0", + ) + + def lookup(ids: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned( + ids, out, module.layer_name + ) + return out + + torch.compile(lookup, fullgraph=True)(ngram_ids, output) + torch.cuda.current_stream().synchronize() + + expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2) + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm") +@pytest.mark.parametrize("fp8_checkpoint", [False, True], ids=["bf16", "fp8"]) +def test_amd_pinned_embedding_output_written_under_cudagraph_capture( + monkeypatch: pytest.MonkeyPatch, + fp8_checkpoint: bool, +) -> None: + """The AMD pinned lookup op must be capture-safe. + + It launches a single kernel on the current stream, so a cudagraph capture + reproduces the lookup into ``output`` without cross-stream work. + """ + module, loaded_weight = _build_amd_ngram_embedding( + monkeypatch, device="cuda:0", cpu_offload=True, fp8_checkpoint=fp8_checkpoint + ) + + ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0") + output = torch.zeros( + ngram_ids.shape[0], + module.embedding_dim, + dtype=module.ngram_embedding.weight.dtype, + device="cuda:0", + ) + + def lookup(ids: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned( + ids, out, module.layer_name + ) + return out + + graph = torch.cuda.CUDAGraph() + warmup_stream = torch.cuda.Stream(device="cuda:0") + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + lookup(ngram_ids, output) + torch.cuda.current_stream().wait_stream(warmup_stream) + + with torch.cuda.graph(graph): + lookup(ngram_ids, output) + for _ in range(2): + ngram_ids.copy_(ngram_ids.flip(-1) + 3) + output.fill_(float("nan")) + graph.replay() + torch.cuda.current_stream().synchronize() + + expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2) + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) diff --git a/tests/test_config.py b/tests/test_config.py index 3c93743756..aff9758a45 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1530,26 +1530,32 @@ def test_engram_dp_shared_memory_config_validation( @pytest.mark.parametrize( - "architecture, ple_layers, accelerator, supported", + "architecture, ple_layers, platform, supported", [ - ("DeepseekV41ForCausalLM", [1], True, True), - ("DeepseekV41ForCausalLM", [], True, False), - ("DeepseekV41ForCausalLM", [1], False, False), - ("Qwen4ExpForCausalLM", [1], True, True), - ("Qwen4ExpForConditionalGeneration", [1], True, True), - ("Qwen4ExpForCausalLM", [], True, False), - ("Qwen4ExpForCausalLM", None, True, False), - ("Qwen4ExpForCausalLM", [1], False, False), - ("LlamaForCausalLM", [1], True, False), - ("Qwen4ExpMTP", [], True, False), - (None, None, True, False), + ("DeepseekV41ForCausalLM", [1], "cuda", True), + ("DeepseekV41ForCausalLM", [1], "rocm", True), + ("DeepseekV41ForCausalLM", [], "cuda", False), + ("DeepseekV41ForCausalLM", [1], "cpu", False), + ("Qwen4ExpForCausalLM", [1], "cuda", True), + ("Qwen4ExpForCausalLM", [1], "rocm", True), + ("Qwen4ExpForConditionalGeneration", [1], "cuda", True), + ("Qwen4ExpForConditionalGeneration", [1], "rocm", True), + ("Qwen4ExpForCausalLM", [], "cuda", False), + ("Qwen4ExpForCausalLM", None, "cuda", False), + ("Qwen4ExpForCausalLM", [1], "cpu", False), + ("LlamaForCausalLM", [1], "cuda", False), + ("Qwen4ExpMTP", [], "cuda", False), + (None, None, "cuda", False), ], ) def test_engram_model_support( - monkeypatch, architecture, ple_layers, accelerator, supported + monkeypatch, architecture, ple_layers, platform, supported ): """A similarly named HF field must not enable unsupported implementations.""" - monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: accelerator) + monkeypatch.setattr(current_platform, "is_cuda", lambda: platform == "cuda") + monkeypatch.setattr( + current_platform, "is_cuda_alike", lambda: platform in ("cuda", "rocm") + ) model = ( cast( ModelConfig, diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 2cd67c0973..fa28d803b5 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -82,7 +82,8 @@ class EngramConfig: ): raise ValueError( "EngramConfig requires a model with supported Engram " - "embeddings, non-empty n-gram layer ids, and CUDA." + "embeddings, non-empty n-gram layer ids, and a CUDA-alike " + "device (CUDA or ROCm)." ) def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None: diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index 2616a3e9e0..0055821d56 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -11,6 +11,7 @@ from torch import nn from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config from vllm.forward_context import get_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mamba.mamba_utils import ( @@ -18,6 +19,9 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateShapeCalculator, is_conv_state_dim_first, ) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.transformers_utils.configs.qwen4_exp import ( Qwen4ExpTextConfig, @@ -30,7 +34,13 @@ from vllm.v1.attention.backends.short_conv_attn import ( ) from vllm.v1.attention.backends.utils import NULL_BLOCK_ID -from ..common.ple import PLEVocabParallelEmbedding +from ..common.ngram_embedding import ( + Qwen4ExpPLEDeviceEmbedding, + Qwen4ExpPLEEmbeddingMethod, + Qwen4ExpPLEPinnedHostEmbedding, +) + +logger = init_logger(__name__) class Qwen4ExpPLEGroupedNorm(nn.Module): @@ -172,6 +182,10 @@ class Qwen4ExpNGramEmbedding(nn.Module): max_num_reqs: int, prefix: str, layer_name: str, + *, + data_parallel_rank: int = 0, + quant_config: QuantizationConfig | None = None, + params_dtype: torch.dtype | None = None, ) -> None: super().__init__() self.embedding_dim = embedding_dim @@ -224,11 +238,39 @@ class Qwen4ExpNGramEmbedding(nn.Module): ) divisor = int(config.make_ngram_vocab_size_divisible_by) padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor - self.ngram_embedding = PLEVocabParallelEmbedding( + if params_dtype is None: + params_dtype = torch.get_default_dtype() + embedding_prefix = f"{prefix}.ngram_embedding" + embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config( + quant_config, + embedding_prefix, + getattr(config, "ple_embedding_dtype", None), + ) + engram_config = get_current_vllm_config().engram_config + embedding_cls = ( + Qwen4ExpPLEPinnedHostEmbedding + if engram_config is not None and engram_config.cpu_offload + else Qwen4ExpPLEDeviceEmbedding + ) + self.ngram_embedding = embedding_cls( padded_vocab_size, self.head_dim, + params_dtype=params_dtype, padding_size=divisor, - prefix=f"{prefix}.ngram_embedding", + prefix=embedding_prefix, + embedding_method=embedding_quant_method, + num_ngram_heads=self.ngram_heads, + max_total_tokens=max_total_tokens, + data_parallel_rank=data_parallel_rank, + ) + logger.info( + "Initialized AMD PLE embedding %s: quantization_method=%s, " + "weight_dtype=%s, weight_device=%s, pinned=%s", + embedding_prefix, + type(embedding_quant_method).__name__, + self.ngram_embedding.weight.dtype, + self.ngram_embedding.weight.device, + self.ngram_embedding.weight.is_pinned(), ) self.register_buffer( "positions_buffer", @@ -344,9 +386,21 @@ class Qwen4ExpNGramEmbedding(nn.Module): ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets id_blocks.append(ids[request_indices, adjusted_columns]) ngram_ids = torch.cat(id_blocks, dim=-1) + embedding = self.ngram_embedding + if embedding.supports_prefetch: + output = ngram_ids.new_empty( + (ngram_ids.shape[0], self.embedding_dim), + dtype=embedding.weight.dtype, + ) + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned( + ngram_ids, + output, + self.layer_name, + ) + return output output = ngram_ids.new_empty( (ngram_ids.shape[0], self.embedding_dim), - dtype=self.ngram_embedding.params_dtype, + dtype=embedding.params_dtype, ) torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding( ngram_ids, @@ -459,6 +513,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): vllm_config.scheduler_config.max_num_seqs, f"{prefix}.ple_embedding", prefix, + data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, + quant_config=quant_config, + params_dtype=model_config.dtype, ) self.key_proj = ReplicatedLinear( int(config.ple_embed_dim), @@ -1044,6 +1101,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): f"{hidden_states.shape[0]}" ) embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context) + embeddings = self.ple_embedding.ngram_embedding.dequantize( + embeddings, + self.ple_embedding.ngram_embedding.params_dtype, + ) key, _ = self.key_proj(embeddings) value, _ = self.value_proj(embeddings) token_count = hidden_states.shape[0] @@ -1081,6 +1142,24 @@ def qwen4_exp_amd_ple_ngram_embedding( output.copy_(result) +def qwen4_exp_amd_ple_ngram_embedding_pinned( + ngram_ids: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> None: + """Run the pinned PLE UVA lookup outside Inductor's FX graph. + + Same rationale as the device-path escape: keeping the large embedding weight + out of the graph prevents AOT compile-time autotuning from materializing a + synthetic copy of the weight. + """ + layer = get_forward_context().no_compile_layers[layer_name] + if not isinstance(layer, Qwen4ExpPLELayer): + raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner") + result = layer.ple_embedding.ngram_embedding.sync_lookup(ngram_ids).flatten(-2) + output.copy_(result) + + def qwen4_exp_ple_short_conv( inputs: torch.Tensor, output: torch.Tensor, @@ -1098,6 +1177,13 @@ direct_register_custom_op( ) +direct_register_custom_op( + op_name="qwen4_exp_amd_ple_ngram_embedding_pinned", + op_func=qwen4_exp_amd_ple_ngram_embedding_pinned, + mutates_args=["output"], +) + + direct_register_custom_op( op_name="qwen4_exp_ple_short_conv", op_func=qwen4_exp_ple_short_conv, diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py new file mode 100644 index 0000000000..4d1d358958 --- /dev/null +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -0,0 +1,542 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared Qwen4Exp n-gram embedding storage with device and pinned-host backends. + +Both the NVIDIA and AMD Qwen4Exp implementations use these classes so the (large) +n-gram embedding table can be kept in pinned host memory and looked up through +Unified Virtual Addressing on any CUDA-alike platform. +""" + +from abc import ABC, abstractmethod +from typing import ClassVar + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.distributed import get_dp_group, get_etp_group, get_tp_group +from vllm.forward_context import DPMetadata, get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptMixedPrecisionConfig, + ModelOptQuantConfigBase, +) +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + create_fp8_scale_parameter, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.parameter import ( + ModelWeightParameter, + PerTensorScaleParameter, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.triton_utils import tl, triton +from vllm.utils.platform_utils import is_uva_available +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + +from .ple import PLEVocabParallelEmbedding + +logger = init_logger(__name__) + + +class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): + """ETP-sharded PLE table shared by device and pinned-host backends.""" + + supports_prefetch: ClassVar[bool] = False + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + *, + params_dtype: torch.dtype, + padding_size: int, + prefix: str, + embedding_method: "Qwen4ExpPLEEmbeddingMethod", + num_ngram_heads: int = 1, + max_total_tokens: int = 0, + data_parallel_rank: int = 0, + ) -> None: + del num_ngram_heads, max_total_tokens + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + padding_size=padding_size, + prefix=prefix, + quant_method=embedding_method, + parallel_group=get_etp_group(), + ) + self.embedding_method = embedding_method + self.data_parallel_rank = data_parallel_rank + tp_size = get_tp_group().world_size + if self.tp_size % tp_size: + raise ValueError( + "ETP size must be divisible by TP size, but got " + f"ETP={self.tp_size} and TP={tp_size}" + ) + self.etp_data_parallel_size = self.tp_size // tp_size + + @abstractmethod + def allocate_embedding_weight( + self, + num_embeddings: int, + embedding_dim: int, + dtype: torch.dtype, + ) -> torch.Tensor: + """Allocate storage for the complete embedding weight.""" + raise NotImplementedError + + def dequantize( + self, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + """Delegate storage-format conversion to the embedding method.""" + return self.embedding_method.dequantize(self, embeddings, output_dtype) + + def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]: + """Return the per-DP slot size and this rank's slot offset.""" + if self.etp_data_parallel_size == 1: + return local_num_tokens, 0 + dp_metadata: DPMetadata | None = get_forward_context().dp_metadata + if dp_metadata is None: + raise RuntimeError("ETP spanning DP requires DP token metadata") + group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * ( + self.etp_data_parallel_size + ) + group_end = group_start + self.etp_data_parallel_size + token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist() + group_counts = token_counts[group_start:group_end] + slot_size = max(group_counts) + dp_rank = get_dp_group().rank_in_group + return slot_size, dp_rank * slot_size + + def _gather_dp_ids( + self, + ngram_ids: torch.Tensor, + slot_size: int, + ) -> torch.Tensor: + """Gather DP-local IDs that share one ETP-sharded PLE table.""" + if self.etp_data_parallel_size == 1: + return ngram_ids + if ngram_ids.shape[0] < slot_size: + padding = ngram_ids.new_zeros( + slot_size - ngram_ids.shape[0], ngram_ids.shape[1] + ) + ngram_ids = torch.cat((ngram_ids, padding), dim=0) + return get_dp_group().all_gather(ngram_ids, dim=0) + + def _select_embeddings( + self, + embeddings: torch.Tensor, + local_num_tokens: int, + slot_offset: int, + ) -> torch.Tensor: + """Select this DP rank's rows from the ETP-reduced embeddings.""" + if self.etp_data_parallel_size == 1: + return embeddings + return embeddings.narrow(0, slot_offset, local_num_tokens) + + @abstractmethod + def start_prefetch( + self, + hidden_states: torch.Tensor, + ngram_ids: torch.Tensor, + ) -> None: + """Start an asynchronous lookup when supported.""" + raise NotImplementedError + + +class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): + """Quantization interface shared by resident and pinned PLE tables.""" + + # PLE post-load processing only validates scales in their current storage. + requires_device_loading: bool = False + + @staticmethod + def from_quant_config( + quant_config: QuantizationConfig | None, + prefix: str, + embedding_dtype: str | None = None, + ) -> "Qwen4ExpPLEEmbeddingMethod": + """Select the concrete PLE embedding format for a layer.""" + if embedding_dtype == "float8_e4m3fn": + return Qwen4ExpPLEFp8EmbeddingMethod() + if quant_config is None: + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if isinstance(quant_config, ModelOptMixedPrecisionConfig): + if quant_config._resolve_quant_algo(prefix) == "FP8": + return Qwen4ExpPLEFp8EmbeddingMethod() + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if isinstance( + quant_config, ModelOptQuantConfigBase + ) and quant_config.is_layer_excluded(prefix): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if not isinstance(quant_config, Fp8Config): + raise NotImplementedError( + "Qwen4Exp PLE embedding does not support quantization config " + f"{type(quant_config).__name__}" + ) + + ignored_layers = quant_config.ignored_layers + if is_layer_skipped( + prefix, + ignored_layers, + quant_config.packed_modules_mapping, + match_mode=quant_config.ignored_layers_match_mode, + ): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + # PLE checkpoint shards form one runtime embedding parameter. + shard_prefix = f"{prefix}.shard_" + if any(name.startswith(shard_prefix) for name in ignored_layers): + return Qwen4ExpPLEUnquantizedEmbeddingMethod() + if not quant_config.is_checkpoint_fp8_serialized: + raise NotImplementedError( + "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints" + ) + return Qwen4ExpPLEFp8EmbeddingMethod() + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + raise NotImplementedError("PLE weights only support embedding lookup") + + def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: + return F.embedding(input_, layer.weight) + + @abstractmethod + def dequantize( + self, + layer: nn.Module, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + """Convert looked-up PLE rows to the activation dtype.""" + raise NotImplementedError + + +class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): + """Unquantized PLE embedding storage and lookup semantics.""" + + def create_weights( + self, + layer: Qwen4ExpPLEEmbedding, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size + weight = nn.Parameter( + layer.allocate_embedding_weight( + sum(output_partition_sizes), + input_size_per_partition, + params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) + set_weight_attrs(weight, extra_weight_attrs) + layer.register_parameter("weight", weight) + + def dequantize( + self, + layer: nn.Module, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + del layer, output_dtype + return embeddings + + +class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): + """FP8 PLE embedding with one global checkpoint scale.""" + + def create_weights( + self, + layer: Qwen4ExpPLEEmbedding, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size, params_dtype + weight_loader = extra_weight_attrs.get("weight_loader") + weight = ModelWeightParameter( + data=layer.allocate_embedding_weight( + sum(output_partition_sizes), + input_size_per_partition, + torch.float8_e4m3fn, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + weight_scale = create_fp8_scale_parameter( + PerTensorScaleParameter, + output_partition_sizes, + input_size_per_partition, + None, + weight_loader, + scale_dtype=torch.float32, + ) + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Reject FP8 PLE checkpoints without a global scale.""" + sentinel = torch.finfo(torch.float32).min + if torch.any(layer.weight_scale == sentinel): + raise ValueError("FP8 PLE checkpoint is missing its global scale") + + def dequantize( + self, + layer: nn.Module, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + weight_scale = getattr(layer, "weight_scale", None) + if weight_scale is None: + raise RuntimeError("FP8 PLE embedding is missing its global scale") + if weight_scale.device != embeddings.device: + raise RuntimeError("FP8 PLE embedding scale must be on the output device") + return embeddings.to(output_dtype) * weight_scale.to(output_dtype) + + +class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding): + """PLE table allocated on the active model device.""" + + def allocate_embedding_weight( + self, + num_embeddings: int, + embedding_dim: int, + dtype: torch.dtype, + ) -> torch.Tensor: + """Allocate the complete PLE weight on the active device.""" + return torch.empty(num_embeddings, embedding_dim, dtype=dtype) + + def start_prefetch( + self, + hidden_states: torch.Tensor, + ngram_ids: torch.Tensor, + ) -> None: + """Resident embedding prefetch is a no-op.""" + return None + + def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor: + """Gather ETP inputs, look up embeddings, and select local rows.""" + slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) + gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) + embeddings = super().forward(gathered_ids) + return self._select_embeddings( + embeddings, + ngram_ids.shape[0], + slot_offset, + ) + + +@triton.jit +def _lookup_ple_embedding_from_pinned_kernel( + weight_ptr, + ids_ptr, + output_ptr, + embedding_dim, + tp_vocab_start, + tp_vocab_end, + BLOCK_D: tl.constexpr, +): + """Look up TP-owned PLE rows through a CUDA view of pinned host memory.""" + row_id = tl.program_id(0) + global_idx = tl.load(ids_ptr + row_id) + in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end) + local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0) + offsets = tl.arange(0, BLOCK_D) + store_mask = offsets < embedding_dim + load_mask = store_mask & in_range + values = tl.load( + weight_ptr + local_idx * embedding_dim + offsets, + mask=load_mask, + other=0.0, + ) + tl.store( + output_ptr + row_id * embedding_dim + offsets, + values, + mask=store_mask, + ) + + +class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): + """PLE table loaded into pinned CPU memory and looked up through UVA.""" + + supports_prefetch: ClassVar[bool] = True + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + *, + params_dtype: torch.dtype, + padding_size: int, + prefix: str, + embedding_method: Qwen4ExpPLEEmbeddingMethod, + num_ngram_heads: int = 1, + max_total_tokens: int = 0, + data_parallel_rank: int = 0, + ) -> None: + if not is_uva_available(): + raise RuntimeError("Engram CPU offload requires UVA support") + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + padding_size=padding_size, + prefix=prefix, + embedding_method=embedding_method, + num_ngram_heads=num_ngram_heads, + max_total_tokens=max_total_tokens, + data_parallel_rank=data_parallel_rank, + ) + self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) + self._block_d = triton.next_power_of_2(self.embedding_dim) + self._prefetch_stream = torch.cuda.Stream(device=self._uva_weight.device) + self._prefetch_buffer = torch.empty( + max_total_tokens * self.etp_data_parallel_size, + num_ngram_heads, + self.embedding_dim, + dtype=self.weight.dtype, + device=self._uva_weight.device, + ) + self._output_dim = num_ngram_heads * self.embedding_dim + + def allocate_embedding_weight( + self, + num_embeddings: int, + embedding_dim: int, + dtype: torch.dtype, + ) -> torch.Tensor: + """Allocate the complete PLE weight directly in pinned CPU memory.""" + return torch.empty( + num_embeddings, + embedding_dim, + dtype=dtype, + device="cpu", + pin_memory=True, + ) + + def _lookup( + self, + input_ids: torch.Tensor, + output: torch.Tensor | None = None, + ) -> torch.Tensor: + """Look up local ETP rows while preserving the weight storage dtype.""" + expected_shape = (*input_ids.shape, self.embedding_dim) + if output is None: + output = torch.empty( + expected_shape, + dtype=self.weight.dtype, + device=input_ids.device, + ) + elif ( + tuple(output.shape) != expected_shape + or output.dtype != self.weight.dtype + or output.device != input_ids.device + ): + raise ValueError( + "PLE prefetch output must match the input shape, weight dtype, " + "and input device" + ) + + flat_ids = input_ids.reshape(-1).long() + if flat_ids.numel(): + _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)]( + self._uva_weight, + flat_ids, + output, + self.embedding_dim, + self.shard_indices.org_vocab_start_index, + self.shard_indices.org_vocab_end_index, + BLOCK_D=self._block_d, + ) + return output + + def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor: + """Synchronous UVA lookup for platforms without prefetch wiring.""" + slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) + gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) + embeddings = self._lookup(gathered_ids) + embeddings = self._reduce_etp_embeddings(embeddings) + return self._select_embeddings( + embeddings, + ngram_ids.shape[0], + slot_offset, + ) + + def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: + """Combine pinned lookup results owned by different ETP ranks.""" + if self.tp_size == 1: + return embeddings + assert self.parallel_group is not None + if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + # Each vocabulary row has one owner, so reduce the raw FP8 bytes. + reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8)) + return reduced.view(embeddings.dtype) + return self.parallel_group.all_reduce(embeddings) + + @eager_break_during_capture + def start_prefetch( + self, + hidden_states: torch.Tensor, + ngram_ids: torch.Tensor, + ) -> None: + """Gather ETP IDs and launch their UVA lookup on the side stream.""" + slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0]) + gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) + active_output = self._prefetch_buffer[: gathered_ids.shape[0]] + prefetch_stream = self._prefetch_stream + prefetch_stream.wait_stream(torch.cuda.current_stream()) + gathered_ids.record_stream(prefetch_stream) + with torch.cuda.stream(prefetch_stream): + self._lookup(gathered_ids, output=active_output) + + @eager_break_during_capture + def _finalize_prefetch( + self, + prefetch_output: torch.Tensor, + output: torch.Tensor, + ) -> None: + """Join the side stream, reduce ETP shards, and select local rows.""" + torch.cuda.current_stream().wait_stream(self._prefetch_stream) + slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0]) + active_output = prefetch_output[: slot_size * self.etp_data_parallel_size] + embeddings = self._reduce_etp_embeddings(active_output) + embeddings = self._select_embeddings( + embeddings, + output.shape[0], + slot_offset, + ) + output.copy_(embeddings.flatten(-2)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Finish the pinned lookup into graph-owned output storage.""" + output = self._prefetch_buffer.new_empty( + (hidden_states.shape[0], self._output_dim) + ) + self._finalize_prefetch(self._prefetch_buffer, output) + return output diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 12050fb077..0b6fd650b1 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -2,534 +2,42 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Qwen4Exp n-gram embeddings with device and pinned-host storage.""" -from abc import ABC, abstractmethod from collections.abc import Iterable -from typing import ClassVar import torch -import torch.nn.functional as F from torch import nn -from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import get_current_vllm_config -from vllm.distributed import get_dp_group, get_etp_group, get_tp_group -from vllm.forward_context import DPMetadata, get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, - QuantizeMethodBase, -) -from vllm.model_executor.layers.quantization.fp8 import Fp8Config -from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptMixedPrecisionConfig, - ModelOptQuantConfigBase, -) -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - create_fp8_scale_parameter, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - is_layer_skipped, ) from vllm.model_executor.models.utils import AutoWeightsLoader -from vllm.model_executor.parameter import ( - ModelWeightParameter, - PerTensorScaleParameter, -) -from vllm.model_executor.utils import set_weight_attrs from vllm.transformers_utils.configs.qwen4_exp import ( Qwen4ExpTextConfig, ) -from vllm.triton_utils import tl, triton -from vllm.utils.platform_utils import is_uva_available -from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor -from ..common.ple import PLEVocabParallelEmbedding +from ..common.ngram_embedding import ( + Qwen4ExpPLEDeviceEmbedding, + Qwen4ExpPLEEmbedding, + Qwen4ExpPLEEmbeddingMethod, + Qwen4ExpPLEFp8EmbeddingMethod, + Qwen4ExpPLEPinnedHostEmbedding, + Qwen4ExpPLEUnquantizedEmbeddingMethod, +) from .ops.ple import ple_ngram_ids logger = init_logger(__name__) - -class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC): - """ETP-sharded PLE table shared by device and pinned-host backends.""" - - supports_prefetch: ClassVar[bool] = False - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - *, - params_dtype: torch.dtype, - padding_size: int, - prefix: str, - embedding_method: "Qwen4ExpPLEEmbeddingMethod", - num_ngram_heads: int = 1, - max_total_tokens: int = 0, - data_parallel_rank: int = 0, - ) -> None: - del num_ngram_heads, max_total_tokens - super().__init__( - num_embeddings, - embedding_dim, - params_dtype=params_dtype, - padding_size=padding_size, - prefix=prefix, - quant_method=embedding_method, - parallel_group=get_etp_group(), - ) - self.embedding_method = embedding_method - self.data_parallel_rank = data_parallel_rank - tp_size = get_tp_group().world_size - if self.tp_size % tp_size: - raise ValueError( - "ETP size must be divisible by TP size, but got " - f"ETP={self.tp_size} and TP={tp_size}" - ) - self.etp_data_parallel_size = self.tp_size // tp_size - - @abstractmethod - def allocate_embedding_weight( - self, - num_embeddings: int, - embedding_dim: int, - dtype: torch.dtype, - ) -> torch.Tensor: - """Allocate storage for the complete embedding weight.""" - raise NotImplementedError - - def dequantize( - self, - embeddings: torch.Tensor, - output_dtype: torch.dtype, - ) -> torch.Tensor: - """Delegate storage-format conversion to the embedding method.""" - return self.embedding_method.dequantize(self, embeddings, output_dtype) - - def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]: - """Return the per-DP slot size and this rank's slot offset.""" - if self.etp_data_parallel_size == 1: - return local_num_tokens, 0 - dp_metadata: DPMetadata | None = get_forward_context().dp_metadata - if dp_metadata is None: - raise RuntimeError("ETP spanning DP requires DP token metadata") - group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * ( - self.etp_data_parallel_size - ) - group_end = group_start + self.etp_data_parallel_size - token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist() - group_counts = token_counts[group_start:group_end] - slot_size = max(group_counts) - dp_rank = get_dp_group().rank_in_group - return slot_size, dp_rank * slot_size - - def _gather_dp_ids( - self, - ngram_ids: torch.Tensor, - slot_size: int, - ) -> torch.Tensor: - """Gather DP-local IDs that share one ETP-sharded PLE table.""" - if self.etp_data_parallel_size == 1: - return ngram_ids - if ngram_ids.shape[0] < slot_size: - padding = ngram_ids.new_zeros( - slot_size - ngram_ids.shape[0], ngram_ids.shape[1] - ) - ngram_ids = torch.cat((ngram_ids, padding), dim=0) - return get_dp_group().all_gather(ngram_ids, dim=0) - - def _select_embeddings( - self, - embeddings: torch.Tensor, - local_num_tokens: int, - slot_offset: int, - ) -> torch.Tensor: - """Select this DP rank's rows from the ETP-reduced embeddings.""" - if self.etp_data_parallel_size == 1: - return embeddings - return embeddings.narrow(0, slot_offset, local_num_tokens) - - @abstractmethod - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: - """Start an asynchronous lookup when supported.""" - raise NotImplementedError - - -class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): - """Quantization interface shared by resident and pinned PLE tables.""" - - # PLE post-load processing only validates scales in their current storage. - requires_device_loading: bool = False - - @staticmethod - def from_quant_config( - quant_config: QuantizationConfig | None, - prefix: str, - embedding_dtype: str | None = None, - ) -> "Qwen4ExpPLEEmbeddingMethod": - """Select the concrete PLE embedding format for a layer.""" - if embedding_dtype == "float8_e4m3fn": - return Qwen4ExpPLEFp8EmbeddingMethod() - if quant_config is None: - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if isinstance(quant_config, ModelOptMixedPrecisionConfig): - if quant_config._resolve_quant_algo(prefix) == "FP8": - return Qwen4ExpPLEFp8EmbeddingMethod() - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if isinstance( - quant_config, ModelOptQuantConfigBase - ) and quant_config.is_layer_excluded(prefix): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if not isinstance(quant_config, Fp8Config): - raise NotImplementedError( - "Qwen4Exp PLE embedding does not support quantization config " - f"{type(quant_config).__name__}" - ) - - ignored_layers = quant_config.ignored_layers - if is_layer_skipped( - prefix, - ignored_layers, - quant_config.packed_modules_mapping, - match_mode=quant_config.ignored_layers_match_mode, - ): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - # PLE checkpoint shards form one runtime embedding parameter. - shard_prefix = f"{prefix}.shard_" - if any(name.startswith(shard_prefix) for name in ignored_layers): - return Qwen4ExpPLEUnquantizedEmbeddingMethod() - if not quant_config.is_checkpoint_fp8_serialized: - raise NotImplementedError( - "Qwen4Exp PLE embedding only supports serialized FP8 checkpoints" - ) - return Qwen4ExpPLEFp8EmbeddingMethod() - - def apply( - self, - layer: nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - raise NotImplementedError("PLE weights only support embedding lookup") - - def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: - return F.embedding(input_, layer.weight) - - @abstractmethod - def dequantize( - self, - layer: nn.Module, - embeddings: torch.Tensor, - output_dtype: torch.dtype, - ) -> torch.Tensor: - """Convert looked-up PLE rows to the activation dtype.""" - raise NotImplementedError - - -class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): - """Unquantized PLE embedding storage and lookup semantics.""" - - def create_weights( - self, - layer: Qwen4ExpPLEEmbedding, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - del input_size, output_size - weight = nn.Parameter( - layer.allocate_embedding_weight( - sum(output_partition_sizes), - input_size_per_partition, - params_dtype, - ), - requires_grad=False, - ) - set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) - set_weight_attrs(weight, extra_weight_attrs) - layer.register_parameter("weight", weight) - - def dequantize( - self, - layer: nn.Module, - embeddings: torch.Tensor, - output_dtype: torch.dtype, - ) -> torch.Tensor: - del layer, output_dtype - return embeddings - - -class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): - """FP8 PLE embedding with one global checkpoint scale.""" - - def create_weights( - self, - layer: Qwen4ExpPLEEmbedding, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - del input_size, output_size, params_dtype - weight_loader = extra_weight_attrs.get("weight_loader") - weight = ModelWeightParameter( - data=layer.allocate_embedding_weight( - sum(output_partition_sizes), - input_size_per_partition, - torch.float8_e4m3fn, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - weight_scale = create_fp8_scale_parameter( - PerTensorScaleParameter, - output_partition_sizes, - input_size_per_partition, - None, - weight_loader, - scale_dtype=torch.float32, - ) - layer.register_parameter("weight_scale", weight_scale) - - def process_weights_after_loading(self, layer: nn.Module) -> None: - """Reject FP8 PLE checkpoints without a global scale.""" - sentinel = torch.finfo(torch.float32).min - if torch.any(layer.weight_scale == sentinel): - raise ValueError("FP8 PLE checkpoint is missing its global scale") - - def dequantize( - self, - layer: nn.Module, - embeddings: torch.Tensor, - output_dtype: torch.dtype, - ) -> torch.Tensor: - weight_scale = getattr(layer, "weight_scale", None) - if weight_scale is None: - raise RuntimeError("FP8 PLE embedding is missing its global scale") - if weight_scale.device != embeddings.device: - raise RuntimeError("FP8 PLE embedding scale must be on the output device") - return embeddings.to(output_dtype) * weight_scale.to(output_dtype) - - -class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding): - """PLE table allocated on the active model device.""" - - def allocate_embedding_weight( - self, - num_embeddings: int, - embedding_dim: int, - dtype: torch.dtype, - ) -> torch.Tensor: - """Allocate the complete PLE weight on the active device.""" - return torch.empty(num_embeddings, embedding_dim, dtype=dtype) - - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: - """Resident embedding prefetch is a no-op.""" - return None - - def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor: - """Gather ETP inputs, look up embeddings, and select local rows.""" - slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0]) - gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) - embeddings = super().forward(gathered_ids) - return self._select_embeddings( - embeddings, - ngram_ids.shape[0], - slot_offset, - ) - - -@triton.jit -def _lookup_ple_embedding_from_pinned_kernel( - weight_ptr, - ids_ptr, - output_ptr, - embedding_dim, - tp_vocab_start, - tp_vocab_end, - BLOCK_D: tl.constexpr, -): - """Look up TP-owned PLE rows through a CUDA view of pinned host memory.""" - row_id = tl.program_id(0) - global_idx = tl.load(ids_ptr + row_id) - in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end) - local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0) - offsets = tl.arange(0, BLOCK_D) - store_mask = offsets < embedding_dim - load_mask = store_mask & in_range - values = tl.load( - weight_ptr + local_idx * embedding_dim + offsets, - mask=load_mask, - other=0.0, - ) - tl.store( - output_ptr + row_id * embedding_dim + offsets, - values, - mask=store_mask, - ) - - -class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): - """PLE table loaded into pinned CPU memory and looked up through UVA.""" - - supports_prefetch: ClassVar[bool] = True - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - *, - params_dtype: torch.dtype, - padding_size: int, - prefix: str, - embedding_method: Qwen4ExpPLEEmbeddingMethod, - num_ngram_heads: int = 1, - max_total_tokens: int = 0, - data_parallel_rank: int = 0, - ) -> None: - if not is_uva_available(): - raise RuntimeError("Engram CPU offload requires UVA support") - super().__init__( - num_embeddings, - embedding_dim, - params_dtype=params_dtype, - padding_size=padding_size, - prefix=prefix, - embedding_method=embedding_method, - num_ngram_heads=num_ngram_heads, - max_total_tokens=max_total_tokens, - data_parallel_rank=data_parallel_rank, - ) - self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) - self._block_d = triton.next_power_of_2(self.embedding_dim) - self._prefetch_stream = torch.cuda.Stream(device=self._uva_weight.device) - self._prefetch_buffer = torch.empty( - max_total_tokens * self.etp_data_parallel_size, - num_ngram_heads, - self.embedding_dim, - dtype=self.weight.dtype, - device=self._uva_weight.device, - ) - self._output_dim = num_ngram_heads * self.embedding_dim - - def allocate_embedding_weight( - self, - num_embeddings: int, - embedding_dim: int, - dtype: torch.dtype, - ) -> torch.Tensor: - """Allocate the complete PLE weight directly in pinned CPU memory.""" - return torch.empty( - num_embeddings, - embedding_dim, - dtype=dtype, - device="cpu", - pin_memory=True, - ) - - def _lookup( - self, - input_ids: torch.Tensor, - output: torch.Tensor | None = None, - ) -> torch.Tensor: - """Look up local ETP rows while preserving the weight storage dtype.""" - expected_shape = (*input_ids.shape, self.embedding_dim) - if output is None: - output = torch.empty( - expected_shape, - dtype=self.weight.dtype, - device=input_ids.device, - ) - elif ( - tuple(output.shape) != expected_shape - or output.dtype != self.weight.dtype - or output.device != input_ids.device - ): - raise ValueError( - "PLE prefetch output must match the input shape, weight dtype, " - "and input device" - ) - - flat_ids = input_ids.reshape(-1).long() - if flat_ids.numel(): - _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)]( - self._uva_weight, - flat_ids, - output, - self.embedding_dim, - self.shard_indices.org_vocab_start_index, - self.shard_indices.org_vocab_end_index, - BLOCK_D=self._block_d, - ) - return output - - def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: - """Combine pinned lookup results owned by different ETP ranks.""" - if self.tp_size == 1: - return embeddings - assert self.parallel_group is not None - if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): - # Each vocabulary row has one owner, so reduce the raw FP8 bytes. - reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8)) - return reduced.view(embeddings.dtype) - return self.parallel_group.all_reduce(embeddings) - - @eager_break_during_capture - def start_prefetch( - self, - hidden_states: torch.Tensor, - ngram_ids: torch.Tensor, - ) -> None: - """Gather ETP IDs and launch their UVA lookup on the side stream.""" - slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0]) - gathered_ids = self._gather_dp_ids(ngram_ids, slot_size) - active_output = self._prefetch_buffer[: gathered_ids.shape[0]] - prefetch_stream = self._prefetch_stream - prefetch_stream.wait_stream(torch.cuda.current_stream()) - gathered_ids.record_stream(prefetch_stream) - with torch.cuda.stream(prefetch_stream): - self._lookup(gathered_ids, output=active_output) - - @eager_break_during_capture - def _finalize_prefetch( - self, - prefetch_output: torch.Tensor, - output: torch.Tensor, - ) -> None: - """Join the side stream, reduce ETP shards, and select local rows.""" - torch.cuda.current_stream().wait_stream(self._prefetch_stream) - slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0]) - active_output = prefetch_output[: slot_size * self.etp_data_parallel_size] - embeddings = self._reduce_etp_embeddings(active_output) - embeddings = self._select_embeddings( - embeddings, - output.shape[0], - slot_offset, - ) - output.copy_(embeddings.flatten(-2)) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Finish the pinned lookup into graph-owned output storage.""" - output = self._prefetch_buffer.new_empty( - (hidden_states.shape[0], self._output_dim) - ) - self._finalize_prefetch(self._prefetch_buffer, output) - return output +__all__ = [ + "Qwen4ExpPLEDeviceEmbedding", + "Qwen4ExpPLEEmbedding", + "Qwen4ExpPLEEmbeddingMethod", + "Qwen4ExpPLEFp8EmbeddingMethod", + "Qwen4ExpPLEPinnedHostEmbedding", + "Qwen4ExpPLEUnquantizedEmbeddingMethod", + "Qwen4ExpNGramEmbedding", +] class Qwen4ExpNGramEmbedding(nn.Module):