Support Review of ConcatKvCache (#3143) and Plan for Future Adoption

[DrJesseGlass]

Summary

PR #3143 introduces ConcatKvCache which provides 2-5x GPU speedup for autoregressive generation with no breaking changes to the API. However, it hasn't been reviewed yet.

This issue aims to:

1. Raise visibility of feat(candle-nn) ConcatKvCache for 2-5x GPU speedup on autoregressive generation #3143 for maintainer review
2. Document the systemic need for this improvement across transformers
3. Propose a path forward once feat(candle-nn) ConcatKvCache for 2-5x GPU speedup on autoregressive generation #3143 is merged

Why This Matters Now

Multiple transformer implementations (including the new SmolLM3 in #3180) currently use workarounds that ConcatKvCache would eliminate:

// Current workaround in Llama, Mistral, Phi, Qwen, SmolLM3, etc:
let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?;

// With ConcatKvCache - clean and faster:
let (k, v) = self.kv_cache.append(&k, &v)?;

This pattern appears across the codebase, indicating a systemic issue that ConcatKvCache would solve.

Background: PR #3143

What It Provides

• 2-5x GPU speedup for autoregressive generation
• No breaking changes - API-compatible with existing KvCache
• Equal or better performance on all backends (GPU, CPU, WASM)
• Cleaner code - eliminates .contiguous() workarounds

Current Status

• Awaiting review - not yet merged
• Implementation complete
• Benchmarks show clear benefits
• Needs maintainer attention

The Broader Problem

Current State

Every transformer implementation manually works around KV-cache limitations:

Current pattern:

let kv_cache = KvCache::new(2, 512);
// ...
let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?;

With ConcatKvCache:

- let kv_cache = KvCache::new(2, 512);
+ let kv_cache = ConcatKvCache::new(2);  // Simpler API - no max_seq_len
  // ...
- let (k, v) = self.kv_cache.append(&k.contiguous()?, &v.contiguous()?)?;
+ let (k, v) = self.kv_cache.append(&k, &v)?;  // No contiguous() needed

Additional benefits:

• Simpler constructor (no max sequence length parameter)
• Built-in reset() method useful for generation applications
• 2-5x faster on GPU

Models affected:

• Qwen family (qwen3.rs)
• SmolLM3 (pending in Add SmolLM3: Full and Quantized Implementation #3180)

Models that would benefit:

• Llama family (llama.rs)
• Mistral family (mistral.rs)
• Phi family (phi.rs, phi3.rs)
• Gemma family (gemma.rs)
• All other transformers using KV-cache

This code pattern across the codebase suggests the underlying abstraction could be improved.

Performance Impact

Users running transformers on GPU are getting 2-5x slower performance than they could be getting, simply because we're using the legacy KvCache implementation.

Proposal: Two-Phase Approach

Phase 1: Get #3143 Reviewed and Merged (Immediate)

1. Request maintainer review of feat(candle-nn) ConcatKvCache for 2-5x GPU speedup on autoregressive generation #3143
2. Validate benchmarks across different models
3. Test integration with existing transformers
4. Merge once validated

Phase 2: Discuss Adoption Strategy (After #3143 Merges)

Once #3143 is merged and proven stable, we should discuss how to adopt it across the codebase:

Questions for Discussion:

1. Migration strategy - Should we batch-update all models or migrate incrementally?
2. Deprecation path - How do we handle the API change (constructor signature)?
3. Timeline - What's a reasonable timeline for migration?
4. Compatibility - Do we need any compatibility layer, or is a clean break acceptable?

Note: ConcatKvCache has a different API:

// Old API
let kv_cache = KvCache::new(num_layers, max_seq_len);

// New API (simpler - no max_seq_len parameter)
let kv_cache = ConcatKvCache::new(num_layers);

This isn't a drop-in replacement, so migration will require updating each transformer's initialization code. However, the benefits seem worth it:

• 2-5x GPU speedup
• Simpler API (no max_seq_len)
• Built-in reset() method (useful for generation)
• No .contiguous() workarounds needed

I'm happy to help test integration approaches and contribute to the migration effort once we have consensus on the strategy.

Why This Needs Attention

User Impact

• Users are getting suboptimal performance (2-5x slower on GPU)
• No indication that performance could be better
• Workaround is invisible to users but affects everyone

Developer Impact

• Code duplication across models
• Current Transformers violate DRY principles

Request for Action

Short Term

Can a maintainer review PR #3143?

• The implementation looks solid
• Would benefit multiple existing and new PRs

Medium Term (After Merge)

Make ConcatKvCache the default

• Via type alias for compatibility
• Update documentation
• Plan migration for existing models

Clean up workarounds

• Update transformers systematically
• Remove .contiguous() calls
• Improve code quality

Long Term

• ConcatKvCache becomes standard approach
• New transformers get optimal performance by default
• Cleaner, more maintainable codebase

The main work is systematic cleanup across transformer implementations, which can be done incrementally.

[ivarflakstad]

The PR is great. I'm actively reviewing it (hence the recent activity).
I'm looking into how best to adopt it (like should we select kv cache impl based on device, etc), and simultaneously investigating if there are improvements to be made to both the original and concat based implementations. (Managed to get the original impl 15% faster with very little effort)

This issue is mostly just noisy in my opinion. And it's precisely this style of writing that made me look at a bunch of other PRs before yours. It comes across as very generative which often implies slop. I review more slop than I'd like to, so when I see typical indicators like repetition, heavy use of ✅, inflated numbers, etc it gets moved down on my list of priorities.

Don't worry, the PR will get merged :)

[DrJesseGlass]

Apologies. I should have been more conscientious of your time. Glad you were able to further improve it's efficiency. The code is overly annotated and I am happy to make any changes to the descriptions, tests, or code. This was my first attempt at contributing to Candle. It seemed like a straightforward gain but I wasn't sure how best to go about integrating it. Thanks for your response. None of the numbers I provided are inflated. I ran all of those experiments. :)

[ivarflakstad]

I meant that I managed to improve the original kv cache implementation with little effort. I wasn't dunking on you hehe

And apologies for implying the numbers were inflated. I should have phrased it as "potentially inflated" or something. I was just trying to describe the experience of a reader where I miscategorized the PR because of the description style :)

[ivarflakstad]

Ref adoption, how about something like this?
#3188

[DrJesseGlass]

Great. A trait will be great as we expand the types of KV caches

What about a more expansive trait?

pub trait KvCache {
    fn new(dim: usize, max_seq_len: usize) -> Self;
    fn append(&mut self, k: &Tensor, v: &Tensor) -> Result<(Tensor, Tensor)>;
    fn reset(&mut self);
    fn current_seq_len(&self) -> usize;
    fn dim(&self) -> usize;
    fn get_k(&self) -> Option<&Tensor>;
    fn get_v(&self) -> Option<&Tensor>;
}

Then we can layer a more specialized trait for the indexed / scattered case:

pub trait IndexedKvCache: KvCache {
    fn append_indexed(
        &mut self,
        k: &Tensor,
        v: &Tensor,
        indices_and_mask: &IndicesAndMask,
    ) -> Result<(Tensor, Tensor)>;
}

This would cover the current ScatteredKvCache, and leave room for future block-sparse / FA2-style fused kernels that need indices + masks.

The only other thought I’d add: from my benchmarks, ConcatKvCache has consistently outperformed the current KvCache (slice-based) on CPU, CUDA, and WASM. Longer term, it might make sense to either:

deprecate/rename the existing KvCache struct (e.g. to SliceKvCache), and/or
promote ConcatKvCache as the default KV cache implementation

[ivarflakstad]

What about a more expansive trait?

We'll expand as needed, keeping it as small as possible until a need arises. For example Qwen only needs the three I added.

ConcatKvCache has consistently outperformed the current KvCache (slice-based) on CPU, CUDA, and WASM

Yes. We've actually been using the same approach in the llama implementation for over 2 years, yet somehow it was never made into the default approach.

Then we can layer a more specialized trait for the indexed / scattered case:

The specialized trait won't let us use append_indexed where we have specified that the generic implements KvCache though. And if we say that we expect it to implement IndexedKvCache then we can only use that.
I was initially thinking of just defining append input as an associated type, but with lifetimes etc I think it would become a hassle. Cloning tensors is very cheap though (Arc), so maybe we should simplify and just remove & from the methods 🤔

rename the existing KvCache struct (e.g. to SliceKvCache)

Decent name. I agree. I'd like to keep it around for a while longer. I managed to make it a bit faster, so I'd like to see if there is more to it.

promote ConcatKvCache as the default KV cache implementation

Yes. I was thinking we could add type DefaultKvCache = ConcatKvCache in kv_cache.rs and then use like Qwen3Attention<KvCache: KvCacheTrait = DefaultKvCache> where needed. It would also make it easier for users to pick a kv cache type if we promote it as a type with that name.