Skip to content

Commit

Permalink
input_lengths -> seqlen
Browse files Browse the repository at this point in the history
  • Loading branch information
@tgaddair
tgaddair committed Oct 18, 2024
1 parent df5ae30 commit 401f1ae
Show file tree
Hide file tree
Showing 13 changed files with 133 additions and 115 deletions.
18 changes: 10 additions & 8 deletions server/lorax_server/models/custom_modeling/flash_cohere_modeling.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@

from typing import List, Optional, Tuple

from lorax_server.utils.attention.common import Seqlen

import dropout_layer_norm
import rotary_emb
import torch
Expand Down Expand Up @@ -252,7 +254,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand Down Expand Up @@ -303,7 +305,7 @@ def forward(
self.kv_head_mapping,
self.softmax_scale,
block_tables,
input_lengths,
seqlen,
max_s,
)

Expand Down Expand Up @@ -387,7 +389,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand All @@ -402,7 +404,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -453,7 +455,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
) -> torch.Tensor:
Expand All @@ -474,7 +476,7 @@ def forward(
kv_cache[i],
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -518,7 +520,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
prefill_cache_indices: Optional[torch.Tensor] = None,
Expand All @@ -531,7 +533,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down
21 changes: 11 additions & 10 deletions server/lorax_server/models/custom_modeling/flash_dbrx_modeling.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

from typing import Any, List, Optional, Tuple

from lorax_server.utils.attention.common import Seqlen
import numpy as np
import torch
import torch.distributed
Expand Down Expand Up @@ -394,7 +395,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand Down Expand Up @@ -441,7 +442,7 @@ def forward(
self.kv_head_mapping,
self.softmax_scale,
block_tables,
input_lengths,
seqlen,
max_s,
)

Expand Down Expand Up @@ -475,7 +476,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand All @@ -490,7 +491,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -892,7 +893,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand All @@ -906,7 +907,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -947,7 +948,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
) -> torch.Tensor:
Expand All @@ -968,7 +969,7 @@ def forward(
kv_cache[i],
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -1003,7 +1004,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
prefill_cache_indices: Optional[torch.Tensor] = None,
Expand All @@ -1016,7 +1017,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down
17 changes: 9 additions & 8 deletions server/lorax_server/models/custom_modeling/flash_gemma2_modeling.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@

from typing import List, Optional, Tuple

from lorax_server.utils.attention.common import Seqlen
import torch
import torch.distributed
from torch import nn
Expand Down Expand Up @@ -239,7 +240,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand Down Expand Up @@ -283,7 +284,7 @@ def forward(
self.kv_head_mapping,
self.softmax_scale,
block_tables,
input_lengths,
seqlen,
max_s,
)

Expand Down Expand Up @@ -379,7 +380,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand All @@ -394,7 +395,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -445,7 +446,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
) -> torch.Tensor:
Expand All @@ -466,7 +467,7 @@ def forward(
kv_cache[i],
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -510,7 +511,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
prefill_cache_indices: Optional[torch.Tensor] = None,
Expand All @@ -524,7 +525,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down
17 changes: 9 additions & 8 deletions server/lorax_server/models/custom_modeling/flash_gemma_modeling.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

from typing import List, Optional, Tuple

from lorax_server.utils.attention.common import Seqlen
import torch
import torch.distributed
from torch import nn
Expand Down Expand Up @@ -245,7 +246,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand Down Expand Up @@ -288,7 +289,7 @@ def forward(
self.kv_head_mapping,
self.softmax_scale,
block_tables,
input_lengths,
seqlen,
max_s,
)

Expand Down Expand Up @@ -393,7 +394,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
):
Expand All @@ -409,7 +410,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down Expand Up @@ -458,7 +459,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
) -> torch.Tensor:
Expand All @@ -482,7 +483,7 @@ def forward(
kv_cache[i],
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand All @@ -509,7 +510,7 @@ def forward(
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
seqlen: Seqlen,
max_s: int,
adapter_data: AdapterBatchData,
prefill_cache_indices: Optional[torch.Tensor] = None,
Expand All @@ -522,7 +523,7 @@ def forward(
kv_cache,
block_tables,
slots,
input_lengths,
seqlen,
max_s,
adapter_data,
)
Expand Down
Loading

0 comments on commit 401f1ae

Please sign in to comment.