flash_causal_lm.py 73.5 KB
Newer Older
1
from contextlib import nullcontext
2
import math
3
import os
4
import time
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
14
15
16
17
18
from transformers import (
    PreTrainedTokenizerBase,
    AutoConfig,
    AutoTokenizer,
    GenerationConfig,
)
19
from typing import Any, ContextManager, Iterable, Optional, Tuple, List, Type, Dict
fxmarty's avatar
fxmarty committed
20

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

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

Nicolas Patry's avatar
Nicolas Patry committed
61
62
tracer = trace.get_tracer(__name__)

63
64
65
66
67
68
69
70
71
72
73
74
75
76

# 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

77

78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
def init_cpu_threads_env(rank_id: int, world_size: int):
    import importlib.util

    if importlib.util.find_spec("numa") is not None:
        import numa
        import psutil

        nodes = numa.get_max_node() + 1
        rank_per_node = math.ceil(world_size / nodes)
        num_cpus_per_nodes = int(psutil.cpu_count(logical=False) / nodes)
        node_id = int(rank_id / rank_per_node)
        rank_offset_per_node = rank_id % rank_per_node
        if os.getenv("OMP_NUM_THREADS") is None:
            num_cpus_per_rank = max(int(num_cpus_per_nodes / rank_per_node), 1)
        else:
            num_cpus_per_rank = int(os.getenv("OMP_NUM_THREADS"))
        if len(numa.get_membind()) == nodes:
            numa.set_membind([node_id])
        torch.set_num_threads(num_cpus_per_rank)
        if len(numa.get_affinity(0)) == psutil.cpu_count(logical=True):
            cpu_start = num_cpus_per_rank * rank_offset_per_node
            numa.set_affinity(
                0,
                list(numa.node_to_cpus(node_id))[
                    cpu_start : cpu_start + num_cpus_per_rank
                ],
            )
        logger.info(f"affinity={numa.get_affinity(0)}, membind = {numa.get_membind()}")


108
109
110
111
@dataclass
class FlashCausalLMBatch(Batch):
    batch_id: int
    requests: List[generate_pb2.Request]
112
113
    # request id -> idx in list mapping
    requests_idx_mapping: Dict[int, int]
114
115

    # Decoder values
116
117
    input_ids: torch.Tensor
    position_ids: torch.Tensor
118
    speculative_ids: Optional[torch.Tensor]
119

120
121
122
123
    # 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]
124
125
126
    # 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]
127
128
129
130
131
132
133
134
135
136

    # 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
137
    block_tables: List[List[int]]
138
    # tensor of size [b, max_total_seqlen // block_size] holding the paged attention block tables for all sequences
139
    block_tables_tensor: torch.Tensor
140
    # tensor of length \sum_{i=0}^{b} max_s_i  holding the paged attention slots for all sequences
141
    slots: torch.Tensor
Nicolas Patry's avatar
Nicolas Patry committed
142
143
144
    # size [b], containing the number of blocks that can be retrieved from the cache
    prefix_lens: List[int]
    prefix_lens_tensor: torch.Tensor
145

146
147
    max_seqlen: int

148
149
150
151
152
    # 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]]

Nicolas Patry's avatar
Nicolas Patry committed
153
154
155
    # Prefixes
    prefix_ids: List[List[int]]

156
157
    # All tokens
    all_input_ids: List[List[int]]
158
    all_input_ids_tensor: torch.Tensor
159
160
161

    # Lengths of all generations present in the batch
    input_lengths: List[int]
162
    input_lengths_tensor: torch.Tensor
163
164
    prefix_offsets: List[Optional[int]]
    read_offsets: List[Optional[int]]
165
166

    # Generation helpers
167
    next_token_chooser: HeterogeneousNextTokenChooser
168
    stopping_criterias: List[StoppingCriteria]
Nicolas Patry's avatar
Nicolas Patry committed
169
170
    top_n_tokens: List[int]
    top_n_tokens_tensor: torch.Tensor
171

drbh's avatar
drbh committed
172
173
174
    # Adapter metadata for each request
    adapter_meta: AdapterBatchMetadata

175
    # Number of blocks in this batch
176
    num_blocks: int
177
178
    # Maximum number of blocks
    max_blocks: int
179

180
181
    def to_pb(self) -> generate_pb2.CachedBatch:
        return generate_pb2.CachedBatch(
182
            id=self.batch_id,
183
            request_ids=[r.id for r in self.requests],
184
            size=len(self),
185
            max_tokens=self.num_blocks * BLOCK_SIZE,
186
187
188
        )

    @classmethod
Daniël de Kok's avatar
Daniël de Kok committed
189
190
191
    def batch_tokenized_inputs(
        cls, requests: Iterable[generate_pb2.Request], tokenizer
    ):
192
193
        batch_inputs = []
        max_truncation = 0
194
        for r in requests:
Daniël de Kok's avatar
Daniël de Kok committed
195
            batch_inputs.append(concat_text_chunks(r.input_chunks.chunks))
196
197
198
199
200
            max_truncation = max(max_truncation, r.truncate)

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

drbh's avatar
drbh committed
203
204
205
206
207
208
209
210
211
    @classmethod
    def from_tokenized(
        cls,
        pb: generate_pb2.Batch,
        tokenizer: PreTrainedTokenizerBase,
        batch_tokenized_inputs,
        dtype: torch.dtype,
        device: torch.device,
    ) -> "FlashCausalLMBatch":
212
        sliding_window = get_sliding_windows()
213
        position_ids = []
214
        cu_seqlen_prefill = [0]
215
216
        start_slots = []
        slot_indices = []
217
        prefill_cache_indices = []
218
219

        input_lengths = []
220
221
        prefix_offsets = []
        read_offsets = []
222
        all_input_ids = []
Nicolas Patry's avatar
Nicolas Patry committed
223
        prefix_ids = []
224
        requests_idx_mapping = {}
225

226
227
228
229
230
231
        all_prefill_logprobs = True
        no_prefill_logprobs = True
        prefill_head_indices = []
        prefill_next_token_indices = []
        prefill_cu_outlens = [0]

232
        next_token_chooser_parameters = []
233
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
234
        top_n_tokens = []
235

drbh's avatar
drbh committed
236
237
238
        adapter_indices_list = []
        adapter_set = set()

239
240
        # Cumulative length
        cumulative_length = 0
Nicolas Patry's avatar
Nicolas Patry committed
241
        cumulative_slot_tokens = 0
242
        prefill_out_cumulative_length = 0
243

244
        num_blocks = 0
245
        max_seqlen = 0
246
        max_length = 0
247
        max_blocks = 0
248

249
250
        block_tables = []
        slots = []
Nicolas Patry's avatar
Nicolas Patry committed
251
        prefix_lens = []
252

253
        # Parse batch
254
255
256
        for i, (r, tokenized_input) in enumerate(
            zip(pb.requests, batch_tokenized_inputs)
        ):
257
258
259
            # request id -> idx in list mapping
            requests_idx_mapping[r.id] = i

260
            tokenized_input = tokenized_input[-r.truncate :]
261
262
263
264
265
            if (
                tokenized_input[0] == tokenizer.bos_token_id
                and tokenized_input[1] == tokenizer.bos_token_id
            ):
                tokenized_input = tokenized_input[1:]
266

Nicolas Patry's avatar
Nicolas Patry committed
267
268
269
270
271
272
273
274
275
276
277
278
279
            orig_input_length = len(tokenized_input)

            if PREFIX_CACHING:
                prefix_len = r.prefix_len
                if prefix_len == orig_input_length:
                    assert prefix_len > 0
                    prefix_len -= 1
            else:
                prefix_len = 0

            prefix_ids.append(tokenized_input[:prefix_len])
            tokenized_input = tokenized_input[prefix_len:]

280
281
            input_length = len(tokenized_input)
            input_lengths.append(input_length)
282

283
            prefix_offsets.append(input_length - 5)
284
            read_offsets.append(input_length)
285

286
            all_input_ids.append(tokenized_input)
287
288

            # Position ids
