flash_causal_lm.py 59 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
from text_generation_server.models.globals import (
    MEM_POOL,
33
34
    FLASH_DECODING,
    BLOCK_SIZE,
Nicolas Patry's avatar
Nicolas Patry committed
35
36
37
38
    CUDA_GRAPHS,
    get_adapter_to_index,
    MODEL_ID,
)
39
from text_generation_server.layers.attention import Seqlen
40
from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
41
from text_generation_server.utils.dist import MEMORY_FRACTION
drbh's avatar
drbh committed
42
from text_generation_server.utils.segments import SegmentConcatBuilder, find_segments
43

Nicolas Patry's avatar
Nicolas Patry committed
44
from text_generation_server.utils.import_utils import (
Nicolas Patry's avatar
Nicolas Patry committed
45
46
47
    empty_cache,
    synchronize,
    get_free_memory,
Nicolas Patry's avatar
Nicolas Patry committed
48
49
)

Nicolas Patry's avatar
Nicolas Patry committed
50
51
tracer = trace.get_tracer(__name__)

52
53
54
55
56
57
58
59
60
61
62
63
64
65

# 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

66

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

    # Decoder values
75
76
    input_ids: torch.Tensor
    position_ids: torch.Tensor
77
    speculative_ids: Optional[torch.Tensor]
78

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

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

102
103
    max_seqlen: int

104
105
106
107
108
    # 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]]

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

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

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

drbh's avatar
drbh committed
125
126
127
    # Adapter metadata for each request
    adapter_meta: AdapterBatchMetadata

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

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

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

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

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

        input_lengths = []
173
174
        prefix_offsets = []
        read_offsets = []
175
        all_input_ids = []
176
        requests_idx_mapping = {}
177

178
179
180
181
182
183
        all_prefill_logprobs = True
        no_prefill_logprobs = True
        prefill_head_indices = []
        prefill_next_token_indices = []
        prefill_cu_outlens = [0]

184
        next_token_chooser_parameters = []
185
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
186
        top_n_tokens = []
187

drbh's avatar
drbh committed
188
189
190
        adapter_indices_list = []
        adapter_set = set()

191
192
        # Cumulative length
        cumulative_length = 0
193
        cumulative_max_length = 0
194
        prefill_out_cumulative_length = 0
195

196
        num_blocks = 0
197
        max_seqlen = 0
198
        max_length = 0
199
        max_blocks = 0
200

201
202
203
        block_tables = []
        slots = []

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

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

218
219
            input_length = len(tokenized_input)
            input_lengths.append(input_length)
220

221
            prefix_offsets.append(input_length - 5)
222
            read_offsets.append(input_length)
223

224
            all_input_ids.append(tokenized_input)
225
226

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

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

233
            next_token_chooser_parameters.append(r.parameters)
234

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

387
388
389
390
391
392
393
394
        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)

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

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

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

452
        device = self.input_ids.device
453

454
455
456
        # New values after filtering
        requests_idx_mapping = {}

457
458
459
        # Used to index into tensors
        indices = []

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

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

469
        requests = []
470
471
        start_slots = []
        block_tables = []
472
473
        all_input_ids = []

474
        input_lengths = []
475
476
        prefix_offsets = []
        read_offsets = []
477

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

482
        num_blocks = 0
483
484
485
486
        max_blocks = 0
        # Cumulative length
        cumulative_max_length = 0

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

            requests.append(self.requests[idx])
493
494
495
496

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

498
499
500
            all_input_ids.append(self.all_input_ids[idx])

            input_lengths.append(request_input_length)
501
502
            prefix_offsets.append(self.prefix_offsets[idx])
            read_offsets.append(self.read_offsets[idx])
503

504
505
            stopping_criteria = self.stopping_criterias[idx]
            stopping_criterias.append(stopping_criteria)
506

Nicolas Patry's avatar
Nicolas Patry committed
507
508
            top_n_tokens.append(self.top_n_tokens[idx])

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

513
            remaining_tokens = (
514
515
                stopping_criteria.max_new_tokens - stopping_criteria.current_tokens
            )
516

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

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

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

            cumulative_max_length += request_input_length + remaining_tokens - 1
534

