RIVET: Ragged Inference with Variable-Granularity Execution for Token-Tree Verification in Speculative Decoding

As agentic applications use larger language models and longer interactions, autoregressive decoding incurs increasingly costly serial target-model passes. Tree-based speculative decoding reduces these passes by verifying multiple candidate continuations at once. This verification requires tree attention: each candidate sees the shared history and its own ancestor path, while the number of candidates and the tree topology vary across rounds. Conventional attention kernels use fixed query tiles, KV partitions, and general masks, leading to padded query work, inefficient KV scheduling, and repeated visibility checks even over the fully visible prefix. We present \method, a variable-granularity execution scheme for token-tree verification. It maps tree visibility to specialized regions, processing regular blocks with dense or causal attention while restricting ancestor checks to irregular blocks. It selects query tiles for the physical query shape and sizes KV partitions to balance parallelism against GPU resource and merge costs. Online-softmax merging preserves exact tree-attention semantics across regions and partitions. In nine kernel configurations, \method reduces attention latency by 51.6% on average relative to FlashInfer. Across 24 model–workload configurations, it improves output throughput by 7.6–50.5% and reduces median time per output token (TPOT) by 3.0–34.2% relative to the strongest baseline selected separately for each configuration and metric.