[Feature Proposal] Runtime-Adaptive Elastic Sparsity for Dynamic Head Recovery and Token Selection

[SpaceGazers]

[Feature Proposal] Runtime-Adaptive Elastic Sparsity for Dynamic Head Recovery and Token Selection

Motivation

Thanks for releasing RedKnot. I really like the idea of decomposing KV-cache recovery along the head dimension instead of treating all tokens and heads uniformly.

One possible extension I would like to discuss is making Elastic Sparsity more runtime-adaptive.

Currently, RedKnot mainly relies on offline profiling to determine:

• global vs. local KV heads;
• local attention windows;
• dense/sparse layer boundaries;
• Sparse-FFN token-selection thresholds.

This is a reasonable design because the behavior of many attention heads is relatively stable across requests. However, the amount of recovery actually required may still vary significantly at runtime.

For example:

• the same head may behave differently across retrieval, summarization, coding, and agent workloads;
• attention concentration changes with context length;
• different reused chunks may have very different dependencies on the new prefix;
• some requests may require only a subset of the profiled global heads to be recomputed;
• some local heads may need a small repair window for one request but a much larger window for another.

Therefore, instead of treating the offline head-class map as a hard execution policy, it may be useful to treat it as a prior, while dynamically deciding the actual recovery sparsity at runtime.

---

Proposal: Runtime-Adaptive Elastic Sparsity

The main idea is:

Use offline profiling to define the candidate sparse structure, but let a lightweight runtime controller decide how much computation should actually be restored for each request, layer, and head.

The runtime could dynamically choose among several recovery modes:

REUSE
  ↓
LOCAL_REPAIR(w)
  ↓
FULL_HEAD_RECOMPUTE
  ↓
DENSE_LAYER_FALLBACK

Instead of using a fixed binary global/local execution decision, RedKnot could gradually increase the recovery budget only when the current request requires it.

---

1. Treat Offline Head Classification as a Prior

The existing offline

(layer, head) -> {global, local}

classification can remain unchanged.

However, rather than directly determining runtime execution, it could be interpreted as a prior:

P(recompute | layer, head)

For example:

offline-local head
        ↓
low recomputation prior

offline-global head
        ↓
high recomputation prior

The runtime controller can then adjust this prior using request-specific signals.

This preserves the benefits of offline profiling while enabling per-request adaptation.

---

2. Lightweight Runtime Signals

I do not think this requires another expensive classifier.

Several signals that are already available, or can be collected cheaply during execution, could potentially be reused:

• attention mass outside the current local window;
• edge/boundary attention mass;
• attention entropy or concentration;
• effective attention span;
• cross-chunk / cross-prefix attention strength;
• context length;
• layer depth;
• prefix-position changes;
• architecture-specific sparse/indexer signals, e.g. for DeepSeek-style architectures.

For each (layer, head), the runtime could estimate a lightweight recovery-risk score:

risk(l, h) =
    f(
        offline_head_prior,
        edge_mass,
        attention_entropy,
        effective_attention_span,
        cross_chunk_mass,
        context_length,
        architecture_specific_signal
    )

A first implementation probably does not need a learned model.

A heuristic or threshold-based controller may already be sufficient.

---

3. Dynamically Select How Many Heads to Recompute

Instead of always recomputing every offline-global head, RedKnot could dynamically select a subset according to runtime risk.

For example:

candidate_heads = offline_global_heads

for h in candidate_heads:
    risk[h] = estimate_runtime_risk(h)

recompute_heads = select_heads(
    risk,
    quality_threshold=epsilon,
    compute_budget=B,
)

The number of recomputed heads then becomes request-dependent.

For an easier request:

12 offline-global heads
        ↓
runtime controller
        ↓
4 heads recomputed

For a more difficult request:

12 offline-global heads
        ↓
runtime controller
        ↓
10–12 heads recomputed

If the runtime signal indicates high uncertainty, RedKnot can simply fall back to the original global-head recomputation policy.

This makes the optimization conservative by construction.

---

4. Adaptive Local-Head Repair Window

The repair window of local heads could also become dynamic.

Instead of using a fixed value such as:

w = 256

the runtime could select from multiple levels:

128 → 256 → 512 → 1024 → full

according to the observed attention behavior.

For example:

w = initial_window

while outside_window_mass > threshold:
    w *= 2

This allows each local head to use the minimum context range required by the current request.

The sparsity therefore becomes adaptive along two dimensions:

Head dimension:
    How many heads need recovery?

Context dimension:
    How much context does each head need?

This seems particularly compatible with SegPagedAttention because different heads already have independent ragged page lists and heterogeneous visible ranges.

---

5. Runtime-Adaptive Sparse FFN

The same idea could also be applied to Sparse FFN.

Instead of using a fixed token-selection threshold, the runtime could dynamically choose the FFN budget based on the concentration of the recovered attention signal.

For example:

importance = token_importance(attention_output)

target_mass = runtime_policy(
    attention_entropy,
    context_length,
    layer_id,
)

selected_tokens = smallest_set_covering(
    importance,
    target_mass,
)

For highly concentrated attention:

small token subset
        ↓
execute FFN

remaining tokens
        ↓
residual path

For diffuse attention:

larger token subset
        ↓
execute FFN

This would allow the FFN sparsity ratio to adapt naturally to the information density of each request.

