[GDN2] Add Hopper SM90 CuTe DSL prefill - #113
Conversation
8d39d85 to
49dc389
Compare
Aharrypotter
changed the title
Shape-driven dispatch on top of the production GDN2PrefillKernel: - N==1, T<=64: exact released preparation/commit schedule - N==1, Hv=16, init+final, T>64: V64 single State-WG with register-resident final-tail carry - all other shapes: unchanged V128 production path (compile-time branches, defaults produce instruction-identical kernel) Compile-cache key extended by the two derived booleans. Public API, state layout, and supported-shape matrix unchanged. Validation (H20, source-bound): - 13 pytest passed incl. explicit t64/t65 boundary cases - 100k-launch determinism, 4-tool sanitizer, codegen zero stack/local/spill all PASS - targeted S2/Q2 rows 0.9362x vs released incumbent (6.4% faster); full S1-S5 all faster than pinned FLA
Stop requiring nvidia-cutlass-dsl==4.5.1. Availability and runtime checks now accept any installed version (matching the KDA backend policy); the version is still recorded in stress reports for diagnostics. Verified: 13 pytest passed on H20 with 4.5.1, and is_available() no longer depends on the exact version constant.
icavan
left a comment
icavan
left a comment
There was a problem hiding this comment.
Thanks for the unusually thorough validation and performance evidence. I found one numerical-correctness issue that should be resolved before merging, plus two contract/compatibility inconsistencies:
- The public contract accepts any finite non-positive
g, but the transformed chunk formulation can overflow for valid inputs even when the direct recurrence remains finite. This needs either a numerically stable implementation or an explicit, validated input bound with adversarial tests. - The documented compile-cache identity does not match the actual key, which also changes across
N == 1/T <= 64dispatch boundaries. Given the roughly 20-second compilation cost reported here, this is observable API behavior rather than a cosmetic documentation issue. - CuTe DSL compatibility is inconsistent: runtime availability now accepts any installed version, the API document still says
==4.5.1, and all supplied GPU evidence is from 4.5.1. Please define and validate one supported range.
The PR also currently conflicts with the latest main in README.md, so it will need a rebase after these issues are addressed.
| raw_stage, | ||
| ], | ||
| ) | ||
| * cute.math.exp2( |
There was a problem hiding this comment.
[P1] Please avoid forming this unbounded inverse-decay factor, or tighten and enforce the public input contract. The wrapper currently accepts every finite g <= 0, but a valid 64-token chunk with g = -2 reaches a prefix of -128; exp(-prefix) exceeds the FP32 range once -prefix > ~88.72 (around token 45), while the direct tokenwise recurrence remains finite because it only multiplies by exp(g) <= 1. The current tests use only g in [-0.05, 0], so they cannot expose this. Please add adversarial decay tests (for example g = -1/-2 and gate endpoints) and either rescale/factor the algebra so intermediates stay bounded or reject values outside a documented safe range.
| and inputs.total_tokens > 64 | ||
| ) | ||
| retain_final_tail = store_final_state and not (inputs.num_sequences == 1 and inputs.total_tokens <= 64) | ||
| key = ( |
There was a problem hiding this comment.
[P2] This cache key contradicts the stated (device, Hv, has_initial_state, store_final_state) identity and the claim that T and N remain fully dynamic. use_n1_hv16_v64 and retain_final_tail are derived from N and the T <= 64 boundary, so moving between N=1,T<=64, N=1,T>64, and N>1 can trigger additional ~20-second compilations for the same four documented fields. Please update the PR description/docs and add a cache-boundary test, or redesign the dispatch so the advertised cache contract is true.
| properties = torch.cuda.get_device_properties(device) | ||
| if (properties.major, properties.minor) != (9, 0): | ||
| return False | ||
| return _installed_cutlass_dsl_version() is not None |
There was a problem hiding this comment.
[P2] Treating every installed CuTe DSL version as available is broader than the evidence and documentation. The API document still requires nvidia-cutlass-dsl==4.5.1, all validation in this PR was performed on 4.5.1, and the current project dependency has its own bounded/excluded range. Please define one supported version range, use it consistently in dependency metadata, docs, is_sm90_gdn2_available(), and the runtime error, and validate the relevant boundary versions before reporting the backend as available.
2 participants
📌 Description
This PR adds the first production Gated DeltaNet-2 (GDN2) prefill backend for
NVIDIA Hopper SM90a, implemented as a fully fused CuTe DSL kernel.
cula.gdn2.chunk_gdn2with direct architecture dispatch and noFLA, Triton, C++, or environment-selected fallback;
recurrent state, MHA, GVA2, and GVA4 head mappings;
longest-processing-time sequence-to-CTA ordering;
while specializing the compile cache only by device, value-head count, and
initial/final-state modes;
boundary correctness tests, a 100,000-launch deterministic stress test,
four-tool Compute Sanitizer coverage, benchmarks, and documentation.
Across the frozen five-row H20 matrix, the fused kernel is faster than the
pinned FLA GDN2 Triton public path on every row:
2.379xequal-weightgeometric-mean speedup and
1.375xminimum row speedup.Supported SM90 contract
Hq=16Hv={16,32,64}1 <= N <= 32, each sequence non-emptyK=V=128[N,Hv,V,K]Kernel and dispatch
(device, Hv, has_initial_state, store_final_state)whileTandNremain dynamic;sm90a_cutedsl_gdn2_prefill_v1;🔍 Related issue and scope
main; it does not depend on PR [GDN] Add Hopper SM90 CuTe DSL prefill #108.cula.gdn2/cula.ops.gdn2and has nosource or import dependency on the existing GDN-v1 kernel.
b3a1ed0bd304970b515ae7004765f0387052be0f.d1ce07369d581813553f30a750af3b6b5f9af6a9.SM100 validation remain a separate follow-up; this PR does not close the
issue.
🧪 Final-source validation
Validation environment
Validation summary
13 passed in 208.10sThe allocating API path can issue one
cudaMallocon an allocator cache miss.The caller-preallocated path avoids that allocation. Setting
validate_inputs=Trueintentionally enables diagnostic device-contentvalidation and may synchronize; the default is
False.Accuracy vs independent PyTorch reference
The independent reference is a pure-PyTorch, tokenwise implementation with no
cuLA imports. Every case requires finite reference and product outputs, BF16
output agreement with
rtol=atol=0.01, and FP32 final-state agreement withrtol=0.001,atol=0.005.mha-single-token[1]mha-tail-and-init[65, 1]mha-initial-no-final[65, 63]mha-production-t1024[1024]mha-max-sequences[1] × 32gva2-packed-tails[1, 63, 65, 2]gva4-init[4]Each case also checks repeat-bitwise-exact product output/state, input
immutability, caller-provided output/state identity, and output/state
redzones. Unsupported
Hq,Hv, andN=33metadata is rejected beforecompilation.
The tokenwise PyTorch implementation is a correctness oracle, not a meaningful
fused-kernel performance denominator. Performance is compared with the pinned
FLA Triton public path.
Stability and Compute Sanitizer
One H20 process completed 100,000 round-robin launches with fixed per-case
inputs and initial states. Every launch checked bitwise output/state equality
against its baseline and finite values. The run completed with zero output or
state mismatches, unchanged inputs and redzones, and no host synchronization
inside the launch loop.
100,000-launch deterministic stress matrix (6 rows)
S1-MHA-T64[64]S2-MHA-T1024[1024]S3-MHA-PACKED-T40967–694tokens)N32-MHA-IRREGULAR[1, 63, 64, 65] × 8GVA2-PACKED[1, 63, 65, 2]GVA4-PACKED[65, 1, 129, 63]The same six-row product matrix was exercised under every applicable NVIDIA
Compute Sanitizer tool:
⚡ Performance vs pinned FLA Triton
Comparator and claim boundary
The release claim is the exact five-row H20 matrix below against the pinned
FLA GDN2 Triton public-logical-call denominator. It is not an arbitrary-shape,
other-GPU, or other-FLA-revision claim.
For GVA2/GVA4, FLA requires Q/K/G/B head expansion. That expansion remains
inside FLA's timed public logical call; it is not moved into setup. Both
implementations consume byte-identical canonical inputs.
The results below come from a fresh paired campaign bound to this independent
GDN2 source manifest. The campaign does not inherit latency receipts from the
previous stacked branch.
Paired methodology
environment, FLA commit/tree, and benchmark protocol before timing.
(row, implementation, replica)in a fresh process with uniqueCUDA, CuTe DSL, TorchInductor, Triton, and XDG caches.
replica 1.
timing.
receipt.
median of 3 process observations as the row point estimate.
seed 6606, and a one-sided 95% upper bound for
product / FLA.source/backend identity, hardware identity, cache uniqueness, pair order,
and contention before accepting a receipt.
The fresh paired campaign accepted all
5 rows x 2 implementations x 3 processes = 30receipts, used 120 unique required cache directories, and required no escalation or requeue.Canonical five-row matrix
All rows use
Hq=16,K=V=128, BF16 Q/K/V/B/W, FP32 G/state, and apublic
[N,Hv,V,K]state layout.[64][1024][1024][1024]Exact S3 lengths:
Performance gates
< 1.00.7271< 1.00.7383< 1.00.428930/3030/3036/36byte-identicalFinal latency results
Every row is faster than pinned FLA at the point estimate. The smallest row speedup is
1.375x; the equal-weight geometric-mean speedup is2.379x.Short-sequence and N=1 V64 specialization
An additive, shape-driven dispatch extension improves the highest-impact
single-sequence routes without changing the public API, state layout, or any
other supported shape:
N == 1, T <= 64: exact released preparation/commit schedule;N == 1, Hv == 16, initial + final state, T > 64: V64 single State-WG withregister-resident final-tail carry;
Paired G3 qualification on the same H20 and frozen matrix measured
candidate/incumbent:
Other rows stay within ±2.4%(S1 short no-state guardrail is
1.0231, insidethe
1.03limit). The full five-row matrix remains faster than pinned FLA asshown above.
Workload-group summary
These groups are derived summaries of the frozen rows, not additional release gates.
Per-process steady-state observations (30 accepted receipts)
Each value below is the arithmetic mean of 100 CUDA-event samples. The row median in the previous table is the median of these three independent process observations.
Spread is
max(process average) / min(process average) - 1.First-call setup and compilation cost
First-call setup/compile is measured in each fresh process but excluded from all steady-state latency and speedup claims above.
Schedule-selection bakeoff (38 cases, 114,000 CUDA-event samples)
The selected load-balanced longest-processing-time sequence-to-CTA
ordering was compared with a source-identical submission-order
baseline. Only sequence-to-CTA ordering changed.
N={1,2,4,8,12,13,20,32}< 1.03N4-T193-DESC-GVA4-I0F0Kernel resources and source-bound codegen
All six product specializations compile to one active CTA per SM with no
stack, local-memory, or spill traffic. The independent source reproduced all
36 captured MLIR, PTX, JIT/reassembled cubin, and JIT/reassembled SASS
artifacts byte-for-byte.
h16-init-final-producth16-init-nofinal-producth16-noinit-final-producth16-noinit-nofinal-producth32-noinit-final-producth64-init-final-productThe normalized source-bound instruction counts are identical across all six
specializations:
Benchmark reproduction
From the repository root:
Use
--list-matrixto inspect the exact five rows without launching CUDAwork. The standalone benchmark is a developer diagnostic; the final release
claim additionally requires the fresh-process, unique-cache, alternating-order,
source/input identity, and replay audits described above.
🚀 Pull Request Checklist
variable-length input covered.
guards and documented dtype tolerances.
pass.
with 30 fresh-process receipts.
steady-state latency.
👀 Reviewer Notes
Suggested review focus:
[N,Hv,V,K]public-state layout.denominator.
addresses [Feature] Add GDN2 support #112; it does not claim SM100 completion.
Known release limits:
Hqother than 16 orHvoutside{16,32,64};