diff --git a/CMakeLists.txt b/CMakeLists.txt index ed51927362..bf7a58e572 100644 --- a/CMakeLists.txt +++ b/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" diff --git a/benchmarks/kernels/benchmark_vocab_parallel_embedding.py b/benchmarks/kernels/benchmark_vocab_parallel_embedding.py new file mode 100644 index 0000000000..4925e6c024 --- /dev/null +++ b/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) diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 75f95e4bfd..3b79fc591f 100644 --- a/csrc/libtorch_stable/ops.h +++ b/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, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index ec1963b8a9..85af95e67e 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/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)); diff --git a/csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu b/csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu new file mode 100644 index 0000000000..863da3cd87 --- /dev/null +++ b/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 +#include + +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 +__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(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t vec_stride = static_cast(gridDim.x) * blockDim.x; + + for (int64_t token = blockIdx.y; token < num_tokens; token += gridDim.y) { + const int64_t id = static_cast(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 +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(sizeof(vec_t)); + const int grid_x = static_cast( + std::min((vecs_per_row + kBlockThreads - 1) / kBlockThreads, + static_cast(kBlockThreads))); + const int grid_y = static_cast(std::min( + num_tokens, static_cast(kTokenWaves) * + get_device_prop()->multiProcessorCount)); + + vocab_parallel_embedding_kernel + <<>>( + reinterpret_cast(out.mutable_data_ptr()), + reinterpret_cast(input_ids.const_data_ptr()), + reinterpret_cast(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 +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(out.mutable_data_ptr()) | + reinterpret_cast(weight.const_data_ptr()) | + static_cast(row_bytes); + +#define VLLM_LAUNCH_VOCAB_EMBEDDING(vec_t) \ + launch(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( + 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( + 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); + } +} diff --git a/tests/kernels/core/test_vocab_parallel_embedding.py b/tests/kernels/core/test_vocab_parallel_embedding.py new file mode 100644 index 0000000000..b0038c084e --- /dev/null +++ b/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) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index badadc162a..78a4c43159 100644 --- a/vllm/_custom_ops.py +++ b/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, diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 37241c8a01..70c6516e2e 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/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}"