---

6. Possible Runtime Controller

A simple first version could look like:

for layer in model.layers:

    signals = collect_runtime_signals(layer)

    for head in layer.kv_heads:

        risk = estimate_risk(
            offline_prior=head.class_prior,
            edge_mass=signals.edge_mass[head],
            entropy=signals.entropy[head],
            span=signals.attention_span[head],
        )

        if risk < tau_reuse:
            policy[head] = REUSE

        elif risk < tau_full:
            w = choose_local_window(risk)
            policy[head] = LOCAL_REPAIR(w)

        else:
            policy[head] = FULL_RECOMPUTE

    if aggregate_risk(policy) > safety_threshold:
        promote_high_risk_heads()

    ffn_budget = choose_ffn_budget(
        attention_concentration=signals.concentration,
        context_length=context_length,
        layer_id=layer.id,
    )

    selected_tokens = select_tokens(ffn_budget)

One important point is that I would probably avoid arbitrary token-level partial recomputation inside a global head, because this could reintroduce the stale-context / cascading-error problem that head-level recovery is intended to avoid.

Therefore, the dynamic decisions can remain structured:

Global head:
    reuse OR full-head recomputation

Local head:
    reuse OR adaptive-window repair

FFN:
    adaptive token subset

This preserves the core design philosophy of RedKnot.

---

7. SLO-Aware Recovery Budget

Another interesting extension would be to expose a runtime compute budget B.

For example:

B = f(
    current_GPU_load,
    batch_size,
    TTFT_SLO,
    available_HBM,
    request_priority
)

The controller could conceptually optimize:

maximize expected output fidelity

subject to

recovery_cost <= B

The recovery budget could then be distributed between attention and FFN:

              Runtime Recovery Budget
                        │
           ┌────────────┴────────────┐
           │                         │
     Head Recovery               Sparse FFN
           │                         │
   select # of heads          select # of tokens
           │                         │
  adaptive windows          adaptive threshold

Under low load, RedKnot could use a more conservative recovery policy.

Under high load or tighter TTFT SLOs, the runtime could increase sparsity while still protecting high-risk global/retrieval heads.

This could turn Elastic Sparsity from a mostly static optimization into a real serving-time control mechanism.

---

Why This Fits RedKnot

I think this direction fits naturally with the existing RedKnot design:

1. Head-aware recovery already exposes the correct execution granularity.
2. SegPagedAttention already supports heterogeneous per-head context ranges.
3. Sparse FFN already provides an independent token-level sparsity dimension.
4. The missing component is mainly a lightweight runtime policy that determines the actual sparsity level.

The architecture could become:

Offline Profiling
       │
       ▼
Head / Layer Sparsity Prior
       │
       ▼
Runtime Signals
       │
       ▼
Adaptive Sparsity Controller
       │
       ├── Heads to recompute
       ├── Local repair window
       ├── FFN token budget
       └── Dense fallback
       │
       ▼
SegPagedAttention + Sparse FFN

This would make the term Elastic Sparsity even more literal: the amount of sparse execution can expand or contract according to the actual recovery demand of each request.

---

Suggested Initial Prototype

To keep the first implementation simple, I would suggest starting with two mechanisms.

A. Dynamic Head Recovery

Use lightweight runtime signals to dynamically determine whether a head really requires full recomputation.

offline classification
        ↓
runtime risk estimation
        ↓
REUSE / LOCAL_REPAIR / FULL_RECOMPUTE

The original offline policy can always remain the safe fallback.

B. Dynamic FFN Sparsity

Instead of using one fixed Sparse-FFN threshold, dynamically determine the token keep ratio from runtime attention concentration.

No learned controller would be required initially.

If these two mechanisms already improve the accuracy/TTFT Pareto frontier, more sophisticated runtime policies could be explored later.

---

Evaluation

It would be interesting to compare:

1. Static RedKnot
2. Dynamic head recovery only
3. Dynamic Sparse-FFN only
4. Dynamic head recovery + Dynamic Sparse-FFN
5. Dense recomputation

across:

• different context lengths;
• retrieval QA vs. summarization vs. agent workloads;
• different numbers and orderings of reused chunks;
• different prefix changes;
• low-load vs. high-load serving conditions.

Useful metrics could include:

• TTFT;
• FLOPs;
• QPS;
• fraction of heads recomputed;
• average local-head repair window;
• Sparse-FFN token ratio;
• logit cosine similarity;
• top-k agreement;
• downstream task accuracy;
• QPS under a fixed TTFT SLO.

The main question would be:

Can runtime adaptation push the accuracy-latency Pareto frontier beyond a single statically profiled RedKnot configuration?

---

Summary

In short, the proposal is to change:

Offline Profiling
       ↓
Fixed Sparse Execution

into:

Offline Profiling
       ↓
Sparsity Prior
       ↓
Runtime Signals
       ↓
Dynamic Recovery Budget
       ↓
Per-Request Sparse Execution

The key idea is to allow RedKnot to dynamically determine how sparse the current request can safely be, and therefore dynamically decide:

• how many heads need to be recomputed;
• how much context each local head needs to repair;
• how many token states need full FFN computation.

Would this direction fit the current Elastic Sparsity / SegPagedAttention roadmap?

I would be interested in helping prototype a lightweight runtime controller if this abstraction makes sense.