kangaroospeculative-decodingearly-exitllm-inferencevllmtutorial

Kangaroo: Self-Speculative Decoding via Shallow Sub-Networks

September 23, 2026

Kangaroo Architecture: Self-Speculative Decoding, Fixed Shallow Sub-Networks, and Double Early Exiting

The Problem: The Memory and Distribution Traps of Speculative Drafting

In production artificial intelligence applications, autonomous web extractors, and interactive agent loops like AiFetchly, foundation model serving is fundamentally constrained by the memory-bandwidth wall.

During autoregressive generation, a 70-billion parameter transformer must load its entire weight matrix (consuming ~140 GB in 16-bit precision) across the GPU memory bus just to generate a single token. Because the arithmetic intensity of single-token generation is near $1.0 \text{ FLOP/byte}$, modern GPU Tensor Cores (such as NVIDIA H100 or H200) spend over 90% of their compute cycles idling, waiting for weights to stream from High-Bandwidth Memory (HBM).

To break this bottleneck, Speculative Decoding (Leviathan et al., Chen et al.) was introduced: a small, fast “draft model” speculates $K$ candidate tokens, which the primary “target model” verifies in parallel in a single forward pass.

However, deploying traditional dual-model speculative decoding in production creates severe operational friction:

The Operational Dilemmas of Dual-Model Speculative Decoding:
┌───────────────────────────┬────────────────────────────────────────────────────────┐
│ Operational Bottleneck    │ Real-World Impact                                      │
├───────────────────────────┼────────────────────────────────────────────────────────┤
│ Dual VRAM Allocation      │ Storing an 8B draft model consumes 16GB-24GB VRAM!     │
│ Distribution Shift        │ Mismatched training data causes acceptance α < 45%.    │
│ Multi-GPU Interconnect    │ Running tiny draft models across TP=8 causes stalls.   │
│ Multi-Head Incoherence    │ Medusa-style heads cannot condition on drafted tokens. │
└───────────────────────────┴────────────────────────────────────────────────────────┘

1. The VRAM Memory Penalty of Secondary Draft Models

In classical speculative sampling, pairing a 70B target model with an 8B draft model requires allocating dedicated GPU memory for:

The draft model’s static weight parameters (~16 GB in FP16).
An entirely separate Key-Value (KV) cache buffer for the draft model.
Dedicated CUDA execution graphs and scheduler queues.

In enterprise serving clusters, sacrificing 16 GB to 24 GB of GPU memory per node severely reduces the maximum concurrency batch size, canceling out the throughput benefits of speculative decoding.

2. Distribution Shift and Acceptance Rate Decay

Even when a small draft model belongs to the same family as the target model (e.g., LLaMA-3 8B drafting for LLaMA-3 70B), their learned output distributions diverge substantially:

On challenging reasoning, mathematical proofs, and complex code syntax, the small model frequently hallucinates tokens that the large target model immediately rejects.
When the draft acceptance rate $\alpha$ falls below 50%, the time spent generating rejected draft tokens exceeds the time saved during verification, causing speculative decoding to run slower than standard autoregression.

3. The Failure of Independent Multi-Head Drafting

To eliminate secondary draft models, draft-free methods like Medusa attach multiple parallel linear heads to the final layer of the base model.

However, Medusa heads predict future tokens simultaneously and independently without conditioning on intermediate tokens: $$P(x_{t+1}, x_{t+2}, x_{t+3} \mid x_{\le t}) \approx P(x_{t+1} \mid x_{\le t}) \cdot P(x_{t+2} \mid x_{\le t}) \cdot P(x_{t+3} \mid x_{\le t})$$

This violates the fundamental chain rule of probability. The second and third heads cannot observe what the first head decided to generate. As a consequence, Medusa’s acceptance rate drops sharply on multi-token phrases that require strict sequential syntactic coherence.

The Architecture: Kangaroo and Double Early Exiting

To achieve true self-speculative acceleration without loading a secondary draft model or sacrificing conditional dependency, researchers introduced Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting (Liu et al., NeurIPS 2024 / arXiv:2404.18911).

Kangaroo fundamentally reconceptualizes the drafter: Instead of loading an external neural network, Kangaroo reuses the first few layers of the target model itself, augmented by a lightweight recurrent adapter and a dynamic confidence exit.

Kangaroo Architecture: Shallow Sub-Network Partitioning, Adapter Feature Drafting, and Parallel Verification

1. Fixed Shallow Sub-Network Drafting

Consider a large target model consisting of $L$ transformer layers (e.g., $L = 32$ in LLaMA-3 8B or $L = 80$ in LLaMA-3 70B).

Kangaroo partitions the model into two functional components:

The Shallow Sub-Network (Layers $1$ to $K$): The first $K$ layers of the target model (typically $K = 2$ or $K = 4$).
The Deep Sub-Network (Layers $K+1$ to $L$): The remaining layers of the target model.

During the drafting phase, Kangaroo executes only the shallow sub-network. Because layers $1 \dots K$ are identical to the target model’s own layers:

There is zero parameter duplication: the drafter requires zero extra model weights in VRAM.
The drafter directly populates the target model’s primary KV cache, eliminating duplicate memory allocations.
Shallow hidden representations $h_K$ already contain rich, grounded representations aligned with the target model’s feature space.

2. The Double Early Exiting Mechanism

To generate coherent multi-token candidate sequences while avoiding doomed speculative branches, Kangaroo introduces Double Early Exiting:

Exit 1: The Lightweight Early Exit Adapter: To predict subsequent tokens autoregressively without re-running the entire model, Kangaroo appends an ultra-lightweight Early Exit Adapter (a single transformer layer or lightweight linear projection, adding $<1%$ parameter overhead) on top of layer $K$:

$$h_{\text{adapt}}^{(t+1)} = \text{Adapter}(h_K^{(t)}, x_t)$$

$$\hat{P}(x_{t+1}) = \text{Softmax}(W_{\text{head}} \cdot h_{\text{adapt}}^{(t+1)})$$

Unlike Medusa’s parallel heads, Kangaroo’s adapter generates candidate tokens sequentially: $$x_{t+1} \sim \hat{P}(x_{t+1}) \quad \longrightarrow \quad x_{t+2} \sim \hat{P}(x_{t+2} \mid x_{t+1}) \quad \longrightarrow \quad x_{t+3} \sim \hat{P}(x_{t+3} \mid x_{t+1}, x_{t+2})$$

By feeding each generated draft token back into the adapter, Kangaroo strictly respects the causal chain rule, ensuring syntactic and grammatical coherence across the entire draft sequence.

Exit 2: Dynamic Confidence-Aware Early Stopping: In traditional speculative decoding, the drafter always produces a static number of candidate tokens (e.g., exactly 4 or 5 tokens), regardless of whether the query is simple or extraordinarily complex. When a token is difficult, the drafter generates low-confidence guesses that the target model will inevitably reject, wasting drafting compute.

Kangaroo introduces a confidence threshold exit: at each drafting step $i$, the adapter calculates the normalized softmax confidence of its top prediction:

$$c_i = \max_{v} \hat{P}(x_{t+i} = v)$$

If $c_i < \tau$ (where $\tau$ is a tuned confidence threshold, e.g., $\tau = 0.65$):

The drafter detects that the current token is ambiguous or difficult.
Kangaroo immediately halts the drafting loop (Early Exit 2) and sends the candidate tokens accumulated so far to the target model for verification.
This dynamic mechanism prevents the system from wasting cycles drafting tokens with a high probability of rejection, optimizing hardware efficiency.

3. Lossless Speculative Verification

Once the shallow sub-network and adapter assemble a candidate sequence of $m$ tokens ($m \le M$), the deep sub-network (layers $K+1 \dots L$) executes a single batched verification pass.

Kangaroo applies standard speculative rejection sampling: For each candidate token $x_i$, let $p_i = P_{\text{target}}(x_i \mid x_{<i})$ and $q_i = P_{\text{draft}}(x_i \mid x_{<i})$:

If a random sample $r \sim U(0, 1) \le \min\left(1, \frac{p_i}{q_i}\right)$, the token is accepted.
If rejected, the token is discarded, the target model emits an exact sample from the adjusted distribution: $$P_{\text{corrected}}(x) = \frac{\max(0, p_i(x) - q_i(x))}{\sum_v \max(0, p_i(v) - q_i(v))}$$ and verification terminates.

Because rejection sampling is provably exact, Kangaroo’s output distribution is 100% mathematically identical to the uncompressed target model: zero perplexity loss, zero accuracy degradation, and zero hallucination risk.

Comparative Breakdown: Speculative Decoding Paradigms

Dimension	Standard Dual-Model SD	Medusa Multi-Head	EAGLE-2 Feature SD	Kangaroo (Self-SD)
Drafter Architecture	Separate Small LLM (e.g. 8B)	Multiple Linear Heads	Small Autoregressive Head	Fixed Shallow Layers (1..K)
Draft Memory Overhead	Huge (+16 GB to 24 GB)	Minimal (~0.5 GB)	Minimal (~1 GB)	Virtually Zero (<0.2 GB)
Causal Conditioning	Yes (Autoregressive)	No (Independent heads)	Yes (Feature recurrent)	Yes (Adapter recurrent)
Draft Length Policy	Static Fixed Length	Static Tree Structure	Dynamic Tree Expansion	Dynamic Confidence Exit
Distribution Alignment	Prone to drift	Tied to final layer	Tied to top-2 layer	Native (Exact Sub-Network)
Inference Speedup	$1.4\times - 1.8\times$	$1.5\times - 2.2\times$	$2.0\times - 3.0\times$	$1.7\times - 2.2\times$