535
536
            max_blocks = max(max_blocks, len(request_block_table))

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

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

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

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

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

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

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

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
636
637
638
639
640
641
642
643
644
645
        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)
646
        )
Nicolas Patry's avatar
Nicolas Patry committed
647
648
649
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
drbh's avatar
drbh committed
650
651
652
653
654
655
656
657
        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()
658

659
660
        start_slots = []
        block_tables = []
661
662
663
        all_input_ids = []

        input_lengths = []
664
665
        prefix_offsets = []
        read_offsets = []
666

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

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

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

            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

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

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

drbh's avatar
drbh committed
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
            # 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
            )

715
716
717
            all_input_ids_tensor[
                start_index:end_index, : batch.all_input_ids_tensor.shape[1]
            ] = batch.all_input_ids_tensor[:, :max_length]
718

719
720
721
            block_tables_tensor[
                start_index:end_index, : batch.block_tables_tensor.shape[1]
            ] = batch.block_tables_tensor[:, :max_blocks]
722

723
724
725
            start_slots.append(batch.start_slots + cumulative_slots)

            block_tables.extend(batch.block_tables)
726
727
            all_input_ids.extend(batch.all_input_ids)

728
            input_lengths.extend(batch.input_lengths)
729
730
            prefix_offsets.extend(batch.prefix_offsets)
            read_offsets.extend(batch.read_offsets)
731

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

Nicolas Patry's avatar
Nicolas Patry committed
736
737
            top_n_tokens.extend(batch.top_n_tokens)

738
            # Update
739
            cumulative_batch_size += len(batch)
740
            cumulative_slots += len(batch.slots)
741

742
        start_slots = torch.concat(start_slots)
743

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

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

drbh's avatar
drbh committed
758
759
        adapter_segments, adapter_segment_indices = adapter_segment_builder.build()

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

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


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

821
        self.cuda_graphs = {}
822
        self.kv_cache = []
823

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

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

840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
    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
856
857
858
859
        if SYSTEM == "ipex" and device.type == "xpu":
            x = 1
        else:
            x = BLOCK_SIZE // element_size
860

861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
        if FLASH_DECODING:
            self.kv_cache = [
                (
                    torch.empty(
                        (num_blocks, BLOCK_SIZE, num_heads, head_size),
                        dtype=dtype,
                        device=device,
                    ),
                    torch.empty(
                        (num_blocks, BLOCK_SIZE, num_heads, head_size),
                        dtype=dtype,
                        device=device,
                    ),
                )
                for _ in range(num_layers)
            ]
        elif SYSTEM == "ipex" and device == torch.device("cpu"):
Wang, Yi's avatar
Wang, Yi committed
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
            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)
            ]
909

910
911
912
    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)
913
        slots = torch.arange(bs, dtype=torch.int64, device=self.device)
914
915
916
917
918
919
920
921
922
923
        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,
924
            "kv_cache": self.kv_cache,
925
926
927
928
            "block_tables": block_tables,
            "slots": slots,
            "input_lengths": input_lengths,
        }
929
        input_lengths_ = Seqlen(input_lengths=input_lengths)
930
931
932
933
934
935
936
937
938
        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,
939
            kv_cache=self.kv_cache,
940
941
            block_tables=block_tables,
            slots=slots,
942
            input_lengths=input_lengths_,
943
            max_s=max_s,
944
            prefill_cache_indices=None,
945
946
947
948
949
            lm_head_indices=None,
        )
        torch.cuda.synchronize()

        with torch.cuda.graph(graph, pool=MEM_POOL):
950
            input_lengths = Seqlen(input_lengths=input_lengths)
951
            logits, speculative_logits = self.model.forward(
952
953
954
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlen_prefill=None,
955
                kv_cache=self.kv_cache,
956
957
958
959
                block_tables=block_tables,
                slots=slots,
                input_lengths=input_lengths,
                max_s=max_s,
960
                prefill_cache_indices=None,
961
962
                lm_head_indices=None,
            )
963
964
            self.cuda_graphs[bs]["logits"] = logits
            self.cuda_graphs[bs]["speculative_logits"] = speculative_logits
965
966
        torch.cuda.synchronize()

967
    def warmup(self, batch: FlashCausalLMBatch):
968
        # The warmup batch is the biggest batch we could ever receive
