flash_causal_lm.py 51.8 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
fxmarty's avatar
fxmarty committed
15
16
17

from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
from text_generation_server.utils.import_utils import SYSTEM
OlivierDehaene's avatar
OlivierDehaene committed
18
from text_generation_server.models import Model
19
from text_generation_server.utils.tokens import batch_top_tokens
20
from text_generation_server.utils.dist import RANK
Nicolas Patry's avatar
Nicolas Patry committed
21
from text_generation_server.utils.speculate import get_speculate
22
23
from text_generation_server.models.types import (
    Batch,
Nicolas Patry's avatar
Nicolas Patry committed
24
    Tokens,
25
26
27
28
    Generation,
    GeneratedText,
)
from text_generation_server.pb import generate_pb2
29
from text_generation_server.models.globals import MEM_POOL, CUDA_GRAPHS
fxmarty's avatar
fxmarty committed
30
import text_generation_server.models.globals as tgi_globals
31
from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
32
from text_generation_server.utils.dist import MEMORY_FRACTION
33

Nicolas Patry's avatar
Nicolas Patry committed
34
from text_generation_server.utils.import_utils import (
Nicolas Patry's avatar
Nicolas Patry committed
35
36
37
    empty_cache,
    synchronize,
    get_free_memory,
Nicolas Patry's avatar
Nicolas Patry committed
38
39
)

Nicolas Patry's avatar
Nicolas Patry committed
40
41
tracer = trace.get_tracer(__name__)

42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
BLOCK_SIZE: int = 16

# Will be set in init
SLIDING_WINDOW: Optional[int] = None


def set_sliding_window(sliding_window: int):
    global SLIDING_WINDOW
    SLIDING_WINDOW = sliding_window


def get_sliding_windows() -> int:
    global SLIDING_WINDOW
    return SLIDING_WINDOW

57

58
59
60
61
@dataclass
class FlashCausalLMBatch(Batch):
    batch_id: int
    requests: List[generate_pb2.Request]
62
63
    # request id -> idx in list mapping
    requests_idx_mapping: Dict[int, int]
64
65

    # Decoder values
66
67
    input_ids: torch.Tensor
    position_ids: torch.Tensor
68
    speculative_ids: Optional[torch.Tensor]
69

70
71
72
73
    # 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]
74
75
76
    # Prefill cache indices is used to slice into the kv tensor before caching it into the paged attention buffers
    # as we only keep SLIDING_WINDOW values instead of the whole tensor
    prefill_cache_indices: Optional[torch.Tensor]
77
78
79
80
81
82
83
84
85
86

    # 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 length b of list of length s_i // block_size
87
    block_tables: List[List[int]]
88
    # tensor of size [b, max_total_seqlen // block_size] holding the paged attention block tables for all sequences
89
    block_tables_tensor: torch.Tensor
90
    # tensor of length \sum_{i=0}^{b} max_s_i  holding the paged attention slots for all sequences
91
    slots: torch.Tensor
92

93
94
    max_seqlen: int

95
96
97
98
99
    # 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]]

100
101
    # All tokens
    all_input_ids: List[List[int]]
102
    all_input_ids_tensor: torch.Tensor
103
104
105

    # Lengths of all generations present in the batch
    input_lengths: List[int]
106
    input_lengths_tensor: torch.Tensor
107
108
    prefix_offsets: List[Optional[int]]
    read_offsets: List[Optional[int]]
109
110

    # Generation helpers
111
    next_token_chooser: HeterogeneousNextTokenChooser
112
    stopping_criterias: List[StoppingCriteria]
Nicolas Patry's avatar
Nicolas Patry committed
113
114
    top_n_tokens: List[int]
    top_n_tokens_tensor: torch.Tensor
115

116
    # Number of blocks in this batch
117
    num_blocks: int
118
119
    # Maximum number of blocks
    max_blocks: int
120

121
122
    def to_pb(self) -> generate_pb2.CachedBatch:
        return generate_pb2.CachedBatch(
123
            id=self.batch_id,
124
            request_ids=[r.id for r in self.requests],
125
            size=len(self),
126
            max_tokens=self.num_blocks * BLOCK_SIZE,
127
128
129
        )

    @classmethod
130
    def batch_tokenized_inputs(cls, requests, tokenizer):
131
132
        batch_inputs = []
        max_truncation = 0
133
        for r in requests:
134
135
136
137
138
139
            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"]
140
        return batch_tokenized_inputs
141

drbh's avatar
drbh committed
142
143
144
145
146
147
148
149
150
    @classmethod
    def from_tokenized(
        cls,
        pb: generate_pb2.Batch,
        tokenizer: PreTrainedTokenizerBase,
        batch_tokenized_inputs,
        dtype: torch.dtype,
        device: torch.device,
    ) -> "FlashCausalLMBatch":
151
        sliding_window = get_sliding_windows()
152
        position_ids = []
153
        cu_seqlen_prefill = [0]
154
155
        start_slots = []
        slot_indices = []
156
        prefill_cache_indices = []
157
158

        input_lengths = []
159
160
        prefix_offsets = []
        read_offsets = []
161
        all_input_ids = []
162
        requests_idx_mapping = {}
163

164
165
166
167
168
169
        all_prefill_logprobs = True
        no_prefill_logprobs = True
        prefill_head_indices = []
        prefill_next_token_indices = []
        prefill_cu_outlens = [0]

170
        next_token_chooser_parameters = []
171
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
172
        top_n_tokens = []
