flash_causal_lm.py 44.3 KB
Newer Older
1
import math
2
import os
3
import time
4
import itertools
5
6
7
import torch
import torch.distributed

8
9
import numpy as np

10
from loguru import logger
11
12
from dataclasses import dataclass
from opentelemetry import trace
13
from transformers import PreTrainedTokenizerBase
14
from typing import Optional, Tuple, List, Type, Dict
15

OlivierDehaene's avatar
OlivierDehaene committed
16
from text_generation_server.models import Model
17
from text_generation_server.utils.tokens import batch_top_tokens
Nicolas Patry's avatar
Nicolas Patry committed
18
from text_generation_server.utils.speculate import get_speculate
19
20
from text_generation_server.models.types import (
    Batch,
Nicolas Patry's avatar
Nicolas Patry committed
21
    Tokens,
22
23
24
    Generation,
    GeneratedText,
)
25
26
27
28
29
from text_generation_server.models.cache_manager import (
    get_cache_manager,
    set_cache_manager,
    BLOCK_SIZE,
)
30
from text_generation_server.pb import generate_pb2
31
from text_generation_server.models.globals import MEM_POOL
32
from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
33
from text_generation_server.utils.dist import MEMORY_FRACTION
34
35
36

tracer = trace.get_tracer(__name__)

37

38
39
40
41
42

@dataclass
class FlashCausalLMBatch(Batch):
    batch_id: int
    requests: List[generate_pb2.Request]
43
44
    # request id -> idx in list mapping
    requests_idx_mapping: Dict[int, int]
45
46

    # Decoder values
47
48
    input_ids: torch.Tensor
    position_ids: torch.Tensor
Nicolas Patry's avatar
Nicolas Patry committed
49
    speculative_ids: torch.Tensor
50

51
52
53
54
    # Flash Attention values

    # tensor of length b containing the cumulative sequence lengths of the sequences in the batch, only used in prefill
    cu_seqlen_prefill: Optional[torch.Tensor]
55
56
57
58
59
60
61
62
63
64
65
66
67
68

    # Paged Attention values

    # Set when creating the batch
    # CPU tensor of length b indicating the start of each sequence in slots
    start_slots: torch.Tensor
    # tensor of indices of the currently used slots, length = \sum_{i=0}^{b} s_i in prefill, length = b in decode
    slot_indices: torch.Tensor
    # List of tuple of ints representing the number of blocks and slots needed by each sequence
    needed_blocks_slots: Optional[List[Tuple[int, int]]]

    # Set in prefill by the CacheManager
    # list of length b of list of length s_i // block_size
    block_tables: Optional[List[List[int]]]
69
    # tensor of size [b, max_total_seqlen // block_size] holding the paged attention block tables for all sequences
70
71
72
73
    block_tables_tensor: Optional[torch.Tensor]
    # tensor of length \sum_{i=0}^{b} max_s_i  holding the paged attention slots for all sequences
    slots: Optional[torch.Tensor]

74
75
    max_seqlen: int

76
77
78
79
80
    # Prefill metadata tensors to efficiently compute logprobs
    prefill_head_indices: Optional[torch.Tensor]
    prefill_next_token_indices: Optional[torch.tensor]
    prefill_cu_outlens: Optional[List[int]]

81
82
    # All tokens
    all_input_ids: List[List[int]]
83
    all_input_ids_tensor: torch.Tensor
84
85
86

    # Lengths of all generations present in the batch
    input_lengths: List[int]
87
    input_lengths_tensor: torch.Tensor
88
89
    prefix_offsets: List[Optional[int]]
    read_offsets: List[Optional[int]]
90
91

    # Generation helpers
92
    next_token_chooser: HeterogeneousNextTokenChooser
93
    stopping_criterias: List[StoppingCriteria]
Nicolas Patry's avatar
Nicolas Patry committed
94
95
    top_n_tokens: List[int]
    top_n_tokens_tensor: torch.Tensor
96

97
98
99
100
    # Number of blocks in this batch
    blocks: int
    # Maximum number of blocks
    max_blocks: int
101

102
103
    def to_pb(self) -> generate_pb2.CachedBatch:
        return generate_pb2.CachedBatch(
104
            id=self.batch_id,
105
            request_ids=[r.id for r in self.requests],
106
            size=len(self),
107
            max_tokens=self.blocks * BLOCK_SIZE,
108
109
110
111
112
113
114
        )

    @classmethod
    def from_pb(
        cls,
        pb: generate_pb2.Batch,
        tokenizer: PreTrainedTokenizerBase,
115
        dtype: torch.dtype,
116
        device: torch.device,
117
    ) -> "FlashCausalLMBatch":
118
119
120
121
122
123
124
125
126
127
        batch_inputs = []
        max_truncation = 0
        for r in pb.requests:
            batch_inputs.append(r.inputs)
            max_truncation = max(max_truncation, r.truncate)

        batch_tokenized_inputs = tokenizer(
            batch_inputs, truncation=True, max_length=max_truncation
        )["input_ids"]

128
        position_ids = []
Nicolas Patry's avatar
Nicolas Patry committed
129
        speculative_ids = []
130
        cu_seqlen_prefill = [0]
131
132
133
        needed_blocks_slots = []
        start_slots = []
        slot_indices = []
134
135

        input_lengths = []
136
137
        prefix_offsets = []
        read_offsets = []
138
        all_input_ids = []
139
        requests_idx_mapping = {}
140

141
142
143
144
145
146
        all_prefill_logprobs = True
        no_prefill_logprobs = True
        prefill_head_indices = []
        prefill_next_token_indices = []
        prefill_cu_outlens = [0]

147
        next_token_chooser_parameters = []
148
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
149
        top_n_tokens = []
150
151
152

        # Cumulative length
        cumulative_length = 0
