ruvllm: real speculative decoding + sparse attention integration

Cargo.toml

@@ -369,5 +369,12 @@ debug = true
 
 # Patch hnsw_rs to use rand 0.8 instead of 0.9 for WASM compatibility
 # This resolves the getrandom version conflict (0.2 vs 0.3)
+#
+# Patch candle-transformers (0.9.2) to add ModelWeights::forward_all_positions,
+# which returns per-position logits for a multi-token forward pass instead of
+# only the last position. Needed by ruvllm's speculative decoding to verify
+# several draft tokens against the main model in a single batched forward
+# pass; upstream's `forward()` always slices to the final position.
 [patch.crates-io]
 hnsw_rs = { path = "./patches/hnsw_rs" }
+candle-transformers = { path = "./patches/candle-transformers" }

crates/ruvllm/CHANGELOG.md

@@ -5,6 +5,29 @@ All notable changes to the ruvllm crate will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+- Real speculative decoding: `speculative.rs`'s `SpeculativeDecoder` now drafts and
+  verifies against actual model logits (single batched forward pass per round to
+  verify K draft tokens) instead of the previous text-roundtrip stub. Requires a new
+  `SpeculativeBackend` trait (implemented for `CandleBackend`) exposing per-position
+  batched logits and cheap KV-cache snapshot/restore. See `examples/speculative_bench.rs`
+  for real measured tokens/sec on Llama-2-7B (main) + TinyLlama-1.1B (draft), GGUF.
+- `ruvllm_sparse_attention` integration: prompt-prefill attention in the GGUF
+  (`QuantizedLlama`) backend can now route through
+  `ruvllm_sparse_attention::SubquadraticSparseAttention` instead of dense causal
+  attention, via the new `sparse-attention` feature and
+  `CandleBackend::enable_sparse_attention`. See `examples/sparse_attention_check.rs`.
+- `patches/candle-transformers`: a patched copy of candle-transformers 0.9.2 (see
+  `patches/README.md`) adding `ModelWeights::forward_all_positions` (per-position
+  logits from a multi-token forward pass), `snapshot_kv_cache`/`restore_kv_cache`,
+  a fix for the causal mask when continuing a non-empty KV cache with more than one
+  new token, and (behind `sparse-attention`) `forward_sparse`/`forward_attn_sparse`.
+- `CandleBackend::forward_multi`, `reset_context`, `context_len`, `snapshot_context`,
+  `restore_context`, `vocab_size` — the primitives speculative decoding needs, usable
+  independently of it.
+
 ## [2.0.0] - 2025-01-19
 
 ### Added

crates/ruvllm/Cargo.toml

@@ -68,6 +68,11 @@ candle-core = { version = "0.9", optional = true }
 candle-nn = { version = "0.9", optional = true }
 candle-transformers = { version = "0.9", optional = true }
 
+# Subquadratic sparse attention, optionally spliced into the candle-transformers
+# GGUF inference path (patches/candle-transformers) for prompt prefill. See the
+# `sparse-attention` feature below.
+ruvllm_sparse_attention = { path = "../ruvllm_sparse_attention", optional = true }
+
 # Direct cudarc access for fused CUDA kernels (fused-act feature).
 # Upgraded to 0.19 alongside candle 0.9 — supports CUDA 13.0 natively.
 cudarc = { version = "0.19", optional = true, default-features = false, features = ["nvrtc", "driver", "std"] }
@@ -156,6 +161,18 @@ parallel = ["dep:rayon"]
 # load works without it; HF Hub fetch needs `hub-download` feature.
 candle = ["candle-core", "candle-nn", "candle-transformers", "tokenizers"]
 
+# Route the candle GGUF backend's prompt-prefill attention through
+# ruvllm_sparse_attention::SubquadraticSparseAttention instead of dense O(n^2)
+# causal attention. See patches/candle-transformers's quantized_llama.rs
+# (forward_sparse / forward_attn_sparse) and CandleBackend::load_gguf's
+# `sparse_attention` config for the wiring; falls back to dense attention for
+# incremental decode steps, where prefill's forward_gqa contract doesn't apply.
+sparse-attention = [
+    "candle",
+    "candle-transformers/sparse-attention",
+    "dep:ruvllm_sparse_attention",
+]
+
 # Lattice backend for LLM inference (pure-Rust, Metal acceleration on Mac).
 # macOS-only today: the dependency is declared under
 # `[target.'cfg(target_os = "macos")'.dependencies]` so `metal`/`objc` FFI
@@ -302,6 +319,16 @@ name = "run_eval"
 path = "examples/run_eval.rs"
 required-features = ["async-runtime"]
 
+[[example]]
+name = "speculative_bench"
+path = "examples/speculative_bench.rs"
+required-features = ["candle"]
+
+[[example]]
+name = "sparse_attention_check"
+path = "examples/sparse_attention_check.rs"
+required-features = ["sparse-attention"]
+
 # Lint configuration for the package
 [lints.clippy]
 # Allow common patterns in test code for readability

crates/ruvllm/README.md

@@ -505,31 +505,37 @@ println!("Batch utilization: {:.1}%", stats.avg_batch_utilization * 100.0);
 
 ## Speculative Decoding
 
