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

7
8
import numpy as np

9
from loguru import logger
10
11
from dataclasses import dataclass
from opentelemetry import trace
12
13
14
15
16
17
from transformers import (
    PreTrainedTokenizerBase,
    AutoConfig,
    AutoTokenizer,
    GenerationConfig,
)
Daniël de Kok's avatar
Daniël de Kok committed
18
from typing import Iterable, Optional, Tuple, List, Type, Dict
fxmarty's avatar
fxmarty committed
19

drbh's avatar
drbh committed
20
from text_generation_server.adapters import AdapterBatchData, AdapterBatchMetadata
fxmarty's avatar
fxmarty committed
21
from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
Daniël de Kok's avatar
Daniël de Kok committed
22
from text_generation_server.utils.chunks import concat_text_chunks
Nicolas Patry's avatar
Nicolas Patry committed
23
from text_generation_server.utils.import_utils import SYSTEM
OlivierDehaene's avatar
OlivierDehaene committed
24
from text_generation_server.models import Model
25
from text_generation_server.utils.log import log_master
26
from text_generation_server.utils.tokens import batch_top_tokens
Nicolas Patry's avatar
Nicolas Patry committed
27
from text_generation_server.utils.speculate import get_speculate
28
29
30
31
32
from text_generation_server.utils import (
    initialize_torch_distributed,
    weight_files,
    Weights,
)
33
34
from text_generation_server.models.types import (
    Batch,
Nicolas Patry's avatar
Nicolas Patry committed
35
    Tokens,
36
37
38
39
    Generation,
    GeneratedText,
)
from text_generation_server.pb import generate_pb2
Nicolas Patry's avatar
Nicolas Patry committed
40
41
from text_generation_server.models.globals import (
    MEM_POOL,
42
43
    FLASH_DECODING,
    BLOCK_SIZE,
Nicolas Patry's avatar
Nicolas Patry committed
44
45
46
    CUDA_GRAPHS,
    get_adapter_to_index,
)
47
from text_generation_server.layers.attention import Seqlen
48
from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
49
from text_generation_server.utils.dist import MEMORY_FRACTION
50
from text_generation_server.utils.quantization import get_loader
drbh's avatar
drbh committed
51
from text_generation_server.utils.segments import SegmentConcatBuilder, find_segments
52

Nicolas Patry's avatar
Nicolas Patry committed
53
from text_generation_server.utils.import_utils import (
Nicolas Patry's avatar
Nicolas Patry committed
54
55
56
    empty_cache,
    synchronize,
    get_free_memory,
Nicolas Patry's avatar
Nicolas Patry committed
57
58
)

Nicolas Patry's avatar
Nicolas Patry committed
59
60
tracer = trace.get_tracer(__name__)

61
62
63
64
65
66
67
68
69
70
71
72
73
74

# 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

75

76
77
78
79
@dataclass
class FlashCausalLMBatch(Batch):
    batch_id: int
    requests: List[generate_pb2.Request]
80
81
    # request id -> idx in list mapping
    requests_idx_mapping: Dict[int, int]
82
83

    # Decoder values
84
85
    input_ids: torch.Tensor
    position_ids: torch.Tensor
86
    speculative_ids: Optional[torch.Tensor]
87

88
89
90
91
    # 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]
92
93
94
    # 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]
95
96
97
98
99
100
101
102
103
104

    # 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
105
    block_tables: List[List[int]]
106
    # tensor of size [b, max_total_seqlen // block_size] holding the paged attention block tables for all sequences
107
    block_tables_tensor: torch.Tensor
108
    # tensor of length \sum_{i=0}^{b} max_s_i  holding the paged attention slots for all sequences
109
    slots: torch.Tensor
110

111
112
    max_seqlen: int

113
114
115
116
117
    # 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]]

118
119
    # All tokens
    all_input_ids: List[List[int]]
120
    all_input_ids_tensor: torch.Tensor
121
122
123

    # Lengths of all generations present in the batch
    input_lengths: List[int]
124
    input_lengths_tensor: torch.Tensor
125
126
    prefix_offsets: List[Optional[int]]
    read_offsets: List[Optional[int]]
127
128

    # Generation helpers
129
    next_token_chooser: HeterogeneousNextTokenChooser
130
    stopping_criterias: List[StoppingCriteria]
Nicolas Patry's avatar
Nicolas Patry committed
131
132
    top_n_tokens: List[int]
    top_n_tokens_tensor: torch.Tensor
133

drbh's avatar
drbh committed
134
135
136
    # Adapter metadata for each request
    adapter_meta: AdapterBatchMetadata

137
    # Number of blocks in this batch
138
    num_blocks: int
139
140
    # Maximum number of blocks
    max_blocks: int
141

142
143
    def to_pb(self) -> generate_pb2.CachedBatch:
        return generate_pb2.CachedBatch(
144
            id=self.batch_id,
145
            request_ids=[r.id for r in self.requests],
146
            size=len(self),
147
            max_tokens=self.num_blocks * BLOCK_SIZE,
148
149
150
        )

    @classmethod
Daniël de Kok's avatar
Daniël de Kok committed
151
152
153
    def batch_tokenized_inputs(
        cls, requests: Iterable[generate_pb2.Request], tokenizer
    ):
154
155
        batch_inputs = []
        max_truncation = 0
156
        for r in requests:
Daniël de Kok's avatar
Daniël de Kok committed
157
            batch_inputs.append(concat_text_chunks(r.input_chunks.chunks))
158
159
160
161
162
            max_truncation = max(max_truncation, r.truncate)

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