153
        cumulative_max_length = 0
154
        prefill_out_cumulative_length = 0
155

156
157
        blocks = 0
        max_seqlen = 0
158
        max_length = 0
159
        max_blocks = 0
160

161
        # Parse batch
162
163
164
        for i, (r, tokenized_input) in enumerate(
            zip(pb.requests, batch_tokenized_inputs)
        ):
165
166
167
            # request id -> idx in list mapping
            requests_idx_mapping[r.id] = i

168
            tokenized_input = tokenized_input[-r.truncate :]
169

170
171
            input_length = len(tokenized_input)
            input_lengths.append(input_length)
172

173
            prefix_offsets.append(input_length - 5)
174
            read_offsets.append(input_length)
175

176
            all_input_ids.append(tokenized_input)
177
178

            # Position ids
179
180
            request_position_ids = torch.arange(0, input_length, dtype=torch.int32)
            position_ids.append(request_position_ids)
181
182

            # Add cumulative lengths of all previous inputs
183
            cu_seqlen_prefill.append(cumulative_length + input_length)
184

185
            next_token_chooser_parameters.append(r.parameters)
186

187
188
189
            stopping_criteria = StoppingCriteria.from_pb(
                r.stopping_parameters, tokenizer
            )
190
            max_new_tokens = stopping_criteria.max_new_tokens
191
            stopping_criterias.append(stopping_criteria)
Nicolas Patry's avatar
Nicolas Patry committed
192
            top_n_tokens.append(r.top_n_tokens)
193

194
195
            # Paged attention
            # Remove one as the first token des not have a past
Nicolas Patry's avatar
Nicolas Patry committed
196
197
            speculative_length = get_speculate()
            total_tokens = input_length + max_new_tokens - 1 + speculative_length
198
199
200
201
202
203
204
205
206
207
208
209
            needed_blocks = math.ceil(total_tokens / BLOCK_SIZE)
            blocks += needed_blocks
            needed_blocks_slots.append((needed_blocks, total_tokens))
            start_slots.append(cumulative_max_length)

            request_slot_indices = torch.arange(
                cumulative_max_length,
                cumulative_max_length + input_length,
                dtype=torch.int64,
            )
            slot_indices.append(request_slot_indices)

210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
            all_prefill_logprobs = all_prefill_logprobs and r.prefill_logprobs
            no_prefill_logprobs = no_prefill_logprobs and not r.prefill_logprobs

            if r.prefill_logprobs:
                prefill_head_indices.append(request_position_ids + cumulative_length)
                prefill_next_token_indices.append(
                    prefill_out_cumulative_length + input_length - 1
                )
                prefill_cu_outlens.append(prefill_out_cumulative_length + input_length)
                prefill_out_cumulative_length += input_length
            else:
                prefill_head_indices.append(
                    torch.tensor(
                        [cumulative_length + input_length - 1], dtype=torch.int32
                    )
                )
                prefill_next_token_indices.append(prefill_out_cumulative_length)
                prefill_cu_outlens.append(prefill_out_cumulative_length + 1)
                prefill_out_cumulative_length += 1

230
231
            # Update
            cumulative_length += input_length
232
233
234
            cumulative_max_length += total_tokens
            max_seqlen = max(max_seqlen, input_length)
            max_blocks = max(max_blocks, needed_blocks)
OlivierDehaene's avatar
OlivierDehaene committed
235
236
237
            max_length = max(
                max_length, input_length + max_new_tokens + speculative_length
            )
238
239
240
241

        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
            next_token_chooser_parameters, dtype, device
        )
242
        start_slots = torch.tensor(start_slots, dtype=torch.int64)
243
244
245
246
247
248
249

        # Padded all_input_ids_tensor
        all_input_ids_tensor = np.zeros(
            (len(all_input_ids), max_length), dtype=np.int64
        )
        for i, input_ids in enumerate(all_input_ids):
            all_input_ids_tensor[i, : len(input_ids)] = input_ids
250

251
252
253
254
255
        # Create tensors on device
        all_input_ids_tensor = torch.tensor(
            all_input_ids_tensor, dtype=torch.int64, device=device
        )

256
257
258
        if len(pb.requests) > 1:
            input_ids = np.concatenate(all_input_ids, dtype=np.int64)
            position_ids = torch.cat(position_ids)
259
            slot_indices = torch.cat(slot_indices)
260
261
262
        else:
            input_ids = all_input_ids[0]
            position_ids = position_ids[0]
263
            slot_indices = slot_indices[0]
264

265
266
        cu_seqlen_prefill = torch.tensor(
            cu_seqlen_prefill, device=device, dtype=torch.int32
267
268
269
        )
        position_ids = position_ids.to(device)
        slot_indices = slot_indices.to(device)
270
        input_ids = torch.tensor(input_ids, dtype=torch.int64, device=device)
271
272
        input_lengths_tensor = torch.tensor(
            input_lengths, dtype=torch.int32, device=device
273
        )
274

275
276
        if all_prefill_logprobs:
            prefill_head_indices = None
277
            prefill_next_token_indices = cu_seqlen_prefill[1:] - 1
278
        elif no_prefill_logprobs:
279
            prefill_head_indices = cu_seqlen_prefill[1:] - 1
280
281
282
283
284
285
286
287
            prefill_next_token_indices = None
        else:
            prefill_head_indices = torch.tensor(
                torch.cat(prefill_head_indices), dtype=torch.int64, device=device
            )
            prefill_next_token_indices = torch.tensor(
                prefill_next_token_indices, dtype=torch.int64, device=device
            )
Nicolas Patry's avatar
Nicolas Patry committed
288
289
290
        top_n_tokens_tensor = torch.tensor(
            top_n_tokens, device=device, dtype=torch.int64
        )
291

