ruvllm: real speculative decoding + sparse attention integration

[ricable]

Summary

• Speculative decoding was previously a stub: speculative.rs's draft/verify phases round-tripped through generate()'s text API (detokenize → retokenize per candidate token) instead of comparing real logits — no actual batched verification, no real speedup mechanism. This replaces it with a real, logit-level implementation:

  • Vendors+patches candle-transformers 0.9.2 (patches/candle-transformers, following the existing patches/hnsw_rs convention) to add ModelWeights::forward_all_positions — per-position logits from a multi-token forward pass, needed to verify K draft tokens against the main model in a single batched forward pass. Also fixes a latent causal-mask bug (a square [t,t] mask fails to broadcast when continuing a non-empty KV cache with more than one new token — a combination nothing in this codebase exercised before forward_all_positions existed), and adds cheap KV-cache snapshot/restore so a rejected draft token rolls back without a full context replay.
  • New SpeculativeBackend trait (implemented for CandleBackend) exposing forward_logits/reset_context/snapshot_context/restore_context; SpeculativeDecoder's draft/verify/generate_tokens rewritten around it with greedy argmax-match verification and real log-probabilities.
  • examples/speculative_bench.rs: end-to-end benchmark on real GGUF weights (Llama-2-7B main + TinyLlama-1.1B draft, same tokenizer — required, see module docs). Measured results are acceptance-rate dependent, as expected from the literature: ~1.1x speedup at 85.9% draft acceptance, ~0.6-0.7x at 62-70% acceptance where per-round forward-call overhead isn't fully amortized. Documented honestly in speculative.rs's module docs — not every prompt shows a win, and that's a property of the technique, not a bug.

• ruvllm_sparse_attention was declared but unused — wired it into the actual inference path instead of leaving it standalone: a new sparse-attention feature routes GGUF prompt-prefill attention through SubquadraticSparseAttention (forward_sparse/forward_attn_sparse in the same patched quantized_llama.rs), falling back to dense attention for incremental decode steps where forward_gqa's q.seq == k.seq contract doesn't hold. CandleBackend gains enable_sparse_attention/disable_sparse_attention/sparse_attention_enabled. examples/sparse_attention_check.rs verifies the tensor bridging end-to-end on real weights.

Test plan

• cargo build -p ruvllm --features candle,metal — clean (regression check, sparse-attention off)
• cargo build -p ruvllm --features candle,metal,sparse-attention — clean
• cargo test -p ruvllm --features candle,metal,sparse-attention --lib — 1589 passed, 2 ignored
• cargo test -p ruvllm --features candle,metal,sparse-attention --test speculative_integration — 29 passed
• examples/speculative_bench.rs run end-to-end against real downloaded GGUF weights (Llama-2-7B + TinyLlama-1.1B) — real measured tok/s, not simulated
• examples/sparse_attention_check.rs run end-to-end — coherent (non-garbled) output on both short and window-exceeding prompts, dense vs sparse

🤖 Generated with Claude Code