[kernel] Fused embedding kernel (#53677)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

CMakeLists.txt

@@ -430,6 +430,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/cuda_utils_kernels.cu"
     "csrc/libtorch_stable/activation_kernels.cu"
     "csrc/libtorch_stable/ngram_embedding_kernels.cu"
+    "csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu"
     "csrc/libtorch_stable/quantization/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
     "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"

benchmarks/kernels/benchmark_vocab_parallel_embedding.py

@@ -0,0 +1,197 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark the TP > 1 input-embedding path of VocabParallelEmbedding.
+
+Two providers, both computing one rank's partial embedding (the all-reduce
+that follows is identical for each and is not measured):
+
+  unfused   mask + id shift + int64 cast + F.embedding + masked_fill_
+  fused     the single fused CUDA kernel, _C::vocab_parallel_embedding
+
+With --use-cuda-graph each provider is captured into a CUDA graph and the
+replay is timed, which is what a decode step actually costs.
+"""
+
+import itertools
+
+import torch
+
+from vllm import _custom_ops as vllm_ops
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+    get_masked_input_and_mask,
+    pad_vocab_size,
+)
+from vllm.triton_utils import triton
+
+
+def get_shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size):
+    return VocabParallelEmbedding._get_indices(
+        pad_vocab_size(vocab_size),
+        pad_vocab_size(org_vocab_size),
+        vocab_size,
+        org_vocab_size,
+        tp_rank,
+        tp_size,
+    )
+
+
+def embedding_reference(input_ids, weight, shard_indices):
+    masked_input, input_mask = get_masked_input_and_mask(
+        input_ids,
+        shard_indices.org_vocab_start_index,
+        shard_indices.org_vocab_end_index,
+        shard_indices.num_org_vocab_padding,
+        shard_indices.added_vocab_start_index,
+        shard_indices.added_vocab_end_index,
+    )
+    output = torch.nn.functional.embedding(masked_input.long(), weight)
+    return output.masked_fill_(input_mask.unsqueeze(-1), 0)
+
+
+def embedding_fused(input_ids, weight, shard_indices):
+    return vllm_ops.vocab_parallel_embedding(
+        input_ids,
+        weight,
+        shard_indices.org_vocab_start_index,
+        shard_indices.org_vocab_end_index,
+        shard_indices.num_org_vocab_padding,
+        shard_indices.added_vocab_start_index,
+        shard_indices.added_vocab_end_index,
+    )
+
+
+PROVIDERS = {
+    "unfused": embedding_reference,
+    "fused": embedding_fused,
+}
+
+
+def make_inputs(vocab_size, hidden_size, tp_size, num_tokens, dtype, tp_rank=0):
+    shard_indices = get_shard_indices(vocab_size, vocab_size, tp_rank, tp_size)
+    weight = torch.randn(
+        pad_vocab_size(vocab_size) // tp_size,
+        hidden_size,
+        dtype=dtype,
+        device="cuda",
+    )
+    input_ids = torch.randint(
+        0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda"
+    )
+    return input_ids, weight, shard_indices
+
+
+def graph_replay(fn):
+    """Capture fn into a CUDA graph and return its replay."""
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            fn()
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fn()
+    return graph.replay
+
+
+def calculate_diff(vocab_size, hidden_size, tp_size, num_tokens):
+    """The fused kernel must be bit-exact with the path it replaces, on every
+    rank of the shard."""
+    dtype = torch.bfloat16
+    for tp_rank in range(tp_size):
+        input_ids, weight, shard_indices = make_inputs(
+            vocab_size, hidden_size, tp_size, num_tokens, dtype, tp_rank
+        )
+        output_reference = embedding_reference(input_ids, weight, shard_indices)
+        output_fused = embedding_fused(input_ids, weight, shard_indices)
+        if not torch.equal(output_reference, output_fused):
+            print(f"❌ rank {tp_rank}: fused output differs")
+            return
+    print(f"✅ fused matches the reference path on all {tp_size} ranks")
+
+
+tp_size_range = [2, 4, 8]
+num_tokens_range = [1, 8, 64, 512, 2048, 8192, 32768]
+configs = list(itertools.product(tp_size_range, num_tokens_range))
+
+
+def get_benchmark(vocab_size, hidden_size, dtype, use_cuda_graph):
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["tp_size", "num_tokens"],
+            x_vals=[list(_) for _ in configs],
+            line_arg="provider",
+            line_vals=list(PROVIDERS),
+            line_names=["Unfused PyTorch", "Fused kernel"],
+            styles=[("green", "-"), ("red", "-")],
+            ylabel="us",
+            plot_name=(
+                f"vocab-parallel-embedding-perf-vocab{vocab_size}-"
+                f"hidden{hidden_size}{'-cudagraph' if use_cuda_graph else ''}"
+            ),
+            args={},
+        )
+    )
+    def benchmark(tp_size, num_tokens, provider):
+        input_ids, weight, shard_indices = make_inputs(
+            vocab_size, hidden_size, tp_size, num_tokens, dtype
+        )
+        fn = PROVIDERS[provider]
+
+        def run():
+            return fn(input_ids, weight, shard_indices)
+
+        if use_cuda_graph:
+            run = graph_replay(run)
+
+        ms, min_ms, max_ms = triton.testing.do_bench(run, quantiles=[0.5, 0.2, 0.8])
+        return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+    return benchmark
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--vocab-size", type=int, default=128256, help="Vocabulary size"
+    )
+    parser.add_argument(
+        "--hidden-size", type=int, default=4096, help="Embedding dimension"
+    )
+    parser.add_argument(
+        "--dtype", type=str, default="bfloat16", help="Embedding table dtype"
+    )
+    parser.add_argument(
+        "--tp-size",
+        type=int,
+        default=4,
+        help="TP size used for the correctness check",
+    )
+    parser.add_argument(
+        "--use-cuda-graph",
+        action="store_true",
+        help="Time CUDA graph replay instead of eager launches",
+    )
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default=None,
+        help="Optional path to save benchmark results",
+    )
+    args = parser.parse_args()
+    dtype = getattr(torch, args.dtype)
+
+    calculate_diff(
+        vocab_size=args.vocab_size,
+        hidden_size=args.hidden_size,
+        tp_size=args.tp_size,
+        num_tokens=1024,
+    )
+
+    benchmark = get_benchmark(
+        args.vocab_size, args.hidden_size, dtype, args.use_cuda_graph
+    )
+    benchmark.run(print_data=True, save_path=args.save_path)

csrc/libtorch_stable/ops.h

@@ -697,6 +697,14 @@ void cp_gather_indexer_k_quant_cache(
     const torch::stable::Tensor& block_table,   // [batch_size, num_blocks]
     const torch::stable::Tensor& cu_seq_lens);  // [batch_size + 1]
 
+// Fused vocab-parallel embedding lookup (see
+// vocab_parallel_embedding_kernels.cu).
+void vocab_parallel_embedding(
+    torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
+    const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
+    int64_t org_vocab_end_index, int64_t num_org_vocab_padding,
+    int64_t added_vocab_start_index, int64_t added_vocab_end_index);
+
 // LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu).
 void ngram_compute_n_gram_ids(
     int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights,

csrc/libtorch_stable/torch_bindings.cpp

@@ -700,6 +700,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // LongCat n-gram embedding index kernel. All tensor args are marked mutable
   // to match the (non-const) stable-Tensor& C++ signature; only ne_token_table
   // and n_gram_ids are actually written in place.
+  // Fused vocab-parallel embedding: gather the rows this rank owns and write
+  // zeros for the rest, so the following all-reduce reconstructs the full
+  // embedding.
+  ops.def(
+      "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, "
+      "int org_vocab_start_index, int org_vocab_end_index, "
+      "int num_org_vocab_padding, int added_vocab_start_index, "
+      "int added_vocab_end_index) -> ()");
   ops.def(
       "ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, "
       "Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, "
@@ -711,6 +719,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // LongCat n-gram embedding index kernel.
   ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids));
+  ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding));
 
   // Per-token group quantization
   ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));

csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu

@@ -0,0 +1,166 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Fused vocab-parallel embedding lookup.
+//
+// With TP > 1 each rank owns a slice of the vocabulary, so the eager path is a
+// range mask, an id shift, an int64 cast, a gather and a masked_fill_ of the
+// gathered rows. This kernel does all of it in one pass: rows owned by another
+// rank are written as zeros, which keeps the following all-reduce exact.
+//
+// The gather is a pure row copy, so the kernel is templated on a byte-vector
+// type instead of the value type and works for any embedding dtype.
+
+#include "torch_utils.h"
+
+#include "ops.h"
+
+#include <algorithm>
+#include <cstdint>
+
+namespace vllm::vocab_embedding {
+
+constexpr int kBlockThreads = 256;
+// Enough token-blocks to fill the device several times over; the token loop
+// strides over the rest.
+constexpr int kTokenWaves = 4;
+
+template <typename idx_t, typename vec_t>
+__global__ void vocab_parallel_embedding_kernel(
+    vec_t* __restrict__ out,              // [num_tokens, vecs_per_row]
+    const idx_t* __restrict__ input_ids,  // [num_tokens]
+    const vec_t* __restrict__ weight,     // [num_rows, vecs_per_row]
+    const int64_t num_tokens, const int64_t vecs_per_row,
+    const int64_t org_vocab_start_index, const int64_t org_vocab_end_index,
+    const int64_t added_vocab_start_index, const int64_t added_vocab_end_index,
+    const int64_t added_row_offset) {
+  const int64_t vec_begin =
+      static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int64_t vec_stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
+
+  for (int64_t token = blockIdx.y; token < num_tokens; token += gridDim.y) {
+    const int64_t id = static_cast<int64_t>(input_ids[token]);
+    const bool in_org = id >= org_vocab_start_index && id < org_vocab_end_index;
+    const bool in_added =
+        id >= added_vocab_start_index && id < added_vocab_end_index;
+    vec_t* out_row = out + token * vecs_per_row;
+
+    if (in_org || in_added) {
+      const int64_t row =
+          in_org ? id - org_vocab_start_index : id + added_row_offset;
+      const vec_t* weight_row = weight + row * vecs_per_row;
+      for (int64_t v = vec_begin; v < vecs_per_row; v += vec_stride) {
+        out_row[v] = weight_row[v];
+      }
+    } else {
+      // Token belongs to another rank's shard: contribute zeros.
+      const vec_t zero = vec_t();
+      for (int64_t v = vec_begin; v < vecs_per_row; v += vec_stride) {
+        out_row[v] = zero;
+      }
+    }
+  }
+}
+
+template <typename idx_t, typename vec_t>
+void launch(torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
+            const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
+            int64_t org_vocab_end_index, int64_t added_vocab_start_index,
+            int64_t added_vocab_end_index, int64_t added_row_offset,
+            int64_t num_tokens, int64_t row_bytes, cudaStream_t stream) {
+  const int64_t vecs_per_row = row_bytes / static_cast<int64_t>(sizeof(vec_t));
+  const int grid_x = static_cast<int>(
+      std::min<int64_t>((vecs_per_row + kBlockThreads - 1) / kBlockThreads,
+                        static_cast<int64_t>(kBlockThreads)));
+  const int grid_y = static_cast<int>(std::min<int64_t>(
+      num_tokens, static_cast<int64_t>(kTokenWaves) *
+                      get_device_prop()->multiProcessorCount));
+
+  vocab_parallel_embedding_kernel<idx_t, vec_t>
+      <<<dim3(grid_x, grid_y), kBlockThreads, 0, stream>>>(
+          reinterpret_cast<vec_t*>(out.mutable_data_ptr()),
+          reinterpret_cast<const idx_t*>(input_ids.const_data_ptr()),
+          reinterpret_cast<const vec_t*>(weight.const_data_ptr()), num_tokens,
+          vecs_per_row, org_vocab_start_index, org_vocab_end_index,
+          added_vocab_start_index, added_vocab_end_index, added_row_offset);
+}
+
+template <typename idx_t>
+void dispatch_vec_width(
+    torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
+    const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
+    int64_t org_vocab_end_index, int64_t added_vocab_start_index,
+    int64_t added_vocab_end_index, int64_t added_row_offset, int64_t num_tokens,
+    int64_t row_bytes, cudaStream_t stream) {
+  const uintptr_t addrs = reinterpret_cast<uintptr_t>(out.mutable_data_ptr()) |
+                          reinterpret_cast<uintptr_t>(weight.const_data_ptr()) |
+                          static_cast<uintptr_t>(row_bytes);
+
+#define VLLM_LAUNCH_VOCAB_EMBEDDING(vec_t)                                  \
+  launch<idx_t, vec_t>(out, input_ids, weight, org_vocab_start_index,       \
+                       org_vocab_end_index, added_vocab_start_index,        \
+                       added_vocab_end_index, added_row_offset, num_tokens, \
+                       row_bytes, stream)
+
+  if (addrs % 16 == 0) {
+    VLLM_LAUNCH_VOCAB_EMBEDDING(uint4);
+  } else if (addrs % 8 == 0) {
+    VLLM_LAUNCH_VOCAB_EMBEDDING(uint2);
+  } else if (addrs % 4 == 0) {
+    VLLM_LAUNCH_VOCAB_EMBEDDING(uint32_t);
+  } else if (addrs % 2 == 0) {
+    VLLM_LAUNCH_VOCAB_EMBEDDING(uint16_t);
+  } else {
+    VLLM_LAUNCH_VOCAB_EMBEDDING(uint8_t);
+  }
+#undef VLLM_LAUNCH_VOCAB_EMBEDDING
+}
+
+}  // namespace vllm::vocab_embedding
+
+void vocab_parallel_embedding(
+    torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
+    const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
+    int64_t org_vocab_end_index, int64_t num_org_vocab_padding,
+    int64_t added_vocab_start_index, int64_t added_vocab_end_index) {
+  STD_TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
+  STD_TORCH_CHECK(out.dim() == 2, "out must be 2D");
+  STD_TORCH_CHECK(out.size(1) == weight.size(1),
+                  "out and weight must share the embedding dim");
+  STD_TORCH_CHECK(out.size(0) == input_ids.numel(),
+                  "out must have one row per input id");
+  STD_TORCH_CHECK(out.scalar_type() == weight.scalar_type(),
+                  "out and weight must have the same dtype");
+  STD_TORCH_CHECK(input_ids.is_contiguous() && weight.is_contiguous() &&
+                      out.is_contiguous(),
+                  "vocab_parallel_embedding requires contiguous tensors");
+
+  const int64_t num_tokens = input_ids.numel();
+  if (num_tokens == 0) return;
+
+  const int64_t row_bytes = weight.size(1) * weight.element_size();
+  // Offset that maps an added-vocab (LoRA) id onto its local row, mirroring
+  // get_masked_input_and_mask().
+  const int64_t added_row_offset = org_vocab_end_index - org_vocab_start_index +
+                                   num_org_vocab_padding -
+                                   added_vocab_start_index;
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      out.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+
+  if (input_ids.scalar_type() == torch::headeronly::ScalarType::Long) {
+    vllm::vocab_embedding::dispatch_vec_width<int64_t>(
+        out, input_ids, weight, org_vocab_start_index, org_vocab_end_index,
+        added_vocab_start_index, added_vocab_end_index, added_row_offset,
+        num_tokens, row_bytes, stream);
+  } else {
+    STD_TORCH_CHECK(
+        input_ids.scalar_type() == torch::headeronly::ScalarType::Int,
+        "input_ids must be int32 or int64");
+    vllm::vocab_embedding::dispatch_vec_width<int32_t>(
+        out, input_ids, weight, org_vocab_start_index, org_vocab_end_index,
+        added_vocab_start_index, added_vocab_end_index, added_row_offset,
+        num_tokens, row_bytes, stream);
+  }
+}

tests/kernels/core/test_vocab_parallel_embedding.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the fused vocab-parallel embedding kernel
+(``vllm._custom_ops.vocab_parallel_embedding``).
+
+The guarantee: the kernel is a drop-in for the TP > 1 embedding path, i.e. it
+returns bit-exactly what mask + shift + gather + ``masked_fill_`` returns on
+each rank, and summing the per-rank outputs (what the all-reduce does)
+reconstructs a full-table lookup.
+"""
+
+import pytest
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+    get_masked_input_and_mask,
+    pad_vocab_size,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+requires_cuda = pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="the fused vocab-parallel embedding kernel is CUDA-only",
+)
+
+
+def _shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size, pad=64):
+    return VocabParallelEmbedding._get_indices(
+        pad_vocab_size(vocab_size, pad),
+        pad_vocab_size(org_vocab_size, pad),
+        vocab_size,
+        org_vocab_size,
+        tp_rank,
+        tp_size,
+    )
+
+
+def _reference(input_ids, weight, si):
+    """The path the kernel replaces."""
+    masked_input, input_mask = get_masked_input_and_mask(
+        input_ids,
+        si.org_vocab_start_index,
+        si.org_vocab_end_index,
+        si.num_org_vocab_padding,
+        si.added_vocab_start_index,
+        si.added_vocab_end_index,
+    )
+    out = torch.nn.functional.embedding(masked_input.long(), weight)
+    return out.masked_fill_(input_mask.unsqueeze(-1), 0)
+
+
+def _fused(input_ids, weight, si):
+    return ops.vocab_parallel_embedding(
+        input_ids,
+        weight,
+        si.org_vocab_start_index,
+        si.org_vocab_end_index,
+        si.num_org_vocab_padding,
+        si.added_vocab_start_index,
+        si.added_vocab_end_index,
+    )
+
+
+def _shard(full_table, si, rows_per_rank, dtype):
+    """Lay a rank's rows out the way the weight loader does."""
+    weight = torch.zeros(rows_per_rank, full_table.shape[1], dtype=dtype, device="cuda")
+    num_org = si.org_vocab_end_index - si.org_vocab_start_index
+    weight[:num_org] = full_table[si.org_vocab_start_index : si.org_vocab_end_index]
+    added_at = num_org + si.num_org_vocab_padding
+    num_added = si.added_vocab_end_index - si.added_vocab_start_index
+    weight[added_at : added_at + num_added] = full_table[
+        si.added_vocab_start_index : si.added_vocab_end_index
+    ]
+    return weight
+
+
+@requires_cuda
+@torch.inference_mode()
+@pytest.mark.parametrize("tp_size", [2, 4, 8])
+@pytest.mark.parametrize(
+    "vocab_size,org_vocab_size", [(128256, 128256), (32256, 32000)]
+)
+@pytest.mark.parametrize("hidden", [4096, 3584])
+@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64])
+def test_matches_masked_gather(tp_size, vocab_size, org_vocab_size, hidden, id_dtype):
+    """Every rank's partial output is bit-exact with the eager path, including
+    the added-vocab (LoRA) rows and the ids owned by another rank."""
+    set_random_seed(7)
+    dtype = torch.bfloat16
+    rows = pad_vocab_size(vocab_size, 64) // tp_size
+    full_table = torch.randn(pad_vocab_size(vocab_size, 64), hidden, dtype=dtype).cuda()
+    ids = torch.randint(0, vocab_size, (257,), dtype=id_dtype, device="cuda")
+
+    for tp_rank in range(tp_size):
+        si = _shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size)
+        weight = _shard(full_table, si, rows, dtype)
+        torch.testing.assert_close(
+            _fused(ids, weight, si), _reference(ids, weight, si), atol=0.0, rtol=0.0
+        )
+
+
+@requires_cuda
+@torch.inference_mode()
+@pytest.mark.parametrize("tp_size", [2, 8])
+def test_rank_sum_is_full_lookup(tp_size):
+    """What the all-reduce sees: the partials add up to a full-table lookup."""
+    set_random_seed(11)
+    vocab_size, org_vocab_size, hidden = 32256, 32000, 2048
+    dtype = torch.bfloat16
+    rows = pad_vocab_size(vocab_size, 64) // tp_size
+    full_table = torch.randn(pad_vocab_size(vocab_size, 64), hidden, dtype=dtype).cuda()
+    # Ids from both the base and the added (LoRA) range.
+    ids = torch.randint(0, vocab_size, (129,), dtype=torch.int32, device="cuda")
+    ids[:16] = torch.randint(
+        org_vocab_size, vocab_size, (16,), dtype=torch.int32, device="cuda"
+    )
+
+    total = torch.zeros(ids.shape[0], hidden, dtype=dtype, device="cuda")
+    for tp_rank in range(tp_size):
+        si = _shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size)
+        total += _fused(ids, _shard(full_table, si, rows, dtype), si)
+
+    torch.testing.assert_close(total, full_table[ids.long()], atol=0.0, rtol=0.0)