292
293
294
        return cls(
            batch_id=pb.id,
            requests=pb.requests,
295
            requests_idx_mapping=requests_idx_mapping,
296
297
            input_ids=input_ids,
            position_ids=position_ids,
298
            cu_seqlen_prefill=cu_seqlen_prefill,
299
300
301
302
303
304
            start_slots=start_slots,
            slot_indices=slot_indices,
            needed_blocks_slots=needed_blocks_slots,
            block_tables=None,
            block_tables_tensor=None,
            slots=None,
305
            max_seqlen=max_seqlen,
306
307
308
            prefill_head_indices=prefill_head_indices,
            prefill_next_token_indices=prefill_next_token_indices,
            prefill_cu_outlens=prefill_cu_outlens,
309
            input_lengths=input_lengths,
310
            input_lengths_tensor=input_lengths_tensor,
311
312
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
313
            all_input_ids=all_input_ids,
314
315
            all_input_ids_tensor=all_input_ids_tensor,
            next_token_chooser=next_token_chooser,
316
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
317
318
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
319
320
            blocks=blocks,
            max_blocks=max_blocks,
Nicolas Patry's avatar
Nicolas Patry committed
321
            speculative_ids=None,
322
323
        )

324
    @tracer.start_as_current_span("filter")
325
326
    def filter(self, request_ids: List[int]) -> "FlashCausalLMBatch":
        if len(request_ids) == 0:
327
328
            raise ValueError("Batch must have at least one request")
        # We assume that if len(requests) == len(self) then the requests are the same
329
        if len(request_ids) == len(self):
330
331
            return self

332
        device = self.input_ids.device
333

334
335
336
        # New values after filtering
        requests_idx_mapping = {}

337
338
339
        # Used to index into tensors
        indices = []

340
341
342
        # slots to keep after filtering
        slot_filtering_indices = torch.zeros(
            self.slots.shape[0], dtype=torch.bool, device=device
343
344
        )

345
        # Create on CPU to only move to GPU once instead of at every copy
346
        slot_indices = torch.empty(len(request_ids), dtype=torch.int64)
347
348
        max_seqlen = 0

349
        requests = []
350
351
        start_slots = []
        block_tables = []
352
353
        all_input_ids = []

354
        input_lengths = []
355
356
        prefix_offsets = []
        read_offsets = []
357

358
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
359
        top_n_tokens = []
360

361
362
363
364
365
        blocks = 0
        max_blocks = 0
        # Cumulative length
        cumulative_max_length = 0

366
367
        for i, request_id in enumerate(request_ids):
            idx = self.requests_idx_mapping[request_id]
368
            indices.append(idx)
369
370
371
            requests_idx_mapping[request_id] = i

            requests.append(self.requests[idx])
372
373
374
375

            # Get length
            request_input_length = self.input_lengths[idx]
            max_seqlen = max(max_seqlen, request_input_length)
376

377
378
379
            all_input_ids.append(self.all_input_ids[idx])

            input_lengths.append(request_input_length)
380
381
            prefix_offsets.append(self.prefix_offsets[idx])
            read_offsets.append(self.read_offsets[idx])
382

383
384
            stopping_criteria = self.stopping_criterias[idx]
            stopping_criterias.append(stopping_criteria)
385

Nicolas Patry's avatar
Nicolas Patry committed
386
387
            top_n_tokens.append(self.top_n_tokens[idx])

388
            remaining_tokens = (
389
390
                stopping_criteria.max_new_tokens - stopping_criteria.current_tokens
            )
391

392
393
394
395
396
            request_block_table = self.block_tables[idx]
            blocks += len(request_block_table)
            block_tables.append(request_block_table)
            start_slots.append(cumulative_max_length)

397
            # Copy to tensor (CPU)
398
            slot_indices[i] = cumulative_max_length + request_input_length - 1
399
400

            # Set slice
401
402
403
404
405
            slot_filtering_indices[
                self.start_slots[idx] : self.start_slots[idx]
                + request_input_length
                + remaining_tokens
                - 1
406
407
408
            ] = True

            cumulative_max_length += request_input_length + remaining_tokens - 1
409

410
411
412
413
414
415
416
417
418
            max_blocks = max(max_blocks, len(request_block_table))

        block_indices_to_free = []
        # Iterate on all requests
        for i, r in enumerate(self.requests):
            # Filter requests that are not part of the new batch
            if r.id not in requests_idx_mapping.keys():
                block_indices_to_free.extend(self.block_tables[i])
        # Free blocks
419
        get_cache_manager().free(block_indices_to_free)
420
421
422
        # Needed to avoid dropping blocks when the batches will go out of scope
        self.block_tables = None

423
424
425
426
        # Index into tensors
        input_ids = self.input_ids[indices]
        position_ids = self.position_ids[indices]
        all_input_ids_tensor = self.all_input_ids_tensor[indices]
427
428
429
        block_tables_tensor = self.block_tables_tensor[indices]
        input_lengths_tensor = self.input_lengths_tensor[indices]
        slots = self.slots[slot_filtering_indices]
430
        next_token_chooser = self.next_token_chooser.filter(indices)
Nicolas Patry's avatar
Nicolas Patry committed
431
        top_n_tokens_tensor = self.top_n_tokens_tensor[indices]
OlivierDehaene's avatar
OlivierDehaene committed
432
433
434
        speculative_ids = (
            self.speculative_ids[indices] if self.speculative_ids is not None else None
        )
435
436

        start_slots = torch.tensor(start_slots, dtype=torch.int64)
437

438
        # Move to GPU now that we have the whole tensor
439
        slot_indices = slot_indices.to(device)
440

