flash_causal_lm.py 58.2 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
Daniël de Kok's avatar
Daniël de Kok committed
14
from typing import Iterable, Optional, Tuple, List, Type, Dict
fxmarty's avatar
fxmarty committed
15

drbh's avatar
drbh committed
16
from text_generation_server.adapters import AdapterBatchData, AdapterBatchMetadata
fxmarty's avatar
fxmarty committed
17
from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
Daniël de Kok's avatar
Daniël de Kok committed
18
from text_generation_server.utils.chunks import concat_text_chunks
Nicolas Patry's avatar
Nicolas Patry committed
19
from text_generation_server.utils.import_utils import SYSTEM
OlivierDehaene's avatar
OlivierDehaene committed
20
from text_generation_server.models import Model
21
from text_generation_server.utils.tokens import batch_top_tokens
22
from text_generation_server.utils.dist import RANK
Nicolas Patry's avatar
Nicolas Patry committed
23
from text_generation_server.utils.speculate import get_speculate
24
25
from text_generation_server.models.types import (
    Batch,
Nicolas Patry's avatar
Nicolas Patry committed
26
    Tokens,
27
28
29
30
    Generation,
    GeneratedText,
)
from text_generation_server.pb import generate_pb2
Nicolas Patry's avatar
Nicolas Patry committed
31
32
33
34
35
36
from text_generation_server.models.globals import (
    MEM_POOL,
    CUDA_GRAPHS,
    get_adapter_to_index,
    MODEL_ID,
)
37
from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
38
from text_generation_server.utils.dist import MEMORY_FRACTION
drbh's avatar
drbh committed
39
from text_generation_server.utils.segments import SegmentConcatBuilder, find_segments
40

Nicolas Patry's avatar
Nicolas Patry committed
41
from text_generation_server.utils.import_utils import (
Nicolas Patry's avatar
Nicolas Patry committed
42
43
44
    empty_cache,
    synchronize,
    get_free_memory,
Nicolas Patry's avatar
Nicolas Patry committed
45
46
)

Nicolas Patry's avatar
Nicolas Patry committed
47
48
tracer = trace.get_tracer(__name__)

49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
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

64

65
66
67
68
@dataclass
class FlashCausalLMBatch(Batch):
    batch_id: int
    requests: List[generate_pb2.Request]
69
70
    # request id -> idx in list mapping
    requests_idx_mapping: Dict[int, int]
71
72

    # Decoder values
73
74
    input_ids: torch.Tensor
    position_ids: torch.Tensor
75
    speculative_ids: Optional[torch.Tensor]
76

77
78
79
80
    # 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]
81
82
83
    # 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]
84
85
86
87
88
89
90
91
92
93

    # 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
94
    block_tables: List[List[int]]
95
    # tensor of size [b, max_total_seqlen // block_size] holding the paged attention block tables for all sequences
96
    block_tables_tensor: torch.Tensor
97
    # tensor of length \sum_{i=0}^{b} max_s_i  holding the paged attention slots for all sequences
98
    slots: torch.Tensor
99

100
101
    max_seqlen: int

102
103
104
105
106
    # 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]]

107
108
    # All tokens
    all_input_ids: List[List[int]]
109
    all_input_ids_tensor: torch.Tensor
110
111
112

    # Lengths of all generations present in the batch
    input_lengths: List[int]
113
    input_lengths_tensor: torch.Tensor
114
115
    prefix_offsets: List[Optional[int]]
    read_offsets: List[Optional[int]]
116
117

    # Generation helpers
118
    next_token_chooser: HeterogeneousNextTokenChooser
119
    stopping_criterias: List[StoppingCriteria]
Nicolas Patry's avatar
Nicolas Patry committed
120
121
    top_n_tokens: List[int]
    top_n_tokens_tensor: torch.Tensor
122

drbh's avatar
drbh committed
123
124
125
    # Adapter metadata for each request
    adapter_meta: AdapterBatchMetadata

126
    # Number of blocks in this batch
127
    num_blocks: int
128
129
    # Maximum number of blocks
    max_blocks: int
130

131
132
    def to_pb(self) -> generate_pb2.CachedBatch:
        return generate_pb2.CachedBatch(
133
            id=self.batch_id,
134
            request_ids=[r.id for r in self.requests],
135
            size=len(self),
136
            max_tokens=self.num_blocks * BLOCK_SIZE,
137
138
139
        )

    @classmethod
Daniël de Kok's avatar
Daniël de Kok committed
140
141
142
    def batch_tokenized_inputs(
        cls, requests: Iterable[generate_pb2.Request], tokenizer
    ):
143
144
        batch_inputs = []
        max_truncation = 0
145
        for r in requests:
Daniël de Kok's avatar
Daniël de Kok committed
146
            batch_inputs.append(concat_text_chunks(r.input_chunks.chunks))
147
148
149
150
151
            max_truncation = max(max_truncation, r.truncate)

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

drbh's avatar
drbh committed
154
155
156
157
158
159
160
161
162
    @classmethod
    def from_tokenized(
        cls,
        pb: generate_pb2.Batch,
        tokenizer: PreTrainedTokenizerBase,
        batch_tokenized_inputs,
        dtype: torch.dtype,
        device: torch.device,
    ) -> "FlashCausalLMBatch":
163
        sliding_window = get_sliding_windows()
164
        position_ids = []
165
        cu_seqlen_prefill = [0]
166
167
        start_slots = []
        slot_indices = []
168
        prefill_cache_indices = []
169
170

        input_lengths = []
171
172
        prefix_offsets = []
        read_offsets = []
173
        all_input_ids = []
174
        requests_idx_mapping = {}
175

