flash_causal_lm.py 45.1 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, CUDA_GRAPHS
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
@dataclass
class FlashCausalLMBatch(Batch):
    batch_id: int
    requests: List[generate_pb2.Request]
42
43
    # request id -> idx in list mapping
    requests_idx_mapping: Dict[int, int]
44
45

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

50
51
52
53
    # 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]
54
55
56
57
58
59
60
61
62
63
64
65
66
67

    # 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]]]
68
    # tensor of size [b, max_total_seqlen // block_size] holding the paged attention block tables for all sequences
69
70
71
72
    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]

73
74
    max_seqlen: int

75
76
77
78
79
    # 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]]

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

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

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

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

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

    @classmethod
    def from_pb(
        cls,
        pb: generate_pb2.Batch,
        tokenizer: PreTrainedTokenizerBase,
114
        dtype: torch.dtype,
115
        device: torch.device,
116
    ) -> "FlashCausalLMBatch":
117
118
119
120
121
122
123
124
125
126
        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"]

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

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

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

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

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

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

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

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

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

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

175
            all_input_ids.append(tokenized_input)
176
177

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

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

184
            next_token_chooser_parameters.append(r.parameters)
185

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

193
194
            # Paged attention
            # Remove one as the first token des not have a past
Nicolas Patry's avatar
Nicolas Patry committed
195
196
            speculative_length = get_speculate()
            total_tokens = input_length + max_new_tokens - 1 + speculative_length
197
198
199
200
201
202
203
204
205
206
207
208
            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)

209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
            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

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

        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
drbh's avatar
drbh committed
239
            next_token_chooser_parameters, dtype, device, tokenizer
240
        )
241
        start_slots = torch.tensor(start_slots, dtype=torch.int64)
242
243
244
245
246
247
248

        # 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
249

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

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

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

274
275
        if all_prefill_logprobs:
            prefill_head_indices = None
276
            prefill_next_token_indices = cu_seqlen_prefill[1:] - 1
277
        elif no_prefill_logprobs:
278
            prefill_head_indices = cu_seqlen_prefill[1:] - 1
279
280
281
282
283
284
285
286
            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
287
288
289
        top_n_tokens_tensor = torch.tensor(
            top_n_tokens, device=device, dtype=torch.int64
        )
290

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

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

331
        device = self.input_ids.device
332

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

            cumulative_max_length += request_input_length + remaining_tokens - 1
408

409
410
411
412
413
414
415
416
417
            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
418
        get_cache_manager().free(block_indices_to_free)
419
420
421
        # Needed to avoid dropping blocks when the batches will go out of scope
        self.block_tables = None

422
423
424
425
        # 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]
426
427
428
        block_tables_tensor = self.block_tables_tensor[indices]
        input_lengths_tensor = self.input_lengths_tensor[indices]
        slots = self.slots[slot_filtering_indices]
429
        next_token_chooser = self.next_token_chooser.filter(indices)
Nicolas Patry's avatar
Nicolas Patry committed
430
        top_n_tokens_tensor = self.top_n_tokens_tensor[indices]
OlivierDehaene's avatar
OlivierDehaene committed
431
432
433
        speculative_ids = (
            self.speculative_ids[indices] if self.speculative_ids is not None else None
        )
434
435

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

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

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

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

479
480
481
482
483
484
485
486
487
488
        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
489
490
491
            speculative_length = (
                b.speculative_ids.shape[1] if b.speculative_ids is not None else 0
            )
492
493
494
495
496
497
498
            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
499
                    + speculative_length
500
501
502
503
504
505
                    - stopping_criteria.current_tokens
                    for input_length, stopping_criteria in zip(
                        b.input_lengths, b.stopping_criterias
                    )
                ),
            )
506
507
508

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
509
510
511
512
513
514
515
516
517
518
        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)
519
        )
Nicolas Patry's avatar
Nicolas Patry committed
520
521
522
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
523

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

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

532
        next_token_chooser_parameters = []
533
        fsm_grammar_states = []
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
            fsm_grammar_states.extend(batch.next_token_chooser.fsm_grammar_states)
583
584
            stopping_criterias.extend(batch.stopping_criterias)

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

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

591
        start_slots = torch.concat(start_slots)
592

593
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
594
595
596
            next_token_chooser_parameters,
            dtype=batches[0].next_token_chooser.dtype,
            device=batches[0].next_token_chooser.device,