441
        return type(self)(
442
443
444
445
446
            batch_id=self.batch_id,
            requests=requests,
            requests_idx_mapping=requests_idx_mapping,
            input_ids=input_ids,
            position_ids=position_ids,
447
            cu_seqlen_prefill=None,
448
449
450
451
452
453
            start_slots=start_slots,
            slot_indices=slot_indices,
            needed_blocks_slots=None,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
454
            max_seqlen=max_seqlen,
455
456
457
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
458
            input_lengths=input_lengths,
459
            input_lengths_tensor=input_lengths_tensor,
460
461
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
462
463
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
464
            next_token_chooser=next_token_chooser,
465
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
466
467
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
468
469
            blocks=blocks,
            max_blocks=max_blocks,
Nicolas Patry's avatar
Nicolas Patry committed
470
            speculative_ids=speculative_ids,
471
472
473
474
475
476
477
478
479
        )

    @classmethod
    @tracer.start_as_current_span("concatenate")
    def concatenate(cls, batches: List["FlashCausalLMBatch"]) -> "FlashCausalLMBatch":
        # Batch attributes
        requests = []
        requests_idx_mapping = {}

480
481
482
483
484
485
486
487
488
489
        blocks = 0
        total_batch_size = 0
        total_slots = 0
        max_blocks = 0
        max_length = 0
        max_seqlen = 0
        for b in batches:
            total_batch_size += len(b)
            total_slots += len(b.slots)
            blocks += b.blocks
OlivierDehaene's avatar
OlivierDehaene committed
490
491
492
            speculative_length = (
                b.speculative_ids.shape[1] if b.speculative_ids is not None else 0
            )
493
494
495
496
497
498
499
            max_blocks = max(max_blocks, b.max_blocks)
            max_seqlen = max(max_seqlen, b.max_seqlen)
            max_length = max(
                max_length,
                max(
                    input_length
                    + stopping_criteria.max_new_tokens
Nicolas Patry's avatar
Nicolas Patry committed
500
                    + speculative_length
501
502
503
504
505
506
                    - stopping_criteria.current_tokens
                    for input_length, stopping_criteria in zip(
                        b.input_lengths, b.stopping_criterias
                    )
                ),
            )
507
508
509

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
510
511
512
513
514
515
516
517
518
519
        slots = batches[0].slots.new_empty(total_slots)
        slot_indices = batches[0].slot_indices.new_empty(total_batch_size)
        input_lengths_tensor = batches[0].input_lengths_tensor.new_empty(
            total_batch_size
        )
        block_tables_tensor = batches[0].block_tables_tensor.new_zeros(
            (total_batch_size, max_blocks)
        )
        all_input_ids_tensor = batches[0].all_input_ids_tensor.new_zeros(
            (total_batch_size, max_length)
520
        )
Nicolas Patry's avatar
Nicolas Patry committed
521
522
523
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
524

525
526
        start_slots = []
        block_tables = []
527
528
529
        all_input_ids = []

        input_lengths = []
530
531
        prefix_offsets = []
        read_offsets = []
532

533
        next_token_chooser_parameters = []
534
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
535
        top_n_tokens = []
536

537
        # Cumulative length
538
        cumulative_batch_size = 0
539
        cumulative_slots = 0
540
541
542

        for i, batch in enumerate(batches):
            requests.extend(batch.requests)
543
544
545
546
547
548
549
550

            if i == 0:
                requests_idx_mapping = batch.requests_idx_mapping
            else:
                # We need to offset the mapping for each batch by the cumulative batch size
                for k, v in batch.requests_idx_mapping.items():
                    requests_idx_mapping[k] = v + cumulative_batch_size

551
552
            start_index = cumulative_batch_size
            end_index = cumulative_batch_size + len(batch)
553
554
            slots_start_index = cumulative_slots
            slots_end_index = cumulative_slots + len(batch.slots)
555
556
557
558

            # Copy tensors (GPU)
            input_ids[start_index:end_index] = batch.input_ids
            position_ids[start_index:end_index] = batch.position_ids
559
560
            slot_indices[start_index:end_index] = batch.slot_indices + cumulative_slots
            input_lengths_tensor[start_index:end_index] = batch.input_lengths_tensor
Nicolas Patry's avatar
Nicolas Patry committed
561
            top_n_tokens_tensor[start_index:end_index] = batch.top_n_tokens_tensor
562
            slots[slots_start_index:slots_end_index] = batch.slots
563

564
565
566
            all_input_ids_tensor[
                start_index:end_index, : batch.all_input_ids_tensor.shape[1]
            ] = batch.all_input_ids_tensor[:, :max_length]
567

568
569
570
            block_tables_tensor[
                start_index:end_index, : batch.block_tables_tensor.shape[1]
            ] = batch.block_tables_tensor[:, :max_blocks]
571

572
573
574
            start_slots.append(batch.start_slots + cumulative_slots)

            block_tables.extend(batch.block_tables)
575
576
            all_input_ids.extend(batch.all_input_ids)

577
            input_lengths.extend(batch.input_lengths)
578
579
            prefix_offsets.extend(batch.prefix_offsets)
            read_offsets.extend(batch.read_offsets)
580

581
            next_token_chooser_parameters.extend([r.parameters for r in batch.requests])
582
583
            stopping_criterias.extend(batch.stopping_criterias)

Nicolas Patry's avatar
Nicolas Patry committed
584
585
            top_n_tokens.extend(batch.top_n_tokens)

586
            # Update
587
            cumulative_batch_size += len(batch)
588
            cumulative_slots += len(batch.slots)
589

590
        start_slots = torch.concat(start_slots)
591

592
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
593
594
595
            next_token_chooser_parameters,
            dtype=batches[0].next_token_chooser.dtype,
            device=batches[0].next_token_chooser.device,
596
597
        )

OlivierDehaene's avatar
OlivierDehaene committed
598
599
600
601
602
        speculative_ids = (
            torch.cat([b.speculative_ids for b in batches], dim=0)
            if batches[0].speculative_ids is not None
            else None
        )