176
177
178
179
180
181
        all_prefill_logprobs = True
        no_prefill_logprobs = True
        prefill_head_indices = []
        prefill_next_token_indices = []
        prefill_cu_outlens = [0]

182
        next_token_chooser_parameters = []
183
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
184
        top_n_tokens = []
185

drbh's avatar
drbh committed
186
187
188
        adapter_indices_list = []
        adapter_set = set()

189
190
        # Cumulative length
        cumulative_length = 0
191
        cumulative_max_length = 0
192
        prefill_out_cumulative_length = 0
193

194
        num_blocks = 0
195
        max_seqlen = 0
196
        max_length = 0
197
        max_blocks = 0
198

199
200
201
        block_tables = []
        slots = []

202
        # Parse batch
203
204
205
        for i, (r, tokenized_input) in enumerate(
            zip(pb.requests, batch_tokenized_inputs)
        ):
206
207
208
            # request id -> idx in list mapping
            requests_idx_mapping[r.id] = i

209
            tokenized_input = tokenized_input[-r.truncate :]
210
211
212
213
214
            if (
                tokenized_input[0] == tokenizer.bos_token_id
                and tokenized_input[1] == tokenizer.bos_token_id
            ):
                tokenized_input = tokenized_input[1:]
215

216
217
            input_length = len(tokenized_input)
            input_lengths.append(input_length)
218

219
            prefix_offsets.append(input_length - 5)
220
            read_offsets.append(input_length)
221

222
            all_input_ids.append(tokenized_input)
223
224

            # Position ids
225
226
            request_position_ids = torch.arange(0, input_length, dtype=torch.int32)
            position_ids.append(request_position_ids)
227
228

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

231
            next_token_chooser_parameters.append(r.parameters)
232

233
234
235
            stopping_criteria = StoppingCriteria.from_pb(
                r.stopping_parameters, tokenizer
            )
236
            max_new_tokens = stopping_criteria.max_new_tokens
237
            stopping_criterias.append(stopping_criteria)
Nicolas Patry's avatar
Nicolas Patry committed
238
            top_n_tokens.append(r.top_n_tokens)
239

Nicolas Patry's avatar
Nicolas Patry committed
240
241
            ADAPTER_TO_INDEX = get_adapter_to_index()
            adapter_index = ADAPTER_TO_INDEX.get(r.adapter_id, 0)
drbh's avatar
drbh committed
242
243
244
            adapter_indices_list.append(torch.full((input_length,), adapter_index))
            adapter_set.add(adapter_index)

245
246
            # Paged attention
            # Remove one as the first token des not have a past
Nicolas Patry's avatar
Nicolas Patry committed
247
            speculative_length = get_speculate()
drbh's avatar
drbh committed
248
            speculative_length = 0 if speculative_length is None else speculative_length
Nicolas Patry's avatar
Nicolas Patry committed
249
            total_tokens = input_length + max_new_tokens - 1 + speculative_length
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268

            # 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)
269
270
271
272
273
274
275
276
277
            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)

278
279
280
281
282
283
284
285
286
            # 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)

287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
            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

307
308
            # Update
            cumulative_length += input_length
309
310
            cumulative_max_length += total_tokens
            max_seqlen = max(max_seqlen, input_length)
311
            max_blocks = max(max_blocks, len(request_blocks))
OlivierDehaene's avatar
OlivierDehaene committed
312
313
314
            max_length = max(
                max_length, input_length + max_new_tokens + speculative_length
            )
315

drbh's avatar
drbh committed
316
317
318
319
        adapter_indices = torch.cat(adapter_indices_list).to(
            dtype=torch.int64, device=device
        )

320
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
drbh's avatar
drbh committed
321
            next_token_chooser_parameters, dtype, device, tokenizer
322
        )
323
        start_slots = torch.tensor(start_slots, dtype=torch.int64)
324
325
326
327
328
329
330

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

332
333
334
335
336
        # Create tensors on device
        all_input_ids_tensor = torch.tensor(
            all_input_ids_tensor, dtype=torch.int64, device=device
        )

337
338
339
        if len(pb.requests) > 1:
            input_ids = np.concatenate(all_input_ids, dtype=np.int64)
            position_ids = torch.cat(position_ids)
340
            slot_indices = torch.cat(slot_indices)
341
342
            if sliding_window is not None:
                prefill_cache_indices = torch.cat(prefill_cache_indices)
343
344
345
        else:
            input_ids = all_input_ids[0]
            position_ids = position_ids[0]
346
            slot_indices = slot_indices[0]
347
348
            if sliding_window is not None:
                prefill_cache_indices = prefill_cache_indices[0]
349

350
351
        cu_seqlen_prefill = torch.tensor(
            cu_seqlen_prefill, device=device, dtype=torch.int32
352
353
354
        )
        position_ids = position_ids.to(device)
        slot_indices = slot_indices.to(device)
355
356
357
        prefill_cache_indices = (
            prefill_cache_indices.to(device) if sliding_window is not None else None
        )
358
        input_ids = torch.tensor(input_ids, dtype=torch.int64, device=device)
359
360
        input_lengths_tensor = torch.tensor(
            input_lengths, dtype=torch.int32, device=device
361
        )
362

drbh's avatar
drbh committed
363
364
365
366
367
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

368
369
        if all_prefill_logprobs:
            prefill_head_indices = None
370
            prefill_next_token_indices = cu_seqlen_prefill[1:] - 1
371
        elif no_prefill_logprobs:
372
            prefill_head_indices = cu_seqlen_prefill[1:] - 1
373
374
375
376
377
378
379
380
            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
381
382
383
        top_n_tokens_tensor = torch.tensor(
            top_n_tokens, device=device, dtype=torch.int64
        )
384

385
386
387
388
389
390
391
392
        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)