drbh's avatar
drbh committed
165
166
167
168
169
170
171
172
173
    @classmethod
    def from_tokenized(
        cls,
        pb: generate_pb2.Batch,
        tokenizer: PreTrainedTokenizerBase,
        batch_tokenized_inputs,
        dtype: torch.dtype,
        device: torch.device,
    ) -> "FlashCausalLMBatch":
174
        sliding_window = get_sliding_windows()
175
        position_ids = []
176
        cu_seqlen_prefill = [0]
177
178
        start_slots = []
        slot_indices = []
179
        prefill_cache_indices = []
180
181

        input_lengths = []
182
183
        prefix_offsets = []
        read_offsets = []
184
        all_input_ids = []
185
        requests_idx_mapping = {}
186

187
188
189
190
191
192
        all_prefill_logprobs = True
        no_prefill_logprobs = True
        prefill_head_indices = []
        prefill_next_token_indices = []
        prefill_cu_outlens = [0]

193
        next_token_chooser_parameters = []
194
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
195
        top_n_tokens = []
196

drbh's avatar
drbh committed
197
198
199
        adapter_indices_list = []
        adapter_set = set()

200
201
        # Cumulative length
        cumulative_length = 0
202
        cumulative_max_length = 0
203
        prefill_out_cumulative_length = 0
204

205
        num_blocks = 0
206
        max_seqlen = 0
207
        max_length = 0
208
        max_blocks = 0
209

210
211
212
        block_tables = []
        slots = []

213
        # Parse batch
214
215
216
        for i, (r, tokenized_input) in enumerate(
            zip(pb.requests, batch_tokenized_inputs)
        ):
217
218
219
            # request id -> idx in list mapping
            requests_idx_mapping[r.id] = i

220
            tokenized_input = tokenized_input[-r.truncate :]
221
222
223
224
225
            if (
                tokenized_input[0] == tokenizer.bos_token_id
                and tokenized_input[1] == tokenizer.bos_token_id
            ):
                tokenized_input = tokenized_input[1:]
226

227
228
            input_length = len(tokenized_input)
            input_lengths.append(input_length)
229

230
            prefix_offsets.append(input_length - 5)
231
            read_offsets.append(input_length)
232

233
            all_input_ids.append(tokenized_input)
234
235

            # Position ids
236
237
            request_position_ids = torch.arange(0, input_length, dtype=torch.int32)
            position_ids.append(request_position_ids)
238
239

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

242
            next_token_chooser_parameters.append(r.parameters)
243

244
245
246
            stopping_criteria = StoppingCriteria.from_pb(
                r.stopping_parameters, tokenizer
            )
247
            max_new_tokens = stopping_criteria.max_new_tokens
248
            stopping_criterias.append(stopping_criteria)
Nicolas Patry's avatar
Nicolas Patry committed
249
            top_n_tokens.append(r.top_n_tokens)
250

Nicolas Patry's avatar
Nicolas Patry committed
251
252
            ADAPTER_TO_INDEX = get_adapter_to_index()
            adapter_index = ADAPTER_TO_INDEX.get(r.adapter_id, 0)
drbh's avatar
drbh committed
253
254
255
            adapter_indices_list.append(torch.full((input_length,), adapter_index))
            adapter_set.add(adapter_index)

256
257
            # Paged attention
            # Remove one as the first token des not have a past
Nicolas Patry's avatar
Nicolas Patry committed
258
            speculative_length = get_speculate()
drbh's avatar
drbh committed
259
            speculative_length = 0 if speculative_length is None else speculative_length
Nicolas Patry's avatar
Nicolas Patry committed
260
            total_tokens = input_length + max_new_tokens - 1 + speculative_length
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279

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

289
290
291
292
293
294
295
296
297
            # 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)

298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
            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

318
319
            # Update
            cumulative_length += input_length
320
321
            cumulative_max_length += total_tokens
            max_seqlen = max(max_seqlen, input_length)
322
            max_blocks = max(max_blocks, len(request_blocks))
OlivierDehaene's avatar
OlivierDehaene committed
323
324
325
            max_length = max(
                max_length, input_length + max_new_tokens + speculative_length
            )
326

drbh's avatar
drbh committed
327
328
329
330
        adapter_indices = torch.cat(adapter_indices_list).to(
            dtype=torch.int64, device=device
        )

331
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
drbh's avatar
drbh committed
332
            next_token_chooser_parameters, dtype, device, tokenizer
333
        )
334
        start_slots = torch.tensor(start_slots, dtype=torch.int64)
335
336
337
338
339
340
341

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

343
344
345
346
347
        # Create tensors on device
        all_input_ids_tensor = torch.tensor(
            all_input_ids_tensor, dtype=torch.int64, device=device
        )

348
349
350
        if len(pb.requests) > 1:
            input_ids = np.concatenate(all_input_ids, dtype=np.int64)
            position_ids = torch.cat(position_ids)
351
            slot_indices = torch.cat(slot_indices)
352
353
            if sliding_window is not None:
                prefill_cache_indices = torch.cat(prefill_cache_indices)
354
355
356
        else:
            input_ids = all_input_ids[0]
            position_ids = position_ids[0]
357
            slot_indices = slot_indices[0]
358
359
            if sliding_window is not None:
                prefill_cache_indices = prefill_cache_indices[0]
360

361
362
        cu_seqlen_prefill = torch.tensor(
            cu_seqlen_prefill, device=device, dtype=torch.int32
363
364
365
        )
        position_ids = position_ids.to(device)
        slot_indices = slot_indices.to(device)
366
367
368
        prefill_cache_indices = (
            prefill_cache_indices.to(device) if sliding_window is not None else None
        )
369
        input_ids = torch.tensor(input_ids, dtype=torch.int64, device=device)
