mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-19 22:02:06 +00:00
e415b690a6
* Making prefix/flashinfer the default and testing the full release tests. * Include flashinfer in the docker. * Using prebuilt. * Allowing window_left_size (dummy version). * Disabling flashinfer/prefix caching on odd head_dim * Disable prefix caching for lora. * More specific codes. * Update lock * Updating integration tests with new values with FI/FD. Remove paged as a default too, and using FD everywhere. * Update cargo lock ? * Upgrade to 1.80 because of bitstream... * Everywhere 1.80 * Forgot last default place. * Apply suggestions from code review Co-authored-by: drbh <david.richard.holtz@gmail.com> * Updated flake lock * Tmp * Upgrade resolution system for less errors in resolution. * Remove lambda for cleaner function. * Handling debugger. * OVerride the env in server tests. * Is this enough to make it work ? * This seems to be working. * Downgrade some logs. * Fixing the default for vlm. * Don't enable prefix caching on VLM just yet. * Change `add_special_tokens` in order to have the correct tokens for chat input and not (since it's super important with the prefixing now) * Fixing prefix caching for flashdecoding. * Update all models. * Fixed flashinfer version. * add_special_tokens is internal only * Fixing seqlen with the new vlms. * Fixing the issue with `add_special_tokens` not being passed around. * Fixing the test. * Removing encoder_decoder (seq2seq). * Update the chat test. * Fixing the batching tokenization in flash causal lm. * Truncating left for radix purposes. * Oops this doesn't belong here. * Put back default pure shell. * Update server tests - Default to throughput test in k6 - Use TGI_WIGGLE_ROOM to adjust wiggle room * Only n_heads / process_group.size() are necessary. * Revert the integrationt tests change (seem linked to head_size modification). * Adding error message when assert is violated. * Fixing the free algorithm to handle times where the common prefix is smaller. * Apply suggestions from code review Co-authored-by: OlivierDehaene <olivier@huggingface.co> * Update server/text_generation_server/layers/attention/common.py Co-authored-by: OlivierDehaene <olivier@huggingface.co> * Fix disabling prefix caching - Fix windowing checks. * Revert the Cohere tokenizer change (for now using a revision instead). * Fmt. --------- Co-authored-by: drbh <david.richard.holtz@gmail.com> Co-authored-by: OlivierDehaene <olivier@huggingface.co>
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
from dataclasses import dataclass
|
|
from text_generation_server.models.globals import ATTENTION
|
|
import torch
|
|
from typing import Optional
|
|
|
|
|
|
if ATTENTION in {"flashinfer", "flashdecoding"}:
|
|
|
|
@dataclass
|
|
class Seqlen:
|
|
input_lengths: torch.Tensor
|
|
prefix_lengths: torch.Tensor
|
|
cu_seqlen_q: Optional[torch.Tensor]
|
|
cu_seqlen_k: Optional[torch.Tensor]
|
|
max_q: int
|
|
max_k: int
|
|
|
|
def __init__(
|
|
self,
|
|
input_lengths,
|
|
prefix_lengths,
|
|
cu_seqlen_q=None,
|
|
max_q=None,
|
|
max_k=None,
|
|
):
|
|
self.input_lengths = input_lengths
|
|
self.prefix_lengths = prefix_lengths
|
|
device = self.input_lengths.device
|
|
shape = self.input_lengths.shape
|
|
if cu_seqlen_q is None:
|
|
cu_seqlen_q = torch.arange(
|
|
shape[0] + 1,
|
|
device=device,
|
|
dtype=torch.int32,
|
|
)
|
|
max_q = 1
|
|
else:
|
|
assert max_q is not None
|
|
assert max_k is not None
|
|
cu_seqlen_k = torch.zeros(shape[-1] + 1, device=device, dtype=torch.int32)
|
|
|
|
# cuda graphs don't like this and this is necessary to clamp within mistral
|
|
# Although FA2 might not want the clamping
|
|
# cu_seqlen_k[0] = 0
|
|
total = self.input_lengths + self.prefix_lengths
|
|
torch.cumsum(total, -1, out=cu_seqlen_k[1:])
|
|
|
|
self.cu_seqlen_q = cu_seqlen_q
|
|
self.cu_seqlen_k = cu_seqlen_k
|
|
self.max_q = max_q
|
|
self.max_k = max_k
|
|
|
|
def clamp(self, max):
|
|
# Flash decoding doesn't need to clamp
|
|
return self
|
|
|
|
else:
|
|
|
|
@dataclass
|
|
class Seqlen:
|
|
input_lengths: torch.Tensor
|
|
prefix_lengths: torch.Tensor
|
|
cu_seqlen_q: torch.Tensor
|
|
max_q: int
|
|
max_k: int
|
|
|
|
def clamp(self, max):
|
|
raise NotImplementedError("Not implemented seqlen for paged")
|
|
return Seqlen(torch.clamp(self.input_lengths, max=max))
|