fix(gdn): handle padded inputs in SM100 CP prefill - #4428
samuellees wants to merge 1 commit into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between a9343a21a77823c8f1e3279683a4a7bc68b8bc41 and 4deb07f. 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe prefill path allocates per-CTA tensor-map workspace for Blackwell context-parallel GDN execution. The kernel updates K/V descriptors to logical sequence bounds before TMA loads. A regression test covers poisoned physical padding. ChangesContext-parallel GDN tensor maps
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Prefill as cp_delta_rule_mn_precompute_dsl_sm100
participant Kernel as CPDeltaRuleMNPrecomputeUtcmma1Sm100
participant TMA as TMA warp
participant Loads as K/V TMA loads
Prefill->>Kernel: Pass tensor-map workspace
Kernel->>TMA: Derive per-CTA descriptor pointers
TMA->>TMA: Bound descriptors by seq_end
TMA->>Loads: Pass updated descriptors
Loads->>Loads: Execute dynamic TMA copies
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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 |
![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_kernels/blackwell/gated_delta_net_cp.py (1)
25-27: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the per-CTA TensorMap workspace design.
Explain why each CTA needs two independent GMEM descriptors. Explain why a shared or static descriptor cannot support sequence-specific K/V bounds. State that the alternative is to retain the original physical TensorMap, which can read undefined padding.
As per coding guidelines, “For performance-critical hot paths, document the rationale for special algorithmic choices and relevant alternatives in comments.”
Also applies to: 210-226
🤖 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 25 - 27, Document the per-CTA TensorMap workspace near bytes_per_tensormap and num_tensormaps, explaining that each CTA requires two independent GMEM descriptors for sequence-specific K/V bounds; clarify that shared or static descriptors cannot represent those bounds, and that the alternative is retaining the original physical TensorMap, which may read undefined padding.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_kernels/blackwell/gdn_cp_prefill.py`:
- Around line 305-314: The TensorMap workspace in the prefill path is shared
across overlapping streams because `_get_cache_buf` uses only a name and device
key. Update the workspace management around `tensormap_workspace` in the
function accepting `_stream` so each in-flight stream or launch receives an
isolated buffer, reclaiming or reusing it only after prior stream work
completes. Add a regression test that overlaps calls on different streams with
different sequence ends and verifies their results.
---
Nitpick comments:
In `@flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.py`:
- Around line 25-27: Document the per-CTA TensorMap workspace near
bytes_per_tensormap and num_tensormaps, explaining that each CTA requires two
independent GMEM descriptors for sequence-specific K/V bounds; clarify that
shared or static descriptors cannot represent those bounds, and that the
alternative is retaining the original physical TensorMap, which may read
undefined padding.
🪄 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: 03d16ce7-f94d-41d2-a9a8-38dcebf5c7ae
📥 Commits
Reviewing files that changed from the base of the PR and between 4fbac49 and a9343a21a77823c8f1e3279683a4a7bc68b8bc41.
📒 Files selected for processing (3)
flashinfer/gdn_kernels/blackwell/gated_delta_net_cp.pyflashinfer/gdn_kernels/blackwell/gdn_cp_prefill.pytests/gdn/test_prefill_cp_delta_rule.py
| tensormap_workspace_size = ( | ||
| num_sab_heads | ||
| * max_cp_chunks_per_seq | ||
| * num_seqs | ||
| * CPDeltaRuleMNPrecomputeUtcmma1Sm100.num_tensormaps | ||
| * CPDeltaRuleMNPrecomputeUtcmma1Sm100.bytes_per_tensormap | ||
| ) | ||
| tensormap_workspace = _get_cache_buf( | ||
| "gdn_cp_sm100_mn_tensormaps", tensormap_workspace_size, device | ||
| )[:tensormap_workspace_size] |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Isolate TensorMap workspaces for overlapping streams.
_get_cache_buf keys this mutable workspace only by name and device. This function accepts _stream, so concurrent calls on different streams can update the same K/V descriptors while another launch still uses them. This can produce incorrect TMA loads or a device fault.
Give each in-flight launch or stream an independent workspace. Reuse a workspace only after its prior stream work completes. Add a regression test with overlapping calls that use different sequence ends.
🤖 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 305 - 314,
The TensorMap workspace in the prefill path is shared across overlapping streams
because `_get_cache_buf` uses only a name and device key. Update the workspace
management around `tensormap_workspace` in the function accepting `_stream` so
each in-flight stream or launch receives an isolated buffer, reclaiming or
reusing it only after prior stream work completes. Add a regression test that
overlaps calls on different streams with different sequence ends and verifies
their results.
Framework inputs may contain physical padding whose values are undefined. cu_seqlens already contains the logical sequence ends. The non-CP path respects those ends, while the SM100 CP MN-precompute path did not. Bound the CP K/V TensorMaps with the device-side logical ends so TMA zero-fills the padded tail. The public API is unchanged. Co-authored-by: Sam Li <lsam@nvidia.com>
a9343a2 to
4deb07f
Compare
|
I think this is a wont fix. Slice the input on host is trivial. |
4 participants
Frameworks may pass padded Q/K/V buffers whose padding values are undefined. The non-CP path already handles this using device-side cu_seqlens, but the SM100 CP path did not.
This change adds the same handling to the CP kernel, with no API change and negligible performance impact.
cc @guangyunh-nv
📌 Description
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ 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
Bug Fixes
Tests