173
174
175

        # Cumulative length
        cumulative_length = 0
176
        cumulative_max_length = 0
177
        prefill_out_cumulative_length = 0
178

179
        num_blocks = 0
180
        max_seqlen = 0
181
        max_length = 0
182
        max_blocks = 0
183

184
185
186
        block_tables = []
        slots = []

187
        # Parse batch
188
189
190
        for i, (r, tokenized_input) in enumerate(
            zip(pb.requests, batch_tokenized_inputs)
        ):
191
192
193
            # request id -> idx in list mapping
            requests_idx_mapping[r.id] = i

194
            tokenized_input = tokenized_input[-r.truncate :]
195
196
197
198
199
            if (
                tokenized_input[0] == tokenizer.bos_token_id
                and tokenized_input[1] == tokenizer.bos_token_id
            ):
                tokenized_input = tokenized_input[1:]
200

201
202
            input_length = len(tokenized_input)
            input_lengths.append(input_length)
203

204
            prefix_offsets.append(input_length - 5)
205
            read_offsets.append(input_length)
206

207
            all_input_ids.append(tokenized_input)
208
209

            # Position ids
210
211
            request_position_ids = torch.arange(0, input_length, dtype=torch.int32)
            position_ids.append(request_position_ids)
212
213

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

216
            next_token_chooser_parameters.append(r.parameters)
217

218
219
220
            stopping_criteria = StoppingCriteria.from_pb(
                r.stopping_parameters, tokenizer
            )
221
            max_new_tokens = stopping_criteria.max_new_tokens
222
            stopping_criterias.append(stopping_criteria)
Nicolas Patry's avatar
Nicolas Patry committed
223
            top_n_tokens.append(r.top_n_tokens)
224

225
226
            # Paged attention
            # Remove one as the first token des not have a past
Nicolas Patry's avatar
Nicolas Patry committed
227
            speculative_length = get_speculate()
drbh's avatar
drbh committed
228
            speculative_length = 0 if speculative_length is None else speculative_length
Nicolas Patry's avatar
Nicolas Patry committed
229
            total_tokens = input_length + max_new_tokens - 1 + speculative_length
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248

            # blocks and slots can be empty (for example in warmup)
            if not r.blocks:
                needed_blocks = math.ceil(total_tokens / BLOCK_SIZE)
                request_blocks = [
                    b for b in range(num_blocks, num_blocks + needed_blocks)
                ]
                request_slots = [
                    s
                    for b in request_blocks
                    for s in range(b * BLOCK_SIZE, (b + 1) * BLOCK_SIZE)
                ]
            else:
                request_blocks = r.blocks
                request_slots = r.slots

            block_tables.append(request_blocks)
            slots.extend(request_slots[:total_tokens])
            num_blocks += len(request_blocks)
249
250
251
252
253
254
255
256
257
            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)

258
259
260
261
262
263
264
265
266
            # Create tensor to slice into the kv tensor in prefill
            if sliding_window is not None:
                request_prefill_cache_indices = torch.arange(
                    cumulative_length + max(0, input_length - sliding_window),
                    cumulative_length + input_length,
                    dtype=torch.int64,
                )
                prefill_cache_indices.append(request_prefill_cache_indices)

267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
            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

287
288
            # Update
            cumulative_length += input_length
289
290
            cumulative_max_length += total_tokens
            max_seqlen = max(max_seqlen, input_length)
291
            max_blocks = max(max_blocks, len(request_blocks))
OlivierDehaene's avatar
OlivierDehaene committed
292
293
294
            max_length = max(
                max_length, input_length + max_new_tokens + speculative_length
            )
295
296

        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
drbh's avatar
drbh committed
297
            next_token_chooser_parameters, dtype, device, tokenizer
298
        )
299
        start_slots = torch.tensor(start_slots, dtype=torch.int64)
300
301
302
303
304
305
306

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

308
309
310
311
312
        # Create tensors on device
        all_input_ids_tensor = torch.tensor(
            all_input_ids_tensor, dtype=torch.int64, device=device
        )

313
314
315
        if len(pb.requests) > 1:
            input_ids = np.concatenate(all_input_ids, dtype=np.int64)
            position_ids = torch.cat(position_ids)
316
            slot_indices = torch.cat(slot_indices)
317
318
            if sliding_window is not None:
                prefill_cache_indices = torch.cat(prefill_cache_indices)
319
320
321
        else:
            input_ids = all_input_ids[0]
            position_ids = position_ids[0]
322
            slot_indices = slot_indices[0]
323
324
            if sliding_window is not None:
                prefill_cache_indices = prefill_cache_indices[0]
325

326
327
        cu_seqlen_prefill = torch.tensor(
            cu_seqlen_prefill, device=device, dtype=torch.int32
328
329
330
        )
        position_ids = position_ids.to(device)
        slot_indices = slot_indices.to(device)
331
332
333
        prefill_cache_indices = (
            prefill_cache_indices.to(device) if sliding_window is not None else None
        )
334
        input_ids = torch.tensor(input_ids, dtype=torch.int64, device=device)
335
336
        input_lengths_tensor = torch.tensor(
            input_lengths, dtype=torch.int32, device=device
337
        )
338

339
340
        if all_prefill_logprobs:
            prefill_head_indices = None
341
            prefill_next_token_indices = cu_seqlen_prefill[1:] - 1
342
        elif no_prefill_logprobs:
343
            prefill_head_indices = cu_seqlen_prefill[1:] - 1