Nicolas Patry's avatar
Nicolas Patry committed
969
970
        empty_cache()

971
        try:
972
973
            self.init_kv_cache(
                batch.num_blocks,
974
975
976
977
978
979
                self.num_layers,
                self.num_kv_heads,
                self.head_size,
                self.dtype,
                self.device,
            )
980
            max_bt = batch.max_blocks
981
            max_s = max_bt * BLOCK_SIZE
fxmarty's avatar
fxmarty committed
982
983
984

            if SYSTEM == "rocm" and os.environ.get("PYTORCH_TUNABLEOP_ENABLED", False):
                torch.cuda.tunable.tuning_enable(False)
985
            _, batch, _ = self.generate_token(batch)
OlivierDehaene's avatar
OlivierDehaene committed
986
        except torch.cuda.OutOfMemoryError as e:
987
            raise RuntimeError(
988
989
                f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. "
                f"You need to decrease `--max-batch-prefill-tokens`"
990
            ) from e
991

Nicolas Patry's avatar
Nicolas Patry committed
992
        synchronize(self.device)
993

994
995
        # 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
996
997
998
999
        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
1000
        free_memory = get_free_memory(self.device, MEMORY_FRACTION)
drbh's avatar
drbh committed
1001
        batch_num_blocks = batch.num_blocks if batch is not None else 0
1002
1003

        num_blocks = (
1004
1005
            # Leave 5% for some wiggle room
            int((free_memory * 0.95) // total_cache_size)
1006
            # Add batch.num_blocks as we allocated it above, so it is included in the peak memory.
drbh's avatar
drbh committed
1007
            + batch_num_blocks
1008
1009
        )

1010
        del batch
1011

1012
        self.init_kv_cache(
1013
1014
1015
1016
1017
1018
1019
1020
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
            self.dtype,
            self.device,
        )

fxmarty's avatar
fxmarty committed
1021
1022
1023
1024
1025
        if SYSTEM == "rocm":
            if (
                os.environ.get("PYTORCH_TUNABLEOP_ENABLED") is None
                or os.environ.get("PYTORCH_TUNABLEOP_ENABLED") == "1"
            ):
1026
1027
                torch.cuda.tunable.enable()

fxmarty's avatar
fxmarty committed
1028
1029
1030
1031
1032
1033
1034
1035
                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(",")
                    ]
1036
                elif CUDA_GRAPHS is not None:
fxmarty's avatar
fxmarty committed
1037
                    tuning_sequences = CUDA_GRAPHS
1038
1039
1040
                else:
                    # For seqlen = 1, we dispatch to LLMM1 kernel.
                    tuning_sequences = [2, 3, 4, 5, 6, 7]
fxmarty's avatar
fxmarty committed
1041
1042
1043

                tunableop_filepath = os.path.join(
                    HUGGINGFACE_HUB_CACHE,
Nicolas Patry's avatar
Nicolas Patry committed
1044
                    f"tunableop_{MODEL_ID.replace('/', '-')}_tp{self.world_size}_rank{self.rank}.csv",
fxmarty's avatar
fxmarty committed
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
                )

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

1069
        if CUDA_GRAPHS:
1070
            try:
1071
                logger.info(f"Cuda Graphs are enabled for sizes {CUDA_GRAPHS}")
1072
                # Warmup cuda graphs
1073
                for bs in CUDA_GRAPHS:
1074
1075
                    if self.speculate is None or self.speculate + 1 <= bs:
                        self.cuda_graph_warmup(bs, max_s, max_bt)
OlivierDehaene's avatar
OlivierDehaene committed
1076
            except torch.cuda.OutOfMemoryError:
1077
                logger.exception(f"Decode cuda graph warmup failed")
1078
1079
        else:
            logger.info(f"Cuda Graphs are disabled (CUDA_GRAPHS={CUDA_GRAPHS}).")
1080

1081
        return int(num_blocks * BLOCK_SIZE)
1082

fxmarty's avatar
fxmarty committed
1083
1084
1085
1086
1087
    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
1088
1089
        # Dummy value, some models (starcoder2) don't accept `None`.
        input_lengths = torch.ones(seqlen, dtype=torch.int32, device=self.device)
1090
        input_lengths = Seqlen(input_lengths=input_lengths)