370
371
        input_lengths_tensor = torch.tensor(
            input_lengths, dtype=torch.int32, device=device
372
        )
373

drbh's avatar
drbh committed
374
375
376
377
378
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

379
380
        if all_prefill_logprobs:
            prefill_head_indices = None
381
            prefill_next_token_indices = cu_seqlen_prefill[1:] - 1
382
        elif no_prefill_logprobs:
383
            prefill_head_indices = cu_seqlen_prefill[1:] - 1
384
385
386
387
388
389
390
391
            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
392
393
394
        top_n_tokens_tensor = torch.tensor(
            top_n_tokens, device=device, dtype=torch.int64
        )
395

396
397
398
399
400
401
402
403
        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)

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

442
443
444
445
446
447
448
449
450
451
452
    @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)

453
    @tracer.start_as_current_span("filter")
454
455
    def filter(self, request_ids: List[int]) -> "FlashCausalLMBatch":
        if len(request_ids) == 0:
456
457
            raise ValueError("Batch must have at least one request")
        # We assume that if len(requests) == len(self) then the requests are the same
458
        if len(request_ids) == len(self):
459
460
            return self

461
        device = self.input_ids.device
462

463
464
465
        # New values after filtering
        requests_idx_mapping = {}

466
467
468
        # Used to index into tensors
        indices = []

469
470
471
        # slots to keep after filtering
        slot_filtering_indices = torch.zeros(
            self.slots.shape[0], dtype=torch.bool, device=device
472
473
        )

474
        # Create on CPU to only move to GPU once instead of at every copy
475
        slot_indices = torch.empty(len(request_ids), dtype=torch.int64)
476
477
        max_seqlen = 0

478
        requests = []
479
480
        start_slots = []
        block_tables = []
481
482
        all_input_ids = []

483
        input_lengths = []
484
485
        prefix_offsets = []
        read_offsets = []
486

487
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
488
        top_n_tokens = []
drbh's avatar
drbh committed
489
        adapter_set = set()
490

491
        num_blocks = 0
492
493
494
495
        max_blocks = 0
        # Cumulative length
        cumulative_max_length = 0

496
497
        for i, request_id in enumerate(request_ids):
            idx = self.requests_idx_mapping[request_id]
498
            indices.append(idx)
499
500
501
            requests_idx_mapping[request_id] = i

            requests.append(self.requests[idx])
502
503
504
505

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

507
508
509
            all_input_ids.append(self.all_input_ids[idx])

            input_lengths.append(request_input_length)
510
511
            prefix_offsets.append(self.prefix_offsets[idx])
            read_offsets.append(self.read_offsets[idx])
512

513
514
            stopping_criteria = self.stopping_criterias[idx]
            stopping_criterias.append(stopping_criteria)
515

Nicolas Patry's avatar
Nicolas Patry committed
516
517
            top_n_tokens.append(self.top_n_tokens[idx])

Nicolas Patry's avatar
Nicolas Patry committed
518
519
            ADAPTER_TO_INDEX = get_adapter_to_index()
            adapter_index = ADAPTER_TO_INDEX.get(self.requests[idx].adapter_id, 0)
drbh's avatar
drbh committed
520
521
            adapter_set.add(adapter_index)

522
            remaining_tokens = (
523
524
                stopping_criteria.max_new_tokens - stopping_criteria.current_tokens
            )
525

526
            request_block_table = self.block_tables[idx]
527
            num_blocks += len(request_block_table)
528
529
530
            block_tables.append(request_block_table)
            start_slots.append(cumulative_max_length)

531
            # Copy to tensor (CPU)
532
            slot_indices[i] = cumulative_max_length + request_input_length - 1
533
534

            # Set slice
535
536
537
538
539
            slot_filtering_indices[
                self.start_slots[idx] : self.start_slots[idx]
                + request_input_length
                + remaining_tokens
                - 1
540
541
542
            ] = True

            cumulative_max_length += request_input_length + remaining_tokens - 1
543

544
545
            max_blocks = max(max_blocks, len(request_block_table))

546
547
548
        # Index into tensors
        input_ids = self.input_ids[indices]
        position_ids = self.position_ids[indices]
drbh's avatar
drbh committed
549
        adapter_indices = self.adapter_meta.adapter_indices[indices]
550
        all_input_ids_tensor = self.all_input_ids_tensor[indices]
551
552
553
        block_tables_tensor = self.block_tables_tensor[indices]
        input_lengths_tensor = self.input_lengths_tensor[indices]
        slots = self.slots[slot_filtering_indices]
554
        next_token_chooser = self.next_token_chooser.filter(indices)
Nicolas Patry's avatar
Nicolas Patry committed
555
        top_n_tokens_tensor = self.top_n_tokens_tensor[indices]
OlivierDehaene's avatar
OlivierDehaene committed
556
557
558
        speculative_ids = (
            self.speculative_ids[indices] if self.speculative_ids is not None else None
        )
559
560

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

562
        # Move to GPU now that we have the whole tensor
563
        slot_indices = slot_indices.to(device)
564

drbh's avatar
drbh committed
565
566
567
568
569
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

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

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

615
        num_blocks = 0
616
617
618
619
620
621
622
623
        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)
624
            num_blocks += b.num_blocks
OlivierDehaene's avatar
OlivierDehaene committed
625
626
627
            speculative_length = (
                b.speculative_ids.shape[1] if b.speculative_ids is not None else 0
            )
628
629
630
631
632
633
634
            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
635
                    + speculative_length
636
637
638
639
640
641
                    - stopping_criteria.current_tokens
                    for input_length, stopping_criteria in zip(
                        b.input_lengths, b.stopping_criterias
                    )
                ),
            )
