[KDA] Add SM90 fused WY-DqKG backward#109
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces SM90 (Hopper) support for the modular chunk KDA forward and backward kernels, integrating Triton-based adapters and fallback paths alongside the existing SM100 (Blackwell) implementations. It also adds benchmarking scripts, comprehensive numerical correctness tests, and updates the test configuration to properly target SM90 and SM100 devices. The review feedback suggests grouping multiple imports from the same module in tests/test_chunk_wy_dqkg_sm90.py into a single import statement to improve readability.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | ||
| BK as BENCHMARK_BK, | ||
| ) | ||
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | ||
| BT as BENCHMARK_BT, | ||
| ) | ||
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | ||
| BV as BENCHMARK_BV, | ||
| ) | ||
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | ||
| MIN_OCC as BENCHMARK_MIN_OCC, | ||
| ) | ||
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | ||
| SEED as BENCHMARK_SEED, | ||
| ) | ||
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | ||
| K as BENCHMARK_K, | ||
| ) | ||
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | ||
| V as BENCHMARK_V, | ||
| ) |
There was a problem hiding this comment.
Multiple import statements from the same module benchmarks.bench_kda_bwd_wy_dqkg_sm90 can be grouped into a single import statement. This improves readability and adheres to PEP 8 guidelines.
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| BK as BENCHMARK_BK, | |
| ) | |
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| BT as BENCHMARK_BT, | |
| ) | |
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| BV as BENCHMARK_BV, | |
| ) | |
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| MIN_OCC as BENCHMARK_MIN_OCC, | |
| ) | |
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| SEED as BENCHMARK_SEED, | |
| ) | |
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| K as BENCHMARK_K, | |
| ) | |
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| V as BENCHMARK_V, | |
| ) | |
| from benchmarks.bench_kda_bwd_wy_dqkg_sm90 import ( | |
| BK as BENCHMARK_BK, | |
| BT as BENCHMARK_BT, | |
| BV as BENCHMARK_BV, | |
| MIN_OCC as BENCHMARK_MIN_OCC, | |
| SEED as BENCHMARK_SEED, | |
| K as BENCHMARK_K, | |
| V as BENCHMARK_V, | |
| ) |
References
- Imports should be grouped and placed at the top of the file, and multiple imports from the same module should be combined. (link)
568789c to
e07a1a2
Compare
e07a1a2 to
9ef302f
Compare
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
1 participant
📌 Description
dq,dk,dv,db,dg, anddAin the same order as FLA.float32/bfloat16beta, and selectableBK/BVtile sizes and occupancy.cula.kda.chunk_kdaautograd path on SM90.This ports the SM90 WY-DqKG implementation from the internal
biz/v0.2line into a focused PR based on the public GitHubmain, without carrying unrelated SM100 or build-system changes.🔍 Related Issues
N/A
🚀 Pull Request Checklist
Thank you for contributing to cuLA! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.pre-commitwas run successfully on every file changed by this PR. The repository-wide command was intentionally not marked complete because existing unrelated and ignored source trees are not fully Ruff-formatted.🧪 Tests
Validation completed locally:
The numerical tests compare all six outputs with the FLA v0.5.0 Triton reference. The benchmark-sized determinism suite is bitwise exact across repeated runs. The CUDA runtime reports compute capability
(9, 0);nvidia-smilabels the provisioned deviceNVIDIA L20X.pip install -e . --no-build-isolationwas also attempted. The extension build is blocked by the host toolchain mismatch (system NVCC 13.0 versus PyTorch cu129), independent of this Python/CuTe change.⚡ Performance
Benchmark target:
The benchmark covers fixed-length
B={1,2}, sequence lengths from 256 through 32768, and uniform/random/skewed variable-length workloads. It reports FLA v0.5.0 latency, cuLA SM90 latency, speedup, and accuracy fordq,dk,dv,db,dg, anddA.Measured with
H=32,K=V=128,BT=64,BK=32,BV=64, occupancy 2, bf16, 25 warmup iterations, and 100 timed iterations. The CUDA runtime reports SM90 capability(9, 0);nvidia-smilabels the provisioned deviceNVIDIA L20X.Fixed-length results
Variable-length results
For every row, the two accuracy columns are the maxima across
dq,dk,dv,db,dg, anddA; both maxima came fromdg. Values are rounded to the six decimal places printed by the benchmark.Reviewer Notes
chunk_fwd.py,chunk_intra.py, andchunk_bwd.py; SM100/SM103 behavior should remain unchanged.