$ show tanishqkumar/ssd
This repository implements Speculative Speculative Decoding (SSD), a novel LLM inference algorithm that enables parallel execution of draft and target models on separate hardware. The system provides a complete inference engine with tensor parallelism support, optimized kernels (FlashInfer, CUDA graphs), and advanced KV caching techniques. It includes benchmarking tools to compare performance against baselines like SGLang and vLLM, and supports both synchronous and asynchronous speculative decoding modes. The codebase is well-structured with clear separation of concerns between model implementations, engine components, and utility functions.
The LLM identified 3 files worth highlighting for reuse:
Implements innovative KV cache management using hash-based block lookup with reference counting. Enables efficient prefix caching by identifying reusable token sequences through content-based hashing rather than positional tracking. The block allocator/deallocator manages memory efficiently with minimal fragmentation, and the design supports dynamic batch sizes. This component could be extracted for use in any LLM inference system needing optimized memory management.
Provides reusable sampling logic with special support for tree decoding in asynchronous speculative decoding. Implements the 'squirrel' sampling technique that combines greedy and stochastic sampling based on temperature. Includes a fix for a FlashInfer RNG state bug that would cause device mismatch errors. The module is self-contained with minimal dependencies and could be easily integrated into any text generation system requiring flexible sampling strategies.
Contains core algorithms for asynchronous speculative decoding, including tree decoding logic and forked recovery token selection. The get_forked_recovery_tokens_from_logits function implements a sophisticated method for selecting tokens based on logits and fan-out strategies, using top-k selection with masking to handle cache hits/misses. These utility functions represent the novel speculative decoding mechanics and could be adapted for use in other speculative decoding implementations with minimal changes to the data structures.