Add GDN benchmarking - #501
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe benchmark image now installs FlashQLA and its runtime dependencies. The linear-attention benchmark supports FlashQLA for GDN workloads. A new CLI utility generates per-batch performance charts from benchmark CSV files. ChangesFlashQLA benchmark integration
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant BenchmarkCLI
participant run_benchmark
participant flash_qla
BenchmarkCLI->>run_benchmark: select and validate flash_qla
run_benchmark->>flash_qla: pass dense GDN tensors
flash_qla-->>run_benchmark: return attention output
run_benchmark-->>BenchmarkCLI: preserve output for validation
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
benchmark/linear_attention/plot_results.py (1)
56-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse explicit optional annotations for
cudnn_versionandbatch_sizes. ImportOptional, then useOptional[str]for bothcudnn_versionparameters andOptional[list]forbatch_sizes.🤖 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 `@benchmark/linear_attention/plot_results.py` around lines 56 - 63, Update the type annotations in get_backend_display_name and generate_charts to use explicit Optional types: import Optional, annotate both cudnn_version parameters as Optional[str], and annotate generate_charts’ batch_sizes as Optional[list].Source: Linters/SAST tools
🤖 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 `@benchmark/linear_attention/Dockerfile`:
- Around line 22-27: Pin the benchmark environment in the Dockerfile around the
FlashQLA installation: check out a specific FlashQLA tag or commit, require
flash-linear-attention==0.5.0, and pin cudnn-frontend to a specific source
revision. Add and apply a constraints or lock file covering these packages and
all transitive dependencies, including tilelang==0.1.9 and
apache-tvm-ffi==0.1.9, while recording the tested versions.
In `@benchmark/linear_attention/plot_results.py`:
- Around line 102-103: Update the chart title in the plotting flow around
fig.suptitle to use the selected variant value instead of hardcoding “GDN”,
while preserving the existing batch, heads, head dimension, and GPU information.
- Around line 19-24: Add pandas, matplotlib, and seaborn to the dependencies
installed by benchmark/linear_attention/Dockerfile so the imports in
plot_results.py resolve successfully when the linear-attention image runs.
---
Nitpick comments:
In `@benchmark/linear_attention/plot_results.py`:
- Around line 56-63: Update the type annotations in get_backend_display_name and
generate_charts to use explicit Optional types: import Optional, annotate both
cudnn_version parameters as Optional[str], and annotate generate_charts’
batch_sizes as Optional[list].
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 15da04bd-6e11-468b-9b7e-bb3d19890418
⛔ Files ignored due to path filters (2)
benchmark/linear_attention/results/gdn/b300/gdn.pngis excluded by!**/*.pngbenchmark/linear_attention/results/gdn/b300/gdn_20260806.csvis excluded by!**/*.csv
📒 Files selected for processing (4)
benchmark/linear_attention/Dockerfilebenchmark/linear_attention/README.mdbenchmark/linear_attention/benchmark_single_linear_attention.pybenchmark/linear_attention/plot_results.py
| # Install the Cutlass DSL runtime (cuDNN FROST engines) and FLA. | ||
| RUN pip install nvidia-cutlass-dsl[cu13]==4.7.0 apache-tvm-ffi flash-linear-attention | ||
|
|
||
| # Install FlashQLA from source. | ||
| RUN git clone https://github.com/QwenLM/FlashQLA.git | ||
| RUN pip install -v /workspace/FlashQLA |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- Dockerfile ---'
cat -n benchmark/linear_attention/Dockerfile
printf '%s\n' '--- dependency and benchmark references ---'
rg -n -i --glob '!node_modules' --glob '!dist' \
'FlashQLA|flash-linear-attention|apache-tvm-ffi|nvidia-cutlass-dsl|constraints|requirements|lock' \
benchmark README* .github 2>/dev/null || true
printf '%s\n' '--- tracked files near the benchmark ---'
git ls-files benchmark/linear_attentionRepository: NVIDIA/cudnn-frontend
Length of output: 36203
🌐 Web query:
QwenLM FlashQLA GitHub README benchmark instructions pinned dependencies release notes API alignment
💡 Result:
FlashQLA is a high-performance linear attention library by QwenLM designed to accelerate Gated Delta Networks (GDN) [1]. Benchmark Instructions To benchmark FlashQLA, you must have the required dependencies installed, including flash-linear-attention (FLA) and flashinfer [2][3]. The benchmark scripts are located in the benchmark directory [2][4]. For example, to run the gated delta rule benchmark, ensure the correct versions are installed and run the script [2][4]: pip install flash_linear_attention==0.5.0 flashinfer-python==0.6.13 python benchmark/bench_gated_delta_rule.py Detailed benchmark results for H200 (SM90) and GB200 (SM100) are available in the repository's benchmark directory [2][5]. Pinned Dependencies The project enforces specific dependency versions for compatibility [1]. As of version 0.1.2, these include: - torch: >=2.8 [1] - tilelang: ==0.1.9 [1] - apache-tvm-ffi: ==0.1.9 [1] - flash_linear_attention: ==0.5.0 (for comparison/API alignment) [2][3] Release Notes - v0.1.2 (July 2026): Added forward pass support for SM120 (Blackwell) architectures and established the library as a backend for the flash-linear-attention GDN, enabling plug-and-play acceleration [1]. - v0.1.1 (June 2026): Introduced intra-card sequence parallelism for the backward pass and added SM100 architecture support [5][1]. This release also upgraded the tilelang dependency to v0.1.9 and aligned entry function signatures with the latest flash-linear-attention interface [5][6][1]. API Alignment FlashQLA is aligned with the flash-linear-attention (FLA) API, allowing it to serve as a high-performance backend [1]. The high-level API is accessed via chunk_gated_delta_rule, and the low-level API provides separate forward (chunk_gated_delta_rule_fwd) and backward (chunk_gated_delta_rule_bwd) functions, mirroring standard interfaces used in linear attention research [1].
Citations:
- 1: https://pypi.org/project/flash-qla/0.1.2/
- 2: https://github.com/QwenLM/FlashQLA/blob/main/README.md
- 3: https://github.com/QwenLM/FlashQLA
- 4: https://context7.com/qwenlm/flashqla/llms.txt
- 5: https://github.com/QwenLM/FlashQLA/releases/tag/v0.1.1
- 6: https://github.com/QwenLM/FlashQLA/releases
🏁 Script executed:
#!/bin/bash
set -eu
cat -n benchmark/linear_attention/Dockerfile
rg -n -i 'FlashQLA|flash-linear-attention|apache-tvm-ffi|nvidia-cutlass-dsl|constraints|requirements|lock' \
benchmark README* .github 2>/dev/null || true
git ls-files benchmark/linear_attentionRepository: NVIDIA/cudnn-frontend
Length of output: 36099
🌐 Web query:
QwenLM FlashQLA GitHub README benchmark pinned dependencies release notes API alignment
💡 Result:
FlashQLA is a high-performance linear attention kernel library developed by QwenLM that leverages the TileLang compiler to optimize GDN (Gated Delta Network) Chunked Prefill operations [1][2]. Key technical aspects include: Benchmarks: The library demonstrates significant performance improvements over the FLA (Flash Linear Attention) Triton kernel, achieving 2-3x speedup in forward passes and 2x speedup in backward passes on NVIDIA Hopper (SM90) and Blackwell (SM100) architectures [2][3]. Detailed benchmark results for H200 and GB200 systems are provided in the repository's benchmark directory [3]. Pinned Dependencies: The project maintains strict dependency requirements to ensure compatibility with its TileLang-based kernels. According to setup.py and README documentation, the library requires: - Python >= 3.10 [4] - PyTorch >= 2.8 [2][4] - TileLang == 0.1.9 [4] - Apache-tvm-ffi == 0.1.9 [4] - CUDA 12.8 or above [2] Release Notes & Alignment: Recent updates, such as v0.1.1, introduced support for the SM100 architecture and the state_v_first parameter [5][3]. A key focus in recent releases has been aligning the entry function signatures with the latest interfaces from the flash-linear-attention library, ensuring better interoperability [5][3]. API: FlashQLA provides both high-level and low-level APIs [1][6]: - High-level: Access via chunk_gated_delta_rule, which handles the core forward/backward computations and state management [1]. - Low-level: Access via separate chunk_gated_delta_rule_fwd and chunk_gated_delta_rule_bwd functions for fine-grained control over forward and backward passes [1]. The project is released under the MIT License [1][4].
Citations:
- 1: https://github.com/QwenLM/FlashQLA
- 2: https://github.com/QwenLM/FlashQLA/blob/main/README.md
- 3: https://github.com/QwenLM/FlashQLA/releases/tag/v0.1.1
- 4: https://github.com/QwenLM/FlashQLA/blob/main/setup.py
- 5: https://github.com/QwenLM/FlashQLA/releases
- 6: https://github.com/qwenlm/flashqla
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import json
import urllib.request
urls = {
"tags": "https://api.github.com/repos/QwenLM/FlashQLA/tags?per_page=20",
"setup": "https://raw.githubusercontent.com/QwenLM/FlashQLA/main/setup.py",
"readme": "https://raw.githubusercontent.com/QwenLM/FlashQLA/main/README.md",
"releases": "https://api.github.com/repos/QwenLM/FlashQLA/releases?per_page=10",
}
for name, url in urls.items():
print(f"--- {name} ---")
request = urllib.request.Request(url, headers={"Accept": "application/vnd.github+json"})
with urllib.request.urlopen(request) as response:
data = response.read().decode()
if name in {"tags", "releases"}:
for item in json.loads(data):
if name == "tags":
print(item["name"], item["commit"]["sha"])
else:
print(item["tag_name"], item["target_commitish"])
else:
for number, line in enumerate(data.splitlines(), 1):
if any(term in line.lower() for term in (
"install_requires", "apache-tvm-ffi",
"flash_linear_attention", "tilelang",
"pip install", "benchmark", "0.1.",
)):
print(f"{number}: {line}")
PY
cat -n benchmark/linear_attention/README.md | sed -n '45,110p'Repository: NVIDIA/cudnn-frontend
Length of output: 2459
🏁 Script executed:
#!/bin/bash
set -eu
fetch() {
curl -k -fsSL --retry 2 "$1"
}
printf '%s\n' '--- tags ---'
fetch 'https://api.github.com/repos/QwenLM/FlashQLA/tags?per_page=20' |
jq -r '.[] | [.name, .commit.sha] | `@tsv`'
printf '%s\n' '--- releases ---'
fetch 'https://api.github.com/repos/QwenLM/FlashQLA/releases?per_page=10' |
jq -r '.[] | [.tag_name, .target_commitish] | `@tsv`'
printf '%s\n' '--- setup.py dependency lines ---'
fetch 'https://raw.githubusercontent.com/QwenLM/FlashQLA/main/setup.py' |
nl -ba |
grep -Ei 'apache-tvm-ffi|flash[_-]linear[_-]attention|tilelang|install_requires|version'
printf '%s\n' '--- README install and benchmark lines ---'
fetch 'https://raw.githubusercontent.com/QwenLM/FlashQLA/main/README.md' |
nl -ba |
grep -Ei -C 3 'apache-tvm-ffi|flash[_-]linear[_-]attention|tilelang|benchmark|pip install'Repository: NVIDIA/cudnn-frontend
Length of output: 406
🏁 Script executed:
#!/bin/bash
set -eu
fetch() {
curl -k -fsSL --retry 2 "$1"
}
for ref in v0.1.2 main; do
printf '%s\n' "--- setup.py at $ref ---"
fetch "https://raw.githubusercontent.com/QwenLM/FlashQLA/$ref/setup.py" |
grep -Ei -C 3 'apache-tvm-ffi|flash[_-]linear[_-]attention|tilelang|install_requires|version' || true
printf '%s\n' "--- README at $ref ---"
fetch "https://raw.githubusercontent.com/QwenLM/FlashQLA/$ref/README.md" |
grep -Ei -C 3 'apache-tvm-ffi|flash[_-]linear[_-]attention|tilelang|benchmark|pip install' || true
doneRepository: NVIDIA/cudnn-frontend
Length of output: 8408
Pin the complete benchmark dependency set.
Pin the FlashQLA tag or commit, flash-linear-attention==0.5.0, and all transitive dependencies. FlashQLA v0.1.2 requires tilelang==0.1.9 and apache-tvm-ffi==0.1.9. Also pin the cudnn-frontend source clone. Record the tested versions in a constraints or lock file to keep benchmark results comparable across rebuilds.
🤖 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 `@benchmark/linear_attention/Dockerfile` around lines 22 - 27, Pin the
benchmark environment in the Dockerfile around the FlashQLA installation: check
out a specific FlashQLA tag or commit, require flash-linear-attention==0.5.0,
and pin cudnn-frontend to a specific source revision. Add and apply a
constraints or lock file covering these packages and all transitive
dependencies, including tilelang==0.1.9 and apache-tvm-ffi==0.1.9, while
recording the tested versions.
Source: MCP tools
Anerudhan
added
cat-feature
3 participants
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*(see label list).Affected area
Summary
Add GDN benchmark results and FlashQLA as a backend for benchmarking.
Why
Related issues
API and compatibility impact
Testing
Summary by CodeRabbit