feat: add sm100 cp delta rule prefill - #4078
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds Blackwell SM100 context-parallel GDN support. The change adds UTCMMA precompute kernels, TCGen05 prefill execution, SM100 dispatch heuristics, optional state handling, public DSL exports, compile-cache support, and architecture-specific tests. ChangesSM100 context-parallel GDN
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant chunk_gated_delta_rule
participant cp_delta_rule_dsl_sm100
participant cp_delta_rule_t_precompute_dsl_sm100
participant cp_delta_rule_mn_precompute_dsl_sm100
participant cp_delta_rule_fixup_dsl_sm100
participant cp_delta_rule_prefill_dsl_sm100
chunk_gated_delta_rule->>cp_delta_rule_dsl_sm100: dispatch SM100 CP workload
cp_delta_rule_dsl_sm100->>cp_delta_rule_t_precompute_dsl_sm100: compute T workspace
cp_delta_rule_t_precompute_dsl_sm100->>cp_delta_rule_mn_precompute_dsl_sm100: pass T workspace
cp_delta_rule_mn_precompute_dsl_sm100->>cp_delta_rule_fixup_dsl_sm100: pass transfer and state workspaces
cp_delta_rule_fixup_dsl_sm100->>cp_delta_rule_prefill_dsl_sm100: pass fixed state
cp_delta_rule_prefill_dsl_sm100-->>chunk_gated_delta_rule: write output and optional final state
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![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 adds native context-parallel (CP) Gated Delta Net (GDN) prefill kernels for Blackwell SM100 GPUs, including MN precompute, state fixup, and prefill DSL implementations, along with corresponding dispatch heuristics and end-to-end tests. The review feedback identifies a potential division-by-zero error in the dispatch heuristics when the number of sequences is zero, and suggests optimizing GPU memory usage by deferring the allocation of tensormaps_t to the reference kernel path instead of allocating it unconditionally.
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.
22d584c to
963fbf9
Compare
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
flashinfer/gdn_prefill.py (1)
317-324: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate the
use_cpdocstring to include SM100.This change enables CP dispatch for
_arch_major == 10, but theuse_cpparameter docstring (lines 196-200) still reads "SM90/SM120 context-parallel DSL implementation." Please sync it to mention SM100. As per coding guidelines: "Keep documentation in sync with code changes."🤖 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 `@flashinfer/gdn_prefill.py` around lines 317 - 324, Update the use_cp parameter docstring to mention SM100 alongside SM90 and SM120, keeping the documentation aligned with the _arch_major support in the cp_heuristic_matches logic.Source: Coding guidelines
🤖 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 `@flashinfer/gdn_prefill.py`:
- Around line 63-73: Update _cp_delta_rule_rejection_reason to reject non-fp32
initial_state when arch_major is 10, returning a rejection reason so auto
routing falls back to the non-CP path. Preserve the existing SM100
kernel-availability and other architecture checks.
---
Nitpick comments:
In `@flashinfer/gdn_prefill.py`:
- Around line 317-324: Update the use_cp parameter docstring to mention SM100
alongside SM90 and SM120, keeping the documentation aligned with the _arch_major
support in the cp_heuristic_matches logic.
🪄 Autofix (Beta)
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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 40b2faf3-6882-4d90-8d05-2adb6c0e131d
📒 Files selected for processing (8)
flashinfer/gdn_kernels/__init__.pyflashinfer/gdn_kernels/blackwell/__init__.pyflashinfer/gdn_kernels/blackwell/gated_delta_net_cp.pyflashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.pyflashinfer/gdn_kernels/blackwell/gdn_cp_prefill.pyflashinfer/gdn_kernels/delta_rule_dsl/varlen_helper.pyflashinfer/gdn_prefill.pytests/gdn/test_prefill_cp_delta_rule.py
|
/bot run tests/gdn |
|
GitLab MR !1030 has been created, and the CI pipeline #59316121 is currently running. I'll report back once the pipeline job completes. |
|
[FAILED] Pipeline #59316121 — 11/18 executed test jobs passed Compared with nightly #59054944 (different CI configuration). Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
Failure detailsPR-related regressions
Could not compare
Timeouts, infrastructure, or incomplete jobs
|
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
flashinfer/gdn_prefill.py (2)
363-372: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUpdate the state-index dispatch documentation.
The comment and error text state that
state_indicesis supported only by the non-CP SM100 kernel. The code now forwardsstate_indicestocp_delta_rule_dsl_sm100.Describe support as SM100/SM103-specific without excluding CP.
As per coding guidelines, keep documentation synchronized with code changes. <coding_guidelines>
🤖 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 `@flashinfer/gdn_prefill.py` around lines 363 - 372, Update the documentation and NotImplementedError text in the state_indices dispatch guard to describe support as SM100/SM103-specific without saying it is non-CP-only. Keep the validation logic unchanged, including the _arch_major check and existing use_cp context.Source: Coding guidelines
437-449: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winConvert
cu_seqlensfor the selected CP architecture.The public API accepts either integer dtype. This call now forwards the tensor unchanged. SM90 and SM120 CP wrappers require
torch.int64, while SM100 accepts its supported index dtypes. An int32 input therefore fails only when automatic routing selects SM90 or SM120 CP.Convert to int64 for SM90 and SM120, or make all selected kernels accept the documented input dtypes.
Proposed fix
+ cp_cu_seqlens = ( + cu_seqlens if _arch_major == 10 else cu_seqlens.to(torch.int64) + ) cp_delta_rule_dsl( output, output_state, q, k, v, _g, _beta, - cu_seqlens, + cp_cu_seqlens,🤖 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 `@flashinfer/gdn_prefill.py` around lines 437 - 449, Update the cp_delta_rule_dsl invocation in the CP prefill path to convert cu_seqlens to torch.int64 when routing to SM90 or SM120 architectures, while preserving supported input dtypes for SM100. Ensure both documented integer input dtypes work with automatic CP architecture selection.
🧹 Nitpick comments (9)
flashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.py (3)
1691-1691: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDelete the two dead locals.
valid_tokensat line 1691 andsum_vat line 1713 are assigned and never read. Both carry# noqa: F841to silence the linter. Remove the assignments and the suppressions instead of keeping them.Also applies to: 1713-1713
🤖 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 `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.py` at line 1691, Remove the unused valid_tokens assignment near batch_end and the unused sum_v assignment in the same function, including their # noqa: F841 suppressions. Leave the surrounding computation unchanged.Source: Linters/SAST tools
1737-1743: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the unused
work_argsand unused unpacked names.
mma_cg0_warpdeclareswork_argsand never reads it;mma_cg0_pair_cppasses an empty tuple at line 2005. Static analysis also reports unused unpacked names at lines 2061 (tiled_mma_kv), 2065 (batch_idx), 2474 (sV), 2687 (sBeta), and 2703 (num_pairs_b). Remove the parameter and prefix the intentionally unused names with an underscore.Also applies to: 2035-2043
🤖 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 `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.py` around lines 1737 - 1743, Remove the unused work_args parameter from mma_cg0_warp and stop passing the empty tuple from mma_cg0_pair_cp. Rename the intentionally unused unpacked names at the referenced sites—tiled_mma_kv, batch_idx, sV, sBeta, and num_pairs_b—with an underscore prefix while preserving tuple unpacking behavior.Source: Linters/SAST tools
2431-2444: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffThe checkpoint blocks are unreachable in this class.
__init__pinsself.enable_checkpoints = Falseat line 195 and ignores theenable_checkpointsargument. Every guarded block therefore compiles out, andmS_checkpoints,checkpoint_offset, andcheckpoint_every_n_tokensthread dead values throughcompute_group_1_chunkand_store_final_state. If CP will never support checkpointing, remove the parameter and the blocks. If it will, add a TODO that states the intent.Also applies to: 2953-2994
🤖 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 `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.py` around lines 2431 - 2444, Update the class initialization and checkpoint flow so checkpoint support is explicit: either honor the constructor’s enable_checkpoints argument and ensure the checkpoint values reach compute_group_1_chunk and _store_final_state, or, if CP will not support checkpoints, remove the parameter, guarded checkpoint blocks, and dead mS_checkpoints, checkpoint_offset, and checkpoint_every_n_tokens plumbing; if retaining the disabled path, add a TODO documenting the intended support.flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.py (3)
906-910: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueRemove the duplicate TMEM store fence.
_scale_tmem_matrixalready ends withcute.arch.fence_view_async_tmem_store()at line 292. The fence at line 909 repeats it with no intervening TMEM store. Drop it unless it guards a hazard that is not visible here.🤖 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 `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.py` around lines 906 - 910, Remove the redundant cute.arch.fence_view_async_tmem_store() call immediately after _scale_tmem_matrix in the update path, since _scale_tmem_matrix already fences its TMEM store. Keep y_ready_producer.acquire_and_advance().commit() and the surrounding synchronization unchanged.
1365-1365: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse the named warp-ID attributes in this dispatch condition.
Every other branch in this dispatch chain compares
warp_idxagainst a named attribute, for exampleself.transfer_mma_warp_idandself.compute_group_1_warp_ids. This branch hardcodes0and3. Ifcompute_group_0_warp_idschanges, this branch silently diverges.♻️ Proposed change
- elif is_valid_chunk and warp_idx >= 0 and warp_idx <= 3: + elif ( + is_valid_chunk + and warp_idx >= self.compute_group_0_warp_ids[0] + and warp_idx <= self.compute_group_0_warp_ids[-1] + ):🤖 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 `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.py` at line 1365, Update the dispatch condition in the branch following the `is_valid_chunk` check to compare `warp_idx` against the bounds represented by `self.compute_group_0_warp_ids` instead of hardcoded `0` and `3`. Preserve the existing validity and inclusive range behavior while deriving both limits from the named attribute.
851-857: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the
tiled_mma_zparameter ofrun_compute_group_1.The caller at line 1399 passes
tiled_mma_yinto this slot. The parameter name statesz. Rename the parameter totiled_mma_yso the reader does not have to trace the call site to learn which MMA is used.🤖 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 `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.py` around lines 851 - 857, Rename the run_compute_group_1 parameter tiled_mma_z to tiled_mma_y, preserving its position and behavior so it matches the tiled_mma_y value passed by the caller.flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py (2)
156-158: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the duplicated
from_dlpacklambda into one module-level helper. All four wrappers define the identicalfrom_dlpacklambda inside their compile branch. Ruff reports E731 at each site. One shareddefat module scope removes the duplication and clears all four findings.
flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py#L156-L158: delete the local lambda incp_delta_rule_t_precompute_dsl_sm100and call the shared helper.flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py#L311-L313: delete the local lambda incp_delta_rule_mn_precompute_dsl_sm100and call the shared helper.flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py#L563-L565: delete the local lambda incp_delta_rule_fixup_dsl_sm100and call the shared helper.flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py#L841-L843: delete the local lambda incp_delta_rule_prefill_dsl_sm100and call the shared helper.♻️ Proposed helper
Add near
_get_cp_workspace:def _from_dlpack(*args, **kwargs): return cute.runtime.from_dlpack(*args, **{**kwargs, "enable_tvm_ffi": True})Then at each site:
- from_dlpack = lambda *args, **kwargs: cute.runtime.from_dlpack( - *args, **{**kwargs, "enable_tvm_ffi": True} - ) + from_dlpack = _from_dlpack🤖 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 `@flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py` around lines 156 - 158, Extract the duplicated from_dlpack lambda into a single module-level _from_dlpack helper near _get_cp_workspace, preserving the enable_tvm_ffi behavior. Remove the local lambda and use the shared helper in cp_delta_rule_t_precompute_dsl_sm100 at flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py:156-158, cp_delta_rule_mn_precompute_dsl_sm100 at :311-313, cp_delta_rule_fixup_dsl_sm100 at :563-565, and cp_delta_rule_prefill_dsl_sm100 at :841-843.Source: Linters/SAST tools
538-550: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAvoid exposing unsupported SM100 fixup kernel kinds.
cp_delta_rule_fixup_dsl_sm100()auto-selects onlysimt_row4,utcmma64, andutcmma128, but_get_fixup_kernel()also acceptssimt_row8andhmma. Those branches can only be reached by passing an unsupported_kernel_kindmanually. Remove the unsupported branches or add SM100 selection paths/tests for them.🤖 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 `@flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py` around lines 538 - 550, Update _get_fixup_kernel() to stop accepting the unsupported simt_row8 and hmma kernel kinds for cp_delta_rule_fixup_dsl_sm100(); remove those branches so dispatch is limited to the three kinds selected by cp_delta_rule_fixup_dsl_sm100(): simt_row4, utcmma64, and utcmma128.flashinfer/utils.py (1)
281-282: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the shared device-property lookup.
get_device_propertiesis not decorated withfunctools.cache. As a result,get_device_sm_countand other direct consumers still query CUDA properties on each call.Proposed fix
+@functools.cache def get_device_properties(device: torch.device): return torch.cuda.get_device_properties(device)🤖 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 `@flashinfer/utils.py` around lines 281 - 282, Decorate get_device_properties with functools.cache so repeated lookups for the same device reuse the cached CUDA properties. Import functools as needed, while leaving get_device_sm_count and other consumers unchanged.
🤖 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 `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.py`:
- Around line 2590-2609: Resolve the redundant `valid_state` branch in the
`num_kphases_qkv` GEMM loop: determine the intended accumulation seed by
matching the sibling GEMM loops, using `valid_state or (kphase_idx != 0)` if the
non-initial-state path requires accumulation from later phases; otherwise
collapse both identical branches into one loop. Ensure the resulting behavior
remains correct if `valid_state` becomes false in the future.
In `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.py`:
- Around line 1533-1565: Run ruff format on the file, allowing it to normalize
the column-aligned assignments in CPDeltaRuleFixupUtcmmaSm100.__init__ and the
compressed signatures below it. Commit the formatter’s complete output without
manual formatting overrides.
In `@flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py`:
- Line 817: Run ruff format on
flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py and commit the resulting
formatting changes, including wrapping the state_dtype ternary chain and the
nearby multi-condition checks to match Ruff’s output.
In `@flashinfer/gdn_kernels/delta_rule_dsl/varlen_helper.py`:
- Around line 150-169: Update the calls to the chunk-length heuristic in
cp_delta_rule_dsl_sm90 and cp_delta_rule_dsl_sm120 to pass each caller’s actual
num_seqs value along with total_seqlen. Preserve the existing heuristic logic in
chunk_bound_for_len and ensure both SM90 and SM120 paths no longer rely on the
default num_seqs=1.
In `@flashinfer/gdn_prefill.py`:
- Around line 353-362: Update the CP dispatch decision around
_cp_delta_rule_rejection_reason to reject SM100 CP when the CUDA version is
below 13, before will_use_cp can select it. For use_cp == "auto", return the
non-CP path when this rejection applies; preserve explicit use_cp=True rejection
behavior and existing supported-version dispatch.
---
Outside diff comments:
In `@flashinfer/gdn_prefill.py`:
- Around line 363-372: Update the documentation and NotImplementedError text in
the state_indices dispatch guard to describe support as SM100/SM103-specific
without saying it is non-CP-only. Keep the validation logic unchanged, including
the _arch_major check and existing use_cp context.
- Around line 437-449: Update the cp_delta_rule_dsl invocation in the CP prefill
path to convert cu_seqlens to torch.int64 when routing to SM90 or SM120
architectures, while preserving supported input dtypes for SM100. Ensure both
documented integer input dtypes work with automatic CP architecture selection.
---
Nitpick comments:
In `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.py`:
- Line 1691: Remove the unused valid_tokens assignment near batch_end and the
unused sum_v assignment in the same function, including their # noqa: F841
suppressions. Leave the surrounding computation unchanged.
- Around line 1737-1743: Remove the unused work_args parameter from mma_cg0_warp
and stop passing the empty tuple from mma_cg0_pair_cp. Rename the intentionally
unused unpacked names at the referenced sites—tiled_mma_kv, batch_idx, sV,
sBeta, and num_pairs_b—with an underscore prefix while preserving tuple
unpacking behavior.
- Around line 2431-2444: Update the class initialization and checkpoint flow so
checkpoint support is explicit: either honor the constructor’s
enable_checkpoints argument and ensure the checkpoint values reach
compute_group_1_chunk and _store_final_state, or, if CP will not support
checkpoints, remove the parameter, guarded checkpoint blocks, and dead
mS_checkpoints, checkpoint_offset, and checkpoint_every_n_tokens plumbing; if
retaining the disabled path, add a TODO documenting the intended support.
In `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.py`:
- Around line 906-910: Remove the redundant
cute.arch.fence_view_async_tmem_store() call immediately after
_scale_tmem_matrix in the update path, since _scale_tmem_matrix already fences
its TMEM store. Keep y_ready_producer.acquire_and_advance().commit() and the
surrounding synchronization unchanged.
- Line 1365: Update the dispatch condition in the branch following the
`is_valid_chunk` check to compare `warp_idx` against the bounds represented by
`self.compute_group_0_warp_ids` instead of hardcoded `0` and `3`. Preserve the
existing validity and inclusive range behavior while deriving both limits from
the named attribute.
- Around line 851-857: Rename the run_compute_group_1 parameter tiled_mma_z to
tiled_mma_y, preserving its position and behavior so it matches the tiled_mma_y
value passed by the caller.
In `@flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py`:
- Around line 156-158: Extract the duplicated from_dlpack lambda into a single
module-level _from_dlpack helper near _get_cp_workspace, preserving the
enable_tvm_ffi behavior. Remove the local lambda and use the shared helper in
cp_delta_rule_t_precompute_dsl_sm100 at
flashinfer/gdn_kernels/blackwell/gdn_cp_prefill.py:156-158,
cp_delta_rule_mn_precompute_dsl_sm100 at :311-313, cp_delta_rule_fixup_dsl_sm100
at :563-565, and cp_delta_rule_prefill_dsl_sm100 at :841-843.
- Around line 538-550: Update _get_fixup_kernel() to stop accepting the
unsupported simt_row8 and hmma kernel kinds for cp_delta_rule_fixup_dsl_sm100();
remove those branches so dispatch is limited to the three kinds selected by
cp_delta_rule_fixup_dsl_sm100(): simt_row4, utcmma64, and utcmma128.
In `@flashinfer/utils.py`:
- Around line 281-282: Decorate get_device_properties with functools.cache so
repeated lookups for the same device reuse the cached CUDA properties. Import
functools as needed, while leaving get_device_sm_count and other consumers
unchanged.
🪄 Autofix (Beta)
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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 2515ecef-f5ba-4323-bf4d-bf531cd78893
📒 Files selected for processing (12)
flashinfer/gdn_kernels/__init__.pyflashinfer/gdn_kernels/blackwell/__init__.pyflashinfer/gdn_kernels/blackwell/gated_delta_net_cp.pyflashinfer/gdn_kernels/blackwell/gated_delta_net_cp_prefill.pyflashinfer/gdn_kernels/blackwell/gdn_cp_prefill.pyflashinfer/gdn_kernels/delta_rule_dsl/alpha.pyflashinfer/gdn_kernels/delta_rule_dsl/custom_compile_cache.pyflashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.pyflashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.pyflashinfer/gdn_kernels/delta_rule_dsl/varlen_helper.pyflashinfer/gdn_prefill.pyflashinfer/utils.py
🚧 Files skipped from review as they are similar to previous changes (2)
- flashinfer/gdn_kernels/blackwell/init.py
- flashinfer/gdn_kernels/init.py
| # Target one wave of MN CTAs. Account for the known longest sequence, then | ||
| # safely bound the chunks contributed by all remaining uneven sequences. | ||
| target_chunks = max(1, num_sms // num_heads) | ||
| min_chunk_len = _ceil_div(max_seqlen, target_chunks) | ||
| return _round_up(min_chunk_len, chunk_len_granularity) | ||
| remaining_seqlen = max(0, total_seqlen - max_seqlen) | ||
| remaining_seqs = max(0, num_seqs - 1) | ||
|
|
||
| def chunk_bound_for_len(chunk_len: int) -> int: | ||
| return _ceil_div(max_seqlen, chunk_len) + chunk_bound_host( | ||
| remaining_seqs, remaining_seqlen, chunk_len | ||
| ) | ||
|
|
||
| lo = 1 | ||
| hi = max(1, _ceil_div(max_seqlen, chunk_len_granularity)) | ||
| while lo < hi: | ||
| mid = (lo + hi) // 2 | ||
| if chunk_bound_for_len(mid * chunk_len_granularity) <= target_chunks: | ||
| hi = mid | ||
| else: | ||
| lo = mid + 1 | ||
| return lo * chunk_len_granularity |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Pass num_seqs from the SM90 and SM120 callers.
Those callers provide total_seqlen but retain the default num_seqs=1. Therefore, remaining_seqs becomes zero and chunk_bound_host ignores the per-sequence ceiling cost. The heuristic can select a chunk length that launches more than the target one-wave CTA count.
Pass the actual sequence count from cp_delta_rule_dsl_sm90 and cp_delta_rule_dsl_sm120.
🤖 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 `@flashinfer/gdn_kernels/delta_rule_dsl/varlen_helper.py` around lines 150 -
169, Update the calls to the chunk-length heuristic in cp_delta_rule_dsl_sm90
and cp_delta_rule_dsl_sm120 to pass each caller’s actual num_seqs value along
with total_seqlen. Preserve the existing heuristic logic in chunk_bound_for_len
and ensure both SM90 and SM120 paths no longer rely on the default num_seqs=1.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/gdn/test_prefill_cp_delta_rule.py`:
- Line 89: Update the CUDA-version skip diagnostics to use consistent
SM100/SM103 (or SM10x) wording: change the message in
tests/gdn/test_prefill_cp_delta_rule.py at lines 89-89 and apply the same
wording in tests/gdn/test_prefill_delta_rule.py at lines 58-58, preserving the
existing CUDA version details.
🪄 Autofix (Beta)
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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4e1d618f-3075-495e-8231-424635f56d7e
📒 Files selected for processing (2)
tests/gdn/test_prefill_cp_delta_rule.pytests/gdn/test_prefill_delta_rule.py
|
/bot run tests/gdn |
|
[FAILED] Pipeline #60988907 — 17/18 executed test jobs passed Compared with nightly #60831563. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
No individual test or infrastructure failures could be extracted. |
jiahanc
left a comment
jiahanc
left a comment
There was a problem hiding this comment.
Thanks for the contribution!
|
@guangyunh-nv |
|
Hi @guangyunh-nv thanks for the reminder and this should never happen again, I have made some changes on the github side and agent side to prohibit such behaviors. |
Mirrors the v0.6.16 revert (#4246, commit 0a27ba3) and its v0.6.17 counterpart (09fd5fc) so that 0.6.18 does not ship the API-breaking ring-buffer cache contract that neither 0.6.16 nor 0.6.17 shipped. Reverts: afd4754 mamba checkpointing SSU: two-kernel split + ring-buffer cache for checkpointing SSU (#3975) f90e9c4 docs(mamba): document checkpointing varlen arguments (#4129) Like release-v0.6.17, this branch carries #4129 as well, so both are reverted; 0.6.16 only needed #3975. main still carries both. Verified no collateral damage: the SM107 change from #4280 in tests/mamba/conftest.py and the #4029/#4078 changes in flashinfer/utils.py are preserved, and the five core reverted files now byte-match release-v0.6.17. NOTE: like both earlier reverts, this also reverts the cvt_rs fix that rode along in #3975 -- is_cvt_rs_supported goes back to `major in (10, 11)` (wrong for SM110a) and the CUDA guard back to SM100_ALL only (B300/sm_103a falls to software emulation). That matches what 0.6.17 shipped, but the hunk remains a candidate to keep.
kahyunnam
added
model: qwen3.5 / 3.6 / 3.8
op: linear attention
6 participants
📌 Description
This PR adds cp impl for delta rule.
🔍 Related Issues
🚀 Pull Request Checklist
✅ 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.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
Tests