Implementation: Simulating Kangaroo Self-Speculative Decoding in Python

Let us construct a reference simulation in Python that demonstrates Kangaroo’s architectural workflow:

Shallow Sub-Network Execution: Reuses early layers for base feature extraction.
Adapter Recurrent Drafting (Exit 1): Generates candidate tokens autoregressively.
Confidence-Aware Early Exit (Exit 2): Halts drafting dynamically when prediction entropy exceeds threshold.
Single-Pass Lossless Verification: Verifies candidate tokens using rejection sampling.
Efficiency & VRAM Profiler: Measures latency speedup and memory savings compared to dual-model baselines.

import numpy as np
from typing import List, Dict, Tuple, Any

class KangarooSimulation:
    """
    Simulates Kangaroo self-speculative decoding with double early exiting.
    Demonstrates shallow sub-network reuse and dynamic confidence stopping.
    """
    def __init__(
        self,
        total_layers: int = 32,
        shallow_layers: int = 4,
        confidence_threshold: float = 0.65,
        max_draft_tokens: int = 5,
    ):
        self.total_layers = total_layers
        self.shallow_layers = shallow_layers
        self.deep_layers = total_layers - shallow_layers
        self.tau = confidence_threshold
        self.max_tokens = max_draft_tokens

        # Normalized execution latencies per layer (milliseconds)
        self.ms_per_layer = 1.0
        self.ms_adapter = 0.35  # Tiny adapter forward pass

    def run_draft_phase(self) -> Tuple[int, float, List[float]]:
        """
        Executes Kangaroo's Double Early Exit drafting loop:
        1. Runs shallow sub-network once.
        2. Adapter autoregressively drafts tokens until confidence drops below tau.
        """
        draft_time = self.shallow_layers * self.ms_per_layer
        drafted_count = 0
        confidences = []

        for step in range(self.max_tokens):
            draft_time += self.ms_adapter
            drafted_count += 1

            # Simulate adapter confidence (beta distribution modeling token entropy)
            # Easy grammatical tokens have high confidence; complex reasoning has lower
            confidence = float(np.random.beta(a=4.5, b=1.8))
            confidences.append(confidence)

            # Exit 2: Dynamic Confidence Early Stopping
            if confidence < self.tau:
                # Halt drafting early to avoid proposing doomed tokens!
                break

        return drafted_count, draft_time, confidences

    def run_verification_phase(self, drafted_count: int) -> Tuple[int, float]:
        """
        Executes target model verification in a single batched pass across deep layers.
        Rejection sampling guarantees lossless output.
        """
        verify_time = self.total_layers * self.ms_per_layer

        # Simulate rejection sampling acceptance
        # High confidence tokens achieve ~85% acceptance rate
        accepted_tokens = 1  # 1 bonus token guaranteed by target verification
        for _ in range(drafted_count):
            if np.random.random() < 0.82:
                accepted_tokens += 1
            else:
                break

        return accepted_tokens, verify_time

    def benchmark(self, iterations: int = 100) -> Dict[str, Any]:
        """Runs a comparative benchmark across decoding paradigms."""
        np.random.seed(42)

        total_accepted = 0
        kangaroo_total_time = 0.0
        dual_model_total_time = 0.0
        autoregressive_total_time = 0.0

        draft_lengths = []

        # Standard dual-model baseline (e.g. 8B draft model takes 6.0ms per token)
        ms_dual_drafter_step = 6.0

        for _ in range(iterations):
            # 1. Kangaroo step
            num_drafted, t_draft, confs = self.run_draft_phase()
            accepted, t_verify = self.run_verification_phase(num_drafted)

            kangaroo_step_time = t_draft + t_verify
            kangaroo_total_time += kangaroo_step_time
            total_accepted += accepted
            draft_lengths.append(num_drafted)

            # 2. Standard Dual-Model SD step (fixed 4 tokens drafted)
            dual_step_time = (4 * ms_dual_drafter_step) + (self.total_layers * self.ms_per_layer)
            dual_model_total_time += dual_step_time

            # 3. Autoregressive baseline step
            autoregressive_total_time += (accepted * self.total_layers * self.ms_per_layer)

        speedup_kangaroo = autoregressive_total_time / kangaroo_total_time
        speedup_dual = autoregressive_total_time / dual_model_total_time

        return {
            "total_tokens": total_accepted,
            "mean_draft_length": float(np.mean(draft_lengths)),
            "tokens_per_step": float(total_accepted / iterations),
            "speedup_kangaroo": float(speedup_kangaroo),
            "speedup_dual": float(speedup_dual),
        }


