fix(gdn): address remaining CodeRabbit feedback from #3001#3165
Conversation
arpera
requested review from
bkryu,
kahyunnam,
yongwww and
yzh119
as code owners
📝 WalkthroughWalkthroughOptimized memory allocation in GDN prefill operations by making Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
![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 refactors the chunk_gated_delta_rule function to optimize memory allocation and ensure numerical consistency across different GPU architectures. The allocation of output_state has been moved into backend-specific logic, allowing the SM100 path to skip allocation when the final state is not requested, thereby reducing memory overhead. Additionally, the scale normalization logic was updated to treat scale=0.0 as the default value, ensuring consistent behavior between SM90 and SM100 kernels. I have no feedback to provide as the changes correctly address the specific requirements of the respective hardware backends.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
flashinfer/gdn_prefill.py (2)
273-273: Docstring drift: documentscale=0.0as sentinel.The normalization now treats both
Noneand exactly0.0as "use the default1/sqrt(head_size)", which aligns SM100 numerics with the SM90 C++ kernel's sentinel convention (as intended by this PR). However, the parameter docstring at lines 139-141 still says the default applies only whenscaleis "not provided", so a caller reading the docs might be surprised that explicitly passing0.0is silently normalized.Consider updating the docstring for clarity.
📝 Proposed docstring tweak
scale (Optional[float]): Scale factor for the attention scores. - If not provided, defaults to ``1 / sqrt(head_size)``. Default: ``None``. + If ``None`` or ``0.0``, defaults to ``1 / sqrt(head_size)`` + (``0.0`` is treated as a sentinel for parity with the SM90 kernel + which interprets ``0.0`` as "use default"). Default: ``None``.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/gdn_prefill.py` at line 273, Update the docstring for the parameter named "scale" in the function that computes _scale to state that both None and an explicit 0.0 are treated as the sentinel to use the default 1/sqrt(head_size); mention the normalization logic (_scale = scale if scale is not None and scale != 0.0 else 1.0 / math.sqrt(head_size)) and clarify that passing 0.0 will be interpreted the same as omitting the parameter so callers aren't surprised.
282-291: SM100 conditional allocation — LGTM, with one behavior note.The change correctly avoids the wasted
[num_seqs, H, 128, 128]float32 scratch buffer on the CuTe-DSL path whenoutput_final_state=False, and theelif output_state is Nonebranch preserves a caller-supplied tensor when they actually want the final state. Matches the PR rationale.One observable asymmetry worth being aware of: if a caller passes
output_state=<tensor>together withoutput_final_state=False, the SM100 path silently drops the tensor (set toNonehere), while the SM90 path at lines 329-334 keeps and writes into it (because the C++ kernel writes unconditionally). The function only returnsoutput_statewhenoutput_final_state=True, so per the documented contract this is fine, but the two backends now have a subtly different treatment of a misuse pattern. Not blocking — just flagging in case a stricter validation (e.g., warn/reject whenoutput_stateis provided butoutput_final_state=False) is preferred for consistency.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@flashinfer/gdn_prefill.py` around lines 282 - 291, The current SM100 conditional allocation silently drops a caller-provided output_state when output_final_state is False, creating an asymmetry with the SM90 path; add a validation in the same function around the output_final_state/output_state logic to detect when output_state is not None but output_final_state is False and either log a warning (processLogger.warn or warnings.warn) or raise a ValueError to reject the misuse, referencing the output_final_state and output_state variables so callers are informed and behavior is consistent across backends.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@flashinfer/gdn_prefill.py`:
- Line 273: Update the docstring for the parameter named "scale" in the function
that computes _scale to state that both None and an explicit 0.0 are treated as
the sentinel to use the default 1/sqrt(head_size); mention the normalization
logic (_scale = scale if scale is not None and scale != 0.0 else 1.0 /
math.sqrt(head_size)) and clarify that passing 0.0 will be interpreted the same
as omitting the parameter so callers aren't surprised.
- Around line 282-291: The current SM100 conditional allocation silently drops a
caller-provided output_state when output_final_state is False, creating an
asymmetry with the SM90 path; add a validation in the same function around the
output_final_state/output_state logic to detect when output_state is not None
but output_final_state is False and either log a warning (processLogger.warn or
warnings.warn) or raise a ValueError to reject the misuse, referencing the
output_final_state and output_state variables so callers are informed and
behavior is consistent across backends.
jiahanc
left a comment
jiahanc
left a comment
There was a problem hiding this comment.
LGTM, thanks for the fix
|
/bot run |
|
Mirroring Failed Failed to mirror PR to GitLab. Check logs for details. |
|
/bot run |
|
@kahyunnam, could you tell me please, what is the problem with CI? Is it due to my change? |
|
4 participants
📌 Description
Addresses the two remaining CodeRabbit findings on #3001 that weren't applied before merge:
Normalize
scale=0.0to the default1/sqrt(d_k)before backend dispatch so the same call gives matching numerics on SM90 and SM100. The SM90 C++ kernel treats0.0as a sentinel for "use default", but the SM100 CuTe-DSL kernel forwarded the literal0.0→ zeroed QK → broken attention.Don't eagerly allocate
output_stateon the SM100 path whenoutput_final_state=False. The CuTe-DSL kernel drops the buffer anyway, so the old code wasted a full[num_seqs, H, 128, 128]float32 scratch per call. SM90 still allocates unconditionally because its C++ kernel always writes intooutput_state.Dispatcher callsites now pass
output_statedirectly on both branches (no inlineoutput_state if output_final_state else None), so SM90 and SM100 read identically.🔍 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
Improvements