Remove target_sparsity_ratio mode

Signed-off-by: Kai Xu <kaix@nvidia.com>

Author: kaix-nv

examples/llm_sparsity/kv_cache_sparsity/hf_triattention.py

@@ -27,9 +27,6 @@
     # Fixed-size budget (retain top-K tokens per head)
     python hf_triattention.py --model Qwen/Qwen3-0.6B --budget 2048
 
-    # Percentile-based eviction (evict 70% at each prune step)
-    python hf_triattention.py --model Qwen/Qwen3-0.6B --target-sparsity-ratio 0.7
-
     # With custom budget and calibration length
     python hf_triattention.py --model Qwen/Qwen3-0.6B --budget 1024 --calib-seq-len 4096
 
@@ -119,20 +116,11 @@ def main(args):
     print(f"  Calibration complete in {elapsed:.1f}s")
 
     # Step 2: Apply KV cache sparsity mode
-    if args.target_sparsity_ratio is not None:
-        print(
-            f"\nApplying TriAttention mode (target_sparsity_ratio={args.target_sparsity_ratio})..."
-        )
-        triattention_config = TriAttentionConfig(
-            target_sparsity_ratio=args.target_sparsity_ratio,
-            prune_interval=args.prune_interval,
-        )
-    else:
-        print(f"\nApplying TriAttention mode (budget={args.budget})...")
-        triattention_config = TriAttentionConfig(
-            budget=args.budget,
-            prune_interval=args.prune_interval,
-        )
+    print(f"\nApplying TriAttention mode (budget={args.budget})...")
+    triattention_config = TriAttentionConfig(
+        budget=args.budget,
+        prune_interval=args.prune_interval,
+    )
     model = mtskv.sparsify(model, triattention_config)
     print("  Mode applied (no-op on weights).")
 
@@ -184,22 +172,12 @@ def main(args):
         default="Qwen/Qwen3-0.6B",
         help="HuggingFace model name or local path.",
     )
-    policy = parser.add_mutually_exclusive_group()
-    policy.add_argument(
+    parser.add_argument(
         "--budget",
         type=int,
-        default=None,
+        default=2048,
         help="KV token budget (tokens to retain per head). "
-        "Mutually exclusive with --target-sparsity-ratio. "
-        "Defaults to 2048 if neither is set.",
-    )
-    policy.add_argument(
-        "--target-sparsity-ratio",
-        type=float,
-        default=None,
-        help="Fraction of tokens to evict at each prune step, in (0, 1). "
-        "Example: 0.7 evicts 70%% of tokens (keeps top 30%% by score). "
-        "Mutually exclusive with --budget.",
+        "Compression triggers after --prune-interval additional tokens.",
     )
     parser.add_argument(
         "--prune-interval",
@@ -229,6 +207,4 @@ def main(args):
     )
 
     args = parser.parse_args()
-    if args.budget is None and args.target_sparsity_ratio is None:
-        args.budget = 2048
     main(args)

examples/speculative_decoding/eagle_utils.py

@@ -25,7 +25,6 @@
 import transformers
 from datasets import load_dataset
 from packaging.version import Version
-from scripts.ar_validate import validate_ar
 from transformers import Trainer, TrainerCallback
 
 import modelopt
@@ -41,6 +40,7 @@
     ShardedDataset,
     VisionLanguageDataCollator,
 )
+from scripts.ar_validate import validate_ar
 
 try:
     import wandb

modelopt/torch/sparsity/kv_cache/config.py

@@ -33,29 +33,19 @@ class TriAttentionConfig(ModeloptBaseConfig):
     pre-RoPE Q/K concentration. Calibration computes per-head frequency statistics;
     at runtime, the serving engine scores and evicts tokens periodically.
 
