Repository navigation
Conversation
Signed-off-by: Kai Xu <kaix@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. |
This was referenced Sep 22, 2026
kaix-nv
added this pull request to stack #2510
September 22, 2026 21:27
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## kaix/linear-attention-qat-m2 #2506 +/- ##
================================================================
+ Coverage 70.49% 70.51% +0.01%
================================================================
Files 612 615 +3
Lines 67875 68025 +150
================================================================
+ Hits 47850 47969 +119
- Misses 20025 20056 +31
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
removed this pull request from stack #2510
September 23, 2026 01:07
Contributor
|
This was referenced Sep 23, 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 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
1 participant
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.
Superseded: KDA prefill is now included in #2503, after the shared GDN/KDA decode infrastructure in #2519. This closed PR and its branch retain the original review and qualification history.
What does this PR do?
Type of change: New feature, new example, new tests, documentation.
Adds experimental training-time numerical emulation to FLA 0.5.1 KimiDeltaAttention. ModelOpt conversion creates disabled-by-default KDA state/W and eight prefill operand-site controls. Enabling a policy routes the existing layer through differentiable FP32 materialized prefill with dynamic FP8 or NVFP4 QDQ; disabling all controls preserves the original FLA forward. This is stacked on #2503.
KDA uses per-key-channel retention. Causal pairwise exponentials are computed from prefix differences, avoiding an overflowing inverse-decay factorization. The adapter retains the model's projections, short convolution, gate formulas, and normalization, and preserves gradients and saved numerical policies. Invocation-local dispatch also works for copied modules without changing global FLA functions.
A pinned public checkpoint/WikiText runner records matched pre/post-training NLL, token hashes, source hashes, optimizer updates, and time/memory. The exact solve remains the default. Approximate inverse and decode/replay training are subsequent milestones.
Usage
Install the example's pinned FLA dependencies only when using this adapter. Core ModelOpt imports do not require FLA.
Testing
Revision:
e3987d01bb2b54188bb18165b478972a101ac54e(signed). Remote staged source hashes matched the implementation.arcee-ai/AFM-4.5B-Base-KDA-Onlyat01ad2e06ee4f1214193c17b69e09105a9b257e80, WikiText-2 raw revisionb08601e04326c79dfdd32d625aee71d232d685c3. All checkpoint keys loaded, all trainable attention gradients were finite, and an optimizer update changed query weights. Runs used identical token hashes/seed, one 128-token training step and four validation blocks (512 predicted tokens), attention-only FP32 AdamW under BF16 autocast with activation checkpointing.These are one-step integration results, not evidence of quality recovery. The model payload has 36 KDA layers; configuration and weight index were checked rather than relying on the model-card layer count. Original FP32 gate parameters are preserved during loading.
H100 environment: NeMo 26.08, Torch 2.13.0a0+8145d630e8.nv26.06, CUDA 13.3, Triton 3.7, FLA/core 0.5.1, Transformers 4.57.1, TileLang 0.1.8, TVM-FFI 0.1.9, Megatron-Core 0.19.1, Transformer Engine 2.18.0. Numerical qualification excludes TF32.
Measured KDA training overhead
One H100, shape
[1,1024,4,64], BF16 Q/K/V, FP32 working math, TF32 off. Three warmups and 20 interleaved trials; median times and paired wall ratios. Extra memory is peak tensor allocation above the pre-forward baseline. Compilation is excluded; samples and source hashes retained.The per-channel interaction tensors account for substantial materialization cost. This is an explicit numerical research path, with no speed claim.
The model-study command and pinned dependencies are in
examples/llm_qat/linear_attention/README.md; run once with no quant config, once withkda_exact_matmul.json, and once withkda_prefill_fp8.json.Before your PR is "Ready for review"
CONTRIBUTING.md: ✅ Uses existing optional FLA APIs; example dependencies are permissively licensed and pinned where needed. No new mandatory core dependency.Additional Information
The materialized path is a numerical research implementation. It does not claim native low-precision execution, serving acceleration, or model-quality recovery. Qualification uses FLA's general gate path; the optional safe_gate TensorCore rounding is not emulated. CP, recurrent inference, and FLA intermediate/recompute flags are rejected. FLA selects recurrent inference for evaluation sequences of at most 64 tokens; these QAT prefill studies use longer sequences.