Nicolas Patry's avatar
Nicolas Patry committed
289
290
291
            request_position_ids = torch.arange(
                prefix_len, orig_input_length, dtype=torch.int32
            )
292
            position_ids.append(request_position_ids)
293
294

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

297
            next_token_chooser_parameters.append(r.parameters)
298

299
300
301
            stopping_criteria = StoppingCriteria.from_pb(
                r.stopping_parameters, tokenizer
            )
302
            max_new_tokens = stopping_criteria.max_new_tokens
303
            stopping_criterias.append(stopping_criteria)
Nicolas Patry's avatar
Nicolas Patry committed
304
            top_n_tokens.append(r.top_n_tokens)
305

Nicolas Patry's avatar
Nicolas Patry committed
306
307
            ADAPTER_TO_INDEX = get_adapter_to_index()
            adapter_index = ADAPTER_TO_INDEX.get(r.adapter_id, 0)
drbh's avatar
drbh committed
308
309
310
            adapter_indices_list.append(torch.full((input_length,), adapter_index))
            adapter_set.add(adapter_index)

311
312
            # Paged attention
            # Remove one as the first token des not have a past
Nicolas Patry's avatar
Nicolas Patry committed
313
            speculative_length = get_speculate()
drbh's avatar
drbh committed
314
            speculative_length = 0 if speculative_length is None else speculative_length
Nicolas Patry's avatar
Nicolas Patry committed
315
316
317
318
319
320
321

            # Tokens that need to be mapped to blocks.
            block_tokens = orig_input_length + max_new_tokens - 1 + speculative_length

            # Tokens that need to be mapped to slots. We don't need slots for the
            # cached prefix (if present).
            slot_tokens = input_length + max_new_tokens - 1 + speculative_length
322
323
324

            # blocks and slots can be empty (for example in warmup)
            if not r.blocks:
Nicolas Patry's avatar
Nicolas Patry committed
325
                needed_blocks = math.ceil(block_tokens / BLOCK_SIZE)
326
327
328
329
330
331
332
333
334
335
                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
Nicolas Patry's avatar
Nicolas Patry committed
336
337
338
                request_slots = r.slots[
                    prefix_len:  #: orig_input_length + max_new_tokens + speculative_length
                ]
339
340

            block_tables.append(request_blocks)
Nicolas Patry's avatar
Nicolas Patry committed
341
342
343

            slots.extend(request_slots)
            prefix_lens.append(prefix_len)
344
            num_blocks += len(request_blocks)
Nicolas Patry's avatar
Nicolas Patry committed
345
            start_slots.append(cumulative_slot_tokens)
346
347

            request_slot_indices = torch.arange(
Nicolas Patry's avatar
Nicolas Patry committed
348
349
                cumulative_slot_tokens,
                cumulative_slot_tokens + input_length,
350
351
352
353
                dtype=torch.int64,
            )
            slot_indices.append(request_slot_indices)

354
355
356
357
358
359
360
361
362
            # 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)

363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
            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

383
384
            # Update
            cumulative_length += input_length
Nicolas Patry's avatar
Nicolas Patry committed
385
            cumulative_slot_tokens += slot_tokens
386
            max_seqlen = max(max_seqlen, input_length)
387
            max_blocks = max(max_blocks, len(request_blocks))
OlivierDehaene's avatar
OlivierDehaene committed
388
389
390
            max_length = max(
                max_length, input_length + max_new_tokens + speculative_length
            )
391

drbh's avatar
drbh committed
392
393
394
395
        adapter_indices = torch.cat(adapter_indices_list).to(
            dtype=torch.int64, device=device
        )

396
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
drbh's avatar
drbh committed
397
            next_token_chooser_parameters, dtype, device, tokenizer
398
        )
399
        start_slots = torch.tensor(start_slots, dtype=torch.int64)
400
401
402
403
404
405
406

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

408
409
410
411
412
        # Create tensors on device
        all_input_ids_tensor = torch.tensor(
            all_input_ids_tensor, dtype=torch.int64, device=device
        )

413
414
415
        if len(pb.requests) > 1:
            input_ids = np.concatenate(all_input_ids, dtype=np.int64)
            position_ids = torch.cat(position_ids)
416
            slot_indices = torch.cat(slot_indices)
417
418
            if sliding_window is not None:
                prefill_cache_indices = torch.cat(prefill_cache_indices)
419
420
421
        else:
            input_ids = all_input_ids[0]
            position_ids = position_ids[0]
422
            slot_indices = slot_indices[0]
423
424
            if sliding_window is not None:
                prefill_cache_indices = prefill_cache_indices[0]
425

426
427
        cu_seqlen_prefill = torch.tensor(
            cu_seqlen_prefill, device=device, dtype=torch.int32
428
429
430
        )
        position_ids = position_ids.to(device)
        slot_indices = slot_indices.to(device)
431
432
433
        prefill_cache_indices = (
            prefill_cache_indices.to(device) if sliding_window is not None else None
        )
434
        input_ids = torch.tensor(input_ids, dtype=torch.int64, device=device)
435
436
        input_lengths_tensor = torch.tensor(
            input_lengths, dtype=torch.int32, device=device
437
        )
438

drbh's avatar
drbh committed
439
440
441
442
443
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

444
445
        if all_prefill_logprobs:
            prefill_head_indices = None
446
            prefill_next_token_indices = cu_seqlen_prefill[1:] - 1
447
        elif no_prefill_logprobs:
448
            prefill_head_indices = cu_seqlen_prefill[1:] - 1
449
450
451
452
453
454
455
456
            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
457
458
459
        top_n_tokens_tensor = torch.tensor(
            top_n_tokens, device=device, dtype=torch.int64
        )
460

461
        slots = torch.tensor(slots, dtype=torch.int64, device=device)
Nicolas Patry's avatar
Nicolas Patry committed
462

463
464
465
466
467
468
        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)
Nicolas Patry's avatar
Nicolas Patry committed
469
        prefix_lens_tensor = torch.tensor(prefix_lens, dtype=torch.int32, device=device)
470

471
472
473
        return cls(
            batch_id=pb.id,
            requests=pb.requests,
474
            requests_idx_mapping=requests_idx_mapping,
475
476
            input_ids=input_ids,
            position_ids=position_ids,
477
            cu_seqlen_prefill=cu_seqlen_prefill,
478
            prefill_cache_indices=prefill_cache_indices,
479
480
            start_slots=start_slots,
            slot_indices=slot_indices,
481
482
483
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
Nicolas Patry's avatar
Nicolas Patry committed
484
485
            prefix_lens=prefix_lens,
            prefix_lens_tensor=prefix_lens_tensor,
486
            max_seqlen=max_seqlen,
487
488
489
            prefill_head_indices=prefill_head_indices,
            prefill_next_token_indices=prefill_next_token_indices,
            prefill_cu_outlens=prefill_cu_outlens,
490
            input_lengths=input_lengths,
491
            input_lengths_tensor=input_lengths_tensor,
492
493
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
494
            all_input_ids=all_input_ids,
495
            all_input_ids_tensor=all_input_ids_tensor,
Nicolas Patry's avatar
Nicolas Patry committed
496
            prefix_ids=prefix_ids,
497
            next_token_chooser=next_token_chooser,
498
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
499
500
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
501
            num_blocks=num_blocks,
502
            max_blocks=max_blocks,
drbh's avatar
drbh committed
503
504
505
506
507
508
            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
509
            speculative_ids=None,
510
511
        )

512
513
514
515
516
517
518
519
520
521
522
    @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)

523
    @tracer.start_as_current_span("filter")
524
525
    def filter(self, request_ids: List[int]) -> "FlashCausalLMBatch":
        if len(request_ids) == 0:
526
527
            raise ValueError("Batch must have at least one request")
        # We assume that if len(requests) == len(self) then the requests are the same
528
        if len(request_ids) == len(self):
529
530
            return self

531
        device = self.input_ids.device
532

533
534
535
        # New values after filtering
        requests_idx_mapping = {}

536
537
538
        # Used to index into tensors
        indices = []

539
540
541
        # slots to keep after filtering
        slot_filtering_indices = torch.zeros(
            self.slots.shape[0], dtype=torch.bool, device=device
542
543
        )

