input_batch.py 17.1 KB
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
2
3
4
5
6
7
8
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from typing import Any

import numpy as np
import torch

9
from vllm.triton_utils import tl, triton
Woosuk Kwon's avatar
Woosuk Kwon committed
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from vllm.utils import random_uuid


class InputBuffers:
    def __init__(
        self,
        max_num_reqs: int,
        max_num_tokens: int,
        device: torch.device,
    ):
        self.max_num_reqs = max_num_reqs
        self.max_num_tokens = max_num_tokens
        self.device = device

24
        self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device)
25
        self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
26
27
        self.query_start_loc = torch.zeros(
            max_num_reqs + 1, dtype=torch.int32, device=device
Woosuk Kwon's avatar
Woosuk Kwon committed
28
        )
29
        self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
30
31
32
33
        # DCP: per-request local seq_lens buffer
        self.dcp_local_seq_lens = torch.zeros(
            max_num_reqs, dtype=torch.int32, device=device
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
34
35
36
37
38
39
40
41
42
43
44


@dataclass
class InputBatch:
    # batch_idx -> req_id
    req_ids: list[str]
    num_reqs: int

    # batch_idx -> req_state_idx
    idx_mapping: torch.Tensor
    idx_mapping_np: np.ndarray
45
46
    # Identical to idx_mapping except for spec decoding.
    expanded_idx_mapping: torch.Tensor
47
48
    # [total_num_logits] position within request for each logit
    expanded_local_pos: torch.Tensor
Woosuk Kwon's avatar
Woosuk Kwon committed
49
50
51
52
53
54
55

    # [num_reqs]
    # batch_idx -> num_scheduled_tokens
    num_scheduled_tokens: np.ndarray
    # sum(num_scheduled_tokens)
    num_tokens: int
    num_tokens_after_padding: int
56
    num_draft_tokens: int
Woosuk Kwon's avatar
Woosuk Kwon committed
57
58
59
60
61
62
63
64
65
66
67

    # [num_reqs + 1]
    query_start_loc: torch.Tensor
    query_start_loc_np: np.ndarray
    # [num_reqs]
    seq_lens: torch.Tensor

    # [num_tokens_after_padding]
    input_ids: torch.Tensor
    # [num_tokens_after_padding]
    positions: torch.Tensor
68
    # [3, num_tokens_after_padding]
69
    mrope_positions: torch.Tensor | None
70
71
    # [num_tokens_after_padding, hidden_size]
    inputs_embeds: torch.Tensor | None
Woosuk Kwon's avatar
Woosuk Kwon committed
72
73
74

    # layer_name -> Metadata
    attn_metadata: dict[str, Any]
75
76
    # layer_name -> slot_mapping
    slot_mappings: dict[str, torch.Tensor]
Woosuk Kwon's avatar
Woosuk Kwon committed
77

78
    # [total_num_logits]
Woosuk Kwon's avatar
Woosuk Kwon committed
79
    logits_indices: torch.Tensor
80
81
    # [num_reqs + 1]
    cu_num_logits: torch.Tensor
82
    cu_num_logits_np: np.ndarray
Woosuk Kwon's avatar
Woosuk Kwon committed
83

84
85
86
    # Whether any requests in batch use structured output.
    has_structured_output_reqs: bool

Woosuk Kwon's avatar
Woosuk Kwon committed
87
88
89
90
91
92
93
94
95
96
97
98
    @classmethod
    def make_dummy(
        cls,
        num_reqs: int,
        num_tokens: int,
        input_buffers: InputBuffers,
        device: torch.device,
    ) -> "InputBatch":
        assert 0 < num_reqs <= num_tokens
        req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
        idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
        idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
99
        expanded_idx_mapping = idx_mapping
100
        expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
Woosuk Kwon's avatar
Woosuk Kwon committed
101
102
103
104
105
        num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
        num_scheduled_tokens[-1] += num_tokens % num_reqs
        assert int(num_scheduled_tokens.sum()) == num_tokens

        # seq_len equals to query_len
106
107
        input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs
        input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs
108
        # Pad for full CUDA graph mode.
109
110
        input_buffers.seq_lens[num_reqs:] = 0
        seq_lens = input_buffers.seq_lens[:num_reqs]
Woosuk Kwon's avatar
Woosuk Kwon committed
111

112
113
114
        query_start_loc_np = np.empty(num_reqs + 1, dtype=np.int32)
        query_start_loc_np[0] = 0
        np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1:])