344
345
346
347
348
349
350
351
            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
352
353
354
        top_n_tokens_tensor = torch.tensor(
            top_n_tokens, device=device, dtype=torch.int64
        )
355

356
357
358
359
360
361
362
363
        slots = torch.tensor(slots, dtype=torch.int64, device=device)
        block_tables_tensor = torch.zeros(
            (len(block_tables), max_blocks), dtype=torch.int32, device="cpu"
        )
        for i, request_blocks in enumerate(block_tables):
            block_tables_tensor[i, : len(request_blocks)] = torch.tensor(request_blocks)
        block_tables_tensor = block_tables_tensor.to(device)

364
365
366
        return cls(
            batch_id=pb.id,
            requests=pb.requests,
367
            requests_idx_mapping=requests_idx_mapping,
368
369
            input_ids=input_ids,
            position_ids=position_ids,
370
            cu_seqlen_prefill=cu_seqlen_prefill,
371
            prefill_cache_indices=prefill_cache_indices,
372
373
            start_slots=start_slots,
            slot_indices=slot_indices,
374
375
376
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
377
            max_seqlen=max_seqlen,
378
379
380
            prefill_head_indices=prefill_head_indices,
            prefill_next_token_indices=prefill_next_token_indices,
            prefill_cu_outlens=prefill_cu_outlens,
381
            input_lengths=input_lengths,
382
            input_lengths_tensor=input_lengths_tensor,
383
384
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
385
            all_input_ids=all_input_ids,
386
387
            all_input_ids_tensor=all_input_ids_tensor,
            next_token_chooser=next_token_chooser,
388
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
389
390
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
391
            num_blocks=num_blocks,
392
            max_blocks=max_blocks,
Nicolas Patry's avatar
Nicolas Patry committed
393
            speculative_ids=None,
394
395
        )

396
397
398
399
400
401
402
403
404
405
406
    @classmethod
    def from_pb(
        cls,
        pb: generate_pb2.Batch,
        tokenizer: PreTrainedTokenizerBase,
        dtype: torch.dtype,
        device: torch.device,
    ) -> "FlashCausalLMBatch":
        batch_tokenized_inputs = cls.batch_tokenized_inputs(pb.requests, tokenizer)
        return cls.from_tokenized(pb, tokenizer, batch_tokenized_inputs, dtype, device)

407
    @tracer.start_as_current_span("filter")
408
409
    def filter(self, request_ids: List[int]) -> "FlashCausalLMBatch":
        if len(request_ids) == 0:
410
411
            raise ValueError("Batch must have at least one request")
        # We assume that if len(requests) == len(self) then the requests are the same
412
        if len(request_ids) == len(self):
413
414
            return self

415
        device = self.input_ids.device
416

417
418
419
        # New values after filtering
        requests_idx_mapping = {}

420
421
422
        # Used to index into tensors
        indices = []

423
424
425
        # slots to keep after filtering
        slot_filtering_indices = torch.zeros(
            self.slots.shape[0], dtype=torch.bool, device=device
426
427
        )

428
        # Create on CPU to only move to GPU once instead of at every copy
429
        slot_indices = torch.empty(len(request_ids), dtype=torch.int64)
430
431
        max_seqlen = 0

432
        requests = []
433
434
        start_slots = []
        block_tables = []
435
436
        all_input_ids = []

437
        input_lengths = []
438
439
        prefix_offsets = []
        read_offsets = []
440

441
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
442
        top_n_tokens = []
443

444
        num_blocks = 0
445
446
447
448
        max_blocks = 0
        # Cumulative length
        cumulative_max_length = 0

449
450
        for i, request_id in enumerate(request_ids):
            idx = self.requests_idx_mapping[request_id]
451
            indices.append(idx)
452
453
454
            requests_idx_mapping[request_id] = i

            requests.append(self.requests[idx])
455
456
457
458

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

460
461
462
            all_input_ids.append(self.all_input_ids[idx])

            input_lengths.append(request_input_length)
463
464
            prefix_offsets.append(self.prefix_offsets[idx])
            read_offsets.append(self.read_offsets[idx])
465

466
467
            stopping_criteria = self.stopping_criterias[idx]
            stopping_criterias.append(stopping_criteria)
468

Nicolas Patry's avatar
Nicolas Patry committed
469
470
            top_n_tokens.append(self.top_n_tokens[idx])

471
            remaining_tokens = (
472
473
                stopping_criteria.max_new_tokens - stopping_criteria.current_tokens
            )
474

475
            request_block_table = self.block_tables[idx]
476
            num_blocks += len(request_block_table)
477
478
479
            block_tables.append(request_block_table)
            start_slots.append(cumulative_max_length)

480
            # Copy to tensor (CPU)
481
            slot_indices[i] = cumulative_max_length + request_input_length - 1
482
483

            # Set slice
484
485
486
487
488
            slot_filtering_indices[
                self.start_slots[idx] : self.start_slots[idx]
                + request_input_length
                + remaining_tokens
                - 1
489
490
491
            ] = True

            cumulative_max_length += request_input_length + remaining_tokens - 1
492

493
494
            max_blocks = max(max_blocks, len(request_block_table))

495
496
497
498
        # 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]
499
500
501
        block_tables_tensor = self.block_tables_tensor[indices]
        input_lengths_tensor = self.input_lengths_tensor[indices]
        slots = self.slots[slot_filtering_indices]
