[kernel] Fused embedding kernel (#53677)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Aug 31, 2026 at 19:34 UTC
e2c8eeac40d01051d6d60b8f42ebbe5842a1c672
8 files changed
+580
-10
CMakeLists.txt
+1
@@ -430,6 +430,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
430
"csrc/libtorch_stable/cuda_utils_kernels.cu"
431
"csrc/libtorch_stable/activation_kernels.cu"
432
"csrc/libtorch_stable/ngram_embedding_kernels.cu"
433
+ "csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu"
434
"csrc/libtorch_stable/quantization/activation_kernels.cu"
435
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
436
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
benchmarks/kernels/benchmark_vocab_parallel_embedding.py
new
+197
@@ -0,0 +1,197 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Benchmark the TP > 1 input-embedding path of VocabParallelEmbedding.
4
+
5
+Two providers, both computing one rank's partial embedding (the all-reduce
6
+that follows is identical for each and is not measured):
7
+
8
+ unfused mask + id shift + int64 cast + F.embedding + masked_fill_
9
+ fused the single fused CUDA kernel, _C::vocab_parallel_embedding
10
+
11
+With --use-cuda-graph each provider is captured into a CUDA graph and the
12
+replay is timed, which is what a decode step actually costs.
13
+"""
14
+
15
+import itertools
16
+
17
+import torch
18
+
19
+from vllm import _custom_ops as vllm_ops
20
+from vllm.model_executor.layers.vocab_parallel_embedding import (
21
+ VocabParallelEmbedding,
22
+ get_masked_input_and_mask,
23
+ pad_vocab_size,
24
+)
25
+from vllm.triton_utils import triton
26
+
27
+
28
+def get_shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size):
29
+ return VocabParallelEmbedding._get_indices(
30
+ pad_vocab_size(vocab_size),
31
+ pad_vocab_size(org_vocab_size),
32
+ vocab_size,
33
+ org_vocab_size,
34
+ tp_rank,
35
+ tp_size,
36
+ )
37
+
38
+
39
+def embedding_reference(input_ids, weight, shard_indices):
40
+ masked_input, input_mask = get_masked_input_and_mask(
41
+ input_ids,
42
+ shard_indices.org_vocab_start_index,
43
+ shard_indices.org_vocab_end_index,
44
+ shard_indices.num_org_vocab_padding,
45
+ shard_indices.added_vocab_start_index,
46
+ shard_indices.added_vocab_end_index,
47
+ )
48
+ output = torch.nn.functional.embedding(masked_input.long(), weight)
49
+ return output.masked_fill_(input_mask.unsqueeze(-1), 0)
50
+
51
+
52
+def embedding_fused(input_ids, weight, shard_indices):
53
+ return vllm_ops.vocab_parallel_embedding(
54
+ input_ids,
55
+ weight,
56
+ shard_indices.org_vocab_start_index,
57
+ shard_indices.org_vocab_end_index,
58
+ shard_indices.num_org_vocab_padding,
59
+ shard_indices.added_vocab_start_index,
60
+ shard_indices.added_vocab_end_index,
61
+ )
62
+
63
+
64
+PROVIDERS = {
65
+ "unfused": embedding_reference,
66
+ "fused": embedding_fused,
67
+}
68
+
69
+
70
+def make_inputs(vocab_size, hidden_size, tp_size, num_tokens, dtype, tp_rank=0):
71
+ shard_indices = get_shard_indices(vocab_size, vocab_size, tp_rank, tp_size)
72
+ weight = torch.randn(
73
+ pad_vocab_size(vocab_size) // tp_size,
74
+ hidden_size,
75
+ dtype=dtype,
76
+ device="cuda",
77
+ )
78
+ input_ids = torch.randint(
79
+ 0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda"
80
+ )
81
+ return input_ids, weight, shard_indices
82
+
83
+
84
+def graph_replay(fn):
85
+ """Capture fn into a CUDA graph and return its replay."""
86
+ stream = torch.cuda.Stream()
87
+ stream.wait_stream(torch.cuda.current_stream())
88
+ with torch.cuda.stream(stream):
89
+ for _ in range(3):
90
+ fn()
91
+ torch.cuda.current_stream().wait_stream(stream)
92
+ graph = torch.cuda.CUDAGraph()
93
+ with torch.cuda.graph(graph):
94
+ fn()
95
+ return graph.replay
96
+
97
+
98
+def calculate_diff(vocab_size, hidden_size, tp_size, num_tokens):
99
+ """The fused kernel must be bit-exact with the path it replaces, on every
100
+ rank of the shard."""
101
+ dtype = torch.bfloat16
102
+ for tp_rank in range(tp_size):
103
+ input_ids, weight, shard_indices = make_inputs(
104
+ vocab_size, hidden_size, tp_size, num_tokens, dtype, tp_rank
105
+ )
106
+ output_reference = embedding_reference(input_ids, weight, shard_indices)
107
+ output_fused = embedding_fused(input_ids, weight, shard_indices)
108
+ if not torch.equal(output_reference, output_fused):
109
+ print(f"❌ rank {tp_rank}: fused output differs")
110
+ return
111
+ print(f"✅ fused matches the reference path on all {tp_size} ranks")
112
+
113
+
114
+tp_size_range = [2, 4, 8]
115
+num_tokens_range = [1, 8, 64, 512, 2048, 8192, 32768]
116
+configs = list(itertools.product(tp_size_range, num_tokens_range))
117
+
118
+
119
+def get_benchmark(vocab_size, hidden_size, dtype, use_cuda_graph):
120
+ @triton.testing.perf_report(
121
+ triton.testing.Benchmark(
122
+ x_names=["tp_size", "num_tokens"],
123
+ x_vals=[list(_) for _ in configs],
124
+ line_arg="provider",
125
+ line_vals=list(PROVIDERS),
126
+ line_names=["Unfused PyTorch", "Fused kernel"],
127
+ styles=[("green", "-"), ("red", "-")],
128
+ ylabel="us",
129
+ plot_name=(
130
+ f"vocab-parallel-embedding-perf-vocab{vocab_size}-"
131
+ f"hidden{hidden_size}{'-cudagraph' if use_cuda_graph else ''}"
132
+ ),
133
+ args={},
134
+ )
135
+ )
136
+ def benchmark(tp_size, num_tokens, provider):
137
+ input_ids, weight, shard_indices = make_inputs(
138
+ vocab_size, hidden_size, tp_size, num_tokens, dtype
139
+ )
140
+ fn = PROVIDERS[provider]
141
+
142
+ def run():
143
+ return fn(input_ids, weight, shard_indices)
144
+
145
+ if use_cuda_graph:
146
+ run = graph_replay(run)
147
+
148
+ ms, min_ms, max_ms = triton.testing.do_bench(run, quantiles=[0.5, 0.2, 0.8])
149
+ return 1000 * ms, 1000 * max_ms, 1000 * min_ms
150
+
151
+ return benchmark
152
+
153
+
154
+if __name__ == "__main__":
155
+ import argparse
156
+
157
+ parser = argparse.ArgumentParser()
158
+ parser.add_argument(
159
+ "--vocab-size", type=int, default=128256, help="Vocabulary size"
160
+ )
161
+ parser.add_argument(
162
+ "--hidden-size", type=int, default=4096, help="Embedding dimension"
163
+ )
164
+ parser.add_argument(
165
+ "--dtype", type=str, default="bfloat16", help="Embedding table dtype"
166
+ )
167
+ parser.add_argument(
168
+ "--tp-size",
169
+ type=int,
170
+ default=4,
171
+ help="TP size used for the correctness check",
172
+ )
173
+ parser.add_argument(
174
+ "--use-cuda-graph",
175
+ action="store_true",
176
+ help="Time CUDA graph replay instead of eager launches",
177
+ )
178
+ parser.add_argument(
179
+ "--save-path",
180
+ type=str,
181
+ default=None,
182
+ help="Optional path to save benchmark results",
183
+ )
184
+ args = parser.parse_args()
185
+ dtype = getattr(torch, args.dtype)
186
+
187
+ calculate_diff(
188
+ vocab_size=args.vocab_size,
189
+ hidden_size=args.hidden_size,
190
+ tp_size=args.tp_size,
191
+ num_tokens=1024,
192
+ )
193
+
194
+ benchmark = get_benchmark(
195
+ args.vocab_size, args.hidden_size, dtype, args.use_cuda_graph
196
+ )
197
+ benchmark.run(print_data=True, save_path=args.save_path)
csrc/libtorch_stable/ops.h
+8
@@ -697,6 +697,14 @@ void cp_gather_indexer_k_quant_cache(
697
const torch::stable::Tensor& block_table, // [batch_size, num_blocks]
698
const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1]
699
700
+// Fused vocab-parallel embedding lookup (see
701
+// vocab_parallel_embedding_kernels.cu).
702
+void vocab_parallel_embedding(
703
+ torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
704
+ const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
705
+ int64_t org_vocab_end_index, int64_t num_org_vocab_padding,
706
+ int64_t added_vocab_start_index, int64_t added_vocab_end_index);
707
+
708
// LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu).
709
void ngram_compute_n_gram_ids(
710
int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights,
csrc/libtorch_stable/torch_bindings.cpp
+9
@@ -700,6 +700,14 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
700
// LongCat n-gram embedding index kernel. All tensor args are marked mutable
701
// to match the (non-const) stable-Tensor& C++ signature; only ne_token_table
702
// and n_gram_ids are actually written in place.
703
+ // Fused vocab-parallel embedding: gather the rows this rank owns and write
704
+ // zeros for the rest, so the following all-reduce reconstructs the full
705
+ // embedding.
706
+ ops.def(
707
+ "vocab_parallel_embedding(Tensor! out, Tensor input_ids, Tensor weight, "
708
+ "int org_vocab_start_index, int org_vocab_end_index, "
709
+ "int num_org_vocab_padding, int added_vocab_start_index, "
710
+ "int added_vocab_end_index) -> ()");
711
ops.def(
712
"ngram_compute_n_gram_ids(int ne_n, int ne_k, Tensor(a!) ne_weights, "
713
"Tensor(b!) ne_mods, Tensor(c!) exclusive_ne_embedder_size_sums, "
@@ -711,6 +719,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
719
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
720
// LongCat n-gram embedding index kernel.
721
ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids));
722
+ ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding));
723
724
// Per-token group quantization
725
ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));
csrc/libtorch_stable/vocab_parallel_embedding_kernels.cu
new
+166
@@ -0,0 +1,166 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+//
4
+// Fused vocab-parallel embedding lookup.
5
+//
6
+// With TP > 1 each rank owns a slice of the vocabulary, so the eager path is a
7
+// range mask, an id shift, an int64 cast, a gather and a masked_fill_ of the
8
+// gathered rows. This kernel does all of it in one pass: rows owned by another
9
+// rank are written as zeros, which keeps the following all-reduce exact.
10
+//
11
+// The gather is a pure row copy, so the kernel is templated on a byte-vector
12
+// type instead of the value type and works for any embedding dtype.
13
+
14
+#include "torch_utils.h"
15
+
16
+#include "ops.h"
17
+
18
+#include <algorithm>
19
+#include <cstdint>
20
+
21
+namespace vllm::vocab_embedding {
22
+
23
+constexpr int kBlockThreads = 256;
24
+// Enough token-blocks to fill the device several times over; the token loop
25
+// strides over the rest.
26
+constexpr int kTokenWaves = 4;
27
+
28
+template <typename idx_t, typename vec_t>
29
+__global__ void vocab_parallel_embedding_kernel(
30
+ vec_t* __restrict__ out, // [num_tokens, vecs_per_row]
31
+ const idx_t* __restrict__ input_ids, // [num_tokens]
32
+ const vec_t* __restrict__ weight, // [num_rows, vecs_per_row]
33
+ const int64_t num_tokens, const int64_t vecs_per_row,
34
+ const int64_t org_vocab_start_index, const int64_t org_vocab_end_index,
35
+ const int64_t added_vocab_start_index, const int64_t added_vocab_end_index,
36
+ const int64_t added_row_offset) {
37
+ const int64_t vec_begin =
38
+ static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
39
+ const int64_t vec_stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
40
+
41
+ for (int64_t token = blockIdx.y; token < num_tokens; token += gridDim.y) {
42
+ const int64_t id = static_cast<int64_t>(input_ids[token]);
43
+ const bool in_org = id >= org_vocab_start_index && id < org_vocab_end_index;
44
+ const bool in_added =
45
+ id >= added_vocab_start_index && id < added_vocab_end_index;
46
+ vec_t* out_row = out + token * vecs_per_row;
47
+
48
+ if (in_org || in_added) {
49
+ const int64_t row =
50
+ in_org ? id - org_vocab_start_index : id + added_row_offset;
51
+ const vec_t* weight_row = weight + row * vecs_per_row;
52
+ for (int64_t v = vec_begin; v < vecs_per_row; v += vec_stride) {
53
+ out_row[v] = weight_row[v];
54
+ }
55
+ } else {
56
+ // Token belongs to another rank's shard: contribute zeros.
57
+ const vec_t zero = vec_t();
58
+ for (int64_t v = vec_begin; v < vecs_per_row; v += vec_stride) {
59
+ out_row[v] = zero;
60
+ }
61
+ }
62
+ }
63
+}
64
+
65
+template <typename idx_t, typename vec_t>
66
+void launch(torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
67
+ const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
68
+ int64_t org_vocab_end_index, int64_t added_vocab_start_index,
69
+ int64_t added_vocab_end_index, int64_t added_row_offset,
70
+ int64_t num_tokens, int64_t row_bytes, cudaStream_t stream) {
71
+ const int64_t vecs_per_row = row_bytes / static_cast<int64_t>(sizeof(vec_t));
72
+ const int grid_x = static_cast<int>(
73
+ std::min<int64_t>((vecs_per_row + kBlockThreads - 1) / kBlockThreads,
74
+ static_cast<int64_t>(kBlockThreads)));
75
+ const int grid_y = static_cast<int>(std::min<int64_t>(
76
+ num_tokens, static_cast<int64_t>(kTokenWaves) *
77
+ get_device_prop()->multiProcessorCount));
78
+
79
+ vocab_parallel_embedding_kernel<idx_t, vec_t>
80
+ <<<dim3(grid_x, grid_y), kBlockThreads, 0, stream>>>(
81
+ reinterpret_cast<vec_t*>(out.mutable_data_ptr()),
82
+ reinterpret_cast<const idx_t*>(input_ids.const_data_ptr()),
83
+ reinterpret_cast<const vec_t*>(weight.const_data_ptr()), num_tokens,
84
+ vecs_per_row, org_vocab_start_index, org_vocab_end_index,
85
+ added_vocab_start_index, added_vocab_end_index, added_row_offset);
86
+}
87
+
88
+template <typename idx_t>
89
+void dispatch_vec_width(
90
+ torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
91
+ const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
92
+ int64_t org_vocab_end_index, int64_t added_vocab_start_index,
93
+ int64_t added_vocab_end_index, int64_t added_row_offset, int64_t num_tokens,
94
+ int64_t row_bytes, cudaStream_t stream) {
95
+ const uintptr_t addrs = reinterpret_cast<uintptr_t>(out.mutable_data_ptr()) |
96
+ reinterpret_cast<uintptr_t>(weight.const_data_ptr()) |
97
+ static_cast<uintptr_t>(row_bytes);
98
+
99
+#define VLLM_LAUNCH_VOCAB_EMBEDDING(vec_t) \
100
+ launch<idx_t, vec_t>(out, input_ids, weight, org_vocab_start_index, \
101
+ org_vocab_end_index, added_vocab_start_index, \
102
+ added_vocab_end_index, added_row_offset, num_tokens, \
103
+ row_bytes, stream)
104
+
105
+ if (addrs % 16 == 0) {
106
+ VLLM_LAUNCH_VOCAB_EMBEDDING(uint4);
107
+ } else if (addrs % 8 == 0) {
108
+ VLLM_LAUNCH_VOCAB_EMBEDDING(uint2);
109
+ } else if (addrs % 4 == 0) {
110
+ VLLM_LAUNCH_VOCAB_EMBEDDING(uint32_t);
111
+ } else if (addrs % 2 == 0) {
112
+ VLLM_LAUNCH_VOCAB_EMBEDDING(uint16_t);
113
+ } else {
114
+ VLLM_LAUNCH_VOCAB_EMBEDDING(uint8_t);
115
+ }
116
+#undef VLLM_LAUNCH_VOCAB_EMBEDDING
117
+}
118
+
119
+} // namespace vllm::vocab_embedding
120
+
121
+void vocab_parallel_embedding(
122
+ torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
123
+ const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
124
+ int64_t org_vocab_end_index, int64_t num_org_vocab_padding,
125
+ int64_t added_vocab_start_index, int64_t added_vocab_end_index) {
126
+ STD_TORCH_CHECK(weight.dim() == 2, "weight must be 2D");
127
+ STD_TORCH_CHECK(out.dim() == 2, "out must be 2D");
128
+ STD_TORCH_CHECK(out.size(1) == weight.size(1),
129
+ "out and weight must share the embedding dim");
130
+ STD_TORCH_CHECK(out.size(0) == input_ids.numel(),
131
+ "out must have one row per input id");
132
+ STD_TORCH_CHECK(out.scalar_type() == weight.scalar_type(),
133
+ "out and weight must have the same dtype");
134
+ STD_TORCH_CHECK(input_ids.is_contiguous() && weight.is_contiguous() &&
135
+ out.is_contiguous(),
136
+ "vocab_parallel_embedding requires contiguous tensors");
137
+
138
+ const int64_t num_tokens = input_ids.numel();
139
+ if (num_tokens == 0) return;
140
+
141
+ const int64_t row_bytes = weight.size(1) * weight.element_size();
142
+ // Offset that maps an added-vocab (LoRA) id onto its local row, mirroring
143
+ // get_masked_input_and_mask().
144
+ const int64_t added_row_offset = org_vocab_end_index - org_vocab_start_index +
145
+ num_org_vocab_padding -
146
+ added_vocab_start_index;
147
+
148
+ const torch::stable::accelerator::DeviceGuard device_guard(
149
+ out.get_device_index());
150
+ const cudaStream_t stream = get_current_cuda_stream();
151
+
152
+ if (input_ids.scalar_type() == torch::headeronly::ScalarType::Long) {
153
+ vllm::vocab_embedding::dispatch_vec_width<int64_t>(
154
+ out, input_ids, weight, org_vocab_start_index, org_vocab_end_index,
155
+ added_vocab_start_index, added_vocab_end_index, added_row_offset,
156
+ num_tokens, row_bytes, stream);
157
+ } else {
158
+ STD_TORCH_CHECK(
159
+ input_ids.scalar_type() == torch::headeronly::ScalarType::Int,
160
+ "input_ids must be int32 or int64");
161
+ vllm::vocab_embedding::dispatch_vec_width<int32_t>(
162
+ out, input_ids, weight, org_vocab_start_index, org_vocab_end_index,
163
+ added_vocab_start_index, added_vocab_end_index, added_row_offset,
164
+ num_tokens, row_bytes, stream);
165
+ }
166
+}
tests/kernels/core/test_vocab_parallel_embedding.py
new
+126
@@ -0,0 +1,126 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the fused vocab-parallel embedding kernel
4
+(``vllm._custom_ops.vocab_parallel_embedding``).
5
+
6
+The guarantee: the kernel is a drop-in for the TP > 1 embedding path, i.e. it
7
+returns bit-exactly what mask + shift + gather + ``masked_fill_`` returns on
8
+each rank, and summing the per-rank outputs (what the all-reduce does)
9
+reconstructs a full-table lookup.
10
+"""
11
+
12
+import pytest
13
+import torch
14
+
15
+from vllm import _custom_ops as ops
16
+from vllm.model_executor.layers.vocab_parallel_embedding import (
17
+ VocabParallelEmbedding,
18
+ get_masked_input_and_mask,
19
+ pad_vocab_size,
20
+)
21
+from vllm.platforms import current_platform
22
+from vllm.utils.torch_utils import set_random_seed
23
+
24
+requires_cuda = pytest.mark.skipif(
25
+ not current_platform.is_cuda(),
26
+ reason="the fused vocab-parallel embedding kernel is CUDA-only",
27
+)
28
+
29
+
30
+def _shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size, pad=64):
31
+ return VocabParallelEmbedding._get_indices(
32
+ pad_vocab_size(vocab_size, pad),
33
+ pad_vocab_size(org_vocab_size, pad),
34
+ vocab_size,
35
+ org_vocab_size,
36
+ tp_rank,
37
+ tp_size,
38
+ )
39
+
40
+
41
+def _reference(input_ids, weight, si):
42
+ """The path the kernel replaces."""
43
+ masked_input, input_mask = get_masked_input_and_mask(
44
+ input_ids,
45
+ si.org_vocab_start_index,
46
+ si.org_vocab_end_index,
47
+ si.num_org_vocab_padding,
48
+ si.added_vocab_start_index,
49
+ si.added_vocab_end_index,
50
+ )
51
+ out = torch.nn.functional.embedding(masked_input.long(), weight)
52
+ return out.masked_fill_(input_mask.unsqueeze(-1), 0)
53
+
54
+
55
+def _fused(input_ids, weight, si):
56
+ return ops.vocab_parallel_embedding(
57
+ input_ids,
58
+ weight,
59
+ si.org_vocab_start_index,
60
+ si.org_vocab_end_index,
61
+ si.num_org_vocab_padding,
62
+ si.added_vocab_start_index,
63
+ si.added_vocab_end_index,
64
+ )
65
+
66
+
67
+def _shard(full_table, si, rows_per_rank, dtype):
68
+ """Lay a rank's rows out the way the weight loader does."""
69
+ weight = torch.zeros(rows_per_rank, full_table.shape[1], dtype=dtype, device="cuda")
70
+ num_org = si.org_vocab_end_index - si.org_vocab_start_index
71
+ weight[:num_org] = full_table[si.org_vocab_start_index : si.org_vocab_end_index]
72
+ added_at = num_org + si.num_org_vocab_padding
73
+ num_added = si.added_vocab_end_index - si.added_vocab_start_index
74
+ weight[added_at : added_at + num_added] = full_table[
75
+ si.added_vocab_start_index : si.added_vocab_end_index
76
+ ]
77
+ return weight
78
+
79
+
80
+@requires_cuda
81
+@torch.inference_mode()
82
+@pytest.mark.parametrize("tp_size", [2, 4, 8])
83
+@pytest.mark.parametrize(
84
+ "vocab_size,org_vocab_size", [(128256, 128256), (32256, 32000)]
85
+)
86
+@pytest.mark.parametrize("hidden", [4096, 3584])
87
+@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64])
88
+def test_matches_masked_gather(tp_size, vocab_size, org_vocab_size, hidden, id_dtype):
89
+ """Every rank's partial output is bit-exact with the eager path, including
90
+ the added-vocab (LoRA) rows and the ids owned by another rank."""
91
+ set_random_seed(7)
92
+ dtype = torch.bfloat16
93
+ rows = pad_vocab_size(vocab_size, 64) // tp_size
94
+ full_table = torch.randn(pad_vocab_size(vocab_size, 64), hidden, dtype=dtype).cuda()
95
+ ids = torch.randint(0, vocab_size, (257,), dtype=id_dtype, device="cuda")
96
+
97
+ for tp_rank in range(tp_size):
98
+ si = _shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size)
99
+ weight = _shard(full_table, si, rows, dtype)
100
+ torch.testing.assert_close(
101
+ _fused(ids, weight, si), _reference(ids, weight, si), atol=0.0, rtol=0.0
102
+ )
103
+
104
+
105
+@requires_cuda
106
+@torch.inference_mode()
107
+@pytest.mark.parametrize("tp_size", [2, 8])
108
+def test_rank_sum_is_full_lookup(tp_size):
109
+ """What the all-reduce sees: the partials add up to a full-table lookup."""
110
+ set_random_seed(11)
111
+ vocab_size, org_vocab_size, hidden = 32256, 32000, 2048
112
+ dtype = torch.bfloat16
113
+ rows = pad_vocab_size(vocab_size, 64) // tp_size
114
+ full_table = torch.randn(pad_vocab_size(vocab_size, 64), hidden, dtype=dtype).cuda()
115
+ # Ids from both the base and the added (LoRA) range.
116
+ ids = torch.randint(0, vocab_size, (129,), dtype=torch.int32, device="cuda")
117
+ ids[:16] = torch.randint(
118
+ org_vocab_size, vocab_size, (16,), dtype=torch.int32, device="cuda"
119
+ )
120
+
121
+ total = torch.zeros(ids.shape[0], hidden, dtype=dtype, device="cuda")
122
+ for tp_rank in range(tp_size):
123
+ si = _shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size)
124
+ total += _fused(ids, _shard(full_table, si, rows, dtype), si)
125
+
126
+ torch.testing.assert_close(total, full_table[ids.long()], atol=0.0, rtol=0.0)
vllm/_custom_ops.py
+43
@@ -260,6 +260,49 @@ def rms_norm(
260
torch.ops._C.rms_norm(out, input, weight, epsilon)
261
262
263
+# Fused vocab-parallel embedding lookup
264
+# (see csrc/.../vocab_parallel_embedding_kernels.cu).
265
+def vocab_parallel_embedding(
266
+ input_ids: torch.Tensor,
267
+ weight: torch.Tensor,
268
+ org_vocab_start_index: int,
269
+ org_vocab_end_index: int,
270
+ num_org_vocab_padding: int,
271
+ added_vocab_start_index: int,
272
+ added_vocab_end_index: int,
273
+) -> torch.Tensor:
274
+ """Gather the embedding rows owned by this rank, zeros for the others.
275
+
276
+ Fuses the range mask, id shift, gather and output masking of the TP > 1
277
+ embedding path into a single kernel. Out-of-shard tokens produce zero rows,
278
+ so an all-reduce over the ranks yields the full embedding.
279
+
280
+ Args:
281
+ input_ids: ``[num_tokens]`` int32 or int64 token ids.
282
+ weight: ``[num_embeddings_per_partition, embedding_dim]`` shard.
283
+
284
+ Returns:
285
+ ``[num_tokens, embedding_dim]`` partial embeddings.
286
+ """
287
+ out = torch.empty(
288
+ input_ids.shape[0],
289
+ weight.shape[1],
290
+ dtype=weight.dtype,
291
+ device=weight.device,
292
+ )
293
+ torch.ops._C.vocab_parallel_embedding(
294
+ out,
295
+ input_ids,
296
+ weight,
297
+ org_vocab_start_index,
298
+ org_vocab_end_index,
299
+ num_org_vocab_padding,
300
+ added_vocab_start_index,
301
+ added_vocab_end_index,
302
+ )
303
+ return out
304
+
305
+
306
# LongCat n-gram embedding index kernel (see csrc/.../ngram_embedding_kernels.cu).
307
def ngram_compute_n_gram_ids(
308
ne_n: int,
vllm/model_executor/layers/vocab_parallel_embedding.py
+30
-10
@@ -9,6 +9,7 @@ import torch.nn.functional as F
9
from torch.nn.parameter import Parameter
10
11
import vllm.envs as envs
12
+from vllm import _custom_ops as ops
13
from vllm.distributed import (
14
divide,
15
get_tensor_model_parallel_rank,
@@ -325,6 +326,14 @@ class VocabParallelEmbedding(PluggableLayer):
326
- self.shard_indices.added_vocab_start_index
327
)
328
329
+ # The fused kernel masks, shifts and gathers in one pass; it assumes a
330
+ # plain row-major shard, so it only applies to unquantized weights.
331
+ self.use_fused_embedding = (
332
+ self.tp_size > 1
333
+ and current_platform.is_cuda()
334
+ and isinstance(quant_method, UnquantizedEmbeddingMethod)
335
+ )
336
+
337
self.quant_method.create_weights(
338
self,
339
self.embedding_dim,
@@ -488,8 +497,24 @@ class VocabParallelEmbedding(PluggableLayer):
497
param[loaded_weight.shape[0] :].data.fill_(0)
498
499
def forward(self, input_):
491
- if self.tp_size > 1:
492
- # Build the mask.
500
+ if self.tp_size == 1:
501
+ return self.quant_method.embedding(self, input_.long())
502
+
503
+ if self.use_fused_embedding:
504
+ output_parallel = ops.vocab_parallel_embedding(
505
+ input_ if input_.ndim == 1 else input_.reshape(-1),
506
+ self.weight,
507
+ self.shard_indices.org_vocab_start_index,
508
+ self.shard_indices.org_vocab_end_index,
509
+ self.shard_indices.num_org_vocab_padding,
510
+ self.shard_indices.added_vocab_start_index,
511
+ self.shard_indices.added_vocab_end_index,
512
+ )
513
+ if input_.ndim != 1:
514
+ output_parallel = output_parallel.view(
515
+ *input_.shape, self.embedding_dim
516
+ )
517
+ else:
518
masked_input, input_mask = get_masked_input_and_mask(
519
input_,
520
self.shard_indices.org_vocab_start_index,
@@ -498,12 +523,7 @@ class VocabParallelEmbedding(PluggableLayer):
523
self.shard_indices.added_vocab_start_index,
524
self.shard_indices.added_vocab_end_index,
525
)
501
- else:
502
- masked_input = input_
503
- # Get the embeddings.
504
- output_parallel = self.quant_method.embedding(self, masked_input.long())
505
- # Mask the output embedding.
506
- if self.tp_size > 1:
526
+ output_parallel = self.quant_method.embedding(self, masked_input.long())
527
if output_parallel.dtype in (
528
torch.float8_e4m3fn,
529
torch.float8_e5m2,
@@ -514,8 +534,8 @@ class VocabParallelEmbedding(PluggableLayer):
534
output = tensor_model_parallel_all_reduce(comm_output)
535
return output.view(output_parallel.dtype)
536
output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
517
- return tensor_model_parallel_all_reduce(output_parallel)
518
- return output_parallel
537
+ # Reduce across all the model parallel GPUs.
538
+ return tensor_model_parallel_all_reduce(output_parallel)
539
540
def extra_repr(self) -> str:
541
s = f"num_embeddings={self.num_embeddings}"