ZoRRo Inference & Forest Cascade Attention

ZoRRo Inference & Forest Cascade Attention

Overview

ZoRRo Inference optimizes the sampling phase of RL training by implementing Forest Cascade Attention (FCA), which deduplicates shared KV cache reads at the attention layer. During RL rollouts, many sequences are generated from the same prompt. In standard attention, the KV cache of those shared prefixes is read once per request, causing the sampler to spend most of its memory bandwidth fetching identical keys and values repeatedly.

FCA removes this waste by discovering groups of requests that share a KV-cache prefix and splitting each attention computation into:

1. A single grouped pass over shared prefix blocks (read once per group instead of once per request)
2. A per-request pass over unique suffix blocks
3. A rigorous weighting reduction of the two partial results

The result is mathematically equivalent to standard attention while dramatically cutting redundant memory accesses—the longer and more-shared the prefixes, the larger the performance gain.

How It Works

Forest Cascade Attention operates transparently in the vLLM sampling engine and activates automatically for decode-heavy batches with shared prefixes.

Core Mechanism

For each decode batch, FCA:

1. Discovers groups of requests sharing a KV-cache prefix
2. Splits attention into:
   • A grouped pass over shared prefix blocks (read once per group)
   • A per-request pass over unique suffix blocks
3. Reduces the two partial results with rigorous weighting

This ensures each shared prefix block is read once per group instead of once per request, eliminating redundant memory accesses while preserving correctness.

Integration

ZoRRo Inference is implemented in the vLLM sampling engine within the Arctic Platform architecture. It integrates into the broader Arctic RL stack, which includes:

• Training engine — a DeepSpeed engine for forward/backward and optimizer steps
• Log-prob / reference engine — a forward-only DeepSpeed engine for reference/old log-prob computation
• Sampling engine — a vLLM engine with ArcticInference and Forest Cascade Attention for fast rollouts

The sampling engine orchestrates these optimizations over Ray, enabling colocated or distributed GPU execution while keeping the sampler in sync with the trainer through NCCL or CUDA-IPC weight transfer.

Performance Characteristics

Forest Cascade Attention delivers the largest gains for:

• Long-context tasks with shared prompts
• Decode-heavy batches where memory bandwidth is the bottleneck
• High-redundancy scenarios where many requests sample from the same prompt prefix

The optimization is mathematically equivalent to standard attention and requires no changes to model weights or inference logic.

Configuration

ZoRRo Inference activates transparently in compatible vLLM sampling configurations. For detailed tuning knobs and performance benchmarks, see the Forest Cascade Attention README in the Arctic Inference repository.

Relationship to ZoRRo Train

While ZoRRo Train eliminates redundant prompt computation during the forward/backward pass of RL training (reducing 80–95% of tokens in batches with shared prompts), ZoRRo Inference optimizes the decode step by eliminating redundant KV cache reads. Together, they form the complete ZoRRo optimization suite for RL workloads.