-    Exactly one of ``budget`` or ``target_sparsity_ratio`` must be set:
-
-    - ``budget``: absolute token count to retain per head (fixed-size cache).
-    - ``target_sparsity_ratio``: fraction of tokens to evict at each pruning step.
-      Cache size auto-scales with generation length. Value in (0, 1).
+    ``budget`` is the absolute token count to retain per head after each pruning
+    round. Runtime compression follows slack/sawtooth semantics: grow until
+    ``budget + prune_interval`` tokens are present, then evict back to
+    ``budget``.
     """
 
-    # Eviction policy (exactly one must be set)
+    # Eviction policy
     budget: int | None = ModeloptField(
         default=None,
         title="KV token budget (absolute).",
         description=(
-            "Number of KV tokens to retain per head after pruning. "
-            "Mutually exclusive with target_sparsity_ratio."
-        ),
-    )
-    target_sparsity_ratio: float | None = ModeloptField(
-        default=None,
-        title="Target sparsity ratio (percentile-based).",
-        description=(
-            "Fraction of tokens to evict at each pruning step, in (0, 1). "
-            "Example: 0.7 means evict 70% of tokens (keep top 30% by score). "
-            "Mutually exclusive with budget."
+            "Number of KV tokens to retain per head after pruning. Runtime "
+            "compression triggers after prune_interval additional tokens."
         ),
     )
 
@@ -127,16 +117,10 @@ def validate_score_aggregation(cls, v: str) -> str:
         return v
 
     @model_validator(mode="after")
-    def validate_budget_or_sparsity(self) -> TriAttentionConfig:
-        """Exactly one of budget or target_sparsity_ratio must be set."""
-        budget_set = self.budget is not None
-        sparsity_set = self.target_sparsity_ratio is not None
-        if not budget_set and not sparsity_set:
-            raise ValueError("Must set exactly one of 'budget' or 'target_sparsity_ratio'")
-        if budget_set and sparsity_set:
-            raise ValueError("Cannot set both 'budget' and 'target_sparsity_ratio'; pick one")
-        if sparsity_set and not (0.0 < self.target_sparsity_ratio < 1.0):
-            raise ValueError(
-                f"target_sparsity_ratio must be in (0, 1), got {self.target_sparsity_ratio}"
-            )
+    def validate_budget(self) -> TriAttentionConfig:
+        """Validate the fixed KV token budget."""
+        if self.budget is None:
+            raise ValueError("TriAttention requires 'budget' to be set")
+        if self.budget <= 0:
+            raise ValueError(f"budget must be positive, got {self.budget}")
         return self

modelopt/torch/sparsity/kv_cache/triattention/scoring.py

@@ -55,10 +55,11 @@ def compute_frequency_statistics_from_means(
 
     Args:
         q_mean_complex: Mean of Q in complex frequency domain, shape (freq_count,).
-        q_abs_mean: Mean of |Q| in frequency domain, shape (freq_count,).
+        q_abs_mean: Mean of ``|Q|`` in frequency domain, shape (freq_count,).
         k_unrot: Unrotated key vectors, shape (num_keys, head_dim).
         style: RoPE pairing style.
-        disable_mlr: If True, use q_abs_mean directly instead of (q_abs_mean - |q_mean|).
+        disable_mlr: If True, use ``q_abs_mean`` directly instead of
+            ``q_abs_mean - |q_mean|``.
 
     Returns:
         amp: Amplitude, shape (num_keys, freq_count).
@@ -135,37 +136,27 @@ def select_keys_to_keep(
     scores: torch.Tensor,
     *,
     kv_budget: int | None = None,
-    target_sparsity_ratio: float | None = None,
 ) -> torch.Tensor:
     """Select which keys to retain based on importance scores.
 
-    Exactly one of ``kv_budget`` or ``target_sparsity_ratio`` must be provided.
-
     Args:
         scores: Importance scores, shape (num_keys,). Higher = more important.
         kv_budget: Absolute number of tokens to retain. Keeps top-K.
             If budget >= num_keys, keeps all.