544
        # Create on CPU to only move to GPU once instead of at every copy
545
        slot_indices = torch.empty(len(request_ids), dtype=torch.int64)
546
547
        max_seqlen = 0

548
        requests = []
549
550
        start_slots = []
        block_tables = []
551
        all_input_ids = []
Nicolas Patry's avatar
Nicolas Patry committed
552
        prefix_ids = []
553

554
        input_lengths = []
Nicolas Patry's avatar
Nicolas Patry committed
555
        prefix_lens = []
556
557
        prefix_offsets = []
        read_offsets = []
558

559
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
560
        top_n_tokens = []
drbh's avatar
drbh committed
561
        adapter_set = set()
562

563
        num_blocks = 0
564
565
566
567
        max_blocks = 0
        # Cumulative length
        cumulative_max_length = 0

568
569
        for i, request_id in enumerate(request_ids):
            idx = self.requests_idx_mapping[request_id]
570
            indices.append(idx)
571
572
573
            requests_idx_mapping[request_id] = i

            requests.append(self.requests[idx])
574
575
576

            # Get length
            request_input_length = self.input_lengths[idx]
Nicolas Patry's avatar
Nicolas Patry committed
577
            prefix_len = self.prefix_lens[idx]
578
            max_seqlen = max(max_seqlen, request_input_length)
579

580
            all_input_ids.append(self.all_input_ids[idx])
Nicolas Patry's avatar
Nicolas Patry committed
581
            prefix_ids.append(self.prefix_ids[idx])
582
583

            input_lengths.append(request_input_length)
Nicolas Patry's avatar
Nicolas Patry committed
584
            prefix_lens.append(prefix_len)
585
586
            prefix_offsets.append(self.prefix_offsets[idx])
            read_offsets.append(self.read_offsets[idx])
587

588
589
            stopping_criteria = self.stopping_criterias[idx]
            stopping_criterias.append(stopping_criteria)
590

Nicolas Patry's avatar
Nicolas Patry committed
591
592
            top_n_tokens.append(self.top_n_tokens[idx])

Nicolas Patry's avatar
Nicolas Patry committed
593
594
            ADAPTER_TO_INDEX = get_adapter_to_index()
            adapter_index = ADAPTER_TO_INDEX.get(self.requests[idx].adapter_id, 0)
drbh's avatar
drbh committed
595
596
            adapter_set.add(adapter_index)

597
            remaining_tokens = (
598
599
                stopping_criteria.max_new_tokens - stopping_criteria.current_tokens
            )
600

601
            request_block_table = self.block_tables[idx]
602
            num_blocks += len(request_block_table)
603
604
605
            block_tables.append(request_block_table)
            start_slots.append(cumulative_max_length)

606
            # Copy to tensor (CPU)
607
            slot_indices[i] = cumulative_max_length + request_input_length - 1
608
609

            # Set slice
610
611
612
613
614
            slot_filtering_indices[
                self.start_slots[idx] : self.start_slots[idx]
                + request_input_length
                + remaining_tokens
                - 1
615
616
617
            ] = True

            cumulative_max_length += request_input_length + remaining_tokens - 1
618

619
620
            max_blocks = max(max_blocks, len(request_block_table))

621
622
623
        # Index into tensors
        input_ids = self.input_ids[indices]
        position_ids = self.position_ids[indices]
drbh's avatar
drbh committed
624
        adapter_indices = self.adapter_meta.adapter_indices[indices]
625
        all_input_ids_tensor = self.all_input_ids_tensor[indices]
626
627
628
        block_tables_tensor = self.block_tables_tensor[indices]
        input_lengths_tensor = self.input_lengths_tensor[indices]
        slots = self.slots[slot_filtering_indices]
Nicolas Patry's avatar
Nicolas Patry committed
629
        prefix_lens_tensor = self.prefix_lens_tensor[indices]
630
        next_token_chooser = self.next_token_chooser.filter(indices)
Nicolas Patry's avatar
Nicolas Patry committed
631
        top_n_tokens_tensor = self.top_n_tokens_tensor[indices]
OlivierDehaene's avatar
OlivierDehaene committed
632
633
634
        speculative_ids = (
            self.speculative_ids[indices] if self.speculative_ids is not None else None
        )
635
636

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

638
        # Move to GPU now that we have the whole tensor
639
        slot_indices = slot_indices.to(device)
640

drbh's avatar
drbh committed
641
642
643
644
645
        adapter_segments, adapter_segment_indices = find_segments(adapter_indices)
        adapter_segments = torch.tensor(
            adapter_segments, dtype=torch.int32, device=device
        )

646
        return type(self)(
647
648
649
650
651
            batch_id=self.batch_id,
            requests=requests,
            requests_idx_mapping=requests_idx_mapping,
            input_ids=input_ids,
            position_ids=position_ids,
652
            cu_seqlen_prefill=None,
653
            prefill_cache_indices=None,
654
655
656
657
658
            start_slots=start_slots,
            slot_indices=slot_indices,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
            slots=slots,
659
            max_seqlen=max_seqlen,
660
661
662
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
663
            input_lengths=input_lengths,
664
            input_lengths_tensor=input_lengths_tensor,
Nicolas Patry's avatar
Nicolas Patry committed
665
666
            prefix_lens=prefix_lens,
            prefix_lens_tensor=prefix_lens_tensor,
667
668
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
669
670
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
Nicolas Patry's avatar
Nicolas Patry committed
671
            prefix_ids=prefix_ids,
672
            next_token_chooser=next_token_chooser,
673
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
674
675
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
676
            num_blocks=num_blocks,
677
            max_blocks=max_blocks,
Nicolas Patry's avatar
Nicolas Patry committed
678
            speculative_ids=speculative_ids,
drbh's avatar
drbh committed
679
680
681
682
683
684
            adapter_meta=AdapterBatchMetadata(
                adapter_indices=adapter_indices,
                adapter_set=adapter_set,
                adapter_segments=adapter_segments,
                segment_indices=adapter_segment_indices,
            ),
685
686
687
688
689
690
691
692
693
        )

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

694
        num_blocks = 0
695
696
697
698
699
700
701
702
        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)
703
            num_blocks += b.num_blocks
OlivierDehaene's avatar
OlivierDehaene committed
704
705
706
            speculative_length = (
                b.speculative_ids.shape[1] if b.speculative_ids is not None else 0
            )
707
708
709
710
711
712
713
            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
714
                    + speculative_length
715
716
717
718
719
720
                    - stopping_criteria.current_tokens
                    for input_length, stopping_criteria in zip(
                        b.input_lengths, b.stopping_criterias
                    )
                ),
            )
721
722
723

        input_ids = batches[0].input_ids.new_empty(total_batch_size)
        position_ids = batches[0].position_ids.new_empty(total_batch_size)
724
725
726
727
728
729
730
731
        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)
        )
Nicolas Patry's avatar
Nicolas Patry committed
732
        prefix_lens_tensor = batches[0].prefix_lens_tensor.new_empty(total_batch_size)
733
734
        all_input_ids_tensor = batches[0].all_input_ids_tensor.new_zeros(
            (total_batch_size, max_length)
735
        )
Nicolas Patry's avatar
Nicolas Patry committed
736
737
738
        top_n_tokens_tensor = batches[0].top_n_tokens_tensor.new_zeros(
            total_batch_size,
        )
drbh's avatar
drbh committed
739
740
741
742
743
744
745
746
        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()
747

748
749
        start_slots = []
        block_tables = []
Nicolas Patry's avatar
Nicolas Patry committed
750
        prefix_lens = []
751
        all_input_ids = []
Nicolas Patry's avatar
Nicolas Patry committed
752
        prefix_ids = []
753
754

        input_lengths = []
755
756
        prefix_offsets = []
        read_offsets = []
757

758
        next_token_chooser_parameters = []
759
        fsm_grammar_states = []
760
        stopping_criterias = []
Nicolas Patry's avatar
Nicolas Patry committed
761
        top_n_tokens = []
762

763
        # Cumulative length
764
        cumulative_batch_size = 0
