Add DASC recurrent state sparsity policy

CHANGELOG.rst

@@ -13,6 +13,7 @@ Changelog
 *Misc*
 
 - A tracked ``examples/hf_ptq/hf_ptq.py`` run now writes ``.experiment.json`` into ``--export_path`` and uploads the same file with the run, so a checkpoint on disk names the experiment and MLflow run id that produced it. The pointer is written only once the export completes, and an export that is not tracked removes one it would otherwise inherit from a reused ``--export_path`` or from a quantized source checkpoint.
+- Add an experimental GDN DASC state-sparsity API that derives whole-head decay policies and selects a checkpoint-specific recovery window from caller-supplied quality and storage measurements. The first release exports restorable policy metadata only; serving runtimes must implement checkpoint packing and DASC-NR or DASC-WR recovery.
 
 **Backward Breaking Changes**
 

docs/source/guides/6_sparsity.rst

@@ -114,6 +114,94 @@ To restore the saved sparse model you can use
     Please see :ref:`saving and restoring of ModelOpt-modified models <save-restore>` to learn
     about all the available options for saving and restoring.
 
+Decay-aware recurrent-state sparsity (experimental)
+---------------------------------------------------
+
+The :mod:`modelopt.torch.sparsity.state_sparsity` package calibrates `DASC
+<https://arxiv.org/abs/2608.30386>`_ policies for persisted Gated DeltaNet (GDN) prefix state. It
+derives one static decay horizon per complete GDN head from
+``A_log`` and ``dt_bias``, then selects the largest caller-evaluated ``Wmax`` that passes every
+configured quality, lifecycle, and physical-storage gate. ``Wmax`` may be any positive integer.
+The measurements are evidence inputs produced by a caller-owned paired evaluation; this API does
+not run the dense-versus-recovery suffix evaluation itself. ``perplexity_retention`` is defined as
+``dense_perplexity / DASC_perplexity`` (equivalently ``exp(dense_NLL - DASC_NLL)``), so higher is
+better and values above one are valid.
+
+The initial API exports policy metadata only. It does not change model execution, quantize state,
+pack ragged checkpoints, replay a suffix, or add a linear-attention kernel. A serving integration
+must implement storage and recovery while preserving convolution state exactly and materializing
+ordinary dense recurrent state before continuation.
+
+.. code-block:: python
+
+    import modelopt.torch.sparsity.state_sparsity as mtss
+
+    config = {
+        "variant": "dasc_wr",  # "dasc_nr" uses zero recovery instead
+        "epsilon": 1e-3,
+        "static_gate_input": -0.3,
+        "wmax_candidates": [32],
+        # Set this to the dtype used to store A_log and dt_bias in the checkpoint.
+        "decay_parameter_storage_dtype": "bfloat16",
+        "model_id": "org/model",
+        "model_revision": "immutable-model-revision",
+        "model_config_id": "sha256:<config-digest>",
+        "calibration_data_id": "sha256:<dataset-and-protocol-digest>",
+    }
+    # Use these storage-canonical horizons to derive the evaluated mask and the
+    # retained_heads/total_heads measurement geometry for every Wmax candidate.
+    horizons = mtss.analyze_gdn_decay(
+        model,
+        epsilon=config["epsilon"],
+        static_gate_input=config["static_gate_input"],
+        decay_parameter_storage_dtype=config["decay_parameter_storage_dtype"],
+    )
+    measurements = [
+        {
+            "variant": "dasc_wr",
+            "wmax": 32,
+            "retained_heads": 40,
+            "total_heads": 96,
+            "checkpoint_savings": 0.21,
+            "quality": [
+                {
+                    "slice_id": "validation-context-1024",
+                    "perplexity_retention": 0.999,
+                    "top1_agreement": 0.99,
+                    "finite_continuation_logits": True,
+                    "retained_state_exact": True,
+                    "omitted_state_matches_recovery": True,
+                    "convolution_state_exact": True,
+                }
+            ],
+        },
+    ]
+
+    model = mtss.calibrate(model, config, measurements)
+    policy = mtss.export_policy(model)
+
+``dasc_nr`` and ``dasc_wr`` remain explicit deployment contracts: DASC-NR restores omitted heads
+from zero, while DASC-WR reconstructs them from a zero-initialized suffix replay of at most the
+selected ``Wmax`` tokens. Both retain whole GDN heads, preserve convolution state, and resume with
+dense recurrence. KDA and serving-runtime integration are not supported by this initial API.
+The initial GDN adapter accepts the ``GatedDeltaNet`` and ``Qwen3NextGatedDeltaNet`` base classes,
+including ModelOpt-generated dynamic subclasses, and fails closed for unrelated implementations
+even when they expose similarly named decay tensors.
+Re-running :func:`~modelopt.torch.sparsity.state_sparsity.calibrate` replaces the existing DASC
+mode-state entry and supersedes its stale policy without growing the checkpoint history. A policy
+with recoverable GDN geometry drift, decay drift, or temporarily unavailable decay tensors remains
+serializable, but removing or replacing the supported GDN architecture fails closed on both save
+and restore. In every stale-policy case,
+:func:`~modelopt.torch.sparsity.state_sparsity.export_policy` rejects it until recalibration.
+Set ``decay_parameter_storage_dtype`` to the checkpoint dtype for ``A_log`` and ``dt_bias`` before
+calibration. Derive the evaluated head masks and reported ``retained_heads``/``total_heads`` from
+:func:`~modelopt.torch.sparsity.state_sparsity.analyze_gdn_decay` using the same ``epsilon``,
+``static_gate_input``, and storage dtype passed to calibration. Policy validation allows only the
+rounding introduced by the declared storage dtype and the live tensor dtype. When they differ,
+distinct lossy inverse rounding bounds are composed in sequence; a duplicate dtype or an exact
+widening cast contributes no additional slack. Decay tensors that are live in BF16 or FP16 are
+still validated against that live dtype's rounding.
+
 .. _sparsity-concepts:
 
 Sparsity Concepts