393
394
395
        return cls(
            batch_id=pb.id,
            requests=pb.requests,
396
            requests_idx_mapping=requests_idx_mapping,
397
398
            input_ids=input_ids,
            position_ids=position_ids,
399
            cu_seqlen_prefill=cu_seqlen_prefill,
400
            prefill_cache_indices=prefill_cache_indices,
401
402
            start_slots=start_slots,
            slot_indices=slot_indices,
403
404
405
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
406
            max_seqlen=max_seqlen,
407
408
409
            prefill_head_indices=prefill_head_indices,
            prefill_next_token_indices=prefill_next_token_indices,
            prefill_cu_outlens=prefill_cu_outlens,
410
            input_lengths=input_lengths,
411
            input_lengths_tensor=input_lengths_tensor,
412
413
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
414
            all_input_ids=all_input_ids,
415
416
            all_input_ids_tensor=all_input_ids_tensor,
            next_token_chooser=next_token_chooser,
417
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
418
419
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
420
            num_blocks=num_blocks,
421
            max_blocks=max_blocks,
drbh's avatar
drbh committed
422
423
424
425
426
427
            adapter_meta=AdapterBatchMetadata(
                adapter_indices=adapter_indices,
                adapter_set=adapter_set,
                adapter_segments=adapter_segments,
                segment_indices=adapter_segment_indices,
            ),
Nicolas Patry's avatar
Nicolas Patry committed
428
            speculative_ids=None,
429
430
        )

431
432
433
434
435
436
437
438
439
440
441
    @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)

442
    @tracer.start_as_current_span("filter")
443
444
    def filter(self, request_ids: List[int]) -> "FlashCausalLMBatch":
        if len(request_ids) == 0:
445
446
            raise ValueError("Batch must have at least one request")
        # We assume that if len(requests) == len(self) then the requests are the same
447
        if len(request_ids) == len(self):
448
449
            return self

450
        device = self.input_ids.device
451

452
453
454
        # New values after filtering
        requests_idx_mapping = {}

455
456
457
        # Used to index into tensors
        indices = []

458
459
460
        # slots to keep after filtering
        slot_filtering_indices = torch.zeros(
            self.slots.shape[0], dtype=torch.bool, device=device
461
462
        )

463
        # Create on CPU to only move to GPU once instead of at every copy
464
        slot_indices = torch.empty(len(request_ids), dtype=torch.int64)
465
466
        max_seqlen = 0

467
        requests = []
468
469
        start_slots = []
        block_tables = []
470
471
        all_input_ids = []

472
        input_lengths = []
473
474
        prefix_offsets = []
        read_offsets = []
475

476
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
477
        top_n_tokens = []
drbh's avatar
drbh committed
478
        adapter_set = set()
479

480
        num_blocks = 0
481
482
483
484
        max_blocks = 0
        # Cumulative length
        cumulative_max_length = 0

485
486
        for i, request_id in enumerate(request_ids):
            idx = self.requests_idx_mapping[request_id]
487
            indices.append(idx)
488
489
490
            requests_idx_mapping[request_id] = i

            requests.append(self.requests[idx])
491
492
493
494

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

496
497
498
            all_input_ids.append(self.all_input_ids[idx])

            input_lengths.append(request_input_length)
499
500
            prefix_offsets.append(self.prefix_offsets[idx])
            read_offsets.append(self.read_offsets[idx])
501

502
503
            stopping_criteria = self.stopping_criterias[idx]
            stopping_criterias.append(stopping_criteria)
504

Nicolas Patry's avatar
Nicolas Patry committed
505
506
            top_n_tokens.append(self.top_n_tokens[idx])

Nicolas Patry's avatar
Nicolas Patry committed
507
508
            ADAPTER_TO_INDEX = get_adapter_to_index()
            adapter_index = ADAPTER_TO_INDEX.get(self.requests[idx].adapter_id, 0)
drbh's avatar
drbh committed
509
510
            adapter_set.add(adapter_index)

511
            remaining_tokens = (
512
513
                stopping_criteria.max_new_tokens - stopping_criteria.current_tokens
            )
514

515
            request_block_table = self.block_tables[idx]
516
            num_blocks += len(request_block_table)
517
518
519
            block_tables.append(request_block_table)
            start_slots.append(cumulative_max_length)

520
            # Copy to tensor (CPU)
521
            slot_indices[i] = cumulative_max_length + request_input_length - 1
522
523

            # Set slice
524
525
526
527
528
            slot_filtering_indices[
                self.start_slots[idx] : self.start_slots[idx]
                + request_input_length
                + remaining_tokens
                - 1
529
530
531
            ] = True

            cumulative_max_length += request_input_length + remaining_tokens - 1
532

533
534
            max_blocks = max(max_blocks, len(request_block_table))

535
536
537
        # Index into tensors
        input_ids = self.input_ids[indices]
        position_ids = self.position_ids[indices]
drbh's avatar
drbh committed
538
        adapter_indices = self.adapter_meta.adapter_indices[indices]
539
        all_input_ids_tensor = self.all_input_ids_tensor[indices]
540
541
542
        block_tables_tensor = self.block_tables_tensor[indices]
        input_lengths_tensor = self.input_lengths_tensor[indices]
        slots = self.slots[slot_filtering_indices]
543
        next_token_chooser = self.next_token_chooser.filter(indices)
Nicolas Patry's avatar
Nicolas Patry committed
544
        top_n_tokens_tensor = self.top_n_tokens_tensor[indices]
OlivierDehaene's avatar
OlivierDehaene committed
545
546
547
        speculative_ids = (
            self.speculative_ids[indices] if self.speculative_ids is not None else None
        )
548
549

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