765
        cumulative_slots = 0
drbh's avatar
drbh committed
766
        cumulative_adapter_indices_size = 0
767
768
769

        for i, batch in enumerate(batches):
            requests.extend(batch.requests)
770
771
772
773
774
775
776
777

            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

778
779
            start_index = cumulative_batch_size
            end_index = cumulative_batch_size + len(batch)
780
781
            slots_start_index = cumulative_slots
            slots_end_index = cumulative_slots + len(batch.slots)
782
783
784
785

            # Copy tensors (GPU)
            input_ids[start_index:end_index] = batch.input_ids
            position_ids[start_index:end_index] = batch.position_ids
786
787
            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
788
            top_n_tokens_tensor[start_index:end_index] = batch.top_n_tokens_tensor
789
            slots[slots_start_index:slots_end_index] = batch.slots
790

drbh's avatar
drbh committed
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
            # 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
            )

806
807
808
            all_input_ids_tensor[
                start_index:end_index, : batch.all_input_ids_tensor.shape[1]
            ] = batch.all_input_ids_tensor[:, :max_length]
809

810
811
812
            block_tables_tensor[
                start_index:end_index, : batch.block_tables_tensor.shape[1]
            ] = batch.block_tables_tensor[:, :max_blocks]
813

Nicolas Patry's avatar
Nicolas Patry committed
814
815
            prefix_lens_tensor[start_index:end_index] = batch.prefix_lens_tensor

816
817
818
            start_slots.append(batch.start_slots + cumulative_slots)

            block_tables.extend(batch.block_tables)
Nicolas Patry's avatar
Nicolas Patry committed
819
            prefix_lens.extend(batch.prefix_lens)
820
            all_input_ids.extend(batch.all_input_ids)
Nicolas Patry's avatar
Nicolas Patry committed
821
            prefix_ids.extend(batch.prefix_ids)
822

823
            input_lengths.extend(batch.input_lengths)
824
825
            prefix_offsets.extend(batch.prefix_offsets)
            read_offsets.extend(batch.read_offsets)
826

827
            next_token_chooser_parameters.extend([r.parameters for r in batch.requests])
828
            fsm_grammar_states.extend(batch.next_token_chooser.fsm_grammar_states)
829
830
            stopping_criterias.extend(batch.stopping_criterias)

Nicolas Patry's avatar
Nicolas Patry committed
831
832
            top_n_tokens.extend(batch.top_n_tokens)

833
            # Update
834
            cumulative_batch_size += len(batch)
835
            cumulative_slots += len(batch.slots)
836

837
        start_slots = torch.concat(start_slots)
838

839
        next_token_chooser = HeterogeneousNextTokenChooser.from_pb(
840
841
842
            next_token_chooser_parameters,
            dtype=batches[0].next_token_chooser.dtype,
            device=batches[0].next_token_chooser.device,
drbh's avatar
drbh committed
843
            tokenizer=batches[0].next_token_chooser.tokenizer,
844
            fsm_grammar_states=fsm_grammar_states,
845
846
        )

OlivierDehaene's avatar
OlivierDehaene committed
847
848
849
850
851
        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
852

drbh's avatar
drbh committed
853
854
        adapter_segments, adapter_segment_indices = adapter_segment_builder.build()

855
        return cls(
856
857
            batch_id=batches[0].batch_id,
            requests=requests,
858
            requests_idx_mapping=requests_idx_mapping,
859
860
            input_ids=input_ids,
            position_ids=position_ids,
861
            cu_seqlen_prefill=None,
862
            prefill_cache_indices=None,
863
864
865
866
            start_slots=start_slots,
            slot_indices=slot_indices,
            block_tables=block_tables,
            block_tables_tensor=block_tables_tensor,
Nicolas Patry's avatar
Nicolas Patry committed
867
868
            prefix_lens=prefix_lens,
            prefix_lens_tensor=prefix_lens_tensor,
869
            slots=slots,
870
            max_seqlen=max_seqlen,
871
872
873
            prefill_head_indices=None,
            prefill_next_token_indices=None,
            prefill_cu_outlens=None,
874
            input_lengths=input_lengths,
875
            input_lengths_tensor=input_lengths_tensor,
876
877
            prefix_offsets=prefix_offsets,
            read_offsets=read_offsets,
878
879
            all_input_ids=all_input_ids,
            all_input_ids_tensor=all_input_ids_tensor,
Nicolas Patry's avatar
Nicolas Patry committed
880
            prefix_ids=prefix_ids,
881
            next_token_chooser=next_token_chooser,
882
            stopping_criterias=stopping_criterias,
Nicolas Patry's avatar
Nicolas Patry committed
883
884
            top_n_tokens=top_n_tokens,
            top_n_tokens_tensor=top_n_tokens_tensor,
885
            num_blocks=num_blocks,
886
            max_blocks=max_blocks,
OlivierDehaene's avatar
OlivierDehaene committed
887
            speculative_ids=speculative_ids,
drbh's avatar
drbh committed
888
889
890
891
892
893
            adapter_meta=AdapterBatchMetadata(
                adapter_indices=adapter_indices,
                adapter_set=adapter_set,
                adapter_segments=adapter_segments,
                segment_indices=adapter_segment_indices,
            ),
894
895
896
897
898
899
        )

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


900
901
902
903
904
905
906
907
908
909
910
911
ADAPTER_LAYERS = [
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
]
ROW_PARALLEL = {"o_proj", "down_proj", "lm_head"}


912
913
914
class FlashCausalLM(Model):
    def __init__(
        self,
drbh's avatar
drbh committed
915
        model_id: str,
916
917
918
919
920
921
922
923
924
925
926
927
        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
928
929
930
        num_kv_heads: Optional[int] = None,
        # Deepseek V2 uses different QK and V dims.
        head_size: Optional[int] = None,
931
        skip_special_tokens: bool = True,
932
    ):
Nicolas Patry's avatar
Nicolas Patry committed
933
        self.quantize = quantize
934
935
936
937
938
939
940
941
942
943
944
945
        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
946
                init_cpu_threads_env(rank_id=rank, world_size=world_size)
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
        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)

975
        weights_loader = get_loader(quantize, model_id, revision)
976
977
        filenames = weight_files(model_id, revision=revision, extension=".safetensors")
        weights = Weights(
978
979
980
981
982
983
            filenames,
            device,
            dtype,
            process_group=self.process_group,
            aliases=aliases,
            weights_loader=weights_loader,
984
985
986
987
988
989
990
991
992
993
        )

        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
994
995
996
997
998
999

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

1000
        self.num_layers = config.num_hidden_layers
1001
        self.num_heads = config.num_attention_heads
1002
1003
        # Validation is done in the model itself
        if num_kv_heads is None:
1004
1005
            num_kv_heads = getattr(config, "num_key_value_heads", None)
            # GPT-2 workaround
1006
            if num_kv_heads is None:
1007
1008
1009
                num_kv_heads = getattr(config, "n_head", None)
        if num_kv_heads is None:
            raise ValueError("Cannot get the number of key/value heads")
1010
1011
1012
1013
1014
1015
        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
1016
1017

        if head_size is None:
Nicolas Patry's avatar
Nicolas Patry committed
1018
1019
1020
1021
1022
1023
            # 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
1024
1025
        else:
            self.head_size = head_size
1026

1027
        self.cuda_graphs = {}
1028
        self.kv_cache = []
1029

1030
        if ATTENTION == "flashinfer":
Nicolas Patry's avatar
Nicolas Patry committed
1031
            from text_generation_server.layers.attention.flashinfer import (
1032
1033
                create_prefill_state,
                create_decode_state,
Nicolas Patry's avatar
Nicolas Patry committed
1034
                create_prefill_with_paged_kv_state,
1035
1036
1037
            )

            self.prefill_state = create_prefill_state(device=device)
Nicolas Patry's avatar
Nicolas Patry committed
1038
1039
1040
            self.prefill_with_paged_kv_state = create_prefill_with_paged_kv_state(
                device=device
            )
1041

