2872 lines
104 KiB
Python
2872 lines
104 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import logging
|
|
import mmap
|
|
import os
|
|
from dataclasses import dataclass
|
|
from itertools import accumulate
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from compressed_tensors.quantization import QuantizationArgs
|
|
from safetensors import safe_open
|
|
from safetensors.torch import load, load_file, save_file
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
from transformers import Qwen4ExpTextConfig
|
|
|
|
import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
|
|
import vllm.model_executor.parameter as parameter_module
|
|
import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module
|
|
import vllm.models.qwen4_exp.nvidia.ngram_embedding as nvidia_ngram_embedding_module
|
|
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
|
|
from vllm.config import SpeculativeConfig, VllmConfig
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.config.quantization import QuantizationConfigArgs
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
|
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
|
|
CompressedTensorsConfig,
|
|
)
|
|
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
|
from vllm.model_executor.layers.quantization.inc import INCConfig
|
|
from vllm.model_executor.layers.quantization.modelopt import (
|
|
ModelOptMixedPrecisionConfig,
|
|
ModelOptNvFp4Config,
|
|
)
|
|
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
|
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
|
|
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,
|
|
PLEVocabParallelEmbedding,
|
|
)
|
|
from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
|
|
Qwen4ExpNGramEmbedding,
|
|
Qwen4ExpPLEDeviceEmbedding,
|
|
Qwen4ExpPLEEmbeddingMethod,
|
|
Qwen4ExpPLEFileGatherEmbedding,
|
|
Qwen4ExpPLEFp8EmbeddingMethod,
|
|
Qwen4ExpPLENvFp4EmbeddingMethod,
|
|
Qwen4ExpPLEPinnedHostEmbedding,
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.torch_utils import weak_ref_tensor
|
|
from vllm.v1.attention.backends.short_conv_attn import (
|
|
PleShortConvAttentionMetadata,
|
|
PleShortConvAttentionMetadataBuilder,
|
|
)
|
|
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
|
from vllm.v1.kv_cache_interface import MambaSpec
|
|
|
|
|
|
@pytest.mark.parametrize("device", ["cpu", "cuda"])
|
|
@pytest.mark.parametrize("full_graph", [False, True])
|
|
@pytest.mark.parametrize(
|
|
"query_lens,drafts,spec_rows,non_spec_rows,token_order",
|
|
[
|
|
pytest.param([4, 2, 0], [3, 1, -1], [0, 1], [], None, id="spec-padding"),
|
|
pytest.param(
|
|
[3, 1, 2, 5, 0],
|
|
[2, -1, 1, -1, -1],
|
|
[0, 2],
|
|
[1, 3],
|
|
[0, 1, 2, 4, 5, 3, 6, 7, 8, 9, 10],
|
|
id="mixed-padding",
|
|
),
|
|
pytest.param([1, 5], [-1, -1], [], [0, 1], None, id="non-spec"),
|
|
],
|
|
)
|
|
def test_ple_metadata_preserves_request_and_token_order(
|
|
device, full_graph, query_lens, drafts, spec_rows, non_spec_rows, token_order
|
|
):
|
|
"""Preserve state routing, acceptance and prefill offsets without GPU sync."""
|
|
if device == "cuda" and not torch.cuda.is_available():
|
|
pytest.skip("CUDA required")
|
|
device = torch.device(device)
|
|
config = VllmConfig()
|
|
config.speculative_config = SpeculativeConfig(
|
|
method="ngram", num_speculative_tokens=3
|
|
)
|
|
config.compilation_config.cudagraph_mode = (
|
|
CUDAGraphMode.FULL if full_graph else CUDAGraphMode.NONE
|
|
)
|
|
config.compilation_config.max_cudagraph_capture_size = 32
|
|
builder = PleShortConvAttentionMetadataBuilder(
|
|
MambaSpec(block_size=16, shapes=((16, 64),), dtypes=(torch.float32,)),
|
|
["ple"],
|
|
config,
|
|
device,
|
|
)
|
|
common = create_common_attn_metadata(
|
|
BatchSpec(seq_lens=[q + 16 for q in query_lens], query_lens=query_lens),
|
|
16,
|
|
device,
|
|
arange_block_indices=True,
|
|
)
|
|
common.is_prefilling = torch.tensor(
|
|
[q > 1 and d < 0 for q, d in zip(query_lens, drafts)]
|
|
)
|
|
# Token padding is independent of the sum of per-request query lengths.
|
|
common.num_actual_tokens += 3
|
|
accepted = torch.arange(1, len(query_lens) + 1, dtype=torch.int32, device=device)
|
|
drafts_cpu = torch.tensor(drafts, dtype=torch.int32)
|
|
sync_mode = torch.cuda.get_sync_debug_mode() if device.type == "cuda" else None
|
|
if sync_mode is not None:
|
|
torch.cuda.set_sync_debug_mode("error")
|
|
try:
|
|
metadata = builder.build(
|
|
0,
|
|
common,
|
|
num_accepted_tokens=accepted,
|
|
num_decode_draft_tokens_cpu=drafts_cpu,
|
|
)
|
|
finally:
|
|
if sync_mode is not None:
|
|
torch.cuda.set_sync_debug_mode(sync_mode)
|
|
|
|
def check(actual, expected):
|
|
torch.testing.assert_close(
|
|
actual, torch.tensor(expected, dtype=actual.dtype, device=device)
|
|
)
|
|
|
|
assert metadata.num_spec_decodes == len(spec_rows)
|
|
assert metadata.num_decodes == sum(query_lens[i] == 1 for i in non_spec_rows)
|
|
assert metadata.num_prefills == sum(query_lens[i] > 1 for i in non_spec_rows)
|
|
if spec_rows:
|
|
torch.testing.assert_close(
|
|
metadata.spec_state_indices_tensor[: len(spec_rows)],
|
|
common.block_table_tensor[spec_rows, 0],
|
|
)
|
|
torch.testing.assert_close(
|
|
metadata.num_accepted_tokens[: len(spec_rows)], accepted[spec_rows]
|
|
)
|
|
check(
|
|
metadata.spec_query_start_loc[: len(spec_rows) + 1],
|
|
[0, *accumulate(query_lens[i] for i in spec_rows)],
|
|
)
|
|
if full_graph and not non_spec_rows:
|
|
check(metadata.spec_state_indices_tensor[len(spec_rows) :], [NULL_BLOCK_ID])
|
|
check(metadata.num_accepted_tokens[len(spec_rows) :], [1])
|
|
check(
|
|
metadata.spec_query_start_loc[len(spec_rows) + 1 :], [sum(query_lens)]
|
|
)
|
|
assert (
|
|
metadata.spec_state_indices_tensor.data_ptr()
|
|
== builder.spec_state_indices_tensor.data_ptr()
|
|
)
|
|
if non_spec_rows:
|
|
torch.testing.assert_close(
|
|
metadata.state_indices_tensor, common.block_table_tensor[non_spec_rows, 0]
|
|
)
|
|
prefill_lens = [query_lens[i] for i in non_spec_rows if query_lens[i] > 1]
|
|
check(metadata.query_start_loc_p, [0, *accumulate(prefill_lens)])
|
|
check(metadata.has_initial_states_p, [True] * len(prefill_lens))
|
|
if token_order is not None:
|
|
check(
|
|
torch.cat((metadata.spec_token_indx, metadata.non_spec_token_indx)),
|
|
token_order,
|
|
)
|
|
else:
|
|
assert metadata.spec_token_indx is None
|
|
assert metadata.non_spec_token_indx is None
|
|
assert metadata.nums_dict is None
|
|
assert metadata.batch_ptr is None
|
|
assert metadata.token_chunk_offset_ptr is None
|
|
|
|
|
|
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=rank,
|
|
world_size=world_size,
|
|
all_reduce=all_reduce,
|
|
)
|
|
monkeypatch.setattr(ngram_embedding_module, "get_etp_group", lambda: group)
|
|
monkeypatch.setattr(
|
|
ngram_embedding_module,
|
|
"get_tp_group",
|
|
lambda: SimpleNamespace(world_size=world_size),
|
|
)
|
|
|
|
|
|
def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
|
|
module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
|
|
nn.Module.__init__(module)
|
|
module.split_ngram_parts = 2
|
|
module.register_buffer("layer_multipliers", torch.zeros(1, dtype=torch.long))
|
|
module.register_buffer("ngram_heads_offsets", torch.zeros(1, dtype=torch.long))
|
|
module.register_buffer("ngram_heads_vocab_sizes", torch.zeros(1, dtype=torch.long))
|
|
module.ngram_embedding = SimpleNamespace(
|
|
org_vocab_size=8,
|
|
embedding_dim=2,
|
|
weight=nn.Parameter(torch.full((4, 2), -1.0)),
|
|
shard_indices=SimpleNamespace(
|
|
org_vocab_start_index=2,
|
|
org_vocab_end_index=6,
|
|
),
|
|
)
|
|
_set_test_embedding_weight_loader(module.ngram_embedding)
|
|
return module
|
|
|
|
|
|
def _make_test_ple_embedding(tp_start: int, tp_end: int) -> PLEVocabParallelEmbedding:
|
|
"""Create an embedding with only the state needed for shard loading."""
|
|
embedding = PLEVocabParallelEmbedding.__new__(PLEVocabParallelEmbedding)
|
|
nn.Module.__init__(embedding)
|
|
embedding.shard_indices = SimpleNamespace(
|
|
org_vocab_start_index=tp_start,
|
|
org_vocab_end_index=tp_end,
|
|
)
|
|
return embedding
|
|
|
|
|
|
def _set_test_embedding_weight_loader(embedding) -> None:
|
|
loader = _make_test_ple_embedding(
|
|
tp_start=embedding.shard_indices.org_vocab_start_index,
|
|
tp_end=embedding.shard_indices.org_vocab_end_index,
|
|
)
|
|
embedding.weight.weight_loader = loader.weight_loader
|
|
|
|
|
|
def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
|
|
module = _make_ngram_embedding_for_load_test()
|
|
embedding = nn.Module()
|
|
embedding.org_vocab_size = 8
|
|
embedding.embedding_dim = 2
|
|
embedding.shard_indices = SimpleNamespace(
|
|
org_vocab_start_index=2,
|
|
org_vocab_end_index=6,
|
|
)
|
|
embedding.register_parameter(
|
|
"weight",
|
|
nn.Parameter(
|
|
torch.full((4, 2), -1.0).to(torch.float8_e4m3fn),
|
|
requires_grad=False,
|
|
),
|
|
)
|
|
embedding.register_parameter(
|
|
"weight_scale",
|
|
nn.Parameter(torch.zeros(1, dtype=torch.float32), requires_grad=False),
|
|
)
|
|
_set_test_embedding_weight_loader(embedding)
|
|
module.ngram_embedding = embedding
|
|
return module
|
|
|
|
|
|
def test_ple_shard_overlap_and_copy() -> None:
|
|
embedding = _make_test_ple_embedding(tp_start=4, tp_end=8)
|
|
overlap = embedding._compute_shard_overlap(checkpoint_start=2, checkpoint_rows=5)
|
|
assert overlap == PLEShardOverlap(source_start=2, destination_start=0, row_count=3)
|
|
|
|
destination = torch.full((4, 2), -1.0)
|
|
loaded = torch.arange(10, dtype=torch.float64).reshape(5, 2)
|
|
copied = embedding._copy_embedding_shard_(
|
|
destination,
|
|
loaded,
|
|
checkpoint_start=2,
|
|
)
|
|
|
|
assert copied == 3
|
|
torch.testing.assert_close(destination[:3], loaded[2:5].float())
|
|
torch.testing.assert_close(destination[3], torch.tensor([-1.0, -1.0]))
|
|
|
|
|
|
def test_ple_shard_copy_is_a_noop_without_overlap() -> None:
|
|
embedding = _make_test_ple_embedding(tp_start=4, tp_end=8)
|
|
destination = torch.ones(4, 2)
|
|
copied = embedding._copy_embedding_shard_(
|
|
destination,
|
|
torch.zeros(2, 2),
|
|
checkpoint_start=10,
|
|
)
|
|
|
|
assert copied == 0
|
|
assert torch.equal(destination, torch.ones_like(destination))
|
|
|
|
|
|
def test_ngram_embedding_loads_shards_and_ignores_legacy_token_lookup() -> None:
|
|
module = _make_ngram_embedding_for_load_test()
|
|
shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2)
|
|
shard_1 = torch.arange(8, 16, dtype=torch.float32).reshape(4, 2)
|
|
|
|
loaded = module.load_weights(
|
|
[
|
|
("ngram_embedding.shard_0.weight", shard_0),
|
|
("ngram_embedding.shard_1.weight", shard_1),
|
|
("token_lookup", torch.tensor([2, 1, 0])),
|
|
]
|
|
)
|
|
|
|
assert loaded == {"ngram_embedding.weight"}
|
|
torch.testing.assert_close(
|
|
module.ngram_embedding.weight,
|
|
torch.cat((shard_0[2:4], shard_1[0:2])),
|
|
)
|
|
|
|
|
|
def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None:
|
|
module = _make_fp8_ngram_embedding_for_load_test()
|
|
shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn)
|
|
shard_1 = (
|
|
torch.arange(8, 16, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn)
|
|
)
|
|
weight_scale = torch.tensor([0.25], dtype=torch.bfloat16)
|
|
|
|
loaded = module.load_weights(
|
|
[
|
|
("ngram_embedding.shard_0.weight", shard_0),
|
|
("ngram_embedding.shard_1.weight", shard_1),
|
|
("ngram_embedding.weight_scale", weight_scale),
|
|
]
|
|
)
|
|
|
|
assert loaded == {"ngram_embedding.weight", "ngram_embedding.weight_scale"}
|
|
assert module.ngram_embedding.weight.dtype == torch.float8_e4m3fn
|
|
assert torch.equal(
|
|
module.ngram_embedding.weight.float(),
|
|
torch.cat((shard_0[2:4], shard_1[0:2])).float(),
|
|
)
|
|
assert module.ngram_embedding.weight_scale.dtype == torch.float32
|
|
torch.testing.assert_close(
|
|
module.ngram_embedding.weight_scale,
|
|
weight_scale.float(),
|
|
)
|
|
|
|
|
|
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,
|
|
dp_rank: int,
|
|
local_tokens: int,
|
|
) -> None:
|
|
embedding = Qwen4ExpPLEDeviceEmbedding.__new__(Qwen4ExpPLEDeviceEmbedding)
|
|
nn.Module.__init__(embedding)
|
|
embedding.tp_size = 2
|
|
embedding.etp_data_parallel_size = 2
|
|
embedding.data_parallel_rank = dp_rank
|
|
gathered_ids = torch.tensor([[10], [11], [0], [20], [21], [22]])
|
|
group = SimpleNamespace(
|
|
rank_in_group=dp_rank,
|
|
all_gather=lambda input_ids, dim: gathered_ids,
|
|
all_reduce=lambda embeddings: embeddings,
|
|
)
|
|
embedding.parallel_group = group
|
|
forward_context = SimpleNamespace(
|
|
dp_metadata=SimpleNamespace(
|
|
num_tokens_across_dp_cpu=torch.tensor([2, 3]),
|
|
)
|
|
)
|
|
monkeypatch.setattr(ngram_embedding_module, "get_dp_group", lambda: group)
|
|
monkeypatch.setattr(
|
|
ngram_embedding_module, "get_forward_context", lambda: forward_context
|
|
)
|
|
local_ids = gathered_ids[dp_rank * 3 : dp_rank * 3 + local_tokens]
|
|
embeddings = torch.arange(12).reshape(6, 2)
|
|
|
|
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
assert torch.equal(input_ids, gathered_ids)
|
|
return embeddings
|
|
|
|
monkeypatch.setattr(
|
|
ngram_embedding_module.PLEVocabParallelEmbedding,
|
|
"forward",
|
|
forward,
|
|
)
|
|
output = embedding(local_ids)
|
|
|
|
expected = embeddings[dp_rank * 3 : dp_rank * 3 + local_tokens]
|
|
assert torch.equal(output, expected)
|
|
|
|
|
|
def _make_fp8_embedding_layer(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
*,
|
|
load_scale: bool = True,
|
|
) -> Qwen4ExpPLEDeviceEmbedding:
|
|
_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
|
|
)
|
|
method = Qwen4ExpPLEFp8EmbeddingMethod()
|
|
layer = Qwen4ExpPLEDeviceEmbedding(
|
|
3,
|
|
2,
|
|
params_dtype=torch.bfloat16,
|
|
padding_size=1,
|
|
prefix="test.ple_embedding",
|
|
embedding_method=method,
|
|
)
|
|
weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]])
|
|
layer.weight.data.copy_(weight.to(torch.float8_e4m3fn))
|
|
if load_scale:
|
|
layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=torch.bfloat16))
|
|
return layer
|
|
|
|
|
|
def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None:
|
|
layer = _make_fp8_embedding_layer(monkeypatch)
|
|
layer.quant_method.process_weights_after_loading(layer)
|
|
quantized_output = layer(torch.tensor([2, 0]))
|
|
ple_layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
|
|
nn.Module.__init__(ple_layer)
|
|
ple_layer.ple_embedding = nn.Module()
|
|
ple_layer.ple_embedding.ngram_embedding = layer
|
|
|
|
output = ple_layer._dequantize_embeddings(
|
|
quantized_output,
|
|
torch.bfloat16,
|
|
)
|
|
|
|
assert layer.weight.dtype == torch.float8_e4m3fn
|
|
assert layer.weight_scale.dtype == torch.float32
|
|
assert quantized_output.dtype == torch.float8_e4m3fn
|
|
assert output.dtype == torch.bfloat16
|
|
weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]])
|
|
torch.testing.assert_close(output, (weight[[2, 0]] * 0.25).bfloat16())
|
|
|
|
|
|
def test_ple_fp8_embedding_rejects_missing_global_scale(monkeypatch) -> None:
|
|
layer = _make_fp8_embedding_layer(monkeypatch, load_scale=False)
|
|
|
|
with pytest.raises(ValueError, match="missing its global scale"):
|
|
layer.quant_method.process_weights_after_loading(layer)
|
|
|
|
|
|
@pytest.mark.parametrize("requires_device_loading", [False, True])
|
|
@pytest.mark.parametrize("load_scale", [False, True])
|
|
def test_pinned_ple_post_load_validates_scale_without_staging(
|
|
monkeypatch: pytest.MonkeyPatch, requires_device_loading: bool, load_scale: bool
|
|
) -> None:
|
|
"""Validate scales through the loader while honoring the staging policy."""
|
|
from contextlib import nullcontext
|
|
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
|
|
from vllm.model_executor.model_loader import utils as loader_utils
|
|
|
|
source = _make_fp8_embedding_layer(monkeypatch, load_scale=load_scale)
|
|
embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
|
|
nn.Module.__init__(embedding)
|
|
embedding.weight = source.weight
|
|
embedding.weight_scale = source.weight_scale
|
|
embedding.embedding_method = source.embedding_method
|
|
embedding.quant_method = source.embedding_method
|
|
embedding.tp_rank = source.tp_rank
|
|
embedding.tp_size = source.tp_size
|
|
assert QuantizeMethodBase.requires_device_loading
|
|
assert not embedding.quant_method.requires_device_loading
|
|
if requires_device_loading:
|
|
embedding.quant_method.requires_device_loading = True
|
|
model = nn.Module()
|
|
model.embedding = embedding
|
|
weight_ptr = embedding.weight.data_ptr()
|
|
staged_modules = []
|
|
|
|
def track_staging(module, target_device):
|
|
staged_modules.append(module)
|
|
return nullcontext()
|
|
|
|
monkeypatch.setattr(loader_utils, "device_loading_context", track_staging)
|
|
monkeypatch.setattr(loader_utils, "maybe_retie_word_embeddings", lambda *args: None)
|
|
monkeypatch.setattr(
|
|
loader_utils, "release_device_memory_under_pressure", lambda *args: None
|
|
)
|
|
config = SimpleNamespace(quantization="fp8")
|
|
if load_scale:
|
|
loader_utils.process_weights_after_loading(model, config, torch.device("cuda"))
|
|
else:
|
|
with pytest.raises(ValueError, match="missing its global scale"):
|
|
loader_utils.process_weights_after_loading(
|
|
model, config, torch.device("cuda")
|
|
)
|
|
assert embedding.weight.device.type == "cpu"
|
|
assert embedding.weight.data_ptr() == weight_ptr
|
|
assert staged_modules == ([embedding] if requires_device_loading else [])
|
|
|
|
|
|
def test_ple_fp8_embedding_uses_int8_for_parallel_reduce(monkeypatch) -> None:
|
|
monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(
|
|
embedding_module, "get_tensor_model_parallel_world_size", lambda: 2
|
|
)
|
|
monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(
|
|
parameter_module, "get_tensor_model_parallel_world_size", lambda: 2
|
|
)
|
|
monkeypatch.setattr(
|
|
embedding_module,
|
|
"get_masked_input_and_mask",
|
|
lambda *args: (
|
|
torch.tensor([0, 0]),
|
|
torch.tensor([False, True]),
|
|
),
|
|
)
|
|
reduced_dtypes = []
|
|
|
|
def all_reduce(tensor: torch.Tensor) -> torch.Tensor:
|
|
reduced_dtypes.append(tensor.dtype)
|
|
return tensor.clone()
|
|
|
|
_mock_etp_group(monkeypatch, world_size=2, all_reduce=all_reduce)
|
|
layer = Qwen4ExpPLEDeviceEmbedding(
|
|
4,
|
|
2,
|
|
params_dtype=torch.bfloat16,
|
|
padding_size=1,
|
|
prefix="test.ple_embedding",
|
|
embedding_method=Qwen4ExpPLEFp8EmbeddingMethod(),
|
|
)
|
|
layer.weight.data.copy_(
|
|
torch.tensor([[1.0, 2.0], [4.0, 8.0]]).to(torch.float8_e4m3fn)
|
|
)
|
|
|
|
output = layer(torch.tensor([0, 2]))
|
|
|
|
assert reduced_dtypes == [torch.int8]
|
|
assert output.dtype == torch.float8_e4m3fn
|
|
torch.testing.assert_close(output[0].float(), layer.weight[0].float())
|
|
assert torch.count_nonzero(output[1].float()) == 0
|
|
|
|
|
|
def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None:
|
|
prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
|
|
quant_config = Fp8Config(
|
|
ignored_layers=[],
|
|
weight_block_size=[128, 128],
|
|
)
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEFp8EmbeddingMethod,
|
|
)
|
|
|
|
quant_config.ignored_layers = [f"{prefix}.shard_0"]
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(None, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"quant_config",
|
|
[
|
|
pytest.param(ModelOptNvFp4Config(exclude_modules=[]), id="modelopt-nvfp4"),
|
|
pytest.param(
|
|
OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_tensor")),
|
|
id="online-fp8",
|
|
),
|
|
pytest.param(
|
|
CompressedTensorsConfig(
|
|
target_scheme_map={
|
|
"re:.*ngram_embedding": {
|
|
"weights": QuantizationArgs(num_bits=8, type="int"),
|
|
},
|
|
},
|
|
ignore=[],
|
|
quant_format="int-quantized",
|
|
),
|
|
id="compressed-tensors",
|
|
),
|
|
],
|
|
)
|
|
def test_ple_embedding_falls_back_for_unsupported_quantization_configs(
|
|
quant_config: QuantizationConfig,
|
|
caplog: pytest.LogCaptureFixture,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Warn when global quantization falls back to an unquantized PLE table."""
|
|
prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
|
|
test_logger = logging.getLogger("test.ple_fallback")
|
|
monkeypatch.setattr(ngram_embedding_module, "logger", test_logger)
|
|
caplog.set_level(logging.WARNING, logger=test_logger.name)
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
assert caplog.records[0].levelno == logging.WARNING
|
|
assert "Falling back to unquantized PLE" in caplog.text
|
|
|
|
|
|
def test_ple_embedding_respects_modelopt_exclusion() -> None:
|
|
prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
|
|
quant_config = ModelOptNvFp4Config(exclude_modules=[prefix])
|
|
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
|
|
|
|
def test_ple_embedding_respects_inc_layer_config() -> None:
|
|
prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
|
|
quant_config = INCConfig(
|
|
weight_bits=4,
|
|
group_size=128,
|
|
block_name_to_quantize="model.layers",
|
|
extra_config={".*ple.*": {"bits": 16, "data_type": "float"}},
|
|
)
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
|
|
quant_config.extra_config = None
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
|
|
|
|
def test_ple_embedding_is_unquantized_under_quark() -> None:
|
|
prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
|
|
quant_config = QuarkConfig({"exclude": [], "global_quant_config": {}})
|
|
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
|
|
|
|
def test_ple_embedding_dtype_overrides_modelopt_exclusion() -> None:
|
|
prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
|
|
quant_config = ModelOptNvFp4Config(exclude_modules=[prefix])
|
|
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(
|
|
quant_config,
|
|
prefix,
|
|
"float8_e4m3fn",
|
|
),
|
|
Qwen4ExpPLEFp8EmbeddingMethod,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
|
|
@pytest.mark.parametrize("reduce_results", [False, True])
|
|
def test_pinned_embedding_finalize_prefetch_optionally_reduces_output(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
dtype: torch.dtype,
|
|
reduce_results: bool,
|
|
) -> None:
|
|
"""Deferring reduction must preserve full token rows and wait for the lookup."""
|
|
total_tokens = 3
|
|
embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
|
|
nn.Module.__init__(embedding)
|
|
values = torch.arange(1, 13).reshape(4, 3)
|
|
embedding._prefetch_buffer = torch.stack(
|
|
(values, torch.zeros_like(values)), dim=1
|
|
).to(dtype)
|
|
embedding._output_dim = 6
|
|
embedding.tp_size = 2
|
|
embedding.reduce_results = reduce_results
|
|
embedding.etp_data_parallel_size = 1
|
|
embedding._prefetch_stream = object()
|
|
waited_streams: list[object] = []
|
|
reduced_dtypes = []
|
|
monkeypatch.setattr(
|
|
torch.cuda,
|
|
"current_stream",
|
|
lambda: SimpleNamespace(wait_stream=waited_streams.append),
|
|
)
|
|
|
|
def all_reduce(tensor: torch.Tensor) -> torch.Tensor:
|
|
assert waited_streams == [embedding._prefetch_stream]
|
|
reduced_dtypes.append(tensor.dtype)
|
|
# The remote rank owns the other head's vocabulary rows.
|
|
return tensor + tensor.flip(1)
|
|
|
|
embedding.parallel_group = SimpleNamespace(all_reduce=all_reduce)
|
|
output = embedding.finalize_prefetch(total_tokens=total_tokens)
|
|
|
|
expected = (
|
|
values.unsqueeze(1).expand(-1, 2, -1).to(dtype)
|
|
if reduce_results
|
|
else embedding._prefetch_buffer
|
|
)
|
|
expected = expected[:total_tokens].flatten(-2)
|
|
comm_dtype = torch.int8 if dtype == torch.float8_e4m3fn else dtype
|
|
assert reduced_dtypes == ([comm_dtype] if reduce_results else [])
|
|
assert waited_streams == [embedding._prefetch_stream]
|
|
assert output.dtype == dtype
|
|
torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
|
|
|
|
|
|
def test_pinned_embedding_finalize_prefetch_requires_prior_start_prefetch() -> None:
|
|
"""Finalizing before start_prefetch must fail instead of reading garbage."""
|
|
embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
|
|
nn.Module.__init__(embedding)
|
|
embedding._prefetch_buffer = None
|
|
embedding._prefetch_stream = None
|
|
embedding._output_dim = 6
|
|
with pytest.raises(RuntimeError, match="prior start_prefetch"):
|
|
embedding.finalize_prefetch(total_tokens=2)
|
|
|
|
|
|
def test_pinned_embedding_finalize_requires_prior_start_prefetch() -> None:
|
|
"""_finalize_prefetch before any start_prefetch must fail loudly."""
|
|
embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
|
|
nn.Module.__init__(embedding)
|
|
embedding._prefetch_buffer = None
|
|
embedding._prefetch_stream = None
|
|
embedding.tp_size = 1
|
|
output = torch.zeros(2, 6, dtype=torch.bfloat16)
|
|
|
|
with pytest.raises(RuntimeError, match="prior start_prefetch"):
|
|
embedding._finalize_prefetch(torch.empty(4, 2, 3), output)
|
|
|
|
|
|
def test_pinned_fp8_embedding_uses_int8_for_parallel_reduce() -> None:
|
|
embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)
|
|
nn.Module.__init__(embedding)
|
|
embedding.tp_size = 2
|
|
reduced_dtypes = []
|
|
|
|
def all_reduce(tensor: torch.Tensor) -> torch.Tensor:
|
|
reduced_dtypes.append(tensor.dtype)
|
|
return tensor.clone()
|
|
|
|
embedding.parallel_group = SimpleNamespace(all_reduce=all_reduce)
|
|
embeddings = torch.arange(8).reshape(2, 4).to(torch.float8_e4m3fn)
|
|
|
|
output = embedding._reduce_etp_embeddings(embeddings)
|
|
|
|
assert reduced_dtypes == [torch.int8]
|
|
assert output.dtype == embeddings.dtype
|
|
torch.testing.assert_close(output.float(), embeddings.float())
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
|
def test_ple_device_embedding_allocates_on_active_device(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
_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
|
|
)
|
|
embedding_method = Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
|
|
|
with torch.device("cuda:0"):
|
|
embedding = Qwen4ExpPLEDeviceEmbedding(
|
|
4,
|
|
3,
|
|
params_dtype=torch.bfloat16,
|
|
padding_size=1,
|
|
prefix="test.ple_embedding",
|
|
embedding_method=embedding_method,
|
|
)
|
|
|
|
assert embedding.weight.device == torch.device("cuda:0")
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
|
@pytest.mark.parametrize("fp8_checkpoint", [False, True])
|
|
def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
fp8_checkpoint: bool,
|
|
) -> None:
|
|
_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
|
|
)
|
|
|
|
with torch.device("cuda:0"):
|
|
embedding_method = (
|
|
Qwen4ExpPLEFp8EmbeddingMethod()
|
|
if fp8_checkpoint
|
|
else Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
|
)
|
|
embedding = Qwen4ExpPLEPinnedHostEmbedding(
|
|
4,
|
|
3,
|
|
params_dtype=torch.bfloat16,
|
|
padding_size=1,
|
|
prefix="test.ple_embedding",
|
|
embedding_method=embedding_method,
|
|
)
|
|
|
|
storage_dtype = torch.float8_e4m3fn if fp8_checkpoint else torch.bfloat16
|
|
loaded_weight = (
|
|
torch.arange(12, dtype=torch.float32).reshape(4, 3).to(storage_dtype)
|
|
)
|
|
copied = embedding._copy_embedding_shard_(
|
|
embedding.weight,
|
|
loaded_weight,
|
|
checkpoint_start=0,
|
|
)
|
|
input_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")
|
|
output = embedding._lookup(input_ids)
|
|
if fp8_checkpoint:
|
|
embedding.weight_scale.data.fill_(0.25)
|
|
weight_ptr = embedding.weight.data_ptr()
|
|
embedding.quant_method.process_weights_after_loading(embedding)
|
|
assert embedding.weight.data_ptr() == weight_ptr
|
|
dequantized = embedding.dequantize(output, torch.bfloat16)
|
|
|
|
assert copied == 4
|
|
assert embedding.weight.device.type == "cpu"
|
|
assert embedding.weight.is_pinned()
|
|
assert output.dtype == storage_dtype
|
|
expected = loaded_weight[input_ids.cpu()].to(device="cuda:0")
|
|
torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
|
|
expected_scale = 0.25 if fp8_checkpoint else 1.0
|
|
torch.testing.assert_close(
|
|
dequantized,
|
|
expected.to(torch.bfloat16) * expected_scale,
|
|
rtol=0,
|
|
atol=0,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
|
def test_pinned_embedding_start_prefetch_allocates_lazily(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Allocate prefetch storage on first use and finalize the lookup results."""
|
|
_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
|
|
)
|
|
|
|
with torch.device("cuda:0"):
|
|
embedding = Qwen4ExpPLEPinnedHostEmbedding(
|
|
4,
|
|
3,
|
|
params_dtype=torch.bfloat16,
|
|
padding_size=1,
|
|
prefix="test.ple_embedding",
|
|
embedding_method=Qwen4ExpPLEUnquantizedEmbeddingMethod(),
|
|
num_ngram_heads=2,
|
|
max_total_tokens=4,
|
|
)
|
|
assert embedding._prefetch_buffer is None
|
|
assert embedding._prefetch_stream is None
|
|
|
|
loaded_weight = torch.arange(12, dtype=torch.float32).reshape(4, 3)
|
|
loaded_weight = loaded_weight.to(torch.bfloat16)
|
|
embedding._copy_embedding_shard_(
|
|
embedding.weight,
|
|
loaded_weight,
|
|
checkpoint_start=0,
|
|
)
|
|
|
|
ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")
|
|
embedding.start_prefetch(ngram_ids)
|
|
|
|
buffer = embedding._prefetch_buffer
|
|
assert buffer is not None
|
|
assert buffer.shape == (4, 2, 3)
|
|
assert buffer.dtype == torch.bfloat16
|
|
assert buffer.device.type == "cuda"
|
|
assert embedding._prefetch_stream is not None
|
|
|
|
output = embedding.finalize_prefetch(total_tokens=ngram_ids.shape[0])
|
|
|
|
expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2)
|
|
torch.testing.assert_close(output, expected, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"exclude_pattern",
|
|
[
|
|
"model.language_model.layers.1.ple.ple_embedding.ngram_embedding",
|
|
"*.ngram_embedding",
|
|
],
|
|
)
|
|
def test_ple_fp8_embedding_supports_mixed_precision_config(
|
|
exclude_pattern: str,
|
|
) -> None:
|
|
prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"
|
|
nvfp4_prefix = "model.language_model.layers.2.ple.ple_embedding.ngram_embedding"
|
|
quant_config = ModelOptMixedPrecisionConfig.from_config(
|
|
{
|
|
"quantization": {
|
|
"quant_algo": "MIXED_PRECISION",
|
|
"exclude_modules": [],
|
|
"group_size": 16,
|
|
"quantized_layers": {
|
|
prefix: {"quant_algo": "FP8"},
|
|
nvfp4_prefix: {"quant_algo": "NVFP4"},
|
|
"model.language_model.layers.2.moe.gate_proj": {
|
|
"quant_algo": "MXFP8"
|
|
},
|
|
},
|
|
}
|
|
}
|
|
)
|
|
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEFp8EmbeddingMethod,
|
|
)
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(
|
|
quant_config,
|
|
"model.language_model.layers.3.ple.ple_embedding.ngram_embedding",
|
|
),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, nvfp4_prefix),
|
|
Qwen4ExpPLENvFp4EmbeddingMethod,
|
|
)
|
|
# An explicitly unsupported PLE format 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",
|
|
)
|
|
|
|
quant_config.exclude_modules = [exclude_pattern]
|
|
assert isinstance(
|
|
Qwen4ExpPLEEmbeddingMethod.from_quant_config(quant_config, prefix),
|
|
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
|
)
|
|
|
|
|
|
def _patch_ple_construction(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
*,
|
|
engram_config: SimpleNamespace | None = None,
|
|
load_format: str = "auto",
|
|
) -> None:
|
|
_mock_etp_group(monkeypatch)
|
|
for module in (embedding_module, parameter_module):
|
|
monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
vllm_config = SimpleNamespace(
|
|
load_config=SimpleNamespace(load_format=load_format),
|
|
engram_config=engram_config,
|
|
)
|
|
monkeypatch.setattr(
|
|
nvidia_ngram_embedding_module, "get_current_vllm_config", lambda: vllm_config
|
|
)
|
|
|
|
|
|
def _make_file_gather_embedding(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
*,
|
|
fp8: bool = False,
|
|
load_format: str = "auto",
|
|
bind: bool = False,
|
|
) -> tuple[Qwen4ExpNGramEmbedding, torch.Tensor]:
|
|
"""Build a two-head, 12-token host-file-gather embedding and its saved table.
|
|
|
|
The 21-row table is saved to ``tmp_path / "model.safetensors"`` as four
|
|
checkpoint shards of stride 6, so the last shard holds 3 rows.
|
|
"""
|
|
_patch_ple_construction(monkeypatch, load_format=load_format)
|
|
method = (
|
|
Qwen4ExpPLEFp8EmbeddingMethod()
|
|
if fp8
|
|
else Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
|
)
|
|
module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
|
|
nn.Module.__init__(module)
|
|
module.use_sequence_parallel = False
|
|
module.split_ngram_parts = 4
|
|
module.eos_token_id = 0
|
|
module.ngram_size = 2
|
|
module.heads_per_ngram = 2
|
|
module.register_buffer("layer_multipliers", torch.tensor([3, 5]))
|
|
module.register_buffer("ngram_heads_vocab_sizes", torch.tensor([13, 7]))
|
|
module.register_buffer("ngram_heads_offsets", torch.tensor([0, 13]))
|
|
module.ngram_embedding = Qwen4ExpPLEFileGatherEmbedding(
|
|
21,
|
|
2,
|
|
params_dtype=torch.bfloat16,
|
|
padding_size=1,
|
|
prefix="test.ple_embedding",
|
|
embedding_method=method,
|
|
num_ngram_heads=2,
|
|
max_total_tokens=12,
|
|
)
|
|
dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16
|
|
table = torch.arange(42, dtype=torch.float32).reshape(21, 2).to(dtype)
|
|
checkpoint = {
|
|
f"ngram_embedding.shard_{index}.weight": shard.clone()
|
|
for index, shard in enumerate(table.split(6))
|
|
}
|
|
if fp8:
|
|
checkpoint["ngram_embedding.weight_scale"] = torch.tensor([0.25]).bfloat16()
|
|
save_file(checkpoint, tmp_path / "model.safetensors")
|
|
if bind:
|
|
_load_file_shards(module, tmp_path / "model.safetensors")
|
|
module.ngram_embedding.bind_file_shards()
|
|
return module, table
|
|
|
|
|
|
def _load_file_shards(module: Qwen4ExpNGramEmbedding, path: Path) -> set[str]:
|
|
"""Load ``path`` through ``load_weights`` as the lazy safetensors iterator does."""
|
|
with safe_open(path, framework="pt") as checkpoint:
|
|
return module.load_weights(
|
|
(name, checkpoint.get_tensor(name))
|
|
for name in checkpoint.keys() # noqa: SIM118
|
|
)
|
|
|
|
|
|
def _dev_zero_shards() -> list[tuple[str, torch.Tensor]]:
|
|
"""Return shards in a ``/dev/zero`` mapping, which is how pinned memory maps."""
|
|
with open("/dev/zero", "r+b") as dev_zero:
|
|
mapping = mmap.mmap(dev_zero.fileno(), mmap.PAGESIZE)
|
|
table = torch.frombuffer(mapping, dtype=torch.uint8)[:84].view(torch.bfloat16)
|
|
return [
|
|
(f"ngram_embedding.shard_{index}.weight", shard)
|
|
for index, shard in enumerate(table.reshape(21, 2).split(6))
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("engram_config", "backend", "rows"),
|
|
[
|
|
(
|
|
SimpleNamespace(host_file_gather=True, cpu_offload=True),
|
|
Qwen4ExpPLEFileGatherEmbedding,
|
|
0,
|
|
),
|
|
(
|
|
SimpleNamespace(host_file_gather=False, cpu_offload=False),
|
|
Qwen4ExpPLEDeviceEmbedding,
|
|
24,
|
|
),
|
|
(None, Qwen4ExpPLEDeviceEmbedding, 24),
|
|
],
|
|
)
|
|
def test_ngram_embedding_selects_backend_from_engram_config(
|
|
monkeypatch, engram_config, backend, rows
|
|
) -> None:
|
|
"""File gather wins with no rows; ``cpu_offload=False`` keeps rows on device."""
|
|
_patch_ple_construction(monkeypatch, engram_config=engram_config)
|
|
config = SimpleNamespace(
|
|
ngram_size=2,
|
|
heads_per_ngram=2,
|
|
eos_token_id=0,
|
|
vocab_size=64,
|
|
ngram_vocab_size_base=8,
|
|
make_ngram_vocab_size_divisible_by=1,
|
|
)
|
|
|
|
module = Qwen4ExpNGramEmbedding(config, 4, 0, 4, data_parallel_rank=0, prefix="t")
|
|
|
|
assert tuple(module.ngram_embedding.weight.shape) == (rows, 2)
|
|
assert isinstance(module.ngram_embedding, backend)
|
|
|
|
|
|
def test_ngram_embedding_default_cpu_offload_selects_pinned_backend(
|
|
monkeypatch,
|
|
) -> None:
|
|
"""Default ``cpu_offload=True`` selects pinned storage and requires UVA."""
|
|
_patch_ple_construction(
|
|
monkeypatch,
|
|
engram_config=SimpleNamespace(host_file_gather=False, cpu_offload=True),
|
|
)
|
|
monkeypatch.setattr(ngram_embedding_module, "is_uva_available", lambda: False)
|
|
config = SimpleNamespace(
|
|
ngram_size=2,
|
|
heads_per_ngram=2,
|
|
eos_token_id=0,
|
|
vocab_size=64,
|
|
ngram_vocab_size_base=8,
|
|
make_ngram_vocab_size_divisible_by=1,
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="CPU offload requires UVA"):
|
|
Qwen4ExpNGramEmbedding(config, 4, 0, 4, data_parallel_rank=0, prefix="t")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("read_shards", "error"),
|
|
[
|
|
pytest.param(
|
|
lambda path: [("ngram_embedding.shard_0.weight", torch.zeros(6, 2))],
|
|
"dtype torch.float32 must match the embedding dtype torch.bfloat16",
|
|
id="wrong_dtype",
|
|
),
|
|
pytest.param(
|
|
lambda path: [*load_file(path).items()] * 2,
|
|
"Duplicate PLE embedding shard 0",
|
|
id="duplicate_shard",
|
|
),
|
|
pytest.param(
|
|
lambda path: [s for s in load_file(path).items() if "_2." not in s[0]],
|
|
"PLE shards cover 15 of 21 rows",
|
|
id="missing_shard",
|
|
),
|
|
pytest.param(lambda path: [], "PLE shards cover 0 of 21 rows", id="no_shards"),
|
|
pytest.param(
|
|
lambda path: [(n, t.clone()) for n, t in load_file(path).items()],
|
|
"not a file-backed view",
|
|
id="anonymous",
|
|
),
|
|
pytest.param(
|
|
lambda path: load(path.read_bytes()).items(),
|
|
"not a file-backed view",
|
|
id="heap",
|
|
),
|
|
pytest.param(
|
|
lambda path: _dev_zero_shards(), "not a file-backed view", id="dev_zero"
|
|
),
|
|
],
|
|
)
|
|
def test_file_gather_rejects_shards_that_cannot_serve_every_row_from_a_file(
|
|
monkeypatch, tmp_path, read_shards, error
|
|
) -> None:
|
|
"""A real load must fail instead of serving zeros or a copied table."""
|
|
module, _ = _make_file_gather_embedding(monkeypatch, tmp_path)
|
|
|
|
with pytest.raises(ValueError, match=error):
|
|
module.load_weights(read_shards(tmp_path / "model.safetensors"))
|
|
module.ngram_embedding.bind_file_shards()
|
|
|
|
|
|
@pytest.mark.parametrize("fp8", [False, True])
|
|
def test_file_gather_binds_zero_copy_shards_through_loader_hook(
|
|
monkeypatch, tmp_path, fp8
|
|
) -> None:
|
|
"""Fails if safetensors stops returning zero-copy views of the file."""
|
|
module, table = _make_file_gather_embedding(monkeypatch, tmp_path, fp8=fp8)
|
|
embedding = module.ngram_embedding
|
|
|
|
loaded = _load_file_shards(module, tmp_path / "model.safetensors")
|
|
embedding.quant_method.process_weights_after_loading(embedding)
|
|
embedding.quant_method.process_weights_after_loading(embedding)
|
|
|
|
assert loaded == {"ngram_embedding.weight"} | (
|
|
{"ngram_embedding.weight_scale"} if fp8 else set()
|
|
)
|
|
assert embedding._bound
|
|
stored = torch.cat([embedding._shards[index] for index in range(4)])
|
|
assert torch.equal(stored, table.view(torch.uint8))
|
|
|
|
|
|
def test_file_gather_binds_shards_opened_through_a_symlink(
|
|
monkeypatch, tmp_path
|
|
) -> None:
|
|
"""HF cache snapshots link to suffix-less blobs, which is the mapped path."""
|
|
module, _ = _make_file_gather_embedding(monkeypatch, tmp_path)
|
|
blob = (tmp_path / "model.safetensors").rename(tmp_path / "blob")
|
|
snapshot = tmp_path / "snapshot.safetensors"
|
|
snapshot.symlink_to(blob)
|
|
_load_file_shards(module, snapshot)
|
|
|
|
module.ngram_embedding.bind_file_shards()
|
|
|
|
assert module.ngram_embedding._bound
|
|
|
|
|
|
def test_file_gather_rejects_reload_and_stays_poisoned(monkeypatch, tmp_path) -> None:
|
|
module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True)
|
|
|
|
with pytest.raises(RuntimeError, match="does not support weight reload"):
|
|
module.load_weights([("ngram_embedding.shard_0.weight", table[:6])])
|
|
with pytest.raises(RuntimeError, match="unbound or was reloaded"):
|
|
module.ngram_embedding.stage_rows(1)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("bad_id", "error"),
|
|
[(21, IndexError), (24, IndexError), (-1, RuntimeError)],
|
|
ids=["past_last_shard_rows", "past_last_shard", "negative"],
|
|
)
|
|
def test_file_gather_out_of_range_id_keeps_previous_staged_rows(
|
|
monkeypatch, tmp_path, bad_id, error
|
|
) -> None:
|
|
"""Out-of-range ids fail before the H2D copy, so no wrong row is served."""
|
|
module, table = _make_file_gather_embedding(monkeypatch, tmp_path, bind=True)
|
|
embedding = module.ngram_embedding
|
|
ids = torch.tensor(
|
|
[
|
|
[20, 0],
|
|
[19, 1],
|
|
[18, 2],
|
|
[17, 3],
|
|
[16, 4],
|
|
[15, 5],
|
|
[14, 6],
|
|
[13, 7],
|
|
[12, 8],
|
|
[11, 9],
|
|
[10, 10],
|
|
]
|
|
)
|
|
embedding._host_ids[:11] = ids
|
|
embedding.stage_rows(11)
|
|
embedding._host_ids[5, 1] = bad_id
|
|
|
|
with pytest.raises(error):
|
|
embedding.stage_rows(11)
|
|
|
|
assert torch.equal(embedding._staging[:11], table[ids])
|
|
|
|
|
|
def test_file_gather_reads_ahead_the_coalesced_pages_of_staged_rows(
|
|
monkeypatch, tmp_path
|
|
) -> None:
|
|
"""Rows start 34 bytes before page 2 of a file mapped from page 1.
|
|
|
|
Rows 0-7 are in page 1, row 8 straddles pages 1 and 2, and the rows of shard 1
|
|
arrive unsorted.
|
|
"""
|
|
module, table = _make_file_gather_embedding(monkeypatch, tmp_path)
|
|
embedding, page = module.ngram_embedding, mmap.PAGESIZE
|
|
start = 2 * page - 34
|
|
path = tmp_path / "table.bin"
|
|
path.write_bytes(bytes(start) + table.view(torch.uint8).numpy().tobytes())
|
|
with open(path, "rb") as file:
|
|
mapping = mmap.mmap(file.fileno(), 0, offset=page, access=mmap.ACCESS_COPY)
|
|
rows = torch.frombuffer(mapping, dtype=torch.uint8)[start - page :]
|
|
module.load_weights(
|
|
(f"ngram_embedding.shard_{index}.weight", shard)
|
|
for index, shard in enumerate(rows.view(table.dtype).reshape(21, 2).split(6))
|
|
)
|
|
embedding.bind_file_shards()
|
|
calls = []
|
|
monkeypatch.setattr(
|
|
nvidia_ngram_embedding_module.os,
|
|
"posix_fadvise",
|
|
lambda fd, *args: calls.append((os.fstat(fd).st_ino, *args)),
|
|
)
|
|
ids = torch.tensor([[0, 1], [8, 7]])
|
|
embedding._host_ids[:2] = ids
|
|
|
|
embedding.stage_rows(2)
|
|
|
|
inode, willneed = path.stat().st_ino, os.POSIX_FADV_WILLNEED
|
|
assert calls == [(inode, page, page, willneed), (inode, page, 2 * page, willneed)]
|
|
assert torch.equal(embedding._staging[:2], table[ids])
|
|
|
|
|
|
def test_file_gather_dummy_load_stages_zeros(monkeypatch, tmp_path) -> None:
|
|
module, _ = _make_file_gather_embedding(monkeypatch, tmp_path, load_format="dummy")
|
|
embedding = module.ngram_embedding
|
|
embedding.quant_method.process_weights_after_loading(embedding)
|
|
embedding._host_rows.view(torch.uint8).fill_(0xFF)
|
|
embedding._staging.view(torch.uint8).fill_(0xFF)
|
|
|
|
embedding.stage_rows(4)
|
|
|
|
staged = embedding._staging.view(torch.uint8)
|
|
assert torch.all(staged[:4] == 0)
|
|
assert torch.all(staged[4:] == 0xFF)
|
|
|
|
|
|
def test_file_gather_forward_reads_staging_without_host_work(
|
|
monkeypatch, tmp_path
|
|
) -> None:
|
|
"""The forward must be a fixed-address read so FULL CUDA graphs replay it."""
|
|
module, _ = _make_file_gather_embedding(monkeypatch, tmp_path)
|
|
staging = module.ngram_embedding._staging
|
|
staging.view(torch.uint8).fill_(0xFF)
|
|
step = (torch.arange(6), torch.tensor([0, 6]), torch.tensor([[1]]))
|
|
hidden_states = torch.zeros(6, 8, dtype=torch.bfloat16)
|
|
|
|
assert module.start_prefetch(*step) is None
|
|
output = module(hidden_states, *step)
|
|
|
|
assert torch.all(staging.view(torch.uint8) == 0xFF)
|
|
assert output.data_ptr() == staging.data_ptr()
|
|
assert output.shape == (6, 4)
|
|
|
|
|
|
def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
|
|
conv_state_len = 6
|
|
dilation = 2
|
|
|
|
conv_weights = torch.tensor([[0.25, -0.5, 0.75, 1.0]])
|
|
conv_state = torch.zeros(2, 1, 9)
|
|
conv_state[1] = torch.arange(1, 10, dtype=torch.float32).reshape(1, 9)
|
|
first_inputs = torch.tensor([[10.0], [20.0], [30.0], [40.0]])
|
|
initial_state = conv_state[1:].clone()
|
|
first_history = torch.cat(
|
|
(initial_state[..., :conv_state_len], first_inputs.T.unsqueeze(0)),
|
|
dim=-1,
|
|
)
|
|
|
|
graph_padded_inputs = F.pad(first_inputs, (0, 0, 0, 4))
|
|
first_output = _short_conv_dilated_spec_pytorch(
|
|
graph_padded_inputs,
|
|
conv_state,
|
|
conv_weights,
|
|
torch.tensor([1, 0]),
|
|
torch.tensor([0, 4, 4]),
|
|
torch.tensor([1, 0]),
|
|
spec_query_len=4,
|
|
conv_state_len=conv_state_len,
|
|
dilation=dilation,
|
|
)
|
|
|
|
expected_first_output = F.silu(
|
|
F.conv1d(
|
|
first_history,
|
|
conv_weights.unsqueeze(1),
|
|
groups=1,
|
|
dilation=dilation,
|
|
)
|
|
).transpose(1, 2)[0]
|
|
expected_first_state = first_history[..., 1:10]
|
|
torch.testing.assert_close(first_output[:4], expected_first_output)
|
|
assert torch.count_nonzero(first_output[4:]) == 0
|
|
assert torch.count_nonzero(conv_state[0]) == 0
|
|
torch.testing.assert_close(conv_state[1:], expected_first_state)
|
|
|
|
second_inputs = torch.tensor([[50.0], [60.0]])
|
|
rollback_state = expected_first_state[..., 1:7]
|
|
padded_second_inputs = F.pad(second_inputs.T.unsqueeze(0), (0, 2))
|
|
second_history = torch.cat((rollback_state, padded_second_inputs), dim=-1)
|
|
expected_second_state = expected_first_state.clone()
|
|
expected_second_state[..., :7] = second_history[..., 1:8]
|
|
|
|
second_output = _short_conv_dilated_spec_pytorch(
|
|
second_inputs,
|
|
conv_state,
|
|
conv_weights,
|
|
torch.tensor([1]),
|
|
torch.tensor([0, 2]),
|
|
torch.tensor([2]),
|
|
spec_query_len=4,
|
|
conv_state_len=conv_state_len,
|
|
dilation=dilation,
|
|
)
|
|
|
|
expected_second_output = F.silu(
|
|
F.conv1d(
|
|
second_history,
|
|
conv_weights.unsqueeze(1),
|
|
groups=1,
|
|
dilation=dilation,
|
|
)
|
|
).transpose(1, 2)[0, :2]
|
|
torch.testing.assert_close(second_output, expected_second_output)
|
|
torch.testing.assert_close(conv_state[1:], expected_second_state)
|
|
|
|
|
|
def _reference_ngram_ids(
|
|
input_ids: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
ngram_context: torch.Tensor,
|
|
layer_multipliers: torch.Tensor,
|
|
ngram_heads_vocab_sizes: torch.Tensor,
|
|
ngram_heads_offsets: torch.Tensor,
|
|
eos_token_id: int,
|
|
heads_per_ngram: int,
|
|
) -> torch.Tensor:
|
|
tokens = input_ids.tolist()
|
|
starts = query_start_loc.tolist()
|
|
contexts = ngram_context.tolist()
|
|
multipliers = layer_multipliers.cpu()
|
|
sizes = ngram_heads_vocab_sizes.cpu()
|
|
offsets = ngram_heads_offsets.cpu()
|
|
rows = []
|
|
context_len = len(contexts[0])
|
|
for req, (start, end) in enumerate(zip(starts, starts[1:])):
|
|
history = contexts[req] + tokens[start:end]
|
|
for pos in range(context_len, len(history)):
|
|
shifted = [history[pos]]
|
|
crossed_eos = False
|
|
for shift in range(1, context_len + 1):
|
|
token = eos_token_id if crossed_eos else history[pos - shift]
|
|
shifted.append(token)
|
|
crossed_eos |= token == eos_token_id
|
|
|
|
row = []
|
|
mixed = torch.tensor(shifted[0], dtype=torch.int64) * multipliers[0]
|
|
for ngram_order in range(2, context_len + 2):
|
|
shift = ngram_order - 1
|
|
mixed ^= shifted[shift] * multipliers[shift]
|
|
head_start = (ngram_order - 2) * heads_per_ngram
|
|
for head in range(head_start, head_start + heads_per_ngram):
|
|
row.append(torch.remainder(mixed, sizes[head]) + offsets[head])
|
|
rows.append(torch.stack(row))
|
|
return torch.stack(rows).to(input_ids.device)
|
|
|
|
|
|
_NGRAM_MULTIPLIERS = (
|
|
18_014_398_509_481_983,
|
|
17_114_398_509_481_981,
|
|
16_214_398_509_481_979,
|
|
15_314_398_509_481_977,
|
|
)
|
|
_NGRAM_HEADS_VOCAB_SIZES = (
|
|
101,
|
|
103,
|
|
107,
|
|
109,
|
|
113,
|
|
127,
|
|
131,
|
|
137,
|
|
139,
|
|
149,
|
|
151,
|
|
157,
|
|
163,
|
|
167,
|
|
173,
|
|
179,
|
|
181,
|
|
191,
|
|
193,
|
|
197,
|
|
199,
|
|
211,
|
|
223,
|
|
227,
|
|
)
|
|
_NGRAM_EOS_TOKEN_ID = 251
|
|
_NGRAM_HEADS_PER_NGRAM = 8
|
|
|
|
|
|
def _ngram_hash_params(device: torch.device, context_len: int) -> dict:
|
|
num_heads = context_len * _NGRAM_HEADS_PER_NGRAM
|
|
sizes = torch.tensor(
|
|
_NGRAM_HEADS_VOCAB_SIZES[:num_heads], dtype=torch.int64, device=device
|
|
)
|
|
offsets = torch.zeros_like(sizes)
|
|
offsets[1:] = torch.cumsum(sizes, dim=0)[:-1]
|
|
return {
|
|
"layer_multipliers": torch.tensor(
|
|
_NGRAM_MULTIPLIERS[: context_len + 1],
|
|
dtype=torch.int64,
|
|
device=device,
|
|
),
|
|
"ngram_heads_vocab_sizes": sizes,
|
|
"ngram_heads_offsets": offsets,
|
|
"eos_token_id": _NGRAM_EOS_TOKEN_ID,
|
|
"heads_per_ngram": _NGRAM_HEADS_PER_NGRAM,
|
|
}
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused PLE needs CUDA")
|
|
@pytest.mark.parametrize(
|
|
("query_lens", "eos_offsets", "contexts", "first_token_id"),
|
|
[
|
|
([1], [], [[11, 12]], 20),
|
|
([3], [1], [[11]], 20),
|
|
([4, 4], [0, 7], [[11, 12], [13, 14]], 20),
|
|
([4, 0, 3], [], [[11, 12], [13, 14], [15, 16]], 20),
|
|
(
|
|
[3, 2, 0, 0],
|
|
[],
|
|
[
|
|
[11, 12],
|
|
[13, 14],
|
|
[_NGRAM_EOS_TOKEN_ID, _NGRAM_EOS_TOKEN_ID],
|
|
[_NGRAM_EOS_TOKEN_ID, _NGRAM_EOS_TOKEN_ID],
|
|
],
|
|
20,
|
|
),
|
|
(
|
|
[1, 33, 2],
|
|
[5, 32],
|
|
[[_NGRAM_EOS_TOKEN_ID, 11], [12, 13], [14, _NGRAM_EOS_TOKEN_ID]],
|
|
20,
|
|
),
|
|
(
|
|
[5, 12, 16, 1, 16, 17],
|
|
[10, 40],
|
|
[[11, 12], [13, 14], [15, 16], [17, 18], [19, 20], [21, 22]],
|
|
20,
|
|
),
|
|
([5, 3], [3], [[11, 12, 13], [14, 15, 16]], 20),
|
|
(
|
|
[4, 0, 3],
|
|
[],
|
|
[
|
|
[200_000, 200_001],
|
|
[250_000, 250_001],
|
|
[300_000, 300_001],
|
|
],
|
|
350_000,
|
|
),
|
|
([3], [], [[1_000_000_000, 1_000_000_001]], 1_000_000_002),
|
|
],
|
|
ids=[
|
|
"single-token",
|
|
"bigram-only",
|
|
"power-of-two",
|
|
"empty-request",
|
|
"trailing-padded-requests",
|
|
"three-requests",
|
|
"six-requests",
|
|
"four-gram",
|
|
"int64-overflow",
|
|
"large-int32-ids",
|
|
],
|
|
)
|
|
def test_fused_ngram_ids_correctness(
|
|
query_lens: list[int],
|
|
eos_offsets: list[int],
|
|
contexts: list[list[int]],
|
|
first_token_id: int,
|
|
) -> None:
|
|
from vllm.models.qwen4_exp.common.ops.ple import ple_ngram_ids
|
|
|
|
device = torch.device("cuda")
|
|
query_start_loc = torch.tensor(
|
|
[0, *accumulate(query_lens)], dtype=torch.int32, device=device
|
|
)
|
|
input_ids = torch.arange(
|
|
first_token_id,
|
|
first_token_id + sum(query_lens),
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
if eos_offsets:
|
|
input_ids[eos_offsets] = _NGRAM_EOS_TOKEN_ID
|
|
ngram_context = torch.tensor(contexts, dtype=torch.int32, device=device)
|
|
params = _ngram_hash_params(device, ngram_context.shape[1])
|
|
|
|
expected = _reference_ngram_ids(input_ids, query_start_loc, ngram_context, **params)
|
|
actual = ple_ngram_ids(input_ids, query_start_loc, ngram_context, **params)
|
|
|
|
assert torch.equal(actual, expected)
|
|
offsets = params["ngram_heads_offsets"]
|
|
sizes = params["ngram_heads_vocab_sizes"]
|
|
assert torch.all((actual >= offsets) & (actual < offsets + sizes))
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused PLE needs CUDA")
|
|
def test_ngram_prefetch_ids_outlive_start_prefetch() -> None:
|
|
"""Eager breaks hand the side-stream lookup weak refs, so the ids must stay
|
|
valid after start_prefetch returns rather than be reused from the pool."""
|
|
device = torch.device("cuda")
|
|
query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
|
|
input_ids = torch.arange(20, 25, dtype=torch.int32, device=device)
|
|
ngram_context = torch.tensor([[11, 12], [13, 14]], dtype=torch.int32, device=device)
|
|
params = _ngram_hash_params(device, ngram_context.shape[1])
|
|
module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
|
|
nn.Module.__init__(module)
|
|
for name in ("layer_multipliers", "ngram_heads_vocab_sizes", "ngram_heads_offsets"):
|
|
module.register_buffer(name, params[name])
|
|
module.eos_token_id = params["eos_token_id"]
|
|
module.heads_per_ngram = params["heads_per_ngram"]
|
|
num_heads = params["ngram_heads_vocab_sizes"].numel()
|
|
module._prefetch_ids = torch.empty(8, num_heads, dtype=torch.long, device=device)
|
|
prefetched: list[torch.Tensor] = []
|
|
module.ngram_embedding = SimpleNamespace(
|
|
supports_prefetch=True,
|
|
start_prefetch=lambda ids: prefetched.append(weak_ref_tensor(ids)),
|
|
)
|
|
expected = _reference_ngram_ids(input_ids, query_start_loc, ngram_context, **params)
|
|
|
|
# A fresh pool makes the next same-size allocation reuse any freed ids,
|
|
# like a later CUDA graph segment would.
|
|
pool = torch.cuda.MemPool()
|
|
with torch.cuda.use_mem_pool(pool):
|
|
module.start_prefetch(input_ids, query_start_loc, ngram_context)
|
|
torch.full_like(expected, -1)
|
|
|
|
assert torch.equal(prefetched[0], expected)
|
|
|
|
|
|
def _short_conv_dilated_decode_pytorch(
|
|
x_d: torch.Tensor,
|
|
conv_state: torch.Tensor,
|
|
conv_weights: torch.Tensor,
|
|
state_indices_tensor_d: torch.Tensor,
|
|
has_initial_states_d: torch.Tensor | None,
|
|
*,
|
|
conv_state_len: int,
|
|
dilation: int,
|
|
) -> torch.Tensor:
|
|
state_indices = state_indices_tensor_d.to(
|
|
device=conv_state.device, dtype=torch.int64
|
|
)
|
|
# FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to
|
|
# slot 0 for a safe gather, then zero output and skip write-back.
|
|
valid_state = state_indices != NULL_BLOCK_ID
|
|
state_indices = torch.where(
|
|
valid_state, state_indices, torch.zeros_like(state_indices)
|
|
)
|
|
if has_initial_states_d is None:
|
|
has_initial_state = valid_state
|
|
else:
|
|
if has_initial_states_d.numel() < state_indices_tensor_d.numel():
|
|
raise ValueError(
|
|
"has_initial_states_d size mismatch: "
|
|
f"got {has_initial_states_d.numel()}, "
|
|
f"need >= {state_indices_tensor_d.numel()}."
|
|
)
|
|
has_initial_state = has_initial_states_d[: state_indices_tensor_d.numel()].to(
|
|
device=conv_state.device, dtype=torch.bool
|
|
)
|
|
has_initial_state = has_initial_state & valid_state
|
|
|
|
cached_state = conv_state.index_select(0, state_indices)
|
|
state = cached_state[..., :conv_state_len].to(x_d.dtype)
|
|
if conv_state_len > 0:
|
|
initial_state = torch.where(
|
|
has_initial_state.view(-1, 1, 1),
|
|
state,
|
|
torch.zeros_like(state),
|
|
)
|
|
history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1)
|
|
else:
|
|
history = x_d.unsqueeze(-1)
|
|
|
|
conv_output = F.conv1d(
|
|
history,
|
|
conv_weights.unsqueeze(1).contiguous(),
|
|
groups=history.size(1),
|
|
dilation=dilation,
|
|
).squeeze(-1)
|
|
output = F.silu(conv_output)
|
|
output = output * valid_state.view(-1, 1).to(output.dtype)
|
|
|
|
if conv_state_len > 0:
|
|
next_state = history[..., -conv_state_len:]
|
|
# Padded rows are remapped to the reserved null slot. Preserve its
|
|
# existing value while writing the new states for valid rows.
|
|
existing_base_state = cached_state[..., :conv_state_len]
|
|
safe_next_state = torch.where(
|
|
valid_state.view(-1, 1, 1),
|
|
next_state.to(conv_state.dtype),
|
|
existing_base_state,
|
|
)
|
|
cached_state[..., :conv_state_len] = safe_next_state
|
|
conv_state.index_copy_(0, state_indices, cached_state)
|
|
|
|
return output
|
|
|
|
|
|
def _short_conv_dilated_prefill_pytorch(
|
|
x_p: torch.Tensor,
|
|
metadata: PleShortConvAttentionMetadata,
|
|
conv_state: torch.Tensor,
|
|
conv_weights: torch.Tensor,
|
|
state_indices_tensor_p: torch.Tensor,
|
|
num_prefills: int,
|
|
num_decode_tokens: int,
|
|
num_prefill_tokens: int,
|
|
*,
|
|
conv_state_len: int,
|
|
dilation: int,
|
|
) -> torch.Tensor:
|
|
# ``non_spec_query_start_loc`` covers the non-spec (decode + prefill)
|
|
# requests and equals ``query_start_loc`` when spec-decode is inactive.
|
|
non_spec_query_start_loc = metadata.non_spec_query_start_loc
|
|
if non_spec_query_start_loc is None:
|
|
raise ValueError("query_start_loc is required for prefill short-conv")
|
|
query_start_loc_p = (
|
|
non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens
|
|
)
|
|
# The metadata builder guarantees that the prefill query offsets start
|
|
# at 0 and end at num_prefill_tokens. Avoid reading those values here,
|
|
# since doing so would force a device-to-host synchronization.
|
|
has_initial_states_p = metadata.has_initial_states_p
|
|
if has_initial_states_p is None:
|
|
raise ValueError("has_initial_states_p is required for prefill short-conv")
|
|
|
|
output = torch.empty_like(x_p)
|
|
q_starts = query_start_loc_p.to(torch.int64)
|
|
if state_indices_tensor_p.numel() < num_prefills:
|
|
raise ValueError(
|
|
"state_indices_tensor_p size mismatch: "
|
|
f"got {state_indices_tensor_p.numel()}, "
|
|
f"need >= {num_prefills}."
|
|
)
|
|
if has_initial_states_p.numel() < num_prefills:
|
|
raise ValueError(
|
|
"has_initial_states_p size mismatch: "
|
|
f"got {has_initial_states_p.numel()}, "
|
|
f"need >= {num_prefills}."
|
|
)
|
|
if num_prefills == 0 or x_p.numel() == 0:
|
|
return output
|
|
lengths = q_starts[1:] - q_starts[:-1]
|
|
# Use the CPU-computed packing width from the metadata builder instead
|
|
# of synchronizing on lengths.max().
|
|
max_len = metadata.max_prefill_query_len
|
|
if max_len <= 0:
|
|
return output
|
|
|
|
hidden_size = x_p.shape[1]
|
|
positions = torch.arange(num_prefill_tokens, device=x_p.device, dtype=torch.int64)
|
|
req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
|
|
col_indices = positions - q_starts[req_indices]
|
|
|
|
packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size))
|
|
packed_tokens[req_indices, col_indices] = x_p
|
|
packed_tokens = packed_tokens.transpose(1, 2).contiguous()
|
|
|
|
state_indices = state_indices_tensor_p[:num_prefills].to(
|
|
device=conv_state.device, dtype=torch.int64
|
|
)
|
|
valid_state = state_indices != NULL_BLOCK_ID
|
|
state_indices = torch.where(
|
|
valid_state, state_indices, torch.zeros_like(state_indices)
|
|
)
|
|
has_initial = has_initial_states_p[:num_prefills].to(
|
|
device=conv_state.device, dtype=torch.bool
|
|
)
|
|
if conv_state_len > 0:
|
|
if conv_state.shape[0] == 0:
|
|
state = conv_state.new_zeros(
|
|
(num_prefills, hidden_size, conv_state_len),
|
|
dtype=x_p.dtype,
|
|
)
|
|
else:
|
|
state = conv_state.index_select(0, state_indices)[..., :conv_state_len].to(
|
|
x_p.dtype
|
|
)
|
|
use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1)
|
|
initial_state = torch.where(
|
|
use_initial_mask,
|
|
state,
|
|
torch.zeros_like(state),
|
|
)
|
|
history = torch.cat((initial_state, packed_tokens), dim=-1)
|
|
else:
|
|
history = packed_tokens
|
|
|
|
conv_output = F.conv1d(
|
|
history,
|
|
conv_weights.unsqueeze(1).contiguous(),
|
|
groups=history.size(1),
|
|
dilation=dilation,
|
|
)
|
|
conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
|
|
|
|
token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64)
|
|
valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1)
|
|
valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view(
|
|
num_prefills, 1
|
|
)
|
|
conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0)
|
|
output.copy_(conv_output[req_indices, col_indices])
|
|
|
|
if conv_state_len > 0 and conv_state.shape[0] > 0:
|
|
state_starts = lengths.to(device=history.device, dtype=torch.int64).view(
|
|
num_prefills, 1, 1
|
|
)
|
|
state_offsets = torch.arange(
|
|
conv_state_len, device=history.device, dtype=torch.int64
|
|
).view(1, 1, conv_state_len)
|
|
next_state = history.gather(
|
|
dim=2,
|
|
index=(state_starts + state_offsets).expand(-1, history.size(1), -1),
|
|
)
|
|
# Write back without a host synchronization. Valid, non-empty rows
|
|
# receive their new state; padding and zero-length rows keep the
|
|
# current cache value.
|
|
existing_state = conv_state.index_select(0, state_indices)
|
|
existing_base_state = existing_state[..., :conv_state_len]
|
|
update_mask = valid_state & (lengths.to(device=conv_state.device) > 0)
|
|
safe_next_state = torch.where(
|
|
update_mask.view(num_prefills, 1, 1),
|
|
next_state.to(conv_state.dtype),
|
|
existing_base_state,
|
|
)
|
|
existing_state[..., :conv_state_len] = safe_next_state
|
|
conv_state.index_copy_(0, state_indices, existing_state)
|
|
return output
|
|
|
|
|
|
def _short_conv_dilated_spec_pytorch(
|
|
x_spec: torch.Tensor,
|
|
conv_state: torch.Tensor,
|
|
conv_weights: torch.Tensor,
|
|
spec_state_indices_tensor: torch.Tensor,
|
|
spec_query_start_loc: torch.Tensor,
|
|
num_accepted_tokens: torch.Tensor,
|
|
spec_query_len: int,
|
|
*,
|
|
conv_state_len: int,
|
|
dilation: int,
|
|
) -> torch.Tensor:
|
|
"""Dilated short-conv for speculative-decode (MTP) requests.
|
|
|
|
Each spec request feeds multiple (draft + 1) query tokens. The conv
|
|
outputs are computed causally after rolling back the previous draft
|
|
state by ``num_accepted_tokens - 1``. The current candidate inputs stay
|
|
in the extended cache for the next forward, matching
|
|
``causal_conv1d_update``.
|
|
|
|
``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query
|
|
length and is a Python int, so no host synchronization is needed; this
|
|
keeps the path safe for full CUDA-graph capture/replay where the buffers
|
|
are padded at the request level.
|
|
"""
|
|
num_reqs = spec_state_indices_tensor.numel()
|
|
hidden_size = x_spec.size(-1)
|
|
# Use a fixed packing width instead of synchronizing on lengths.max().
|
|
max_len = spec_query_len
|
|
# Full CUDA graphs can pad these buffers. Only the first num_reqs
|
|
# accepted-token counts belong to actual speculative requests.
|
|
num_accepted_tokens = num_accepted_tokens[:num_reqs]
|
|
q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64)
|
|
# Keep the number of real speculative tokens on the device.
|
|
total_real_tokens = q_starts[num_reqs]
|
|
|
|
state_indices = spec_state_indices_tensor.to(
|
|
device=conv_state.device, dtype=torch.int64
|
|
)
|
|
valid_state = state_indices != NULL_BLOCK_ID
|
|
state_indices = torch.where(
|
|
valid_state, state_indices, torch.zeros_like(state_indices)
|
|
)
|
|
positions = torch.arange(x_spec.size(0), device=x_spec.device, dtype=torch.int64)
|
|
# Route graph-padded token rows to the discarded dummy request so that
|
|
# they cannot overwrite real packed data.
|
|
req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
|
|
valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs)
|
|
clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0))
|
|
col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1)
|
|
pack_req_indices = torch.where(
|
|
valid_tokens,
|
|
clamped_req_indices,
|
|
torch.full_like(req_indices, num_reqs),
|
|
)
|
|
pack_col_indices = torch.where(
|
|
valid_tokens, col_indices, torch.zeros_like(col_indices)
|
|
)
|
|
|
|
# The last request row is the dummy sink for graph padding.
|
|
packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size))
|
|
packed[pack_req_indices, pack_col_indices] = x_spec
|
|
packed = packed.transpose(1, 2).contiguous()
|
|
|
|
if conv_state_len > 0:
|
|
cached_state = conv_state.index_select(0, state_indices)
|
|
rollback_offsets = num_accepted_tokens.to(
|
|
device=conv_state.device, dtype=torch.int64
|
|
).sub(1)
|
|
rollback_offsets = torch.where(
|
|
valid_state,
|
|
rollback_offsets.clamp_(0, max_len - 1),
|
|
torch.zeros_like(rollback_offsets),
|
|
)
|
|
state_offsets = torch.arange(
|
|
conv_state_len, device=conv_state.device, dtype=torch.int64
|
|
).view(1, 1, conv_state_len)
|
|
rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets
|
|
state = cached_state.gather(2, rollback_indices.expand(-1, hidden_size, -1)).to(
|
|
x_spec.dtype
|
|
)
|
|
state = torch.where(
|
|
valid_state.view(num_reqs, 1, 1),
|
|
state,
|
|
torch.zeros_like(state),
|
|
)
|
|
# Append a zeroed dummy-row state to match the [num_reqs + 1] pack.
|
|
dummy_state = state.new_zeros((1, hidden_size, conv_state_len))
|
|
state_full = torch.cat((state, dummy_state), dim=0)
|
|
history = torch.cat((state_full, packed), dim=-1)
|
|
else:
|
|
history = packed
|
|
|
|
conv_output = F.conv1d(
|
|
history,
|
|
conv_weights.unsqueeze(1).contiguous(),
|
|
groups=history.size(1),
|
|
dilation=dilation,
|
|
)
|
|
conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
|
|
|
|
output = conv_output[pack_req_indices, pack_col_indices]
|
|
output = output * valid_tokens.view(-1, 1).to(output.dtype)
|
|
|
|
# Keep all current candidate inputs in the extended state. On the next
|
|
# target forward, ``num_accepted_tokens - 1`` selects the rollback
|
|
# window before processing the newly scheduled tokens.
|
|
if conv_state_len > 0:
|
|
state_capacity = conv_state_len + max_len - 1
|
|
if conv_state.size(-1) < state_capacity:
|
|
raise RuntimeError(
|
|
"PLE short-conv cache cannot retain speculative tokens: "
|
|
f"got {conv_state.size(-1)}, need {state_capacity}."
|
|
)
|
|
candidate_state = history[:num_reqs, :, 1 : state_capacity + 1]
|
|
query_lengths = q_starts[1:] - q_starts[:-1]
|
|
state_positions = torch.arange(
|
|
state_capacity, device=history.device, dtype=torch.int64
|
|
).view(1, 1, state_capacity)
|
|
update_lengths = (conv_state_len + query_lengths - 1).view(num_reqs, 1, 1)
|
|
update_mask = valid_state.view(num_reqs, 1, 1) & (
|
|
state_positions < update_lengths
|
|
)
|
|
existing_state = cached_state[..., :state_capacity]
|
|
next_state = torch.where(
|
|
update_mask,
|
|
candidate_state.to(conv_state.dtype),
|
|
existing_state,
|
|
)
|
|
cached_state[..., :state_capacity] = next_state
|
|
conv_state.index_copy_(0, state_indices, cached_state)
|
|
|
|
return output
|
|
|
|
|
|
def _short_conv_dilated_dispatch_pytorch(
|
|
inputs: torch.Tensor,
|
|
residual: torch.Tensor,
|
|
metadata: PleShortConvAttentionMetadata,
|
|
conv_state: torch.Tensor,
|
|
conv_weights: torch.Tensor,
|
|
*,
|
|
conv_state_len: int,
|
|
dilation: int,
|
|
) -> None:
|
|
"""PyTorch reference for fused short-convolution dispatch."""
|
|
num_prefills = metadata.num_prefills
|
|
num_decodes = metadata.num_decodes
|
|
num_decode_tokens = metadata.num_decode_tokens
|
|
num_prefill_tokens = metadata.num_prefill_tokens
|
|
has_prefill = num_prefills > 0
|
|
has_decode = num_decodes > 0
|
|
has_spec = metadata.spec_sequence_masks is not None
|
|
x = inputs[: metadata.num_actual_tokens]
|
|
residual = residual[: metadata.num_actual_tokens]
|
|
|
|
if has_spec:
|
|
if has_prefill or has_decode:
|
|
assert metadata.spec_token_indx is not None
|
|
assert metadata.non_spec_token_indx is not None
|
|
spec_token_indices = metadata.spec_token_indx.long()
|
|
non_spec_token_indices = metadata.non_spec_token_indx.long()
|
|
x_spec = x.index_select(0, spec_token_indices)
|
|
x_non_spec = x.index_select(0, non_spec_token_indices)
|
|
residual_spec = residual.index_select(0, spec_token_indices)
|
|
residual_non_spec = residual.index_select(0, non_spec_token_indices)
|
|
else:
|
|
spec_token_indices = None
|
|
non_spec_token_indices = None
|
|
x_spec = x
|
|
x_non_spec = None
|
|
residual_spec = residual
|
|
residual_non_spec = None
|
|
else:
|
|
spec_token_indices = None
|
|
non_spec_token_indices = None
|
|
x_spec = None
|
|
x_non_spec = x
|
|
residual_spec = None
|
|
residual_non_spec = residual
|
|
|
|
if has_spec:
|
|
assert metadata.spec_state_indices_tensor is not None
|
|
assert metadata.spec_query_start_loc is not None
|
|
assert metadata.num_accepted_tokens is not None
|
|
assert x_spec is not None
|
|
assert residual_spec is not None
|
|
spec_state_indices = metadata.spec_state_indices_tensor[
|
|
: metadata.num_spec_decodes
|
|
]
|
|
conv_output = _short_conv_dilated_spec_pytorch(
|
|
x_spec=x_spec,
|
|
conv_state=conv_state,
|
|
conv_weights=conv_weights,
|
|
spec_state_indices_tensor=spec_state_indices,
|
|
spec_query_start_loc=metadata.spec_query_start_loc,
|
|
num_accepted_tokens=metadata.num_accepted_tokens,
|
|
spec_query_len=metadata.spec_query_len,
|
|
conv_state_len=conv_state_len,
|
|
dilation=dilation,
|
|
)
|
|
residual_spec.add_(conv_output)
|
|
|
|
state_indices = metadata.state_indices_tensor
|
|
if x_non_spec is not None:
|
|
assert state_indices is not None
|
|
assert residual_non_spec is not None
|
|
if has_prefill:
|
|
state_indices_d, state_indices_p = torch.split(
|
|
state_indices, [num_decodes, num_prefills], dim=0
|
|
)
|
|
x_d, x_p = torch.split(
|
|
x_non_spec, [num_decode_tokens, num_prefill_tokens], dim=0
|
|
)
|
|
residual_d, residual_p = torch.split(
|
|
residual_non_spec,
|
|
[num_decode_tokens, num_prefill_tokens],
|
|
dim=0,
|
|
)
|
|
if has_decode:
|
|
conv_output = _short_conv_dilated_decode_pytorch(
|
|
x_d=x_d,
|
|
conv_state=conv_state,
|
|
conv_weights=conv_weights,
|
|
state_indices_tensor_d=state_indices_d,
|
|
has_initial_states_d=metadata.has_initial_states_d,
|
|
conv_state_len=conv_state_len,
|
|
dilation=dilation,
|
|
)
|
|
residual_d.add_(conv_output)
|
|
conv_output = _short_conv_dilated_prefill_pytorch(
|
|
x_p=x_p,
|
|
metadata=metadata,
|
|
conv_state=conv_state,
|
|
conv_weights=conv_weights,
|
|
state_indices_tensor_p=state_indices_p,
|
|
num_prefills=num_prefills,
|
|
num_decode_tokens=num_decode_tokens,
|
|
num_prefill_tokens=num_prefill_tokens,
|
|
conv_state_len=conv_state_len,
|
|
dilation=dilation,
|
|
)
|
|
residual_p.add_(conv_output)
|
|
else:
|
|
conv_output = _short_conv_dilated_decode_pytorch(
|
|
x_d=x_non_spec,
|
|
conv_state=conv_state,
|
|
conv_weights=conv_weights,
|
|
state_indices_tensor_d=state_indices[: x_non_spec.size(0)],
|
|
has_initial_states_d=metadata.has_initial_states_d,
|
|
conv_state_len=conv_state_len,
|
|
dilation=dilation,
|
|
)
|
|
residual_non_spec.add_(conv_output)
|
|
|
|
if has_spec and residual_non_spec is not None:
|
|
assert spec_token_indices is not None
|
|
assert non_spec_token_indices is not None
|
|
assert residual_spec is not None
|
|
residual.index_copy_(0, spec_token_indices, residual_spec)
|
|
residual.index_copy_(0, non_spec_token_indices, residual_non_spec)
|
|
|
|
|
|
_KERNEL_STATE_BLOCKS = 64
|
|
|
|
|
|
def _make_conv_case(
|
|
device: torch.device,
|
|
seed: int,
|
|
channels: int,
|
|
kernel_size: int,
|
|
dilation: int,
|
|
state_layout: str,
|
|
spec_query_len: int = 1,
|
|
) -> tuple[torch.Generator, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
rng = torch.Generator(device=device).manual_seed(seed)
|
|
state_len = (kernel_size - 1) * dilation
|
|
state_width = state_len + spec_query_len - 1
|
|
state_shape = (
|
|
(_KERNEL_STATE_BLOCKS, state_width, channels)
|
|
if state_layout == "SD"
|
|
else (_KERNEL_STATE_BLOCKS, channels, state_width)
|
|
)
|
|
state_kernel = torch.randn(
|
|
state_shape, device=device, dtype=torch.bfloat16, generator=rng
|
|
)
|
|
conv_state = (
|
|
state_kernel.transpose(-1, -2) if state_layout == "SD" else state_kernel
|
|
)
|
|
state_reference = conv_state.clone()
|
|
weights = torch.randn(
|
|
channels, kernel_size, device=device, dtype=torch.bfloat16, generator=rng
|
|
)
|
|
return rng, state_reference, conv_state, weights
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _ConvBatchCase:
|
|
spec_query_lens: tuple[int, ...] = ()
|
|
num_accepted: tuple[int, ...] = ()
|
|
num_decodes: int = 0
|
|
prefill_query_lens: tuple[int, ...] = ()
|
|
channels: int = 2048
|
|
kernel_size: int = 4
|
|
dilation: int = 3
|
|
spec_query_len: int = 1
|
|
graph_padding: int = 0
|
|
state_index_stride: int = 1
|
|
include_null_state: bool = True
|
|
|
|
|
|
def _make_conv_metadata(
|
|
case: _ConvBatchCase,
|
|
device: torch.device,
|
|
) -> tuple[SimpleNamespace, int]:
|
|
request_groups: dict[str, list[int]] = {
|
|
"spec": [],
|
|
"decode": [],
|
|
"prefill": [],
|
|
}
|
|
token_offset = 0
|
|
max_num_reqs = max(
|
|
len(case.spec_query_lens),
|
|
case.num_decodes,
|
|
len(case.prefill_query_lens),
|
|
)
|
|
request_kinds = []
|
|
for req_idx in range(max_num_reqs):
|
|
if req_idx < len(case.spec_query_lens):
|
|
request_kinds.append(("spec", case.spec_query_lens[req_idx]))
|
|
if req_idx < case.num_decodes:
|
|
request_kinds.append(("decode", 1))
|
|
if req_idx < len(case.prefill_query_lens):
|
|
request_kinds.append(("prefill", case.prefill_query_lens[req_idx]))
|
|
for kind, query_len in request_kinds:
|
|
request_groups[kind].extend(range(token_offset, token_offset + query_len))
|
|
token_offset += query_len
|
|
|
|
num_spec_reqs = len(case.spec_query_lens)
|
|
num_prefills = len(case.prefill_query_lens)
|
|
has_spec = num_spec_reqs > 0
|
|
has_non_spec = case.num_decodes > 0 or num_prefills > 0
|
|
mixed_batch = has_spec and has_non_spec
|
|
|
|
spec_state_indices = torch.arange(
|
|
1, num_spec_reqs + 1, dtype=torch.int32, device=device
|
|
)
|
|
if num_spec_reqs >= 3:
|
|
spec_state_indices[1] = NULL_BLOCK_ID
|
|
non_spec_state_indices = torch.arange(
|
|
num_spec_reqs + 1,
|
|
num_spec_reqs + case.num_decodes + num_prefills + 1,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
if case.state_index_stride > 1:
|
|
strided_indices = torch.full(
|
|
(non_spec_state_indices.numel(), case.state_index_stride),
|
|
NULL_BLOCK_ID,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
strided_indices[:, 0] = non_spec_state_indices
|
|
non_spec_state_indices = strided_indices[:, 0]
|
|
if case.num_decodes > 1:
|
|
non_spec_state_indices[case.num_decodes - 1] = NULL_BLOCK_ID
|
|
if num_prefills > 1 and case.include_null_state:
|
|
non_spec_state_indices[case.num_decodes + 1] = NULL_BLOCK_ID
|
|
|
|
spec_query_start_loc = torch.tensor(
|
|
[0, *accumulate(case.spec_query_lens)], dtype=torch.int32, device=device
|
|
)
|
|
non_spec_query_lens = (1,) * case.num_decodes + case.prefill_query_lens
|
|
non_spec_query_start_loc = torch.tensor(
|
|
[0, *accumulate(non_spec_query_lens)], dtype=torch.int32, device=device
|
|
)
|
|
spec_token_indices = torch.tensor(
|
|
request_groups["spec"], dtype=torch.int64, device=device
|
|
)
|
|
non_spec_token_indices = torch.tensor(
|
|
request_groups["decode"] + request_groups["prefill"],
|
|
dtype=torch.int64,
|
|
device=device,
|
|
)
|
|
|
|
metadata = SimpleNamespace(
|
|
num_prefills=num_prefills,
|
|
num_decodes=case.num_decodes,
|
|
num_decode_tokens=case.num_decodes,
|
|
num_prefill_tokens=sum(case.prefill_query_lens),
|
|
spec_sequence_masks=(
|
|
torch.tensor([kind == "spec" for kind, _ in request_kinds], device=device)
|
|
if has_spec
|
|
else None
|
|
),
|
|
spec_token_indx=spec_token_indices if mixed_batch else None,
|
|
non_spec_token_indx=non_spec_token_indices if mixed_batch else None,
|
|
num_actual_tokens=token_offset + case.graph_padding,
|
|
spec_state_indices_tensor=spec_state_indices if has_spec else None,
|
|
spec_query_start_loc=spec_query_start_loc if has_spec else None,
|
|
num_accepted_tokens=(
|
|
torch.tensor(case.num_accepted, dtype=torch.int32, device=device)
|
|
if has_spec
|
|
else None
|
|
),
|
|
spec_query_len=case.spec_query_len,
|
|
state_indices_tensor=non_spec_state_indices if has_non_spec else None,
|
|
has_initial_states_d=(
|
|
torch.arange(case.num_decodes, device=device) % 2 == 0
|
|
if case.num_decodes
|
|
else None
|
|
),
|
|
non_spec_query_start_loc=(non_spec_query_start_loc if has_non_spec else None),
|
|
query_start_loc_p=(
|
|
non_spec_query_start_loc[case.num_decodes :] - case.num_decodes
|
|
if num_prefills
|
|
else None
|
|
),
|
|
has_initial_states_p=(
|
|
torch.arange(num_prefills, device=device) % 2 == 0 if num_prefills else None
|
|
),
|
|
max_prefill_query_len=max(case.prefill_query_lens, default=0),
|
|
num_spec_decodes=num_spec_reqs,
|
|
)
|
|
return metadata, token_offset
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused conv needs CUDA")
|
|
@pytest.mark.parametrize("state_layout", ["SD", "DS"])
|
|
@pytest.mark.parametrize(
|
|
"case",
|
|
[
|
|
pytest.param(
|
|
_ConvBatchCase(num_decodes=1, channels=512),
|
|
id="decode-single",
|
|
),
|
|
pytest.param(
|
|
_ConvBatchCase(
|
|
num_decodes=33,
|
|
channels=2048,
|
|
kernel_size=5,
|
|
dilation=1,
|
|
),
|
|
id="decode-padded",
|
|
),
|
|
pytest.param(
|
|
_ConvBatchCase(prefill_query_lens=(37, 0, 5, 128)),
|
|
id="prefill",
|
|
),
|
|
pytest.param(
|
|
_ConvBatchCase(
|
|
prefill_query_lens=(16, 676),
|
|
channels=10240,
|
|
state_index_stride=4,
|
|
include_null_state=False,
|
|
),
|
|
id="prefill-strided-state-indices",
|
|
),
|
|
pytest.param(
|
|
_ConvBatchCase(
|
|
spec_query_lens=(1, 4, 4, 0),
|
|
num_accepted=(1, 2, 4, 1),
|
|
spec_query_len=4,
|
|
graph_padding=4,
|
|
),
|
|
id="spec-graph-padded",
|
|
),
|
|
pytest.param(
|
|
_ConvBatchCase(
|
|
spec_query_lens=(3, 4),
|
|
num_accepted=(2, 4),
|
|
num_decodes=1,
|
|
prefill_query_lens=(5,),
|
|
spec_query_len=4,
|
|
),
|
|
id="mixed",
|
|
),
|
|
pytest.param(
|
|
_ConvBatchCase(
|
|
spec_query_lens=(1, 2),
|
|
num_accepted=(0, 1),
|
|
num_decodes=2,
|
|
prefill_query_lens=(0, 8, 3),
|
|
channels=640,
|
|
kernel_size=3,
|
|
dilation=2,
|
|
spec_query_len=3,
|
|
),
|
|
id="mixed-varied",
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
("layer_cls", "fuse_residual"),
|
|
[
|
|
pytest.param(Qwen4ExpPLELayer, True, id="nvidia-fused"),
|
|
pytest.param(Qwen4ExpPLELayer, False, id="nvidia-unfused"),
|
|
pytest.param(Qwen4ExpPLELayerAMD, True, id="amd-fused"),
|
|
],
|
|
)
|
|
def test_fused_conv_correctness(
|
|
case: _ConvBatchCase,
|
|
state_layout: str,
|
|
layer_cls: type[nn.Module],
|
|
fuse_residual: bool,
|
|
) -> None:
|
|
device = torch.device("cuda")
|
|
metadata, num_real_tokens = _make_conv_metadata(case, device)
|
|
if case.state_index_stride > 1:
|
|
assert metadata.state_indices_tensor.stride(0) == case.state_index_stride
|
|
# layer_cls is type[nn.Module], so layer_cls.__new__ is type.__new__.
|
|
module = object.__new__(layer_cls)
|
|
nn.Module.__init__(module)
|
|
module.conv_state_len = (case.kernel_size - 1) * case.dilation
|
|
module.short_conv_dilation = case.dilation
|
|
|
|
rng, state_reference, conv_state, weights = _make_conv_case(
|
|
device,
|
|
seed=num_real_tokens + case.channels,
|
|
channels=case.channels,
|
|
kernel_size=case.kernel_size,
|
|
dilation=case.dilation,
|
|
state_layout=state_layout,
|
|
spec_query_len=case.spec_query_len,
|
|
)
|
|
inputs = torch.randn(
|
|
metadata.num_actual_tokens,
|
|
case.channels,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
generator=rng,
|
|
)
|
|
residual = torch.randn(
|
|
inputs.shape,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
generator=rng,
|
|
)
|
|
outer_residual = torch.randn(
|
|
inputs.shape,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
generator=rng,
|
|
)
|
|
null_state = conv_state[NULL_BLOCK_ID].clone()
|
|
residual_kernel = residual.clone()
|
|
# The unfused kernel must overwrite the buffer without reading its contents.
|
|
residual_reference = (
|
|
residual.clone() if fuse_residual else torch.zeros_like(residual)
|
|
)
|
|
|
|
module._short_conv_dilated_dispatch(
|
|
inputs=inputs,
|
|
residual=residual_kernel,
|
|
outer_residual=outer_residual if fuse_residual else None,
|
|
metadata=metadata,
|
|
conv_state=conv_state,
|
|
conv_weights=weights,
|
|
)
|
|
_short_conv_dilated_dispatch_pytorch(
|
|
inputs=inputs,
|
|
residual=residual_reference,
|
|
metadata=metadata,
|
|
conv_state=state_reference,
|
|
conv_weights=weights,
|
|
conv_state_len=module.conv_state_len,
|
|
dilation=module.short_conv_dilation,
|
|
)
|
|
if fuse_residual:
|
|
residual_reference = outer_residual + residual_reference
|
|
|
|
assert torch.equal(
|
|
residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens]
|
|
)
|
|
assert torch.equal(conv_state, state_reference)
|
|
assert torch.equal(conv_state[NULL_BLOCK_ID], null_state)
|
|
if case.graph_padding:
|
|
assert torch.equal(
|
|
residual_kernel[num_real_tokens:],
|
|
residual_reference[num_real_tokens:],
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused gate needs CUDA")
|
|
@pytest.mark.parametrize(("num_tokens", "strided_kv"), [(1, False), (64, True)])
|
|
def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None:
|
|
import math
|
|
|
|
from vllm.models.qwen4_exp.common.ops.ple import ple_gate
|
|
|
|
device = torch.device("cuda")
|
|
hc, h = 4, 2560
|
|
generator = torch.Generator(device=device).manual_seed(num_tokens)
|
|
kv = torch.randn(
|
|
num_tokens,
|
|
hc * h + h,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
generator=generator,
|
|
)
|
|
key = kv[:, : hc * h]
|
|
value = kv[:, hc * h :]
|
|
if not strided_kv:
|
|
key = key.contiguous()
|
|
value = value.contiguous()
|
|
hidden = torch.randn(
|
|
num_tokens,
|
|
hc * h,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
generator=generator,
|
|
)
|
|
|
|
def make_norm_weight() -> torch.Tensor:
|
|
weight = torch.empty(hc * h, device=device, dtype=torch.bfloat16)
|
|
return weight.normal_(mean=-0.1, std=0.1, generator=generator)
|
|
|
|
norm_key = make_norm_weight()
|
|
norm_query = make_norm_weight()
|
|
norm_conv = make_norm_weight()
|
|
|
|
def grouped_norm(inputs: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
|
grouped = inputs.float().unflatten(-1, (hc, h))
|
|
var = grouped.square().mean(dim=-1, keepdim=True)
|
|
normalized = grouped * torch.rsqrt(var + 1e-6)
|
|
return (normalized.flatten(-2) * (1.0 + weight.float())).to(inputs.dtype)
|
|
|
|
gated, normed = ple_gate(
|
|
key,
|
|
value,
|
|
hidden,
|
|
norm_key,
|
|
norm_query,
|
|
norm_conv,
|
|
1e-6,
|
|
)
|
|
key_normalized = grouped_norm(key, norm_key).reshape(num_tokens, hc, h)
|
|
query_normalized = grouped_norm(hidden, norm_query).reshape(num_tokens, hc, h)
|
|
dot = (key_normalized * query_normalized).sum(dim=-1, keepdim=True)
|
|
dot = (dot / math.sqrt(h)).to(torch.bfloat16)
|
|
gate = torch.sigmoid(dot.sign() * dot.abs().clamp_min(1e-6).sqrt()).to(
|
|
torch.bfloat16
|
|
)
|
|
expected_gated = (gate * value.unsqueeze(-2)).flatten(-2)
|
|
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,
|
|
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")
|
|
def test_amd_pinned_embedding_prefetch_machinery_is_lazy(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""The pinned embedding must not allocate the prefetch stream/buffer in init.
|
|
|
|
Paths that never prefetch (AMD uses the synchronous pinned-lookup op)
|
|
keep the stream/buffer unallocated for the process lifetime. NVIDIA's
|
|
first ``start_prefetch`` — the eager profile run, always before cudagraph
|
|
capture — allocates them lazily.
|
|
"""
|
|
module, _ = _build_amd_ngram_embedding(
|
|
monkeypatch, device="cuda:0", cpu_offload=True, fp8_checkpoint=False
|
|
)
|
|
embedding = module.ngram_embedding
|
|
assert embedding._prefetch_buffer is None
|
|
assert embedding._prefetch_stream is None
|
|
|
|
|
|
@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)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_rocm(),
|
|
reason="only the ROCm backend runs the n-gram embedding under torch.compile",
|
|
)
|
|
@pytest.mark.parametrize(
|
|
("query_lens", "contexts"),
|
|
[([3, 1], [[5, 6], [0, 7]]), ([1, 0, 2], [[1, 2], [3, 4], [0, 0]])],
|
|
ids=["two-requests", "empty-request"],
|
|
)
|
|
def test_amd_ngram_embedding_matches_reference_under_compile(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
query_lens: list[int],
|
|
contexts: list[list[int]],
|
|
) -> None:
|
|
"""The AMD n-gram ID op must run under Dynamo with a symbolic request count."""
|
|
module, loaded_weight = _build_amd_ngram_embedding(
|
|
monkeypatch, device="cuda:0", cpu_offload=False, fp8_checkpoint=False
|
|
)
|
|
device = torch.device("cuda:0")
|
|
query_start_loc = torch.tensor(
|
|
[0, *accumulate(query_lens)], dtype=torch.int32, device=device
|
|
)
|
|
input_ids = torch.tensor(
|
|
[9, 0, 11, 12][: sum(query_lens)], dtype=torch.int32, device=device
|
|
)
|
|
ngram_context = torch.tensor(contexts, dtype=torch.int32, device=device)
|
|
|
|
compiled = torch.compile(module, fullgraph=True, dynamic=True)
|
|
output = compiled(input_ids, query_start_loc, ngram_context)
|
|
|
|
ngram_ids = _reference_ngram_ids(
|
|
input_ids,
|
|
query_start_loc,
|
|
ngram_context,
|
|
module.layer_multipliers,
|
|
module.ngram_heads_vocab_sizes,
|
|
module.ngram_heads_offsets,
|
|
module.eos_token_id,
|
|
module.heads_per_ngram,
|
|
)
|
|
expected = loaded_weight[ngram_ids.cpu()].flatten(-2).to(device)
|
|
torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
|