Skip to content

Support logits_soft_cap parameter in paged_attention#7704

Merged
wonjoo-wj merged 2 commits intomasterfrom
wonjoo/paged-attention-logits-soft-cap
Jul 17, 2024
Merged

Support logits_soft_cap parameter in paged_attention#7704
wonjoo-wj merged 2 commits intomasterfrom
wonjoo/paged-attention-logits-soft-cap

Conversation

@wonjoo-wj
Copy link
Copy Markdown
Collaborator

@wonjoo-wj wonjoo-wj commented Jul 17, 2024

Support logits_soft_cap parameter in paged_attention

JAX reference: jax-ml/jax@f0e36d5#diff-f1713aabe72665edd0f07fe591b1a9ff03327ca4f8d229489fa14297c25bab08

Test plan:

root@t1v-n-4989e8c7-w-0:~/pytorch/xla# python test/test_pallas.py PallasTest.test_paged_attention_wrapper_with_attn_logits_soft_cap
INFO:jax._src.xla_bridge:Unable to initialize backend 'cuda': 
INFO:jax._src.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
.
----------------------------------------------------------------------
Ran 1 test in 2.782s

OK

@wonjoo-wj wonjoo-wj merged commit 2a20d1e into master Jul 17, 2024
This was referenced Aug 28, 2024
Closed
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@JackCaoG JackCaoG JackCaoG approved these changes

@alanwaketan alanwaketan Awaiting requested review from alanwaketan

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@wonjoo-wj @JackCaoG