fxmarty's avatar
fxmarty committed
1091

fxmarty's avatar
fxmarty committed
1092
1093
1094
1095
1096
1097
1098
        # 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
            ),
1099
            kv_cache=self.kv_cache,
fxmarty's avatar
fxmarty committed
1100
            block_tables=None,
fxmarty's avatar
fxmarty committed
1101
            input_lengths=input_lengths,
fxmarty's avatar
fxmarty committed
1102
1103
1104
            slots=slots,
            max_s=seqlen,
            lm_head_indices=None,
1105
            prefill_cache_indices=None,
fxmarty's avatar
fxmarty committed
1106
1107
        )

1108
    def forward(
drbh's avatar
drbh committed
1109
        self, batch: FlashCausalLMBatch, adapter_data: AdapterBatchData
1110
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1111
        # Model Forward
Nicolas Patry's avatar
Nicolas Patry committed
1112
        if batch.speculative_ids is not None:
OlivierDehaene's avatar
OlivierDehaene committed
1113
1114
1115
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1116
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1117
1118
1119
1120
1121
            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
1122
1123
1124

            speculative_ids = batch.speculative_ids

OlivierDehaene's avatar
OlivierDehaene committed
1125
            B, speculative_length = speculative_ids.shape
Nicolas Patry's avatar
Nicolas Patry committed
1126
            new_length = speculative_length + 1
OlivierDehaene's avatar
OlivierDehaene committed
1127
1128
1129
            new_input_ids = torch.cat(
                [input_ids.unsqueeze(-1), speculative_ids], dim=1
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1130
1131
            arange = torch.arange(new_length, device=position_ids.device).unsqueeze(0)
            arange_int = arange.to(dtype=torch.int32)
OlivierDehaene's avatar
OlivierDehaene committed
1132
1133
1134
            new_position_ids = (
                position_ids.unsqueeze(-1).expand(B, new_length) + arange
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1135
            slots = (slots.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1)
OlivierDehaene's avatar
OlivierDehaene committed
1136
1137
1138
            input_lengths = (
                input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1139
1140

            # Add Copy the block tables for all members
OlivierDehaene's avatar
OlivierDehaene committed
1141
1142
1143
1144
1145
1146
            block_tables = (
                block_tables.unsqueeze(1)
                .expand(B, new_length, -1)
                .reshape(B * new_length, -1)
                .contiguous()
            )
Nicolas Patry's avatar
Nicolas Patry committed
1147
1148
1149
1150
1151
            max_s = max_s + speculative_length

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
1152
1153
1154
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1155
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1156
1157
1158
1159
1160
            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
1161

1162
1163
1164
1165
1166
1167
        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)

1168
        bs = input_ids.shape[0]
OlivierDehaene's avatar
OlivierDehaene committed
1169
1170
1171
1172
1173
1174
1175
1176
        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:
1177
            input_lengths = Seqlen(input_lengths=input_lengths)
1178
            logits, speculative_logits = self.model.forward(
1179
1180
1181
1182
1183
1184
1185
1186
                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,
1187
                prefill_cache_indices=batch.prefill_cache_indices,
1188
                lm_head_indices=lm_head_indices,
drbh's avatar
drbh committed
1189
                adapter_data=adapter_data,
1190
            )
1191
1192
1193
            if batch.prefill_cache_indices is not None:
                batch.prefill_cache_indices = None
            return logits, speculative_logits
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209

        # 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
1210
1211
1212
1213
1214
1215
1216
        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
1217
1218
1219
1220

    @tracer.start_as_current_span("generate_token")
    def generate_token(
        self, batch: FlashCausalLMBatch
1221
1222
    ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]:
        start = time.time_ns()
1223
        prefill = batch.cu_seqlen_prefill is not None
1224
        prefill_logprobs = batch.prefill_next_token_indices is not None
1225

drbh's avatar
drbh committed
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
        # 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)
1254

1255
1256
        if prefill:
            next_token_logits = (
1257
                out[batch.prefill_next_token_indices] if prefill_logprobs else out
1258
            )
Nicolas Patry's avatar
Nicolas Patry committed
1259
1260
            if speculative_logits is not None:
                speculative_logits = (
OlivierDehaene's avatar
OlivierDehaene committed
1261
1262
1263
                    speculative_logits[batch.prefill_next_token_indices]
                    if prefill_logprobs
                    else speculative_logits
Nicolas Patry's avatar
Nicolas Patry committed
1264
                )
drbh's avatar
drbh committed
1265
1266
1267
1268
            next_adapter_indices = batch.adapter_meta.adapter_indices.new_empty(
                len(batch)
            )

1269
1270
        else:
            next_token_logits = out
drbh's avatar
drbh committed
1271
            next_adapter_indices = batch.adapter_meta.adapter_indices
1272

Nicolas Patry's avatar
Nicolas Patry committed
1273
        speculate = get_speculate()
OlivierDehaene's avatar
OlivierDehaene committed
1274
1275
1276
1277
1278
1279
1280
1281
1282
        (
            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
1283
            speculate,
OlivierDehaene's avatar
OlivierDehaene committed
1284
1285
            batch.speculative_ids,
            speculative_logits,
1286
1287
        )

Nicolas Patry's avatar
Nicolas Patry committed
1288
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
1289
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1290
1291
        )

1292
        if prefill:
1293
            if len(batch) > 1 and prefill_logprobs:
1294
1295
                # 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
1296
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
1297
1298

            next_position_ids = batch.position_ids.new_empty(len(batch))
1299
1300
1301
            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
1302
1303
1304
1305
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

1306
1307
1308
1309
1310
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
1311
        stopped = True
1312
1313

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

1316
1317
1318
1319
        # 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

1320
        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1321
        index = 0
OlivierDehaene's avatar
OlivierDehaene committed
1322
        for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator):
1323
            # Indexing metadata
1324
1325
1326
            start_index = cumulative_length
            end_index = cumulative_length + input_length

1327
            if prefill:
1328
1329
1330
1331
1332
                # 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

1333
1334
1335
1336
                # 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
1337
1338
1339
1340
1341
1342
                # 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
                ]