drbh's avatar
drbh committed
597
            tokenizer=batches[0].next_token_chooser.tokenizer,
598
            fsm_grammar_states=fsm_grammar_states,
599
600
        )

OlivierDehaene's avatar
OlivierDehaene committed
601
602
603
604
605
        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
606

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

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

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

651
652
653
654
655
656
657
    def __len__(self):
        return len(self.requests)


class FlashCausalLM(Model):
    def __init__(
        self,
658
659
660
661
662
663
664
665
666
        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,
667
        sliding_window: Optional[int] = None,
668
    ):
669
670
671
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_size = head_size
672

673
674
        self.cuda_graphs = {}

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

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

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
    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):
729
            logits, speculative_logits = self.model.forward(
730
731
732
733
734
735
736
737
738
739
                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,
            )
740
741
            self.cuda_graphs[bs]["logits"] = logits
            self.cuda_graphs[bs]["speculative_logits"] = speculative_logits
742
743
        torch.cuda.synchronize()

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

        torch.cuda.synchronize(self.device)

768
769
        # 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
770
771
772
773
        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

774
775
776
777
778
779
        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
        )
780
781

        num_blocks = (
782
783
            # Leave 5% for some wiggle room
            int((free_memory * 0.95) // total_cache_size)
784
            # Add batch.blocks as we allocated it above, so it is included in the peak memory.
785
            + cache_manager.num_blocks
786
787
        )

788
        del batch
789
        del cache_manager
790

791
        set_cache_manager(
792
793
794
795
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
796
            self.sliding_window is not None,
797
798
799
800
            self.dtype,
            self.device,
        )

801
        if CUDA_GRAPHS:
802
            try:
803
                logger.info(f"Cuda Graphs are enabled for sizes {CUDA_GRAPHS}")
804
                # Warmup cuda graphs
805
                for bs in CUDA_GRAPHS:
806
807
808
809
810
                    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")

811
        return int(num_blocks * BLOCK_SIZE)
812

813
814
815
    def forward(
        self, batch: FlashCausalLMBatch
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
816
        # Model Forward
Nicolas Patry's avatar
Nicolas Patry committed
817
        if batch.speculative_ids is not None:
OlivierDehaene's avatar
OlivierDehaene committed
818
819
820
821
822
823
824
825
826
            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
827
828
829

            speculative_ids = batch.speculative_ids

OlivierDehaene's avatar
OlivierDehaene committed
830
            B, speculative_length = speculative_ids.shape
Nicolas Patry's avatar
Nicolas Patry committed
831
            new_length = speculative_length + 1
OlivierDehaene's avatar
OlivierDehaene committed
832
833
834
            new_input_ids = torch.cat(
                [input_ids.unsqueeze(-1), speculative_ids], dim=1
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
835
836
            arange = torch.arange(new_length, device=position_ids.device).unsqueeze(0)
            arange_int = arange.to(dtype=torch.int32)
OlivierDehaene's avatar
OlivierDehaene committed
837
838
839
            new_position_ids = (
                position_ids.unsqueeze(-1).expand(B, new_length) + arange
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
840
            slots = (slots.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1)
OlivierDehaene's avatar
OlivierDehaene committed
841
842
843
            input_lengths = (
                input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
844
845

            # Add Copy the block tables for all members
OlivierDehaene's avatar
OlivierDehaene committed
846
847
848
849
850
851
            block_tables = (
                block_tables.unsqueeze(1)
                .expand(B, new_length, -1)
                .reshape(B * new_length, -1)
                .contiguous()
            )
Nicolas Patry's avatar
Nicolas Patry committed
852
853
854
855
856
            max_s = max_s + speculative_length

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
857
858
859
860
861
862
863
864
865
            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
866

867
868
869
870
871
872
873
874
875
876
877
878
        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)

drbh's avatar
drbh committed
879
880
881
882
883
        if (
            cu_seqlen_prefill is not None
            or cuda_graph is None
            or batch.speculative_ids is not None
        ):
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
            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
911
912
913
914
915
916
917
        speculative_logits = (
            cuda_graph["speculative_logits"][:bs]
            if cuda_graph["speculative_logits"] is not None
            else None
        )
        logits = cuda_graph["logits"][:bs]
        return logits, speculative_logits
918
919
920
921

    @tracer.start_as_current_span("generate_token")
    def generate_token(
        self, batch: FlashCausalLMBatch
922
923
    ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]:
        start = time.time_ns()
924
        prefill = batch.cu_seqlen_prefill is not None
925
        prefill_logprobs = batch.prefill_next_token_indices is not None
926

927
928
        if batch.needed_blocks_slots:
            # Allocate blocks to this batch
929
930
931
932
933
934
935
936
937
938
            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
939

940
        try:
941
            out, speculative_logits = self.forward(batch)
942
943
944
        except Exception as e:
            del batch
            raise e
945

946
947
        if prefill:
            next_token_logits = (
948
                out[batch.prefill_next_token_indices] if prefill_logprobs else out
949
            )
Nicolas Patry's avatar
Nicolas Patry committed
950
951
            if speculative_logits is not None:
                speculative_logits = (
OlivierDehaene's avatar
OlivierDehaene committed
952
953
954
                    speculative_logits[batch.prefill_next_token_indices]
                    if prefill_logprobs
                    else speculative_logits
Nicolas Patry's avatar
Nicolas Patry committed
955
                )
956
957
958
        else:
            next_token_logits = out

Nicolas Patry's avatar
Nicolas Patry committed
959
        speculate = get_speculate()
OlivierDehaene's avatar
OlivierDehaene committed
960
961
962
963
964
965
966
967
968
        (
            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
969
            speculate,
OlivierDehaene's avatar
OlivierDehaene committed
970
971
            batch.speculative_ids,
            speculative_logits,
972
973
        )

Nicolas Patry's avatar
Nicolas Patry committed
974
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
975
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
976
977
        )

978
        if prefill:
979
            if len(batch) > 1 and prefill_logprobs:
980
981
                # 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
982
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
983
984

            next_position_ids = batch.position_ids.new_empty(len(batch))
985
986
987
            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
988
989
990
991
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

992
993
994
995
996
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
997
        stopped = True
998
999

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

1002
1003
1004
1005
        # 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

1006
        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1007
        index = 0
OlivierDehaene's avatar
OlivierDehaene committed
1008
        for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator):
1009
            # Indexing metadata
1010
1011
1012
            start_index = cumulative_length
            end_index = cumulative_length + input_length

1013
            if prefill:
1014
1015
1016
1017
1018
                # 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

1019
1020
1021
1022
1023
1024
                # 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
1025
1026
                if prefill_logprobs:
                    if len(batch) > 1:
drbh's avatar
drbh committed
1027
1028
1029
                        prefill_tokens_indices[out_start_index : out_end_index - 1] = (
                            batch.input_ids[start_index + 1 : start_index + out_length]
                        )
1030
1031
1032
1033
1034
                    else:
                        # Set prefill_tokens_indices to the correct slice
                        prefill_tokens_indices = batch.input_ids[
                            start_index + 1 : start_index + out_length
                        ]
1035

Nicolas Patry's avatar
Nicolas Patry committed
1036
1037
1038
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1039
1040
1041

            cumulative_length += input_length

drbh's avatar
drbh committed
1042
        # Update values
Nicolas Patry's avatar
Nicolas Patry committed
1043
1044
1045
1046
1047
        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
1048

1049
        if prefill and prefill_logprobs:
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
            # 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
1060
        next_token_ids = next_input_ids.tolist()
1061
1062
        accepted_ids = accepted_ids.tolist()
        start_decode = time.time_ns()
1063
1064
1065
1066
1067

        # Zipped iterator
        iterator = zip(
            batch.requests,
            batch.input_lengths,
1068
1069
            batch.prefix_offsets,
            batch.read_offsets,
1070
1071
            batch.stopping_criterias,
            batch.all_input_ids,
1072
1073
            batch.next_token_chooser.do_sample,
            batch.next_token_chooser.seeds,
Nicolas Patry's avatar
Nicolas Patry committed
1074
            batch.top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1075
            accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1076
1077
            batch_top_token_ids,
            batch_top_token_logprobs,
1078
1079
1080
        )

        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1081
        index = 0
1082
1083
1084
        for i, (
            request,
            input_length,
1085
1086
            prefix_offset,
            read_offset,
1087
1088
            stopping_criteria,
            all_input_ids,
1089
1090
            do_sample,
            seed,
Nicolas Patry's avatar
Nicolas Patry committed
1091
            top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1092
            n_accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1093
1094
            top_token_ids,
            top_token_logprobs,
1095
        ) in enumerate(iterator):
1096
            # Append next token to all tokens
Nicolas Patry's avatar
Nicolas Patry committed
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
            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)
1111

Nicolas Patry's avatar
Nicolas Patry committed
1112
1113
1114
1115
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1116

Nicolas Patry's avatar
Nicolas Patry committed
1117
1118
1119
1120
1121
1122
1123
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1124

OlivierDehaene's avatar
OlivierDehaene committed
1125
1126
1127
1128
            _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
1129
            index += n_accepted_ids
1130

1131
1132
1133
1134
1135
            # Shard generations
            # All generations will be appended in the rust sharded client
            if i % self.world_size == self.rank:
                if stop:
                    # Decode generated tokens
1136
1137
                    output_text, _, _ = self.decode_token(
                        all_input_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1138
1139
1140
1141
1142
1143
                        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,
1144
1145
                    )
                    generated_text = GeneratedText(
1146
1147
1148
1149
                        output_text,
                        stopping_criteria.current_tokens,
                        reason,
                        seed if do_sample else None,
1150
1151
1152
1153
1154
                    )
                else:
                    generated_text = None

                # Prefill
1155
1156
1157
1158
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

1159
1160
                    # Remove generated token to only have prefill and add nan for first prompt token
                    request_prefill_logprobs = [float("nan")] + prefill_logprobs[
1161
                        out_start_index : out_end_index - 1
1162
1163
1164
1165
1166
1167
1168
                    ]
                    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
1169
1170

                    prefill_tokens = Tokens(
OlivierDehaene's avatar
OlivierDehaene committed
1171
1172
1173
1174
                        prefill_token_ids,
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1175
1176
1177
1178
                    )
                else:
                    prefill_tokens = None

Nicolas Patry's avatar
Nicolas Patry committed
1179
                if top_n_tokens > 0:
Nicolas Patry's avatar
Nicolas Patry committed
1180
                    all_top_tokens = []
drbh's avatar
drbh committed
1181
                    for top_token_ids, top_token_logprobs in zip(
1182
1183
                        top_token_ids, top_token_logprobs
                    ):
Nicolas Patry's avatar
Nicolas Patry committed
1184
1185
1186
1187
1188
1189
                        toptoken_texts = self.tokenizer.batch_decode(
                            top_token_ids,
                            clean_up_tokenization_spaces=False,
                            skip_special_tokens=False,
                        )
                        special_toptokens = [
1190
1191
                            token_id in self.all_special_ids
                            for token_id in top_token_ids
Nicolas Patry's avatar
Nicolas Patry committed
1192
1193
1194
1195
1196
1197
1198
1199
1200
                        ]
                        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
1201
1202
1203
                else:
                    top_tokens = None

1204
1205
1206
                generation = Generation(
                    request.id,
                    prefill_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1207
1208
1209
1210
1211
1212
                    Tokens(
                        _next_token_ids,
                        _next_token_logprobs,
                        next_token_texts,
                        [nid in self.all_special_ids for nid in _next_token_ids],
                    ),
1213
                    generated_text,
Nicolas Patry's avatar
Nicolas Patry committed
1214
                    top_tokens,
1215
1216
                )

1217
                generations.append(generation)
1218

drbh's avatar
drbh committed
1219
1220
1221
            # accept each new token for this specific request since we may
            # have more than one new token per request with speculative decoding
            for next_token_id in _next_token_ids:
OlivierDehaene's avatar
OlivierDehaene committed
1222
1223
1224
                batch.next_token_chooser = (
                    batch.next_token_chooser.advance_grammar_single(i, next_token_id)
                )
drbh's avatar
drbh committed
1225

1226
            # Update values
1227
            batch.input_lengths[i] = input_length + n_accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1228
1229
            if batch.input_lengths[i] > batch.max_seqlen:
                batch.max_seqlen = batch.input_lengths[i]
1230
1231
            batch.prefix_offsets[i] = prefix_offset
            batch.read_offsets[i] = read_offset
1232
1233
            batch.all_input_ids[i] = all_input_ids

1234
1235
1236
        if stopped:
            del batch
            # No need to return a batch if we know that all requests stopped
1237
1238
1239
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1240

1241
1242
1243
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1244

1245
1246
1247
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)