115
        input_buffers.query_start_loc[:1] = 0
116
117
118
119
120
121
122
        torch.cumsum(
            seq_lens, dim=0, out=input_buffers.query_start_loc[1 : num_reqs + 1]
        )
        # Pad for full CUDA graph mode.
        input_buffers.query_start_loc[num_reqs + 1 :] = num_tokens
        query_start_loc = input_buffers.query_start_loc[: num_reqs + 1]

123
124
125
        input_ids = input_buffers.input_ids[:num_tokens].zero_()
        positions = input_buffers.positions[:num_tokens].zero_()

Woosuk Kwon's avatar
Woosuk Kwon committed
126
127
        # attn_metadata = defaultdict(lambda: None)
        logits_indices = query_start_loc[1:] - 1
128
        cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
129
        cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
Woosuk Kwon's avatar
Woosuk Kwon committed
130
131
132
133
134
        return cls(
            req_ids=req_ids,
            num_reqs=num_reqs,
            idx_mapping=idx_mapping,
            idx_mapping_np=idx_mapping_np,
135
            expanded_idx_mapping=expanded_idx_mapping,
136
            expanded_local_pos=expanded_local_pos,
Woosuk Kwon's avatar
Woosuk Kwon committed
137
138
139
            num_scheduled_tokens=num_scheduled_tokens,
            num_tokens=num_tokens,
            num_tokens_after_padding=num_tokens,
140
            num_draft_tokens=0,
Woosuk Kwon's avatar
Woosuk Kwon committed
141
142
143
144
145
            query_start_loc=query_start_loc,
            query_start_loc_np=query_start_loc_np,
            seq_lens=seq_lens,
            input_ids=input_ids,
            positions=positions,
146
            mrope_positions=None,
147
            inputs_embeds=None,
Woosuk Kwon's avatar
Woosuk Kwon committed
148
            attn_metadata=None,  # type: ignore
149
            slot_mappings=None,  # type: ignore
Woosuk Kwon's avatar
Woosuk Kwon committed
150
            logits_indices=logits_indices,
151
            cu_num_logits=cu_num_logits,
152
            cu_num_logits_np=cu_num_logits_np,
153
            has_structured_output_reqs=False,
Woosuk Kwon's avatar
Woosuk Kwon committed
154
155
156
        )


157
158
159
160
161
162
@triton.jit
def _prepare_prefill_inputs_kernel(
    input_ids_ptr,
    next_prefill_tokens_ptr,
    idx_mapping_ptr,
    query_start_loc_ptr,
163
164
    all_token_ids_ptr,
    all_token_ids_stride,
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
    prefill_lens_ptr,
    num_computed_tokens_ptr,
    BLOCK_SIZE: tl.constexpr,
):
    batch_idx = tl.program_id(0)
    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
    prefill_len = tl.load(prefill_lens_ptr + req_state_idx)
    num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
    if num_computed >= prefill_len:
        # Not prefill.
        return

    query_start = tl.load(query_start_loc_ptr + batch_idx)
    query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
    query_len = query_end - query_start

181
    request_ptr = all_token_ids_ptr + req_state_idx * all_token_ids_stride
182
183
184
    for i in range(0, query_len, BLOCK_SIZE):
        block = i + tl.arange(0, BLOCK_SIZE)
        mask = block < query_len
185
        tokens = tl.load(request_ptr + num_computed + block, mask=mask)
186
187
188
189
        tl.store(input_ids_ptr + query_start + block, tokens, mask=mask)

    next_pos = num_computed + query_len
    if next_pos < prefill_len:
190
        next_token = tl.load(request_ptr + next_pos)
191
        tl.store(next_prefill_tokens_ptr + req_state_idx, next_token)
Woosuk Kwon's avatar
Woosuk Kwon committed
192
193


194
def prepare_prefill_inputs(
195
196
197
198
    input_ids: torch.Tensor,
    next_prefill_tokens: torch.Tensor,
    idx_mapping: torch.Tensor,
    query_start_loc: torch.Tensor,
199
    all_token_ids: torch.Tensor,
200
201
    prefill_len: torch.Tensor,
    num_computed_tokens: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
202
) -> None:
203
204
205
206
    num_reqs = idx_mapping.shape[0]
    _prepare_prefill_inputs_kernel[(num_reqs,)](
        input_ids,
        next_prefill_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
207
        idx_mapping,
208
        query_start_loc,
209
210
        all_token_ids,
        all_token_ids.stride(0),
211
212
213
        prefill_len,
        num_computed_tokens,
        BLOCK_SIZE=1024,
Woosuk Kwon's avatar
Woosuk Kwon committed
214
215
216
217
    )