551
        # Move to GPU now that we have the whole tensor
552
        slot_indices = slot_indices.to(device)
553

drbh's avatar
drbh committed
554
555
556
557
558
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

559
        return type(self)(
560
561
562
563
564
            batch_id=self.batch_id,
            requests=requests,
            requests_idx_mapping=requests_idx_mapping,
            input_ids=input_ids,
            position_ids=position_ids,
565
            cu_seqlen_prefill=None,
566
            prefill_cache_indices=None,
567
568
569
570
571
            start_slots=start_slots,
            slot_indices=slot_indices,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
572
            max_seqlen=max_seqlen,
573
574
575
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
576
            input_lengths=input_lengths,
577
            input_lengths_tensor=input_lengths_tensor,
578
579
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
580
581
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
582
            next_token_chooser=next_token_chooser,
583
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
584
585
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
586
            num_blocks=num_blocks,
587
            max_blocks=max_blocks,
Nicolas Patry's avatar
Nicolas Patry committed
588
            speculative_ids=speculative_ids,
drbh's avatar
drbh committed
589
590
591
592
593
594
            adapter_meta=AdapterBatchMetadata(
                adapter_indices=adapter_indices,
                adapter_set=adapter_set,
                adapter_segments=adapter_segments,
                segment_indices=adapter_segment_indices,
            ),
595
596
597
598
599
600
601
602
603
        )

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

604
        num_blocks = 0
605
606
607
608
609
610
611
612
        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)
613
            num_blocks += b.num_blocks
OlivierDehaene's avatar
OlivierDehaene committed
614
615
616
            speculative_length = (
                b.speculative_ids.shape[1] if b.speculative_ids is not None else 0
            )
617
618
619
620
621
622
623
            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
624
                    + speculative_length
625
626
627
628
629
630
                    - stopping_criteria.current_tokens
                    for input_length, stopping_criteria in zip(
                        b.input_lengths, b.stopping_criterias
                    )
                ),
            )
631
632
633

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
634
635
636
637
638
639
640
641
642
643
        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)
644
        )
Nicolas Patry's avatar
Nicolas Patry committed
645
646
647
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
drbh's avatar
drbh committed
648
649
650
651
652
653
654
655
        total_indices_size = sum(
            b.adapter_meta.adapter_indices.shape[0] for b in batches
        )
        adapter_indices = batches[0].adapter_meta.adapter_indices.new_empty(
            total_indices_size
        )
        adapter_set = set()
        adapter_segment_builder = SegmentConcatBuilder()
656

657
658
        start_slots = []
        block_tables = []
659
660
661
        all_input_ids = []

        input_lengths = []
662
663
        prefix_offsets = []
        read_offsets = []
664

665
        next_token_chooser_parameters = []
666
        fsm_grammar_states = []
667
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
668
        top_n_tokens = []
669

670
        # Cumulative length
671
        cumulative_batch_size = 0
672
        cumulative_slots = 0
drbh's avatar
drbh committed
673
        cumulative_adapter_indices_size = 0
674
675
676

        for i, batch in enumerate(batches):
            requests.extend(batch.requests)
677
678
679
680
681
682
683
684

            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

685
686
            start_index = cumulative_batch_size
            end_index = cumulative_batch_size + len(batch)
687
688
            slots_start_index = cumulative_slots
            slots_end_index = cumulative_slots + len(batch.slots)
689
690
691
692

            # Copy tensors (GPU)
            input_ids[start_index:end_index] = batch.input_ids
            position_ids[start_index:end_index] = batch.position_ids
693
694
            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
695
            top_n_tokens_tensor[start_index:end_index] = batch.top_n_tokens_tensor
696
            slots[slots_start_index:slots_end_index] = batch.slots
697

drbh's avatar
drbh committed
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
            # Copy over adapter indices
            adapter_start_index = cumulative_adapter_indices_size
            adapter_end_index = (
                cumulative_adapter_indices_size
                + batch.adapter_meta.adapter_indices.shape[0]
            )
            adapter_indices[adapter_start_index:adapter_end_index] = (
                batch.adapter_meta.adapter_indices
            )
            cumulative_adapter_indices_size = adapter_end_index
            adapter_set.update(batch.adapter_meta.adapter_set)
            adapter_segment_builder.concat(
                batch.adapter_meta.adapter_segments, batch.adapter_meta.segment_indices
            )

713
714
715
            all_input_ids_tensor[
                start_index:end_index, : batch.all_input_ids_tensor.shape[1]
            ] = batch.all_input_ids_tensor[:, :max_length]
716

717
718
719
            block_tables_tensor[
                start_index:end_index, : batch.block_tables_tensor.shape[1]
            ] = batch.block_tables_tensor[:, :max_blocks]
720

721
722
723
            start_slots.append(batch.start_slots + cumulative_slots)

            block_tables.extend(batch.block_tables)
724
725
            all_input_ids.extend(batch.all_input_ids)

726
            input_lengths.extend(batch.input_lengths)
727
728
            prefix_offsets.extend(batch.prefix_offsets)
            read_offsets.extend(batch.read_offsets)
729

730
            next_token_chooser_parameters.extend([r.parameters for r in batch.requests])
731
            fsm_grammar_states.extend(batch.next_token_chooser.fsm_grammar_states)
732
733
            stopping_criterias.extend(batch.stopping_criterias)

Nicolas Patry's avatar
Nicolas Patry committed
734
735
            top_n_tokens.extend(batch.top_n_tokens)

736
            # Update
737
            cumulative_batch_size += len(batch)
738
            cumulative_slots += len(batch.slots)
739

740
        start_slots = torch.concat(start_slots)
741

