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

8
9
import numpy as np

10
from loguru import logger
11
12
from dataclasses import dataclass
from opentelemetry import trace
13
from transformers import PreTrainedTokenizerBase
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
31
from text_generation_server.models.globals import MEM_POOL, CUDA_GRAPHS
fxmarty's avatar
fxmarty committed
32
import text_generation_server.models.globals as tgi_globals
33
from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
34
from text_generation_server.utils.dist import MEMORY_FRACTION
drbh's avatar
drbh committed
35
from text_generation_server.utils.segments import SegmentConcatBuilder, find_segments
36

Nicolas Patry's avatar
Nicolas Patry committed
37
from text_generation_server.utils.import_utils import (
Nicolas Patry's avatar
Nicolas Patry committed
38
39
40
    empty_cache,
    synchronize,
    get_free_memory,
Nicolas Patry's avatar
Nicolas Patry committed
41
42
)

Nicolas Patry's avatar
Nicolas Patry committed
43
44
tracer = trace.get_tracer(__name__)

45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
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

60

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

    # Decoder values
69
70
    input_ids: torch.Tensor
    position_ids: torch.Tensor
71
    speculative_ids: Optional[torch.Tensor]
72

73
74
75
76
    # 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]
77
78
79
    # 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]
80
81
82
83
84
85
86
87
88
89

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

96
97
    max_seqlen: int

98
99
100
101
102
    # 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]]

103
104
    # All tokens
    all_input_ids: List[List[int]]
105
    all_input_ids_tensor: torch.Tensor
106
107
108

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

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

drbh's avatar
drbh committed
119
120
121
    # Adapter metadata for each request
    adapter_meta: AdapterBatchMetadata

122
    # Number of blocks in this batch
123
    num_blocks: int
124
125
    # Maximum number of blocks
    max_blocks: int
126

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

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

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

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

        input_lengths = []
167
168
        prefix_offsets = []
        read_offsets = []
169
        all_input_ids = []
170
        requests_idx_mapping = {}
171

172
173
174
175
176
177
        all_prefill_logprobs = True
        no_prefill_logprobs = True
        prefill_head_indices = []
        prefill_next_token_indices = []
        prefill_cu_outlens = [0]

178
        next_token_chooser_parameters = []
179
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
180
        top_n_tokens = []
181

drbh's avatar
drbh committed
182
183
184
        adapter_indices_list = []
        adapter_set = set()

185
186
        # Cumulative length
        cumulative_length = 0
187
        cumulative_max_length = 0
188
        prefill_out_cumulative_length = 0
189

190
        num_blocks = 0
191
        max_seqlen = 0
192
        max_length = 0
193
        max_blocks = 0
194

195
196
197
        block_tables = []
        slots = []

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

205
            tokenized_input = tokenized_input[-r.truncate :]
206
207
208
209
210
            if (
                tokenized_input[0] == tokenizer.bos_token_id
                and tokenized_input[1] == tokenizer.bos_token_id
            ):
                tokenized_input = tokenized_input[1:]
211

212
213
            input_length = len(tokenized_input)
            input_lengths.append(input_length)
214

215
            prefix_offsets.append(input_length - 5)
216
            read_offsets.append(input_length)
217

218
            all_input_ids.append(tokenized_input)
219
220

            # Position ids
221
222
            request_position_ids = torch.arange(0, input_length, dtype=torch.int32)
            position_ids.append(request_position_ids)
223
224

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

227
            next_token_chooser_parameters.append(r.parameters)
228

229
230
231
            stopping_criteria = StoppingCriteria.from_pb(
                r.stopping_parameters, tokenizer
            )
232
            max_new_tokens = stopping_criteria.max_new_tokens
233
            stopping_criterias.append(stopping_criteria)
Nicolas Patry's avatar
Nicolas Patry committed
234
            top_n_tokens.append(r.top_n_tokens)
235

drbh's avatar
drbh committed
236
237
238
239
            adapter_index = tgi_globals.ADAPTER_TO_INDEX.get(r.adapter_id, 0)
            adapter_indices_list.append(torch.full((input_length,), adapter_index))
            adapter_set.add(adapter_index)

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

            # 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)
264
265
266
267
268
269
270
271
272
            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)

273
274
275
276
277
278
279
280
281
            # 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)

282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
            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

302
303
            # Update
            cumulative_length += input_length
304
305
            cumulative_max_length += total_tokens
            max_seqlen = max(max_seqlen, input_length)
306
            max_blocks = max(max_blocks, len(request_blocks))
OlivierDehaene's avatar
OlivierDehaene committed
307
308
309
            max_length = max(
                max_length, input_length + max_new_tokens + speculative_length
            )
310

drbh's avatar
drbh committed
311
312
313
314
        adapter_indices = torch.cat(adapter_indices_list).to(
            dtype=torch.int64, device=device
        )