@triton.jit
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
def _prepare_pos_seq_lens_kernel(
    pos_ptr,
    seq_lens_ptr,
    idx_mapping_ptr,
    query_start_loc_ptr,
    num_computed_tokens_ptr,
    max_num_reqs,
    BLOCK_SIZE: tl.constexpr,
):
    req_id = tl.program_id(0)
    num_reqs = tl.num_programs(0) - 1
    if req_id == num_reqs:
        # Pad unused seq_lens as 0 for full CUDA graphs.
        for i in tl.range(num_reqs, max_num_reqs, BLOCK_SIZE):
            block = i + tl.arange(0, BLOCK_SIZE)
            mask = block < max_num_reqs
            tl.store(seq_lens_ptr + block, 0, mask=mask)
        return

    req_state_idx = tl.load(idx_mapping_ptr + req_id)
    num_computed_tokens = tl.load(num_computed_tokens_ptr + req_state_idx)

    start = tl.load(query_start_loc_ptr + req_id)
    end = tl.load(query_start_loc_ptr + req_id + 1)
    query_len = end - start

    seq_len = num_computed_tokens + query_len
    tl.store(seq_lens_ptr + req_id, seq_len)

    for i in tl.range(0, query_len, BLOCK_SIZE):
        block = i + tl.arange(0, BLOCK_SIZE)
        mask = block < query_len
        pos = num_computed_tokens + block
        tl.store(pos_ptr + start + block, pos, mask=mask)


def prepare_pos_seq_lens(
    idx_mapping: torch.Tensor,
    query_start_loc: torch.Tensor,
    num_computed_tokens: torch.Tensor,
    pos: torch.Tensor,
    seq_lens: torch.Tensor,
) -> None:
    num_reqs = idx_mapping.shape[0]
    # NOTE(woosuk): We do +1 because the last thread block is used
    # to pad unused seq_lens as 0 for full CUDA graphs.
    _prepare_pos_seq_lens_kernel[(num_reqs + 1,)](
        pos,
        seq_lens,
        idx_mapping,
        query_start_loc,
        num_computed_tokens,
        seq_lens.shape[0],
        BLOCK_SIZE=1024,
    )


@triton.jit
def _combine_sampled_and_draft_tokens_kernel(
Woosuk Kwon's avatar
Woosuk Kwon committed
277
278
    input_ids_ptr,
    idx_mapping_ptr,
279
    last_sampled_tokens_ptr,
Woosuk Kwon's avatar
Woosuk Kwon committed
280
281
282
    query_start_loc_ptr,
    seq_lens_ptr,
    prefill_len_ptr,
283
284
285
286
287
    draft_tokens_ptr,
    draft_tokens_stride,
    cu_num_logits_ptr,
    logits_indices_ptr,
    BLOCK_SIZE: tl.constexpr,
Woosuk Kwon's avatar
Woosuk Kwon committed
288
289
290
291
):
    batch_idx = tl.program_id(0)
    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)

292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
    # Get the number of logits and draft tokens.
    cu_num_logits_start = tl.load(cu_num_logits_ptr + batch_idx)
    cu_num_logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
    num_logits = cu_num_logits_end - cu_num_logits_start
    num_draft_tokens = num_logits - 1

    # Compute the logits indices.
    block = tl.arange(0, BLOCK_SIZE)
    query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
    logits_start = query_end - num_logits
    tl.store(
        logits_indices_ptr + cu_num_logits_start + block,
        logits_start + block,
        mask=block < num_logits,
    )

Woosuk Kwon's avatar
Woosuk Kwon committed
308
309
310
    seq_len = tl.load(seq_lens_ptr + batch_idx)
    prefill_len = tl.load(prefill_len_ptr + req_state_idx)
    if seq_len <= prefill_len:
311
        # Handling prefill tokens. No sampled or draft tokens.
Woosuk Kwon's avatar
Woosuk Kwon committed
312
313
        return

314
    # Write the last sampled token ID to input_ids.
315
    last_token_id = tl.load(last_sampled_tokens_ptr + req_state_idx)