642
643
644

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
645
646
647
648
649
650
651
652
653
654
        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)
655
        )
Nicolas Patry's avatar
Nicolas Patry committed
656
657
658
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
drbh's avatar
drbh committed
659
660
661
662
663
664
665
666
        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()
667

668
669
        start_slots = []
        block_tables = []
670
671
672
        all_input_ids = []

        input_lengths = []
673
674
        prefix_offsets = []
        read_offsets = []
675

676
        next_token_chooser_parameters = []
677
        fsm_grammar_states = []
678
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
679
        top_n_tokens = []
680

681
        # Cumulative length
682
        cumulative_batch_size = 0
683
        cumulative_slots = 0
drbh's avatar
drbh committed
684
        cumulative_adapter_indices_size = 0
685
686
687

        for i, batch in enumerate(batches):
            requests.extend(batch.requests)
688
689
690
691
692
693
694
695

            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

696
697
            start_index = cumulative_batch_size
            end_index = cumulative_batch_size + len(batch)
698
699
            slots_start_index = cumulative_slots
            slots_end_index = cumulative_slots + len(batch.slots)
700
701
702
703

            # Copy tensors (GPU)
            input_ids[start_index:end_index] = batch.input_ids
            position_ids[start_index:end_index] = batch.position_ids
704
705
            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
706
            top_n_tokens_tensor[start_index:end_index] = batch.top_n_tokens_tensor
707
            slots[slots_start_index:slots_end_index] = batch.slots
708

drbh's avatar
drbh committed
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
            # 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
            )

724
725
726
            all_input_ids_tensor[
                start_index:end_index, : batch.all_input_ids_tensor.shape[1]
            ] = batch.all_input_ids_tensor[:, :max_length]
727

728
729
730
            block_tables_tensor[
                start_index:end_index, : batch.block_tables_tensor.shape[1]
            ] = batch.block_tables_tensor[:, :max_blocks]
731

732
733
734
            start_slots.append(batch.start_slots + cumulative_slots)

            block_tables.extend(batch.block_tables)
735
736
            all_input_ids.extend(batch.all_input_ids)

737
            input_lengths.extend(batch.input_lengths)
738
739
            prefix_offsets.extend(batch.prefix_offsets)
            read_offsets.extend(batch.read_offsets)
740

741
            next_token_chooser_parameters.extend([r.parameters for r in batch.requests])
742
            fsm_grammar_states.extend(batch.next_token_chooser.fsm_grammar_states)
743
744
            stopping_criterias.extend(batch.stopping_criterias)

Nicolas Patry's avatar
Nicolas Patry committed
745
746
            top_n_tokens.extend(batch.top_n_tokens)

747
            # Update
748
            cumulative_batch_size += len(batch)
749
            cumulative_slots += len(batch.slots)
750

751
        start_slots = torch.concat(start_slots)
752

753
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
754
755
756
            next_token_chooser_parameters,
            dtype=batches[0].next_token_chooser.dtype,
            device=batches[0].next_token_chooser.device,
drbh's avatar
drbh committed
757
            tokenizer=batches[0].next_token_chooser.tokenizer,
758
            fsm_grammar_states=fsm_grammar_states,
759
760
        )

OlivierDehaene's avatar
OlivierDehaene committed
761
762
763
764
765
        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
766

drbh's avatar
drbh committed
767
768
        adapter_segments, adapter_segment_indices = adapter_segment_builder.build()

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

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


811
812
813
814
815
816
817
818
819
820
821
822
ADAPTER_LAYERS = [
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
]
ROW_PARALLEL = {"o_proj", "down_proj", "lm_head"}


823
824
825
class FlashCausalLM(Model):
    def __init__(
        self,
drbh's avatar
drbh committed
826
        model_id: str,
827
828
829
830
831
832
833
834
835
836
837
838
        model_class,
        revision: Optional[str] = None,
        quantize: Optional[str] = None,
        speculator: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
        trust_remote_code: bool = False,
        lora_adapter_ids: Optional[list] = [],
        tokenizer_class: PreTrainedTokenizerBase = AutoTokenizer,
        config_class: PreTrainedTokenizerBase = AutoConfig,
        default_dtype=torch.float16,
        aliases=None,
        # Used for Santacoder override of config
839
840
841
        num_kv_heads: Optional[int] = None,
        # Deepseek V2 uses different QK and V dims.
        head_size: Optional[int] = None,
842
        skip_special_tokens: bool = True,
843
    ):
844
845
846
847
848
849
850
851
852
853
854
855
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
        self.process_group, rank, world_size = initialize_torch_distributed()
        if torch.cuda.is_available():
            device = torch.device(f"cuda:{rank}")
            dtype = default_dtype if dtype is None else dtype
        elif SYSTEM == "ipex":
            if hasattr(torch, "xpu") and torch.xpu.is_available():
                device = torch.device(f"xpu:{rank}")
                dtype = default_dtype if dtype is None else dtype
            else:
                device = torch.device("cpu")
                # Float16 doesn't exist on target.
                dtype = torch.bfloat16 if dtype is None else dtype
        else:
            raise NotImplementedError(f"{model_class} is only available on GPU")

        tokenizer = tokenizer_class.from_pretrained(
            model_id,
            revision=revision,
            padding_side="left",
            truncation_side="left",
            trust_remote_code=trust_remote_code,
        )
        try:
            generation_config = GenerationConfig.from_pretrained(
                model_id, revision=revision, trust_remote_code=trust_remote_code
            )
            if isinstance(generation_config.eos_token_id, (list, set)):
                # TODO Huge hack
                tokenizer._eos_token_ids = set(generation_config.eos_token_id)
        except Exception:
            pass

        config = config_class.from_pretrained(
            model_id, revision=revision, trust_remote_code=trust_remote_code
        )
        config.quantize = quantize
        config.speculator = speculator

        torch.distributed.barrier(group=self.process_group)