315
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
drbh's avatar
drbh committed
316
            next_token_chooser_parameters, dtype, device, tokenizer
317
        )
318
        start_slots = torch.tensor(start_slots, dtype=torch.int64)
319
320
321
322
323
324
325

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

327
328
329
330
331
        # Create tensors on device
        all_input_ids_tensor = torch.tensor(
            all_input_ids_tensor, dtype=torch.int64, device=device
        )

332
333
334
        if len(pb.requests) > 1:
            input_ids = np.concatenate(all_input_ids, dtype=np.int64)
            position_ids = torch.cat(position_ids)
335
            slot_indices = torch.cat(slot_indices)
336
337
            if sliding_window is not None:
                prefill_cache_indices = torch.cat(prefill_cache_indices)
338
339
340
        else:
            input_ids = all_input_ids[0]
            position_ids = position_ids[0]
341
            slot_indices = slot_indices[0]
342
343
            if sliding_window is not None:
                prefill_cache_indices = prefill_cache_indices[0]
344

345
346
        cu_seqlen_prefill = torch.tensor(
            cu_seqlen_prefill, device=device, dtype=torch.int32
347
348
349
        )
        position_ids = position_ids.to(device)
        slot_indices = slot_indices.to(device)
350
351
352
        prefill_cache_indices = (
            prefill_cache_indices.to(device) if sliding_window is not None else None
        )
353
        input_ids = torch.tensor(input_ids, dtype=torch.int64, device=device)
354
355
        input_lengths_tensor = torch.tensor(
            input_lengths, dtype=torch.int32, device=device
356
        )
357

drbh's avatar
drbh committed
358
359
360
361
362
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

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

380
381
382
383
384
385
386
387
        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)

388
389
390
        return cls(
            batch_id=pb.id,
            requests=pb.requests,
391
            requests_idx_mapping=requests_idx_mapping,
392
393
            input_ids=input_ids,
            position_ids=position_ids,
394
            cu_seqlen_prefill=cu_seqlen_prefill,
395
            prefill_cache_indices=prefill_cache_indices,
396
397
            start_slots=start_slots,
            slot_indices=slot_indices,
398
399
400
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
401
            max_seqlen=max_seqlen,
402
403
404
            prefill_head_indices=prefill_head_indices,
            prefill_next_token_indices=prefill_next_token_indices,
            prefill_cu_outlens=prefill_cu_outlens,
405
            input_lengths=input_lengths,
406
            input_lengths_tensor=input_lengths_tensor,
407
408
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
409
            all_input_ids=all_input_ids,
410
411
            all_input_ids_tensor=all_input_ids_tensor,
            next_token_chooser=next_token_chooser,
412
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
413
414
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
415
            num_blocks=num_blocks,
416
            max_blocks=max_blocks,
drbh's avatar
drbh committed
417
418
419
420
421
422
            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
423
            speculative_ids=None,
424
425
        )

426
427
428
429
430
431
432
433
434
435
436
    @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)

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

445
        device = self.input_ids.device
446

447
448
449
        # New values after filtering
        requests_idx_mapping = {}

450
451
452
        # Used to index into tensors
        indices = []

453
454
455
        # slots to keep after filtering
        slot_filtering_indices = torch.zeros(
            self.slots.shape[0], dtype=torch.bool, device=device
456
457
        )

458
        # Create on CPU to only move to GPU once instead of at every copy
459
        slot_indices = torch.empty(len(request_ids), dtype=torch.int64)
460
461
        max_seqlen = 0

462
        requests = []
463
464
        start_slots = []
        block_tables = []
465
466
        all_input_ids = []

467
        input_lengths = []
468
469
        prefix_offsets = []
        read_offsets = []
470

471
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
472
        top_n_tokens = []
drbh's avatar
drbh committed
473
        adapter_set = set()
474

475
        num_blocks = 0
476
477
478
479
        max_blocks = 0
        # Cumulative length
        cumulative_max_length = 0

480
481
        for i, request_id in enumerate(request_ids):
            idx = self.requests_idx_mapping[request_id]
482
            indices.append(idx)
483
484
485
            requests_idx_mapping[request_id] = i

            requests.append(self.requests[idx])
486
487
488
489

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

491
492
493
            all_input_ids.append(self.all_input_ids[idx])

            input_lengths.append(request_input_length)
494
495
            prefix_offsets.append(self.prefix_offsets[idx])
            read_offsets.append(self.read_offsets[idx])
496

497
498
            stopping_criteria = self.stopping_criterias[idx]
            stopping_criterias.append(stopping_criteria)
499

Nicolas Patry's avatar
Nicolas Patry committed
500
501
            top_n_tokens.append(self.top_n_tokens[idx])