502
        next_token_chooser = self.next_token_chooser.filter(indices)
Nicolas Patry's avatar
Nicolas Patry committed
503
        top_n_tokens_tensor = self.top_n_tokens_tensor[indices]
OlivierDehaene's avatar
OlivierDehaene committed
504
505
506
        speculative_ids = (
            self.speculative_ids[indices] if self.speculative_ids is not None else None
        )
507
508

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

510
        # Move to GPU now that we have the whole tensor
511
        slot_indices = slot_indices.to(device)
512

513
        return type(self)(
514
515
516
517
518
            batch_id=self.batch_id,
            requests=requests,
            requests_idx_mapping=requests_idx_mapping,
            input_ids=input_ids,
            position_ids=position_ids,
519
            cu_seqlen_prefill=None,
520
            prefill_cache_indices=None,
521
522
523
524
525
            start_slots=start_slots,
            slot_indices=slot_indices,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
526
            max_seqlen=max_seqlen,
527
528
529
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
530
            input_lengths=input_lengths,
531
            input_lengths_tensor=input_lengths_tensor,
532
533
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
534
535
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
536
            next_token_chooser=next_token_chooser,
537
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
538
539
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
540
            num_blocks=num_blocks,
541
            max_blocks=max_blocks,
Nicolas Patry's avatar
Nicolas Patry committed
542
            speculative_ids=speculative_ids,
543
544
545
546
547
548
549
550
551
        )

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

552
        num_blocks = 0
553
554
555
556
557
558
559
560
        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)
561
            num_blocks += b.num_blocks
OlivierDehaene's avatar
OlivierDehaene committed
562
563
564
            speculative_length = (
                b.speculative_ids.shape[1] if b.speculative_ids is not None else 0
            )
565
566
567
568
569
570
571
            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
572
                    + speculative_length
573
574
575
576
577
578
                    - stopping_criteria.current_tokens
                    for input_length, stopping_criteria in zip(
                        b.input_lengths, b.stopping_criterias
                    )
                ),
            )
579
580
581

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
582
583
584
585
586
587
588
589
590
591
        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)
592
        )
Nicolas Patry's avatar
Nicolas Patry committed
593
594
595
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
596

597
598
        start_slots = []
        block_tables = []
599
600
601
        all_input_ids = []

        input_lengths = []
602
603
        prefix_offsets = []
        read_offsets = []
604

605
        next_token_chooser_parameters = []
606
        fsm_grammar_states = []
607
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
608
        top_n_tokens = []
609

610
        # Cumulative length
611
        cumulative_batch_size = 0
612
        cumulative_slots = 0
613
614
615

        for i, batch in enumerate(batches):
            requests.extend(batch.requests)
616
617
618
619
620
621
622
623

            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

624
625
            start_index = cumulative_batch_size
            end_index = cumulative_batch_size + len(batch)
626
627
            slots_start_index = cumulative_slots
            slots_end_index = cumulative_slots + len(batch.slots)
628
629
630
631

            # Copy tensors (GPU)
            input_ids[start_index:end_index] = batch.input_ids
            position_ids[start_index:end_index] = batch.position_ids
632
633
            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
634
            top_n_tokens_tensor[start_index:end_index] = batch.top_n_tokens_tensor
635
            slots[slots_start_index:slots_end_index] = batch.slots
636

637
638
639
            all_input_ids_tensor[
                start_index:end_index, : batch.all_input_ids_tensor.shape[1]
            ] = batch.all_input_ids_tensor[:, :max_length]
640

641
642
643
            block_tables_tensor[
                start_index:end_index, : batch.block_tables_tensor.shape[1]
            ] = batch.block_tables_tensor[:, :max_blocks]
644

645
646
647
            start_slots.append(batch.start_slots + cumulative_slots)

            block_tables.extend(batch.block_tables)
648
649
            all_input_ids.extend(batch.all_input_ids)

650
            input_lengths.extend(batch.input_lengths)
651
652
            prefix_offsets.extend(batch.prefix_offsets)
            read_offsets.extend(batch.read_offsets)
653

654
            next_token_chooser_parameters.extend([r.parameters for r in batch.requests])
655
            fsm_grammar_states.extend(batch.next_token_chooser.fsm_grammar_states)
656
657
            stopping_criterias.extend(batch.stopping_criterias)

Nicolas Patry's avatar
Nicolas Patry committed
658
659
            top_n_tokens.extend(batch.top_n_tokens)

660
            # Update
661
            cumulative_batch_size += len(batch)
662
            cumulative_slots += len(batch.slots)
663

664
        start_slots = torch.concat(start_slots)
665

666
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
667
668
669
            next_token_chooser_parameters,
            dtype=batches[0].next_token_chooser.dtype,
            device=batches[0].next_token_chooser.device,
drbh's avatar
drbh committed
670
            tokenizer=batches[0].next_token_chooser.tokenizer,
671
            fsm_grammar_states=fsm_grammar_states,
672
673
        )

OlivierDehaene's avatar
OlivierDehaene committed
674
675
676
677
678
        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
679

680
        return cls(
681
682
            batch_id=batches[0].batch_id,
            requests=requests,
683
            requests_idx_mapping=requests_idx_mapping,
684
685
            input_ids=input_ids,
            position_ids=position_ids,
686
            cu_seqlen_prefill=None,
687
            prefill_cache_indices=None,
688
689
690
691
692
            start_slots=start_slots,
            slot_indices=slot_indices,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
693
            max_seqlen=max_seqlen,
694
695
696
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
697
            input_lengths=input_lengths,
698
            input_lengths_tensor=input_lengths_tensor,
699
700
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
701
702
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
703
            next_token_chooser=next_token_chooser,
704
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
705
706
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
707
            num_blocks=num_blocks,
708
            max_blocks=max_blocks,
OlivierDehaene's avatar
OlivierDehaene committed
709
            speculative_ids=speculative_ids,
710
711
712
713
714
715
716
717
718
        )

    def __len__(self):
        return len(self.requests)


class FlashCausalLM(Model):
    def __init__(
        self,
719
720
721
722
723
724
725
726
727
        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,
728
        sliding_window: Optional[int] = None,
729
    ):
730
731
732
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_size = head_size
733

734
        self.cuda_graphs = {}
735
        self.kv_cache = []
736

737
        super(FlashCausalLM, self).__init__(
738
            model=model,
739
740
741
742
            tokenizer=tokenizer,
            requires_padding=False,
            dtype=dtype,
            device=device,
743
744
            rank=rank,
            world_size=world_size,
745
            sliding_window=sliding_window,
746
747
748
749
750
751
        )

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

752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
    def max_past(self) -> int:
        return getattr(self.model, "max_past", None)

    def init_kv_cache(
        self,
        num_blocks: int,
        num_layers: int,
        num_heads: int,
        head_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ):
        self.kv_cache = []
        empty_cache()

        element_size = torch.tensor([], dtype=dtype).element_size()
        if SYSTEM == "xpu":
            x = 1
        else:
            x = BLOCK_SIZE // element_size

        self.kv_cache = [
            (
                torch.empty(
                    (num_blocks, num_heads, head_size // x, BLOCK_SIZE, x),
                    dtype=dtype,
                    device=device,
                ),
                torch.empty(
                    (num_blocks, num_heads, head_size, BLOCK_SIZE),
                    dtype=dtype,
                    device=device,
                ),
            )
            for _ in range(num_layers)
        ]

789
790
791
    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)
792
        slots = torch.arange(bs, dtype=torch.int64, device=self.device)
793
794
795
796
797
798
799
800
801
802
        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))
        )

        self.cuda_graphs[bs] = {
            "input_ids": input_ids,
            "position_ids": position_ids,
803
            "kv_cache": self.kv_cache,
804
805
806
807
808
809
810
811
812
813
814
815
816
            "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,
817
            kv_cache=self.kv_cache,
818
819
820
821
            block_tables=block_tables,
            slots=slots,
            input_lengths=input_lengths,
            max_s=max_s,
822
            prefill_cache_indices=None,
823
824
825
826
827
            lm_head_indices=None,
        )
        torch.cuda.synchronize()

        with torch.cuda.graph(graph, pool=MEM_POOL):
828
            logits, speculative_logits = self.model.forward(
829
830
831
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlen_prefill=None,
832
                kv_cache=self.kv_cache,
833
834
835
836
                block_tables=block_tables,
                slots=slots,
                input_lengths=input_lengths,
                max_s=max_s,
837
                prefill_cache_indices=None,
838
839
                lm_head_indices=None,
            )
840
841
            self.cuda_graphs[bs]["logits"] = logits
            self.cuda_graphs[bs]["speculative_logits"] = speculative_logits
842
843
        torch.cuda.synchronize()

844
    def warmup(self, batch: FlashCausalLMBatch):
845
        # The warmup batch is the biggest batch we could ever receive
Nicolas Patry's avatar
Nicolas Patry committed
846
847
        empty_cache()

848
        try:
849
850
            self.init_kv_cache(
                batch.num_blocks,
851
852
853
854
855
856
                self.num_layers,
                self.num_kv_heads,
                self.head_size,
                self.dtype,
                self.device,
            )
857
            max_bt = batch.max_blocks
858
            max_s = max_bt * BLOCK_SIZE
fxmarty's avatar
fxmarty committed
859
860
861

            if SYSTEM == "rocm" and os.environ.get("PYTORCH_TUNABLEOP_ENABLED", False):
                torch.cuda.tunable.tuning_enable(False)
862
            _, batch, _ = self.generate_token(batch)
OlivierDehaene's avatar
OlivierDehaene committed
863
        except torch.cuda.OutOfMemoryError as e:
864
            raise RuntimeError(
865
866
                f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. "
                f"You need to decrease `--max-batch-prefill-tokens`"
867
            ) from e
868

Nicolas Patry's avatar
Nicolas Patry committed
869
        synchronize(self.device)
870

871
872
        # 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
873
874
875
876
        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

Nicolas Patry's avatar
Nicolas Patry committed
877
        free_memory = get_free_memory(self.device, MEMORY_FRACTION)
