Fix batch prefill example script for ragged kv cache#73
Merged
demandal25 merged 2 commits intoROCm:amd-integrationfrom Nov 21, 2025
Merged
Fix batch prefill example script for ragged kv cache#73demandal25 merged 2 commits intoROCm:amd-integrationfrom
demandal25 merged 2 commits intoROCm:amd-integrationfrom
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR fixes the batch prefill example script to properly support ragged KV cache functionality. The changes remove naive attention reference comparisons, add support for non-contiguous KV cache testing, improve LSE (log-sum-exp) output verification, and update test configurations.
Key changes:
- Added
contiguous_kvparameter to test non-contiguous memory layouts for KV cache - Enhanced verification to include LSE output comparison when
return_lse=True - Removed naive attention comparisons and unused backend parameter
- Updated test cases with new parameter values
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| tests/attention_reference.py | Added clarifying comments about NHD layout for K and V tensor parameters |
| examples/batch_prefill_example.py | Refactored batch prefill functions to support non-contiguous KV cache, added LSE verification, removed naive attention comparisons, and updated test configurations |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Basic ragged KV cache test with causal masking | ||
| batch_prefill_with_ragged_kv_cache_example( | ||
| 12, 54, 37, 8, 8, 128, True, "NONE", 0.0, False | ||
| 12, 54, 37, 8, 8, 64, False, "NONE", 0.0, False |
There was a problem hiding this comment.
The comment on line 427 states 'Basic ragged KV cache test with causal masking' but the function call passes causal=False. Either update the comment to match the actual parameter or change the parameter to True.
Suggested change
| 12, 54, 37, 8, 8, 64, False, "NONE", 0.0, False | |
| 12, 54, 37, 8, 8, 64, True, "NONE", 0.0, False |
diptorupd
pushed a commit
that referenced
this pull request
Dec 5, 2025
Upgrade both the base and CI dockers in `docker/Dockerfile.rocm_ci` to the new version: - rocm6.4 - ubuntu24.04 - py3.12 - torch2.7.1
diptorupd
pushed a commit
that referenced
this pull request
Dec 5, 2025
This PR fixes the batch prefill example script for ragged kv cache. The examples for batch prefill with ragged kv cache in the script are now passing. This reinforces the basic correctness of the PR #50 More exhaustive tests are in the pytest script for batch prefill and should pass to call it a full victory.
zhenhantech
pushed a commit
to zhenhantech/flashinfer
that referenced
this pull request
Jan 9, 2026
Upgrade both the base and CI dockers in `docker/Dockerfile.rocm_ci` to the new version: - rocm6.4 - ubuntu24.04 - py3.12 - torch2.7.1
zhenhantech
pushed a commit
to zhenhantech/flashinfer
that referenced
this pull request
Jan 9, 2026
This PR fixes the batch prefill example script for ragged kv cache. The examples for batch prefill with ragged kv cache in the script are now passing. This reinforces the basic correctness of the PR ROCm#50 More exhaustive tests are in the pytest script for batch prefill and should pass to call it a full victory.
diptorupd
pushed a commit
to diptorupd/flashinfer
that referenced
this pull request
Jan 28, 2026
Upgrade both the base and CI dockers in `docker/Dockerfile.rocm_ci` to the new version: - rocm6.4 - ubuntu24.04 - py3.12 - torch2.7.1
diptorupd
pushed a commit
to diptorupd/flashinfer
that referenced
this pull request
Jan 28, 2026
This PR fixes the batch prefill example script for ragged kv cache. The examples for batch prefill with ragged kv cache in the script are now passing. This reinforces the basic correctness of the PR ROCm#50 More exhaustive tests are in the pytest script for batch prefill and should pass to call it a full victory.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR fixes the batch prefill example script for ragged kv cache. The examples for batch prefill with ragged kv cache in the script are now passing. This reinforces the basic correctness of the PR #50
More exhaustive tests are in the pytest script for batch prefill and should pass to call it a full victory.