Repository navigation
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
kaix-nv
added a commit
that referenced
this pull request
Oct 2, 2026
<!-- linear-attention-stack:start --> **Linear-attention PR stack — 6 PRs** | Order | PR | Depends on | | --- | --- | --- | | 1/6 | [#2497 GDN state/W QAT foundation](#2497) | main | | 2/6 | [#2519 Torch GDN/KDA decode QAT + INT8](#2519) | #2497 | | 3/6 | [#2562 Fused Triton GDN/KDA decode QAT](#2562) | #2519 | | 4/6 | [#2541 vLLM GDN/KDA state-only fake quantization](#2541) | #2562 | | 5/6 | [#2503 GDN/KDA prefill GEMM quantization](#2503) | #2541 | | 6/6 | [#2507 Experimental GDN/KDA approximate inverse](#2507) | #2503 | All six PRs form native GitHub stack #2563 in the order shown above. #2541 applies TensorQuantizer before native vLLM prefill/decode calls. A separate vLLM prefill-GEMM PR waits for an optimized fused kernel. #2506 and #2509 are superseded and closed. <!-- linear-attention-stack:end --> ### What does this PR do? Type of change: new feature GatedDeltaNet training keeps recurrent states inside a chunked kernel, so projection quantizers cannot emulate rounding at state boundaries. This PR adds dynamic per-tile FP8 E4M3 fake QDQ to the recurrent state and independent dynamic FP8 fake QDQ to WY-transformed W activations, with identity straight-through gradients for QAT/QAD. Both sites use the standard `quant_cfg` interface and start disabled. State QDQ uses 64-token chunks and recomputes `amax` at each boundary over each full-key by 64-value-column tile, independently per sequence and head. Each tile has its own scalar scale (`amax / 448`, with a zero guard); `fp8_scalar_qdq` applies that supplied scale rather than choosing tensor-wide grouping. W grouping is applied by `TensorQuantizer`. Quantizer settings use normal ModelOpt checkpoint state. There is no `QuantizeConfig.linear_attention` field in this PR; #2519 introduces execution policies for decode and ReplaySSM, and later PRs extend them for prefill and approximate inverse. Configurations or checkpoints from earlier experimental drafts that use those execution policies require #2519; those selecting Triton decode also require #2562. The Megatron adapter supports the direct-forward and older split-forward call layouts, restores the original kernel when disabled, and removes temporary quantizer attributes on export. Independent recurrent/chunk numerical references live under `tests/_test_utils/torch/quantization/`; shared runtime capability checks live in `linear_attention/utils.py`. The fused path requires `fla-core==0.5.1` and chunk size 64. State FP8 emulation requires SM89 or newer. The Hopper path has additional dtype/TileLang restrictions enforced before launch. This PR simulates numerical error; it does not add compressed state storage or faster inference. ### Usage ```python import modelopt.torch.quantization as mtq model = mtq.quantize(model, { "quant_cfg": [ {"quantizer_name": "*", "enable": False}, {"quantizer_name": "*gdn_state_quantizer", "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1)}}, {"quantizer_name": "*gdn_w_quantizer", "cfg": {"num_bits": (4, 3), "type": "dynamic", "axis": (0, 1, 2)}}, ], "algorithm": None, }) # Continue with the framework's normal forward/backward/optimizer steps. ``` Dynamic scales require no calibration. ### Testing The focused GPU suite contains four cases: three BF16 numerical forward/backward checks (disabled, W QDQ, and state+W QDQ) using one shared shape, plus one single-rank, one-layer Megatron QAT/checkpoint test. The Megatron test checks quantizer enable/disable behavior, checkpoint restore, gradients, and an optimizer update; it enables state QDQ when the GPU supports native FP8 conversion. Compilation runs in setup fixtures, and functional calls retain the normal 120-second timeout. There are no dtype, layout, tile-width, or parallelism sweeps. The pinned FLA/TileLang/TVM-FFI dependencies live in the `dev-fla` optional extra, installed by both GPU nox sessions. Validation of the consolidated changes on RTX A6000 (SM86), Python 3.12.8, Torch 2.9.1+cu128, Triton 3.5.1, fla-core 0.5.1, TileLang 0.1.8, Megatron Core 0.19.2, and Transformer Engine 2.16.0: - Cold and warm focused runs: **3 passed, 1 hardware skip** each. The state+W numerical case requires SM89+; the local Megatron test exercised W QDQ. - Fresh Triton/TileLang cache: **363.09s total**, including setup and teardown. Kernel setup took 66.38s + 44.46s; Megatron setup, including shared extension setup and worker startup, took 245.26s. Functional calls totaled about 2.56s. - Same cache, new pytest process: **38.20s total**, with about **2.41s in functional calls**. - Pre-commit checks passed for the four changed files. Dependency-group wiring and installed pinned versions were checked. ```bash PYTHONPATH=. python -m pytest -q \ tests/gpu/torch/kernels/quantization/linear_attention/test_fla_chunk_gated_delta_rule.py \ tests/gpu_megatron/torch/quantization/plugins/test_megatron_gated_delta_net.py \ --durations=0 ``` These timings describe local test setup and execution, not inference performance. Native FP8 state QDQ and Hopper still require suitable GPU/CI runs. This minimal suite does not qualify tensor/context/pipeline parallelism, checkpoint resharding, or model-quality recovery. Mamba compilation coverage is tracked separately in #2572. ### Before your PR is "*Ready for review*" Contributor and security guidance reviewed. Commits are signed and signed off. - Is this change backward compatible?: ✅ Disabled-by-default quantizers, standard-recipe exclusions, and legacy-checkpoint coverage; enabled experimental configurations have explicit capability restrictions. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ❌ Internal third-party approval tracking still needs confirmation. Upstream attribution, MIT/Apache headers, `LICENSE` notice, and license-hook exclusions are included. FLA/TileLang and TVM-FFI license files were reviewed. - Did you write any new necessary tests?: ✅ Numerical, gradient, conversion/checkpoint, and real framework tests. - Did you update Changelog?: ✅ Experimental quantization feature entry. - Did you get Claude approval on this PR?: ❌ Bot feedback addressed or discussed; renewed approval pending. ### Additional Information Related: #2455. This is the first integration slice and does not assume #2455 has merged. Later milestones will extend the numerical boundaries after choosing their approximation contracts. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added experimental dynamic FP8 fake quantization for GatedDeltaNet recurrent states and WY activations during training. * Added PTQ configuration options for state and WY activation quantization. State quantization requires an SM89-or-newer GPU; the fused path requires `fla-core==0.5.1` and a chunk size of 64. * **Bug Fixes** * Improved quantizer configuration validation and restoration for linear-attention models. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Kai Xu <kaix@nvidia.com>
This branch has not been deployed
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