[GDN2] Add Hopper SM90 CuTe DSL prefill - #113
Conversation
8d39d85 to
49dc389
Compare
Aharrypotter
changed the title
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.
Fully fused packed variable-length Gated DeltaNet-2 forward prefill for SM90a: one 384-thread CTA per sequence, TMA/WGMMA, and a load-balanced longest-processing-time sequence order. Supports MHA, GVA2, and GVA4 with independently optional initial and final recurrent state. The intra-chunk factorization is blockwise rebased per 16-token sub-block, so the only operand carrying a positive exponent spans at most 15 in-block token gaps. The separable chunk-start alternative is cheaper -- one scaling of each operand serves the whole 64x64 tile -- but it overflows FP32 for inputs the public contract accepts, which is why it is not used. The supported decay contract is g in [-5, 0], matching the pinned FLA GDN2 safe_gate range. Private MLIR dialects are taken through cula.ops._mlir_compat, so this kernel is covered by the repository-wide CuTeDSL contract at import.
Direct SM90a dispatch with no fallback. Tensor metadata is rejected before compilation or launch; validate_inputs=True additionally checks device-side value preconditions, including the [-5, 0] decay bound. Backend availability and the dispatch error share one nvidia-cutlass-dsl range, >=4.5.1,<4.7. The installed version is read through cula.ops._mlir_compat, so this backend and the shared gateway cannot disagree about which toolchain is in use; the upper bound is the gateway's own, and GDN2 only raises the floor to 4.5.1 because 4.4.x lacks cutlass.cute.nvgpu.OperandMajorMode and cannot import the kernel at all. Membership uses standard version ordering, so a local or post release of a supported version stays supported, and pre-release handling is pinned explicitly rather than inherited from the installed packaging, whose default inference has changed between releases.
… coverage An independent tokenwise PyTorch reference with no cuLA imports; production, boundary, and irregular packed shapes across MHA/GVA2/GVA4 and all state modes; adversarial decay at g = -1, -2, and the -5 contract boundary, plus mixed strong decay and both erase-gate endpoints; out-of-contract rejection; compile-cache route boundaries; the DSL version gate; a deterministic bitwise stress matrix; and a four-tool Compute Sanitizer runner.
Five immutable rows covering MHA, GVA2, GVA4, packed variable-length input,
all four state modes, T={64,1024,4096}, and N={1,20}. Compilation is
recorded separately and excluded from CUDA-event timing.
Covers the public contract and its preconditions, the warp-group schedule, the real six-field compile-cache key and its route boundaries, the audited register and shared-memory figures, the relationship to the repository-wide CuTeDSL contract, and the blockwise-rebased factorization with the bound that produces the [-5, 0] decay contract. Registers cula/gdn2 and cula/ops/gdn2 in REPO_LAYOUT.md, and adds the README quick-start and reproduction commands.
b3a1ed0 to
36a4946
Compare
|
Thanks — all three are addressed, and the first one changed the kernel rather than just the contract. Pushed as 1. Unbounded inverse decay factor. Fixed in the algebra. The intra-chunk factorization now rebases per 16-token sub-block, so the only operand carrying a positive exponent spans at most 15 in-block token gaps instead of the full 64-token chunk. The overflow cliff moves from a uniform Adversarial tests are added: uniform This costs performance and I want to be explicit about why, because the cost is largely intrinsic. The separable chunk-start form is exactly what lets a single WGMMA pair cover the whole tile; the stable correction 2. Compile-cache identity. The documentation was wrong, not the dispatch. 3. CuTeDSL version range. One range, Final gates on |
icavan
left a comment
icavan
left a comment
There was a problem hiding this comment.
Thank you for the update and for documenting the new blockwise factorization in detail. I compared the updated implementation with the pinned FLA GDN2 source. The full-chunk inverse-decay overflow is addressed, but I believe two numerical edge cases remain in the new factor path.
The benchmark invokes FLA with safe_gate=False and use_gate_in_kernel=False. On that path, FLA keeps the chunk-local prefix in FP32, uses non-positive effective exponents for causal score construction, and does not materialize a tiny cross-subblock scale in FP16. Therefore, the two cases below appear specific to this implementation rather than inherited limitations of the reference path.
Would you mind addressing these before merge? Directed regression tests for the sparse exp(-32) case and the finite-BF16 k * exp(75) case would also help protect the intended input contract. Keeping the pair products in FP32 (or recomputing them from the FP32 deltas), and avoiding the positive-exponent BF16 key operand, would align more closely with FLA's numerical structure.
Thanks again for the substantial work on this kernel.
| key_channel, | ||
| cutlass.Int32(0), | ||
| prepare_stage, | ||
| ] = cutlass.Float16(product_two_low) |
There was a problem hiding this comment.
Could we keep these pair products in FP32, or recompute them from the FP32 delta rows at the point of use?
With an allowed uniform g = -1, the distance-two product is exp(-32) ≈ 1.27e-14, which becomes exactly zero when stored as FP16. That can erase a finite causal contribution when q and k compensate for the small decay. For example, with q[32] = k[0] ≈ 1e7, b = 0, and w[0] = v[0] = 1, the direct recurrence and pinned FLA path produce approximately 0.1125 at token 32, while this stored scale makes that contribution zero.
In the pinned FLA safe_gate=False cross-subblock kernel, the corresponding decay factors remain FP32 until accumulation, so this underflow does not occur there.
| ], | ||
| ) | ||
| * cute.math.exp2( | ||
| (block_start_g - g_value) * cutlass.Float32(_INV_LN2), |
There was a problem hiding this comment.
Could we avoid materializing this positive-exponent key operand in BF16, or otherwise add and enforce a corresponding q/k magnitude bound?
The g >= -5 restriction bounds the exponent to 75, but it does not bound the product with k. A finite BF16 value such as k = 999424 gives k * exp(75) ≈ 3.73e38, which exceeds the FP32 maximum before the BF16 conversion and becomes Inf. Inputs with zero q, b, w, and initial state have an exact zero recurrence result, but this intermediate can introduce 0 * Inf -> NaN.
The benchmarked FLA path (safe_gate=False) forms causal decay terms using non-positive exponents. FLA's separate safe-gate kernel also uses a midpoint rebase with FP32 operands rather than this block-start exp(+75) BF16 operand.
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;
with the compile cache keyed by device, value-head count, the two state
modes, and two shape-derived route booleans;
whose stored exponents are bounded, so no valid input overflows FP32;
and adversarial-decay correctness tests, a deterministic stress matrix,
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.165xequal-weightgeometric-mean speedup and
1.319xminimum row speedup.Supported SM90 contract
Hq=16Hv={16,32,64}1 <= N <= 32, each sequence non-emptyK=V=128g[-5, 0]nvidia-cutlass-dsl>=4.5.1,<4.7(both endpoints exercised on H20)[N,Hv,V,K]Kernel and dispatch
(device, Hv, has_initial_state, store_final_state, use_n1_hv16_v64, retain_final_tail); the last two areshape-derived route booleans, so
TandNstay dynamic within adispatch route while crossing the
N=1/T<=64route boundariescompiles a new specialization (at most three per final-state mode). See
docs/gdn2_sm90_pipeline.mdandtest_compile_cache_boundaries;spans at most 15 in-block token gaps
(
docs/gdn2_sm90_stable_factor.md);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.
36a4946.d1ce07369d581813553f30a750af3b6b5f9af6a9.SM100 validation remain a separate follow-up; this PR does not close the
issue.
🧪 Final-source validation
Validation environment
Validation summary
22 passed in 386.92sg = -1,-2,-5(contract boundary), mixed strong decay, both erase-gate endpoints, out-of-contract rejectionT/Ndynamic within a routenvidia-cutlass-dsl4.5.1The 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
Both revisions are divided by the same frozen FLA medians. Same-run FLA
timings on this host drift by up to 20%, which would otherwise make the two
revisions incomparable.
9.407x8.800x1.319x1.491x1.680x1.900x1.534x1.756x1.489x1.733x2.165x2.377xEvery row is faster than pinned FLA. The smallest row speedup is
1.319x;the equal-weight geometric-mean speedup is
2.165x. Against the firstrevision the product itself is
9.7%slower in geometric mean, and thesingle-chunk row is
5%faster.Numerical stability cost
The first revision computed the intra-chunk matrices from a chunk-start
split,
q_i exp(G_i)againstk_j exp(-G_j). That form is separable: onescaling of each operand serves all sixteen block pairs, so a single pair of
m64n64k128WGMMA chains covers the whole64x64tile. It is also the formthat overflows FP32 for valid inputs, which is what this revision fixes.
The stable form rebases per 16-token sub-block. Its correction factor
exp(Gs(I) - Gs(J))is per channel and depends on both block indices, soit cannot be expressed as a per-row scaling of two operands — it has to be
folded into an operand per block pair. Because Hopper WGMMA has a minimum M
of 64, a 16-row band cannot be issued as its own WGMMA without computing and
discarding three quarters of the tile. The per-band basis therefore forces a
choice between per-pair warp-level MMAs (this revision) and 4x redundant
tensor work.
Both alternatives were implemented and measured on H20, not just reasoned
about. Redundant WGMMA regressed multi-chunk rows by 30%. Redistributing the
factor pairs across the idle state warps was correct but 1.8% slower,
because the only deadlock-free arrangement serializes the inverse against
the state recurrence's V/W consumption, which the current schedule overlaps.
Nsight Compute attributes the remaining gap on multi-chunk rows to three
roughly equal legs: the factor computation itself (intrinsic to the stable
form), mbarrier waiting during preparation, and state-warp-group spin. The
TMA wait region is faster than the first revision. With 168 registers and
232,192 B of shared memory the kernel runs one CTA per SM, so there is no
second CTA to hide the ring latency, and the 256 B of shared memory still
free rules out deeper pipelining on the V128 routes.
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
Changes since the first review round (all three points from @icavan):
the intra-chunk factorization was rebased blockwise so no valid input
overflows FP32, and the decay contract is documented and validated as
g in [-5, 0]; the compile-cache key is documented as it actually is andpinned by a test; and one CuTeDSL range is enforced everywhere, with both
endpoints exercised on H20. Rebased onto current
main.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};and this revision's rows were not re-measured with the full paired
30-receipt protocol;
goutside[-5, 0]is outside the contract;validate_inputs=Truerejects it, the default path treats it as a caller precondition;
>=4.5.1,<4.7reports the backend unavailable; 4.4.xcannot import it at all;