Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[python] Make paged attention configurable #986

Merged
merged 1 commit into from
Aug 1, 2023
Merged

[python] Make paged attention configurable #986

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
[python] Make paged attention configurable
  • Loading branch information
@xyang16
xyang16 committed Jul 31, 2023
commit d7560394bd9b50d5aa34044d1277f9f4071cb4f6
9 changes: 6 additions & 3 deletions engines/python/setup/djl_python/rolling_batch/lmi_dist_rolling_batch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,9 @@ def __init__(self, model_id_or_path, device, properties, **kwargs):
self.properties = properties
self.batch_cls = None
self._init_model(kwargs, model_id_or_path)
self._warmup(**kwargs)
self.paged_attention = self.properties.get("paged_attention", "true").lower() == "true"
if self.paged_attention:
self._warmup(**kwargs)
self.batch_id_counter = 0
self.cache: Batch = None

Expand All @@ -67,7 +69,8 @@ def _init_model(self, kwargs, model_id_or_path):
sharded=sharded,
quantize=quantize,
dtype=dtype,
trust_remote_code=kwargs.get("trust_remote_code"))
trust_remote_code=kwargs.get("trust_remote_code"),
paged_attention=self.paged_attention)
self.batch_cls = self.model.batch_type

def _warmup(self, **kwargs):
Expand All @@ -84,7 +87,7 @@ def _warmup(self, **kwargs):
max_prefill_tokens = int(
self.properties.get(
"max_rolling_batch_prefill_tokens",
int(self.properties.get("max_rolling_batch_size", 4)) * 512))
int(self.properties.get("max_rolling_batch_size", 4)) * 272))
requests = [
lmi_dist.utils.types.Request(id=0,
inputs='_test ' * max_prefill_tokens,
Expand Down