742
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
743
744
745
            next_token_chooser_parameters,
            dtype=batches[0].next_token_chooser.dtype,
            device=batches[0].next_token_chooser.device,
drbh's avatar
drbh committed
746
            tokenizer=batches[0].next_token_chooser.tokenizer,
747
            fsm_grammar_states=fsm_grammar_states,
748
749
        )

OlivierDehaene's avatar
OlivierDehaene committed
750
751
752
753
754
        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
755

drbh's avatar
drbh committed
756
757
        adapter_segments, adapter_segment_indices = adapter_segment_builder.build()

758
        return cls(
759
760
            batch_id=batches[0].batch_id,
            requests=requests,
761
            requests_idx_mapping=requests_idx_mapping,
762
763
            input_ids=input_ids,
            position_ids=position_ids,
764
            cu_seqlen_prefill=None,
765
            prefill_cache_indices=None,
766
767
768
769
770
            start_slots=start_slots,
            slot_indices=slot_indices,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
771
            max_seqlen=max_seqlen,
772
773
774
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
775
            input_lengths=input_lengths,
776
            input_lengths_tensor=input_lengths_tensor,
777
778
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
779
780
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
781
            next_token_chooser=next_token_chooser,
782
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
783
784
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
785
            num_blocks=num_blocks,
786
            max_blocks=max_blocks,
OlivierDehaene's avatar
OlivierDehaene committed
787
            speculative_ids=speculative_ids,
drbh's avatar
drbh committed
788
789
790
791
792
793
            adapter_meta=AdapterBatchMetadata(
                adapter_indices=adapter_indices,
                adapter_set=adapter_set,
                adapter_segments=adapter_segments,
                segment_indices=adapter_segment_indices,
            ),
794
795
796
797
798
799
800
801
802
        )

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


class FlashCausalLM(Model):
    def __init__(
        self,
drbh's avatar
drbh committed
803
        model_id: str,
804
805
806
807
808
809
810
811
812
        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,
813
        sliding_window: Optional[int] = None,
814
    ):
815
816
817
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_size = head_size
818

819
        self.cuda_graphs = {}
820
        self.kv_cache = []
821

822
        super(FlashCausalLM, self).__init__(
drbh's avatar
drbh committed
823
            model_id=model_id,
824
            model=model,
825
826
827
828
            tokenizer=tokenizer,
            requires_padding=False,
            dtype=dtype,
            device=device,
829
830
            rank=rank,
            world_size=world_size,
831
            sliding_window=sliding_window,
832
833
834
835
836
837
        )

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

838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
    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()
Wang, Yi's avatar
Wang, Yi committed
854
855
856
857
        if SYSTEM == "ipex" and device.type == "xpu":
            x = 1
        else:
            x = BLOCK_SIZE // element_size
858

Nicolas Patry's avatar
Nicolas Patry committed
859
        if SYSTEM == "ipex" and device == torch.device("cpu"):
Wang, Yi's avatar
Wang, Yi committed
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
            self.kv_cache = [
                (
                    torch.empty(
                        (num_blocks, num_heads, BLOCK_SIZE, head_size),
                        dtype=dtype,
                        device=device,
                    ),
                    torch.empty(
                        (num_blocks, num_heads, BLOCK_SIZE, head_size),
                        dtype=dtype,
                        device=device,
                    ),
                )
                for _ in range(num_layers)
            ]
        else:
            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)
            ]
891

892
893
894
    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)
895
        slots = torch.arange(bs, dtype=torch.int64, device=self.device)
896
897
898
899
900
901
902
903
904
905
        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,
906
            "kv_cache": self.kv_cache,
907
908
909
910
911
912
913
914
915
916
917
918
919
            "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,
920
            kv_cache=self.kv_cache,
921
922
923
924
            block_tables=block_tables,
            slots=slots,
            input_lengths=input_lengths,
            max_s=max_s,
925
            prefill_cache_indices=None,
926
927
928
929
930
            lm_head_indices=None,
        )
        torch.cuda.synchronize()

        with torch.cuda.graph(graph, pool=MEM_POOL):
931
            logits, speculative_logits = self.model.forward(
932
933
934
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlen_prefill=None,
935
                kv_cache=self.kv_cache,
936
937
938
939
                block_tables=block_tables,
                slots=slots,
                input_lengths=input_lengths,
                max_s=max_s,
940
                prefill_cache_indices=None,
941
942
                lm_head_indices=None,
            )
943
944
            self.cuda_graphs[bs]["logits"] = logits
            self.cuda_graphs[bs]["speculative_logits"] = speculative_logits
945
946
        torch.cuda.synchronize()

947
    def warmup(self, batch: FlashCausalLMBatch):
948
        # The warmup batch is the biggest batch we could ever receive
Nicolas Patry's avatar
Nicolas Patry committed
949
950
        empty_cache()

951
        try:
952
953
            self.init_kv_cache(
                batch.num_blocks,
954
955
956
957
958
959
                self.num_layers,
                self.num_kv_heads,
                self.head_size,
                self.dtype,
                self.device,
            )
960
            max_bt = batch.max_blocks
961
            max_s = max_bt * BLOCK_SIZE
fxmarty's avatar
fxmarty committed
962
963
964

            if SYSTEM == "rocm" and os.environ.get("PYTORCH_TUNABLEOP_ENABLED", False):
                torch.cuda.tunable.tuning_enable(False)
965
            _, batch, _ = self.generate_token(batch)
OlivierDehaene's avatar
OlivierDehaene committed
966
        except torch.cuda.OutOfMemoryError as e:
967
            raise RuntimeError(
968
969
                f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. "
                f"You need to decrease `--max-batch-prefill-tokens`"
970
            ) from e
971

Nicolas Patry's avatar
Nicolas Patry committed
972
        synchronize(self.device)