Nicolas Patry's avatar
Nicolas Patry committed
1042
1043
1044
1045
1046
            self.decode_state = create_decode_state(
                device=device,
                num_heads=self.num_heads,
                num_kv_heads=self.num_kv_heads,
            )
1047

1048
        super().__init__(
drbh's avatar
drbh committed
1049
            model_id=model_id,
1050
            model=model,
1051
1052
1053
1054
            tokenizer=tokenizer,
            requires_padding=False,
            dtype=dtype,
            device=device,
1055
1056
            rank=rank,
            world_size=world_size,
1057
            sliding_window=config.sliding_window,
1058
1059
1060
1061
1062
1063
        )

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

1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
    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
1080
1081
1082
1083
        if SYSTEM == "ipex" and device.type == "xpu":
            x = 1
        else:
            x = BLOCK_SIZE // element_size
1084

1085
        if ATTENTION in {"flashdecoding", "flashinfer"}:
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
            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
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
            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)
            ]
1133

1134
1135
1136
    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)
1137
        slots = torch.arange(bs, dtype=torch.int64, device=self.device)
Nicolas Patry's avatar
Nicolas Patry committed
1138
1139
1140
1141
        input_lengths = [max_s] * bs
        prefix_lengths = [0] * bs
        input_lengths_tensor = (
            torch.ones(bs, dtype=torch.int32, device=self.device) * max_s
1142
        )
Nicolas Patry's avatar
Nicolas Patry committed
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
        prefix_lengths_tensor = torch.zeros(bs, dtype=torch.int32, device=self.device)
        block_tables = torch.arange(
            max_bt, dtype=torch.int32, device=self.device
        ).repeat(bs)
        block_tables = block_tables.reshape((bs, max_bt))

        if ATTENTION == "flashinfer":
            block_tables = block_tables_to_ragged(
                block_tables=block_tables,
                input_lengths=input_lengths,
                prefix_lens=prefix_lengths,
            )
1155
1156
1157
1158

        self.cuda_graphs[bs] = {
            "input_ids": input_ids,
            "position_ids": position_ids,
1159
            "kv_cache": self.kv_cache,
1160
1161
            "block_tables": block_tables,
            "slots": slots,
Nicolas Patry's avatar
Nicolas Patry committed
1162
            "input_lengths": input_lengths_tensor,
1163
        }
Nicolas Patry's avatar
Nicolas Patry committed
1164
        input_lengths_ = Seqlen(input_lengths=input_lengths_tensor)
1165
1166
1167
        graph = torch.cuda.CUDAGraph()
        self.cuda_graphs[bs]["graph"] = graph

1168
        if ATTENTION == "flashinfer":
Nicolas Patry's avatar
Nicolas Patry committed
1169
            from text_generation_server.layers.attention.flashinfer import (
1170
1171
1172
1173
1174
1175
1176
1177
1178
                create_decode_state_cuda_graphs,
            )

            block_tables_ptr = torch.zeros(
                bs + 1, dtype=torch.int32, device=self.device
            )
            last_page_len = torch.ones(bs, dtype=torch.int32, device=self.device)
            state = create_decode_state_cuda_graphs(
                device=input_ids.device,
Nicolas Patry's avatar
Nicolas Patry committed
1179
                block_tables=block_tables,
1180
1181
1182
1183
1184
1185
1186
1187
1188
                block_tables_ptr=block_tables_ptr,
                last_page_len=last_page_len,
                num_heads=self.num_heads,
                num_kv_heads=self.num_kv_heads,
            )
            self.cuda_graphs[bs]["state"] = state
        else:
            state = None

1189
1190
        torch.cuda.synchronize()
        # Run once outside to warmup
1191
        with self._forward_context(
1192
            block_tables=block_tables,
1193
1194
            cu_seqlen_prefill=None,
            input_lengths=input_lengths,
Nicolas Patry's avatar
Nicolas Patry committed
1195
            input_lengths_tensor=input_lengths_tensor,
1196
            state=state,
Nicolas Patry's avatar
Nicolas Patry committed
1197
1198
            prefix_lens=prefix_lengths,
            prefix_lens_tensor=prefix_lengths_tensor,
1199
1200
        ):
            self.model.forward(
1201
1202
1203
                input_ids=input_ids,
                position_ids=position_ids,
                cu_seqlen_prefill=None,
1204
                kv_cache=self.kv_cache,
1205
1206
                block_tables=block_tables,
                slots=slots,
1207
                input_lengths=input_lengths_,
1208
                max_s=max_s,
1209
                prefill_cache_indices=None,
1210
1211
                lm_head_indices=None,
            )
1212
1213
1214
1215

            torch.cuda.synchronize()

            with torch.cuda.graph(graph, pool=MEM_POOL):
Nicolas Patry's avatar
Nicolas Patry committed
1216
                input_lengths_tensor = Seqlen(input_lengths=input_lengths_tensor)
1217
1218
1219
1220
1221
1222
1223
                logits, speculative_logits = self.model.forward(
                    input_ids=input_ids,
                    position_ids=position_ids,
                    cu_seqlen_prefill=None,
                    kv_cache=self.kv_cache,
                    block_tables=block_tables,
                    slots=slots,
Nicolas Patry's avatar
Nicolas Patry committed
1224
                    input_lengths=input_lengths_tensor,
1225
1226
1227
1228
1229
1230
                    max_s=max_s,
                    prefill_cache_indices=None,
                    lm_head_indices=None,
                )
                self.cuda_graphs[bs]["logits"] = logits
                self.cuda_graphs[bs]["speculative_logits"] = speculative_logits
1231
1232
        torch.cuda.synchronize()

1233
    def warmup(self, batch: FlashCausalLMBatch):
1234
        # The warmup batch is the biggest batch we could ever receive
Nicolas Patry's avatar
Nicolas Patry committed
1235
1236
        empty_cache()

1237
        try:
1238
1239
            self.init_kv_cache(
                batch.num_blocks,
1240
1241
1242
1243
1244
1245
                self.num_layers,
                self.num_kv_heads,
                self.head_size,
                self.dtype,
                self.device,
            )
1246
            max_bt = batch.max_blocks
1247
            max_s = max_bt * BLOCK_SIZE
fxmarty's avatar
fxmarty committed
1248
1249
1250

            if SYSTEM == "rocm" and os.environ.get("PYTORCH_TUNABLEOP_ENABLED", False):
                torch.cuda.tunable.tuning_enable(False)
1251
            _, batch, _ = self.generate_token(batch)
OlivierDehaene's avatar
OlivierDehaene committed
1252
        except torch.cuda.OutOfMemoryError as e:
1253
            raise RuntimeError(
1254
1255
                f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. "
                f"You need to decrease `--max-batch-prefill-tokens`"
1256
            ) from e
1257

Nicolas Patry's avatar
Nicolas Patry committed
1258
        synchronize(self.device)
1259

1260
1261
        # 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
1262
1263
1264
1265
        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
1266
        free_memory = get_free_memory(self.device, MEMORY_FRACTION)
drbh's avatar
drbh committed
1267
        batch_num_blocks = batch.num_blocks if batch is not None else 0
1268
1269

        num_blocks = (
1270
1271
            # Leave 5% for some wiggle room
            int((free_memory * 0.95) // total_cache_size)
1272
            # Add batch.num_blocks as we allocated it above, so it is included in the peak memory.
drbh's avatar
drbh committed
1273
            + batch_num_blocks
1274
1275
        )

1276
        del batch
1277

1278
        self.init_kv_cache(
1279
1280
1281
1282
1283
1284
1285
1286
            num_blocks,
            self.num_layers,
            self.num_kv_heads,
            self.head_size,
            self.dtype,
            self.device,
        )

fxmarty's avatar
fxmarty committed
1287
1288
1289
1290
1291
        if SYSTEM == "rocm":
            if (
                os.environ.get("PYTORCH_TUNABLEOP_ENABLED") is None
                or os.environ.get("PYTORCH_TUNABLEOP_ENABLED") == "1"
            ):
1292
1293
                torch.cuda.tunable.enable()

fxmarty's avatar
fxmarty committed
1294
1295
1296
1297
1298
1299
1300
1301
                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(",")
                    ]
