From d5de41ec557b3570b2a384ff6f29bcc787ea96a1 Mon Sep 17 00:00:00 2001 From: BlivionIaG Date: Sun, 11 Oct 2026 06:21:12 +0200 Subject: [PATCH 1/2] feat(qwen4_exp): INT4 group-16 PLE table from VLLM_PLE_QUANT_DIR on the upstream PLE path Port of the fork's int4 PLE sidecar table onto upstream's PLE embedding. VLLM_PLE_QUANT_DIR pointing at a ples_int4 directory (META.json layout group16_int4_fp16scale_lownibblefirst, shard_.safetensors with weight_i4 and weight_scale) selects Qwen4ExpPLEInt4EmbeddingMethod: each rank copies its vocabulary rows of the packed table and group scales into the existing pinned-host (or device) storage, the checkpoint's BF16 table is skipped, and lookups dequantize to the activation dtype (Triton UVA kernel for the pinned table, torch gather for the device table). Flash-Next: 30 GB table instead of 95 GB BF16. Unset, nothing changes; other layouts are rejected. Tests: tests/models/qwen4_exp/test_ple.py 73/73 (load across shard boundaries + dequant, layout rejection, pinned UVA kernel incl. rows owned by another rank), test_ple_amd.py 8/8 on gfx1030. --- tests/models/qwen4_exp/test_ple.py | 105 ++++++++++ vllm/models/qwen4_exp/amd/ple_layer.py | 9 +- .../qwen4_exp/common/ngram_embedding.py | 196 +++++++++++++++++- .../qwen4_exp/nvidia/ngram_embedding.py | 5 + 4 files changed, 308 insertions(+), 7 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index b3038524eb96..ce5f18797547 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -28,6 +28,9 @@ from vllm.models.qwen4_exp.amd.ple_layer import ( Qwen4ExpPLELayer as Qwen4ExpPLELayerAMD, ) +from vllm.models.qwen4_exp.common.ngram_embedding import ( + Qwen4ExpPLEInt4EmbeddingMethod, +) from vllm.models.qwen4_exp.common.ple import ( PLEShardOverlap, compute_ple_shard_overlap, @@ -419,6 +422,108 @@ def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None: torch.testing.assert_close(output, (weight[[2, 0]] * 0.25).bfloat16()) +def _write_int4_ple_dir(path, rows: int, width: int, shards: int): + """Write a ples_int4-format table; return its fp32 dequantized rows.""" + import json + + from safetensors.torch import save_file + + generator = torch.Generator().manual_seed(0) + values = torch.randint(0, 16, (rows, width), generator=generator) + scale = torch.rand(rows, width // 16, generator=generator).to(torch.float16) + packed = (values[:, 0::2] | (values[:, 1::2] << 4)).to(torch.uint8) + per_shard = rows // shards + for n in range(shards): + rows_n = slice(n * per_shard, (n + 1) * per_shard) + save_file( + {"weight_i4": packed[rows_n].contiguous(), "weight_scale": scale[rows_n]}, + str(path / f"shard_{n}.safetensors"), + ) + meta = { + "layout": Qwen4ExpPLEInt4EmbeddingMethod.LAYOUT, + "shards": shards, + "rows": rows, + "width": width, + } + (path / "META.json").write_text(json.dumps(meta)) + return (values.float() - 8.0) * scale.float().repeat_interleave(16, dim=1) + + +def test_ple_int4_table_loads_from_quant_dir_and_dequantizes( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + reference = _write_int4_ple_dir(tmp_path, rows=8, width=32, shards=2) + monkeypatch.setattr( + ngram_embedding_module.envs, "VLLM_PLE_QUANT_DIR", str(tmp_path) + ) + method = Qwen4ExpPLEEmbeddingMethod.from_quant_config(None, "test.ple_embedding") + assert isinstance(method, Qwen4ExpPLEInt4EmbeddingMethod) + _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 + ) + layer = Qwen4ExpPLEDeviceEmbedding( + 8, + 32, + params_dtype=torch.float16, + padding_size=1, + prefix="test.ple_embedding", + embedding_method=method, + ) + assert layer.weight.dtype == torch.uint8 and layer.lookup_dtype == torch.float16 + + method.process_weights_after_loading(layer) + ids = torch.tensor([7, 0, 3, 4]) + output = layer(ids) + + assert output.dtype == torch.float16 + torch.testing.assert_close(output, reference[ids].half()) + assert layer.dequantize(output, torch.float16) is output + + +def test_ple_int4_rejects_other_table_layouts(monkeypatch, tmp_path) -> None: + (tmp_path / "META.json").write_text( + '{"layout": "per_row_e4m3", "shards": 1, "rows": 4, "width": 16}' + ) + monkeypatch.setattr( + ngram_embedding_module.envs, "VLLM_PLE_QUANT_DIR", str(tmp_path) + ) + with pytest.raises(NotImplementedError, match="ples_int4"): + Qwen4ExpPLEEmbeddingMethod.from_quant_config(None, "test.ple_embedding") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU for UVA") +def test_ple_int4_pinned_uva_lookup_matches_reference(tmp_path) -> None: + from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + + reference = _write_int4_ple_dir(tmp_path, rows=8, width=160, shards=1) + from safetensors.torch import load_file + + shard = load_file(str(tmp_path / "shard_0.safetensors")) + # This rank owns rows [2, 6); ids outside it must come back as zeros. + packed = shard["weight_i4"][2:6].contiguous().pin_memory() + scale = shard["weight_scale"][2:6].contiguous().pin_memory() + ids = torch.tensor([5, 0, 2, 7, 3], device="cuda") + output = torch.empty(5, 160, dtype=torch.float16, device="cuda") + + ngram_embedding_module._lookup_ple_int4_from_pinned_kernel[(5,)]( + get_accelerator_view_from_cpu_tensor(packed), + get_accelerator_view_from_cpu_tensor(scale), + ids, + output, + 160, + 2, + 6, + GROUP=16, + BLOCK_D=256, + ) + + expected = reference[ids.cpu()].half() + expected[[1, 3]] = 0 + torch.testing.assert_close(output.cpu(), expected) + + def test_ple_fp8_embedding_rejects_missing_global_scale(monkeypatch) -> None: layer = _make_fp8_embedding_layer(monkeypatch, load_scale=False) diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index 8bd1b982253d..fbef15304fc2 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -407,7 +407,7 @@ def forward( if embedding.supports_prefetch: output = ngram_ids.new_empty( (ngram_ids.shape[0], self.embedding_dim), - dtype=embedding.weight.dtype, + dtype=embedding.lookup_dtype, ) torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned( ngram_ids, @@ -452,6 +452,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded.add(name) continue if name.startswith(shard_prefix) and name.endswith(".weight"): + method = getattr(self.ngram_embedding, "embedding_method", None) + if getattr(method, "lookup_dequantizes", False): + # The table comes from VLLM_PLE_QUANT_DIR after loading. + loaded.add("ngram_embedding.weight") + continue shard_text = name[len(shard_prefix) : -len(".weight")] if not shard_text.isdigit(): regular_weights.append((name, loaded_weight)) @@ -1248,7 +1253,7 @@ def forward( if self._rdna_opaque_embedding: ngram = self.ple_embedding.ngram_embedding dtype = ( - ngram.weight.dtype if ngram.supports_prefetch else ngram.params_dtype + ngram.lookup_dtype if ngram.supports_prefetch else ngram.params_dtype ) embeddings = hidden_states.new_empty( (input_ids.shape[0], self.ple_embedding.embedding_dim), dtype=dtype diff --git a/vllm/models/qwen4_exp/common/ngram_embedding.py b/vllm/models/qwen4_exp/common/ngram_embedding.py index 7a32e9a8dbc3..874c2f180c25 100644 --- a/vllm/models/qwen4_exp/common/ngram_embedding.py +++ b/vllm/models/qwen4_exp/common/ngram_embedding.py @@ -7,6 +7,8 @@ Unified Virtual Addressing on any CUDA-alike platform. """ +import json +import os from abc import ABC, abstractmethod from typing import ClassVar @@ -14,6 +16,7 @@ import torch.nn.functional as F from torch import nn +import vllm.envs as envs 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 @@ -98,6 +101,13 @@ def allocate_embedding_weight( """Allocate storage for the complete embedding weight.""" raise NotImplementedError + @property + def lookup_dtype(self) -> torch.dtype: + """Dtype of looked-up rows: storage dtype unless the lookup dequantizes.""" + if self.embedding_method.lookup_dequantizes: + return self.params_dtype + return self.weight.dtype + def dequantize( self, embeddings: torch.Tensor, @@ -164,6 +174,9 @@ class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase): # PLE post-load processing only validates scales in their current storage. requires_device_loading: bool = False + # True when lookups return activation-dtype rows (per-row/group scales are + # applied inside the lookup); dequantize() is then the identity. + lookup_dequantizes: bool = False @staticmethod def from_quant_config( @@ -172,6 +185,8 @@ def from_quant_config( embedding_dtype: str | None = None, ) -> "Qwen4ExpPLEEmbeddingMethod": """Select the concrete PLE embedding format for a layer.""" + if envs.VLLM_PLE_QUANT_DIR: + return Qwen4ExpPLEInt4EmbeddingMethod(envs.VLLM_PLE_QUANT_DIR) if embedding_dtype == "float8_e4m3fn": return Qwen4ExpPLEFp8EmbeddingMethod() if quant_config is None: @@ -328,6 +343,121 @@ def dequantize( return embeddings.to(output_dtype) * weight_scale.to(output_dtype) +class Qwen4ExpPLEInt4EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod): + """INT4 group-16 PLE table loaded from a quantized-table directory. + + ``VLLM_PLE_QUANT_DIR`` points at a ``ples_int4`` directory (e.g. from + primitive-ai/Qwen3.8-Flash-Next-PLE-quant): ``META.json`` plus + ``shard_.safetensors`` with ``weight_i4`` (uint8, two values per byte, + low nibble first, stored as value + 8) and ``weight_scale`` (fp16, one per + 16 values). The checkpoint's own table is skipped; each rank copies only + its vocabulary rows, so the table takes about a third of its BF16 size. + """ + + LAYOUT = "group16_int4_fp16scale_lownibblefirst" + GROUP_SIZE = 16 + lookup_dequantizes = True + + def __init__(self, quant_dir: str) -> None: + with open(os.path.join(quant_dir, "META.json")) as fh: + meta = json.load(fh) + if meta.get("layout") != self.LAYOUT: + raise NotImplementedError( + f"VLLM_PLE_QUANT_DIR={quant_dir} has layout {meta.get('layout')!r}; " + f"only {self.LAYOUT!r} (ples_int4) is supported" + ) + if meta["rows"] % meta["shards"]: + raise ValueError(f"{quant_dir}: rows are not evenly split across shards") + self.quant_dir = quant_dir + self.rows = int(meta["rows"]) + self.width = int(meta["width"]) + self.rows_per_shard = self.rows // int(meta["shards"]) + + 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, extra_weight_attrs + width = input_size_per_partition + if width != self.width or width % self.GROUP_SIZE: + raise ValueError( + f"INT4 PLE table width {self.width} does not match the model " + f"embedding dim {width}" + ) + rows = sum(output_partition_sizes) + for name, cols, dtype in ( + ("weight", width // 2, torch.uint8), + ("weight_scale", width // self.GROUP_SIZE, torch.float16), + ): + param = nn.Parameter( + layer.allocate_embedding_weight(rows, cols, dtype), + requires_grad=False, + ) + set_weight_attrs(param, {"input_dim": 1, "output_dim": 0}) + layer.register_parameter(name, param) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + """Copy this rank's rows of the INT4 table into the allocated storage.""" + from safetensors import safe_open + + start = layer.shard_indices.org_vocab_start_index + end = layer.shard_indices.org_vocab_end_index + if layer.org_vocab_size != self.rows: + raise ValueError( + f"INT4 PLE table has {self.rows} rows, the model expects " + f"{layer.org_vocab_size}" + ) + layer.weight.data[end - start :].zero_() + layer.weight_scale.data[end - start :].zero_() + row = start + while row < end: + shard = row // self.rows_per_shard + first = row - shard * self.rows_per_shard + count = min(end - row, self.rows_per_shard - first) + path = os.path.join(self.quant_dir, f"shard_{shard}.safetensors") + with safe_open(path, framework="pt") as f: + for name, key in ( + ("weight", "weight_i4"), + ("weight_scale", "weight_scale"), + ): + rows = f.get_slice(key)[first : first + count] + getattr(layer, name).data[row - start : row - start + count].copy_( + rows + ) + row += count + logger.info("Loaded INT4 PLE rows [%d, %d) from %s", start, end, self.quant_dir) + + def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: + return _dequantize_int4_rows( + layer.weight[input_], layer.weight_scale[input_], layer.params_dtype + ) + + def dequantize( + self, + layer: nn.Module, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + del layer, output_dtype + return embeddings + + +def _dequantize_int4_rows( + packed: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + """Unpack low-nibble-first INT4 rows (value + 8) and apply group-16 scales.""" + values = torch.stack((packed & 0xF, packed >> 4), dim=-1).flatten(-2) + group = values.shape[-1] // scale.shape[-1] + scale = scale.repeat_interleave(group, dim=-1) + return ((values.to(torch.float32) - 8.0) * scale.to(torch.float32)).to(dtype) + + class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding): """PLE table allocated on the active model device.""" @@ -390,6 +520,45 @@ def _lookup_ple_embedding_from_pinned_kernel( ) +@triton.jit +def _lookup_ple_int4_from_pinned_kernel( + packed_ptr, + scale_ptr, + ids_ptr, + output_ptr, + embedding_dim, + tp_vocab_start, + tp_vocab_end, + GROUP: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Look up and dequantize TP-owned INT4 PLE rows through UVA.""" + 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).to(tl.int64) + offsets = tl.arange(0, BLOCK_D) + store_mask = offsets < embedding_dim + load_mask = store_mask & in_range + packed = tl.load( + packed_ptr + local_idx * (embedding_dim // 2) + offsets // 2, + mask=load_mask, + other=0, + ).to(tl.int32) + nibble = tl.where(offsets % 2 == 0, packed & 0xF, packed >> 4) + scale = tl.load( + scale_ptr + local_idx * (embedding_dim // GROUP) + offsets // GROUP, + mask=load_mask, + other=0.0, + ).to(tl.float32) + values = tl.where(load_mask, (nibble.to(tl.float32) - 8.0) * scale, 0.0) + tl.store( + output_ptr + row_id * embedding_dim + offsets, + values.to(output_ptr.dtype.element_ty), + mask=store_mask, + ) + + class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding): """PLE table loaded into pinned CPU memory and looked up through UVA.""" @@ -422,13 +591,18 @@ def __init__( data_parallel_rank=data_parallel_rank, ) self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight) + self._uva_scale = ( + get_accelerator_view_from_cpu_tensor(self.weight_scale) + if embedding_method.lookup_dequantizes + else None + ) 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, + dtype=self.lookup_dtype, device=self._uva_weight.device, ) self._output_dim = num_ngram_heads * self.embedding_dim @@ -453,17 +627,17 @@ 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 ``lookup_dtype``.""" expected_shape = (*input_ids.shape, self.embedding_dim) if output is None: output = torch.empty( expected_shape, - dtype=self.weight.dtype, + dtype=self.lookup_dtype, device=input_ids.device, ) elif ( tuple(output.shape) != expected_shape - or output.dtype != self.weight.dtype + or output.dtype != self.lookup_dtype or output.device != input_ids.device ): raise ValueError( @@ -472,7 +646,19 @@ def _lookup( ) flat_ids = input_ids.reshape(-1).long() - if flat_ids.numel(): + if flat_ids.numel() and self._uva_scale is not None: + _lookup_ple_int4_from_pinned_kernel[(flat_ids.numel(),)]( + self._uva_weight, + self._uva_scale, + flat_ids, + output, + self.embedding_dim, + self.shard_indices.org_vocab_start_index, + self.shard_indices.org_vocab_end_index, + GROUP=Qwen4ExpPLEInt4EmbeddingMethod.GROUP_SIZE, + BLOCK_D=self._block_d, + ) + elif flat_ids.numel(): _lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)]( self._uva_weight, flat_ids, diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 0b6fd650b101..5a21b2b41a6a 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py @@ -408,6 +408,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded.add(name) continue if name.startswith(shard_prefix) and name.endswith(".weight"): + method = getattr(self.ngram_embedding, "embedding_method", None) + if getattr(method, "lookup_dequantizes", False): + # The table comes from VLLM_PLE_QUANT_DIR after loading. + loaded.add("ngram_embedding.weight") + continue shard_text = name[len(shard_prefix) : -len(".weight")] if not shard_text.isdigit(): regular_weights.append((name, loaded_weight)) From 05c9ed0fd425333871afe88c31084cc48176f007 Mon Sep 17 00:00:00 2001 From: BlivionIaG Date: Sun, 11 Oct 2026 06:31:39 +0200 Subject: [PATCH 2/2] tools(rdna): serve_validate LOAD_HOOK for real-traffic workloads Optional script run against the live server after the bench cells (URL, PORT, SERVED, MODEL, OUT exported), e.g. tools/agentic_bench coding-agent traffic; its PASS/FAIL line lands in the summary. --- tools/rdna/port_v031/serve_validate.sh | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tools/rdna/port_v031/serve_validate.sh b/tools/rdna/port_v031/serve_validate.sh index e5a1f40a666f..cad73d351d1d 100755 --- a/tools/rdna/port_v031/serve_validate.sh +++ b/tools/rdna/port_v031/serve_validate.sh @@ -13,6 +13,10 @@ # CELLS=none skips the bench cells. WARMUP=0 skips the capture-ladder warm-up # (tools/rdna/port_v031/ladder_warmup.py) that runs before the stall probe. # +# LOAD_HOOK=/path/script.sh runs after the bench cells (before the after-load +# probe) with URL, PORT, SERVED, MODEL and OUT exported, e.g. a real-traffic +# workload; its output goes to $OUT/load_hook.txt and must end in PASS/FAIL. +# # STALL_PROBE=1 (default 0) runs tools/rdna/port_v031/stall_probe.py (mixed # prefill/decode stall probe, docs/rdna2/mixed-batch-stall-probe.md) after the # probes, with STALL_ARGS passed through (gates, lengths, repeats). It also @@ -190,6 +194,13 @@ for cell in $CELLS; do break fi done +if [[ -n ${LOAD_HOOK:-} ]] && kill -0 "$SPID" 2>/dev/null; then + log "load hook: $LOAD_HOOK" + URL=$URL PORT=$PORT SERVED=$SERVED MODEL=$MODEL OUT=$OUT \ + bash "$LOAD_HOOK" > "$OUT/load_hook.txt" 2>&1 + log "LOAD HOOK rc=$? $(grep -aoE "(PASS|FAIL)[^\n]*" "$OUT/load_hook.txt" | tail -1)" + grep -aE "^\[|PASS|FAIL|garbage" "$OUT/load_hook.txt" | tail -12 | sed 's/^/ /' | tee -a "$SUM" +fi # One more probe after load: corruption under load shows up here. if "$VENV/bin/python" "$TREE/tools/probe_greedy_correctness.py" --url "$URL" \ --model "$SERVED" --max-tokens 48 > "$OUT/probe-after.txt" 2>&1; then