Nicolas Patry's avatar
Nicolas Patry committed
603

604
605
606
        # Needed to avoid dropping blocks when the batches will go out of scope
        for b in batches:
            b.block_tables = None
607
            del b
608

609
        return cls(
610
611
            batch_id=batches[0].batch_id,
            requests=requests,
612
            requests_idx_mapping=requests_idx_mapping,
613
614
            input_ids=input_ids,
            position_ids=position_ids,
615
            cu_seqlen_prefill=None,
616
617
618
619
620
621
            start_slots=start_slots,
            slot_indices=slot_indices,
            needed_blocks_slots=None,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
622
            max_seqlen=max_seqlen,
623
624
625
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
626
            input_lengths=input_lengths,
627
            input_lengths_tensor=input_lengths_tensor,
628
629
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
630
631
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
632
            next_token_chooser=next_token_chooser,
633
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
634
635
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
636
637
            blocks=blocks,
            max_blocks=max_blocks,
OlivierDehaene's avatar
OlivierDehaene committed
638
            speculative_ids=speculative_ids,
639
640
        )

641
642
643
    def __del__(self):
        if self.block_tables is not None and self.block_tables:
            # Free blocks
644
645
646
            get_cache_manager().free(
                list(itertools.chain.from_iterable(self.block_tables))
            )
647

648
649
650
651
652
653
654
    def __len__(self):
        return len(self.requests)


class FlashCausalLM(Model):
    def __init__(
        self,
655
656
657
658
659
660
661
662
663
        model: torch.nn.Module,
        tokenizer: PreTrainedTokenizerBase,
        num_layers: int,
        num_kv_heads: int,
        head_size: int,
        dtype: torch.dtype,
        device: torch.device,
        rank: int = 0,
        world_size: int = 1,
664
        sliding_window: Optional[int] = None,
665
    ):
666
667
668
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_size = head_size
669

670
671
        self.cuda_graphs = {}

672
        super(FlashCausalLM, self).__init__(
673
            model=model,
674
675
676
677
            tokenizer=tokenizer,
            requires_padding=False,
            dtype=dtype,
            device=device,
678
679
            rank=rank,
            world_size=world_size,
680
            sliding_window=sliding_window,
681
682
683
684
685
686
        )

    @property
    def batch_type(self) -> Type[FlashCausalLMBatch]:
        return FlashCausalLMBatch

687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
    def cuda_graph_warmup(self, bs: int, max_s: int, max_bt: int):
        input_ids = torch.zeros(bs, dtype=torch.int64, device=self.device)
        position_ids = torch.zeros(bs, dtype=torch.int32, device=self.device)
        slots = torch.arange(bs, dtype=torch.int32, device=self.device)
        input_lengths = torch.ones(bs, dtype=torch.int32, device=self.device) * max_s
        block_tables = (
            torch.arange(max_bt, dtype=torch.int32, device=self.device)
            .repeat(bs)
            .reshape((bs, max_bt))
        )
        kv_cache = get_cache_manager().kv_cache

        self.cuda_graphs[bs] = {
            "input_ids": input_ids,
            "position_ids": position_ids,
            "kv_cache": kv_cache,
            "block_tables": block_tables,
            "slots": slots,
            "input_lengths": input_lengths,
        }
        graph = torch.cuda.CUDAGraph()
        self.cuda_graphs[bs]["graph"] = graph

        torch.cuda.synchronize()
        # Run once outside to warmup
        self.model.forward(
            input_ids=input_ids,
            position_ids=position_ids,
            cu_seqlen_prefill=None,
            kv_cache=kv_cache,
            block_tables=block_tables,
            slots=slots,
            input_lengths=input_lengths,
            max_s=max_s,
            lm_head_indices=None,
        )
        torch.cuda.synchronize()

        with torch.cuda.graph(graph, pool=MEM_POOL):
            self.cuda_graphs[bs]["logits"] = self.model.forward(
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlen_prefill=None,
                kv_cache=kv_cache,
                block_tables=block_tables,
                slots=slots,
                input_lengths=input_lengths,
                max_s=max_s,
                lm_head_indices=None,
            )
        torch.cuda.synchronize()

739
    def warmup(self, batch: FlashCausalLMBatch):
740
        # The warmup batch is the biggest batch we could ever receive
741
742
        torch.cuda.empty_cache()
        try:
743
            cache_manager = set_cache_manager(
744
                batch.blocks,
745
746
747
                self.num_layers,
                self.num_kv_heads,
                self.head_size,
748
                self.sliding_window is not None,
749
750
751
                self.dtype,
                self.device,
            )
752
753
            max_bt = batch.max_blocks
            max_s = max_bt * get_cache_manager().block_size
754
            _, batch, _ = self.generate_token(batch)
OlivierDehaene's avatar
OlivierDehaene committed
755
        except torch.cuda.OutOfMemoryError as e:
756
            raise RuntimeError(
757
758
                f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. "
                f"You need to decrease `--max-batch-prefill-tokens`"
759
            ) from e
760
761
762

        torch.cuda.synchronize(self.device)

763
764
        # Inspired by the original implementation in [vllm](https://github.com/vllm-project/vllm)
        # Calculate the number of blocks that can be allocated with the free memory
765
766
767
768
        dtype_size = torch.tensor([], dtype=self.dtype).element_size()
        cache_block_size = BLOCK_SIZE * self.num_kv_heads * self.head_size
        total_cache_size = self.num_layers * cache_block_size * 2 * dtype_size

769
770
771
772
773
774
        total_free_memory, _ = torch.cuda.mem_get_info(self.device)
        total_gpu_memory = torch.cuda.get_device_properties(self.device).total_memory

        free_memory = max(
            0, total_free_memory - (1 - MEMORY_FRACTION) * total_gpu_memory
        )