316
317
318
319
320
321
322
323
324
325
326
327
328
329
    tl.store(input_ids_ptr + query_end - num_logits, last_token_id)

    # Write the draft tokens (if any) to input_ids.
    if num_draft_tokens > 0:
        mask = block < num_draft_tokens
        draft_tokens = tl.load(
            draft_tokens_ptr + req_state_idx * draft_tokens_stride + block,
            mask=mask,
        )
        tl.store(
            input_ids_ptr + query_end - num_draft_tokens + block,
            draft_tokens,
            mask=mask,
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
330
331


332
def combine_sampled_and_draft_tokens(
Woosuk Kwon's avatar
Woosuk Kwon committed
333
334
    input_ids: torch.Tensor,
    idx_mapping: torch.Tensor,
335
    last_sampled_tokens: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
336
337
338
    query_start_loc: torch.Tensor,
    seq_lens: torch.Tensor,
    prefill_len: torch.Tensor,
339
340
341
    draft_tokens: torch.Tensor,
    cu_num_logits: torch.Tensor,
    num_logits: int,
Woosuk Kwon's avatar
Woosuk Kwon committed
342
343
) -> torch.Tensor:
    num_reqs = seq_lens.shape[0]
344
345
346
347
348
349
350
    num_speculative_steps = draft_tokens.shape[-1]

    logits_indices = torch.empty(
        num_logits,
        dtype=torch.int64,
        device=input_ids.device,
    )
351
    _combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
Woosuk Kwon's avatar
Woosuk Kwon committed
352
353
        input_ids,
        idx_mapping,
354
        last_sampled_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
355
356
357
        query_start_loc,
        seq_lens,
        prefill_len,
358
359
360
361
362
363
364
        draft_tokens,
        draft_tokens.stride(0),
        cu_num_logits,
        logits_indices,
        # NOTE(woosuk): Add 1 to ensure the block can cover the last sampled token
        # in addition to all draft tokens.
        BLOCK_SIZE=triton.next_power_of_2(num_speculative_steps + 1),
Woosuk Kwon's avatar
Woosuk Kwon committed
365
    )
366
    return logits_indices
367
368


369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
@triton.jit
def _get_num_sampled_and_rejected_kernel(
    num_sampled_ptr,
    num_rejected_ptr,
    seq_lens_ptr,
    cu_num_logits_ptr,
    idx_mapping_ptr,
    prefill_len_ptr,
):
    batch_idx = tl.program_id(0)
    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)

    seq_len = tl.load(seq_lens_ptr + batch_idx)
    prefill_len = tl.load(prefill_len_ptr + req_state_idx)
    is_chunked_prefilling = seq_len < prefill_len

    num_sampled = tl.load(num_sampled_ptr + batch_idx)
    num_sampled = tl.where(is_chunked_prefilling, 0, num_sampled)
    tl.store(num_sampled_ptr + batch_idx, num_sampled)

    logits_start = tl.load(cu_num_logits_ptr + batch_idx)
    logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
    num_logits = logits_end - logits_start

    num_rejected = num_logits - num_sampled
    num_rejected = tl.where(is_chunked_prefilling, 0, num_rejected)
    tl.store(num_rejected_ptr + batch_idx, num_rejected)