drbh's avatar
drbh committed
502
503
504
505
506
            adapter_index = tgi_globals.ADAPTER_TO_INDEX.get(
                self.requests[idx].adapter_id, 0
            )
            adapter_set.add(adapter_index)

507
            remaining_tokens = (
508
509
                stopping_criteria.max_new_tokens - stopping_criteria.current_tokens
            )
510

511
            request_block_table = self.block_tables[idx]
512
            num_blocks += len(request_block_table)
513
514
515
            block_tables.append(request_block_table)
            start_slots.append(cumulative_max_length)

516
            # Copy to tensor (CPU)
517
            slot_indices[i] = cumulative_max_length + request_input_length - 1
518
519

            # Set slice
520
521
522
523
524
            slot_filtering_indices[
                self.start_slots[idx] : self.start_slots[idx]
                + request_input_length
                + remaining_tokens
                - 1
525
526
527
            ] = True

            cumulative_max_length += request_input_length + remaining_tokens - 1
528

529
530
            max_blocks = max(max_blocks, len(request_block_table))

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

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

547
        # Move to GPU now that we have the whole tensor
548
        slot_indices = slot_indices.to(device)
549

drbh's avatar
drbh committed
550
551
552
553
554
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

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

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

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

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
630
631
632
633
634
635
636
637
638
639
        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)
640
        )
Nicolas Patry's avatar
Nicolas Patry committed
641
642
643
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
drbh's avatar
drbh committed
644
645
646
647
648
649
650
651
        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()
652

653
654
        start_slots = []
        block_tables = []
655
656
657
        all_input_ids = []

        input_lengths = []
658
659
        prefix_offsets = []
        read_offsets = []
660

661
        next_token_chooser_parameters = []
662
        fsm_grammar_states = []
663
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
664
        top_n_tokens = []
665

666
        # Cumulative length
667
        cumulative_batch_size = 0
668
        cumulative_slots = 0
drbh's avatar
drbh committed
669
        cumulative_adapter_indices_size = 0
670
671
672

        for i, batch in enumerate(batches):
            requests.extend(batch.requests)
673
674
675
676
677
678
679
680

            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

681
682
            start_index = cumulative_batch_size
            end_index = cumulative_batch_size + len(batch)
683
684
            slots_start_index = cumulative_slots
            slots_end_index = cumulative_slots + len(batch.slots)
685
686
687
688

            # Copy tensors (GPU)
            input_ids[start_index:end_index] = batch.input_ids
            position_ids[start_index:end_index] = batch.position_ids
689
690
            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
691
            top_n_tokens_tensor[start_index:end_index] = batch.top_n_tokens_tensor
692
            slots[slots_start_index:slots_end_index] = batch.slots
693

drbh's avatar
drbh committed
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
            # 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
            )

709
710
711
            all_input_ids_tensor[
                start_index:end_index, : batch.all_input_ids_tensor.shape[1]
            ] = batch.all_input_ids_tensor[:, :max_length]
712

713
714
715
            block_tables_tensor[
                start_index:end_index, : batch.block_tables_tensor.shape[1]
            ] = batch.block_tables_tensor[:, :max_blocks]
716

717
718
719
            start_slots.append(batch.start_slots + cumulative_slots)

            block_tables.extend(batch.block_tables)
720
721
            all_input_ids.extend(batch.all_input_ids)

722
            input_lengths.extend(batch.input_lengths)
723
724
            prefix_offsets.extend(batch.prefix_offsets)
            read_offsets.extend(batch.read_offsets)
725

726
            next_token_chooser_parameters.extend([r.parameters for r in batch.requests])
727
            fsm_grammar_states.extend(batch.next_token_chooser.fsm_grammar_states)
728
729
            stopping_criterias.extend(batch.stopping_criterias)

Nicolas Patry's avatar
Nicolas Patry committed
730
731
            top_n_tokens.extend(batch.top_n_tokens)

732
            # Update
733
            cumulative_batch_size += len(batch)
734
            cumulative_slots += len(batch.slots)
735

736
        start_slots = torch.concat(start_slots)
737

738
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
739
740
741
            next_token_chooser_parameters,
            dtype=batches[0].next_token_chooser.dtype,
            device=batches[0].next_token_chooser.device,
drbh's avatar
drbh committed
742
            tokenizer=batches[0].next_token_chooser.tokenizer,
743
            fsm_grammar_states=fsm_grammar_states,
744
745
        )

OlivierDehaene's avatar
OlivierDehaene committed
746
747
748
749
750
        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
751

drbh's avatar
drbh committed
752
753
        adapter_segments, adapter_segment_indices = adapter_segment_builder.build()

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

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


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

815
        self.cuda_graphs = {}
816
        self.kv_cache = []
817

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

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

834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
    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
850
851
852
853
        if SYSTEM == "ipex" and device.type == "xpu":
            x = 1
        else:
            x = BLOCK_SIZE // element_size
854

Nicolas Patry's avatar
Nicolas Patry committed
855
        if SYSTEM == "ipex" and device == torch.device("cpu"):
Wang, Yi's avatar
Wang, Yi committed
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
            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)
            ]
887

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

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

943
    def warmup(self, batch: FlashCausalLMBatch):
944
        # The warmup batch is the biggest batch we could ever receive
Nicolas Patry's avatar
Nicolas Patry committed
945
946
        empty_cache()

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

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

Nicolas Patry's avatar
Nicolas Patry committed
968
        synchronize(self.device)
969

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

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

986
        del batch
987

988
        self.init_kv_cache(
989
990
991
992
993
994
995
996
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
            self.dtype,
            self.device,
        )

fxmarty's avatar
fxmarty committed
997
998
999
1000
1001
        if SYSTEM == "rocm":
            if (
                os.environ.get("PYTORCH_TUNABLEOP_ENABLED") is None
                or os.environ.get("PYTORCH_TUNABLEOP_ENABLED") == "1"
            ):
1002
1003
                torch.cuda.tunable.enable()

fxmarty's avatar
fxmarty committed
1004
1005
1006
1007
1008
1009
1010
1011
                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(",")
                    ]
1012
                elif CUDA_GRAPHS is not None:
fxmarty's avatar
fxmarty committed
1013
                    tuning_sequences = CUDA_GRAPHS
1014
1015
1016
                else:
                    # For seqlen = 1, we dispatch to LLMM1 kernel.
                    tuning_sequences = [2, 3, 4, 5, 6, 7]
fxmarty's avatar
fxmarty committed
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044

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

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

1057
        return int(num_blocks * BLOCK_SIZE)
1058

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

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

            speculative_ids = batch.speculative_ids

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

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

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
1127
1128
1129
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1130
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1131
1132
1133
1134
1135
            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
1136

1137
1138
1139
1140
1141
1142
        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)

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

        # 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
1184
1185
1186
1187
1188
1189
1190
        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
1191
1192
1193
1194

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

drbh's avatar
drbh committed
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
        # 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)
1228

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

1243
1244
        else:
            next_token_logits = out
drbh's avatar
drbh committed
1245
            next_adapter_indices = batch.adapter_meta.adapter_indices
1246

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

Nicolas Patry's avatar
Nicolas Patry committed
1262
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
1263
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1264
1265
        )

1266
        if prefill:
1267
            if len(batch) > 1 and prefill_logprobs:
1268
1269
                # 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
1270
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
1271
1272

            next_position_ids = batch.position_ids.new_empty(len(batch))
1273
1274
1275
            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
1276
1277
1278
1279
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

1280
1281
1282
1283
1284
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
1285
        stopped = True
1286
1287

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

1290
1291
1292
1293
        # 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

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

1301
            if prefill:
1302
1303
1304
1305
1306
                # 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

1307
1308
1309
1310
                # 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
1311
1312
1313
1314
1315
1316
                # 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
                ]

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

Nicolas Patry's avatar
Nicolas Patry committed
1330
1331
1332
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1333
1334
1335

            cumulative_length += input_length

drbh's avatar
drbh committed
1336
        # Update values
Nicolas Patry's avatar
Nicolas Patry committed
1337
1338
1339
1340
1341
        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
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
        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,
            )
1352

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

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

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

1404
1405
1406
1407
            if n_accepted_ids > 1:
                if RANK == 0:
                    logger.debug(f"Speculated ids {n_accepted_ids - 1}")

Nicolas Patry's avatar
Nicolas Patry committed
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
            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)
1419

Nicolas Patry's avatar
Nicolas Patry committed
1420
1421
1422
1423
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1424

Nicolas Patry's avatar
Nicolas Patry committed
1425
1426
1427
1428
1429
1430
1431
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1432

OlivierDehaene's avatar
OlivierDehaene committed
1433
1434
1435
1436
            _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
1437
            index += n_accepted_ids
1438

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

                # Prefill
1463
1464
1465
1466
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

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

                    prefill_tokens = Tokens(
OlivierDehaene's avatar
OlivierDehaene committed
1479
1480
1481
1482
                        prefill_token_ids,
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1483
1484
1485
1486
                    )
                else:
                    prefill_tokens = None

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

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

1525
                generations.append(generation)
1526

drbh's avatar
drbh committed
1527
1528
1529
            # 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
1530
1531
1532
                batch.next_token_chooser = (
                    batch.next_token_chooser.advance_grammar_single(i, next_token_id)
                )
drbh's avatar
drbh committed
1533

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

1542
1543
        if stopped:
            # No need to return a batch if we know that all requests stopped
1544
1545
1546
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1547

1548
1549
1550
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1551

1552
1553
1554
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)