-        target_sparsity_ratio: Fraction of tokens to evict, in (0, 1).
-            Keeps top (1 - ratio) fraction. Example: 0.7 → keep top 30%.
 
     Returns:
         Boolean mask, shape (num_keys,). True = keep, False = evict.
     """
-    budget_set = kv_budget is not None
-    sparsity_set = target_sparsity_ratio is not None
-    if budget_set == sparsity_set:
-        raise ValueError(
-            "select_keys_to_keep requires exactly one of kv_budget or target_sparsity_ratio"
-        )
+    if kv_budget is None:
+        raise ValueError("select_keys_to_keep requires kv_budget")
+    if kv_budget <= 0:
+        raise ValueError(f"kv_budget must be positive, got {kv_budget}")
 
     num_keys = scores.shape[0]
     if num_keys == 0:
         return torch.zeros(0, dtype=torch.bool, device=scores.device)
 
-    if budget_set:
-        k = min(kv_budget, num_keys)
-    else:
-        k = max(1, round(num_keys * (1.0 - target_sparsity_ratio)))
+    k = min(kv_budget, num_keys)
 
     if k >= num_keys:
         return torch.ones(num_keys, dtype=torch.bool, device=scores.device)

tests/unit/torch/sparsity/kv_cache/test_config.py

@@ -25,44 +25,24 @@ def test_budget_only():
     """Setting only budget is valid."""
     config = TriAttentionConfig(budget=2048)
     assert config.budget == 2048
-    assert config.target_sparsity_ratio is None
 
 
-def test_target_sparsity_only():
-    """Setting only target_sparsity_ratio is valid."""
-    config = TriAttentionConfig(target_sparsity_ratio=0.7)
-    assert config.budget is None
-    assert config.target_sparsity_ratio == 0.7
-
-
-def test_both_budget_and_sparsity_raises():
-    """Setting both budget and target_sparsity_ratio raises."""
-    with pytest.raises(ValidationError, match="Cannot set both"):
+def test_target_sparsity_ratio_is_not_supported():
+    """Ratio-based eviction is not part of the TriAttention API."""
+    with pytest.raises(ValidationError):
         TriAttentionConfig(budget=2048, target_sparsity_ratio=0.7)
 
 
-def test_neither_budget_nor_sparsity_raises():
-    """Setting neither budget nor target_sparsity_ratio raises."""
-    with pytest.raises(ValidationError, match="Must set exactly one"):
+def test_missing_budget_raises():
+    """TriAttention requires an explicit budget."""
+    with pytest.raises(ValidationError, match="requires 'budget'"):
         TriAttentionConfig()
 
 
-def test_target_sparsity_out_of_range_low():
-    """target_sparsity_ratio <= 0 raises."""
-    with pytest.raises(ValidationError, match="must be in"):
-        TriAttentionConfig(target_sparsity_ratio=0.0)
-
-
-def test_target_sparsity_out_of_range_high():
-    """target_sparsity_ratio >= 1 raises."""
-    with pytest.raises(ValidationError, match="must be in"):
-        TriAttentionConfig(target_sparsity_ratio=1.0)
-
-
-def test_target_sparsity_negative():
-    """Negative target_sparsity_ratio raises."""
-    with pytest.raises(ValidationError):
-        TriAttentionConfig(target_sparsity_ratio=-0.1)
+def test_non_positive_budget_raises():
+    """Budget must be positive."""
+    with pytest.raises(ValidationError, match="budget must be positive"):
+        TriAttentionConfig(budget=0)
 
 
 def test_config_custom_values():
@@ -91,17 +71,6 @@ def test_config_serialization_roundtrip_budget():
     data = config.model_dump()
     restored = TriAttentionConfig(**data)
     assert restored.budget == 1024
-    assert restored.target_sparsity_ratio is None
-    assert restored.prune_interval == 64
-
-
-def test_config_serialization_roundtrip_sparsity():
-    """Config with target_sparsity_ratio survives serialization roundtrip."""
-    config = TriAttentionConfig(target_sparsity_ratio=0.5, prune_interval=64)
-    data = config.model_dump()
-    restored = TriAttentionConfig(**data)
-    assert restored.budget is None
-    assert restored.target_sparsity_ratio == 0.5
     assert restored.prune_interval == 64
 
 

tests/unit/torch/sparsity/kv_cache/test_conversion.py

@@ -92,37 +92,12 @@ def test_update_metadata():
     assert metadata["triattention_config"]["budget"] == 512
 
 
-def test_convert_metadata_with_sparsity_ratio():
-    """Metadata serializes target_sparsity_ratio when set."""
-    model = nn.Linear(16, 16)
-    config = TriAttentionConfig(target_sparsity_ratio=0.7)
-
-    _, metadata = convert_triattention(model, config)
-
-    serialized = metadata["triattention_config"]
-    assert serialized["target_sparsity_ratio"] == 0.7
-    assert serialized["budget"] is None
-
-
 def test_convert_metadata_with_budget():
-    """Metadata has budget set and target_sparsity_ratio None."""
+    """Metadata has budget set."""
     model = nn.Linear(16, 16)
     config = TriAttentionConfig(budget=1024)
 
     _, metadata = convert_triattention(model, config)
 
     serialized = metadata["triattention_config"]
     assert serialized["budget"] == 1024
-    assert serialized["target_sparsity_ratio"] is None
-
-
-def test_update_metadata_with_sparsity_ratio():
-    """update_triattention_metadata serializes target_sparsity_ratio."""
-    model = nn.Linear(16, 16)
-    config = TriAttentionConfig(target_sparsity_ratio=0.5)
-    metadata = {}
-
-    update_triattention_metadata(model, config, metadata)
-
-    assert metadata["triattention_config"]["target_sparsity_ratio"] == 0.5
-    assert metadata["triattention_config"]["budget"] is None

tests/unit/torch/sparsity/kv_cache/test_scoring.py

@@ -207,38 +207,18 @@ def test_select_keys_top_k_exceeds_size():
     assert mask.shape == scores.shape
 
 
-def test_select_keys_percentile_basic():
-    """Percentile selection evicts target fraction."""
-    scores = torch.arange(10, dtype=torch.float32)
-    # sparsity=0.7 → evict 70%, keep top 30% (3 tokens)
-    mask = select_keys_to_keep(scores, target_sparsity_ratio=0.7)
-    assert mask.dtype == torch.bool
-    assert mask.sum().item() == 3
-    # Top 3 by score are indices 7, 8, 9
-    assert mask[7].item() is True
-    assert mask[8].item() is True
-    assert mask[9].item() is True
-
-
-def test_select_keys_percentile_half():
-    """50% sparsity keeps half the tokens."""
-    scores = torch.arange(20, dtype=torch.float32)
-    mask = select_keys_to_keep(scores, target_sparsity_ratio=0.5)
-    assert mask.sum().item() == 10
-
-
-def test_select_keys_both_raises():
-    """Setting both budget and target_sparsity_ratio raises."""
+def test_select_keys_missing_budget_raises():
+    """Budget is required for selection."""
     scores = torch.rand(10)
-    with pytest.raises(ValueError, match="exactly one"):
-        select_keys_to_keep(scores, kv_budget=5, target_sparsity_ratio=0.5)
+    with pytest.raises(ValueError, match="requires kv_budget"):
+        select_keys_to_keep(scores)
 
 
-def test_select_keys_neither_raises():
-    """Setting neither budget nor target_sparsity_ratio raises."""
+def test_select_keys_non_positive_budget_raises():
+    """Budget must be positive."""
     scores = torch.rand(10)
-    with pytest.raises(ValueError, match="exactly one"):
-        select_keys_to_keep(scores)
+    with pytest.raises(ValueError, match="must be positive"):
+        select_keys_to_keep(scores, kv_budget=0)
 
 
 def test_select_keys_empty_scores():