1302
                elif CUDA_GRAPHS is not None:
fxmarty's avatar
fxmarty committed
1303
                    tuning_sequences = CUDA_GRAPHS
1304
1305
1306
                else:
                    # For seqlen = 1, we dispatch to LLMM1 kernel.
                    tuning_sequences = [2, 3, 4, 5, 6, 7]
fxmarty's avatar
fxmarty committed
1307
1308
1309

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

1313
1314
1315
                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
1316
1317
1318
                )

                if os.path.isfile(tunableop_filepath):
1319
1320
1321
                    log_master(
                        logger.info,
                        f"The file {tunableop_filepath} already exists and will be reused.",
fxmarty's avatar
fxmarty committed
1322
1323
1324
1325
1326
1327
                    )
                    torch.cuda.tunable.read_file(tunableop_filepath)

                os.makedirs(HUGGINGFACE_HUB_CACHE, exist_ok=True)

                for seqlen in tuning_sequences:
1328
                    log_master(logger.info, f"Warming up TunableOp for seqlen={seqlen}")
fxmarty's avatar
fxmarty committed
1329
1330
1331
1332
                    self.tunableop_warmup(seqlen)
                    torch.cuda.tunable.write_file(tunableop_filepath)
                torch.cuda.tunable.tuning_enable(False)
            else:
1333
1334
1335
                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
1336
1337
                )

1338
        if CUDA_GRAPHS:
1339
            try:
1340
1341
1342
                log_master(
                    logger.info, f"Cuda Graphs are enabled for sizes {CUDA_GRAPHS}"
                )
1343
                # Warmup cuda graphs
1344
                for bs in CUDA_GRAPHS:
1345
1346
                    if self.speculate is None or self.speculate + 1 <= bs:
                        self.cuda_graph_warmup(bs, max_s, max_bt)
OlivierDehaene's avatar
OlivierDehaene committed
1347
            except torch.cuda.OutOfMemoryError:
1348
                logger.exception("Decode cuda graph warmup failed")
1349
        else:
1350
1351
1352
            log_master(
                logger.info, f"Cuda Graphs are disabled (CUDA_GRAPHS={CUDA_GRAPHS})."
            )
1353

1354
        return int(num_blocks * BLOCK_SIZE)
1355

fxmarty's avatar
fxmarty committed
1356
1357
1358
1359
1360
    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
1361
1362
        # Dummy value, some models (starcoder2) don't accept `None`.
        input_lengths = torch.ones(seqlen, dtype=torch.int32, device=self.device)
1363
        input_lengths = Seqlen(input_lengths=input_lengths)
fxmarty's avatar
fxmarty committed
1364

fxmarty's avatar
fxmarty committed
1365
1366
1367
1368
1369
1370
1371
        # 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
            ),
1372
            kv_cache=self.kv_cache,
fxmarty's avatar
fxmarty committed
1373
            block_tables=None,
fxmarty's avatar
fxmarty committed
1374
            input_lengths=input_lengths,
fxmarty's avatar
fxmarty committed
1375
1376
1377
            slots=slots,
            max_s=seqlen,
            lm_head_indices=None,
1378
            prefill_cache_indices=None,
fxmarty's avatar
fxmarty committed
1379
1380
        )

1381
    def forward(
drbh's avatar
drbh committed
1382
        self, batch: FlashCausalLMBatch, adapter_data: AdapterBatchData
1383
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
1384
        # Model Forward
Nicolas Patry's avatar
Nicolas Patry committed
1385
        if batch.speculative_ids is not None:
OlivierDehaene's avatar
OlivierDehaene committed
1386
1387
1388
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1389
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1390
1391
1392
1393
1394
            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
1395
1396
1397

            speculative_ids = batch.speculative_ids

OlivierDehaene's avatar
OlivierDehaene committed
1398
            B, speculative_length = speculative_ids.shape
Nicolas Patry's avatar
Nicolas Patry committed
1399
            new_length = speculative_length + 1
OlivierDehaene's avatar
OlivierDehaene committed
1400
1401
1402
            new_input_ids = torch.cat(
                [input_ids.unsqueeze(-1), speculative_ids], dim=1
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1403
1404
            arange = torch.arange(new_length, device=position_ids.device).unsqueeze(0)
            arange_int = arange.to(dtype=torch.int32)
OlivierDehaene's avatar
OlivierDehaene committed
1405
1406
1407
            new_position_ids = (
                position_ids.unsqueeze(-1).expand(B, new_length) + arange
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1408
            slots = (slots.unsqueeze(-1).expand(B, new_length) + arange_int).view(-1)
OlivierDehaene's avatar
OlivierDehaene committed
1409
1410
1411
            input_lengths = (
                input_lengths.unsqueeze(-1).expand(B, new_length) + arange_int
            ).view(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1412
1413
1414
            prefix_lens_tensor = (
                batch.prefix_lens_tensor.unsqueeze(-1).expand(B, new_length)
            ).reshape(-1)
Nicolas Patry's avatar
Nicolas Patry committed
1415
1416

            # Add Copy the block tables for all members
OlivierDehaene's avatar
OlivierDehaene committed
1417
1418
1419
1420
1421
1422
            block_tables = (
                block_tables.unsqueeze(1)
                .expand(B, new_length, -1)
                .reshape(B * new_length, -1)
                .contiguous()
            )
Nicolas Patry's avatar
Nicolas Patry committed
1423
1424
1425
1426
1427
            max_s = max_s + speculative_length

            input_ids = new_input_ids
            position_ids = new_position_ids
        else:
OlivierDehaene's avatar
OlivierDehaene committed
1428
1429
1430
            input_ids = batch.input_ids
            position_ids = batch.position_ids
            cu_seqlen_prefill = batch.cu_seqlen_prefill
1431
            kv_cache = self.kv_cache
OlivierDehaene's avatar
OlivierDehaene committed
1432
1433
1434
            block_tables = batch.block_tables_tensor
            slots = batch.slots[batch.slot_indices]
            input_lengths = batch.input_lengths_tensor
Nicolas Patry's avatar
Nicolas Patry committed
1435
            prefix_lens_tensor = batch.prefix_lens_tensor
OlivierDehaene's avatar
OlivierDehaene committed
1436
1437
            max_s = batch.max_seqlen
            lm_head_indices = batch.prefill_head_indices
Nicolas Patry's avatar
Nicolas Patry committed
1438

1439
1440
1441
1442
1443
1444
        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)

1445
        bs = input_ids.shape[0]
OlivierDehaene's avatar
OlivierDehaene committed
1446
1447
1448
1449
1450
1451
1452
1453
        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:
Nicolas Patry's avatar
Nicolas Patry committed
1454
1455
1456
1457
1458
1459
1460
            input_lengths = input_lengths + prefix_lens_tensor
            if PREFIX_CACHING:
                block_tables = block_tables_to_ragged(
                    block_tables=block_tables,
                    input_lengths=batch.input_lengths,
                    prefix_lens=batch.prefix_lens,
                )
1461
            with self._forward_context(
1462
                block_tables=block_tables,
1463
                cu_seqlen_prefill=cu_seqlen_prefill,
Nicolas Patry's avatar
Nicolas Patry committed
1464
1465
1466
1467
                input_lengths=batch.input_lengths,
                input_lengths_tensor=input_lengths,
                prefix_lens=batch.prefix_lens,
                prefix_lens_tensor=prefix_lens_tensor,
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
            ):
                input_lengths = Seqlen(input_lengths=input_lengths)
                logits, speculative_logits = self.model.forward(
                    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,
                    prefill_cache_indices=batch.prefill_cache_indices,
                    lm_head_indices=lm_head_indices,
                    adapter_data=adapter_data,
                )
                if batch.prefill_cache_indices is not None:
                    batch.prefill_cache_indices = None
                return logits, speculative_logits
1486
1487
1488
1489
1490

        # 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
Nicolas Patry's avatar
Nicolas Patry committed
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
        if ATTENTION == "flashinfer":
            block_tables = block_tables_to_ragged(
                block_tables=block_tables,
                input_lengths=batch.input_lengths,
                prefix_lens=batch.prefix_lens,
            )
            cuda_graph["block_tables"][: block_tables.shape[0]] = block_tables
        else:
            cuda_graph["block_tables"][
                : block_tables.shape[0], : block_tables.shape[1]
            ] = block_tables
1502
1503
1504
        cuda_graph["slots"].fill_(-1)
        cuda_graph["slots"][: slots.shape[0]] = slots
        cuda_graph["input_lengths"].zero_()
Nicolas Patry's avatar
Nicolas Patry committed
1505
1506
1507
        cuda_graph["input_lengths"][: input_lengths.shape[0]] = (
            input_lengths + prefix_lens_tensor
        )
1508

1509
        with self._forward_context(
Nicolas Patry's avatar
Nicolas Patry committed
1510
            block_tables=cuda_graph["block_tables"],
1511
            cu_seqlen_prefill=None,
Nicolas Patry's avatar
Nicolas Patry committed
1512
1513
1514
1515
1516
            input_lengths=batch.input_lengths,
            input_lengths_tensor=cuda_graph["input_lengths"],
            prefix_lens=batch.prefix_lens,
            prefix_lens_tensor=prefix_lens_tensor,
            state=cuda_graph.get("state"),
1517
1518
1519
1520
        ):
            # Replay the graph
            cuda_graph["graph"].replay()

1521
        # Slice output to the correct shape
1522
1523
1524
1525
1526
1527
1528
        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
1529
1530
1531
1532

    @tracer.start_as_current_span("generate_token")
    def generate_token(
        self, batch: FlashCausalLMBatch
1533
1534
    ) -> Tuple[List[Generation], Optional[FlashCausalLMBatch], Tuple[int, int]]:
        start = time.time_ns()
1535
        prefill = batch.cu_seqlen_prefill is not None
1536
        prefill_logprobs = batch.prefill_next_token_indices is not None
1537

drbh's avatar
drbh committed
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
        # 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)
1566

1567
1568
        if prefill:
            next_token_logits = (
1569
                out[batch.prefill_next_token_indices] if prefill_logprobs else out
1570
            )
Nicolas Patry's avatar
Nicolas Patry committed
1571
1572
            if speculative_logits is not None:
                speculative_logits = (
OlivierDehaene's avatar
OlivierDehaene committed
1573
1574
1575
                    speculative_logits[batch.prefill_next_token_indices]
                    if prefill_logprobs
                    else speculative_logits
Nicolas Patry's avatar
Nicolas Patry committed
1576
                )
drbh's avatar
drbh committed
1577
1578
1579
1580
            next_adapter_indices = batch.adapter_meta.adapter_indices.new_empty(
                len(batch)
            )

1581
1582
        else:
            next_token_logits = out
drbh's avatar
drbh committed
1583
            next_adapter_indices = batch.adapter_meta.adapter_indices
1584

Nicolas Patry's avatar
Nicolas Patry committed
1585
        speculate = get_speculate()
OlivierDehaene's avatar
OlivierDehaene committed
1586
1587
1588
1589
1590
1591
1592
1593
1594
        (
            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
1595
            speculate,
OlivierDehaene's avatar
OlivierDehaene committed
1596
1597
            batch.speculative_ids,
            speculative_logits,
1598
1599
        )

Nicolas Patry's avatar
Nicolas Patry committed
1600
        batch_top_token_ids, batch_top_token_logprobs = batch_top_tokens(
Nicolas Patry's avatar
Nicolas Patry committed
1601
            batch.top_n_tokens, batch.top_n_tokens_tensor, logprobs, accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1602
1603
        )

1604
        if prefill:
1605
            if len(batch) > 1 and prefill_logprobs:
1606
1607
                # 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
1608
                prefill_tokens_indices = batch.input_ids.new_zeros(len(out))
1609
1610

            next_position_ids = batch.position_ids.new_empty(len(batch))
1611
1612
1613
            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
1614
1615
1616
1617
        else:
            prefill_logprobs = None
            next_position_ids = batch.position_ids

1618
1619
1620
1621
1622
        # Cumulative length
        cumulative_length = 0

        # Results
        generations: List[Generation] = []
1623
        stopped = True
1624
1625

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

1628
1629
1630
1631
        # 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

1632
        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1633
        index = 0
OlivierDehaene's avatar
OlivierDehaene committed
1634
        for i, (input_length, all_input_ids, n_accepted_ids) in enumerate(iterator):
1635
            # Indexing metadata
1636
1637
1638
            start_index = cumulative_length
            end_index = cumulative_length + input_length

1639
            if prefill:
1640
1641
1642
1643
1644
                # 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

1645
1646
1647
1648
                # 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
1649
1650
1651
1652
1653
1654
                # 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
                ]

1655
1656
                # Used to gather prefill logprobs
                # Copy batch.input_ids to prefill_token_indices
1657
1658
                if prefill_logprobs:
                    if len(batch) > 1:
drbh's avatar
drbh committed
1659
1660
1661
                        prefill_tokens_indices[out_start_index : out_end_index - 1] = (
                            batch.input_ids[start_index + 1 : start_index + out_length]
                        )
1662
1663
1664
1665
1666
                    else:
                        # Set prefill_tokens_indices to the correct slice
                        prefill_tokens_indices = batch.input_ids[
                            start_index + 1 : start_index + out_length
                        ]
1667

Nicolas Patry's avatar
Nicolas Patry committed
1668
1669
1670
            for j in range(n_accepted_ids):
                batch.all_input_ids_tensor[i, input_length + j] = next_input_ids[index]
                index += 1
1671
1672
1673

            cumulative_length += input_length

drbh's avatar
drbh committed
1674
        # Update values
Nicolas Patry's avatar
Nicolas Patry committed
1675
1676
1677
1678
1679
        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
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
        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,
            )
1690

1691
        if prefill and prefill_logprobs:
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
            # 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
1702
        next_token_ids = next_input_ids.tolist()
1703
1704
        accepted_ids = accepted_ids.tolist()
        start_decode = time.time_ns()
1705
1706
1707
1708
1709

        # Zipped iterator
        iterator = zip(
            batch.requests,
            batch.input_lengths,
1710
1711
            batch.prefix_offsets,
            batch.read_offsets,
1712
1713
            batch.stopping_criterias,
            batch.all_input_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1714
            batch.prefix_ids,
1715
1716
            batch.next_token_chooser.do_sample,
            batch.next_token_chooser.seeds,
Nicolas Patry's avatar
Nicolas Patry committed
1717
            batch.top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1718
            accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1719
1720
            batch_top_token_ids,
            batch_top_token_logprobs,
1721
1722
1723
        )

        # For each member of the batch
Nicolas Patry's avatar
Nicolas Patry committed
1724
        index = 0
1725
1726
1727
        for i, (
            request,
            input_length,
1728
1729
            prefix_offset,
            read_offset,
1730
1731
            stopping_criteria,
            all_input_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1732
            prefix_ids,
1733
1734
            do_sample,
            seed,
Nicolas Patry's avatar
Nicolas Patry committed
1735
            top_n_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1736
            n_accepted_ids,
Nicolas Patry's avatar
Nicolas Patry committed
1737
1738
            top_token_ids,
            top_token_logprobs,
1739
        ) in enumerate(iterator):
1740
            # Append next token to all tokens
Nicolas Patry's avatar
Nicolas Patry committed
1741
1742
1743
            next_token_texts = []
            left = 0

1744
            if n_accepted_ids > 1:
1745
                log_master(logger.debug, f"Speculated ids {n_accepted_ids - 1}")
1746

Nicolas Patry's avatar
Nicolas Patry committed
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
            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)
1758

Nicolas Patry's avatar
Nicolas Patry committed
1759
1760
1761
1762
                stop, reason = stopping_criteria(
                    next_token_id,
                    next_token_text,
                )
1763

Nicolas Patry's avatar
Nicolas Patry committed
1764
1765
1766
1767
1768
1769
1770
                if stop:
                    left = index + n_accepted_ids - j - 1
                    current_stopped = True
                    break
                else:
                    current_stopped = False
            stopped = stopped and current_stopped
1771

