Conversation
Copy fla/ops/common/chunk_delta_h.py (state recurrence kernels) and fla/ops/gated_delta_rule/chunk.py (autograd wrapper) from fla-org/flash-linear-attention@516143e31fce into modelopt/torch/kernels/quantization/linear_attention/ as a base for a fake-quantized recurrent state. Apart from the license headers and ruff formatting, the only edit is that the wrapper imports the state kernels from the vendored sibling module instead of fla; the remaining fla operators are still imported from the installed package. Add the MIT copyright holder to LICENSE, exclude the copies from the license-header hook, and exempt them from docstring lint and strict mypy. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
…nked kernel Add STATE_QDQ to the vendored fla state-recurrence kernel: at the end of every chunk (and on a provided initial state) the carried fp32 state tiles are quantized to E4M3 and dequantized with one dynamic scale per [K, BV] tile of a head, so the next chunk, the per-chunk h buffer and the final state all see the FP8-rounded state. The backward recomputes the same quantized states and passes the state gradient straight through. BV is an explicit launch argument (default 64: two scales per 128-wide head; 128 gives one per head but spills where the kernel is limited to two warps), so the granularity no longer depends on the autotuner; fla's backend dispatch is bypassed so no other backend can take a quantized call. Expose it through a default-disabled gdn_state_quantizer on Megatron-Core GatedDeltaNet modules (recipe unit configs/ptq/units/gdn_state_fp8_dynamic: e4m3, dynamic, axis [0, 1]); the quantizer only carries the configuration and the vendored kernel is imported on first use. Requires flash-linear-attention >= 0.5.1 and Triton. CPU unit tests pass; GPU kernel and Megatron tests are added but were not run here. The kernel compiles offline for sm_90a and sm_100a with the same register and spill footprint as the unquantized fla kernel. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
Let chunk_gated_delta_rule take a callable that fake-quantizes the WY tensor w ([B, T, HV, K]) once per forward, right after it is materialized and before it multiplies the state, so the recurrence emulates an FP8 x FP8 matmul against the FP8 state. The backward recompute applies the same quantizer and the gradient passes straight through. Because w lands in memory between kernels this needs no Triton change; a ModelOpt TensorQuantizer (e.g. dynamic per-token E4M3) works as the callable, as the added GPU test shows. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (15)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe change adds chunked GatedDeltaNet Triton kernels with optional FP8 fake quantization for recurrent state and WY values. It integrates the kernels with ModelOpt and Megatron, adds PTQ recipes and documentation, and introduces unit and GPU tests. ChangesGatedDeltaNet FP8 quantization
Priority: ⬇️ Low Estimated code review effort: 4 (Complex) | ~60 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant MegatronGatedDeltaNet
participant GatedDeltaNetStateQuantMixin
participant chunk_gated_delta_rule
participant fla_chunk_delta_h
MegatronGatedDeltaNet->>GatedDeltaNetStateQuantMixin: route forward execution
GatedDeltaNetStateQuantMixin->>chunk_gated_delta_rule: pass state and WY quantizers
chunk_gated_delta_rule->>fla_chunk_delta_h: dispatch quantized forward kernel
fla_chunk_delta_h-->>chunk_gated_delta_rule: return outputs and final state
chunk_gated_delta_rule-->>MegatronGatedDeltaNet: return quantized results
Merge Risk: ⚪ Minimal · up to No concrete current-head issue remains that should block merging after normal checks. 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 34.15% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 41 functions across 9 files. (6 skipped: 6 unsupported.)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2455 +/- ##
==========================================
+ Coverage 71.45% 77.71% +6.26%
==========================================
Files 590 593 +3
Lines 64754 66456 +1702
==========================================
+ Hits 46267 51645 +5378
+ Misses 18487 14811 -3676
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kaix-nv
approved these changes
Sep 18, 2026
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
2 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.
What does this PR do?
Type of change: ?
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information
Summary by CodeRabbit
New Features
Documentation
Tests