diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 6845c2ebfb..848ab85ef7 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -212,7 +212,7 @@ steps: - tests/models/qwen4_exp/ commands: - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_prepare.py models/qwen4_exp/test_qsa_reference.py - - pytest -v -s models/qwen4_exp/test_ple.py -k test_fused_conv_correctness + - pytest -v -s models/qwen4_exp/test_ple.py mirror: amd: label: ":amd: (MI355 DPX) Qwen4 Exp" diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index a4a428eccc..119ec3b1ac 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -41,6 +41,7 @@ from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, Qwen4ExpPLEEmbeddingMethod, Qwen4ExpPLEFp8EmbeddingMethod, + Qwen4ExpPLENvFp4EmbeddingMethod, Qwen4ExpPLEPinnedHostEmbedding, Qwen4ExpPLEUnquantizedEmbeddingMethod, ) @@ -174,9 +175,10 @@ def _mock_etp_group( monkeypatch: pytest.MonkeyPatch, world_size: int = 1, all_reduce=lambda tensor: tensor, + rank: int = 0, ) -> None: group = SimpleNamespace( - rank_in_group=0, + rank_in_group=rank, world_size=world_size, all_reduce=all_reduce, ) @@ -325,6 +327,162 @@ def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None: ) +def _make_nvfp4_ngram_embedding(monkeypatch, *, rank=0, device="cpu", pinned=False): + _mock_etp_group(monkeypatch, world_size=2, rank=rank) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_rank", lambda: rank + ) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 2 + ) + cls = Qwen4ExpPLEPinnedHostEmbedding if pinned else Qwen4ExpPLEDeviceEmbedding + with torch.device(device): + embedding = cls( + 8, + 160, + params_dtype=torch.bfloat16, + padding_size=2, + prefix="test.ple_embedding.ngram_embedding", + embedding_method=Qwen4ExpPLENvFp4EmbeddingMethod(), + num_ngram_heads=2, + max_total_tokens=4, + ) + module = _make_ngram_embedding_for_load_test() + module.ngram_embedding = embedding + module.split_ngram_parts = 3 + codes = torch.arange(8 * 80).reshape(8, 80).to(torch.uint8) + scales = ((torch.arange(80).reshape(8, 10) + 1) / 4).to(torch.float8_e4m3fn) + tensors = [] + # Checkpoint shards cross ETP boundaries; scales can arrive before weights. + for shard in (2, 0, 1): + for suffix, tensor in (("weight_scale", scales), ("weight", codes)): + tensors.append( + (f"ngram_embedding.shard_{shard}.{suffix}", tensor[shard * 3 :][:3]) + ) + tensors.append(("ngram_embedding.weight_scale_2", torch.tensor(0.37))) + return module, tensors, codes, scales + + +def _reference_nvfp4_ple(codes, scales): + nibbles = torch.stack((codes & 15, codes >> 4), dim=-1).reshape(8, 160) + values = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) + decoded = values[(nibbles & 7).long()] * torch.where(nibbles < 8, 1.0, -1.0) + return (decoded * scales.float().repeat_interleave(16, -1) * 0.37).bfloat16() + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_nvfp4_ple_loads_streamed_shards_across_etp_boundaries(monkeypatch, rank): + module, tensors, codes, scales = _make_nvfp4_ngram_embedding(monkeypatch, rank=rank) + loaded = set() + for start in range(0, len(tensors), 3): + loaded.update(module.load_weights(iter(tensors[start : start + 3]))) + layer = module.ngram_embedding + layer.embedding_method.process_weights_after_loading(layer) + + assert loaded == { + "ngram_embedding.weight", + "ngram_embedding.weight_scale", + "ngram_embedding.weight_scale_2", + } + assert layer.weight.dtype == torch.uint8 + assert layer.weight_scale.dtype == torch.float8_e4m3fn + assert layer.weight.numel() + layer.weight_scale.numel() == 4 * (80 + 10) + torch.testing.assert_close(layer.weight, codes[rank * 4 :][:4]) + torch.testing.assert_close( + layer.weight_scale.float(), scales[rank * 4 :][:4].float() + ) + ids = torch.tensor([[3, 4], [7, 0], [4, 3]]) + output = layer(ids) + expected = _reference_nvfp4_ple(codes, scales)[ids] + expected[(ids < rank * 4) | (ids >= (rank + 1) * 4)] = 0 + torch.testing.assert_close(output, expected, atol=0, rtol=0) + + +def test_nvfp4_ple_rejects_missing_global_scale(monkeypatch): + module, tensors, _, _ = _make_nvfp4_ngram_embedding(monkeypatch) + module.load_weights(tensors[:-1]) + layer = module.ngram_embedding + with pytest.raises(ValueError, match="missing its global scale"): + layer.embedding_method.process_weights_after_loading(layer) + + +def test_nvfp4_ple_resolves_from_mixed_precision_quantized_layers(): + prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding" + config = ModelOptMixedPrecisionConfig.from_config( + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "quantized_layers": { + prefix: {"quant_algo": "NVFP4", "group_size": 16}, + "model.language_model.layers.0.mlp.experts": { + "quant_algo": "NVFP4", + "group_size": 16, + }, + }, + } + } + ) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(config, prefix), + Qwen4ExpPLENvFp4EmbeddingMethod, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("pinned", [False, True]) +def test_nvfp4_ple_cuda_lookup_and_graph_replay(monkeypatch, pinned): + # Rank 1 exercises both the vocab offset and out-of-shard masking. + rank = 1 + module, tensors, codes, scales = _make_nvfp4_ngram_embedding( + monkeypatch, rank=rank, device="cuda", pinned=pinned + ) + module.load_weights(iter(tensors)) + layer = module.ngram_embedding + layer.embedding_method.process_weights_after_loading(layer) + ids = torch.tensor([[3, 4], [7, 0]], device="cuda") + lookup_fn = layer._lookup if pinned else torch.compile(layer, fullgraph=True) + + def lookup(): + return lookup_fn(ids) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + lookup() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = lookup() + for new_ids in ([[3, 4], [7, 0]], [[4, 3], [0, 7]]): + ids.copy_(torch.tensor(new_ids, device="cuda")) + graph.replay() + cpu_ids = ids.cpu() + expected = _reference_nvfp4_ple(codes, scales)[cpu_ids] + expected[(cpu_ids < rank * 4) | (cpu_ids >= (rank + 1) * 4)] = 0 + torch.testing.assert_close(output.cpu(), expected, atol=0, rtol=0) + if pinned: + layer.start_prefetch(ids) + torch.testing.assert_close( + layer.finalize_prefetch(total_tokens=ids.shape[0]).cpu(), + expected.flatten(-2), + atol=0, + rtol=0, + ) + assert layer.weight.device.type == ("cpu" if pinned else "cuda") + assert layer.weight_scale.device == layer.weight.device + assert layer.weight_scale_2.device.type == "cuda" + assert lookup_fn(ids[:0]).shape == (0, 2, 160) + padded_ids = torch.tensor([[8, -1]], device="cuda") + torch.testing.assert_close( + lookup_fn(padded_ids), + torch.zeros(1, 2, 160, device="cuda", dtype=torch.bfloat16), + atol=0, + rtol=0, + ) + + @pytest.mark.parametrize(("dp_rank", "local_tokens"), [(0, 2), (1, 3)]) def test_etp_lookup_gathers_and_returns_dp_local_rows( monkeypatch: pytest.MonkeyPatch, @@ -865,7 +1023,7 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: "quantized_layers": { prefix: {"quant_algo": "FP8"}, "model.language_model.layers.2.moe.gate_proj": { - "quant_algo": "NVFP4" + "quant_algo": "MXFP8" }, }, } @@ -876,13 +1034,12 @@ def test_ple_fp8_embedding_supports_mixed_precision_config() -> None: Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix), Qwen4ExpPLEFp8EmbeddingMethod, ) - assert isinstance( + # A quant_algo the PLE table cannot load must not fall back to unquantized. + with pytest.raises(ValueError, match="Unsupported quant_algo 'MXFP8'"): Qwen4ExpPLEEmbeddingMethod.from_quant_config( quant_config, "model.language_model.layers.2.moe.gate_proj", - ), - Qwen4ExpPLEUnquantizedEmbeddingMethod, - ) + ) def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None: diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py index a66d481159..568ae3fd1b 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py @@ -45,6 +45,22 @@ def _e2m1_inline(nibble): return tl.where(sign == 1, -val, val) +@triton.jit +def _e4m3_inline(byte): + """Decode an FP8 E4M3FN byte to float32 without an FP8 pointer type. + + Typed FP8 loads need SM89+ in Triton; reading the scale as uint8 and + decoding here keeps the caller portable. + """ + bits = byte.to(tl.int32) + exponent = (bits >> 3) & 0xF + mantissa = (bits & 0x7).to(tl.float32) + normal = (1.0 + mantissa * 0.125) * tl.exp2((exponent - 7).to(tl.float32)) + val = tl.where(exponent == 0, mantissa * 0.001953125, normal) + val = tl.where((bits & 0x7F) == 0x7F, float("nan"), val) + return tl.where((bits & 0x80) != 0, -val, val) + + @triton.jit def _nvfp4_gathered_bias_kernel( markov_ptr, diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index fbc07b73dd..30595df32d 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -32,6 +32,7 @@ from vllm.v1.attention.backends.short_conv_attn import ( from ..common.ngram_embedding import ( Qwen4ExpPLEDeviceEmbedding, Qwen4ExpPLEEmbeddingMethod, + Qwen4ExpPLENvFp4EmbeddingMethod, Qwen4ExpPLEPinnedHostEmbedding, ) from ..common.ops.ple import ple_conv, ple_gate, ple_ngram_ids @@ -241,6 +242,11 @@ class Qwen4ExpNGramEmbedding(nn.Module): embedding_prefix, getattr(config, "ple_embedding_dtype", None), ) + if isinstance(embedding_quant_method, Qwen4ExpPLENvFp4EmbeddingMethod): + raise NotImplementedError( + "NVFP4 Qwen4Exp PLE tables are not supported on ROCm yet; use a " + "checkpoint with a BF16 or FP8 PLE table." + ) engram_config = get_current_vllm_config().engram_config embedding_cls = ( Qwen4ExpPLEPinnedHostEmbedding diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 85a382c5ec..5bf75cb551 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -36,6 +36,11 @@ from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig from vllm.model_executor.layers.quantization.utils.fp8_utils import ( create_fp8_scale_parameter, ) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + _e2m1_inline, + _e4m3_inline, + dequantize_to_dtype, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) @@ -178,13 +183,28 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): embedding_dtype: str | None = None, ) -> "Qwen4ExpPLEEmbeddingMethod": """Select the concrete PLE embedding format for a layer.""" + if embedding_dtype == "nvfp4": + return Qwen4ExpPLENvFp4EmbeddingMethod() if embedding_dtype == "float8_e4m3fn": return Qwen4ExpPLEFp8EmbeddingMethod() + if embedding_dtype is not None: + raise ValueError( + "ple_embedding_dtype must be 'nvfp4' or 'float8_e4m3fn', " + f"got {embedding_dtype!r}" + ) if quant_config is None: return Qwen4ExpPLEUnquantizedEmbeddingMethod() if isinstance(quant_config, ModelOptMixedPrecisionConfig): - if quant_config._resolve_quant_algo(prefix) == "FP8": + quant_algo = quant_config._resolve_quant_algo(prefix) + if quant_algo == "FP8": return Qwen4ExpPLEFp8EmbeddingMethod() + if quant_algo == "NVFP4": + return Qwen4ExpPLENvFp4EmbeddingMethod() + if quant_algo is not None: + raise ValueError( + f"Unsupported quant_algo {quant_algo!r} for the PLE embedding " + f"{prefix!r}; expected FP8 or NVFP4" + ) return Qwen4ExpPLEUnquantizedEmbeddingMethod() if isinstance( quant_config, ModelOptQuantConfigBase @@ -238,6 +258,28 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: return F.embedding(input_, layer.weight) + def lookup_dtype(self, layer: nn.Module) -> torch.dtype: + """Return the dtype of rows that a pinned-host lookup writes.""" + return layer.weight.dtype + + def lookup_from_pinned( + self, + layer: "Qwen4ExpPLEPinnedHostEmbedding", + ids: torch.Tensor, + output: torch.Tensor, + ) -> None: + """Copy the ETP-owned rows for ``ids`` from pinned host memory.""" + row_bytes = layer.embedding_dim * layer.weight.element_size() + _lookup_ple_embedding_from_pinned_kernel[(ids.numel(),)]( + layer._uva_weight, + ids, + output, + row_bytes, + layer.shard_indices.org_vocab_start_index, + layer.shard_indices.org_vocab_end_index, + BLOCK_D=triton.next_power_of_2(row_bytes), + ) + @abstractmethod def dequantize( self, @@ -342,6 +384,177 @@ class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): return embeddings.to(output_dtype) * weight_scale.to(output_dtype) +# NVFP4 packs two E2M1 values per byte and shares one E4M3 scale per block. +_NVFP4_VALUES_PER_BYTE = 2 +_NVFP4_BLOCK_SIZE = 16 + + +class Qwen4ExpPLENvFp4EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): + """NVFP4 PLE embedding: packed E2M1 rows, block scales, one global scale. + + Rows stay packed in storage. A lookup decodes only the selected rows to + the activation dtype, before the ETP reduction. + """ + + 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 + if input_size_per_partition % _NVFP4_BLOCK_SIZE: + raise ValueError( + "NVFP4 PLE embedding dimension must be divisible by " + f"{_NVFP4_BLOCK_SIZE}, got {input_size_per_partition}" + ) + weight_loader = extra_weight_attrs.get("weight_loader") + num_rows = sum(output_partition_sizes) + for name, width, dtype in ( + ( + "weight", + input_size_per_partition // _NVFP4_VALUES_PER_BYTE, + torch.uint8, + ), + ( + "weight_scale", + input_size_per_partition // _NVFP4_BLOCK_SIZE, + torch.float8_e4m3fn, + ), + ): + parameter = ModelWeightParameter( + data=layer.allocate_embedding_weight(num_rows, width, dtype), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter(name, parameter) + weight_scale_2 = create_fp8_scale_parameter( + PerTensorScaleParameter, + output_partition_sizes, + input_size_per_partition, + None, + weight_loader, + scale_dtype=torch.float32, + ) + layer.register_parameter("weight_scale_2", weight_scale_2) + if layer.weight_scale.is_pinned(): + self._uva_weight_scale = get_accelerator_view_from_cpu_tensor( + layer.weight_scale + ) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Reject NVFP4 PLE checkpoints without a global scale.""" + sentinel = torch.finfo(torch.float32).min + if torch.any(layer.weight_scale_2 == sentinel): + raise ValueError("NVFP4 PLE checkpoint is missing its global scale") + + def lookup_dtype(self, layer: nn.Module) -> torch.dtype: + return layer.params_dtype + + def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: + ids = input_.reshape(-1) + if not input_.is_cuda: + rows = dequantize_to_dtype( + F.embedding(ids, layer.weight), + F.embedding(ids, layer.weight_scale), + layer.weight_scale_2, + layer.params_dtype, + swizzle=False, + ) + return rows.reshape(*input_.shape, layer.embedding_dim) + output = torch.empty( + (*input_.shape, layer.embedding_dim), + dtype=layer.params_dtype, + device=input_.device, + ) + if ids.numel(): + # Device rows are already local, so every ID in range is owned. + _lookup_nvfp4_ple_embedding_kernel[(ids.numel(),)]( + layer.weight, + layer.weight_scale.view(torch.uint8), + layer.weight_scale_2, + ids, + output, + layer.embedding_dim, + 0, + layer.weight.shape[0], + BLOCK_D=triton.next_power_of_2(layer.embedding_dim), + ) + return output + + def lookup_from_pinned( + self, + layer: "Qwen4ExpPLEPinnedHostEmbedding", + ids: torch.Tensor, + output: torch.Tensor, + ) -> None: + _lookup_nvfp4_ple_embedding_kernel[(ids.numel(),)]( + layer._uva_weight, + self._uva_weight_scale.view(torch.uint8), + layer.weight_scale_2, + ids, + output, + layer.embedding_dim, + layer.shard_indices.org_vocab_start_index, + layer.shard_indices.org_vocab_end_index, + BLOCK_D=triton.next_power_of_2(layer.embedding_dim), + ) + + def dequantize( + self, + layer: nn.Module, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + del layer + # The lookup already decoded the rows to the activation dtype. + return embeddings.to(output_dtype) + + +@triton.jit +def _lookup_nvfp4_ple_embedding_kernel( + weight_ptr, + scale_ptr, + global_scale_ptr, + ids_ptr, + output_ptr, + embedding_dim, + vocab_start, + vocab_end, + BLOCK_D: tl.constexpr, +): + """Decode owned NVFP4 PLE rows; write zeros for rows another rank owns.""" + row = tl.program_id(0).to(tl.int64) + idx = tl.load(ids_ptr + row).to(tl.int64) + owned = (idx >= vocab_start) & (idx < vocab_end) + local_idx = tl.where(owned, idx - vocab_start, 0) + offsets = tl.arange(0, BLOCK_D) + in_row = offsets < embedding_dim + mask = owned & in_row + packed = tl.load( + weight_ptr + local_idx * (embedding_dim // 2) + offsets // 2, + mask=mask, + other=0, + ) + # The low nibble holds the even element and the high nibble the odd one. + codes = (packed >> ((offsets % 2) * 4)) & 0xF + scales = _e4m3_inline( + tl.load( + scale_ptr + local_idx * (embedding_dim // 16) + offsets // 16, + mask=mask, + other=0, + ) + ) + global_scale = tl.load(global_scale_ptr).to(tl.float32) + values = _e2m1_inline(codes) * scales * global_scale + tl.store(output_ptr + row * embedding_dim + offsets, values, mask=in_row) + + class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding): """PLE table allocated on the active model device.""" @@ -442,8 +655,6 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): reduce_results=reduce_results, ) self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) - self._row_bytes = self.embedding_dim * self.weight.element_size() - self._block_d = triton.next_power_of_2(self._row_bytes) self._prefetch_stream: torch.cuda.Stream | None = None self._prefetch_buffer: torch.Tensor | None = None self._prefetch_alloc_lock = threading.Lock() @@ -471,35 +682,28 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): input_ids: torch.Tensor, output: torch.Tensor | None = None, ) -> torch.Tensor: - """Look up local ETP rows while preserving the weight storage dtype.""" + """Look up local ETP rows in the embedding method's lookup dtype.""" expected_shape = (*input_ids.shape, self.embedding_dim) + lookup_dtype = self.embedding_method.lookup_dtype(self) if output is None: output = torch.empty( expected_shape, - dtype=self.weight.dtype, + dtype=lookup_dtype, device=input_ids.device, ) elif ( tuple(output.shape) != expected_shape - or output.dtype != self.weight.dtype + or output.dtype != lookup_dtype or output.device != input_ids.device ): raise ValueError( - "PLE prefetch output must match the input shape, weight dtype, " + "PLE prefetch output must match the input shape, lookup 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._row_bytes, - self.shard_indices.org_vocab_start_index, - self.shard_indices.org_vocab_end_index, - BLOCK_D=self._block_d, - ) + self.embedding_method.lookup_from_pinned(self, flat_ids, output) return output def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor: @@ -549,7 +753,7 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): self._prefetch_rows, self._num_ngram_heads, self.embedding_dim, - dtype=self.weight.dtype, + dtype=self.embedding_method.lookup_dtype(self), device=self._uva_weight.device, ) self._prefetch_buffer = buffer diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 794eea87a6..1a36551046 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -23,6 +23,7 @@ from ..common.ngram_embedding import ( Qwen4ExpPLEEmbedding, Qwen4ExpPLEEmbeddingMethod, Qwen4ExpPLEFp8EmbeddingMethod, + Qwen4ExpPLENvFp4EmbeddingMethod, Qwen4ExpPLEPinnedHostEmbedding, Qwen4ExpPLEUnquantizedEmbeddingMethod, ) @@ -35,6 +36,7 @@ __all__ = [ "Qwen4ExpPLEEmbedding", "Qwen4ExpPLEEmbeddingMethod", "Qwen4ExpPLEFp8EmbeddingMethod", + "Qwen4ExpPLENvFp4EmbeddingMethod", "Qwen4ExpPLEPinnedHostEmbedding", "Qwen4ExpPLEUnquantizedEmbeddingMethod", "Qwen4ExpNGramEmbedding", @@ -420,6 +422,16 @@ class Qwen4ExpNGramEmbedding(nn.Module): loaded: set[str] = set() regular_weights: list[tuple[str, torch.Tensor]] = [] shard_prefix = "ngram_embedding.shard_" + embedding = self.ngram_embedding + method = getattr(embedding, "embedding_method", None) + nvfp4_method = ( + method if isinstance(method, Qwen4ExpPLENvFp4EmbeddingMethod) else None + ) + # NVFP4 tables split their block scales into shards like the rows. + shard_parameters = ("weight", "weight_scale") if nvfp4_method else ("weight",) + shard_size = ( + embedding.org_vocab_size + self.split_ngram_parts - 1 + ) // self.split_ngram_parts for name, loaded_weight in weights: leaf_name = name.rsplit(".", 1)[-1] @@ -435,9 +447,9 @@ class Qwen4ExpNGramEmbedding(nn.Module): buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) loaded.add(name) continue - if name.startswith(shard_prefix) and name.endswith(".weight"): - shard_text = name[len(shard_prefix) : -len(".weight")] - if not shard_text.isdigit(): + if name.startswith(shard_prefix): + shard_text, _, suffix = name[len(shard_prefix) :].partition(".") + if not shard_text.isdigit() or suffix not in shard_parameters: regular_weights.append((name, loaded_weight)) continue shard_index = int(shard_text) @@ -446,29 +458,40 @@ class Qwen4ExpNGramEmbedding(nn.Module): f"PLE embedding shard index {shard_index} exceeds " f"split_ngram_parts={self.split_ngram_parts}" ) - embedding = self.ngram_embedding - shard_size = ( - embedding.org_vocab_size + self.split_ngram_parts - 1 - ) // self.split_ngram_parts checkpoint_start = shard_index * shard_size expected_rows = max( 0, min(shard_size, embedding.org_vocab_size - checkpoint_start), ) - expected_shape = (expected_rows, embedding.embedding_dim) + parameter = getattr(embedding, suffix) + expected_shape = (expected_rows, parameter.shape[1]) if tuple(loaded_weight.shape) != expected_shape: raise ValueError( - f"Shape mismatch for PLE embedding shard {shard_index}: " - f"expected {expected_shape}, got " + f"Shape mismatch for PLE embedding shard {shard_index} " + f"{suffix}: expected {expected_shape}, got " f"{tuple(loaded_weight.shape)}" ) - embedding.weight.weight_loader( - embedding.weight, + # A cast would silently reinterpret packed E2M1 bytes. + if nvfp4_method is not None and loaded_weight.dtype != parameter.dtype: + raise ValueError( + f"NVFP4 PLE shard {shard_index} {suffix} requires " + f"{parameter.dtype}, got {loaded_weight.dtype}" + ) + parameter.weight_loader( + parameter, loaded_weight, checkpoint_start=checkpoint_start, ) - loaded.add("ngram_embedding.weight") + loaded.add(f"ngram_embedding.{suffix}") continue + if nvfp4_method is not None and name in ( + "ngram_embedding.weight", + "ngram_embedding.weight_scale", + ): + raise ValueError( + f"NVFP4 PLE tables must be stored as {shard_prefix}.weight " + f"and {shard_prefix}.weight_scale shards, got unsharded {name}" + ) regular_weights.append((name, loaded_weight)) if regular_weights: