Repository navigation
Pipeline long ragged GQA4 KV loads - #114
Merged
Merged
Conversation
📝 WalkthroughWalkthroughThe PR changes tiled GQA4 prefill to use separate asynchronous K and V tiles. It adds ragged performance evidence and updates repository and website reporting with refreshed latency results, qualification scope, and evidence links. ChangesGQA4 prefill optimization and qualification
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PrefillWrapper
participant TiledGQA4Kernel
participant SharedMemoryTiles
participant AttentionMMA
PrefillWrapper->>TiledGQA4Kernel: pass separate K and V tile pointers
TiledGQA4Kernel->>SharedMemoryTiles: issue asynchronous K and V copies
SharedMemoryTiles-->>TiledGQA4Kernel: expose K tile before V completion
TiledGQA4Kernel->>AttentionMMA: compute QK from K tile
SharedMemoryTiles-->>TiledGQA4Kernel: complete V tile copy
TiledGQA4Kernel->>AttentionMMA: compute PV from V tile
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
feichai0017
marked this pull request as ready for review
August 12, 2026 09:41
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/results/README.md`:
- Around line 83-84: Update the benchmark text in the dual-tile results section
to distinguish the cohorts: either state the 36.94 µs provider-comparison result
before the 1.68x FlashInfer gap, or change the displayed gap to 1.69x so it
matches the shown 37.04 µs and 21.93 µs values.
In `@website/src/data/site.ts`:
- Around line 123-125: Use the optimized-candidate measurement consistently for
the ragged-GQA4 row: update website/src/data/site.ts lines 123-125 and README.md
line 127 from 36.94 µs and “FlashInfer 1.68× lower” to 37.04 µs and “FlashInfer
1.69× lower”; no other measurement rows require changes.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 5039a73a-4f5c-4678-80d7-ad2f7e8ce64a
📒 Files selected for processing (11)
README.mdcrates/oxide-infer-cuda/README.mdcrates/oxide-infer-cuda/src/attention/prefill.rsdocs/flashinfer-parity.mddocs/operator-catalog.mddocs/results/README.mddocs/results/h20-flashinfer-v0.6.17-ragged-prefill-dual-tile-gqa4-eager-performance-f9b95b0-20260812.jsondocs/roadmap.mdwebsite/src/data/site.tswebsite/src/pages/evidence.astrowebsite/src/pages/index.astro
Comment on lines
+83
to
+84
| dual-tile path reduces the parent median from 39.42 to 37.04 microseconds. | ||
| FlashInfer remains lower at 21.93 microseconds, a 1.68x gap. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Separate the two candidate measurement cohorts.
Line 83 reports the 37.04 µs parent/candidate cohort. The 1.68x gap at Line 84 uses the separate 36.94 µs provider-comparison cohort. The displayed 37.04 µs and 21.93 µs values produce a 1.69x ratio.
State the 36.94 µs provider-comparison result before the 1.68x claim, or change the displayed gap to 1.69x.
Proposed correction
- dual-tile path reduces the parent median from 39.42 to 37.04 microseconds.
- FlashInfer remains lower at 21.93 microseconds, a 1.68x gap.
+ dual-tile path reduces the parent median from 39.42 to 37.04 microseconds.
+ In the separate provider-comparison cohort, Oxide measures 36.94 microseconds.
+ FlashInfer remains lower at 21.93 microseconds, a 1.68x gap.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| dual-tile path reduces the parent median from 39.42 to 37.04 microseconds. | |
| FlashInfer remains lower at 21.93 microseconds, a 1.68x gap. | |
| dual-tile path reduces the parent median from 39.42 to 37.04 microseconds. | |
| In the separate provider-comparison cohort, Oxide measures 36.94 microseconds. | |
| FlashInfer remains lower at 21.93 microseconds, a 1.68x gap. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/results/README.md` around lines 83 - 84, Update the benchmark text in
the dual-tile results section to distinguish the cohorts: either state the 36.94
µs provider-comparison result before the 1.68x FlashInfer gap, or change the
displayed gap to 1.69x so it matches the shown 37.04 µs and 21.93 µs values.
Comment on lines
+123
to
+125
| oxideUs: "36.94", | ||
| flashinferUs: "21.93", | ||
| result: "FlashInfer 1.68× lower", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
files=(
"docs/results/h20-flashinfer-v0.6.17-ragged-prefill-dual-tile-gqa4-eager-performance-f9b95b0-20260812.json"
"docs/results/README.md"
"README.md"
"website/src/data/site.ts"
"website/src/pages/evidence.astro"
)
for file in "${files[@]}"; do
printf '\n== %s ==\n' "$file"
rg -n -C 2 '36\.94|37\.04|21\.93|39\.42|source-bound|oxideUs|FlashInfer' "$file" || true
doneRepository: feichai0017/oxide-infer
Length of output: 11607
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
json=$(fd -t f -g 'h20-flashinfer-v0.6.17-ragged-prefill-dual-tile-gqa4-eager-performance-*.json' | head -n 1)
printf '== candidate record: %s ==\n' "$json"
sed -n '1,80p' "$json"
sed -n '180,230p' "$json"
sed -n '480,550p' "$json"
printf '\n== all relevant public values ==\n'
rg -n -C 3 '36\.94|37\.04|21\.93|39\.42' README.md docs/results/README.md website/src/data/site.ts website/src/pages/evidence.astroRepository: feichai0017/oxide-infer
Length of output: 9864
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
python3 - <<'PY'
import json
from pathlib import Path
paths = list(Path(".").glob("docs/results/h20-flashinfer-v0.6.17-ragged-prefill-dual-tile-gqa4-eager-performance-*.json"))
if len(paths) != 1:
raise SystemExit(f"expected one ragged-GQA4 record, found {len(paths)}")
data = json.loads(paths[0].read_text())
def walk(obj, path=()):
if isinstance(obj, dict):
for key, value in obj.items():
yield from walk(value, path + (str(key),))
elif isinstance(obj, list):
for i, value in enumerate(obj):
yield from walk(value, path + (str(i),))
else:
if isinstance(obj, (int, float)) and ("median" in ".".join(path).lower() or obj in (21.93, 36.94, 37.04, 39.42)):
print(".".join(path), repr(obj))
print("top-level keys:", list(data))
print("median-like numeric leaves:")
walk(data)
PYRepository: feichai0017/oxide-infer
Length of output: 430
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
python3 - <<'PY'
import json
from pathlib import Path
record = next(Path(".").glob(
"docs/results/h20-flashinfer-v0.6.17-ragged-prefill-dual-tile-gqa4-eager-performance-*.json"
))
data = json.loads(record.read_text())
qualification = data["qualification"]
provider_oxide = qualification["providers"]["oxide-infer"]["combined"]["median_us"]
provider_flashinfer = qualification["providers"]["flashinfer"]["combined"]["median_us"]
candidate = data["optimization"]["candidate"]["combined"]["median_us"]
parent = data["optimization"]["parent"]["combined"]["median_us"]
print(f"matched Oxide median: {provider_oxide:.12f} us")
print(f"matched FlashInfer median: {provider_flashinfer:.12f} us")
print(f"optimization candidate median: {candidate:.12f} us")
print(f"optimization parent median: {parent:.12f} us")
print(f"candidate / FlashInfer: {candidate / provider_flashinfer:.12f}x")
print(f"matched Oxide / FlashInfer: {provider_oxide / provider_flashinfer:.12f}x")
print(f"candidate improvement over parent: {(1 - candidate / parent) * 100:.6f}%")
PY
printf '\n== surrounding publication claims ==\n'
sed -n '113,140p' README.md
sed -n '116,130p' website/src/data/site.ts
sed -n '77,87p' docs/results/README.mdRepository: feichai0017/oxide-infer
Length of output: 257
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
python3 - <<'PY'
import json
from pathlib import Path
record = next(Path(".").glob(
"docs/results/h20-flashinfer-v0.6.17-ragged-prefill-dual-tile-gqa4-eager-performance-*.json"
))
data = json.loads(record.read_text())
def summarize(obj, path=(), depth=0):
if depth > 5:
return
if isinstance(obj, dict):
print(".".join(path) or "<root>", "keys=", list(obj.keys()))
for key, value in obj.items():
if key in {"qualification", "optimization", "providers", "combined", "candidate", "parent", "baseline"}:
summarize(value, path + (key,), depth + 1)
elif isinstance(obj, list):
print(".".join(path), f"list[{len(obj)}]")
summarize(data)
def find_key(obj, wanted, path=()):
if isinstance(obj, dict):
for key, value in obj.items():
if key == wanted:
print("FOUND", wanted, "at", ".".join(path + (key,)))
find_key(value, wanted, path + (key,))
elif isinstance(obj, list):
for i, value in enumerate(obj):
find_key(value, wanted, path + (str(i),))
find_key(data, "providers")
find_key(data, "combined")
PYRepository: feichai0017/oxide-infer
Length of output: 1592
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
python3 - <<'PY'
import json
from pathlib import Path
record = next(Path(".").glob(
"docs/results/h20-flashinfer-v0.6.17-ragged-prefill-dual-tile-gqa4-eager-performance-*.json"
))
data = json.loads(record.read_text())
matched_oxide = data["case"]["providers"]["oxide-infer"]["combined"]["median_us"]
flashinfer = data["case"]["providers"]["flashinfer"]["combined"]["median_us"]
candidate = data["optimization"]["candidate"]["combined"]["median_us"]
parent = data["optimization"]["parent"]["combined"]["median_us"]
print(f"matched Oxide median: {matched_oxide:.12f} us")
print(f"matched FlashInfer median: {flashinfer:.12f} us")
print(f"optimization candidate median: {candidate:.12f} us")
print(f"optimization parent median: {parent:.12f} us")
print(f"candidate / FlashInfer: {candidate / flashinfer:.12f}x")
print(f"matched Oxide / FlashInfer: {matched_oxide / flashinfer:.12f}x")
print(f"candidate improvement over parent: {(1 - candidate / parent) * 100:.6f}%")
print("\noptimization publication fields:")
for key in (
"previously_published_record",
"previously_published_oxide_median_us",
"previously_published_flashinfer_median_us",
"previously_published_gap_ratio",
"current_gap_ratio",
"candidate_speedup_over_parent",
"candidate_latency_reduction_percent",
):
print(f"{key}: {data['optimization'].get(key)!r}")
PYRepository: feichai0017/oxide-infer
Length of output: 923
Use one measurement basis for the ragged-GQA4 row.
36.94 µs is the matched provider median. The optimized candidate median is 37.04 µs. Since both public rows describe the optimized path, update website/src/data/site.ts and README.md to 37.04 µs and FlashInfer 1.69× lower. Otherwise, label both rows as matched-provider measurements.
📍 Affects 2 files
website/src/data/site.ts#L123-L125(this comment)README.md#L127-L127
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@website/src/data/site.ts` around lines 123 - 125, Use the optimized-candidate
measurement consistently for the ragged-GQA4 row: update
website/src/data/site.ts lines 123-125 and README.md line 127 from 36.94 µs and
“FlashInfer 1.68× lower” to 37.04 µs and “FlashInfer 1.69× lower”; no other
measurement rows require changes.
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.
Summary
Performance
Exact BF16 ragged GQA4 eager contract: batch 2, query lengths 32/64, KV lengths 256/1024, 16 query heads, 4 KV heads, head dimension 128. Every run uses CPU 20, 200 warmups, 50 samples, and 100 launches per sample.
All four order deltas are below 1.16%, within the 5% acceptance gate. Provider and source-bound comparisons use separate double-order runs, so their candidate medians are recorded independently.
ptxas reports 236 registers, 32 KiB static shared memory, zero stack, zero spill stores, zero spill loads, and zero local memory for the accepted ragged kernel.
Validation
The evidence record includes 400 raw samples, source and runner hashes, fixture parity, resource usage, rejected experiments, and explicit exclusions. It makes no isolated-kernel, Graph-performance, model, serving, other-shape, other-GPU, or overall-library claim.
Summary by CodeRabbit
Performance Improvements
Documentation
Website