1343
1344
                # Used to gather prefill logprobs
                # Copy batch.input_ids to prefill_token_indices
1345
1346
                if prefill_logprobs:
                    if len(batch) > 1:
drbh's avatar
drbh committed
1347
1348
1349
                        prefill_tokens_indices[out_start_index : out_end_index - 1] = (
                            batch.input_ids[start_index + 1 : start_index + out_length]
                        )
1350
1351
1352
1353
1354
                    else:
                        # Set prefill_tokens_indices to the correct slice
                        prefill_tokens_indices = batch.input_ids[
                            start_index + 1 : start_index + out_length
                        ]
1355

Nicolas Patry's avatar
Nicolas Patry committed
1356
1357
1358
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1359
1360
1361

            cumulative_length += input_length

drbh's avatar
drbh committed
1362
        # Update values
Nicolas Patry's avatar
Nicolas Patry committed
1363
1364
1365
1366
1367
        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
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
        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,
            )
1378

1379
        if prefill and prefill_logprobs:
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
            # 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
1390
        next_token_ids = next_input_ids.tolist()
1391
1392
        accepted_ids = accepted_ids.tolist()
        start_decode = time.time_ns()
1393
1394
1395
1396
1397

        # Zipped iterator
        iterator = zip(
            batch.requests,
            batch.input_lengths,
1398
1399
            batch.prefix_offsets,
            batch.read_offsets,
1400
1401
            batch.stopping_criterias,
            batch.all_input_ids,
1402
1403
            batch.next_token_chooser.do_sample,
            batch.next_token_chooser.seeds,
Nicolas Patry's avatar
Nicolas Patry committed
1404
            batch.top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1405
            accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1406
1407
            batch_top_token_ids,
            batch_top_token_logprobs,
1408
1409
1410
        )

        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1411
        index = 0
1412
1413
1414
        for i, (
            request,
            input_length,
1415
1416
            prefix_offset,
            read_offset,
1417
1418
            stopping_criteria,
            all_input_ids,
1419
1420
            do_sample,
            seed,
Nicolas Patry's avatar
Nicolas Patry committed
1421
            top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1422
            n_accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1423
1424
            top_token_ids,
            top_token_logprobs,
1425
        ) in enumerate(iterator):
1426
            # Append next token to all tokens
