Skip to content

Address compatibility issues arising from the removal of the XLA dependency #1423

Merged
kyuyeunk merged 14 commits intomainfrom
weiyu/test-vllm-change
Jan 9, 2026
Merged

Address compatibility issues arising from the removal of the XLA dependency #1423
kyuyeunk merged 14 commits intomainfrom
weiyu/test-vllm-change

Conversation

@weiyu0824
Copy link
Copy Markdown
Collaborator

Description

This PR addresses compatibility issues arising from the removal of the XLA dependency in upstream vLLM (reference.

Changes:

  • Moved PallasAttention static function from vllm to tpu-inference.
  • Update TpuPlatform in tpu_platform.py to use local PallasAttentionBackend.
  • Add create_weights in compressed_tensors_w8a8_int8.py because vllm's implementation would introduce another error.
  • Register PallasAttention using decorator (see line:32 and line:37 in tpu_inference/layers/vllm/attention.py) because we already remove Pallas enum in upstream.

Tests

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@weiyu0824 weiyu0824 changed the title Weiyu/test vllm change Addresses compatibility issues arising from the removal of the XLA dependency Jan 8, 2026
@weiyu0824 weiyu0824 changed the title Addresses compatibility issues arising from the removal of the XLA dependency Address compatibility issues arising from the removal of the XLA dependency Jan 8, 2026
@github-actions
Copy link
Copy Markdown

github-actions bot commented Jan 8, 2026

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

  • why is this change being made,
  • the problem being solved and any relevant context,
  • why this is a good solution,
  • some information about the specific implementation,
  • shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

Copy link
Copy Markdown
Collaborator

@QiliangCui QiliangCui left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you explain more about why

Add create_weights in compressed_tensors_w8a8_int8.py because vllm's implementation would introduce another error.

@weiyu0824 weiyu0824 requested a review from jrplatin as a code owner January 8, 2026 17:07
@weiyu0824 weiyu0824 force-pushed the weiyu/test-vllm-change branch from 54c01b2 to 38f8531 Compare January 9, 2026 02:14
@weiyu0824 weiyu0824 requested a review from jcyang43 as a code owner January 9, 2026 02:14
@weiyu0824 weiyu0824 force-pushed the weiyu/test-vllm-change branch from 5102888 to b3c04e9 Compare January 9, 2026 02:41
@weiyu0824 weiyu0824 added the ready ONLY add when PR is ready to merge/full CI is needed label Jan 9, 2026
@kyuyeunk
Copy link
Copy Markdown
Collaborator

kyuyeunk commented Jan 9, 2026

seems like ci is failing due to upstream change: vllm-project/vllm#30519

I think this should be a simple fix. @weiyu0824 can you add a fix commit to this branch as well?

@weiyu0824 weiyu0824 force-pushed the weiyu/test-vllm-change branch from b3c04e9 to bce6db3 Compare January 9, 2026 03:34
@weiyu0824 weiyu0824 force-pushed the weiyu/test-vllm-change branch from 8ff7761 to 5977a14 Compare January 9, 2026 06:32
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
@weiyu0824 weiyu0824 force-pushed the weiyu/test-vllm-change branch from 5977a14 to ec15832 Compare January 9, 2026 06:35
@kyuyeunk
Copy link
Copy Markdown
Collaborator

kyuyeunk commented Jan 9, 2026

hmmm we have another error due to upstream change 1 hour ago: vllm-project/vllm#32003

yeah, feel free to ignore all my other comments and let's just try to get this pr merged asap.

Signed-off-by: Wei-Yu Lin <weiyulin@google.com>
@kyuyeunk kyuyeunk requested a review from Lumosis as a code owner January 9, 2026 08:52
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
@kyuyeunk kyuyeunk force-pushed the weiyu/test-vllm-change branch from 4d84445 to e5deac5 Compare January 9, 2026 08:52
This reverts commit fece762.

Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
@kyuyeunk kyuyeunk force-pushed the weiyu/test-vllm-change branch 4 times, most recently from 77c71d5 to 90480a1 Compare January 9, 2026 09:32
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
@kyuyeunk kyuyeunk force-pushed the weiyu/test-vllm-change branch 2 times, most recently from 317a769 to 4719933 Compare January 9, 2026 10:41
Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>
@kyuyeunk kyuyeunk force-pushed the weiyu/test-vllm-change branch from 4719933 to 9fbda58 Compare January 9, 2026 10:59
@kyuyeunk kyuyeunk merged commit 67b2082 into main Jan 9, 2026
39 checks passed
@kyuyeunk kyuyeunk deleted the weiyu/test-vllm-change branch January 30, 2026 09:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants