perf(gdn): optimize cp host launch overhead for sm90 and sm120 - #4374
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughSM90 and SM120 CP delta-rule kernels now cache workspaces, kernel instances, and compiled kernels. Public wrappers accept optional device and stream overrides. Top-level pipelines share one device and stream across precompute, fixup, and prefill stages. ChangesCP device-aware caching
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant CPDeltaRulePipeline
participant DeviceUtilities
participant PrecomputeWrappers
participant FixupWrapper
participant PrefillWrapper
CPDeltaRulePipeline->>DeviceUtilities: obtain device metadata and stream
CPDeltaRulePipeline->>PrecomputeWrappers: pass device and stream
PrecomputeWrappers->>FixupWrapper: pass precompute outputs
FixupWrapper->>PrefillWrapper: pass fixed state
PrefillWrapper->>CPDeltaRulePipeline: return final outputs
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 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 |
guangyunh-nv
marked this pull request as ready for review
August 6, 2026 07:52
guangyunh-nv
requested review from
bkryu,
jiahanc,
kahyunnam,
yongwww and
yzh119
as code owners
August 6, 2026 07:52
guangyunh-nv
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py (2)
2732-2738: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRecord the rationale for the SM90 fixup-kernel thresholds.
The head-count thresholds here (8 and 16) match the SM120 file. The SM120 file documents them with a comment stating the values come from NCU measurements on SM120. This file has no rationale. If the thresholds were measured on SM90, add a comment that states so. If they were copied from SM120, state that they are unverified for SM90.
🤖 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/delta_rule_cp_sm90.py` around lines 2732 - 2738, Add a concise comment immediately above the _kernel_kind head-count thresholds documenting their provenance: state that they were determined from SM90 NCU measurements if applicable, or explicitly mark them as copied from SM120 and unverified for SM90. Leave the existing selection logic unchanged.
4186-4187: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPrefix the cp tensormap buffer name with
sm90.The local cache key table uses
(name, device), and existing cp tensormap buffers in this workflow usesm120-specific names. Usegdn_cp_sm90_prefill_tensormapsso future code changes cannot reuse this generic name and share the buffer instead of allocating ansm90-specific tensormap workspace.♻️ Proposed rename
workspace_size = get_device_sm_count(device) * 128 - tensormaps_t = _get_cache_buf("gdn_cp_prefill_tensormaps", workspace_size, device) + tensormaps_t = _get_cache_buf( + "gdn_cp_sm90_prefill_tensormaps", workspace_size, 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/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py` around lines 4186 - 4187, Rename the _get_cache_buf key in the cp prefill tensormap allocation to gdn_cp_sm90_prefill_tensormaps, preserving the existing workspace_size and device arguments.flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py (2)
56-58: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the reuse contract of
_get_cp_workspaceand confirm the buffer is fully written.
_get_cp_workspacereturns a view into a process-global buffer keyed by(name, device). The buffer is not zero-initialized. Two conditions must hold for correctness:
- Each kernel must write every element of its workspace before any read. Previously
torch.emptyalso gave uninitialized memory, so this is unchanged, but thetworkspace is now reused across calls with differenttotal_t_blocks. Stale values from a previous larger call remain in the buffer.- The returned tensors escape to the caller.
cp_delta_rule_t_precompute_dsl_sm120,cp_delta_rule_mn_precompute_dsl_sm120, andcp_delta_rule_fixup_dsl_sm120are module-level entry points. A second call with the same name and device overwrites the tensor a previous caller still holds.Add a short docstring on
_get_cp_workspacethat states the buffer is shared per(name, device)and that returned tensors are valid only until the next call on the same 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/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py` around lines 56 - 58, Add a concise docstring to _get_cp_workspace documenting that it returns a shared, non-zero-initialized buffer keyed by (name, device), requiring callers to fully write workspace elements before reading, and that returned views remain valid only until the next call using the same name and device.
583-585: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHoist the
from_dlpacklambda to a module-level function.Ruff reports E731 at this location and at lines 1770-1772, 2959-2961, and 5123-5125. The same three-line lambda is defined four times in this file and four times in
flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py. Define it once at module scope and import or reuse it.♻️ Proposed module-level helper
+def _from_dlpack_tvm(*args, **kwargs): + return cute.runtime.from_dlpack(*args, **{**kwargs, "enable_tvm_ffi": True}) + + def _get_cp_workspace(name, shape, dtype, device):Then at each call site:
if compiled is None: - from_dlpack = lambda *args, **kwargs: cute.runtime.from_dlpack( - *args, **{**kwargs, "enable_tvm_ffi": True} - ) + from_dlpack = _from_dlpack_tvm kernel_args = (🤖 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/delta_rule_cp_sm120.py` around lines 583 - 585, Define a module-level helper function for the `from_dlpack` wrapper around `cute.runtime.from_dlpack`, then replace all duplicated lambda assignments in this module—including the locations near lines 583, 1770, 2959, and 5123—with that helper. Reuse the same module-level helper in the corresponding SM90 module instead of redefining it, while preserving the `enable_tvm_ffi=True` behavior and call signature.Source: Linters/SAST tools
🤖 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/delta_rule_dsl/delta_rule_cp_sm120.py`:
- Around line 56-58: Prevent process-wide scratch-buffer aliasing in
_get_cp_workspace and the related CP kernel paths: key _get_cache_buf
allocations by the caller stream or an equivalent caller slot so concurrent
cp_delta_rule_dsl_sm90() and cp_delta_rule_dsl_sm120() invocations cannot share
TMA or state buffers. Apply this to
flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py lines 56-58 and
5106-5109, and flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py lines
52-54 and 4186-4187; preserve the existing per-SM TMA indexing while ensuring
each concurrent invocation has isolated storage.
---
Nitpick comments:
In `@flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py`:
- Around line 56-58: Add a concise docstring to _get_cp_workspace documenting
that it returns a shared, non-zero-initialized buffer keyed by (name, device),
requiring callers to fully write workspace elements before reading, and that
returned views remain valid only until the next call using the same name and
device.
- Around line 583-585: Define a module-level helper function for the
`from_dlpack` wrapper around `cute.runtime.from_dlpack`, then replace all
duplicated lambda assignments in this module—including the locations near lines
583, 1770, 2959, and 5123—with that helper. Reuse the same module-level helper
in the corresponding SM90 module instead of redefining it, while preserving the
`enable_tvm_ffi=True` behavior and call signature.
In `@flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py`:
- Around line 2732-2738: Add a concise comment immediately above the
_kernel_kind head-count thresholds documenting their provenance: state that they
were determined from SM90 NCU measurements if applicable, or explicitly mark
them as copied from SM120 and unverified for SM90. Leave the existing selection
logic unchanged.
- Around line 4186-4187: Rename the _get_cache_buf key in the cp prefill
tensormap allocation to gdn_cp_sm90_prefill_tensormaps, preserving the existing
workspace_size and device arguments.
🪄 Autofix
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: aa463673-e137-4876-8b39-4d68187ea2ac
📒 Files selected for processing (2)
flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.pyflashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py
Comment on lines
+56
to
+58
| def _get_cp_workspace(name, shape, dtype, device): | ||
| nbytes = math.prod(shape) * dtype.itemsize | ||
| return _get_cache_buf(name, nbytes, device)[:nbytes].view(dtype).view(shape) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether the CP delta-rule entry points are documented or tested as single-threaded per device.
set -euo pipefail
# Find callers of the top-level CP pipelines.
rg -nP --type=py -C5 '\bcp_delta_rule_dsl_sm(90|120)\s*\('
# Look for concurrency or multi-stream tests touching the CP path.
rg -nP --type=py -C5 'cp_delta_rule' -g 'tests/**' | rg -n -C5 'Thread|stream|Stream|concurren' || echo "no concurrency usage found in tests"
# Check whether other cached workspaces in the repo scope their key beyond (name, device).
rg -nP --type=py -C4 'def _get_cache_buf'Repository: flashinfer-ai/flashinfer
Length of output: 162
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate files =="
git ls-files | rg 'flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm(90|120)\.py$' || true
echo "== _get_cache_buf definitions/usages =="
rg -n "def _get_cache_buf|_get_cache_buf\(" flashinfer -S || true
echo "== relevant snippets =="
for f in flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py; do
if [ -f "$f" ]; then
echo "--- $f lines 1-80 ---"
sed -n '1,80p' "$f"
echo "--- $f lines 4180-4192 ---"
sed -n '4180,4192p' "$f"
echo "--- $f lines 5100-5112 ---"
sed -n '5100,5112p' "$f"
fi
done
echo "== cache key implementation search =="
rg -n "_CACHE|cache.*buf|thread|stream|device|smid|CollectiveStoreTma|tail_tensormap_gmem_ptr" flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm(90|120).py -S || true
echo "== entrypoint docstring areas =="
rg -n "def cp_delta_rule_dsl_sm(90|120)|\"\"\"|'''|cache buf|cached|workspace" flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm(90|120).py -S || trueRepository: flashinfer-ai/flashinfer
Length of output: 9585
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== _get_cache_buf implementation =="
sed -n '220,260p' flashinfer/utils.py
echo "== CP entrypoint signatures and workspace/tensormap construction =="
rg -n "def cp_delta_rule_dsl_sm(90|120)|_stream|tensormaps_t|workspace_size|_get_cp_workspace|stream =" flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py -S
echo "== CollectiveStoreTma tail_tensormap_gmem_ptr definition and call sites =="
rg -n "class CollectiveStoreTma|tail_tensormap_gmem_ptr|gmem_ptr" flashinfer/gdn_kernels/delta_rule_dsl -S
echo "== behavioral probe: cache key shape from source =="
python3 - <<'PY'
import ast, pathlib
for path in ["flashinfer/utils.py","flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py","flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py"]:
p=pathlib.Path(path)
src=p.read_text()
print(f"--- {path} ---")
try:
tree=ast.parse(src)
except SyntaxError:
print("SYNTAX ERROR")
continue
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
if node.name in ("_get_cache_buf","_get_cp_workspace","cp_delta_rule_dsl_sm90","cp_delta_rule_dsl_sm120"):
print(f"function {node.name} args={[n.name for n in node.args.args]}, lineno={node.lineno}")
for stmt in node.body:
if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
print(" stmt:", ast.unparse(stmt))
PYRepository: flashinfer-ai/flashinfer
Length of output: 9560
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== _get_cache_buf implementation with context =="
sed -n '210,275p' flashinfer/utils.py
echo "== CP workspace creation contexts =="
sed -n '540,565p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '1736,1752p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '2918,2936p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '5096,5148p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
echo "== SM90 workspace creation contexts =="
sed -n '542,556p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py
sed -n '1736,1752p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py
sed -n '2704,2720p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py
sed -n '4182,4226p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py
echo "== CollectiveStoreTma implementation =="
sed -n '1,90p' flashinfer/gdn_kernels/delta_rule_dsl/collective_store_tma.py
sed -n '108,145p' flashinfer/gdn_kernels/delta_rule_dsl/collective_store_tma.py
echo "== CP entrypoint docstrings =="
sed -n '5184,5220p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '5260,5290p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '5320,5345p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '5355,5385p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '5048,5080p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
sed -n '4520,4580p' flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py
echo "== simple AST extraction for cache key =="
python3 - <<'PY'
import ast, pathlib
for path in ["flashinfer/utils.py","flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py","flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py"]:
p=pathlib.Path(path)
src=p.read_text()
print(f"--- {path} ---")
tree=ast.parse(src)
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in {"_get_cache_buf","_get_cp_workspace"}:
print(f"function {node.name} line {node.lineno}")
for stmt in node.body:
if not isinstance(stmt, ast.FunctionDef):
print(" ", ast.unparse(stmt))
PYRepository: flashinfer-ai/flashinfer
Length of output: 23249
Avoid process-wide CP scratch buffers on the same device.
_get_cache_buf returns the same (name, device) tensor each call, so cp_delta_rule_dsl_sm90() and cp_delta_rule_dsl_sm120() can share t, transfer_t, state_t, fixed_state, and the prefill TMA buffer across concurrent invocations on one device or across caller-supplied streams. The TMA allocation is sized by num_sms * 128 and indexed by smid() * 128, so two kernels running on the same SM can overwrite each other’s TMA descriptor. Document a single-device, single-stream, or single-threaded-only contract, or key these buffers by stream/caller slot.
📍 Affects 2 files
flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py#L56-L58(this comment)flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py#L52-L54flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py#L5106-L5109flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py#L4186-L4187
🤖 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/delta_rule_cp_sm120.py` around lines 56
- 58, Prevent process-wide scratch-buffer aliasing in _get_cp_workspace and the
related CP kernel paths: key _get_cache_buf allocations by the caller stream or
an equivalent caller slot so concurrent cp_delta_rule_dsl_sm90() and
cp_delta_rule_dsl_sm120() invocations cannot share TMA or state buffers. Apply
this to flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm120.py lines 56-58
and 5106-5109, and flashinfer/gdn_kernels/delta_rule_dsl/delta_rule_cp_sm90.py
lines 52-54 and 4186-4187; preserve the existing per-SM TMA indexing while
ensuring each concurrent invocation has isolated storage.
Collaborator
|
/bot run tests/gdn |
Collaborator
19 tasks
Collaborator
|
[FAILED] Pipeline #61359033 — 16/18 executed test jobs passed Compared with nightly #61182354. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
guangyunh-nv
force-pushed
the
gdn-cp-host-launch-opt
branch
from
August 10, 2026 07:26
5030942 to
4257159
Compare
Collaborator
|
/bot run tests/gdn |
Collaborator
Collaborator
|
[FAILED] Pipeline #61920858 — 17/18 executed test jobs passed Compared with nightly #61799324. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
This was referenced Aug 12, 2026
5 tasks
guangyunh-nv
added a commit
that referenced
this pull request
Aug 24, 2026
Cache SM90 and SM120 non-CP kernel objects and compile options, and construct CuTe tensor wrappers only during compilation. Replay compiled kernels through raw TVM-FFI arguments to remove repeated launch preparation from the eager path. follow up of #4374 for non-CP launch path.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
This PR improves gdn cutedsl kernel launch warpper, the overhead is reduced by ~80%, on SM120 machine with CPU
AMD EPYC 9124 16-Core Processor, the total host luanch latency is reduced from 434.7 us to 98.2 us. This also applies for SM90 machine.What it does:
🚀 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.).Summary by CodeRabbit