From e251861aab69d3eb438d1a76ddec036ae45cfd4c Mon Sep 17 00:00:00 2001 From: Stefano Castagnetta Date: Thu, 10 Sep 2026 15:34:37 +0200 Subject: [PATCH] [Quantization] Support packed NVFP4 Qwen4Exp PLE embeddings Honor the explicit NVFP4 PLE dtype in mixed-precision checkpoints. Retain packed rows and block scales in device or pinned host storage, and decode selected rows before ETP reduction. Validate local shard coverage after all checkpoint streams finish. Cover streamed loading, numerical decoding, pinned prefetch, and CUDA graph replay in the existing PLE tests and H200 CI lane. Signed-off-by: Stefano Castagnetta --- .buildkite/test_areas/models_basic.yaml | 1 + tests/models/qwen4_exp/test_ple.py | 182 +++++++++++- .../qwen4_exp/nvidia/ngram_embedding.py | 276 ++++++++++++++++-- 3 files changed, 431 insertions(+), 28 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 93d6962868a4..5ab51fad00e3 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -108,6 +108,7 @@ steps: - tests/models/qwen4_exp/ commands: - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_pre_indexer.py models/qwen4_exp/test_qsa_reference.py + - pytest -v -s models/qwen4_exp/test_ple.py -k nvfp4 mirror: amd: label: ":amd: (MI300) Qwen4 Exp" diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 020e1cc8fc78..7db504af7c6a 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -29,6 +29,7 @@ Qwen4ExpPLEDeviceEmbedding, Qwen4ExpPLEEmbeddingMethod, Qwen4ExpPLEFp8EmbeddingMethod, + Qwen4ExpPLENvFp4EmbeddingMethod, Qwen4ExpPLEPinnedHostEmbedding, Qwen4ExpPLEUnquantizedEmbeddingMethod, ) @@ -43,9 +44,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, ) @@ -194,6 +196,184 @@ 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) + + +@pytest.mark.parametrize("missing", ["shard_1.weight", "shard_1.weight_scale", "all"]) +def test_nvfp4_ple_rejects_missing_local_shards(monkeypatch, missing): + module, tensors, _, _ = _make_nvfp4_ngram_embedding(monkeypatch) + module.load_weights( + (name, tensor) + for name, tensor in tensors + if missing != name.removeprefix("ngram_embedding.") + and not (missing == "all" and ".shard_" in name) + ) + layer = module.ngram_embedding + with pytest.raises(ValueError, match="checkpoint is missing"): + layer.embedding_method.process_weights_after_loading(layer) + + +@pytest.mark.parametrize("scale", [None, 0.0, -1.0, float("nan"), float("inf")]) +def test_nvfp4_ple_rejects_missing_or_invalid_global_scale(monkeypatch, scale): + module, tensors, _, _ = _make_nvfp4_ngram_embedding(monkeypatch) + module.load_weights(tensors[:-1]) + layer = module.ngram_embedding + if scale is not None: + layer.weight_scale_2.data.fill_(scale) + with pytest.raises(ValueError, match="finite positive global scale"): + layer.embedding_method.process_weights_after_loading(layer) + + +@pytest.mark.parametrize("suffix", ["weight", "weight_scale"]) +def test_nvfp4_ple_rejects_incorrect_packed_shape_or_dtype(monkeypatch, suffix): + module, tensors, _, _ = _make_nvfp4_ngram_embedding(monkeypatch) + name, tensor = next( + pair for pair in tensors if pair[0] == f"ngram_embedding.shard_0.{suffix}" + ) + with pytest.raises(ValueError, match="Shape mismatch"): + module.load_weights([(name, tensor[:, :-1])]) + with pytest.raises(ValueError, match="requires torch"): + module.load_weights([(name, tensor.float())]) + + +def test_nvfp4_ple_uses_explicit_storage_dtype_in_mixed_checkpoint(): + 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": { + "model.layers.0.mlp.experts": {"quant_algo": "NVFP4"} + }, + } + } + ) + assert isinstance( + Qwen4ExpPLEEmbeddingMethod.from_quant_config(config, prefix, "nvfp4"), + Qwen4ExpPLENvFp4EmbeddingMethod, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("pinned", [False, True]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_nvfp4_ple_cuda_lookup_and_graph_replay(monkeypatch, pinned, rank): + 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: + hidden_states = torch.empty(2, 320, device="cuda", dtype=torch.bfloat16) + layer.start_prefetch(hidden_states, ids) + torch.testing.assert_close( + layer(hidden_states).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, diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 84c4e7c5e1bb..5a7c147572d2 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -27,6 +27,10 @@ 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, + dequantize_to_dtype, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) @@ -43,7 +47,7 @@ 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.ple import PLEVocabParallelEmbedding, compute_ple_shard_overlap from .ops.ple import ple_ngram_ids logger = init_logger(__name__) @@ -171,6 +175,8 @@ def from_quant_config( 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 quant_config is None: @@ -218,6 +224,25 @@ def apply( 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 layer.weight.dtype + + def lookup_from_pinned( + self, + layer: "Qwen4ExpPLEPinnedHostEmbedding", + ids: torch.Tensor, + output: torch.Tensor, + ) -> None: + _lookup_ple_embedding_from_pinned_kernel[(ids.numel(),)]( + layer._uva_weight, + ids, + output, + layer.embedding_dim, + layer.shard_indices.org_vocab_start_index, + layer.shard_indices.org_vocab_end_index, + BLOCK_D=layer._block_d, + ) + @abstractmethod def dequantize( self, @@ -322,6 +347,194 @@ def dequantize( return embeddings.to(output_dtype) * weight_scale.to(output_dtype) +class Qwen4ExpPLENvFp4EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): + """Packed E2M1 rows, E4M3 scales per 16 values, and one global 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 + if input_size_per_partition % 16: + raise ValueError("NVFP4 PLE embedding dimension must be divisible by 16") + self._loaded_ranges: dict[str, set[tuple[int, int]]] = { + "weight": set(), + "weight_scale": set(), + } + weight_loader = extra_weight_attrs.get("weight_loader") + for name, width, dtype in ( + ("weight", input_size_per_partition // 2, torch.uint8), + ("weight_scale", input_size_per_partition // 16, torch.float8_e4m3fn), + ): + layer.register_parameter( + name, + ModelWeightParameter( + data=layer.allocate_embedding_weight( + sum(output_partition_sizes), width, dtype + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + layer.register_parameter( + "weight_scale_2", + create_fp8_scale_parameter( + PerTensorScaleParameter, + output_partition_sizes, + input_size_per_partition, + None, + weight_loader, + scale_dtype=torch.float32, + ), + ) + if layer.weight_scale.is_pinned(): + layer._uva_weight_scale = get_accelerator_view_from_cpu_tensor( + layer.weight_scale + ) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + scale = layer.weight_scale_2 + if not torch.all(torch.isfinite(scale) & (scale > 0)): + raise ValueError( + "NVFP4 PLE checkpoint requires a finite positive global scale" + ) + expected_rows = ( + layer.shard_indices.org_vocab_end_index + - layer.shard_indices.org_vocab_start_index + ) + for name, ranges in self._loaded_ranges.items(): + loaded_end = 0 + for start, end in sorted(ranges): + if start > loaded_end: + break + loaded_end = max(loaded_end, end) + if loaded_end != expected_rows: + raise ValueError( + f"NVFP4 PLE checkpoint is missing {name} rows " + f"starting at local row {loaded_end}" + ) + + def record_loaded_rows( + self, + layer: Qwen4ExpPLEEmbedding, + name: str, + checkpoint_start: int, + checkpoint_rows: int, + ) -> None: + overlap = compute_ple_shard_overlap( + checkpoint_start=checkpoint_start, + checkpoint_rows=checkpoint_rows, + tp_start=layer.shard_indices.org_vocab_start_index, + tp_end=layer.shard_indices.org_vocab_end_index, + ) + if overlap is not None: + start = overlap.destination_start + self._loaded_ranges[name].add((start, start + overlap.row_count)) + + 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(): + _lookup_nvfp4_ple_embedding_kernel[(ids.numel(),)]( + layer.weight, + layer.weight_scale, + 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, + layer._uva_weight_scale, + 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=layer._block_d, + ) + + def dequantize( + self, + layer: nn.Module, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + # NVFP4 rows are dequantized during lookup, before the ETP reduction. + 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, +): + row = tl.program_id(0) + 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) + mask = owned & (offsets < embedding_dim) + packed = tl.load( + weight_ptr + local_idx * (embedding_dim // 2) + offsets // 2, + mask=mask, + other=0, + ) + codes = (packed >> ((offsets % 2) * 4)) & 0xF + scales = tl.load( + scale_ptr + local_idx * (embedding_dim // 16) + offsets // 16, + mask=mask, + other=0.0, + ).to(tl.float32) + 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, offsets < embedding_dim + ) + + class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding): """PLE table allocated on the active model device.""" @@ -422,7 +635,7 @@ def __init__( max_total_tokens * self.etp_data_parallel_size, num_ngram_heads, self.embedding_dim, - dtype=self.weight.dtype, + dtype=self.embedding_method.lookup_dtype(self), device=self._uva_weight.device, ) self._output_dim = num_ngram_heads * self.embedding_dim @@ -447,35 +660,28 @@ def _lookup( 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 output format.""" expected_shape = (*input_ids.shape, self.embedding_dim) + output_dtype = self.embedding_method.lookup_dtype(self) if output is None: output = torch.empty( expected_shape, - dtype=self.weight.dtype, + dtype=output_dtype, device=input_ids.device, ) elif ( tuple(output.shape) != expected_shape - or output.dtype != self.weight.dtype + or output.dtype != output_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.embedding_dim, - 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 _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor: @@ -878,6 +1084,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: 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 + ) + 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] @@ -893,9 +1108,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: 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) @@ -904,28 +1119,35 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: 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"Shape mismatch for PLE embedding shard {shard_index} " + f"{suffix}: " f"expected {expected_shape}, got " f"{tuple(loaded_weight.shape)}" ) - embedding.weight.weight_loader( - embedding.weight, + if nvfp4_method 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") + if nvfp4_method: + nvfp4_method.record_loaded_rows( + embedding, suffix, checkpoint_start, expected_rows + ) + loaded.add(f"ngram_embedding.{suffix}") continue regular_weights.append((name, loaded_weight))