775
776

        num_blocks = (
777
778
            # Leave 5% for some wiggle room
            int((free_memory * 0.95) // total_cache_size)
779
            # Add batch.blocks as we allocated it above, so it is included in the peak memory.
780
            + cache_manager.num_blocks
781
782
        )

783
        del batch
784
        del cache_manager
785

786
        set_cache_manager(
787
788
789
790
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
791
            self.sliding_window is not None,
792
793
794
795
            self.dtype,
            self.device,
        )

796
797
798
799
800
801
802
803
804
805
        if os.getenv("ENABLE_CUDA_GRAPHS", "False") == "True":
            try:
                logger.info("Experimental support for Cuda Graphs is enabled")
                # Warmup cuda graphs
                for bs in [1, 2, 4] + [8 * i for i in range(8)]:
                    if self.speculate is None or self.speculate + 1 <= bs:
                        self.cuda_graph_warmup(bs, max_s, max_bt)
            except Exception:
                logger.exception(f"Decode cuda graph warmup failed")

806
        return int(num_blocks * BLOCK_SIZE)
807

808
    def forward(self, batch: FlashCausalLMBatch) -> torch.Tensor:
809
        # Model Forward
Nicolas Patry's avatar
Nicolas Patry committed
810
        if batch.speculative_ids is not None:
OlivierDehaene's avatar
OlivierDehaene committed
811
812
813
814
815
816
817
818
819
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
            kv_cache = get_cache_manager().kv_cache
            block_tables = batch.block_tables_tensor
            slots = batch.slots[batch.slot_indices]
            input_lengths = batch.input_lengths_tensor
            max_s = batch.max_seqlen
            lm_head_indices = batch.prefill_head_indices
Nicolas Patry's avatar
Nicolas Patry committed
820
821
822

            speculative_ids = batch.speculative_ids

OlivierDehaene's avatar
OlivierDehaene committed
823
            B, speculative_length = speculative_ids.shape
Nicolas Patry's avatar
Nicolas Patry committed
824
            new_length = speculative_length + 1
OlivierDehaene's avatar
OlivierDehaene committed
825
826
827
            new_input_ids = torch.cat(
                [input_ids.unsqueeze(-1), speculative_ids], dim=1
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
828
829
            arange = torch.arange(new_length, device=position_ids.device).unsqueeze(0)
            arange_int = arange.to(dtype=torch.int32)
OlivierDehaene's avatar
OlivierDehaene committed
830
831
832
            new_position_ids = (
                position_ids.unsqueeze(-1).expand(B, new_length) + arange
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
833
            slots = (slots.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1)
OlivierDehaene's avatar
OlivierDehaene committed
834
835
836
            input_lengths = (
                input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
837
838

            # Add Copy the block tables for all members
OlivierDehaene's avatar
OlivierDehaene committed
839
840
841
842
843
844
            block_tables = (
                block_tables.unsqueeze(1)
                .expand(B, new_length, -1)
                .reshape(B * new_length, -1)
                .contiguous()
            )
Nicolas Patry's avatar
Nicolas Patry committed
845
846
847
848
849
            max_s = max_s + speculative_length

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
850
851
852
853
854
855
856
857
858
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
            kv_cache = get_cache_manager().kv_cache
            block_tables = batch.block_tables_tensor
            slots = batch.slots[batch.slot_indices]
            input_lengths = batch.input_lengths_tensor
            max_s = batch.max_seqlen
            lm_head_indices = batch.prefill_head_indices
Nicolas Patry's avatar
Nicolas Patry committed
859

860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
        bs = input_ids.shape[0]
        padded_bs = bs
        if bs == 3:
            padded_bs = 4
        elif 3 < bs <= 8:
            padded_bs = 8
        elif bs > 8:
            padded_bs = (bs + 7) // 8 * 8

        # Try to find an associated cuda graph
        cuda_graph = self.cuda_graphs.get(padded_bs, None)

        if cu_seqlen_prefill is not None or cuda_graph is None or batch.speculative_ids is not None:
            return self.model.forward(
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlen_prefill=cu_seqlen_prefill,
                kv_cache=kv_cache,
                block_tables=block_tables,
                slots=slots,
                input_lengths=input_lengths,
                max_s=max_s,
                lm_head_indices=lm_head_indices,
            )

        # Copy inputs to the static inputs of the cuda graph
        # Static inputs are potentially padded
        cuda_graph["input_ids"][: input_ids.shape[0]] = input_ids
        cuda_graph["position_ids"][: position_ids.shape[0]] = position_ids
        cuda_graph["block_tables"][
            : block_tables.shape[0], : block_tables.shape[1]
        ] = block_tables
        cuda_graph["slots"].fill_(-1)
        cuda_graph["slots"][: slots.shape[0]] = slots
        cuda_graph["input_lengths"].zero_()
        cuda_graph["input_lengths"][: input_lengths.shape[0]] = input_lengths

        # Replay the graph
        cuda_graph["graph"].replay()

        # Slice output to the correct shape
        return cuda_graph["logits"][:bs]
902
903
904
905

    @tracer.start_as_current_span("generate_token")
    def generate_token(
        self, batch: FlashCausalLMBatch
906
907
    ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]:
        start = time.time_ns()
908
        prefill = batch.cu_seqlen_prefill is not None
909
        prefill_logprobs = batch.prefill_next_token_indices is not None
910

911
912
        if batch.needed_blocks_slots:
            # Allocate blocks to this batch
913
914
915
916
917
918
919
920
921
922
            block_tables, block_tables_tensor, slots = get_cache_manager().allocate(
                batch.needed_blocks_slots,
                batch.blocks,
                batch.max_blocks,
                batch.input_ids.device,
            )
            batch.needed_blocks_slots = None
            batch.block_tables = block_tables
            batch.block_tables_tensor = block_tables_tensor
            batch.slots = slots
923

924
        try:
925
            out = self.forward(batch)
926
927
928
        except Exception as e:
            del batch
            raise e
929

Nicolas Patry's avatar
Nicolas Patry committed
930
931
932
933
934
        if isinstance(out, tuple):
            out, speculative_logits = out
        else:
            speculative_logits = None

935
936
        if prefill:
            next_token_logits = (
937
                out[batch.prefill_next_token_indices] if prefill_logprobs else out
938
            )
Nicolas Patry's avatar
Nicolas Patry committed
939
940
            if speculative_logits is not None:
                speculative_logits = (
OlivierDehaene's avatar
OlivierDehaene committed
941
942
943
                    speculative_logits[batch.prefill_next_token_indices]
                    if prefill_logprobs
                    else speculative_logits
Nicolas Patry's avatar
Nicolas Patry committed
944
                )
945
946
947
        else:
            next_token_logits = out

Nicolas Patry's avatar
Nicolas Patry committed
948
        speculate = get_speculate()
OlivierDehaene's avatar
OlivierDehaene committed
949
950
951
952
953
954
955
956
957
        (
            next_input_ids,
            next_token_logprobs,
            logprobs,
            accepted_ids,
            speculative_ids,
        ) = batch.next_token_chooser(
            batch.all_input_ids_tensor[:, : batch.max_seqlen],
            next_token_logits,
Nicolas Patry's avatar
Nicolas Patry committed
958
            speculate,
OlivierDehaene's avatar
OlivierDehaene committed
959
960
            batch.speculative_ids,
            speculative_logits,
961
962
        )

Nicolas Patry's avatar
Nicolas Patry committed
963
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
964
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
965
966
        )

967
        if prefill:
968
            if len(batch) > 1 and prefill_logprobs:
969
970
                # We create the prefill_tokens_indices tensor that will be used to gather prefill logprobs
                # When batch == 1, we will just use the batch.input_ids values directly
971
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
972
973

            next_position_ids = batch.position_ids.new_empty(len(batch))
974
975
976
            batch.slot_indices = batch.slot_indices[batch.cu_seqlen_prefill[1:] - 1]
            # We do not need cu_seqlen_prefill anymore
            batch.cu_seqlen_prefill = None
977
978
979
980
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

981
982
983
984
985
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
986
        stopped = True
987
988

        # Zipped iterator
OlivierDehaene's avatar
OlivierDehaene committed
989
        iterator = zip(batch.input_lengths, batch.all_input_ids, accepted_ids)
990

991
992
993
994
        # We do two for loops as the first one can run completely asynchronously from the GPU while for the second
        # one, we need to first do a GPU <-> CPU sync
        # It is faster if we delay this sync for the maximum amount of time

995
        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
996
        index = 0
OlivierDehaene's avatar
OlivierDehaene committed
997
        for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator):
998
            # Indexing metadata
999
1000
1001
            start_index = cumulative_length
            end_index = cumulative_length + input_length

1002
            if prefill:
1003
1004
1005
1006
1007
                # Indexing metadata
                out_start_index = batch.prefill_cu_outlens[i]
                out_end_index = batch.prefill_cu_outlens[i + 1]
                out_length = out_end_index - out_start_index

1008
1009
1010
1011
1012
1013
                # Initialize position_ids
                # In decode, we do not need this as we can just increment position ids
                next_position_ids[i] = batch.position_ids[end_index - 1]

                # Used to gather prefill logprobs
                # Copy batch.input_ids to prefill_token_indices
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
                if prefill_logprobs:
                    if len(batch) > 1:
                        prefill_tokens_indices[
                            out_start_index : out_end_index - 1
                        ] = batch.input_ids[start_index + 1 : start_index + out_length]
                    else:
                        # Set prefill_tokens_indices to the correct slice
                        prefill_tokens_indices = batch.input_ids[
                            start_index + 1 : start_index + out_length
                        ]
1024

Nicolas Patry's avatar
Nicolas Patry committed
1025
1026
1027
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1028
1029
1030

            cumulative_length += input_length

Nicolas Patry's avatar
Nicolas Patry committed
1031
1032
1033
1034
1035
        batch.input_ids = next_input_ids[accepted_ids.cumsum(dim=-1) - 1]
        batch.speculative_ids = speculative_ids
        batch.position_ids = next_position_ids + accepted_ids
        batch.input_lengths_tensor += accepted_ids
        batch.slot_indices += accepted_ids
1036

1037
        if prefill and prefill_logprobs:
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
            # Get prefill logprobs
            prefill_logprobs_tensor = torch.log_softmax(out, -1)
            prefill_logprobs = torch.gather(
                prefill_logprobs_tensor, 1, prefill_tokens_indices.view(-1, 1)
            )
            # GPU <-> CPU sync
            prefill_logprobs = prefill_logprobs.view(-1).tolist()

        # GPU <-> CPU sync
        next_token_logprobs = next_token_logprobs.tolist()
Nicolas Patry's avatar
Nicolas Patry committed
1048
        next_token_ids = next_input_ids.tolist()
1049
1050
        accepted_ids = accepted_ids.tolist()
        start_decode = time.time_ns()
1051
1052
1053
1054
1055

        # Zipped iterator
        iterator = zip(
            batch.requests,
            batch.input_lengths,
1056
1057
            batch.prefix_offsets,
            batch.read_offsets,
1058
1059
            batch.stopping_criterias,
            batch.all_input_ids,
1060
1061
            batch.next_token_chooser.do_sample,
            batch.next_token_chooser.seeds,
Nicolas Patry's avatar
Nicolas Patry committed
1062
            batch.top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1063
            accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1064
1065
            batch_top_token_ids,
            batch_top_token_logprobs,
1066
1067
1068
        )

        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1069
        index = 0
1070
1071
1072
        for i, (
            request,
            input_length,
1073
1074
            prefix_offset,
            read_offset,
1075
1076
            stopping_criteria,
            all_input_ids,
1077
1078
            do_sample,
            seed,
Nicolas Patry's avatar
Nicolas Patry committed
1079
            top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1080
            n_accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1081
1082
            top_token_ids,
            top_token_logprobs,
1083
        ) in enumerate(iterator):
1084
            # Append next token to all tokens
Nicolas Patry's avatar
Nicolas Patry committed
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
            next_token_texts = []
            left = 0

            current_stopped = False
            for j in range(index, index + n_accepted_ids):
                # Generated token
                next_token_id = next_token_ids[j]
                all_input_ids.append(next_token_id)
                next_token_text, prefix_offset, read_offset = self.decode_token(
                    all_input_ids,
                    prefix_offset,
                    read_offset,
                )
                next_token_texts.append(next_token_text)
1099

Nicolas Patry's avatar
Nicolas Patry committed
1100
1101
1102
1103
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1104

Nicolas Patry's avatar
Nicolas Patry committed
1105
1106
1107
1108
1109
1110
1111
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1112

OlivierDehaene's avatar
OlivierDehaene committed
1113
1114
1115
1116
            _next_token_ids = next_token_ids[index : index + n_accepted_ids - left]
            _next_token_logprobs = next_token_logprobs[
                index : index + n_accepted_ids - left
            ]
Nicolas Patry's avatar
Nicolas Patry committed
1117
            index += n_accepted_ids
1118

1119
1120
1121
1122
1123
            # Shard generations
            # All generations will be appended in the rust sharded client
            if i % self.world_size == self.rank:
                if stop:
                    # Decode generated tokens
1124
1125
                    output_text, _, _ = self.decode_token(
                        all_input_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1126
1127
1128
1129
1130
1131
                        prefix_offset=len(all_input_ids)
                        - stopping_criteria.current_tokens
                        - 1,
                        read_offset=len(all_input_ids)
                        - stopping_criteria.current_tokens,
                        skip_special_tokens=True,
1132
1133
                    )
                    generated_text = GeneratedText(
1134
1135
1136
1137
                        output_text,
                        stopping_criteria.current_tokens,
                        reason,
                        seed if do_sample else None,
1138
1139
1140
1141
1142
                    )
                else:
                    generated_text = None

                # Prefill
1143
1144
1145
1146
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

1147
1148
                    # Remove generated token to only have prefill and add nan for first prompt token
                    request_prefill_logprobs = [float("nan")] + prefill_logprobs[
1149
                        out_start_index : out_end_index - 1
1150
1151
1152
1153
1154
1155
1156
                    ]
                    prefill_token_ids = all_input_ids[:-1]
                    prefill_texts = self.tokenizer.batch_decode(
                        prefill_token_ids,
                        clean_up_tokenization_spaces=False,
                        skip_special_tokens=False,
                    )
Nicolas Patry's avatar
Nicolas Patry committed
1157
1158

                    prefill_tokens = Tokens(
OlivierDehaene's avatar
OlivierDehaene committed
1159
1160
1161
1162
                        prefill_token_ids,
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1163
1164
1165
1166
                    )
                else:
                    prefill_tokens = None

Nicolas Patry's avatar
Nicolas Patry committed
1167
                if top_n_tokens > 0:
Nicolas Patry's avatar
Nicolas Patry committed
1168
                    all_top_tokens = []
1169
1170
1171
                    for (top_token_ids, top_token_logprobs) in zip(
                        top_token_ids, top_token_logprobs
                    ):
Nicolas Patry's avatar
Nicolas Patry committed
1172
1173
1174
1175
1176
1177
                        toptoken_texts = self.tokenizer.batch_decode(
                            top_token_ids,
                            clean_up_tokenization_spaces=False,
                            skip_special_tokens=False,
                        )
                        special_toptokens = [
1178
1179
                            token_id in self.all_special_ids
                            for token_id in top_token_ids
Nicolas Patry's avatar
Nicolas Patry committed
1180
1181
1182
1183
1184
1185
1186
1187
1188
                        ]
                        top_tokens = Tokens(
                            top_token_ids,
                            top_token_logprobs,
                            toptoken_texts,
                            special_toptokens,
                        )
                        all_top_tokens.append(top_tokens)
                    top_tokens = all_top_tokens
Nicolas Patry's avatar
Nicolas Patry committed
1189
1190
1191
                else:
                    top_tokens = None

1192
1193
1194
                generation = Generation(
                    request.id,
                    prefill_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1195
1196
1197
1198
1199
1200
                    Tokens(
                        _next_token_ids,
                        _next_token_logprobs,
                        next_token_texts,
                        [nid in self.all_special_ids for nid in _next_token_ids],
                    ),
1201
                    generated_text,
Nicolas Patry's avatar
Nicolas Patry committed
1202
                    top_tokens,
1203
1204
                )

1205
                generations.append(generation)
1206

1207
            # Update values
1208
            batch.input_lengths[i] = input_length + n_accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1209
1210
            if batch.input_lengths[i] > batch.max_seqlen:
                batch.max_seqlen = batch.input_lengths[i]
1211
1212
            batch.prefix_offsets[i] = prefix_offset
            batch.read_offsets[i] = read_offset
1213
1214
            batch.all_input_ids[i] = all_input_ids

1215
1216
1217
        if stopped:
            del batch
            # No need to return a batch if we know that all requests stopped
1218
1219
1220
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1221

1222
1223
1224
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1225

1226
1227
1228
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)