vllm/_custom_ops.py

@@ -260,6 +260,49 @@ def rms_norm(
     torch.ops._C.rms_norm(out, input, weight, epsilon)
 
 
+# Fused vocab-parallel embedding lookup
+# (see csrc/.../vocab_parallel_embedding_kernels.cu).
+def vocab_parallel_embedding(
+    input_ids: torch.Tensor,
+    weight: torch.Tensor,
+    org_vocab_start_index: int,
+    org_vocab_end_index: int,
+    num_org_vocab_padding: int,
+    added_vocab_start_index: int,
+    added_vocab_end_index: int,
+) -> torch.Tensor:
+    """Gather the embedding rows owned by this rank, zeros for the others.
+
+    Fuses the range mask, id shift, gather and output masking of the TP > 1
+    embedding path into a single kernel. Out-of-shard tokens produce zero rows,
+    so an all-reduce over the ranks yields the full embedding.
+
+    Args:
+        input_ids: ``[num_tokens]`` int32 or int64 token ids.
+        weight: ``[num_embeddings_per_partition, embedding_dim]`` shard.
+
+    Returns:
+        ``[num_tokens, embedding_dim]`` partial embeddings.
+    """
+    out = torch.empty(
+        input_ids.shape[0],
+        weight.shape[1],
+        dtype=weight.dtype,
+        device=weight.device,
+    )
+    torch.ops._C.vocab_parallel_embedding(
+        out,
+        input_ids,
+        weight,
+        org_vocab_start_index,
+        org_vocab_end_index,
+        num_org_vocab_padding,
+        added_vocab_start_index,
+        added_vocab_end_index,
+    )
+    return out
+
+
 # LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu).
 def ngram_compute_n_gram_ids(
     ne_n: int,

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -9,6 +9,7 @@ import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
 import vllm.envs as envs
+from vllm import _custom_ops as ops
 from vllm.distributed import (
     divide,
     get_tensor_model_parallel_rank,
@@ -325,6 +326,14 @@ class VocabParallelEmbedding(PluggableLayer):
             - self.shard_indices.added_vocab_start_index
         )
 
+        # The fused kernel masks, shifts and gathers in one pass; it assumes a
+        # plain row-major shard, so it only applies to unquantized weights.
+        self.use_fused_embedding = (
+            self.tp_size > 1
+            and current_platform.is_cuda()
+            and isinstance(quant_method, UnquantizedEmbeddingMethod)
+        )
+
         self.quant_method.create_weights(
             self,
             self.embedding_dim,
@@ -488,8 +497,24 @@ class VocabParallelEmbedding(PluggableLayer):
         param[loaded_weight.shape[0] :].data.fill_(0)
 
     def forward(self, input_):
-        if self.tp_size > 1:
-            # Build the mask.
+        if self.tp_size == 1:
+            return self.quant_method.embedding(self, input_.long())
+
+        if self.use_fused_embedding:
+            output_parallel = ops.vocab_parallel_embedding(
+                input_ if input_.ndim == 1 else input_.reshape(-1),
+                self.weight,
+                self.shard_indices.org_vocab_start_index,
+                self.shard_indices.org_vocab_end_index,
+                self.shard_indices.num_org_vocab_padding,
+                self.shard_indices.added_vocab_start_index,
+                self.shard_indices.added_vocab_end_index,
+            )
+            if input_.ndim != 1:
+                output_parallel = output_parallel.view(
+                    *input_.shape, self.embedding_dim
+                )
+        else:
             masked_input, input_mask = get_masked_input_and_mask(
                 input_,
                 self.shard_indices.org_vocab_start_index,
@@ -498,12 +523,7 @@ class VocabParallelEmbedding(PluggableLayer):
                 self.shard_indices.added_vocab_start_index,
                 self.shard_indices.added_vocab_end_index,
             )
-        else:
-            masked_input = input_
-        # Get the embeddings.
-        output_parallel = self.quant_method.embedding(self, masked_input.long())
-        # Mask the output embedding.
-        if self.tp_size > 1:
+            output_parallel = self.quant_method.embedding(self, masked_input.long())
             if output_parallel.dtype in (
                 torch.float8_e4m3fn,
                 torch.float8_e5m2,
@@ -514,8 +534,8 @@ class VocabParallelEmbedding(PluggableLayer):
                 output = tensor_model_parallel_all_reduce(comm_output)
                 return output.view(output_parallel.dtype)
             output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
-            return tensor_model_parallel_all_reduce(output_parallel)
-        return output_parallel
+        # Reduce across all the model parallel GPUs.
+        return tensor_model_parallel_all_reduce(output_parallel)
 
     def extra_repr(self) -> str:
         s = f"num_embeddings={self.num_embeddings}"