[hw-agnostic] Adding embedding layer and lm_head

tests/models/transformers/test_backend.py

@@ -400,6 +400,23 @@ def test_replace_plain_embedding(vpe):
     assert type(replace(nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE))) is vpe
 
 
+def test_replace_plain_embedding_hw_agnostic(monkeypatch, tp_init):
+    """With `VLLM_USE_HW_AGNOSTIC=1` the plain path builds the hw-agnostic class,
+    the same one `CausalMixin` resolves for its tie-weights check.
+
+    The MRO-rebasing path stays on vLLM's class (covered by the inherited/nested
+    tests above), so only the plain path is asserted here.
+    """
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
+    from vllm.model_executor.models.transformers.layers import (
+        get_vocab_parallel_embedding_cls,
+    )
+
+    hw_vpe = get_vocab_parallel_embedding_cls()
+    assert "hw_agnostic.layers" in hw_vpe.__module__
+    assert type(replace(nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE))) is hw_vpe
+
+
 def test_replace_infers_shape_and_dtype(tp_init):
     """Shape and dtype come from the replaced module, not from the config."""
     embedding = nn.Embedding(VOCAB_SIZE * 2, HIDDEN_SIZE + 1, dtype=torch.float16)

tests/models/transformers/test_layer_registry.py

@@ -9,6 +9,7 @@
 logging that reports which source was used.
 """
 
+import importlib
 import logging
 import sys
 import types
@@ -78,8 +79,62 @@ def test_act_and_mul_falls_back_for_unknown_activation(
     assert isinstance(layers.get_act_and_mul_fn("gelu"), GeluAndMul)
 
 
-@pytest.fixture(scope="module")
-def tiny_llama_path(tmp_path_factory):
+# Each getter and the module/class name it resolves between the two trees.
+_CLASS_GETTERS = (
+    (
+        "get_vocab_parallel_embedding_cls",
+        "vocab_parallel_embedding",
+        "VocabParallelEmbedding",
+    ),
+    ("get_parallel_lm_head_cls", "vocab_parallel_embedding", "ParallelLMHead"),
+    ("get_logits_processor_cls", "logits_processor", "LogitsProcessor"),
+)
+
+
+@pytest.mark.parametrize("getter,module,name", _CLASS_GETTERS)
+def test_class_getter_falls_back_when_disabled(monkeypatch, getter, module, name):
+    """Disabled: each getter returns the vLLM class."""
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0")
+    vllm_cls = getattr(
+        importlib.import_module(f"vllm.model_executor.layers.{module}"), name
+    )
+    assert getattr(layers, getter)() is vllm_cls
+
+
+@pytest.mark.parametrize("getter,module,name", _CLASS_GETTERS)
+def test_class_getter_uses_hw_agnostic_when_enabled(
+    monkeypatch, caplog, getter, module, name
+):
+    """Enabled: each getter returns the hw-agnostic class and logs it."""
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
+    hw_cls = getattr(
+        importlib.import_module(f"vllm.model_executor.hw_agnostic.layers.{module}"),
+        name,
+    )
+    with caplog.at_level(logging.INFO):
+        resolved = getattr(layers, getter)()
+    assert resolved is hw_cls
+    assert f"Using hw-agnostic layer: {name}" in caplog.text
+
+
+@pytest.mark.parametrize("getter,module,name", _CLASS_GETTERS)
+def test_class_getter_falls_back_when_symbol_missing(
+    monkeypatch, caplog, getter, module, name
+):
+    """Enabled but the symbol is not ported: fall back to vLLM and warn."""
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
+    hw_module = f"vllm.model_executor.hw_agnostic.layers.{module}"
+    monkeypatch.setitem(sys.modules, hw_module, types.ModuleType(hw_module))
+    vllm_cls = getattr(
+        importlib.import_module(f"vllm.model_executor.layers.{module}"), name
+    )
+    with caplog.at_level(logging.WARNING):
+        resolved = getattr(layers, getter)()
+    assert resolved is vllm_cls
+    assert "falling back to default" in caplog.text
+
+
+def _save_tiny_llama(tmp_path_factory, name: str, *, tie_word_embeddings: bool) -> str:
     """A randomly-initialized microscopic Llama saved to disk (with an ungated
     tokenizer) so vLLM can load it like any local checkpoint."""
     from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
@@ -93,19 +148,40 @@ def tiny_llama_path(tmp_path_factory):
         num_attention_heads=4,
         rms_norm_eps=1e-6,
         hidden_act="silu",
+        tie_word_embeddings=tie_word_embeddings,
     )
     torch.manual_seed(0)
     model = LlamaForCausalLM(config)
 
-    path = tmp_path_factory.mktemp("tiny_llama")
+    path = tmp_path_factory.mktemp(name)
     model.save_pretrained(path)
     tokenizer.save_pretrained(path)
     return str(path)
 
 
+@pytest.fixture(scope="module")
+def tiny_llama_path(tmp_path_factory):
+    """A tiny Llama with an untied `lm_head`."""
+    return _save_tiny_llama(tmp_path_factory, "tiny_llama", tie_word_embeddings=False)
+
+
+@pytest.fixture(scope="module")
+def tiny_llama_tied_path(tmp_path_factory):
+    """A tiny Llama whose `lm_head` is tied to the input embedding."""
+    return _save_tiny_llama(
+        tmp_path_factory, "tiny_llama_tied", tie_word_embeddings=True
+    )
+
+
 # Registered names of the layers the backend can
 # currently route to hw-agnostic implementations.
-_COVERED_LAYERS = ("rms_norm", "silu_and_mul")
+_COVERED_LAYERS = (
+    "rms_norm",
+    "silu_and_mul",
+    "vocab_parallel_embedding",
+    "parallel_lm_head",
+    "logits_processor",
+)
 
 
 def _layer_providers(model) -> dict[str, str]:
@@ -162,12 +238,42 @@ def test_hw_agnostic_matches_vllm_end_to_end(monkeypatch, vllm_runner, tiny_llam
 
     monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0")
     vllm_providers, vllm_outputs = _serve(vllm_runner, tiny_llama_path, prompts)
-    # Both replaceable layers present in a Llama block must be vLLM's here.
-    assert vllm_providers == {"rms_norm": "vllm", "silu_and_mul": "vllm"}
+    # Every replaceable layer present in the model must be vLLM's here.
+    assert vllm_providers == dict.fromkeys(_COVERED_LAYERS, "vllm")
 
     monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
     hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_path, prompts)
-    assert hw_providers == {"rms_norm": "hw_agnostic", "silu_and_mul": "hw_agnostic"}
+    assert hw_providers == dict.fromkeys(_COVERED_LAYERS, "hw_agnostic")
+
+    check_logprobs_close(
+        outputs_0_lst=vllm_outputs,
+        outputs_1_lst=hw_outputs,
+        name_0="vllm",
+        name_1="hw_agnostic",
+    )
+
+
+def test_hw_agnostic_matches_vllm_with_tied_lm_head(
+    monkeypatch, vllm_runner, tiny_llama_tied_path
+):
+    """Tied `lm_head`: the hw-agnostic embedding and head still match vLLM.
+
+    Exercises `ParallelLMHead.tie_weights` across the hw-agnostic classes and the
+    `isinstance` check that decides whether to tie; a class mismatch there would
+    silently drop the tie, so this guards it end to end.
+    """
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    from ..utils import check_logprobs_close
+
+    prompts = ["The capital of France is", "vLLM is"]
+
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0")
+    _, vllm_outputs = _serve(vllm_runner, tiny_llama_tied_path, prompts)
+
+    monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
+    hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_tied_path, prompts)
+    assert hw_providers == dict.fromkeys(_COVERED_LAYERS, "hw_agnostic")
 
     check_logprobs_close(
         outputs_0_lst=vllm_outputs,

vllm/model_executor/hw_agnostic/layers/logits_processor.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+import torch.nn.functional as F
+
+from vllm.config import get_current_vllm_config
+from vllm.distributed import (
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_gather,
+)
+from vllm.model_executor.hw_agnostic.custom_op import PluggableLayer
+from vllm.model_executor.hw_agnostic.layers.vocab_parallel_embedding import (
+    UnquantizedEmbeddingMethod,
+    VocabParallelEmbedding,
+)
+from vllm.platforms import current_platform
+
+
+@PluggableLayer.register("logits_processor")
+class LogitsProcessor(PluggableLayer):
+    """Process logits and apply logits processors from sampling metadata.
+
+    1. Gather logits from model hidden_states.
+    2. Scale logits if needed.
+    3. Apply logits processors (if any).
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        org_vocab_size: int | None = None,
+        scale: float = 1.0,
+        logits_as_input: bool = False,
+        soft_cap: float | None = None,
+    ) -> None:
+        super().__init__()
+        self.scale = scale
+        self.vocab_size = vocab_size
+        self.logits_as_input = logits_as_input
+        self.org_vocab_size = org_vocab_size or vocab_size
+        # Soft cap the logits. Used in Gemma 2.
+        self.soft_cap = soft_cap
+        self.use_all_gather = current_platform.use_all_gather()
+        # Dtype of the lm_head projection; an fp32 head (via
+        # `--hf-overrides '{"head_dtype": "float32"}'`) is required for
+        # RL training-inference consistency. Defaults to the model dtype.
+        model_config = get_current_vllm_config().model_config
+        self.head_dtype = model_config.head_dtype if model_config is not None else None
+
+    def forward(
+        self,
+        lm_head: VocabParallelEmbedding,
+        hidden_states: torch.Tensor,
+        embedding_bias: torch.Tensor | None = None,
+    ) -> torch.Tensor | None:
+        if self.logits_as_input:
+            logits = hidden_states
+        else:
+            logits = self._get_logits(hidden_states, lm_head, embedding_bias)
+        if logits is not None:
+            if self.soft_cap is not None:
+                logits = logits / self.soft_cap
+                logits = torch.tanh(logits)
+                logits = logits * self.soft_cap
+            if self.scale != 1.0:
+                logits *= self.scale
+        return logits
+
+    def _gather_logits(self, logits: torch.Tensor) -> torch.Tensor:
+        if self.use_all_gather:
+            # Gather isn't supported for some devices (e.g. TPUs); use
+            # all-gather to keep all ranks in lockstep.
+            logits = tensor_model_parallel_all_gather(logits)
+        else:
+            # None may be returned for rank > 0.
+            logits = tensor_model_parallel_gather(logits)
+        return logits
+
+    def _apply_head(
+        self,
+        lm_head: VocabParallelEmbedding,
+        hidden_states: torch.Tensor,
+        embedding_bias: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """Project hidden states through the lm_head, honoring head_dtype."""
+        if self.head_dtype is None or self.head_dtype == hidden_states.dtype:
+            return lm_head.quant_method.apply(
+                lm_head, hidden_states, bias=embedding_bias
+            )
+
+        if not isinstance(lm_head.quant_method, UnquantizedEmbeddingMethod):
+            raise ValueError(
+                "A head_dtype different from the model dtype is only "
+                "supported for an unquantized lm_head."
+            )
+        if (
+            self.head_dtype == torch.float32
+            and (current_platform.is_cuda() or current_platform.is_rocm())
+            and hidden_states.is_cuda
+        ):
+            # Accumulate directly into fp32 to avoid materializing an fp32 copy
+            # of the lm_head weight each step. `torch.mm(out_dtype=...)` supports
+            # fp32 output for fp16/bf16 inputs only on CUDA and ROCm; other
+            # platforms fall back to the cast path below.
+            flat = hidden_states.reshape(-1, hidden_states.shape[-1])
+            logits = torch.mm(flat, lm_head.weight.t(), out_dtype=self.head_dtype)
+            if embedding_bias is not None:
+                logits = logits + embedding_bias.to(self.head_dtype)
+            return logits.reshape(*hidden_states.shape[:-1], -1)
+        return F.linear(
+            hidden_states.to(self.head_dtype),
+            lm_head.weight.to(self.head_dtype),
+            embedding_bias.to(self.head_dtype) if embedding_bias is not None else None,
+        )
+
+    def _get_logits(
+        self,
+        hidden_states: torch.Tensor,
+        lm_head: VocabParallelEmbedding,
+        embedding_bias: torch.Tensor | None,
+    ) -> torch.Tensor | None:
+        logits = self._apply_head(lm_head, hidden_states, embedding_bias)
+        if lm_head.tp_size > 1:
+            logits = self._gather_logits(logits)
+        if logits is not None:
+            logits = logits[..., : self.org_vocab_size]
+        return logits
+
+    def get_top_tokens(
+        self,
+        lm_head: VocabParallelEmbedding,
+        hidden_states: torch.Tensor,
+        embedding_bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Vocab-parallel argmax without all-gathering full logits.
+
+        Each TP rank computes local argmax, then only the (value, index) pairs
+        are gathered and reduced: O(batch * 2 * tp_size) vs O(batch * vocab_size).
+        """
+        if self.scale <= 0.0 and self.scale != 1.0:
+            raise ValueError(
+                "The local argmax reduction optimization is not supported for "
+                "non-positive logit scaling factors."
+            )
+        tp_size = lm_head.tp_size
+
+        logits = self._apply_head(lm_head, hidden_states, embedding_bias)
+        if self.soft_cap is not None:
+            logits = torch.tanh(logits / self.soft_cap) * self.soft_cap
+        if self.scale != 1.0:
+            logits = logits * self.scale
+
+        num_pad = lm_head.shard_indices.num_org_vocab_padding
+        if num_pad > 0:
+            logits[..., -num_pad:] = -float("inf")
+
+        local_max_vals, local_max_indices = logits.max(dim=-1)
+        vocab_start = lm_head.shard_indices.org_vocab_start_index
+        global_indices = local_max_indices + vocab_start
+
+        if tp_size == 1:
+            return global_indices
+
+        # All-gather (value, index) pairs, then reduce to global argmax.
+        # float32 avoids bf16 precision loss on large vocab indices.
+        local_pair = torch.stack(
+            [local_max_vals.float(), global_indices.float()], dim=-1
+        )
+        gathered = tensor_model_parallel_all_gather(local_pair, dim=-1)
+        gathered = gathered.view(hidden_states.shape[0], tp_size, 2)
+        max_rank_idx = gathered[:, :, 0].argmax(dim=-1, keepdim=True)
+        top_tokens = gathered[:, :, 1].gather(dim=-1, index=max_rank_idx)
+        return top_tokens.squeeze(-1).to(torch.int64)
+
+    def extra_repr(self) -> str:
+        s = f"vocab_size={self.vocab_size}"
+        s += f", org_vocab_size={self.org_vocab_size}"
+        s += f", scale={self.scale}, logits_as_input={self.logits_as_input}"
+        return s