A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[V0 Deprecation] Remove advance_step (#22969)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 15, 2025 at 08:22 UTC
1c859a1387286cf650c3bc24fdeac706b97999e8
16 files changed
+9
-892
CMakeLists.txt
-1
@@ -249,7 +249,6 @@ set(VLLM_EXT_SRC
249
"csrc/quantization/gguf/gguf_kernel.cu"
250
"csrc/quantization/activation_kernels.cu"
251
"csrc/cuda_utils_kernels.cu"
252
- "csrc/prepare_inputs/advance_step.cu"
252
"csrc/custom_all_reduce.cu"
253
"csrc/torch_bindings.cpp")
254
csrc/ops.h
-16
@@ -145,22 +145,6 @@ void gelu_fast(torch::Tensor& out, torch::Tensor& input);
145
146
void gelu_quick(torch::Tensor& out, torch::Tensor& input);
147
148
-void advance_step_flashattn(int64_t num_seqs, int64_t num_queries,
149
- int64_t block_size, torch::Tensor& input_tokens,
150
- torch::Tensor& sampled_token_ids,
151
- torch::Tensor& input_positions,
152
- torch::Tensor& seq_lens,
153
- torch::Tensor& slot_mapping,
154
- torch::Tensor& block_tables);
155
-
156
-void advance_step_flashinfer(
157
- int64_t num_seqs, int64_t num_queries, int64_t block_size,
158
- torch::Tensor& input_tokens, torch::Tensor& sampled_token_ids,
159
- torch::Tensor& input_positions, torch::Tensor& seq_lens,
160
- torch::Tensor& slot_mapping, torch::Tensor& block_tables,
161
- torch::Tensor& paged_kv_indices, torch::Tensor& paged_kv_indptr,
162
- torch::Tensor& paged_kv_last_page_len, torch::Tensor& block_table_bounds);
163
-
148
void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
149
torch::Tensor const& q_pe,
150
torch::Tensor const& kv_c_and_k_pe_cache,
csrc/prepare_inputs/advance_step.cu
deleted
-336
@@ -1,336 +0,0 @@
1
-/*
2
- * The goal of this GPU kernel is to advance input tensors on the GPU directly
3
- * PR: https://github.com/vllm-project/vllm/pull/6338
4
- * Current restrictions:
5
- * 1. Specialized for DraftModelRunner
6
- * 2. Supports flash_attn only
7
- */
8
-
9
-#include "advance_step.cuh"
10
-
11
-namespace prepare_inputs {
12
-
13
-//
14
-template <int const num_threads>
15
-__global__ void advance_step_flashattn_kernel(
16
- int num_seqs, int num_queries, int block_size, long* input_tokens_ptr,
17
- long const* sampled_token_ids_ptr, long* input_positions_ptr,
18
- int* seq_lens_ptr, long* slot_mapping_ptr, int const* block_tables_ptr,
19
- int64_t const block_tables_stride) {
20
- int const n_pad = num_seqs - num_queries;
21
- if (n_pad && blockIdx.x == 0) {
22
- // Handle cuda graph padding
23
- int const offset = num_queries;
24
- for (int i = threadIdx.x; i < n_pad; i += blockDim.x) {
25
- input_tokens_ptr[offset + i] = 0;
26
- input_positions_ptr[offset + i] = 0;
27
- slot_mapping_ptr[offset + i] = -1;
28
- }
29
- }
30
-
31
- int num_query_blocks = div_ceil(num_queries, num_threads);
32
-
33
- if (blockIdx.x >= num_query_blocks) {
34
- return;
35
- }
36
-
37
- int cur_query_id = blockIdx.x * num_threads + threadIdx.x;
38
-
39
- if (cur_query_id >= num_queries) {
40
- return;
41
- }
42
-
43
- // Update input_tokens
44
- input_tokens_ptr[cur_query_id] = sampled_token_ids_ptr[cur_query_id];
45
-
46
- int seq_len = seq_lens_ptr[cur_query_id];
47
- int next_seq_len = seq_len + 1;
48
- int next_input_pos = next_seq_len - 1;
49
-
50
- // Update seq_lens
51
- seq_lens_ptr[cur_query_id] = next_seq_len;
52
- // Update input_positions
53
- input_positions_ptr[cur_query_id] = next_input_pos;
54
-
55
- int const* seq_block_tables_ptr =
56
- block_tables_ptr + block_tables_stride * cur_query_id;
57
-
58
- int block_index = next_input_pos / block_size;
59
- int block_offset = next_input_pos % block_size;
60
-
61
- int slot_num = seq_block_tables_ptr[block_index] * block_size + block_offset;
62
- // Update slot_mapping
63
- slot_mapping_ptr[cur_query_id] = slot_num;
64
-}
65
-
66
-inline void verify_tensor(std::string const& name, torch::Tensor const& t,
67
- int64_t const size_0, int64_t const size_1,
68
- c10::ScalarType const type) {
69
- bool size_0_cond = true;
70
- if (size_0 != -1) {
71
- size_0_cond = t.size(0) == size_0;
72
- }
73
-
74
- bool size_1_cond = true;
75
- if (size_1 != -1) {
76
- size_1_cond = t.size(1) == size_1;
77
- }
78
-
79
- bool is_contiguous = t.is_contiguous();
80
- bool same_type = t.dtype() == type;
81
-
82
- bool pass = size_0_cond && size_1_cond && is_contiguous && same_type;
83
- if (!pass) {
84
- TORCH_CHECK(false, "tensor: name = ", name, ", shape = ", t.sizes(),
85
- " is_cont = ", t.is_contiguous(), ", type = ", t.dtype(),
86
- " is not as expected: shape = [", size_0, ", ", size_1,
87
- "], type = ", type);
88
- }
89
-}
90
-
91
-/// each thread processes a block per query
92
-__global__ void advance_step_flashinfer_kernel(
93
- int num_threads, int num_seqs, int num_queries, int block_size,
94
- long* input_tokens_ptr, long const* sampled_token_ids_ptr,
95
- long* input_positions_ptr, int* seq_lens_ptr, long* slot_mapping_ptr,
96
- int const* block_tables_ptr, int64_t const block_tables_stride,
97
- int* paged_kv_last_page_len_ptr, int* block_table_bound_ptr) {
98
- int const n_pad = num_seqs - num_queries;
99
- if (n_pad && blockIdx.x == 0) {
100
- // Handle cuda graph padding
101
- int const offset = num_queries;
102
- for (int i = threadIdx.x; i < n_pad; i += blockDim.x) {
103
- input_tokens_ptr[offset + i] = 0;
104
- input_positions_ptr[offset + i] = 0;
105
- slot_mapping_ptr[offset + i] = -1;
106
- }
107
- }
108
- int num_query_blocks = div_ceil(num_queries, num_threads);
109
-
110
- if (blockIdx.x < num_query_blocks) {
111
- int cur_query_id = blockIdx.x * num_threads + threadIdx.x;
112
-
113
- if (cur_query_id < num_queries) {
114
- // Update input_tokens
115
- input_tokens_ptr[cur_query_id] = sampled_token_ids_ptr[cur_query_id];
116
-
117
- int seq_len = seq_lens_ptr[cur_query_id];
118
- int next_seq_len = seq_len + 1;
119
- int next_input_pos = next_seq_len - 1;
120
-
121
- // Update seq_lens
122
- seq_lens_ptr[cur_query_id] = next_seq_len;
123
- // Update input_positions
124
- input_positions_ptr[cur_query_id] = next_input_pos;
125
-
126
- int const* seq_block_tables_ptr =
127
- block_tables_ptr + block_tables_stride * cur_query_id;
128
-
129
- int block_index = next_input_pos / block_size;
130
- int block_offset = next_input_pos % block_size;
131
-
132
- // Update paged_kv_last_page_len
133
- paged_kv_last_page_len_ptr[cur_query_id] = block_offset + 1;
134
-
135
- int slot_num =
136
- seq_block_tables_ptr[block_index] * block_size + block_offset;
137
- // Update slot_mapping
138
- slot_mapping_ptr[cur_query_id] = slot_num;
139
- block_table_bound_ptr[cur_query_id] = div_ceil(next_seq_len, block_size);
140
- }
141
- }
142
-}
143
-
144
-__global__ void advance_step_flashinfer_indptr_kernel(
145
- int num_threads, int num_seqs, int num_queries, int* paged_kv_indptr_ptr,
146
- int* block_table_bound_ptr) {
147
- int idx = blockIdx.x * num_threads + threadIdx.x;
148
- // Update paged_kv_indptr
149
- if (idx == 0) {
150
- paged_kv_indptr_ptr[idx] = 0;
151
- }
152
- if (idx < num_queries) {
153
- int sum = 0;
154
- for (int i = 0; i <= idx; ++i) {
155
- sum += block_table_bound_ptr[i];
156
- }
157
- paged_kv_indptr_ptr[idx + 1] = sum;
158
- }
159
-}
160
-
161
-__global__ void advance_step_flashinfer_indices_kernel(
162
- int num_seqs, int num_queries, int const* block_tables_ptr,
163
- int64_t const max_num_blocks_per_seq, int* paged_kv_indices_ptr,
164
- int* paged_kv_indptr_ptr, int* block_table_bound_ptr) {
165
- // note: max_num_blocks_per_seq = block_tables.stride(0)
166
- int tid = blockIdx.x * blockDim.x + threadIdx.x;
167
-
168
- // when cuda graphs are enabled, paged_kv_indptr tensor
169
- // has to be updated for the padded queries
170
- // tid represents a query# for paged_kv_indptr tensor
171
- if (num_queries < tid && tid <= num_seqs) {
172
- paged_kv_indptr_ptr[tid] = paged_kv_indptr_ptr[num_queries];
173
- }
174
-
175
- // each thread processes a block_ptr in block_tables
176
- // block_tables shape: [num_queries, max_num_blocks_per_seq]
177
- // paged_kv_indices is flattened block_tables.
178
- for (int idx = tid; idx < (num_seqs * max_num_blocks_per_seq);
179
- idx += (gridDim.x * blockDim.x)) {
180
- // block_tables-row = paged_kv_indptr[queryNum]
181
- int queryNum = idx / max_num_blocks_per_seq;
182
- int col = idx % max_num_blocks_per_seq;
183
- if (queryNum < num_queries && col < block_table_bound_ptr[queryNum]) {
184
- int indices_arr_idx = paged_kv_indptr_ptr[queryNum] + col;
185
- int block_tables_idx = queryNum * max_num_blocks_per_seq + col;
186
- paged_kv_indices_ptr[indices_arr_idx] =
187
- block_tables_ptr[block_tables_idx];
188
- }
189
- }
190
-}
191
-
192
-void advance_step_flashattn(int num_seqs, int num_queries, int block_size,
193
- torch::Tensor& input_tokens, // type: long
194
- torch::Tensor& sampled_token_ids, // type: long
195
- torch::Tensor& input_positions, // type: long
196
- torch::Tensor& seq_lens, // type: int
197
- torch::Tensor& slot_mapping, // type: long
198
- torch::Tensor& block_tables) { // type: int
199
-
200
- if (logging) {
201
- printf("advance_step_flashattn:\n");
202
- printf(" num_seqs = %d\n", num_seqs);
203
- printf(" num_queries = %d\n", num_queries);
204
- printf(" block_size = %d\n", block_size);
205
- }
206
- // Verify all tensors
207
- verify_tensor("input_tokens", input_tokens, num_seqs, -1, at::kLong);
208
- verify_tensor("sampled_token_ids", sampled_token_ids, num_queries, 1,
209
- at::kLong);
210
- verify_tensor("input_positions", input_positions, num_seqs, -1, at::kLong);
211
- verify_tensor("seq_lens", seq_lens, num_seqs, -1, at::kInt);
212
- verify_tensor("slot_mapping", slot_mapping, num_seqs, -1, at::kLong);
213
- verify_tensor("block_tables", block_tables, num_seqs, -1, at::kInt);
214
-
215
- int dev = sampled_token_ids.get_device();
216
- cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
217
-
218
- int blocks;
219
- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
220
-
221
- advance_step_flashattn_kernel<max_threads>
222
- <<<blocks, max_threads, 0, stream>>>(
223
- num_seqs, num_queries, block_size,
224
- reinterpret_cast<long*>(input_tokens.data_ptr()),
225
- reinterpret_cast<long const*>(sampled_token_ids.data_ptr()),
226
- reinterpret_cast<long*>(input_positions.data_ptr()),
227
- reinterpret_cast<int*>(seq_lens.data_ptr()),
228
- reinterpret_cast<long*>(slot_mapping.data_ptr()),
229
- reinterpret_cast<int const*>(block_tables.data_ptr()),
230
- block_tables.stride(0));
231
-}
232
-
233
-void advance_step_flashinfer(
234
- int num_seqs, int num_queries, int block_size,
235
- torch::Tensor& input_tokens, // type: long
236
- torch::Tensor& sampled_token_ids, // type: long
237
- torch::Tensor& input_positions, // type: long
238
- torch::Tensor& seq_lens, // type: int
239
- torch::Tensor& slot_mapping, // type: long
240
- torch::Tensor& block_tables, // type: int
241
- torch::Tensor& paged_kv_indices, // type: int
242
- torch::Tensor& paged_kv_indptr, // type: int
243
- torch::Tensor& paged_kv_last_page_len, // type: int
244
- torch::Tensor& block_table_bound) { // type: int
245
-
246
- if (logging) {
247
- printf("advance_step_flashinfer:\n");
248
- printf(" num_seqs = %d\n", num_seqs);
249
- printf(" num_queries = %d\n", num_queries);
250
- printf(" block_size = %d\n", block_size);
251
- printf(" block_tables.stride(0) = %zu\n", block_tables.stride(0));
252
- }
253
- // Verify all tensors
254
- verify_tensor("input_tokens", input_tokens, num_seqs, -1, at::kLong);
255
- // verify_tensor("sampled_token_ids", sampled_token_ids, num_queries, 1,
256
- // at::kLong);
257
- verify_tensor("input_positions", input_positions, num_seqs, -1, at::kLong);
258
- verify_tensor("seq_lens", seq_lens, num_seqs, -1, at::kInt);
259
- verify_tensor("slot_mapping", slot_mapping, num_seqs, -1, at::kLong);
260
- verify_tensor("block_tables", block_tables, num_seqs, -1, at::kInt);
261
-
262
- verify_tensor("paged_kv_indices", paged_kv_indices, -1, -1, at::kInt);
263
- verify_tensor("paged_kv_indptr", paged_kv_indptr, num_seqs + 1, -1, at::kInt);
264
- verify_tensor("paged_kv_last_page_len", paged_kv_last_page_len, num_seqs, -1,
265
- at::kInt);
266
-
267
- verify_tensor("block_table_bound", block_table_bound, num_seqs, -1, at::kInt);
268
-
269
- int dev = sampled_token_ids.get_device();
270
- cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
271
-
272
- int blocks;
273
- int threads;
274
- cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
275
- cudaDeviceGetAttribute(&threads, cudaDevAttrMaxThreadsPerBlock, dev);
276
-
277
- TORCH_CHECK((blocks * threads > num_queries),
278
- "multi-step: not enough threads to map to num_queries = ",
279
- num_queries, " block_tables.stride(0) = ", block_tables.stride(0),
280
- " blocks = ", blocks, " max_threads = ", threads);
281
- if (logging) {
282
- printf("launching kernels with %d blocks and %d threads\n", blocks,
283
- threads);
284
- }
285
- advance_step_flashinfer_kernel<<<blocks, threads, 0, stream>>>(
286
- threads, num_seqs, num_queries, block_size,
287
- reinterpret_cast<long*>(input_tokens.data_ptr()),
288
- reinterpret_cast<long const*>(sampled_token_ids.data_ptr()),
289
- reinterpret_cast<long*>(input_positions.data_ptr()),
290
- reinterpret_cast<int*>(seq_lens.data_ptr()),
291
- reinterpret_cast<long*>(slot_mapping.data_ptr()),
292
- reinterpret_cast<int const*>(block_tables.data_ptr()),
293
- block_tables.stride(0),
294
- reinterpret_cast<int*>(paged_kv_last_page_len.data_ptr()),
295
- reinterpret_cast<int*>(block_table_bound.data_ptr()));
296
-
297
- advance_step_flashinfer_indptr_kernel<<<blocks, threads, 0, stream>>>(
298
- threads, num_seqs, num_queries,
299
- reinterpret_cast<int*>(paged_kv_indptr.data_ptr()),
300
- reinterpret_cast<int*>(block_table_bound.data_ptr()));
301
-
302
- advance_step_flashinfer_indices_kernel<<<blocks, threads, 0, stream>>>(
303
- num_seqs, num_queries,
304
- reinterpret_cast<int const*>(block_tables.data_ptr()),
305
- block_tables.stride(0),
306
- reinterpret_cast<int*>(paged_kv_indices.data_ptr()),
307
- reinterpret_cast<int*>(paged_kv_indptr.data_ptr()),
308
- reinterpret_cast<int*>(block_table_bound.data_ptr()));
309
-}
310
-
311
-} // namespace prepare_inputs
312
-
313
-void advance_step_flashattn(int64_t num_seqs, int64_t num_queries,
314
- int64_t block_size, torch::Tensor& input_tokens,
315
- torch::Tensor& sampled_token_ids,
316
- torch::Tensor& input_positions,
317
- torch::Tensor& seq_lens,
318
- torch::Tensor& slot_mapping,
319
- torch::Tensor& block_tables) {
320
- prepare_inputs::advance_step_flashattn(
321
- num_seqs, num_queries, block_size, input_tokens, sampled_token_ids,
322
- input_positions, seq_lens, slot_mapping, block_tables);
323
-}
324
-
325
-void advance_step_flashinfer(
326
- int64_t num_seqs, int64_t num_queries, int64_t block_size,
327
- torch::Tensor& input_tokens, torch::Tensor& sampled_token_ids,
328
- torch::Tensor& input_positions, torch::Tensor& seq_lens,
329
- torch::Tensor& slot_mapping, torch::Tensor& block_tables,
330
- torch::Tensor& paged_kv_indices, torch::Tensor& paged_kv_indptr,
331
- torch::Tensor& paged_kv_last_page_len, torch::Tensor& block_table_bound) {
332
- prepare_inputs::advance_step_flashinfer(
333
- num_seqs, num_queries, block_size, input_tokens, sampled_token_ids,
334
- input_positions, seq_lens, slot_mapping, block_tables, paged_kv_indices,
335
- paged_kv_indptr, paged_kv_last_page_len, block_table_bound);
336
-}
csrc/prepare_inputs/advance_step.cuh
deleted
-19
@@ -1,19 +0,0 @@
1
-#pragma once
2
-
3
-#include <torch/all.h>
4
-
5
-#include <ATen/cuda/CUDAContext.h>
6
-#include <c10/cuda/CUDAGuard.h>
7
-#include <cuda.h>
8
-#include <cuda_fp16.h>
9
-#include <cuda_runtime.h>
10
-#include <iostream>
11
-
12
-namespace prepare_inputs {
13
-
14
-static constexpr int max_threads = 256;
15
-static constexpr bool logging = false;
16
-
17
-constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
18
-
19
-} // namespace prepare_inputs
csrc/torch_bindings.cpp
-19
@@ -142,25 +142,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
142
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
143
ops.impl("gelu_quick", torch::kCUDA, &gelu_quick);
144
145
- // prepare_inputs advance_step
146
- ops.def(
147
- "advance_step_flashattn(int num_seqs, int num_queries, int block_size, "
148
- "Tensor! input_tokens, Tensor sampled_token_ids, "
149
- "Tensor! input_positions, Tensor! seq_lens, Tensor! slot_mapping, "
150
- "Tensor block_tables) -> ()");
151
- ops.impl("advance_step_flashattn", torch::kCUDA, &advance_step_flashattn);
152
-
153
- ops.def(
154
- "advance_step_flashinfer("
155
- " int num_seqs, int num_queries, int block_size,"
156
- " Tensor! input_tokens, Tensor sampled_token_ids,"
157
- " Tensor! input_positions, Tensor! seq_lens, Tensor! slot_mapping,"
158
- " Tensor block_tables, Tensor! paged_kv_indices,"
159
- " Tensor! paged_kv_indptr, Tensor! paged_kv_last_page_len,"
160
- " Tensor! block_table_bounds"
161
- ") -> ()");
162
- ops.impl("advance_step_flashinfer", torch::kCUDA, &advance_step_flashinfer);
163
-
145
// Layernorm
146
// Apply Root Mean Square (RMS) Normalization to the input tensor.
147
ops.def(
vllm/_custom_ops.py
-32
@@ -319,38 +319,6 @@ def apply_repetition_penalties(logits: torch.Tensor, prompt_mask: torch.Tensor,
319
repetition_penalties)
320
321
322
-def advance_step_flashattn(num_seqs: int, num_queries: int, block_size: int,
323
- input_tokens: torch.Tensor,
324
- sampled_token_ids: torch.Tensor,
325
- input_positions: torch.Tensor,
326
- seq_lens: torch.Tensor, slot_mapping: torch.Tensor,
327
- block_tables: torch.Tensor) -> None:
328
- """Advance a step on GPU for existing inputs for a multi-step runner"""
329
- return torch.ops._C.advance_step_flashattn(num_seqs, num_queries,
330
- block_size, input_tokens,
331
- sampled_token_ids,
332
- input_positions, seq_lens,
333
- slot_mapping, block_tables)
334
-
335
-
336
-def advance_step_flashinfer(num_seqs: int, num_queries: int, block_size: int,
337
- input_tokens: torch.Tensor,
338
- sampled_token_ids: torch.Tensor,
339
- input_positions: torch.Tensor,
340
- seq_lens: torch.Tensor, slot_mapping: torch.Tensor,
341
- block_tables: torch.Tensor,
342
- paged_kv_indices: torch.Tensor,
343
- paged_kv_indptr: torch.Tensor,
344
- paged_kv_last_page_len: torch.Tensor,
345
- block_table_bound: torch.Tensor) -> None:
346
-
347
- return torch.ops._C.advance_step_flashinfer(
348
- num_seqs, num_queries, block_size, input_tokens, sampled_token_ids,
349
- input_positions, seq_lens, slot_mapping, block_tables,
350
- paged_kv_indices, paged_kv_indptr, paged_kv_last_page_len,
351
- block_table_bound)
352
-
353
-
322
# fused quant layer norm ops
323
def rms_norm_dynamic_per_token_quant(
324
input: torch.Tensor,
vllm/attention/backends/abstract.py
-5
@@ -101,11 +101,6 @@ class AttentionBackend(ABC):
101
) -> None:
102
raise NotImplementedError
103
104
- def advance_step(self, model_input: "ModelRunnerInputBase",
105
- sampled_token_ids: Optional[torch.Tensor],
106
- block_size: int, num_seqs: int, num_queries: int) -> None:
107
- raise NotImplementedError
108
-
104
@classmethod
105
def full_cls_name(cls) -> tuple[str, str]:
106
return (cls.__module__, cls.__qualname__)
vllm/attention/backends/differential_flash_attn.py
+1
-75
@@ -35,8 +35,7 @@ from vllm.vllm_flash_attn import (flash_attn_varlen_func,
35
flash_attn_with_kvcache)
36
37
if TYPE_CHECKING:
38
- from vllm.worker.model_runner import (ModelInputForGPUBuilder,
39
- ModelInputForGPUWithSamplingMetadata)
38
+ from vllm.worker.model_runner import ModelInputForGPUBuilder
39
40
logger = init_logger(__name__)
41
@@ -326,79 +325,6 @@ class DifferentialFlashAttentionMetadata(AttentionMetadata):
325
cross_block_tables=self.cross_block_tables)
326
return self._cached_decode_metadata
327
329
- def advance_step(self,
330
- model_input: "ModelInputForGPUWithSamplingMetadata",
331
- sampled_token_ids: Optional[torch.Tensor],
332
- block_size: int,
333
- num_seqs: int,
334
- num_queries: int,
335
- turn_prefills_into_decodes: bool = False):
336
- """
337
- Update metadata in-place to advance one decode step.
338
- """
339
- # When using cudagraph, the num_seqs is padded to the next captured
340
- # batch sized, but num_queries tracks the actual number of requests in
341
- # the batch. For --enforce-eager mode, num_seqs == num_queries
342
- if num_seqs != num_queries:
343
- assert num_seqs > num_queries
344
- assert self.use_cuda_graph
345
-
346
- if turn_prefills_into_decodes:
347
- # When Multi-Step is enabled with Chunked-Prefill, prefills and
348
- # decodes are scheduled together. In the first step, all the
349
- # prefills turn into decodes. This update reflects that
350
- # conversion.
351
- assert self.num_decode_tokens + self.num_prefills == num_seqs
352
- self.num_decode_tokens += self.num_prefills
353
- self.num_prefills = 0
354
- self.num_prefill_tokens = 0
355
- self.max_prefill_seq_len = 0
356
- self.max_query_len = 1
357
-
358
- self.slot_mapping = self.slot_mapping[:num_seqs]
359
- else:
360
- assert self.seq_lens is not None
361
- assert self.max_decode_seq_len == max(self.seq_lens)
362
-
363
- assert self.num_prefills == 0
364
- assert self.num_prefill_tokens == 0
365
- assert self.num_decode_tokens == num_seqs
366
- assert self.slot_mapping.shape == (num_seqs, )
367
-
368
- assert self.seq_lens is not None
369
- assert len(self.seq_lens) == num_seqs
370
- assert self.seq_lens_tensor is not None
371
- assert self.seq_lens_tensor.shape == (num_seqs, )
372
- assert self.max_query_len == 1
373
- assert self.max_prefill_seq_len == 0
374
-
375
- assert self.query_start_loc is not None
376
- assert self.query_start_loc.shape == (num_queries + 1, )
377
- assert self.seq_start_loc is not None
378
- assert self.seq_start_loc.shape == (num_seqs + 1, )
379
-
380
- assert self.context_lens_tensor is not None
381
- assert self.context_lens_tensor.shape == (num_queries, )
382
-
383
- assert self.block_tables is not None
384
- assert self.block_tables.shape[0] == num_seqs
385
-
386
- # Update query lengths. Note that we update only queries and not seqs,
387
- # since tensors may be padded due to captured cuda graph batch size
388
- for i in range(num_queries):
389
- self.seq_lens[i] += 1
390
- self.max_decode_seq_len = max(self.seq_lens)
391
-
392
- ops.advance_step_flashattn(num_seqs=num_seqs,
393
- num_queries=num_queries,
394
- block_size=block_size,
395
- input_tokens=model_input.input_tokens,
396
- sampled_token_ids=sampled_token_ids,
397
- input_positions=model_input.input_positions,
398
- seq_lens=self.seq_lens_tensor,
399
- slot_mapping=self.slot_mapping,
400
- block_tables=self.block_tables)
401
-
328
329
class DifferentialFlashAttentionMetadataBuilder(
330
AttentionMetadataBuilder[DifferentialFlashAttentionMetadata]):
vllm/attention/backends/flash_attn.py
+1
-75
@@ -32,8 +32,7 @@ from vllm.vllm_flash_attn import (flash_attn_varlen_func,
32
flash_attn_with_kvcache)
33
34
if TYPE_CHECKING:
35
- from vllm.worker.model_runner import (ModelInputForGPUBuilder,
36
- ModelInputForGPUWithSamplingMetadata)
35
+ from vllm.worker.model_runner import ModelInputForGPUBuilder
36
37
logger = init_logger(__name__)
38
@@ -309,79 +308,6 @@ class FlashAttentionMetadata(AttentionMetadata):
308
cross_block_tables=self.cross_block_tables)
309
return self._cached_decode_metadata
310
312
- def advance_step(self,
313
- model_input: "ModelInputForGPUWithSamplingMetadata",
314
- sampled_token_ids: Optional[torch.Tensor],
315
- block_size: int,
316
- num_seqs: int,
317
- num_queries: int,
318
- turn_prefills_into_decodes: bool = False):
319
- """
320
- Update metadata in-place to advance one decode step.
321
- """
322
- # When using cudagraph, the num_seqs is padded to the next captured
323
- # batch sized, but num_queries tracks the actual number of requests in
324
- # the batch. For --enforce-eager mode, num_seqs == num_queries
325
- if num_seqs != num_queries:
326
- assert num_seqs > num_queries
327
- assert self.use_cuda_graph
328
-
329
- if turn_prefills_into_decodes:
330
- # When Multi-Step is enabled with Chunked-Prefill, prefills and
331
- # decodes are scheduled together. In the first step, all the
332
- # prefills turn into decodes. This update reflects that
333
- # conversion.
334
- assert self.num_decode_tokens + self.num_prefills == num_seqs
335
- self.num_decode_tokens += self.num_prefills
336
- self.num_prefills = 0
337
- self.num_prefill_tokens = 0
338
- self.max_prefill_seq_len = 0
339
- self.max_query_len = 1
340
-
341
- self.slot_mapping = self.slot_mapping[:num_seqs]
342
- else:
343
- assert self.seq_lens is not None
344
- assert self.max_decode_seq_len == max(self.seq_lens)
345
-
346
- assert self.num_prefills == 0
347
- assert self.num_prefill_tokens == 0
348
- assert self.num_decode_tokens == num_seqs
349
- assert self.slot_mapping.shape == (num_seqs, )
350
-
351
- assert self.seq_lens is not None
352
- assert len(self.seq_lens) == num_seqs
353
- assert self.seq_lens_tensor is not None
354
- assert self.seq_lens_tensor.shape == (num_seqs, )
355
- assert self.max_query_len == 1
356
- assert self.max_prefill_seq_len == 0
357
-
358
- assert self.query_start_loc is not None
359
- assert self.query_start_loc.shape == (num_queries + 1, )
360
- assert self.seq_start_loc is not None
361
- assert self.seq_start_loc.shape == (num_seqs + 1, )
362
-
363
- assert self.context_lens_tensor is not None
364
- assert self.context_lens_tensor.shape == (num_queries, )
365
-
366
- assert self.block_tables is not None
367
- assert self.block_tables.shape[0] == num_seqs
368
-
369
- # Update query lengths. Note that we update only queries and not seqs,
370
- # since tensors may be padded due to captured cuda graph batch size
371
- for i in range(num_queries):
372
- self.seq_lens[i] += 1
373
- self.max_decode_seq_len = max(self.seq_lens)
374
-
375
- ops.advance_step_flashattn(num_seqs=num_seqs,
376
- num_queries=num_queries,
377
- block_size=block_size,
378
- input_tokens=model_input.input_tokens,
379
- sampled_token_ids=sampled_token_ids,
380
- input_positions=model_input.input_positions,
381
- seq_lens=self.seq_lens_tensor,
382
- slot_mapping=self.slot_mapping,
383
- block_tables=self.block_tables)
384
-
311
312
class FlashAttentionMetadataBuilder(
313
AttentionMetadataBuilder[FlashAttentionMetadata]):
vllm/attention/backends/flashinfer.py
+2
-63
@@ -51,8 +51,7 @@ from vllm.utils.flashinfer import use_trtllm_attention
51
logger = init_logger(__name__)
52
53
if TYPE_CHECKING:
54
- from vllm.worker.model_runner import (ModelInputForGPUBuilder,
55
- ModelInputForGPUWithSamplingMetadata)
54
+ from vllm.worker.model_runner import ModelInputForGPUBuilder
55
56
57
class FlashInferBackend(AttentionBackend):
@@ -428,7 +427,7 @@ class FlashInferMetadata(AttentionMetadata):
427
query_start_loc: Optional[torch.Tensor] = None
428
block_tables: Optional[torch.Tensor] = None
429
431
- # used for GPU in-place advance_step
430
+ # used for GPU operations
431
seq_lens_tensor: Optional[torch.Tensor] = None
432
block_table_bound: Optional[torch.Tensor] = None
433
@@ -587,66 +586,6 @@ class FlashInferMetadata(AttentionMetadata):
586
return None
587
return self
588
590
- def advance_step(self,
591
- model_input: "ModelInputForGPUWithSamplingMetadata",
592
- sampled_token_ids: Optional[torch.Tensor],
593
- block_size: int,
594
- num_seqs: int,
595
- num_queries: int,
596
- turn_prefills_into_decodes: bool = False):
597
- """
598
- Update metadata in-place to advance one decode step.
599
- """
600
-
601
- if turn_prefills_into_decodes:
602
- # When Multi-Step is enabled with Chunked-Prefill, prefills and
603
- # decodes are scheduled together. In the first step, all the
604
- # prefills turn into decodes. This update reflects that
605
- # conversion.
606
- assert self.num_decode_tokens + self.num_prefills == num_seqs
607
- # Flashinfer doesn't support speculative decoding + chunked-prefill
608
- # + multi-step scheduling yet.
609
- assert self.decode_query_len == 1
610
- self.num_decode_tokens += self.num_prefills
611
- self.num_prefills = 0
612
- self.num_prefill_tokens = 0
613
- self.max_prefill_seq_len = 0
614
- self.max_query_len = 1
615
-
616
- self.slot_mapping = self.slot_mapping[:num_seqs]
617
- else:
618
- assert self.seq_lens_tensor is not None
619
-
620
- assert num_seqs > 0
621
- assert num_queries > 0
622
- assert model_input.attn_metadata is not None
623
- assert sampled_token_ids is not None
624
-
625
- # When using cudagraph, the num_seqs is padded to the next captured
626
- # batch sized, but num_queries tracks the actual number of requests in
627
- # the batch. For --enforce-eager mode, num_seqs == num_queries
628
- if num_seqs != num_queries:
629
- assert num_seqs > num_queries
630
- assert self.use_cuda_graph
631
-
632
- model_input.input_tokens[:num_queries] = sampled_token_ids.flatten()
633
-
634
- # Update GPU tensors
635
- ops.advance_step_flashinfer(
636
- num_seqs=num_seqs,
637
- num_queries=num_queries,
638
- block_size=block_size,
639
- input_tokens=model_input.input_tokens,
640
- sampled_token_ids=model_input.input_tokens,
641
- input_positions=model_input.input_positions,
642
- seq_lens=self.seq_lens_tensor,
643
- slot_mapping=self.slot_mapping,
644
- block_tables=self.block_tables,
645
- paged_kv_indices=self.paged_kv_indices,
646
- paged_kv_indptr=self.paged_kv_indptr,
647
- paged_kv_last_page_len=self.paged_kv_last_page_len,
648
- block_table_bound=self.block_table_bound)
649
-
589
590
class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
591
vllm/attention/backends/flashmla.py
+1
-14
@@ -3,7 +3,7 @@
3
4
from contextlib import contextmanager
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, List, Optional, Tuple, Type
6
+from typing import List, Optional, Tuple, Type
7
8
import torch
9
@@ -18,9 +18,6 @@ from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
18
get_mla_metadata,
19
is_flashmla_supported)
20
21
-if TYPE_CHECKING:
22
- from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
23
-
21
22
class FlashMLABackend(MLACommonBackend):
23
@@ -62,16 +59,6 @@ class FlashMLAMetadata(MLACommonMetadata):
59
self.decode_num_splits
60
return decode_metadata
61
65
- def advance_step(self,
66
- model_input: "ModelInputForGPUWithSamplingMetadata",
67
- sampled_token_ids: Optional[torch.Tensor],
68
- block_size: int,
69
- num_seqs: int,
70
- num_queries: int,
71
- turn_prefills_into_decodes: bool = False):
72
- raise NotImplementedError(
73
- "advance_step is not implemented for FlashMLA")
74
-
62
63
class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
64
vllm/attention/backends/mla/common.py
+1
-86
@@ -234,8 +234,7 @@ except ImportError:
234
flash_attn_varlen_func = None
235
236
if TYPE_CHECKING:
237
- from vllm.worker.model_runner import (ModelInputForGPUBuilder,
238
- ModelInputForGPUWithSamplingMetadata)
237
+ from vllm.worker.model_runner import ModelInputForGPUBuilder
238
239
is_hip = current_platform.is_rocm()
240
@@ -631,90 +630,6 @@ class MLACommonMetadata(AttentionMetadata):
630
is_profile_run=self.is_profile_run)
631
return self._cached_decode_metadata
632
634
- def advance_step(self,
635
- model_input: "ModelInputForGPUWithSamplingMetadata",
636
- sampled_token_ids: Optional[torch.Tensor],
637
- block_size: int,
638
- num_seqs: int,
639
- num_queries: int,
640
- turn_prefills_into_decodes: bool = False):
641
- """
642
- Update metadata in-place to advance one decode step.
643
- """
644
- # When using cudagraph, the num_seqs is padded to the next captured
645
- # batch sized, but num_queries tracks the actual number of requests in
646
- # the batch. For --enforce-eager mode, num_seqs == num_queries
647
- if num_seqs != num_queries:
648
- assert num_seqs > num_queries
649
-
650
- if turn_prefills_into_decodes:
651
- # When Multi-Step is enabled with Chunked-Prefill, prefills and
652
- # decodes are scheduled together. In the first step, all the
653
- # prefills turn into decodes. This update reflects that
654
- # conversion.
655
- assert self.num_decode_tokens + self.num_prefills == num_seqs
656
- self.num_decode_tokens += self.num_prefills
657
- self.num_prefills = 0
658
- self.num_prefill_tokens = 0
659
- self.max_prefill_seq_len = 0
660
- self.max_query_len = 1
661
-
662
- self.slot_mapping = self.slot_mapping[:num_seqs]
663
- else:
664
- assert self.seq_lens is not None
665
- assert self.max_decode_seq_len == max(self.seq_lens)
666
-
667
- assert self.num_prefills == 0
668
- assert self.num_prefill_tokens == 0
669
- assert self.num_decode_tokens == num_seqs
670
- assert self.slot_mapping.shape == (num_seqs, )
671
-
672
- assert self.seq_lens is not None
673
- assert len(self.seq_lens) == num_seqs
674
- assert self.seq_lens_tensor is not None
675
- assert self.seq_lens_tensor.shape == (num_seqs, )
676
- assert self.max_query_len == 1
677
- assert self.max_prefill_seq_len == 0
678
-
679
- assert self.query_start_loc is not None
680
- assert self.query_start_loc.shape == (num_queries + 1, )
681
- assert self.seq_start_loc is not None
682
- assert self.seq_start_loc.shape == (num_seqs + 1, )
683
-
684
- assert self.context_lens_tensor is not None
685
- assert self.context_lens_tensor.shape == (num_queries, )
686
-
687
- assert self.block_tables is not None
688
- assert self.block_tables.shape[0] == num_seqs
689
-
690
- # Update query lengths. Note that we update only queries and not seqs,
691
- # since tensors may be padded due to captured cuda graph batch size
692
- for i in range(num_queries):
693
- self.seq_lens[i] += 1
694
- self.max_decode_seq_len = max(self.seq_lens)
695
-
696
- self._ops_advance_step(num_seqs=num_seqs,
697
- num_queries=num_queries,
698
- block_size=block_size,
699
- input_tokens=model_input.input_tokens,
700
- sampled_token_ids=sampled_token_ids,
701
- input_positions=model_input.input_positions)
702
-
703
- def _ops_advance_step(self, num_seqs: int, num_queries: int,
704
- block_size: int, input_tokens: torch.Tensor,
705
- sampled_token_ids: torch.Tensor,
706
- input_positions: torch.Tensor) -> None:
707
- # here we use advance_step_flashinfo to update the paged_kv_* tensors
708
- ops.advance_step_flashattn(num_seqs=num_seqs,
709
- num_queries=num_queries,
710
- block_size=block_size,
711
- input_tokens=input_tokens,
712
- sampled_token_ids=sampled_token_ids,
713
- input_positions=input_positions,
714
- seq_lens=self.seq_lens_tensor,
715
- slot_mapping=self.slot_mapping,
716
- block_tables=self.block_tables)
717
-
633
634
class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
635
"""
vllm/attention/backends/placeholder_attn.py
+1
-61
@@ -15,8 +15,7 @@ from vllm.attention.backends.utils import CommonAttentionState
15
from vllm.multimodal import MultiModalPlaceholderMap
16
17
if TYPE_CHECKING:
18
- from vllm.worker.model_runner import (ModelInputForGPUBuilder,
19
- ModelInputForGPUWithSamplingMetadata)
18
+ from vllm.worker.model_runner import (ModelInputForGPUBuilder)
19
from vllm.utils import async_tensor_h2d
20
21
# Placeholder attention backend for models like Mamba and pooling models that
@@ -201,65 +200,6 @@ class PlaceholderAttentionMetadata(AttentionMetadata):
200
)
201
return self._cached_decode_metadata
202
204
- def advance_step(self,
205
- model_input: "ModelInputForGPUWithSamplingMetadata",
206
- sampled_token_ids: Optional[torch.Tensor],
207
- block_size: int,
208
- num_seqs: int,
209
- num_queries: int,
210
- turn_prefills_into_decodes: bool = False):
211
- """
212
- Update metadata in-place to advance one decode step.
213
- """
214
- # When using cudagraph, the num_seqs is padded to the next captured
215
- # batch sized, but num_queries tracks the actual number of requests in
216
- # the batch. For --enforce-eager mode, num_seqs == num_queries
217
- if num_seqs != num_queries:
218
- assert num_seqs > num_queries
219
- assert self.use_cuda_graph
220
-
221
- assert not turn_prefills_into_decodes, \
222
- ("Multi-Step + Chunked-Prefill is not supported for attention-free"
223
- "models. turn_prefills_into_decodes is a "
224
- "Multi-Step + Chunked-Prefill specific parameter.")
225
-
226
- assert self.seq_lens is not None
227
- assert self.max_decode_seq_len == max(self.seq_lens)
228
-
229
- assert self.num_prefills == 0
230
- assert self.num_prefill_tokens == 0
231
- assert self.num_decode_tokens == num_seqs
232
-
233
- assert self.seq_lens is not None
234
- assert len(self.seq_lens) == num_seqs
235
- assert self.seq_lens_tensor is not None
236
- assert self.seq_lens_tensor.shape == (num_seqs, )
237
- assert self.max_query_len == 1
238
- assert self.max_prefill_seq_len == 0
239
-
240
- assert self.query_start_loc is not None
241
- assert self.query_start_loc.shape == (num_queries + 1, )
242
- assert self.seq_start_loc is not None
243
- assert self.seq_start_loc.shape == (num_seqs + 1, )
244
-
245
- assert self.context_lens_tensor is not None
246
- assert self.context_lens_tensor.shape == (num_queries, )
247
-
248
- # Update query lengths. Note that we update only queries and not seqs,
249
- # since tensors may be padded due to captured cuda graph batch size
250
- for i in range(num_queries):
251
- self.seq_lens[i] += 1
252
- self.max_decode_seq_len = max(self.seq_lens)
253
-
254
- # Update sequences, masking off entries greater than num_queries
255
- device = self.seq_lens_tensor.device
256
- mask = torch.arange(self.seq_lens_tensor.size(0),
257
- device=device) < num_queries
258
- self.seq_lens_tensor += mask.to(self.seq_lens_tensor.dtype)
259
- if sampled_token_ids is not None:
260
- model_input.input_tokens.masked_scatter_(
261
- mask, sampled_token_ids[:num_queries])
262
-
203
204
class PlaceholderAttentionMetadataBuilder(
205
AttentionMetadataBuilder[PlaceholderAttentionMetadata]):
vllm/attention/backends/rocm_aiter_mla.py
-21
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Optional, Type, Union
7
8
import torch
9
10
-import vllm._custom_ops as ops
10
import vllm.envs as envs
11
from vllm.attention.backends.mla.common import (MLACommonBackend,
12
MLACommonImpl,
@@ -107,26 +106,6 @@ class AiterMLAMetadata(MLACommonMetadata):
106
107
return self._cached_decode_metadata
108
110
- def _ops_advance_step(self, num_seqs: int, num_queries: int,
111
- block_size: int, input_tokens: torch.Tensor,
112
- sampled_token_ids: torch.Tensor,
113
- input_positions: torch.Tensor) -> None:
114
-
115
- ops.advance_step_flashinfer(
116
- num_seqs=num_seqs,
117
- num_queries=num_queries,
118
- block_size=block_size,
119
- input_tokens=input_tokens,
120
- sampled_token_ids=sampled_token_ids,
121
- input_positions=input_positions,
122
- seq_lens=self.seq_lens_tensor,
123
- slot_mapping=self.slot_mapping,
124
- block_tables=self.block_tables,
125
- paged_kv_indices=self.paged_kv_indices,
126
- paged_kv_indptr=self.paged_kv_indptr,
127
- paged_kv_last_page_lens=self.paged_kv_last_page_lens,
128
- block_table_bound=self.block_table_bound)
129
-
109
110
class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
111
BLOCK_TABLE_EXTENDER: list[list[int]] = [[]]
vllm/attention/backends/rocm_flash_attn.py
+1
-67
@@ -4,7 +4,7 @@
4
import itertools
5
from dataclasses import dataclass
6
from functools import cache
7
-from typing import TYPE_CHECKING, List, Optional, Tuple, Type
7
+from typing import List, Optional, Tuple, Type
8
9
import torch
10
@@ -23,9 +23,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
GroupShape)
24
from vllm.platforms import current_platform
25
26
-if TYPE_CHECKING:
27
- from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
28
-
26
logger = init_logger(__name__)
27
_PARTITION_SIZE_ROCM = 256
28
@@ -261,69 +258,6 @@ class ROCmFlashAttentionMetadata(AttentionMetadata, PagedAttentionMetadata):
258
self._cached_decode_metadata.query_start_loc = qs - qs[0]
259
return self._cached_decode_metadata
260
264
- def advance_step(self,
265
- model_input: "ModelInputForGPUWithSamplingMetadata",
266
- sampled_token_ids: Optional[torch.Tensor],
267
- block_size: int,
268
- num_seqs: int,
269
- num_queries: int,
270
- turn_prefills_into_decodes: bool = False):
271
- """
272
- Update metadata in-place to advance one decode step.
273
- """
274
-
275
- assert not turn_prefills_into_decodes, \
276
- ("Chunked prefill is not supported with rocm_flash_attn yet."
277
- "turn_prefills_into_decodes is a Multi-Step + Chunked-Prefill "
278
- "specific parameter.")
279
-
280
- # When using cudagraph, the num_seqs is padded to the next captured
281
- # batch sized, but num_queries tracks the actual number of requests in
282
- # the batch. For --enforce-eager mode, num_seqs == num_queries
283
- if num_seqs != num_queries:
284
- assert num_seqs > num_queries
285
- assert self.use_cuda_graph
286
-
287
- assert self.num_prefills == 0
288
- assert self.num_prefill_tokens == 0
289
- assert self.num_decode_tokens == num_seqs
290
- assert self.slot_mapping.shape == (num_seqs, )
291
-
292
- assert self.seq_lens is not None
293
- assert len(self.seq_lens) == num_seqs
294
- assert self.seq_lens_tensor is not None
295
- assert self.seq_lens_tensor.shape == (num_seqs, )
296
- assert self.max_query_len == 1
297
- assert self.max_prefill_seq_len == 0
298
- assert self.max_decode_seq_len == max(self.seq_lens)
299
-
300
- assert self.query_start_loc is not None
301
- assert self.query_start_loc.shape == (num_queries + 1, )
302
- assert self.seq_start_loc is not None
303
- assert self.seq_start_loc.shape == (num_seqs + 1, )
304
-
305
- assert self.context_lens_tensor is not None
306
- assert self.context_lens_tensor.shape == (num_queries, )
307
-
308
- assert self.block_tables is not None
309
- assert self.block_tables.shape[0] == num_seqs
310
-
311
- # Update query lengths. Note that we update only queries and not seqs,
312
- # since tensors may be padded due to captured cuda graph batch size
313
- for i in range(num_queries):
314
- self.seq_lens[i] += 1
315
- self.max_decode_seq_len = max(self.seq_lens)
316
-
317
- ops.advance_step_flashattn(num_seqs=num_seqs,
318
- num_queries=num_queries,
319
- block_size=block_size,
320
- input_tokens=model_input.input_tokens,
321
- sampled_token_ids=sampled_token_ids,
322
- input_positions=model_input.input_positions,
323
- seq_lens=self.seq_lens_tensor,
324
- slot_mapping=self.slot_mapping,
325
- block_tables=self.block_tables)
326
-
261
262
class ROCmFlashAttentionMetadataBuilder(
263
CommonMetadataBuilder[ROCmFlashAttentionMetadata]):
vllm/worker/model_runner.py
+1
-2
@@ -762,8 +762,7 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
762
has Prefills (if any). The rest of the steps are guaranteed to be all
763
decodes. In this case, we set up the padding as if all the sequences
764
are decodes so we may run all steps except the first step in CUDA graph
765
- mode. The padding is accounted for in the multi-step `advance_step`
766
- family of functions.
765
+ mode.
766
767
Args:
768
num_seqs (int): Number of sequences scheduled to run.