def get_num_sampled_and_rejected(
    num_sampled: torch.Tensor,
    seq_lens: torch.Tensor,
    cu_num_logits: torch.Tensor,
    idx_mapping: torch.Tensor,
    prefill_len: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    num_reqs = idx_mapping.shape[0]
    num_rejected = torch.empty_like(num_sampled)
    _get_num_sampled_and_rejected_kernel[(num_reqs,)](
        num_sampled,
        num_rejected,
        seq_lens,
        cu_num_logits,
        idx_mapping,
        prefill_len,
    )
    return num_sampled, num_rejected


418
@triton.jit
419
def _post_update_kernel(
420
421
    idx_mapping_ptr,
    num_computed_tokens_ptr,
422
    last_sampled_tokens_ptr,
423
424
    output_bin_counts_ptr,
    output_bin_counts_stride,
425
426
427
    sampled_tokens_ptr,
    sampled_tokens_stride,
    num_sampled_ptr,
428
    num_rejected_ptr,
429
    query_start_loc_ptr,
430
431
432
    all_token_ids_ptr,
    all_token_ids_stride,
    total_len_ptr,
433
434
435
436
):
    req_id = tl.program_id(0)
    req_state_idx = tl.load(idx_mapping_ptr + req_id)

437
    total_len = tl.load(total_len_ptr + req_state_idx)
438
439
440
441
442
443
    num_sampled = tl.load(num_sampled_ptr + req_id)
    if num_sampled > 0:
        token_id = tl.load(
            sampled_tokens_ptr + req_id * sampled_tokens_stride + num_sampled - 1
        )
        tl.store(last_sampled_tokens_ptr + req_state_idx, token_id)
444
        tl.store(total_len_ptr + req_state_idx, total_len + num_sampled)
445

446
447
448
449
450
451
452
453
    for i in range(num_sampled):
        token_id = tl.load(sampled_tokens_ptr + req_id * sampled_tokens_stride + i)
        token_ptr = (
            output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + token_id
        )
        count = tl.load(token_ptr)
        count += 1
        tl.store(token_ptr, count)
454
455
456
457
        tl.store(
            all_token_ids_ptr + req_state_idx * all_token_ids_stride + total_len + i,
            token_id,
        )
458

459
460
461
    query_start = tl.load(query_start_loc_ptr + req_id)
    query_end = tl.load(query_start_loc_ptr + req_id + 1)
    query_len = query_end - query_start
462
    num_rejected = tl.load(num_rejected_ptr + req_id)
463
464

    num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
465
    num_computed += query_len - num_rejected
466
467
468
469
470
    tl.store(num_computed_tokens_ptr + req_state_idx, num_computed)


def post_update(
    # [num_reqs]
471
    idx_mapping: torch.Tensor,
472
    # [max_num_reqs]
473
    num_computed_tokens: torch.Tensor,
474
475
    # [max_num_reqs]
    last_sampled_tokens: torch.Tensor,
476
477
    # [max_num_reqs, vocab_size]
    output_bin_counts: torch.Tensor,
478
479
480
481
    # [num_reqs, num_speculative_steps + 1]
    sampled_tokens: torch.Tensor,
    # [num_reqs]
    num_sampled: torch.Tensor,
482
483
    # [num_reqs]
    num_rejected: torch.Tensor,
484
    # [num_reqs + 1]
485
    query_start_loc: torch.Tensor,
486
487
488
489
    # [max_num_reqs, max_model_len]
    all_token_ids: torch.Tensor,
    # [max_num_reqs]
    total_len: torch.Tensor,
490
491
) -> None:
    num_reqs = idx_mapping.shape[0]
492
    _post_update_kernel[(num_reqs,)](
493
494
        idx_mapping,
        num_computed_tokens,
495
        last_sampled_tokens,
496
497
        output_bin_counts,
        output_bin_counts.stride(0),
498
499
500
        sampled_tokens,
        sampled_tokens.stride(0),
        num_sampled,
501
        num_rejected,
502
        query_start_loc,
503
504
505
        all_token_ids,
        all_token_ids.stride(0),
        total_len,
506
        num_warps=1,
507
    )
508
509
510
511
512
513


@triton.jit
def _expand_idx_mapping_kernel(
    idx_mapping_ptr,
    expanded_idx_mapping_ptr,
514
    expanded_local_pos_ptr,
515
516
517
518
519
520
521
522
523
524
525
526
    cu_num_logits_ptr,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)
    start_idx = tl.load(cu_num_logits_ptr + req_idx)
    end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
    num_tokens = end_idx - start_idx

    block = tl.arange(0, BLOCK_SIZE)
    mask = block < num_tokens
    req_state_idx = tl.load(idx_mapping_ptr + req_idx)
    tl.store(expanded_idx_mapping_ptr + start_idx + block, req_state_idx, mask=mask)
527
    tl.store(expanded_local_pos_ptr + start_idx + block, block, mask=mask)
528
529
530
531
532
533
534


def expand_idx_mapping(
    idx_mapping: torch.Tensor,
    total_num_logits: int,
    cu_num_logits: torch.Tensor,
    max_expand_len: int,
535
) -> tuple[torch.Tensor, torch.Tensor]:
536
537
    num_reqs = idx_mapping.shape[0]
    expanded_idx_mapping = idx_mapping.new_empty(total_num_logits)
538
539
540
    expanded_local_pos = torch.empty(
        total_num_logits, dtype=torch.int32, device=idx_mapping.device
    )
541
542
543
    _expand_idx_mapping_kernel[(num_reqs,)](
        idx_mapping,
        expanded_idx_mapping,
544
        expanded_local_pos,
545
546
547
        cu_num_logits,
        BLOCK_SIZE=triton.next_power_of_2(max_expand_len),
    )
548
    return expanded_idx_mapping, expanded_local_pos