modelopt/torch/sparsity/__init__.py

@@ -15,10 +15,12 @@
 
 """API for sparsification algorithms.
 
-This module provides access to both weight sparsity and attention sparsity algorithms.
+This module provides access to weight, attention, and recurrent-state sparsity algorithms.
 For backward compatibility, weight sparsity APIs are re-exported at the module level.
 """
 
+from . import state_sparsity
+
 # Import weight sparsity for backward compatibility
 from .weight_sparsity import mode, module, plugins
 from .weight_sparsity.sparsification import *

modelopt/torch/sparsity/state_sparsity/__init__.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Decay-aware sparsity policies for persisted recurrent state."""
+
+from . import mode  # imported for mode-registration side effects
+from .api import calibrate, export_policy
+from .config import (
+    DASCCalibrationMeasurement,
+    DASCConfig,
+    DASCLayerPolicy,
+    DASCPolicy,
+    DASCQualityMeasurement,
+)
+from .policy import analyze_gdn_decay, compute_gdn_decay_horizons
+
+__all__ = [
+    "DASCCalibrationMeasurement",
+    "DASCConfig",
+    "DASCLayerPolicy",
+    "DASCPolicy",
+    "DASCQualityMeasurement",
+    "analyze_gdn_decay",
+    "calibrate",
+    "compute_gdn_decay_horizons",
+    "export_policy",
+]

modelopt/torch/sparsity/state_sparsity/api.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Public DASC state-sparsity APIs."""
+
+import copy
+from collections.abc import Iterable
+from typing import Any
+
+from torch import nn
+
+from modelopt.torch.opt.conversion import ModeloptStateManager, apply_mode
+from modelopt.torch.utils import unwrap_model
+
+from .config import DASCCalibrationMeasurement, DASCConfig
+from .conversion import get_attached_dasc_policy, replace_dasc_mode
+from .mode import DASCModeRegistry
+from .policy import validate_dasc_decay_parameters, validate_dasc_model_structure
+
+__all__ = ["calibrate", "export_policy"]
+
+
+def calibrate(
+    model: nn.Module,
+    config: dict[str, Any] | DASCConfig,
+    measurements: Iterable[DASCCalibrationMeasurement | dict],
+) -> nn.Module:
+    """Calibrate and attach a DASC policy without changing model execution.
+
+    ``measurements`` must contain exactly one entry for every configured ``Wmax`` candidate.
+    The largest candidate passing every quality, lifecycle, and storage gate is selected.
+    Recalibrating replaces the existing DASC mode-state entry in place.
+
+    Example::
+
+        import modelopt.torch.sparsity.state_sparsity as mtss
+
+        model = mtss.calibrate(model, config, measurements)
+        deployment_policy = mtss.export_policy(model)
+
+    Args:
+        model: Model containing GatedDeltaNet modules with one-dimensional ``A_log`` and
+            ``dt_bias`` tensors.
+        config: Checkpoint provenance, candidate windows, and quality gates.
+        measurements: Quality and checkpoint-storage results produced by the caller's paired
+            dense-versus-DASC calibration workflow.
+
+    Returns:
+        The input model with a serializable DASC policy attached through ModelOpt state.
+    """
+    model = unwrap_model(model, force_unwrap=True)
+    config_object = config if isinstance(config, DASCConfig) else DASCConfig(**config)
+    if ModeloptStateManager.is_converted(model, is_root=True) and any(
+        mode == "dasc" for mode, _ in ModeloptStateManager(model).state_dict()
+    ):
+        return replace_dasc_mode(model, config_object, measurements)
+
+    return apply_mode(
+        model,
+        mode=[("dasc", config_object.model_dump())],
+        registry=DASCModeRegistry,
+        mode_kwargs={"measurements": measurements},
+    )
+
+
+def export_policy(model: nn.Module) -> dict[str, Any]:
+    """Export a JSON-safe DASC policy after validating model structure and decay parameters.
+
+    This policy does not implement checkpoint packing or recovery. A serving backend must preserve
+    convolution state, store retained complete GDN heads, recover omitted heads according to the
+    declared variant, and materialize the ordinary dense runtime state before continuation.
+    """
+    policy = get_attached_dasc_policy(model)
+    validate_dasc_model_structure(model, policy)
+    validate_dasc_decay_parameters(model, policy)
+    return copy.deepcopy(policy.model_dump(mode="json"))