[GDN] Add Hopper SM90 CuTe DSL prefill#108
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 a packed variable-length Gated DeltaNet (GDN) prefill operator for NVIDIA Hopper (SM90) GPUs, implemented in CuTe DSL. The feedback highlights a critical concurrency hazard in the workspace allocation for tensormaps_t that could lead to data corruption, as well as high-severity performance bottlenecks caused by synchronous host-device transfers during tensor validation. Additionally, a defensive check is recommended to prevent crashes when checking prefill availability on non-CUDA devices.
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.
| workspace_size = get_device_sm_count(q.device) * 128 | ||
| tensormaps_t = _get_cache_buf("gdn_prefill_tensormaps", workspace_size, q.device) |
There was a problem hiding this comment.
Using a shared global/cached buffer _get_cache_buf for tensormaps_t introduces a critical concurrency hazard. If multiple streams or concurrent threads launch the GDN kernel on the same device, they will share the same tensormaps_t buffer. Since the kernel indexes into this buffer using smid(), concurrent execution on the same Streaming Multiprocessors (SMs) will cause silent data corruption or illegal memory access crashes. Replacing this with a direct allocation via PyTorch's caching allocator is thread-safe, stream-safe, and has virtually zero overhead for such a small buffer (~16 KB).
| workspace_size = get_device_sm_count(q.device) * 128 | |
| tensormaps_t = _get_cache_buf("gdn_prefill_tensormaps", workspace_size, q.device) | |
| workspace_size = get_device_sm_count(q.device) * 128 | |
| tensormaps_t = torch.empty(workspace_size, dtype=torch.uint8, device=q.device) |
There was a problem hiding this comment.
Thanks for flagging this. I audited the exact specialization used by this PR and do not believe a per-launch workspace is required here.
The TensorMap workspace is indexed by the physical smid() with one 128-byte slot per SM. The source-bound H20 resource and occupancy capture for this 512-thread kernel reports 185,728 bytes of dynamic shared memory and exactly one active CTA per SM. In addition, the output S2G path executes cp_async_bulk_wait_group(0) before the CTA can exit.
Consequently, CTAs from different streams cannot be resident concurrently on the same SM and write the same slot. The next CTA can reuse that SM's slot only after the previous CTA's bulk operation has retired and the CTA has exited.
I added a comment in 6bc5c4d to make this specialization-bound invariant explicit. If a future resource layout permits more than one resident CTA per SM, the cache design will need to be revisited.
| def validate() -> tuple[int, ...]: | ||
| offsets = tuple(int(value) for value in cu_seqlens.detach().cpu().tolist()) |
There was a problem hiding this comment.
Calling cu_seqlens.detach().cpu().tolist() forces a synchronous host-device transfer and CPU-GPU synchronization. In variable-length workloads where sequence lengths change dynamically, cu_seqlens is recreated at each step, causing the cache to miss and introducing a major performance bottleneck. Consider refactoring _GDNPrefillInputs to store num_seqs: int (which is simply cu_seqlens.numel() - 1 and can be retrieved without synchronization) instead of seq_lens: tuple[int, ...], and bypassing or making these synchronous validation checks optional in production/steady-state execution.
There was a problem hiding this comment.
Fixed in 6bc5c4d.
The default path no longer calls cu_seqlens.detach().cpu().tolist(). It derives num_seqs from tensor metadata (numel() - 1) and retains the existing rank, dtype, device, contiguity, and alignment checks.
CUDA-resident value checks are now an explicit validate_inputs=True diagnostic path, so production calls do not force a device-to-host synchronization.
An H20 CPU-profiler trace covered five calls with a new set of Q/K/V/g/beta/cu_seqlens/output tensors on every call and observed none of the following operations:
aten::_local_scalar_denseaten::_to_copyaten::allaten::gtaten::isfiniteaten::item
The complete SM90 test file also passes: 25 tests on H20 with CuTe DSL 4.5.1.
| def validate_alpha() -> None: | ||
| if not bool(torch.isfinite(alpha).all()) or not bool((alpha > 0).all()): | ||
| raise ValueError("g must contain finite, strictly positive forget factors") | ||
|
|
||
| _cached_content_validation("alpha", alpha, validate_alpha) |
There was a problem hiding this comment.
Calling bool(tensor.all()) on CUDA tensors forces a synchronous device-to-host transfer to retrieve the boolean value, blocking the CPU. When gate factors (g and beta) are dynamically computed by the network at each step, they are new tensors, causing cache misses and introducing constant pipeline stalls. Bypassing or making these validation checks optional in production/steady-state execution will significantly improve throughput.
There was a problem hiding this comment.
I removed the default bool(torch.isfinite(...).all()) and bool((g > 0).all()) reductions, along with the object-identity content cache.
The default fast path now treats finite and positive CUDA values as caller preconditions, while validate_inputs=True retains the synchronous checks for debugging and input-pipeline validation. The API documentation and regression tests describe this contract.
The fresh-object Hopper profiler gate observed no scalar extraction or validation-reduction operations across five calls, and the complete SM90 test file passes 25/25.
6bc5c4d to
7f7619d
Compare
7f7619d to
88737e9
Compare
1 participant
📌 Description
This PR adds a Hopper SM90 Gated DeltaNet prefill path implemented entirely in
CuTe DSL.
cula.ops.gdn.sm90while preserving upstream authorship and provenance;cula.gdn.chunk_gated_delta_rule, including optional gates, initial/finalstate, MHA, GQA, and GVA head mappings;
headers/submodules at runtime;
tests, safety coverage, and Hopper user documentation.
🔍 Related Issues
Closes #76.
Issue checklist
cula/ops/gdn/sm90/implements the SM90a 512-thread TMA/WGMMA kernel and direct TVM-FFI launch path.cula.gdn.chunk_gated_delta_rulevalidates packed inputs and exposes optional gates and recurrent state.benchmarks/bench_gdn_prefill.pydefaults to all 10 fixed and 18 packed-varlen rows from the issue.0.993329822xgeometric mean and1.147581166xworst row versus frozen upstream C++; all 28 rows pass the1.10x/1.25xgates.🚀 Pull Request Checklist
✅ Pre-commit Checks
git diff --checkpasses.🧪 Tests
22 passed.repeat/continuation, and output/state redzones pass.
zero errors; racecheck reports zero hazards.
182/182files).128 registers, zero stack/spill/local memory, and one active CTA/SM.
⚡ Performance
Measured on one idle NVIDIA H20-3e with BF16
Hq=Hk=Hv=64,D=128, 20warmups and 100 CUDA-event timed iterations per row. Each backend ran in three
independent processes; each table entry is the median of the three process
averages. Compilation and first-call setup are excluded. A ratio below
1.0xmeans the CuTe DSL kernel is faster.
<= 1.10x0.993329822x<= 1.25x1.147581166x(fixed-b1-t512)28/281.003484x0.972185x1.147581x0.987733x0.976909x0.999400x0.993330x0.975219x1.147581xFull 28-row median latency table
fixed-b1-t512fixed-b1-t1024fixed-b1-t4096fixed-b1-t8192fixed-b1-t16384fixed-b2-t512fixed-b2-t1024fixed-b2-t4096fixed-b2-t8192fixed-b2-t16384varlen-n10-t4096-uniformvarlen-n10-t4096-randomvarlen-n10-t4096-skewedvarlen-n10-t8192-uniformvarlen-n10-t8192-randomvarlen-n10-t8192-skewedvarlen-n10-t16384-uniformvarlen-n10-t16384-randomvarlen-n10-t16384-skewedvarlen-n20-t4096-uniformvarlen-n20-t4096-randomvarlen-n20-t4096-skewedvarlen-n20-t8192-uniformvarlen-n20-t8192-randomvarlen-n20-t8192-skewedvarlen-n20-t16384-uniformvarlen-n20-t16384-randomvarlen-n20-t16384-skewedAll backend/process inputs are byte-identical. The canonical benchmark can be
reproduced with:
Reviewer Notes
original rewrite
eb31811856e07f3c5fc94a7aaafeaeaabd954f4f, imported source97996a4123924140d28a3079292535aa534eba92, and frozen C++ oracle41e5aa29a0e218bde2e2045316ed88e3f4dc8ca2.s < B and t < Bpredicate and explicit compile-cache key; cuLA supplies the public API,
packaging, tests, benchmark, and documentation around the upstream kernel.
nvidia-cutlass-dsl==4.5.1and fails explicitly on anunsupported GPU or DSL version; it never silently falls back to C++.
Q/K/V, head sizes other than 128, decode, and backward are intentionally out
of scope and documented as unsupported.
correctness/benchmark coverage, and documentation.