[CI/Build] tests(v1): feed Triton attention the (num_blocks, 2, …) KV cache layout in backend-correctness tests

[hl475]

Purpose

This PR makes K/V cache unbinding robust across cache layouts by detecting the axis of size 2 at runtime instead of assuming it sits at dim=1. This fixes unpacking errors seen when kv_cache is shaped with the K/V dimension elsewhere (e.g., dim=0).

When running tests tests/v1/attention/test_attention_backends.py on H100, the following line

key_cache, value_cache = kv_cache.unbind(1)

failed with ValueError: too many values to unpack (expected 2)

This change avoids relying on a specific layout and works with both older and newer cache shapes.

Per discussion below, we updated the test to provide Triton the same (num_blocks, 2, …) KV cache layout that FlashInfer uses.

Test Plan

pytest -v -s tests/v1/attention/test_attention_backends.py

on H100

Test Result

============================================================= 15 passed, 1 warning in 48.49s ==============================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

[tdoublep]

Could you explain exactly which test is failing? Could it just be that the test is using the "wrong" layout?

[hl475]

Could you explain exactly which test is failing? Could it just be that the test is using the "wrong" layout?

Thanks @tdoublep ! The failed test is from https://github.com/vllm-project/vllm/blob/a6049be73cb965bad04f6657de6c4d98261a5237/tests/v1/attention/test_attention_backends.py where all 15 tests are failing on H100 with the same unpack error

[tdoublep]

OK. I think that the test should be modified to provide the KV cache with the correct layout. e.g., we can look at how it works for Flashinfer which has the same (num_blocks, 2, ...) layout.

[hl475]

Thanks @tdoublep for suggestions! I updated my PR to only change the test. PTAL

[tdoublep]

Thank you. Curious now why this test failure wasn't caught as part of CI. Are we failing to trigger this test when we change the attention backend code?

[hl475]

Thank you. Curious now why this test failure wasn't caught as part of CI. Are we failing to trigger this test when we change the attention backend code?

I don't think this test is currently running in CI. We are trying to enabling them from #26649 , so I found the failure.

[yeqcharlotte]

@tdoublep there's no H100 in the test queue. so pretty much most attention tests that are relevant have not been running. @hl475 and @simon-mo got a new one just set up this week.

[seindum]

This seems to be covered already by #26597.

[yeqcharlotte]

LGTM

[tdoublep]

there's no H100 in the test queue. so pretty much most attention tests that are relevant have not been running.

How come we can't run the attention tests on L4 where the other tests run?

[tdoublep]

LGTM (needs follow-up to enable test to run in CI)

[tdoublep]

Hmm CI failure suggests we are trying to run this test on CPU now?? Looking at the test job definition, I don't understand why the attention test is running?

[yeqcharlotte]

cpu tests pick up attention tests after this pr :(

[tdoublep]

@yeqcharlotte Do you understand how that can be happening? I'm a bit baffled tbh

[tdoublep]

It's like the CI job is trying to execute commands that are different to what is checked into the branch. It's really weird. I tried to create a clean branch (added you as co-author @hl475) with this change and triggered the failing job in CI to see if it is reproducible.

[tdoublep]

So after some investigation it looks like we are now generating the test pipeline automatically based on the files that have changed vllm-project/ci-infra#184

This PR changes a single test that should be run on GPU, and it trying to run it in the CPU jobs.

[rzabarazesh]

I have been investigating this. This isn't really a test filtering issue. The main problem is that these tests are orphaned and not being run anywhere in the first place. Test filtering is doing a "best guess" but ends up putting it in the wrong test group.

[rzabarazesh]

I fixed the CI issue above is now resolved in vllm-project/ci-infra#194.
As far as the signals are concerned this test is still orphaned however

[tdoublep]

LGTM

[tdoublep]

@rzabarazesh Thank you for the investigation and fix!

@yeqcharlotte It looks like we can't merge until you approve since you requested changes.