878
879

        num_blocks = (
880
881
            # Leave 5% for some wiggle room
            int((free_memory * 0.95) // total_cache_size)
882
883
            # Add batch.num_blocks as we allocated it above, so it is included in the peak memory.
            + batch.num_blocks
884
885
        )

886
        del batch
887

888
        self.init_kv_cache(
889
890
891
892
893
894
895
896
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
            self.dtype,
            self.device,
        )

fxmarty's avatar
fxmarty committed
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
        if SYSTEM == "rocm":
            if (
                os.environ.get("PYTORCH_TUNABLEOP_ENABLED") is None
                or os.environ.get("PYTORCH_TUNABLEOP_ENABLED") == "1"
            ):
                if os.environ.get("PYTORCH_TUNABLEOP_TUNING") != "0":
                    torch.cuda.tunable.tuning_enable(True)

                if os.environ.get("PYTORCH_TUNABLEOP_SEQLENS") is not None:
                    tuning_sequences = [
                        int(val)
                        for val in os.environ["PYTORCH_TUNABLEOP_SEQLENS"].split(",")
                    ]
                else:
                    tuning_sequences = CUDA_GRAPHS

                tunableop_filepath = os.path.join(
                    HUGGINGFACE_HUB_CACHE,
                    f"tunableop_{tgi_globals.MODEL_ID.replace('/', '-')}_tp{self.world_size}_rank{self.rank}.csv",
                )

                logger.info(
                    f"PyTorch TunableOp (https://github.com/fxmarty/pytorch/tree/2.3-patched/aten/src/ATen/cuda/tunable) is enabled. The warmup may take several minutes, picking the ROCm optimal matrix multiplication kernel for the target lengths {', '.join([str(seqlen) for seqlen in tuning_sequences])}, with typical 5-8% latency improvement for small sequence lengths. The picked GEMMs are saved in the file {tunableop_filepath}. To disable TunableOp, please launch TGI with `PYTORCH_TUNABLEOP_ENABLED=0`."
                )

                if os.path.isfile(tunableop_filepath):
                    logger.info(
                        f"The file {tunableop_filepath} already exists and will be reused."
                    )
                    torch.cuda.tunable.read_file(tunableop_filepath)

                os.makedirs(HUGGINGFACE_HUB_CACHE, exist_ok=True)

                for seqlen in tuning_sequences:
                    logger.info(f"Warming up TunableOp for seqlen={seqlen}")
                    self.tunableop_warmup(seqlen)
                    torch.cuda.tunable.write_file(tunableop_filepath)
                torch.cuda.tunable.tuning_enable(False)
            else:
                logger.info(
                    "PyTorch ROCm TunableOp (https://github.com/pytorch/pytorch/tree/main/aten/src/ATen/cuda/tunable) is disabled. TunableOp brings an additional 5-8% latency improvement for small sequence lengths but requires a warmup. If necessary, please use the environment variable PYTORCH_TUNABLEOP_ENABLED=1 to enable TunableOp."
                )

940
        if CUDA_GRAPHS:
941
            try:
942
                logger.info(f"Cuda Graphs are enabled for sizes {CUDA_GRAPHS}")
943
                # Warmup cuda graphs
944
                for bs in CUDA_GRAPHS:
945
946
                    if self.speculate is None or self.speculate + 1 <= bs:
                        self.cuda_graph_warmup(bs, max_s, max_bt)
OlivierDehaene's avatar
OlivierDehaene committed
947
            except torch.cuda.OutOfMemoryError:
948
                logger.exception(f"Decode cuda graph warmup failed")
949
950
        else:
            logger.info(f"Cuda Graphs are disabled (CUDA_GRAPHS={CUDA_GRAPHS}).")
951

952
        return int(num_blocks * BLOCK_SIZE)
953

fxmarty's avatar
fxmarty committed
954
955
956
957
958
    def tunableop_warmup(self, seqlen: int):
        input_ids = torch.zeros(seqlen, dtype=torch.int64, device=self.device)
        position_ids = torch.zeros(seqlen, dtype=torch.int32, device=self.device)
        slots = torch.arange(seqlen, dtype=torch.int64, device=self.device)

fxmarty's avatar
fxmarty committed
959
960
961
        # Dummy value, some models (starcoder2) don't accept `None`.
        input_lengths = torch.ones(seqlen, dtype=torch.int32, device=self.device)

fxmarty's avatar
fxmarty committed
962
963
964
965
966
967
968
        # We pass a `cu_seqlen_prefill` in order not to have to deal with paged attention cache allocation/deallocation.
        self.model.forward(
            input_ids=input_ids,
            position_ids=position_ids,
            cu_seqlen_prefill=torch.tensor(
                [0, seqlen], device=self.device, dtype=torch.int32
            ),
969
            kv_cache=self.kv_cache,
fxmarty's avatar
fxmarty committed
970
            block_tables=None,
fxmarty's avatar
fxmarty committed
971
            input_lengths=input_lengths,
fxmarty's avatar
fxmarty committed
972
973
974
            slots=slots,
            max_s=seqlen,
            lm_head_indices=None,
975
            prefill_cache_indices=None,
fxmarty's avatar
fxmarty committed
976
977
        )

978
979
980
    def forward(
        self, batch: FlashCausalLMBatch
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
981
        # Model Forward
Nicolas Patry's avatar
Nicolas Patry committed
982
        if batch.speculative_ids is not None:
OlivierDehaene's avatar
OlivierDehaene committed
983
984
985
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
986
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
987
988
989
990
991
            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
992
993
994

            speculative_ids = batch.speculative_ids

OlivierDehaene's avatar
OlivierDehaene committed
995
            B, speculative_length = speculative_ids.shape
Nicolas Patry's avatar
Nicolas Patry committed
996
            new_length = speculative_length + 1
OlivierDehaene's avatar
OlivierDehaene committed
997
998
999
            new_input_ids = torch.cat(
                [input_ids.unsqueeze(-1), speculative_ids], dim=1
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1000
1001
            arange = torch.arange(new_length, device=position_ids.device).unsqueeze(0)
            arange_int = arange.to(dtype=torch.int32)
OlivierDehaene's avatar
OlivierDehaene committed
1002
1003
1004
            new_position_ids = (
                position_ids.unsqueeze(-1).expand(B, new_length) + arange
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1005
            slots = (slots.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1)
OlivierDehaene's avatar
OlivierDehaene committed
1006
1007
1008
            input_lengths = (
                input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1009
1010

            # Add Copy the block tables for all members
OlivierDehaene's avatar
OlivierDehaene committed
1011
1012
1013
1014
1015
1016
            block_tables = (
                block_tables.unsqueeze(1)
                .expand(B, new_length, -1)
                .reshape(B * new_length, -1)
                .contiguous()
            )
Nicolas Patry's avatar
Nicolas Patry committed
1017
1018
1019
1020
1021
            max_s = max_s + speculative_length

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
1022
1023
1024
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1025
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1026
1027
1028
1029
1030
            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
1031

1032
1033
1034
1035
1036
1037
        if cu_seqlen_prefill is None and self.max_past() is not None:
            # In decode, not prefill, we're actually overwriting the KV-cache
            # in a circular buffer mode.
            # This makes sure the max_s for the decode pass is correct.
            max_s = min(self.max_past(), max_s)

1038
        bs = input_ids.shape[0]
OlivierDehaene's avatar
OlivierDehaene committed
1039
1040
1041
1042
1043
1044
1045
1046
        sorted_padded_bs = sorted([k for k in self.cuda_graphs.keys() if k >= bs])
        if sorted_padded_bs:
            # Get associated cuda graph
            cuda_graph = self.cuda_graphs[sorted_padded_bs[0]]
        else:
            cuda_graph = None

        if cu_seqlen_prefill is not None or cuda_graph is None:
1047
            logits, speculative_logits = self.model.forward(
1048
1049
1050
1051
1052
1053
1054
1055
                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,
1056
                prefill_cache_indices=batch.prefill_cache_indices,
1057
1058
                lm_head_indices=lm_head_indices,
            )
1059
1060
1061
            if batch.prefill_cache_indices is not None:
                batch.prefill_cache_indices = None
            return logits, speculative_logits
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077

        # 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
1078
1079
1080
1081
1082
1083
1084
        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
1085
1086
1087
1088

    @tracer.start_as_current_span("generate_token")
    def generate_token(
        self, batch: FlashCausalLMBatch
1089
1090
    ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]:
        start = time.time_ns()
1091
        prefill = batch.cu_seqlen_prefill is not None
1092
        prefill_logprobs = batch.prefill_next_token_indices is not None
1093

1094
        out, speculative_logits = self.forward(batch)
1095

1096
1097
        if prefill:
            next_token_logits = (
1098
                out[batch.prefill_next_token_indices] if prefill_logprobs else out
1099
            )
Nicolas Patry's avatar
Nicolas Patry committed
1100
1101
            if speculative_logits is not None:
                speculative_logits = (
OlivierDehaene's avatar
OlivierDehaene committed
1102
1103
1104
                    speculative_logits[batch.prefill_next_token_indices]
                    if prefill_logprobs
                    else speculative_logits
Nicolas Patry's avatar
Nicolas Patry committed
1105
                )
1106
1107
1108
        else:
            next_token_logits = out

Nicolas Patry's avatar
Nicolas Patry committed
1109
        speculate = get_speculate()
OlivierDehaene's avatar
OlivierDehaene committed
1110
1111
1112
1113
1114
1115
1116
1117
1118
        (
            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
1119
            speculate,
OlivierDehaene's avatar
OlivierDehaene committed
1120
1121
            batch.speculative_ids,
            speculative_logits,
1122
1123
        )

Nicolas Patry's avatar
Nicolas Patry committed
1124
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
1125
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1126
1127
        )

1128
        if prefill:
1129
            if len(batch) > 1 and prefill_logprobs:
1130
1131
                # 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
1132
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
1133
1134

            next_position_ids = batch.position_ids.new_empty(len(batch))
1135
1136
1137
            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
1138
1139
1140
1141
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

1142
1143
1144
1145
1146
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
1147
        stopped = True
1148
1149

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

1152
1153
1154
1155
        # 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

1156
        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1157
        index = 0
OlivierDehaene's avatar
OlivierDehaene committed
1158
        for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator):
1159
            # Indexing metadata
1160
1161
1162
            start_index = cumulative_length
            end_index = cumulative_length + input_length

1163
            if prefill:
1164
1165
1166
1167
1168
                # 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

1169
1170
1171
1172
1173
1174
                # 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
1175
1176
                if prefill_logprobs:
                    if len(batch) > 1:
drbh's avatar
drbh committed
1177
1178
1179
                        prefill_tokens_indices[out_start_index : out_end_index - 1] = (
                            batch.input_ids[start_index + 1 : start_index + out_length]
                        )
1180
1181
1182
1183
1184
                    else:
                        # Set prefill_tokens_indices to the correct slice
                        prefill_tokens_indices = batch.input_ids[
                            start_index + 1 : start_index + out_length
                        ]
1185

Nicolas Patry's avatar
Nicolas Patry committed
1186
1187
1188
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1189
1190
1191

            cumulative_length += input_length

drbh's avatar
drbh committed
1192
        # Update values
Nicolas Patry's avatar
Nicolas Patry committed
1193
1194
1195
1196
1197
        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
1198

1199
        if prefill and prefill_logprobs:
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
            # 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
1210
        next_token_ids = next_input_ids.tolist()
1211
1212
        accepted_ids = accepted_ids.tolist()
        start_decode = time.time_ns()
1213
1214
1215
1216
1217

        # Zipped iterator
        iterator = zip(
            batch.requests,
            batch.input_lengths,
1218
1219
            batch.prefix_offsets,
            batch.read_offsets,
1220
1221
            batch.stopping_criterias,
            batch.all_input_ids,
1222
1223
            batch.next_token_chooser.do_sample,
            batch.next_token_chooser.seeds,
Nicolas Patry's avatar
Nicolas Patry committed
1224
            batch.top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1225
            accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1226
1227
            batch_top_token_ids,
            batch_top_token_logprobs,
1228
1229
1230
        )

        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1231
        index = 0
1232
1233
1234
        for i, (
            request,
            input_length,
1235
1236
            prefix_offset,
            read_offset,
1237
1238
            stopping_criteria,
            all_input_ids,
1239
1240
            do_sample,
            seed,
Nicolas Patry's avatar
Nicolas Patry committed
1241
            top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1242
            n_accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1243
1244
            top_token_ids,
            top_token_logprobs,
1245
        ) in enumerate(iterator):
1246
            # Append next token to all tokens
Nicolas Patry's avatar
Nicolas Patry committed
1247
1248
1249
            next_token_texts = []
            left = 0

1250
1251
1252
1253
            if n_accepted_ids > 1:
                if RANK == 0:
                    logger.debug(f"Speculated ids {n_accepted_ids - 1}")

Nicolas Patry's avatar
Nicolas Patry committed
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
            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)
1265

Nicolas Patry's avatar
Nicolas Patry committed
1266
1267
1268
1269
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1270

Nicolas Patry's avatar
Nicolas Patry committed
1271
1272
1273
1274
1275
1276
1277
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1278

OlivierDehaene's avatar
OlivierDehaene committed
1279
1280
1281
1282
            _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
1283
            index += n_accepted_ids
1284

1285
1286
1287
1288
1289
            # Shard generations
            # All generations will be appended in the rust sharded client
            if i % self.world_size == self.rank:
                if stop:
                    # Decode generated tokens
1290
1291
                    output_text, _, _ = self.decode_token(
                        all_input_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1292
1293
1294
1295
1296
1297
                        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,
1298
1299
                    )
                    generated_text = GeneratedText(
1300
1301
1302
1303
                        output_text,
                        stopping_criteria.current_tokens,
                        reason,
                        seed if do_sample else None,
1304
1305
1306
1307
1308
                    )
                else:
                    generated_text = None

                # Prefill
1309
1310
1311
1312
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

1313
1314
                    # Remove generated token to only have prefill and add nan for first prompt token
                    request_prefill_logprobs = [float("nan")] + prefill_logprobs[
1315
                        out_start_index : out_end_index - 1
1316
1317
1318
1319
1320
1321
1322
                    ]
                    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
1323
1324

                    prefill_tokens = Tokens(
OlivierDehaene's avatar
OlivierDehaene committed
1325
1326
1327
1328
                        prefill_token_ids,
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1329
1330
1331
1332
                    )
                else:
                    prefill_tokens = None

Nicolas Patry's avatar
Nicolas Patry committed
1333
                if top_n_tokens > 0:
Nicolas Patry's avatar
Nicolas Patry committed
1334
                    all_top_tokens = []
drbh's avatar
drbh committed
1335
                    for top_token_ids, top_token_logprobs in zip(
1336
1337
                        top_token_ids, top_token_logprobs
                    ):
Nicolas Patry's avatar
Nicolas Patry committed
1338
1339
1340
1341
1342
1343
                        toptoken_texts = self.tokenizer.batch_decode(
                            top_token_ids,
                            clean_up_tokenization_spaces=False,
                            skip_special_tokens=False,
                        )
                        special_toptokens = [
1344
1345
                            token_id in self.all_special_ids
                            for token_id in top_token_ids
Nicolas Patry's avatar
Nicolas Patry committed
1346
1347
1348
1349
1350
1351
1352
1353
1354
                        ]
                        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
1355
1356
1357
                else:
                    top_tokens = None

1358
1359
1360
                generation = Generation(
                    request.id,
                    prefill_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1361
1362
1363
1364
1365
1366
                    Tokens(
                        _next_token_ids,
                        _next_token_logprobs,
                        next_token_texts,
                        [nid in self.all_special_ids for nid in _next_token_ids],
                    ),
1367
                    generated_text,
Nicolas Patry's avatar
Nicolas Patry committed
1368
                    top_tokens,
1369
1370
                )

1371
                generations.append(generation)
1372

drbh's avatar
drbh committed
1373
1374
1375
            # 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
1376
1377
1378
                batch.next_token_chooser = (
                    batch.next_token_chooser.advance_grammar_single(i, next_token_id)
                )
drbh's avatar
drbh committed
1379

1380
            # Update values
1381
            batch.input_lengths[i] = input_length + n_accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1382
1383
            if batch.input_lengths[i] > batch.max_seqlen:
                batch.max_seqlen = batch.input_lengths[i]
1384
1385
            batch.prefix_offsets[i] = prefix_offset
            batch.read_offsets[i] = read_offset
1386
1387
            batch.all_input_ids[i] = all_input_ids

1388
1389
        if stopped:
            # No need to return a batch if we know that all requests stopped
1390
1391
1392
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1393

1394
1395
1396
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1397

1398
1399
1400
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)