884
        weights_loader = get_loader(quantize, model_id, revision)
885
886
        filenames = weight_files(model_id, revision=revision, extension=".safetensors")
        weights = Weights(
887
888
889
890
891
892
            filenames,
            device,
            dtype,
            process_group=self.process_group,
            aliases=aliases,
            weights_loader=weights_loader,
893
894
895
896
897
898
899
900
901
902
        )

        prefix = ""
        model = model_class(prefix, config, weights)
        torch.distributed.barrier(group=self.process_group)

        # VLM models define the config we care about in their text_config
        text_config = getattr(config, "text_config", None)
        if text_config is not None:
            config = text_config
903
904
905
906
907
908

        if getattr(config, "sliding_window", None) is not None:
            set_sliding_window(config.sliding_window)
        else:
            config.sliding_window = None

909
910
911
        self.num_layers = config.num_hidden_layers
        # Validation is done in the model itself
        if num_kv_heads is None:
912
913
            num_kv_heads = getattr(config, "num_key_value_heads", None)
            # GPT-2 workaround
914
            if num_kv_heads is None:
915
916
917
                num_kv_heads = getattr(config, "n_head", None)
        if num_kv_heads is None:
            raise ValueError("Cannot get the number of key/value heads")
918
919
920
921
922
923
        self.num_kv_heads = (
            num_kv_heads // self.process_group.size()
            if num_kv_heads > 1
            else num_kv_heads
        )
        assert self.num_kv_heads > 0
924
925

        if head_size is None:
Nicolas Patry's avatar
Nicolas Patry committed
926
927
928
929
930
931
            # Some models use GQA and different sizes for o_proj
            # and q_proj, that allows for that.
            if hasattr(config, "head_dim"):
                self.head_size = config.head_dim
            else:
                self.head_size = config.hidden_size // config.num_attention_heads
932
933
        else:
            self.head_size = head_size
934

935
        self.cuda_graphs = {}
936
        self.kv_cache = []
937

938
        super().__init__(
drbh's avatar
drbh committed
939
            model_id=model_id,
940
            model=model,
941
942
943
944
            tokenizer=tokenizer,
            requires_padding=False,
            dtype=dtype,
            device=device,
945
946
            rank=rank,
            world_size=world_size,
947
            sliding_window=config.sliding_window,
948
949
950
951
952
953
        )

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

954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
    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
970
971
972
973
        if SYSTEM == "ipex" and device.type == "xpu":
            x = 1
        else:
            x = BLOCK_SIZE // element_size
974

975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
        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
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
            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)
            ]
1023

1024
1025
1026
    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)
1027
        slots = torch.arange(bs, dtype=torch.int64, device=self.device)
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
        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,
1038
            "kv_cache": self.kv_cache,
1039
1040
1041
1042
            "block_tables": block_tables,
            "slots": slots,
            "input_lengths": input_lengths,
        }
1043
        input_lengths_ = Seqlen(input_lengths=input_lengths)
1044
1045
1046
1047
1048
1049
1050
1051
1052
        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,
1053
            kv_cache=self.kv_cache,
1054
1055
            block_tables=block_tables,
            slots=slots,
1056
            input_lengths=input_lengths_,
1057
            max_s=max_s,
1058
            prefill_cache_indices=None,
1059
1060
1061
1062
1063
            lm_head_indices=None,
        )
        torch.cuda.synchronize()

        with torch.cuda.graph(graph, pool=MEM_POOL):
1064
            input_lengths = Seqlen(input_lengths=input_lengths)
1065
            logits, speculative_logits = self.model.forward(
1066
1067
1068
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlen_prefill=None,
1069
                kv_cache=self.kv_cache,
1070
1071
1072
1073
                block_tables=block_tables,
                slots=slots,
                input_lengths=input_lengths,
                max_s=max_s,
1074
                prefill_cache_indices=None,
1075
1076
                lm_head_indices=None,
            )
1077
1078
            self.cuda_graphs[bs]["logits"] = logits
            self.cuda_graphs[bs]["speculative_logits"] = speculative_logits
1079
1080
        torch.cuda.synchronize()

1081
    def warmup(self, batch: FlashCausalLMBatch):
1082
        # The warmup batch is the biggest batch we could ever receive
Nicolas Patry's avatar
Nicolas Patry committed
1083
1084
        empty_cache()

1085
        try:
1086
1087
            self.init_kv_cache(
                batch.num_blocks,
1088
1089
1090
1091
1092
1093
                self.num_layers,
                self.num_kv_heads,
                self.head_size,
                self.dtype,
                self.device,
            )
1094
            max_bt = batch.max_blocks
1095
            max_s = max_bt * BLOCK_SIZE
fxmarty's avatar
fxmarty committed
1096
1097
1098

            if SYSTEM == "rocm" and os.environ.get("PYTORCH_TUNABLEOP_ENABLED", False):
                torch.cuda.tunable.tuning_enable(False)
1099
            _, batch, _ = self.generate_token(batch)
OlivierDehaene's avatar
OlivierDehaene committed
1100
        except torch.cuda.OutOfMemoryError as e:
1101
            raise RuntimeError(
1102
1103
                f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. "
                f"You need to decrease `--max-batch-prefill-tokens`"
1104
            ) from e
1105

Nicolas Patry's avatar
Nicolas Patry committed
1106
        synchronize(self.device)
1107

1108
1109
        # 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
1110
1111
1112
1113
        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
1114
        free_memory = get_free_memory(self.device, MEMORY_FRACTION)
drbh's avatar
drbh committed
1115
        batch_num_blocks = batch.num_blocks if batch is not None else 0
1116
1117

        num_blocks = (
1118
1119
            # Leave 5% for some wiggle room
            int((free_memory * 0.95) // total_cache_size)
1120
            # Add batch.num_blocks as we allocated it above, so it is included in the peak memory.
drbh's avatar
drbh committed
1121
            + batch_num_blocks
1122
1123
        )

1124
        del batch
1125

1126
        self.init_kv_cache(
1127
1128
1129
1130
1131
1132
1133
1134
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
            self.dtype,
            self.device,
        )

fxmarty's avatar
fxmarty committed
1135
1136
1137
1138
1139
        if SYSTEM == "rocm":
            if (
                os.environ.get("PYTORCH_TUNABLEOP_ENABLED") is None
                or os.environ.get("PYTORCH_TUNABLEOP_ENABLED") == "1"
            ):
1140
1141
                torch.cuda.tunable.enable()

fxmarty's avatar
fxmarty committed
1142
1143
1144
1145
1146
1147
1148
1149
                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(",")
                    ]
1150
                elif CUDA_GRAPHS is not None:
fxmarty's avatar
fxmarty committed
1151
                    tuning_sequences = CUDA_GRAPHS
1152
1153
1154
                else:
                    # For seqlen = 1, we dispatch to LLMM1 kernel.
                    tuning_sequences = [2, 3, 4, 5, 6, 7]
fxmarty's avatar
fxmarty committed
1155
1156
1157

                tunableop_filepath = os.path.join(
                    HUGGINGFACE_HUB_CACHE,
drbh's avatar
drbh committed
1158
                    f"tunableop_{self.model_id.replace('/', '-')}_tp{self.world_size}_rank{self.rank}.csv",
fxmarty's avatar
fxmarty committed
1159
1160
                )