973

974
975
        # 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
976
977
978
979
        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
980
        free_memory = get_free_memory(self.device, MEMORY_FRACTION)
drbh's avatar
drbh committed
981
        batch_num_blocks = batch.num_blocks if batch is not None else 0
982
983

        num_blocks = (
984
985
            # Leave 5% for some wiggle room
            int((free_memory * 0.95) // total_cache_size)
986
            # Add batch.num_blocks as we allocated it above, so it is included in the peak memory.
drbh's avatar
drbh committed
987
            + batch_num_blocks
988
989
        )

990
        del batch
991

992
        self.init_kv_cache(
993
994
995
996
997
998
999
1000
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
            self.dtype,
            self.device,
        )

fxmarty's avatar
fxmarty committed
1001
1002
1003
1004
1005
        if SYSTEM == "rocm":
            if (
                os.environ.get("PYTORCH_TUNABLEOP_ENABLED") is None
                or os.environ.get("PYTORCH_TUNABLEOP_ENABLED") == "1"
            ):
1006
1007
                torch.cuda.tunable.enable()

fxmarty's avatar
fxmarty committed
1008
1009
1010
1011
1012
1013
1014
1015
                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(",")
                    ]
1016
                elif CUDA_GRAPHS is not None:
fxmarty's avatar
fxmarty committed
1017
                    tuning_sequences = CUDA_GRAPHS
1018
1019
1020
                else:
                    # For seqlen = 1, we dispatch to LLMM1 kernel.
                    tuning_sequences = [2, 3, 4, 5, 6, 7]
fxmarty's avatar
fxmarty committed
1021
1022
1023

                tunableop_filepath = os.path.join(
                    HUGGINGFACE_HUB_CACHE,
Nicolas Patry's avatar
Nicolas Patry committed
1024
                    f"tunableop_{MODEL_ID.replace('/', '-')}_tp{self.world_size}_rank{self.rank}.csv",
fxmarty's avatar
fxmarty committed
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
                )

                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."
                )

1049
        if CUDA_GRAPHS:
1050
            try:
1051
                logger.info(f"Cuda Graphs are enabled for sizes {CUDA_GRAPHS}")
1052
                # Warmup cuda graphs
1053
                for bs in CUDA_GRAPHS:
1054
1055
                    if self.speculate is None or self.speculate + 1 <= bs:
                        self.cuda_graph_warmup(bs, max_s, max_bt)
OlivierDehaene's avatar
OlivierDehaene committed
1056
            except torch.cuda.OutOfMemoryError:
1057
                logger.exception(f"Decode cuda graph warmup failed")
1058
1059
        else:
            logger.info(f"Cuda Graphs are disabled (CUDA_GRAPHS={CUDA_GRAPHS}).")
1060

1061
        return int(num_blocks * BLOCK_SIZE)
1062

fxmarty's avatar
fxmarty committed
1063
1064
1065
1066
1067
    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
1068
1069
1070
        # 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
1071
1072
1073
1074
1075
1076
1077
        # 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
            ),
1078
            kv_cache=self.kv_cache,
fxmarty's avatar
fxmarty committed
1079
            block_tables=None,
fxmarty's avatar
fxmarty committed
1080
            input_lengths=input_lengths,
fxmarty's avatar
fxmarty committed
1081
1082
1083
            slots=slots,
            max_s=seqlen,
            lm_head_indices=None,
1084
            prefill_cache_indices=None,
fxmarty's avatar
fxmarty committed
1085
1086
        )

1087
    def forward(
drbh's avatar
drbh committed
1088
        self, batch: FlashCausalLMBatch, adapter_data: AdapterBatchData
1089
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1090
        # Model Forward
Nicolas Patry's avatar
Nicolas Patry committed
1091
        if batch.speculative_ids is not None:
OlivierDehaene's avatar
OlivierDehaene committed
1092
1093
1094
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1095
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1096
1097
1098
1099
1100
            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
1101
1102
1103

            speculative_ids = batch.speculative_ids

OlivierDehaene's avatar
OlivierDehaene committed
1104
            B, speculative_length = speculative_ids.shape
Nicolas Patry's avatar
Nicolas Patry committed
1105
            new_length = speculative_length + 1
OlivierDehaene's avatar
OlivierDehaene committed
1106
1107
1108
            new_input_ids = torch.cat(
                [input_ids.unsqueeze(-1), speculative_ids], dim=1
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1109
1110
            arange = torch.arange(new_length, device=position_ids.device).unsqueeze(0)
            arange_int = arange.to(dtype=torch.int32)
OlivierDehaene's avatar
OlivierDehaene committed
1111
1112
1113
            new_position_ids = (
                position_ids.unsqueeze(-1).expand(B, new_length) + arange
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1114
            slots = (slots.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1)
OlivierDehaene's avatar
OlivierDehaene committed
1115
1116
1117
            input_lengths = (
                input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1118
1119

            # Add Copy the block tables for all members
OlivierDehaene's avatar
OlivierDehaene committed
1120
1121
1122
1123
1124
1125
            block_tables = (
                block_tables.unsqueeze(1)
                .expand(B, new_length, -1)
                .reshape(B * new_length, -1)
                .contiguous()
            )
Nicolas Patry's avatar
Nicolas Patry committed
1126
1127
1128
1129
1130
            max_s = max_s + speculative_length

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
1131
1132
1133
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1134
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1135
1136
1137
1138
1139
            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
1140

1141
1142
1143
1144
1145
1146
        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)

1147
        bs = input_ids.shape[0]
OlivierDehaene's avatar
OlivierDehaene committed
1148
1149
1150
1151
1152
1153
1154
1155
        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:
1156
            logits, speculative_logits = self.model.forward(
1157
1158
1159
1160
1161
1162
1163
1164
                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,
1165
                prefill_cache_indices=batch.prefill_cache_indices,
1166
                lm_head_indices=lm_head_indices,
drbh's avatar
drbh committed
1167
                adapter_data=adapter_data,
1168
            )
1169
1170
1171
            if batch.prefill_cache_indices is not None:
                batch.prefill_cache_indices = None
            return logits, speculative_logits
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187

        # 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
1188
1189
1190
1191
1192
1193
1194
        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
1195
1196
1197
1198

    @tracer.start_as_current_span("generate_token")
    def generate_token(
        self, batch: FlashCausalLMBatch
1199
1200
    ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]:
        start = time.time_ns()
1201
        prefill = batch.cu_seqlen_prefill is not None
1202
        prefill_logprobs = batch.prefill_next_token_indices is not None
1203

drbh's avatar
drbh committed
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
        # Update adapter indices for speculative tokens (if present)
        adapter_meta = batch.adapter_meta
        if batch.speculative_ids is not None:
            B, speculative_length = batch.speculative_ids.shape
            new_length = speculative_length + 1
            adapter_indices = (
                adapter_meta.adapter_indices.unsqueeze(-1)
                .expand(B, new_length)
                .reshape(-1)
            )
            adapter_segments = adapter_meta.adapter_segments * new_length
            adapter_meta = AdapterBatchMetadata(
                adapter_indices=adapter_indices,
                adapter_set=adapter_meta.adapter_set,
                adapter_segments=adapter_segments,
                segment_indices=adapter_meta.segment_indices,
            )

        # Assign pointers to adapter weights
        # TODO(travis): don't update this if indices haven't changed
        adapter_data = AdapterBatchData.from_meta(
            adapter_meta,
            self.layer_to_adapter_weights,
            prefill,
            batch.prefill_head_indices,
        )

        out, speculative_logits = self.forward(batch, adapter_data)
1232

1233
1234
        if prefill:
            next_token_logits = (
1235
                out[batch.prefill_next_token_indices] if prefill_logprobs else out
1236
            )
Nicolas Patry's avatar
Nicolas Patry committed
1237
1238
            if speculative_logits is not None:
                speculative_logits = (
OlivierDehaene's avatar
OlivierDehaene committed
1239
1240
1241
                    speculative_logits[batch.prefill_next_token_indices]
                    if prefill_logprobs
                    else speculative_logits
Nicolas Patry's avatar
Nicolas Patry committed
1242
                )
drbh's avatar
drbh committed
1243
1244
1245
1246
            next_adapter_indices = batch.adapter_meta.adapter_indices.new_empty(
                len(batch)
            )

1247
1248
        else:
            next_token_logits = out
drbh's avatar
drbh committed
1249
            next_adapter_indices = batch.adapter_meta.adapter_indices
1250

Nicolas Patry's avatar
Nicolas Patry committed
1251
        speculate = get_speculate()
OlivierDehaene's avatar
OlivierDehaene committed
1252
1253
1254
1255
1256
1257
1258
1259
1260
        (
            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
1261
            speculate,
OlivierDehaene's avatar
OlivierDehaene committed
1262
1263
            batch.speculative_ids,
            speculative_logits,
1264
1265
        )

Nicolas Patry's avatar
Nicolas Patry committed
1266
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
1267
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1268
1269
        )

1270
        if prefill:
1271
            if len(batch) > 1 and prefill_logprobs:
1272
1273
                # 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
1274
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
1275
1276

            next_position_ids = batch.position_ids.new_empty(len(batch))
1277
1278
1279
            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
1280
1281
1282
1283
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

1284
1285
1286
1287
1288
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
1289
        stopped = True
1290
1291

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

1294
1295
1296
1297
        # 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

1298
        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1299
        index = 0
OlivierDehaene's avatar
OlivierDehaene committed
1300
        for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator):
1301
            # Indexing metadata
1302
1303
1304
            start_index = cumulative_length
            end_index = cumulative_length + input_length

1305
            if prefill:
1306
1307
1308
1309
1310
                # 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

1311
1312
1313
1314
                # 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]

drbh's avatar
drbh committed
1315
1316
1317
1318
1319
1320
                # Initialize adapter indices
                # In decode, we only have one token per row in the batch, so grab last index
                next_adapter_indices[i] = batch.adapter_meta.adapter_indices[
                    end_index - 1
                ]

1321
1322
                # Used to gather prefill logprobs
                # Copy batch.input_ids to prefill_token_indices
1323
1324
                if prefill_logprobs:
                    if len(batch) > 1:
drbh's avatar
drbh committed
1325
1326
1327
                        prefill_tokens_indices[out_start_index : out_end_index - 1] = (
                            batch.input_ids[start_index + 1 : start_index + out_length]
                        )
1328
1329
1330
1331
1332
                    else:
                        # Set prefill_tokens_indices to the correct slice
                        prefill_tokens_indices = batch.input_ids[
                            start_index + 1 : start_index + out_length
                        ]
1333

Nicolas Patry's avatar
Nicolas Patry committed
1334
1335
1336
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1337
1338
1339

            cumulative_length += input_length

drbh's avatar
drbh committed
1340
        # Update values
Nicolas Patry's avatar
Nicolas Patry committed
1341
1342
1343
1344
1345
        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
drbh's avatar
drbh committed
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
        batch.adapter_meta.adapter_indices = next_adapter_indices

        if prefill:
            # adjust segment lengths to account for all request lengths being 1 during decoding
            adapter_segments, _ = find_segments(batch.adapter_meta.adapter_indices)
            batch.adapter_meta.adapter_segments = torch.tensor(
                adapter_segments,
                dtype=torch.int32,
                device=batch.adapter_meta.adapter_segments.device,
            )
1356

1357
        if prefill and prefill_logprobs:
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
            # 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
1368
        next_token_ids = next_input_ids.tolist()
1369
1370
        accepted_ids = accepted_ids.tolist()
        start_decode = time.time_ns()
1371
1372
1373
1374
1375

        # Zipped iterator
        iterator = zip(
            batch.requests,
            batch.input_lengths,
1376
1377
            batch.prefix_offsets,
            batch.read_offsets,
1378
1379
            batch.stopping_criterias,
            batch.all_input_ids,
1380
1381
            batch.next_token_chooser.do_sample,
            batch.next_token_chooser.seeds,
Nicolas Patry's avatar
Nicolas Patry committed
1382
            batch.top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1383
            accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1384
1385
            batch_top_token_ids,
            batch_top_token_logprobs,
1386
1387
1388
        )

        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1389
        index = 0
1390
1391
1392
        for i, (
            request,
            input_length,
1393
1394
            prefix_offset,
            read_offset,
1395
1396
            stopping_criteria,
            all_input_ids,
1397
1398
            do_sample,
            seed,
Nicolas Patry's avatar
Nicolas Patry committed
1399
            top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1400
            n_accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1401
1402
            top_token_ids,
            top_token_logprobs,
1403
        ) in enumerate(iterator):
1404
            # Append next token to all tokens
Nicolas Patry's avatar
Nicolas Patry committed
1405
1406
1407
            next_token_texts = []
            left = 0

1408
1409
1410
1411
            if n_accepted_ids > 1:
                if RANK == 0:
                    logger.debug(f"Speculated ids {n_accepted_ids - 1}")

Nicolas Patry's avatar
Nicolas Patry committed
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
            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)
1423

Nicolas Patry's avatar
Nicolas Patry committed
1424
1425
1426
1427
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1428

Nicolas Patry's avatar
Nicolas Patry committed
1429
1430
1431
1432
1433
1434
1435
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1436

OlivierDehaene's avatar
OlivierDehaene committed
1437
1438
1439
1440
            _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
1441
            index += n_accepted_ids
1442

1443
1444
1445
1446
1447
            # Shard generations
            # All generations will be appended in the rust sharded client
            if i % self.world_size == self.rank:
                if stop:
                    # Decode generated tokens
1448
1449
                    output_text, _, _ = self.decode_token(
                        all_input_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1450
1451
1452
1453
1454
1455
                        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,
1456
1457
                    )
                    generated_text = GeneratedText(
1458
1459
1460
1461
                        output_text,
                        stopping_criteria.current_tokens,
                        reason,
                        seed if do_sample else None,
1462
1463
1464
1465
1466
                    )
                else:
                    generated_text = None

                # Prefill
1467
1468
1469
1470
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

1471
1472
                    # Remove generated token to only have prefill and add nan for first prompt token
                    request_prefill_logprobs = [float("nan")] + prefill_logprobs[
1473
                        out_start_index : out_end_index - 1
1474
1475
1476
1477
1478
1479
1480
                    ]
                    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
1481
1482

                    prefill_tokens = Tokens(
OlivierDehaene's avatar
OlivierDehaene committed
1483
1484
1485
1486
                        prefill_token_ids,
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1487
1488
1489
1490
                    )
                else:
                    prefill_tokens = None

Nicolas Patry's avatar
Nicolas Patry committed
1491
                if top_n_tokens > 0:
Nicolas Patry's avatar
Nicolas Patry committed
1492
                    all_top_tokens = []
drbh's avatar
drbh committed
1493
                    for top_token_ids, top_token_logprobs in zip(
1494
1495
                        top_token_ids, top_token_logprobs
                    ):
Nicolas Patry's avatar
Nicolas Patry committed
1496
1497
1498
1499
1500
1501
                        toptoken_texts = self.tokenizer.batch_decode(
                            top_token_ids,
                            clean_up_tokenization_spaces=False,
                            skip_special_tokens=False,
                        )
                        special_toptokens = [
1502
1503
                            token_id in self.all_special_ids
                            for token_id in top_token_ids
Nicolas Patry's avatar
Nicolas Patry committed
1504
1505
1506
1507
1508
1509
1510
1511
1512
                        ]
                        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
1513
1514
1515
                else:
                    top_tokens = None

1516
1517
1518
                generation = Generation(
                    request.id,
                    prefill_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1519
1520
1521
1522
1523
1524
                    Tokens(
                        _next_token_ids,
                        _next_token_logprobs,
                        next_token_texts,
                        [nid in self.all_special_ids for nid in _next_token_ids],
                    ),
1525
                    generated_text,
Nicolas Patry's avatar
Nicolas Patry committed
1526
                    top_tokens,
1527
1528
                )

1529
                generations.append(generation)
1530

drbh's avatar
drbh committed
1531
1532
1533
            # 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
1534
1535
1536
                batch.next_token_chooser = (
                    batch.next_token_chooser.advance_grammar_single(i, next_token_id)
                )
drbh's avatar
drbh committed
1537

1538
            # Update values
1539
            batch.input_lengths[i] = input_length + n_accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1540
1541
            if batch.input_lengths[i] > batch.max_seqlen:
                batch.max_seqlen = batch.input_lengths[i]
1542
1543
            batch.prefix_offsets[i] = prefix_offset
            batch.read_offsets[i] = read_offset
1544
1545
            batch.all_input_ids[i] = all_input_ids

1546
1547
        if stopped:
            # No need to return a batch if we know that all requests stopped
1548
1549
1550
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1551

1552
1553
1554
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1555

1556
1557
1558
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)