-Accelerate generation with draft models:
+Accelerate generation with a small draft model, verified against the main model in
+batched forward passes. Main and draft models **must share a tokenizer/vocabulary**
+(e.g. TinyLlama-1.1B drafting for Llama-2-7B/13B, which use the same tokenizer) — see
+`src/speculative.rs` module docs for the full design, its KV-cache-recovery approach,
+and measured tokens/sec on real weights (`examples/speculative_bench.rs`).
 
 ```rust
 use ruvllm::speculative::{SpeculativeDecoder, SpeculativeConfig};
+use std::sync::Arc;
 
 let config = SpeculativeConfig {
-    draft_tokens: 4,           // Tokens to draft per step
-    acceptance_threshold: 0.8, // Min probability for acceptance
+    lookahead: 5,               // Draft tokens to speculate per round
     ..Default::default()
 };
 
 let decoder = SpeculativeDecoder::new(
-    target_model,
-    draft_model,
+    Arc::new(target_model),
+    Arc::new(draft_model),
     config,
-)?;
+);
 
-// Generate with speculation
+// Generate with speculation (greedy/low-temperature — see should_use_speculative)
 let output = decoder.generate(prompt, GenerateParams {
     max_tokens: 256,
+    temperature: 0.0,
     ..Default::default()
 })?;
 
-println!("Acceptance rate: {:.1}%", output.stats.acceptance_rate * 100.0);
-println!("Speedup: {:.2}x", output.stats.speedup);
+let stats = decoder.stats();
+println!("Acceptance rate: {:.1}%", stats.acceptance_rate * 100.0);
+println!("Avg tokens per main-model forward pass: {:.2}", stats.avg_tokens_per_main_pass);
 ```
 
 ## GGUF Model Loading

crates/ruvllm/examples/sparse_attention_check.rs

@@ -0,0 +1,108 @@
+#![allow(clippy::all)]
+//! Sanity check for the sparse-attention integration
+//! (`patches/candle-transformers`'s `forward_sparse` / `forward_attn_sparse`,
+//! wired via `CandleBackend::enable_sparse_attention`).
+//!
+//! Loads one real GGUF model, generates the same prompt with dense attention
+//! and then with sparse attention enabled, and prints both outputs so a human
+//! can eyeball that sparse attention produces coherent (not garbled) text —
+//! this is not a numerical equivalence check (sparse attention is a genuinely
+//! different, approximate algorithm; the tensor bridging is the thing under
+//! test here, not exact output parity).
+//!
+//! ```bash
+//! cargo run -p ruvllm --release --features candle,metal,sparse-attention \
+//!     --example sparse_attention_check -- --model ./model.gguf
+//! ```
+
+use std::env;
+use std::path::PathBuf;
+
+#[cfg(feature = "sparse-attention")]
+fn main() {
+    use ruvllm::{CandleBackend, GenerateParams, LlmBackend, ModelConfig};
+
+    let args: Vec<String> = env::args().collect();
+    let mut model_path = PathBuf::new();
+    let mut i = 1;
+    while i < args.len() {
+        if args[i] == "--model" || args[i] == "-m" {
+            i += 1;
+            if i < args.len() {
+                model_path = PathBuf::from(&args[i]);
+            }
+        }
+        i += 1;
+    }
+    if model_path.as_os_str().is_empty() {
+        eprintln!("Usage: --model <GGUF_PATH> (tokenizer.json expected alongside it)");
+        std::process::exit(1);
+    }
+    let tokenizer_path = model_path
+        .parent()
+        .expect("model path has no parent")
+        .join("tokenizer.json");
+
+    let params = GenerateParams {
+        max_tokens: 40,
+        temperature: 0.0,
+        top_p: 1.0,
+        top_k: 1,
+        ..Default::default()
+    };
+    let long = args.iter().any(|a| a == "--long");
+    let prompt = if long {
+        // Deliberately longer than SparseAttentionConfig::default()'s
+        // 128-token window, so the local-window/landmark sparsity actually
+        // activates instead of degenerating to full attention.
+        "The history of the city of Paris stretches back over two thousand years. \
+         Founded by a Celtic tribe known as the Parisii around the 3rd century BC, \
+         the settlement was originally located on the Ile de la Cite, an island in \
+         the Seine river. The Romans conquered the area in 52 BC and renamed it \
+         Lutetia, building roads, baths, and an amphitheater. Over the following \
+         centuries the city grew steadily, becoming the capital of the Kingdom of \
+         France under the Capetian dynasty. During the Middle Ages, Paris became a \
+         major center of learning, commerce, and religion in Europe. The capital of \
+         France is"
+    } else {
+        "The capital of France is"
+    };
+
+    let mut dense = CandleBackend::new().expect("create backend");
+    dense
+        .load_gguf(&model_path, &ModelConfig::default())
+        .expect("load gguf (dense)");
+    dense
+        .load_tokenizer(&tokenizer_path)
+        .expect("load tokenizer (dense)");
+    let dense_output = dense
+        .generate(prompt, params.clone())
+        .expect("dense generate");
+    println!("[dense]  {dense_output:?}");
+
+    let mut sparse = CandleBackend::new().expect("create backend");
+    sparse
+        .load_gguf(&model_path, &ModelConfig::default())
+        .expect("load gguf (sparse)");
+    sparse
+        .load_tokenizer(&tokenizer_path)
+        .expect("load tokenizer (sparse)");
+    sparse
+        .enable_sparse_attention(ruvllm_sparse_attention::SparseAttentionConfig::default())
+        .expect("enable sparse attention");
+    assert!(sparse.sparse_attention_enabled());
+    let sparse_output = sparse.generate(prompt, params).expect("sparse generate");
+    println!("[sparse] {sparse_output:?}");
+
+    if dense_output.trim().is_empty() || sparse_output.trim().is_empty() {
+        eprintln!("FAIL: empty output from one of the paths");
+        std::process::exit(1);
+    }
+    println!("\nBoth paths produced non-empty output; inspect them above for coherence.");
+}
+
+#[cfg(not(feature = "sparse-attention"))]
+fn main() {
+    eprintln!("This example requires --features candle,metal,sparse-attention");
+    std::process::exit(1);
+}