1161
1162
1163
                log_master(
                    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`.",
fxmarty's avatar
fxmarty committed
1164
1165
1166
                )

                if os.path.isfile(tunableop_filepath):
1167
1168
1169
                    log_master(
                        logger.info,
                        f"The file {tunableop_filepath} already exists and will be reused.",
fxmarty's avatar
fxmarty committed
1170
1171
1172
1173
1174
1175
                    )
                    torch.cuda.tunable.read_file(tunableop_filepath)

                os.makedirs(HUGGINGFACE_HUB_CACHE, exist_ok=True)

                for seqlen in tuning_sequences:
1176
                    log_master(logger.info, f"Warming up TunableOp for seqlen={seqlen}")
fxmarty's avatar
fxmarty committed
1177
1178
1179
1180
                    self.tunableop_warmup(seqlen)
                    torch.cuda.tunable.write_file(tunableop_filepath)
                torch.cuda.tunable.tuning_enable(False)
            else:
1181
1182
1183
                log_master(
                    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.",
fxmarty's avatar
fxmarty committed
1184
1185
                )

1186
        if CUDA_GRAPHS:
1187
            try:
1188
1189
1190
                log_master(
                    logger.info, f"Cuda Graphs are enabled for sizes {CUDA_GRAPHS}"
                )
1191
                # Warmup cuda graphs
1192
                for bs in CUDA_GRAPHS:
1193
1194
                    if self.speculate is None or self.speculate + 1 <= bs:
                        self.cuda_graph_warmup(bs, max_s, max_bt)
OlivierDehaene's avatar
OlivierDehaene committed
1195
            except torch.cuda.OutOfMemoryError:
1196
                logger.exception("Decode cuda graph warmup failed")
1197
        else:
1198
1199
1200
            log_master(
                logger.info, f"Cuda Graphs are disabled (CUDA_GRAPHS={CUDA_GRAPHS})."
            )
1201

1202
        return int(num_blocks * BLOCK_SIZE)
1203

fxmarty's avatar
fxmarty committed
1204
1205
1206
1207
1208
    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
1209
1210
        # Dummy value, some models (starcoder2) don't accept `None`.
        input_lengths = torch.ones(seqlen, dtype=torch.int32, device=self.device)
1211
        input_lengths = Seqlen(input_lengths=input_lengths)
fxmarty's avatar
fxmarty committed
1212

fxmarty's avatar
fxmarty committed
1213
1214
1215
1216
1217
1218
1219
        # 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
            ),
1220
            kv_cache=self.kv_cache,
fxmarty's avatar
fxmarty committed
1221
            block_tables=None,
fxmarty's avatar
fxmarty committed
1222
            input_lengths=input_lengths,
fxmarty's avatar
fxmarty committed
1223
1224
1225
            slots=slots,
            max_s=seqlen,
            lm_head_indices=None,
1226
            prefill_cache_indices=None,
fxmarty's avatar
fxmarty committed
1227
1228
        )

1229
    def forward(
drbh's avatar
drbh committed
1230
        self, batch: FlashCausalLMBatch, adapter_data: AdapterBatchData
1231
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1232
        # Model Forward
Nicolas Patry's avatar
Nicolas Patry committed
1233
        if batch.speculative_ids is not None:
OlivierDehaene's avatar
OlivierDehaene committed
1234
1235
1236
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1237
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1238
1239
1240
1241
1242
            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
1243
1244
1245

            speculative_ids = batch.speculative_ids

OlivierDehaene's avatar
OlivierDehaene committed
1246
            B, speculative_length = speculative_ids.shape
Nicolas Patry's avatar
Nicolas Patry committed
1247
            new_length = speculative_length + 1
OlivierDehaene's avatar
OlivierDehaene committed
1248
1249
1250
            new_input_ids = torch.cat(
                [input_ids.unsqueeze(-1), speculative_ids], dim=1
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1251
1252
            arange = torch.arange(new_length, device=position_ids.device).unsqueeze(0)
            arange_int = arange.to(dtype=torch.int32)
OlivierDehaene's avatar
OlivierDehaene committed
1253
1254
1255
            new_position_ids = (
                position_ids.unsqueeze(-1).expand(B, new_length) + arange
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1256
            slots = (slots.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1)
OlivierDehaene's avatar
OlivierDehaene committed
1257
1258
1259
            input_lengths = (
                input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1260
1261

            # Add Copy the block tables for all members
OlivierDehaene's avatar
OlivierDehaene committed
1262
1263
1264
1265
1266
1267
            block_tables = (
                block_tables.unsqueeze(1)
                .expand(B, new_length, -1)
                .reshape(B * new_length, -1)
                .contiguous()
            )
Nicolas Patry's avatar
Nicolas Patry committed
1268
1269
1270
1271
1272
            max_s = max_s + speculative_length

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
1273
1274
1275
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1276
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1277
1278
1279
1280
1281
            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
1282

1283
1284
1285
1286
1287
1288
        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)

1289
        bs = input_ids.shape[0]
OlivierDehaene's avatar
OlivierDehaene committed
1290
1291
1292
1293
1294
1295
1296
1297
        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:
1298
            input_lengths = Seqlen(input_lengths=input_lengths)
1299
            logits, speculative_logits = self.model.forward(
1300
1301
1302
1303
1304
1305
1306
1307
                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,
1308
                prefill_cache_indices=batch.prefill_cache_indices,
1309
                lm_head_indices=lm_head_indices,
drbh's avatar
drbh committed
1310
                adapter_data=adapter_data,
1311
            )
1312
1313
1314
            if batch.prefill_cache_indices is not None:
                batch.prefill_cache_indices = None
            return logits, speculative_logits
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330

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

    @tracer.start_as_current_span("generate_token")
    def generate_token(
        self, batch: FlashCausalLMBatch
1342
1343
    ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]:
        start = time.time_ns()
1344
        prefill = batch.cu_seqlen_prefill is not None
1345
        prefill_logprobs = batch.prefill_next_token_indices is not None
1346

drbh's avatar
drbh committed
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
        # 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)
1375

1376
1377
        if prefill:
            next_token_logits = (
1378
                out[batch.prefill_next_token_indices] if prefill_logprobs else out
1379
            )
Nicolas Patry's avatar
Nicolas Patry committed
1380
1381
            if speculative_logits is not None:
                speculative_logits = (
OlivierDehaene's avatar
OlivierDehaene committed
1382
1383
1384
                    speculative_logits[batch.prefill_next_token_indices]
                    if prefill_logprobs
                    else speculative_logits
Nicolas Patry's avatar
Nicolas Patry committed
1385
                )
drbh's avatar
drbh committed
1386
1387
1388
1389
            next_adapter_indices = batch.adapter_meta.adapter_indices.new_empty(
                len(batch)
            )

1390
1391
        else:
            next_token_logits = out
drbh's avatar
drbh committed
1392
            next_adapter_indices = batch.adapter_meta.adapter_indices
1393

Nicolas Patry's avatar
Nicolas Patry committed
1394
        speculate = get_speculate()
OlivierDehaene's avatar
OlivierDehaene committed
1395
1396
1397
1398
1399
1400
1401
1402
1403
        (
            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
1404
            speculate,
OlivierDehaene's avatar
OlivierDehaene committed
1405
1406
            batch.speculative_ids,
            speculative_logits,
1407
1408
        )

Nicolas Patry's avatar
Nicolas Patry committed
1409
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
1410
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1411
1412
        )

1413
        if prefill:
1414
            if len(batch) > 1 and prefill_logprobs:
1415
1416
                # 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
1417
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
1418
1419

            next_position_ids = batch.position_ids.new_empty(len(batch))
1420
1421
1422
            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
1423
1424
1425
1426
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

1427
1428
1429
1430
1431
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
1432
        stopped = True
1433
1434

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

1437
1438
1439
1440
        # 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

1441
        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1442
        index = 0
OlivierDehaene's avatar
OlivierDehaene committed
1443
        for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator):
1444
            # Indexing metadata
1445
1446
1447
            start_index = cumulative_length
            end_index = cumulative_length + input_length

1448
            if prefill:
1449
1450
1451
1452
1453
                # 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

1454
1455
1456
1457
                # 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
1458
1459
1460
1461
1462
1463
                # 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
                ]

1464
1465
                # Used to gather prefill logprobs
                # Copy batch.input_ids to prefill_token_indices
1466
1467
                if prefill_logprobs:
                    if len(batch) > 1:
drbh's avatar
drbh committed
1468
1469
1470
                        prefill_tokens_indices[out_start_index : out_end_index - 1] = (
                            batch.input_ids[start_index + 1 : start_index + out_length]
                        )
1471
1472
1473
1474
1475
                    else:
                        # Set prefill_tokens_indices to the correct slice
                        prefill_tokens_indices = batch.input_ids[
                            start_index + 1 : start_index + out_length
                        ]
1476

Nicolas Patry's avatar
Nicolas Patry committed
1477
1478
1479
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1480
1481
1482

            cumulative_length += input_length

drbh's avatar
drbh committed
1483
        # Update values
Nicolas Patry's avatar
Nicolas Patry committed
1484
1485
1486
1487
1488
        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
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
        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,
            )
1499

1500
        if prefill and prefill_logprobs:
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
            # 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
1511
        next_token_ids = next_input_ids.tolist()
1512
1513
        accepted_ids = accepted_ids.tolist()
        start_decode = time.time_ns()
1514
1515
1516
1517
1518

        # Zipped iterator
        iterator = zip(
            batch.requests,
            batch.input_lengths,
1519
1520
            batch.prefix_offsets,
            batch.read_offsets,
1521
1522
            batch.stopping_criterias,
            batch.all_input_ids,
1523
1524
            batch.next_token_chooser.do_sample,
            batch.next_token_chooser.seeds,
Nicolas Patry's avatar
Nicolas Patry committed
1525
            batch.top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1526
            accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1527
1528
            batch_top_token_ids,
            batch_top_token_logprobs,
1529
1530
1531
        )

        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1532
        index = 0
1533
1534
1535
        for i, (
            request,
            input_length,
1536
1537
            prefix_offset,
            read_offset,
1538
1539
            stopping_criteria,
            all_input_ids,
1540
1541
            do_sample,
            seed,
Nicolas Patry's avatar
Nicolas Patry committed
1542
            top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1543
            n_accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1544
1545
            top_token_ids,
            top_token_logprobs,
1546
        ) in enumerate(iterator):
1547
            # Append next token to all tokens
Nicolas Patry's avatar
Nicolas Patry committed
1548
1549
1550
            next_token_texts = []
            left = 0

1551
            if n_accepted_ids > 1:
1552
                log_master(logger.debug, f"Speculated ids {n_accepted_ids - 1}")
1553

Nicolas Patry's avatar
Nicolas Patry committed
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
            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)
1565

Nicolas Patry's avatar
Nicolas Patry committed
1566
1567
1568
1569
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1570

Nicolas Patry's avatar
Nicolas Patry committed
1571
1572
1573
1574
1575
1576
1577
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1578

OlivierDehaene's avatar
OlivierDehaene committed
1579
1580
1581
1582
            _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
1583
            index += n_accepted_ids
1584

1585
1586
1587
1588
1589
            # Shard generations
            # All generations will be appended in the rust sharded client
            if i % self.world_size == self.rank:
                if stop:
                    # Decode generated tokens
1590
1591
                    output_text, _, _ = self.decode_token(
                        all_input_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1592
1593
1594
1595
1596
1597
                        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,
1598
1599
                    )
                    generated_text = GeneratedText(
1600
1601
1602
1603
                        output_text,
                        stopping_criteria.current_tokens,
                        reason,
                        seed if do_sample else None,
1604
1605
1606
1607
1608
                    )
                else:
                    generated_text = None

                # Prefill
1609
1610
1611
1612
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

1613
1614
                    # Remove generated token to only have prefill and add nan for first prompt token
                    request_prefill_logprobs = [float("nan")] + prefill_logprobs[
1615
                        out_start_index : out_end_index - 1
1616
1617
1618
1619
1620
1621
1622
                    ]
                    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
1623
1624

                    prefill_tokens = Tokens(
OlivierDehaene's avatar
OlivierDehaene committed
1625
1626
1627
1628
                        prefill_token_ids,
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1629
1630
1631
1632
                    )
                else:
                    prefill_tokens = None

Nicolas Patry's avatar
Nicolas Patry committed
1633
                if top_n_tokens > 0:
Nicolas Patry's avatar
Nicolas Patry committed
1634
                    all_top_tokens = []
drbh's avatar
drbh committed
1635
                    for top_token_ids, top_token_logprobs in zip(
1636
1637
                        top_token_ids, top_token_logprobs
                    ):
Nicolas Patry's avatar
Nicolas Patry committed
1638
1639
1640
1641
1642
1643
                        toptoken_texts = self.tokenizer.batch_decode(
                            top_token_ids,
                            clean_up_tokenization_spaces=False,
                            skip_special_tokens=False,
                        )
                        special_toptokens = [
1644
1645
                            token_id in self.all_special_ids
                            for token_id in top_token_ids
Nicolas Patry's avatar
Nicolas Patry committed
1646
1647
1648
1649
1650
1651
1652
1653
1654
                        ]
                        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
1655
1656
1657
                else:
                    top_tokens = None

1658
1659
1660
                generation = Generation(
                    request.id,
                    prefill_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1661
1662
1663
1664
1665
1666
                    Tokens(
                        _next_token_ids,
                        _next_token_logprobs,
                        next_token_texts,
                        [nid in self.all_special_ids for nid in _next_token_ids],
                    ),
1667
                    generated_text,
Nicolas Patry's avatar
Nicolas Patry committed
1668
                    top_tokens,
1669
1670
                )

1671
                generations.append(generation)
1672

drbh's avatar
drbh committed
1673
1674
1675
            # 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
1676
1677
1678
                batch.next_token_chooser = (
                    batch.next_token_chooser.advance_grammar_single(i, next_token_id)
                )
drbh's avatar
drbh committed
1679

1680
            # Update values
1681
            batch.input_lengths[i] = input_length + n_accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1682
1683
            if batch.input_lengths[i] > batch.max_seqlen:
                batch.max_seqlen = batch.input_lengths[i]
1684
1685
            batch.prefix_offsets[i] = prefix_offset
            batch.read_offsets[i] = read_offset
1686
1687
            batch.all_input_ids[i] = all_input_ids

1688
1689
        if stopped:
            # No need to return a batch if we know that all requests stopped
1690
1691
1692
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1693

1694
1695
1696
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1697

1698
1699
1700
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)