Motivated by the observation that memory-bound decode attention is the dominant bottleneck in RLMs inference, SparseSpec uses sparse attention as draft model followed by a full attention as target model, which dramtically reduces the memory loading thus boosts the throughput (consider check H2O and Quest if not familiar with sparse attention). Thus, an accurate and efficient sparse attention mechanism is crucial for SparseSpec.

To improve accuracy, SparseSpec introduces PillarAttn, which periodically updates the sparsity patterns (i.e., crucial tokens), to adapt to the context dynamics of the RLMs' long-generation paradigm. To avoid additional storage overhead (e.g., storing KV-Cache metadata) and computational overhead, SparseSpec co-designs the identification of sparsity patterns with the speculative decoding. Specifically, SparseSpec leverages the intermidates results of the full attention from verification phase (i.e., every $k$ draft phases), to obtain the crucial tokens with Top-K attention scores. We visualize this process in the figure above.

Resource usage fluctuation

Draft and verify phases in sparse self-speculation have heterogeneous resource usages. Specifically, draft requests lead to $1$ input token while verify requests have $(k+1)$ input tokens. Besides, draft requests use sparse attention, with a sparsity up to $5$% over KV-Cache, while verify requests use full attention. Such heterogenity leads to resource usage fluctuation if both phases are scheduled in sequential manner, as visualized in the below figure.

To mitigate this issue, SparseSpec introduces a unified batch scheduler to collocate draft and verify requests within a batch, and balance the resource usages across iterations. This is feasible as the sparse and full attention essentially share the same control and data flow, except for the sparse loading over KV-Cache, which can be easily manipulated by the PageAttention. Check serve/scheduler/spec_scheduler.py L68 for detailed implementation.

Fused draft and verify attention

To increase bandwidth utilization, SparseSpec introduces a customized attention kernel that automatically dispatches tiles from draft and verify to the best kernel configuration via a persistent kernel style. Such a fused prefill/decode attention also benefit the CUDA graph capture since a single graph can cover different combination of prefill and decode configurations. We upstreamed this kernel to FlashInfer in #1137 and #1200.

In standard speculative decoding, each draft depends on the previous verification step to evict rejected tokens and update token-ID metadata. Because of this dependency, the next draft iteration cannot begin until the host finishes verification, leaving the GPU idle. SparseSpec resolves this by extracting verification requests from the batch and deferring them to the next iteration, allowing the remaining requests to proceed without waiting. The overall workflow is illustrated in the figure below.

CUDA graphs require fixed workspace buffers, which prevents concurrently reusing the same buffer across consecutive iterations to avoid race condition. This create a synchronization bubble to wait for the previous graph to finish. SparseSpec eliminates this gap with a double-buffered CUDA graph scheme: it maintains two identical graphs, each backed by its own workspace, allowing one graph to run while the other is being prepared. This is practical because the fused attention kernel significantly reduces the number of CUDA graphs, keeping memory usage manageable. An nsys profile is shown below:

To keep offloading off the critical path, SparseSpec performs KV-Cache offloading asynchronously in a separate CUDA stream. Since each iteration only produces a bounded number of tokens (limited by the GEMM batch size), the amount of data needing offload per step is also bounded. This makes chunk-wise offloading feasible, and SparseSpec implements it with a simple PyTorch API, shown below.

def offload(self, idx_gpu, idx_cpu):
    self._physical_mem_cpu[idx_cpu[0] : idx_cpu[-1] + 1, ...].copy_(
        self._physical_mem_gpu[:, idx_gpu, ...].transpose(0, 1), non_blocking=True
    )

def restore(self, idx_cpu, idx_gpu):
    self._physical_mem_gpu[:, idx_gpu, ...] = (
        self._physical_mem_cpu[idx_cpu[0] : idx_cpu[-1] + 1, ...]
        .to(self.device, non_blocking=True)
        .transpose(0, 1)
    )

SparseSpec also offloads the page indices that encode the sparsity pattern of draft attention. Because physical device pages are freed and reassigned after offloading, these indices must be converted into logical (virtual) indices that remain consistent across the sequence. SparseSpec performs this conversion incrementally using binary search, amortizing the cost across chunks. The detailed implementation can be found in serve/request/kv_cache_ptr/base.py L17.