[Qwen4Exp][ROCm] PLE n-gram table CPU offload (#57497)
Signed-off-by: Mathew Odden <modden@redhat.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: opencode+deepseek-v4-flash+vllm <opencode+deepseek-v4-flash+vllm@example.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Mathew Odden
co-authored by
opencode+deepseek-v4-flash+vllm
Andreas Karatzas
OpenAI Codex
opencode+deepseek-v4-flash+vllm
Andreas Karatzas
OpenAI Codex
parent
fa94450b74
commit
267eee54bf
@@ -169,6 +169,8 @@ steps:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- pytest -v -s models/qwen4_exp/test_qsa_amd.py
|
||||
# PLE and QSA import conflicting HC registrations in a shared process.
|
||||
- pytest -v -s models/qwen4_exp/test_ple.py
|
||||
|
||||
- label: ":computer: (CPU) Qwen4 Exp"
|
||||
key: qwen4-exp-unit-tests-cpu
|
||||
|
||||
@@ -13,7 +13,7 @@ from torch.nn import functional as F
|
||||
|
||||
import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
|
||||
import vllm.model_executor.parameter as parameter_module
|
||||
import vllm.models.qwen4_exp.nvidia.ngram_embedding as ngram_embedding_module
|
||||
import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module
|
||||
from vllm.config.quantization import QuantizationConfigArgs
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.layers.quantization.modelopt import (
|
||||
@@ -21,11 +21,16 @@ from vllm.model_executor.layers.quantization.modelopt import (
|
||||
ModelOptNvFp4Config,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
||||
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,
|
||||
compute_ple_shard_overlap,
|
||||
copy_ple_embedding_shard_,
|
||||
)
|
||||
from vllm.models.qwen4_exp.config import Qwen4ExpTextConfig
|
||||
from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
|
||||
Qwen4ExpNGramEmbedding,
|
||||
Qwen4ExpPLEDeviceEmbedding,
|
||||
@@ -1750,3 +1755,211 @@ def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None:
|
||||
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,
|
||||
max_num_reqs=2,
|
||||
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")
|
||||
@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)
|
||||
|
||||
+20
-14
@@ -1530,26 +1530,32 @@ def test_engram_dp_shared_memory_config_validation(
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"architecture, ple_layers, accelerator, supported",
|
||||
"architecture, ple_layers, platform, supported",
|
||||
[
|
||||
("DeepseekV41ForCausalLM", [1], True, True),
|
||||
("DeepseekV41ForCausalLM", [], True, False),
|
||||
("DeepseekV41ForCausalLM", [1], False, False),
|
||||
("Qwen4ExpForCausalLM", [1], True, True),
|
||||
("Qwen4ExpForConditionalGeneration", [1], True, True),
|
||||
("Qwen4ExpForCausalLM", [], True, False),
|
||||
("Qwen4ExpForCausalLM", None, True, False),
|
||||
("Qwen4ExpForCausalLM", [1], False, False),
|
||||
("LlamaForCausalLM", [1], True, False),
|
||||
("Qwen4ExpMTP", [], True, False),
|
||||
(None, None, True, False),
|
||||
("DeepseekV41ForCausalLM", [1], "cuda", True),
|
||||
("DeepseekV41ForCausalLM", [1], "rocm", True),
|
||||
("DeepseekV41ForCausalLM", [], "cuda", False),
|
||||
("DeepseekV41ForCausalLM", [1], "cpu", False),
|
||||
("Qwen4ExpForCausalLM", [1], "cuda", True),
|
||||
("Qwen4ExpForCausalLM", [1], "rocm", True),
|
||||
("Qwen4ExpForConditionalGeneration", [1], "cuda", True),
|
||||
("Qwen4ExpForConditionalGeneration", [1], "rocm", True),
|
||||
("Qwen4ExpForCausalLM", [], "cuda", False),
|
||||
("Qwen4ExpForCausalLM", None, "cuda", False),
|
||||
("Qwen4ExpForCausalLM", [1], "cpu", False),
|
||||
("LlamaForCausalLM", [1], "cuda", False),
|
||||
("Qwen4ExpMTP", [], "cuda", False),
|
||||
(None, None, "cuda", False),
|
||||
],
|
||||
)
|
||||
def test_engram_model_support(
|
||||
monkeypatch, architecture, ple_layers, accelerator, supported
|
||||
monkeypatch, architecture, ple_layers, platform, supported
|
||||
):
|
||||
"""A similarly named HF field must not enable unsupported implementations."""
|
||||
monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: accelerator)
|
||||
monkeypatch.setattr(current_platform, "is_cuda", lambda: platform == "cuda")
|
||||
monkeypatch.setattr(
|
||||
current_platform, "is_cuda_alike", lambda: platform in ("cuda", "rocm")
|
||||
)
|
||||
model = (
|
||||
cast(
|
||||
ModelConfig,
|
||||
|
||||
@@ -82,7 +82,8 @@ class EngramConfig:
|
||||
):
|
||||
raise ValueError(
|
||||
"EngramConfig requires a model with supported Engram "
|
||||
"embeddings, non-empty n-gram layer ids, and CUDA."
|
||||
"embeddings, non-empty n-gram layer ids, and a CUDA-alike "
|
||||
"device (CUDA or ROCm)."
|
||||
)
|
||||
|
||||
def resolve_dp_shared_memory(self, parallel_config: "ParallelConfig") -> None:
|
||||
|
||||
@@ -11,6 +11,7 @@ from torch import nn
|
||||
|
||||
from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
from vllm.model_executor.layers.mamba.abstract import MambaBase
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
@@ -18,6 +19,9 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
MambaStateShapeCalculator,
|
||||
is_conv_state_dim_first,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
)
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
from vllm.transformers_utils.configs.qwen4_exp import (
|
||||
Qwen4ExpTextConfig,
|
||||
@@ -30,7 +34,13 @@ from vllm.v1.attention.backends.short_conv_attn import (
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
|
||||
from ..common.ple import PLEVocabParallelEmbedding
|
||||
from ..common.ngram_embedding import (
|
||||
Qwen4ExpPLEDeviceEmbedding,
|
||||
Qwen4ExpPLEEmbeddingMethod,
|
||||
Qwen4ExpPLEPinnedHostEmbedding,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Qwen4ExpPLEGroupedNorm(nn.Module):
|
||||
@@ -172,6 +182,10 @@ class Qwen4ExpNGramEmbedding(nn.Module):
|
||||
max_num_reqs: int,
|
||||
prefix: str,
|
||||
layer_name: str,
|
||||
*,
|
||||
data_parallel_rank: int = 0,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
params_dtype: torch.dtype | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.embedding_dim = embedding_dim
|
||||
@@ -224,11 +238,39 @@ class Qwen4ExpNGramEmbedding(nn.Module):
|
||||
)
|
||||
divisor = int(config.make_ngram_vocab_size_divisible_by)
|
||||
padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
|
||||
self.ngram_embedding = PLEVocabParallelEmbedding(
|
||||
if params_dtype is None:
|
||||
params_dtype = torch.get_default_dtype()
|
||||
embedding_prefix = f"{prefix}.ngram_embedding"
|
||||
embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config(
|
||||
quant_config,
|
||||
embedding_prefix,
|
||||
getattr(config, "ple_embedding_dtype", None),
|
||||
)
|
||||
engram_config = get_current_vllm_config().engram_config
|
||||
embedding_cls = (
|
||||
Qwen4ExpPLEPinnedHostEmbedding
|
||||
if engram_config is not None and engram_config.cpu_offload
|
||||
else Qwen4ExpPLEDeviceEmbedding
|
||||
)
|
||||
self.ngram_embedding = embedding_cls(
|
||||
padded_vocab_size,
|
||||
self.head_dim,
|
||||
params_dtype=params_dtype,
|
||||
padding_size=divisor,
|
||||
prefix=f"{prefix}.ngram_embedding",
|
||||
prefix=embedding_prefix,
|
||||
embedding_method=embedding_quant_method,
|
||||
num_ngram_heads=self.ngram_heads,
|
||||
max_total_tokens=max_total_tokens,
|
||||
data_parallel_rank=data_parallel_rank,
|
||||
)
|
||||
logger.info(
|
||||
"Initialized AMD PLE embedding %s: quantization_method=%s, "
|
||||
"weight_dtype=%s, weight_device=%s, pinned=%s",
|
||||
embedding_prefix,
|
||||
type(embedding_quant_method).__name__,
|
||||
self.ngram_embedding.weight.dtype,
|
||||
self.ngram_embedding.weight.device,
|
||||
self.ngram_embedding.weight.is_pinned(),
|
||||
)
|
||||
self.register_buffer(
|
||||
"positions_buffer",
|
||||
@@ -344,9 +386,21 @@ class Qwen4ExpNGramEmbedding(nn.Module):
|
||||
ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
|
||||
id_blocks.append(ids[request_indices, adjusted_columns])
|
||||
ngram_ids = torch.cat(id_blocks, dim=-1)
|
||||
embedding = self.ngram_embedding
|
||||
if embedding.supports_prefetch:
|
||||
output = ngram_ids.new_empty(
|
||||
(ngram_ids.shape[0], self.embedding_dim),
|
||||
dtype=embedding.weight.dtype,
|
||||
)
|
||||
torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned(
|
||||
ngram_ids,
|
||||
output,
|
||||
self.layer_name,
|
||||
)
|
||||
return output
|
||||
output = ngram_ids.new_empty(
|
||||
(ngram_ids.shape[0], self.embedding_dim),
|
||||
dtype=self.ngram_embedding.params_dtype,
|
||||
dtype=embedding.params_dtype,
|
||||
)
|
||||
torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding(
|
||||
ngram_ids,
|
||||
@@ -459,6 +513,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
vllm_config.scheduler_config.max_num_seqs,
|
||||
f"{prefix}.ple_embedding",
|
||||
prefix,
|
||||
data_parallel_rank=vllm_config.parallel_config.data_parallel_rank,
|
||||
quant_config=quant_config,
|
||||
params_dtype=model_config.dtype,
|
||||
)
|
||||
self.key_proj = ReplicatedLinear(
|
||||
int(config.ple_embed_dim),
|
||||
@@ -1044,6 +1101,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
f"{hidden_states.shape[0]}"
|
||||
)
|
||||
embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
|
||||
embeddings = self.ple_embedding.ngram_embedding.dequantize(
|
||||
embeddings,
|
||||
self.ple_embedding.ngram_embedding.params_dtype,
|
||||
)
|
||||
key, _ = self.key_proj(embeddings)
|
||||
value, _ = self.value_proj(embeddings)
|
||||
token_count = hidden_states.shape[0]
|
||||
@@ -1081,6 +1142,24 @@ def qwen4_exp_amd_ple_ngram_embedding(
|
||||
output.copy_(result)
|
||||
|
||||
|
||||
def qwen4_exp_amd_ple_ngram_embedding_pinned(
|
||||
ngram_ids: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> None:
|
||||
"""Run the pinned PLE UVA lookup outside Inductor's FX graph.
|
||||
|
||||
Same rationale as the device-path escape: keeping the large embedding weight
|
||||
out of the graph prevents AOT compile-time autotuning from materializing a
|
||||
synthetic copy of the weight.
|
||||
"""
|
||||
layer = get_forward_context().no_compile_layers[layer_name]
|
||||
if not isinstance(layer, Qwen4ExpPLELayer):
|
||||
raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner")
|
||||
result = layer.ple_embedding.ngram_embedding.sync_lookup(ngram_ids).flatten(-2)
|
||||
output.copy_(result)
|
||||
|
||||
|
||||
def qwen4_exp_ple_short_conv(
|
||||
inputs: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
@@ -1098,6 +1177,13 @@ direct_register_custom_op(
|
||||
)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="qwen4_exp_amd_ple_ngram_embedding_pinned",
|
||||
op_func=qwen4_exp_amd_ple_ngram_embedding_pinned,
|
||||
mutates_args=["output"],
|
||||
)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="qwen4_exp_ple_short_conv",
|
||||
op_func=qwen4_exp_ple_short_conv,
|
||||
|
||||
@@ -0,0 +1,542 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Shared Qwen4Exp n-gram embedding storage with device and pinned-host backends.
|
||||
|
||||
Both the NVIDIA and AMD Qwen4Exp implementations use these classes so the (large)
|
||||
n-gram embedding table can be kept in pinned host memory and looked up through
|
||||
Unified Virtual Addressing on any CUDA-alike platform.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import ClassVar
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
||||
from vllm.distributed import get_dp_group, get_etp_group, get_tp_group
|
||||
from vllm.forward_context import DPMetadata, get_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.layers.quantization.modelopt import (
|
||||
ModelOptMixedPrecisionConfig,
|
||||
ModelOptQuantConfigBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
create_fp8_scale_parameter,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
is_layer_skipped,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.utils.platform_utils import is_uva_available
|
||||
from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
|
||||
|
||||
from .ple import PLEVocabParallelEmbedding
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC):
|
||||
"""ETP-sharded PLE table shared by device and pinned-host backends."""
|
||||
|
||||
supports_prefetch: ClassVar[bool] = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
*,
|
||||
params_dtype: torch.dtype,
|
||||
padding_size: int,
|
||||
prefix: str,
|
||||
embedding_method: "Qwen4ExpPLEEmbeddingMethod",
|
||||
num_ngram_heads: int = 1,
|
||||
max_total_tokens: int = 0,
|
||||
data_parallel_rank: int = 0,
|
||||
) -> None:
|
||||
del num_ngram_heads, max_total_tokens
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
params_dtype=params_dtype,
|
||||
padding_size=padding_size,
|
||||
prefix=prefix,
|
||||
quant_method=embedding_method,
|
||||
parallel_group=get_etp_group(),
|
||||
)
|
||||
self.embedding_method = embedding_method
|
||||
self.data_parallel_rank = data_parallel_rank
|
||||
tp_size = get_tp_group().world_size
|
||||
if self.tp_size % tp_size:
|
||||
raise ValueError(
|
||||
"ETP size must be divisible by TP size, but got "
|
||||
f"ETP={self.tp_size} and TP={tp_size}"
|
||||
)
|
||||
self.etp_data_parallel_size = self.tp_size // tp_size
|
||||
|
||||
@abstractmethod
|
||||
def allocate_embedding_weight(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate storage for the complete embedding weight."""
|
||||
raise NotImplementedError
|
||||
|
||||
def dequantize(
|
||||
self,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Delegate storage-format conversion to the embedding method."""
|
||||
return self.embedding_method.dequantize(self, embeddings, output_dtype)
|
||||
|
||||
def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]:
|
||||
"""Return the per-DP slot size and this rank's slot offset."""
|
||||
if self.etp_data_parallel_size == 1:
|
||||
return local_num_tokens, 0
|
||||
dp_metadata: DPMetadata | None = get_forward_context().dp_metadata
|
||||
if dp_metadata is None:
|
||||
raise RuntimeError("ETP spanning DP requires DP token metadata")
|
||||
group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * (
|
||||
self.etp_data_parallel_size
|
||||
)
|
||||
group_end = group_start + self.etp_data_parallel_size
|
||||
token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist()
|
||||
group_counts = token_counts[group_start:group_end]
|
||||
slot_size = max(group_counts)
|
||||
dp_rank = get_dp_group().rank_in_group
|
||||
return slot_size, dp_rank * slot_size
|
||||
|
||||
def _gather_dp_ids(
|
||||
self,
|
||||
ngram_ids: torch.Tensor,
|
||||
slot_size: int,
|
||||
) -> torch.Tensor:
|
||||
"""Gather DP-local IDs that share one ETP-sharded PLE table."""
|
||||
if self.etp_data_parallel_size == 1:
|
||||
return ngram_ids
|
||||
if ngram_ids.shape[0] < slot_size:
|
||||
padding = ngram_ids.new_zeros(
|
||||
slot_size - ngram_ids.shape[0], ngram_ids.shape[1]
|
||||
)
|
||||
ngram_ids = torch.cat((ngram_ids, padding), dim=0)
|
||||
return get_dp_group().all_gather(ngram_ids, dim=0)
|
||||
|
||||
def _select_embeddings(
|
||||
self,
|
||||
embeddings: torch.Tensor,
|
||||
local_num_tokens: int,
|
||||
slot_offset: int,
|
||||
) -> torch.Tensor:
|
||||
"""Select this DP rank's rows from the ETP-reduced embeddings."""
|
||||
if self.etp_data_parallel_size == 1:
|
||||
return embeddings
|
||||
return embeddings.narrow(0, slot_offset, local_num_tokens)
|
||||
|
||||
@abstractmethod
|
||||
def start_prefetch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
ngram_ids: torch.Tensor,
|
||||
) -> None:
|
||||
"""Start an asynchronous lookup when supported."""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
|
||||
"""Quantization interface shared by resident and pinned PLE tables."""
|
||||
|
||||
# PLE post-load processing only validates scales in their current storage.
|
||||
requires_device_loading: bool = False
|
||||
|
||||
@staticmethod
|
||||
def from_quant_config(
|
||||
quant_config: QuantizationConfig | None,
|
||||
prefix: str,
|
||||
embedding_dtype: str | None = None,
|
||||
) -> "Qwen4ExpPLEEmbeddingMethod":
|
||||
"""Select the concrete PLE embedding format for a layer."""
|
||||
if embedding_dtype == "float8_e4m3fn":
|
||||
return Qwen4ExpPLEFp8EmbeddingMethod()
|
||||
if quant_config is None:
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if isinstance(quant_config, ModelOptMixedPrecisionConfig):
|
||||
if quant_config._resolve_quant_algo(prefix) == "FP8":
|
||||
return Qwen4ExpPLEFp8EmbeddingMethod()
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if isinstance(
|
||||
quant_config, ModelOptQuantConfigBase
|
||||
) and quant_config.is_layer_excluded(prefix):
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if not isinstance(quant_config, Fp8Config):
|
||||
raise NotImplementedError(
|
||||
"Qwen4Exp PLE embedding does not support quantization config "
|
||||
f"{type(quant_config).__name__}"
|
||||
)
|
||||
|
||||
ignored_layers = quant_config.ignored_layers
|
||||
if is_layer_skipped(
|
||||
prefix,
|
||||
ignored_layers,
|
||||
quant_config.packed_modules_mapping,
|
||||
match_mode=quant_config.ignored_layers_match_mode,
|
||||
):
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
# PLE checkpoint shards form one runtime embedding parameter.
|
||||
shard_prefix = f"{prefix}.shard_"
|
||||
if any(name.startswith(shard_prefix) for name in ignored_layers):
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if not quant_config.is_checkpoint_fp8_serialized:
|
||||
raise NotImplementedError(
|
||||
"Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
|
||||
)
|
||||
return Qwen4ExpPLEFp8EmbeddingMethod()
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError("PLE weights only support embedding lookup")
|
||||
|
||||
def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
|
||||
return F.embedding(input_, layer.weight)
|
||||
|
||||
@abstractmethod
|
||||
def dequantize(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Convert looked-up PLE rows to the activation dtype."""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
|
||||
"""Unquantized PLE embedding storage and lookup semantics."""
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: Qwen4ExpPLEEmbedding,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
del input_size, output_size
|
||||
weight = nn.Parameter(
|
||||
layer.allocate_embedding_weight(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
|
||||
set_weight_attrs(weight, extra_weight_attrs)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
def dequantize(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
del layer, output_dtype
|
||||
return embeddings
|
||||
|
||||
|
||||
class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
|
||||
"""FP8 PLE embedding with one global checkpoint scale."""
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: Qwen4ExpPLEEmbedding,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
del input_size, output_size, params_dtype
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
weight = ModelWeightParameter(
|
||||
data=layer.allocate_embedding_weight(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
torch.float8_e4m3fn,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
weight_scale = create_fp8_scale_parameter(
|
||||
PerTensorScaleParameter,
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
None,
|
||||
weight_loader,
|
||||
scale_dtype=torch.float32,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
def process_weights_after_loading(self, layer: nn.Module) -> None:
|
||||
"""Reject FP8 PLE checkpoints without a global scale."""
|
||||
sentinel = torch.finfo(torch.float32).min
|
||||
if torch.any(layer.weight_scale == sentinel):
|
||||
raise ValueError("FP8 PLE checkpoint is missing its global scale")
|
||||
|
||||
def dequantize(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
weight_scale = getattr(layer, "weight_scale", None)
|
||||
if weight_scale is None:
|
||||
raise RuntimeError("FP8 PLE embedding is missing its global scale")
|
||||
if weight_scale.device != embeddings.device:
|
||||
raise RuntimeError("FP8 PLE embedding scale must be on the output device")
|
||||
return embeddings.to(output_dtype) * weight_scale.to(output_dtype)
|
||||
|
||||
|
||||
class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding):
|
||||
"""PLE table allocated on the active model device."""
|
||||
|
||||
def allocate_embedding_weight(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate the complete PLE weight on the active device."""
|
||||
return torch.empty(num_embeddings, embedding_dim, dtype=dtype)
|
||||
|
||||
def start_prefetch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
ngram_ids: torch.Tensor,
|
||||
) -> None:
|
||||
"""Resident embedding prefetch is a no-op."""
|
||||
return None
|
||||
|
||||
def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Gather ETP inputs, look up embeddings, and select local rows."""
|
||||
slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
|
||||
gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
|
||||
embeddings = super().forward(gathered_ids)
|
||||
return self._select_embeddings(
|
||||
embeddings,
|
||||
ngram_ids.shape[0],
|
||||
slot_offset,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _lookup_ple_embedding_from_pinned_kernel(
|
||||
weight_ptr,
|
||||
ids_ptr,
|
||||
output_ptr,
|
||||
embedding_dim,
|
||||
tp_vocab_start,
|
||||
tp_vocab_end,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
"""Look up TP-owned PLE rows through a CUDA view of pinned host memory."""
|
||||
row_id = tl.program_id(0)
|
||||
global_idx = tl.load(ids_ptr + row_id)
|
||||
in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end)
|
||||
local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0)
|
||||
offsets = tl.arange(0, BLOCK_D)
|
||||
store_mask = offsets < embedding_dim
|
||||
load_mask = store_mask & in_range
|
||||
values = tl.load(
|
||||
weight_ptr + local_idx * embedding_dim + offsets,
|
||||
mask=load_mask,
|
||||
other=0.0,
|
||||
)
|
||||
tl.store(
|
||||
output_ptr + row_id * embedding_dim + offsets,
|
||||
values,
|
||||
mask=store_mask,
|
||||
)
|
||||
|
||||
|
||||
class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):
|
||||
"""PLE table loaded into pinned CPU memory and looked up through UVA."""
|
||||
|
||||
supports_prefetch: ClassVar[bool] = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
*,
|
||||
params_dtype: torch.dtype,
|
||||
padding_size: int,
|
||||
prefix: str,
|
||||
embedding_method: Qwen4ExpPLEEmbeddingMethod,
|
||||
num_ngram_heads: int = 1,
|
||||
max_total_tokens: int = 0,
|
||||
data_parallel_rank: int = 0,
|
||||
) -> None:
|
||||
if not is_uva_available():
|
||||
raise RuntimeError("Engram CPU offload requires UVA support")
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
params_dtype=params_dtype,
|
||||
padding_size=padding_size,
|
||||
prefix=prefix,
|
||||
embedding_method=embedding_method,
|
||||
num_ngram_heads=num_ngram_heads,
|
||||
max_total_tokens=max_total_tokens,
|
||||
data_parallel_rank=data_parallel_rank,
|
||||
)
|
||||
self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight)
|
||||
self._block_d = triton.next_power_of_2(self.embedding_dim)
|
||||
self._prefetch_stream = torch.cuda.Stream(device=self._uva_weight.device)
|
||||
self._prefetch_buffer = torch.empty(
|
||||
max_total_tokens * self.etp_data_parallel_size,
|
||||
num_ngram_heads,
|
||||
self.embedding_dim,
|
||||
dtype=self.weight.dtype,
|
||||
device=self._uva_weight.device,
|
||||
)
|
||||
self._output_dim = num_ngram_heads * self.embedding_dim
|
||||
|
||||
def allocate_embedding_weight(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate the complete PLE weight directly in pinned CPU memory."""
|
||||
return torch.empty(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
dtype=dtype,
|
||||
device="cpu",
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
def _lookup(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
output: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Look up local ETP rows while preserving the weight storage dtype."""
|
||||
expected_shape = (*input_ids.shape, self.embedding_dim)
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
expected_shape,
|
||||
dtype=self.weight.dtype,
|
||||
device=input_ids.device,
|
||||
)
|
||||
elif (
|
||||
tuple(output.shape) != expected_shape
|
||||
or output.dtype != self.weight.dtype
|
||||
or output.device != input_ids.device
|
||||
):
|
||||
raise ValueError(
|
||||
"PLE prefetch output must match the input shape, weight dtype, "
|
||||
"and input device"
|
||||
)
|
||||
|
||||
flat_ids = input_ids.reshape(-1).long()
|
||||
if flat_ids.numel():
|
||||
_lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
|
||||
self._uva_weight,
|
||||
flat_ids,
|
||||
output,
|
||||
self.embedding_dim,
|
||||
self.shard_indices.org_vocab_start_index,
|
||||
self.shard_indices.org_vocab_end_index,
|
||||
BLOCK_D=self._block_d,
|
||||
)
|
||||
return output
|
||||
|
||||
def sync_lookup(self, ngram_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Synchronous UVA lookup for platforms without prefetch wiring."""
|
||||
slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
|
||||
gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
|
||||
embeddings = self._lookup(gathered_ids)
|
||||
embeddings = self._reduce_etp_embeddings(embeddings)
|
||||
return self._select_embeddings(
|
||||
embeddings,
|
||||
ngram_ids.shape[0],
|
||||
slot_offset,
|
||||
)
|
||||
|
||||
def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
|
||||
"""Combine pinned lookup results owned by different ETP ranks."""
|
||||
if self.tp_size == 1:
|
||||
return embeddings
|
||||
assert self.parallel_group is not None
|
||||
if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
||||
# Each vocabulary row has one owner, so reduce the raw FP8 bytes.
|
||||
reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
|
||||
return reduced.view(embeddings.dtype)
|
||||
return self.parallel_group.all_reduce(embeddings)
|
||||
|
||||
@eager_break_during_capture
|
||||
def start_prefetch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
ngram_ids: torch.Tensor,
|
||||
) -> None:
|
||||
"""Gather ETP IDs and launch their UVA lookup on the side stream."""
|
||||
slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0])
|
||||
gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
|
||||
active_output = self._prefetch_buffer[: gathered_ids.shape[0]]
|
||||
prefetch_stream = self._prefetch_stream
|
||||
prefetch_stream.wait_stream(torch.cuda.current_stream())
|
||||
gathered_ids.record_stream(prefetch_stream)
|
||||
with torch.cuda.stream(prefetch_stream):
|
||||
self._lookup(gathered_ids, output=active_output)
|
||||
|
||||
@eager_break_during_capture
|
||||
def _finalize_prefetch(
|
||||
self,
|
||||
prefetch_output: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
) -> None:
|
||||
"""Join the side stream, reduce ETP shards, and select local rows."""
|
||||
torch.cuda.current_stream().wait_stream(self._prefetch_stream)
|
||||
slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0])
|
||||
active_output = prefetch_output[: slot_size * self.etp_data_parallel_size]
|
||||
embeddings = self._reduce_etp_embeddings(active_output)
|
||||
embeddings = self._select_embeddings(
|
||||
embeddings,
|
||||
output.shape[0],
|
||||
slot_offset,
|
||||
)
|
||||
output.copy_(embeddings.flatten(-2))
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
"""Finish the pinned lookup into graph-owned output storage."""
|
||||
output = self._prefetch_buffer.new_empty(
|
||||
(hidden_states.shape[0], self._output_dim)
|
||||
)
|
||||
self._finalize_prefetch(self._prefetch_buffer, output)
|
||||
return output
|
||||
@@ -2,534 +2,42 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Qwen4Exp n-gram embeddings with device and pinned-host storage."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Iterable
|
||||
from typing import ClassVar
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.distributed import get_dp_group, get_etp_group, get_tp_group
|
||||
from vllm.forward_context import DPMetadata, get_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.layers.quantization.modelopt import (
|
||||
ModelOptMixedPrecisionConfig,
|
||||
ModelOptQuantConfigBase,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
create_fp8_scale_parameter,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
is_layer_skipped,
|
||||
)
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader
|
||||
from vllm.model_executor.parameter import (
|
||||
ModelWeightParameter,
|
||||
PerTensorScaleParameter,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.transformers_utils.configs.qwen4_exp import (
|
||||
Qwen4ExpTextConfig,
|
||||
)
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.utils.platform_utils import is_uva_available
|
||||
from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
|
||||
|
||||
from ..common.ple import PLEVocabParallelEmbedding
|
||||
from ..common.ngram_embedding import (
|
||||
Qwen4ExpPLEDeviceEmbedding,
|
||||
Qwen4ExpPLEEmbedding,
|
||||
Qwen4ExpPLEEmbeddingMethod,
|
||||
Qwen4ExpPLEFp8EmbeddingMethod,
|
||||
Qwen4ExpPLEPinnedHostEmbedding,
|
||||
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
||||
)
|
||||
from .ops.ple import ple_ngram_ids
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Qwen4ExpPLEEmbedding(PLEVocabParallelEmbedding, ABC):
|
||||
"""ETP-sharded PLE table shared by device and pinned-host backends."""
|
||||
|
||||
supports_prefetch: ClassVar[bool] = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
*,
|
||||
params_dtype: torch.dtype,
|
||||
padding_size: int,
|
||||
prefix: str,
|
||||
embedding_method: "Qwen4ExpPLEEmbeddingMethod",
|
||||
num_ngram_heads: int = 1,
|
||||
max_total_tokens: int = 0,
|
||||
data_parallel_rank: int = 0,
|
||||
) -> None:
|
||||
del num_ngram_heads, max_total_tokens
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
params_dtype=params_dtype,
|
||||
padding_size=padding_size,
|
||||
prefix=prefix,
|
||||
quant_method=embedding_method,
|
||||
parallel_group=get_etp_group(),
|
||||
)
|
||||
self.embedding_method = embedding_method
|
||||
self.data_parallel_rank = data_parallel_rank
|
||||
tp_size = get_tp_group().world_size
|
||||
if self.tp_size % tp_size:
|
||||
raise ValueError(
|
||||
"ETP size must be divisible by TP size, but got "
|
||||
f"ETP={self.tp_size} and TP={tp_size}"
|
||||
)
|
||||
self.etp_data_parallel_size = self.tp_size // tp_size
|
||||
|
||||
@abstractmethod
|
||||
def allocate_embedding_weight(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate storage for the complete embedding weight."""
|
||||
raise NotImplementedError
|
||||
|
||||
def dequantize(
|
||||
self,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Delegate storage-format conversion to the embedding method."""
|
||||
return self.embedding_method.dequantize(self, embeddings, output_dtype)
|
||||
|
||||
def _get_dp_gather_slot(self, local_num_tokens: int) -> tuple[int, int]:
|
||||
"""Return the per-DP slot size and this rank's slot offset."""
|
||||
if self.etp_data_parallel_size == 1:
|
||||
return local_num_tokens, 0
|
||||
dp_metadata: DPMetadata | None = get_forward_context().dp_metadata
|
||||
if dp_metadata is None:
|
||||
raise RuntimeError("ETP spanning DP requires DP token metadata")
|
||||
group_start = (self.data_parallel_rank // self.etp_data_parallel_size) * (
|
||||
self.etp_data_parallel_size
|
||||
)
|
||||
group_end = group_start + self.etp_data_parallel_size
|
||||
token_counts = dp_metadata.num_tokens_across_dp_cpu.tolist()
|
||||
group_counts = token_counts[group_start:group_end]
|
||||
slot_size = max(group_counts)
|
||||
dp_rank = get_dp_group().rank_in_group
|
||||
return slot_size, dp_rank * slot_size
|
||||
|
||||
def _gather_dp_ids(
|
||||
self,
|
||||
ngram_ids: torch.Tensor,
|
||||
slot_size: int,
|
||||
) -> torch.Tensor:
|
||||
"""Gather DP-local IDs that share one ETP-sharded PLE table."""
|
||||
if self.etp_data_parallel_size == 1:
|
||||
return ngram_ids
|
||||
if ngram_ids.shape[0] < slot_size:
|
||||
padding = ngram_ids.new_zeros(
|
||||
slot_size - ngram_ids.shape[0], ngram_ids.shape[1]
|
||||
)
|
||||
ngram_ids = torch.cat((ngram_ids, padding), dim=0)
|
||||
return get_dp_group().all_gather(ngram_ids, dim=0)
|
||||
|
||||
def _select_embeddings(
|
||||
self,
|
||||
embeddings: torch.Tensor,
|
||||
local_num_tokens: int,
|
||||
slot_offset: int,
|
||||
) -> torch.Tensor:
|
||||
"""Select this DP rank's rows from the ETP-reduced embeddings."""
|
||||
if self.etp_data_parallel_size == 1:
|
||||
return embeddings
|
||||
return embeddings.narrow(0, slot_offset, local_num_tokens)
|
||||
|
||||
@abstractmethod
|
||||
def start_prefetch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
ngram_ids: torch.Tensor,
|
||||
) -> None:
|
||||
"""Start an asynchronous lookup when supported."""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class Qwen4ExpPLEEmbeddingMethod(QuantizeMethodBase):
|
||||
"""Quantization interface shared by resident and pinned PLE tables."""
|
||||
|
||||
# PLE post-load processing only validates scales in their current storage.
|
||||
requires_device_loading: bool = False
|
||||
|
||||
@staticmethod
|
||||
def from_quant_config(
|
||||
quant_config: QuantizationConfig | None,
|
||||
prefix: str,
|
||||
embedding_dtype: str | None = None,
|
||||
) -> "Qwen4ExpPLEEmbeddingMethod":
|
||||
"""Select the concrete PLE embedding format for a layer."""
|
||||
if embedding_dtype == "float8_e4m3fn":
|
||||
return Qwen4ExpPLEFp8EmbeddingMethod()
|
||||
if quant_config is None:
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if isinstance(quant_config, ModelOptMixedPrecisionConfig):
|
||||
if quant_config._resolve_quant_algo(prefix) == "FP8":
|
||||
return Qwen4ExpPLEFp8EmbeddingMethod()
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if isinstance(
|
||||
quant_config, ModelOptQuantConfigBase
|
||||
) and quant_config.is_layer_excluded(prefix):
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if not isinstance(quant_config, Fp8Config):
|
||||
raise NotImplementedError(
|
||||
"Qwen4Exp PLE embedding does not support quantization config "
|
||||
f"{type(quant_config).__name__}"
|
||||
)
|
||||
|
||||
ignored_layers = quant_config.ignored_layers
|
||||
if is_layer_skipped(
|
||||
prefix,
|
||||
ignored_layers,
|
||||
quant_config.packed_modules_mapping,
|
||||
match_mode=quant_config.ignored_layers_match_mode,
|
||||
):
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
# PLE checkpoint shards form one runtime embedding parameter.
|
||||
shard_prefix = f"{prefix}.shard_"
|
||||
if any(name.startswith(shard_prefix) for name in ignored_layers):
|
||||
return Qwen4ExpPLEUnquantizedEmbeddingMethod()
|
||||
if not quant_config.is_checkpoint_fp8_serialized:
|
||||
raise NotImplementedError(
|
||||
"Qwen4Exp PLE embedding only supports serialized FP8 checkpoints"
|
||||
)
|
||||
return Qwen4ExpPLEFp8EmbeddingMethod()
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError("PLE weights only support embedding lookup")
|
||||
|
||||
def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
|
||||
return F.embedding(input_, layer.weight)
|
||||
|
||||
@abstractmethod
|
||||
def dequantize(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Convert looked-up PLE rows to the activation dtype."""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class Qwen4ExpPLEUnquantizedEmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
|
||||
"""Unquantized PLE embedding storage and lookup semantics."""
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: Qwen4ExpPLEEmbedding,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
del input_size, output_size
|
||||
weight = nn.Parameter(
|
||||
layer.allocate_embedding_weight(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
|
||||
set_weight_attrs(weight, extra_weight_attrs)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
def dequantize(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
del layer, output_dtype
|
||||
return embeddings
|
||||
|
||||
|
||||
class Qwen4ExpPLEFp8EmbeddingMethod(Qwen4ExpPLEEmbeddingMethod):
|
||||
"""FP8 PLE embedding with one global checkpoint scale."""
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: Qwen4ExpPLEEmbedding,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
) -> None:
|
||||
del input_size, output_size, params_dtype
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
weight = ModelWeightParameter(
|
||||
data=layer.allocate_embedding_weight(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
torch.float8_e4m3fn,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
weight_scale = create_fp8_scale_parameter(
|
||||
PerTensorScaleParameter,
|
||||
output_partition_sizes,
|
||||
input_size_per_partition,
|
||||
None,
|
||||
weight_loader,
|
||||
scale_dtype=torch.float32,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
def process_weights_after_loading(self, layer: nn.Module) -> None:
|
||||
"""Reject FP8 PLE checkpoints without a global scale."""
|
||||
sentinel = torch.finfo(torch.float32).min
|
||||
if torch.any(layer.weight_scale == sentinel):
|
||||
raise ValueError("FP8 PLE checkpoint is missing its global scale")
|
||||
|
||||
def dequantize(
|
||||
self,
|
||||
layer: nn.Module,
|
||||
embeddings: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
weight_scale = getattr(layer, "weight_scale", None)
|
||||
if weight_scale is None:
|
||||
raise RuntimeError("FP8 PLE embedding is missing its global scale")
|
||||
if weight_scale.device != embeddings.device:
|
||||
raise RuntimeError("FP8 PLE embedding scale must be on the output device")
|
||||
return embeddings.to(output_dtype) * weight_scale.to(output_dtype)
|
||||
|
||||
|
||||
class Qwen4ExpPLEDeviceEmbedding(Qwen4ExpPLEEmbedding):
|
||||
"""PLE table allocated on the active model device."""
|
||||
|
||||
def allocate_embedding_weight(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate the complete PLE weight on the active device."""
|
||||
return torch.empty(num_embeddings, embedding_dim, dtype=dtype)
|
||||
|
||||
def start_prefetch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
ngram_ids: torch.Tensor,
|
||||
) -> None:
|
||||
"""Resident embedding prefetch is a no-op."""
|
||||
return None
|
||||
|
||||
def forward(self, ngram_ids: torch.Tensor) -> torch.Tensor:
|
||||
"""Gather ETP inputs, look up embeddings, and select local rows."""
|
||||
slot_size, slot_offset = self._get_dp_gather_slot(ngram_ids.shape[0])
|
||||
gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
|
||||
embeddings = super().forward(gathered_ids)
|
||||
return self._select_embeddings(
|
||||
embeddings,
|
||||
ngram_ids.shape[0],
|
||||
slot_offset,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _lookup_ple_embedding_from_pinned_kernel(
|
||||
weight_ptr,
|
||||
ids_ptr,
|
||||
output_ptr,
|
||||
embedding_dim,
|
||||
tp_vocab_start,
|
||||
tp_vocab_end,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
"""Look up TP-owned PLE rows through a CUDA view of pinned host memory."""
|
||||
row_id = tl.program_id(0)
|
||||
global_idx = tl.load(ids_ptr + row_id)
|
||||
in_range = (global_idx >= tp_vocab_start) & (global_idx < tp_vocab_end)
|
||||
local_idx = tl.where(in_range, global_idx - tp_vocab_start, 0)
|
||||
offsets = tl.arange(0, BLOCK_D)
|
||||
store_mask = offsets < embedding_dim
|
||||
load_mask = store_mask & in_range
|
||||
values = tl.load(
|
||||
weight_ptr + local_idx * embedding_dim + offsets,
|
||||
mask=load_mask,
|
||||
other=0.0,
|
||||
)
|
||||
tl.store(
|
||||
output_ptr + row_id * embedding_dim + offsets,
|
||||
values,
|
||||
mask=store_mask,
|
||||
)
|
||||
|
||||
|
||||
class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):
|
||||
"""PLE table loaded into pinned CPU memory and looked up through UVA."""
|
||||
|
||||
supports_prefetch: ClassVar[bool] = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
*,
|
||||
params_dtype: torch.dtype,
|
||||
padding_size: int,
|
||||
prefix: str,
|
||||
embedding_method: Qwen4ExpPLEEmbeddingMethod,
|
||||
num_ngram_heads: int = 1,
|
||||
max_total_tokens: int = 0,
|
||||
data_parallel_rank: int = 0,
|
||||
) -> None:
|
||||
if not is_uva_available():
|
||||
raise RuntimeError("Engram CPU offload requires UVA support")
|
||||
super().__init__(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
params_dtype=params_dtype,
|
||||
padding_size=padding_size,
|
||||
prefix=prefix,
|
||||
embedding_method=embedding_method,
|
||||
num_ngram_heads=num_ngram_heads,
|
||||
max_total_tokens=max_total_tokens,
|
||||
data_parallel_rank=data_parallel_rank,
|
||||
)
|
||||
self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight)
|
||||
self._block_d = triton.next_power_of_2(self.embedding_dim)
|
||||
self._prefetch_stream = torch.cuda.Stream(device=self._uva_weight.device)
|
||||
self._prefetch_buffer = torch.empty(
|
||||
max_total_tokens * self.etp_data_parallel_size,
|
||||
num_ngram_heads,
|
||||
self.embedding_dim,
|
||||
dtype=self.weight.dtype,
|
||||
device=self._uva_weight.device,
|
||||
)
|
||||
self._output_dim = num_ngram_heads * self.embedding_dim
|
||||
|
||||
def allocate_embedding_weight(
|
||||
self,
|
||||
num_embeddings: int,
|
||||
embedding_dim: int,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Allocate the complete PLE weight directly in pinned CPU memory."""
|
||||
return torch.empty(
|
||||
num_embeddings,
|
||||
embedding_dim,
|
||||
dtype=dtype,
|
||||
device="cpu",
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
def _lookup(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
output: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Look up local ETP rows while preserving the weight storage dtype."""
|
||||
expected_shape = (*input_ids.shape, self.embedding_dim)
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
expected_shape,
|
||||
dtype=self.weight.dtype,
|
||||
device=input_ids.device,
|
||||
)
|
||||
elif (
|
||||
tuple(output.shape) != expected_shape
|
||||
or output.dtype != self.weight.dtype
|
||||
or output.device != input_ids.device
|
||||
):
|
||||
raise ValueError(
|
||||
"PLE prefetch output must match the input shape, weight dtype, "
|
||||
"and input device"
|
||||
)
|
||||
|
||||
flat_ids = input_ids.reshape(-1).long()
|
||||
if flat_ids.numel():
|
||||
_lookup_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
|
||||
self._uva_weight,
|
||||
flat_ids,
|
||||
output,
|
||||
self.embedding_dim,
|
||||
self.shard_indices.org_vocab_start_index,
|
||||
self.shard_indices.org_vocab_end_index,
|
||||
BLOCK_D=self._block_d,
|
||||
)
|
||||
return output
|
||||
|
||||
def _reduce_etp_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
|
||||
"""Combine pinned lookup results owned by different ETP ranks."""
|
||||
if self.tp_size == 1:
|
||||
return embeddings
|
||||
assert self.parallel_group is not None
|
||||
if embeddings.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
||||
# Each vocabulary row has one owner, so reduce the raw FP8 bytes.
|
||||
reduced = self.parallel_group.all_reduce(embeddings.view(torch.int8))
|
||||
return reduced.view(embeddings.dtype)
|
||||
return self.parallel_group.all_reduce(embeddings)
|
||||
|
||||
@eager_break_during_capture
|
||||
def start_prefetch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
ngram_ids: torch.Tensor,
|
||||
) -> None:
|
||||
"""Gather ETP IDs and launch their UVA lookup on the side stream."""
|
||||
slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0])
|
||||
gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)
|
||||
active_output = self._prefetch_buffer[: gathered_ids.shape[0]]
|
||||
prefetch_stream = self._prefetch_stream
|
||||
prefetch_stream.wait_stream(torch.cuda.current_stream())
|
||||
gathered_ids.record_stream(prefetch_stream)
|
||||
with torch.cuda.stream(prefetch_stream):
|
||||
self._lookup(gathered_ids, output=active_output)
|
||||
|
||||
@eager_break_during_capture
|
||||
def _finalize_prefetch(
|
||||
self,
|
||||
prefetch_output: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
) -> None:
|
||||
"""Join the side stream, reduce ETP shards, and select local rows."""
|
||||
torch.cuda.current_stream().wait_stream(self._prefetch_stream)
|
||||
slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0])
|
||||
active_output = prefetch_output[: slot_size * self.etp_data_parallel_size]
|
||||
embeddings = self._reduce_etp_embeddings(active_output)
|
||||
embeddings = self._select_embeddings(
|
||||
embeddings,
|
||||
output.shape[0],
|
||||
slot_offset,
|
||||
)
|
||||
output.copy_(embeddings.flatten(-2))
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
"""Finish the pinned lookup into graph-owned output storage."""
|
||||
output = self._prefetch_buffer.new_empty(
|
||||
(hidden_states.shape[0], self._output_dim)
|
||||
)
|
||||
self._finalize_prefetch(self._prefetch_buffer, output)
|
||||
return output
|
||||
__all__ = [
|
||||
"Qwen4ExpPLEDeviceEmbedding",
|
||||
"Qwen4ExpPLEEmbedding",
|
||||
"Qwen4ExpPLEEmbeddingMethod",
|
||||
"Qwen4ExpPLEFp8EmbeddingMethod",
|
||||
"Qwen4ExpPLEPinnedHostEmbedding",
|
||||
"Qwen4ExpPLEUnquantizedEmbeddingMethod",
|
||||
"Qwen4ExpNGramEmbedding",
|
||||
]
|
||||
|
||||
|
||||
class Qwen4ExpNGramEmbedding(nn.Module):
|
||||
|
||||
Reference in New Issue
Block a user