Nicolas Patry's avatar
Nicolas Patry committed
1427
1428
1429
            next_token_texts = []
            left = 0

1430
1431
1432
1433
            if n_accepted_ids > 1:
                if RANK == 0:
                    logger.debug(f"Speculated ids {n_accepted_ids - 1}")

Nicolas Patry's avatar
Nicolas Patry committed
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
            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)
1445

Nicolas Patry's avatar
Nicolas Patry committed
1446
1447
1448
1449
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1450

Nicolas Patry's avatar
Nicolas Patry committed
1451
1452
1453
1454
1455
1456
1457
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1458

OlivierDehaene's avatar
OlivierDehaene committed
1459
1460
1461
1462
            _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
1463
            index += n_accepted_ids
1464

1465
1466
1467
1468
1469
            # Shard generations
            # All generations will be appended in the rust sharded client
            if i % self.world_size == self.rank:
                if stop:
                    # Decode generated tokens
1470
1471
                    output_text, _, _ = self.decode_token(
                        all_input_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1472
1473
1474
1475
1476
1477
                        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,
1478
1479
                    )
                    generated_text = GeneratedText(
1480
1481
1482
1483
                        output_text,
                        stopping_criteria.current_tokens,
                        reason,
                        seed if do_sample else None,
1484
1485
1486
1487
1488
                    )
                else:
                    generated_text = None

                # Prefill
1489
1490
1491
1492
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

1493
1494
                    # Remove generated token to only have prefill and add nan for first prompt token
                    request_prefill_logprobs = [float("nan")] + prefill_logprobs[
1495
                        out_start_index : out_end_index - 1
1496
1497
1498
1499
1500
1501
1502
                    ]
                    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
1503
1504

                    prefill_tokens = Tokens(
OlivierDehaene's avatar
OlivierDehaene committed
1505
1506
1507
1508
                        prefill_token_ids,
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1509
1510
1511
1512
                    )
                else:
                    prefill_tokens = None

Nicolas Patry's avatar
Nicolas Patry committed
1513
                if top_n_tokens > 0:
Nicolas Patry's avatar
Nicolas Patry committed
1514
                    all_top_tokens = []
drbh's avatar
drbh committed
1515
                    for top_token_ids, top_token_logprobs in zip(
1516
1517
                        top_token_ids, top_token_logprobs
                    ):
Nicolas Patry's avatar
Nicolas Patry committed
1518
1519
1520
1521
1522
1523
                        toptoken_texts = self.tokenizer.batch_decode(
                            top_token_ids,
                            clean_up_tokenization_spaces=False,
                            skip_special_tokens=False,
                        )
                        special_toptokens = [
1524
1525
                            token_id in self.all_special_ids
                            for token_id in top_token_ids
Nicolas Patry's avatar
Nicolas Patry committed
1526
1527
1528
1529
1530
1531
1532
1533
1534
                        ]
                        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
1535
1536
1537
                else:
                    top_tokens = None

1538
1539
1540
                generation = Generation(
                    request.id,
                    prefill_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1541
1542
1543
1544
1545
1546
                    Tokens(
                        _next_token_ids,
                        _next_token_logprobs,
                        next_token_texts,
                        [nid in self.all_special_ids for nid in _next_token_ids],
                    ),
1547
                    generated_text,
Nicolas Patry's avatar
Nicolas Patry committed
1548
                    top_tokens,
1549
1550
                )

1551
                generations.append(generation)
1552

drbh's avatar
drbh committed
1553
1554
1555
            # 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
1556
1557
1558
                batch.next_token_chooser = (
                    batch.next_token_chooser.advance_grammar_single(i, next_token_id)
                )
drbh's avatar
drbh committed
1559

1560
            # Update values
1561
            batch.input_lengths[i] = input_length + n_accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1562
1563
            if batch.input_lengths[i] > batch.max_seqlen:
                batch.max_seqlen = batch.input_lengths[i]
1564
1565
            batch.prefix_offsets[i] = prefix_offset
            batch.read_offsets[i] = read_offset
1566
1567
            batch.all_input_ids[i] = all_input_ids

1568
1569
        if stopped:
            # No need to return a batch if we know that all requests stopped
1570
1571
1572
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1573

1574
1575
1576
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1577

1578
1579
1580
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)