def run_kangaroo_verification():
    print("=" * 70)
    print("KANGAROO SELF-SPECULATIVE DECODING BENCHMARK SIMULATION")
    print("=" * 70)

    sim = KangarooSimulation(
        total_layers=32,
        shallow_layers=4,
        confidence_threshold=0.65,
        max_draft_tokens=5,
    )

    results = sim.benchmark(iterations=100)

    print(f"Target Model:               32-Layer LLM (e.g. LLaMA-3 8B/70B)")
    print(f"Shallow Sub-Network Size:   4 Layers (Layers 1..4 reused natively)")
    print(f"Adapter Parameter Footprint: < 0.8% parameter overhead\n")

    print("DYNAMIC EARLY EXIT METRICS:")
    print(f"  Average Drafted Length:   {results['mean_draft_length']:.2f} tokens (dynamically bounded)")
    print(f"  Effective Accepted Tokens: {results['tokens_per_step']:.2f} tokens per step\n")

    print("LATENCY & SPEEDUP EVALUATION:")
    print("  Autoregressive Baseline:  1.00x (Standard decoding)")
    print(f"  Dual-Model Speculative:   {results['speedup_dual']:.2f}x speedup")
    print(f"  Kangaroo Self-Speculative: {results['speedup_kangaroo']:.2f}x speedup\n")

    # VRAM comparison (70B parameter model)
    target_vram = 140.0
    dual_draft_vram = 16.0
    kangaroo_adapter_vram = 0.15

    print("VRAM MEMORY ALLOCATION FOR DRAFTING:")
    print(f"  Dual-Model Draft LLM:     +{dual_draft_vram:.1f} GB VRAM (+11.4% memory tax)")
    print(f"  Kangaroo Sub-Net + Adapter: +{kangaroo_adapter_vram:.2f} GB VRAM (+0.1% memory tax)")
    print(f"  VRAM Savings:             {((1 - kangaroo_adapter_vram/dual_draft_vram)*100):.1f}% reduction in drafter memory!")
    print("=" * 70)
    assert results["speedup_kangaroo"] > 1.5, "Kangaroo must deliver >1.5x speedup"
    print("  [SUCCESS] Kangaroo self-speculative decoding verified successfully!")


if __name__ == "__main__":
    run_kangaroo_verification()

Production Advantages and Deployment Considerations

Deploying Kangaroo in enterprise inference engines (such as vLLM, TensorRT-LLM, or SGLang) delivers critical systems benefits:

1. 99% Reduction in Drafter Memory Footprint

In memory-constrained production environments:

Standard speculative decoding requires allocating up to 24 GB of High-Bandwidth Memory solely to host the secondary draft model and its independent KV cache.
Kangaroo requires only a lightweight adapter checkpoint (<150 MB), liberating virtually the entire GPU memory pool for larger batch sizes, longer context prefixes, and increased concurrent server throughput.

2. Elimination of Multi-GPU Communication Overheads

In distributed deployments with Tensor Parallelism ($TP=4$ or $TP=8$):

Running a tiny 1B or 8B draft model across 8 GPUs is notoriously inefficient: inter-GPU All-Reduce communication over NVLink dominates runtime because the mathematical workload is too small to saturate Tensor Cores.
Because Kangaroo executes within the target model’s own sharded layers, tensor-parallel operations remain naturally balanced, eliminating communication bottlenecks.

3. Native Integration with PagedAttention and KV Cache Reuse

Because Kangaroo’s shallow sub-network shares the exact same token embedding weights, layer normalization scales, and hidden dimension dimensions as the base model:

The drafter can directly read and write to the target model’s existing PagedAttention virtual memory blocks.
There is zero need to synchronize or copy KV cache states between separate memory pools, drastically simplifying inference scheduler architecture.

Conclusion

Kangaroo solves the central operational paradox of speculative decoding: achieving substantial latency acceleration without the crippling VRAM overhead and distribution drift of a secondary draft model.

By coupling a fixed shallow sub-network with double early exiting—recurrent adapter drafting and dynamic confidence stopping—Kangaroo doubles LLM generation speed while adding less than 1% parameter overhead. For AI infrastructure engineers scaling high-throughput serving clusters, interactive agents, and real-time reasoning pipelines, Kangaroo establishes the definitive architectural blueprint for efficient, lossless self-speculative inference.