@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 166 lines 7.34 KB
Raw
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 }