[RDNA2] INT4 PLE table on the upstream PLE path (VLLM_PLE_QUANT_DIR)

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)
 

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

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