OlivierDehaene's avatar
OlivierDehaene committed
1772
1773
1774
1775
            _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
1776
            index += n_accepted_ids
1777

1778
1779
1780
1781
1782
            # Shard generations
            # All generations will be appended in the rust sharded client
            if i % self.world_size == self.rank:
                if stop:
                    # Decode generated tokens
1783
1784
                    output_text, _, _ = self.decode_token(
                        all_input_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1785
1786
1787
1788
1789
1790
                        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,
1791
1792
                    )
                    generated_text = GeneratedText(
1793
1794
1795
1796
                        output_text,
                        stopping_criteria.current_tokens,
                        reason,
                        seed if do_sample else None,
1797
1798
1799
1800
1801
                    )
                else:
                    generated_text = None

                # Prefill
1802
1803
1804
1805
                if prefill and request.prefill_logprobs:
                    out_start_index = batch.prefill_cu_outlens[i]
                    out_end_index = batch.prefill_cu_outlens[i + 1]

1806
                    # Remove generated token to only have prefill and add nan for first prompt token
Nicolas Patry's avatar
Nicolas Patry committed
1807
1808
1809
                    request_prefill_logprobs = (
                        [float("nan")] * (len(prefix_ids) + 1)
                    ) + prefill_logprobs[out_start_index : out_end_index - 1]
1810
1811
                    prefill_token_ids = all_input_ids[:-1]
                    prefill_texts = self.tokenizer.batch_decode(
Nicolas Patry's avatar
Nicolas Patry committed
1812
                        prefix_ids + prefill_token_ids,
1813
1814
1815
                        clean_up_tokenization_spaces=False,
                        skip_special_tokens=False,
                    )
Nicolas Patry's avatar
Nicolas Patry committed
1816
1817

                    prefill_tokens = Tokens(
Nicolas Patry's avatar
Nicolas Patry committed
1818
                        prefix_ids + prefill_token_ids,
OlivierDehaene's avatar
OlivierDehaene committed
1819
1820
1821
                        request_prefill_logprobs,
                        prefill_texts,
                        is_special=[],
1822
1823
1824
1825
                    )
                else:
                    prefill_tokens = None

Nicolas Patry's avatar
Nicolas Patry committed
1826
                if top_n_tokens > 0:
Nicolas Patry's avatar
Nicolas Patry committed
1827
                    all_top_tokens = []
drbh's avatar
drbh committed
1828
                    for top_token_ids, top_token_logprobs in zip(
1829
1830
                        top_token_ids, top_token_logprobs
                    ):
Nicolas Patry's avatar
Nicolas Patry committed
1831
1832
1833
1834
1835
1836
                        toptoken_texts = self.tokenizer.batch_decode(
                            top_token_ids,
                            clean_up_tokenization_spaces=False,
                            skip_special_tokens=False,
                        )
                        special_toptokens = [
1837
1838
                            token_id in self.all_special_ids
                            for token_id in top_token_ids
Nicolas Patry's avatar
Nicolas Patry committed
1839
1840
1841
1842
1843
1844
1845
1846
1847
                        ]
                        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
1848
1849
1850
                else:
                    top_tokens = None

1851
1852
1853
                generation = Generation(
                    request.id,
                    prefill_tokens,
Nicolas Patry's avatar
Nicolas Patry committed
1854
1855
1856
1857
1858
1859
                    Tokens(
                        _next_token_ids,
                        _next_token_logprobs,
                        next_token_texts,
                        [nid in self.all_special_ids for nid in _next_token_ids],
                    ),
1860
                    generated_text,
Nicolas Patry's avatar
Nicolas Patry committed
1861
                    top_tokens,
1862
1863
                )

1864
                generations.append(generation)
1865

drbh's avatar
drbh committed
1866
1867
1868
            # 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
1869
1870
1871
                batch.next_token_chooser = (
                    batch.next_token_chooser.advance_grammar_single(i, next_token_id)
                )
drbh's avatar
drbh committed
1872

1873
            # Update values
1874
            batch.input_lengths[i] = input_length + n_accepted_ids
Nicolas Patry's avatar
Nicolas Patry committed
1875
1876
            if batch.input_lengths[i] > batch.max_seqlen:
                batch.max_seqlen = batch.input_lengths[i]
1877
1878
            batch.prefix_offsets[i] = prefix_offset
            batch.read_offsets[i] = read_offset
1879
1880
            batch.all_input_ids[i] = all_input_ids

1881
1882
        if stopped:
            # No need to return a batch if we know that all requests stopped
1883
1884
1885
            forward_ns = start_decode - start
            decode_ns = time.time_ns() - start_decode
            return generations, None, (forward_ns, decode_ns)
1886

1887
1888
1889
        batch.prefill_cu_outlens = None
        batch.prefill_head_indices = None
        batch.prefill_next_token_indices = None
1890

1891
1892
1893
        forward_ns = start_decode - start
        decode_ns = time.time_ns() - start_decode
        return generations, batch, (forward_ns, decode_ns)
1894
1895
1896
1897
1898
1899

    def _forward_context(
        self,
        *,
        block_tables: torch.Tensor,
        cu_seqlen_prefill: Optional[torch.Tensor],
Nicolas Patry's avatar
Nicolas Patry committed
1900
1901
1902
1903
        input_lengths: List[int],
        input_lengths_tensor: torch.Tensor,
        prefix_lens: List[int],
        prefix_lens_tensor: torch.Tensor,
1904
1905
        state: Optional[Any] = None,
    ) -> ContextManager:
1906
        if ATTENTION != "flashinfer":
1907
1908
            return nullcontext()

Nicolas Patry's avatar
Nicolas Patry committed
1909
        from text_generation_server.layers.attention.flashinfer import (
1910
            use_decode_state,
Nicolas Patry's avatar
Nicolas Patry committed
1911
            use_prefill_with_paged_kv_state,
1912
1913
        )

Nicolas Patry's avatar
Nicolas Patry committed
1914
1915
        # has_prefix_lens = any(prefix_len > 0 for prefix_len in prefix_lens)

1916
        if cu_seqlen_prefill is not None:
Nicolas Patry's avatar
Nicolas Patry committed
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
            return use_prefill_with_paged_kv_state(
                state=(
                    state if state is not None else self.prefill_with_paged_kv_state
                ),
                # block_tables=block_tables_to_ragged(
                #     block_tables=block_tables,
                #     input_lengths=input_lengths,
                #     prefix_lens=prefix_lens,
                # ),
                block_tables=block_tables,
1927
                cu_seqlens=cu_seqlen_prefill,
Nicolas Patry's avatar
Nicolas Patry committed
1928
                input_lengths=input_lengths_tensor,
1929
1930
1931
                num_heads=self.num_heads,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
Nicolas Patry's avatar
Nicolas Patry committed
1932
                page_size=BLOCK_SIZE,
1933
1934
            )
        else:
Nicolas Patry's avatar
Nicolas Patry committed
1935
            assert input_lengths_tensor is not None
1936
1937
            return use_decode_state(
                state=state if state is not None else self.decode_state,
Nicolas Patry's avatar
Nicolas Patry committed
1938
1939
                input_lengths=input_lengths_tensor,
                block_tables=block_tables,
1940
1941
1942
1943
1944
                num_heads=self.num_heads,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
                page_size=BLOCK_SIZE,
            )
Nicolas Patry's avatar
Nicolas Patry committed
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964


def block_tables_to_ragged(
    *, block_tables: torch.Tensor, input_lengths: List[int], prefix_lens: List[int]
) -> torch.Tensor:
    """Convert block table to ragged format compatible with FlashInfer."""
    assert len(input_lengths) == len(prefix_lens)

    total_len = sum(input_lengths) + sum(prefix_lens)
    block_tables_ragged = torch.empty(
        total_len, dtype=torch.int32, device=block_tables.device
    )

    offset = 0
    for i, (input_length, prefix_len) in enumerate(zip(input_lengths, prefix_lens)):
        seq_len = prefix_len + input_length
        block_tables_ragged[offset : offset + seq_len] = block_tables[i][:seq_len]
        offset += seq_len

    return block_tables_ragged