Struct BucketedCompileCache 

pub struct BucketedCompileCache { /* private fields */ }

Implementations

impl BucketedCompileCache

pub fn new(device: Device, buckets: Vec<Range<u64>>) -> Self

pub fn power_of_two_ladder(device: Device, min: u64, max: u64) -> Self

Power-of-two ladder over [1, max], with extents [min_pow2, 2·min_pow2, 4·min_pow2, …, max_pow2] where min_pow2 = min.next_power_of_two() and max_pow2 is the smallest power of two ≥ max. Each bucket compiles at its upper-bound extent, so an actual value in bucket (prev_extent .. ext] runs kernels at extent ext (not at the worst case of the whole range). Guarantees compute waste from padding ≤2× — actual > ext / 2 for every bucket except possibly the smallest.

Example: power_of_two_ladder(Device::Cpu, 8, 256) yields buckets 1..9, 9..17, 17..33, 33..65, 65..129, 129..257 with compile extents 8, 16, 32, 64, 128, 256. An actual = 17 runs at extent 32 instead of the 255 a single wide 1..256 bucket would compile at — that’s the “skip compute” win, paid for with O(log max) compiled artifacts instead of one.

pub fn power_of_two_ladder_with_policy( device: Device, min: u64, max: u64, policy: Option<PrecisionPolicy>, ) -> Self

pub fn with_policy( device: Device, buckets: Vec<Range<u64>>, policy: Option<PrecisionPolicy>, ) -> Self

pub fn get_or_compile<F: FnOnce(u64) -> Graph>( &mut self, key: u64, build: F, ) -> Option<(u64, &mut CompiledGraph)>

Find the bucket containing key, compile if needed, return (upper, &mut CompiledGraph) where upper = range.end - 1 is the extent the graph was compiled for. Caller pads inputs to upper before calling run. Returns None if key is outside every bucket — caller decides whether to fall back to a one-off compile.

build receives upper and must return a Graph specialized for that extent.

pub fn get_or_compile_with_options<F: FnOnce(u64) -> Graph>( &mut self, key: u64, build: F, options: &CompileOptions, ) -> Option<(u64, &mut CompiledGraph)>

Like Self::get_or_compile with explicit [CompileOptions].

pub fn get_or_compile_hir<F: FnOnce(u64) -> HirModule>( &mut self, key: u64, build: F, ) -> Option<(u64, &mut CompiledGraph)>

Like Self::get_or_compile but builds and compiles HIR directly through the fusion-first pipeline (Session::compile_hir).

pub fn get_or_compile_hir_with_options<F: FnOnce(u64) -> HirModule>( &mut self, key: u64, build: F, options: &CompileOptions, ) -> Option<(u64, &mut CompiledGraph)>

Like Self::get_or_compile_hir with explicit [CompileOptions] (tier-1 profile, fusion target, …).

pub fn bucket_for(&self, key: u64) -> Option<usize>

Index of the bucket containing key, or None if out of range. Linear scan — bucket counts are small in practice.

pub fn bucket_upper_for_key(&self, key: u64) -> Option<u64>

Upper compile extent for key’s bucket (range.end - 1), without compiling.

pub fn buckets(&self) -> impl Iterator<Item = &Range<u64>>

pub fn compiled_count(&self) -> usize

Number of buckets that have been compiled so far (≤ total buckets).

pub fn compiled_for_key_mut(&mut self, key: u64) -> Option<&mut CompiledGraph>

Mutable compiled graph for key’s bucket, if already compiled.

pub fn compiled_for_upper(&self, upper: u64) -> Option<&CompiledGraph>

Immutable compiled graph for a bucket with compile upper bound upper.

pub fn try_copy_params_between_uppers( &mut self, dst_upper: u64, src_upper: u64, ) -> bool

Copy parameter storage from the bucket compiled at src_upper into dst_upper.

pub fn weight_donor(&self) -> Option<&CompiledGraph>

The bucket currently designated as the shared weight donor, if any.

pub fn set_weight_donor(&mut self, upper: u64)

Record the bucket at upper as the canonical shared weight donor. The first bucket to receive a real weight upload calls this; later buckets share its buffer via Self::try_share_params_from_donor.

pub fn try_share_params_from_donor(&mut self, dst_upper: u64) -> bool

Retain the donor bucket’s weight buffer for the bucket at dst_upper (one GPU copy for both) when their layout matches exactly. Returns false — caller must upload — when there is no donor or the layout differs.

If the recorded donor was evicted, tries any other live compiled bucket as a share source (needed when RLX_KV_CACHE_MAX_RESIDENT=1 climbs).

pub fn seed_resident_kv_prefix_from_keys( &mut self, src_key: u64, dst_key: u64, prefix_tokens: usize, outgoing_upper: usize, kv_dim: usize, n_layers: usize, ) -> bool

D2D seed resident KV from src_key’s bucket into dst_key’s bucket. See rlx-cuda::CudaExecutable::copy_resident_kv_rows_from and rlx-llama32/docs/cuda-gguf-decode.md.

pub fn rebind_resident_kv_hybrid_from_keys( &mut self, src_key: u64, dst_key: u64, host_k: &[Vec<f32>], host_v: &[Vec<f32>], prefix_tokens: usize, outgoing_upper: usize, upper: usize, kv_dim: usize, n_layers: usize, ) -> bool

Hybrid bucket rollover: H2D the host-known prefix once, then copy only rows the host cache does not already have from src (via copy_resident_kv_rows_from). Not used by rlx-llama32 today (generator flush+bind path); kept for experiments.

pub fn total_buckets(&self) -> usize

pub fn evict_except(&mut self, keep: usize)

Drop compiled graphs for all buckets except keep, freeing weight params.

pub fn evict_one_except(&mut self, protect: usize) -> Option<usize>

Drop the single highest-index compiled bucket that is not protect, freeing its weight params, and return its index (or None if no other bucket is compiled). Pass usize::MAX for protect to allow evicting any bucket.

This is the incremental counterpart to Self::evict_except: rather than collapsing to a single resident bucket, callers can free one bucket at a time (highest index first) until a memory budget is met, keeping the rest of the ladder cached. Highest-index-first is the right victim order for a monotonically increasing past_seq sweep — the top bucket is reached last within an utterance and the low buckets recur first at the start of the next, so evicting from the top maximizes cross-utterance reuse.

pub fn clear_compiled(&mut self)

Drop all compiled graphs (free param storage).

pub fn run_padded<F: FnOnce(u64) -> Graph>( &mut self, key: u64, actual_rows: usize, build: F, inputs: &[(&str, &[f32], usize)], output_inners: &[usize], ) -> Option<(u64, Vec<Vec<f32>>)>

“Compile at max, run at less” convenience for inputs and outputs whose outer dimension is the bucket key:

1. Find or compile the bucket containing key.
2. For each input, pad to upper rows along the outer dim using pad_rows (caller passes the inner-dim stride per input; inner = 1 for purely 1D inputs).
3. Run the compiled graph at full extent.
4. Slice each output back to actual_rows along its outer dim. Outputs flagged with inner = 0 in output_inners are returned unsliced (use this for extent-independent outputs like a pooled [hidden] embedding). Missing entries past the end of output_inners are also returned unsliced.

Returns (upper, outputs). Returns None if key falls outside every bucket.

Compute scope: kernels execute at the bucket’s compile extent (upper), not at actual_rows. This means smaller buckets directly translate to less padded compute. With power_of_two_ladder the worst- case waste is bounded at 2×; with hand-tuned buckets it can be arbitrarily tight. True active-extent dispatch — one big compile, kernels short-circuit at runtime — is a separate per-backend change.

pub fn ensure_graph_with_params<F>( &mut self, key: u64, build: F, options: &CompileOptions, ) -> Option<(u64, &mut CompiledGraph)>

where F: FnOnce(u64) -> (Graph, HashMap<String, Vec<f32>>),

Like Self::get_or_compile_with_options but also uploads params on first compile.

pub fn ensure_hir_with_params<F>( &mut self, key: u64, build: F, options: &CompileOptions, ) -> Option<(u64, &mut CompiledGraph)>

where F: FnOnce(u64) -> (HirModule, HashMap<String, Vec<f32>>),

HIR variant of Self::ensure_graph_with_params.

pub fn note_resident_bytes(&mut self, key: u64, extra_bytes: usize)

Record additional resident weight bytes after a packed GGUF upload.

F32 params alone under-count packed Q1_0 / Q4_K arenas (~3–4 GiB), so without this call discrete-VRAM eviction never classifies buckets as “large” and keeps every rung resident → OOM on 16 GB cards.

pub fn run_padded_mixed<F>( &mut self, key: u64, actual_rows: usize, build: F, inputs: &[CacheRunInput<'_>], output_inners: &[usize], ) -> Option<(u64, Vec<Vec<f32>>)>

where F: FnOnce(u64) -> Graph,

Self::run_padded with per-input optional row padding (CacheRunInput).

pub fn sync_all(&mut self)

Drain in-flight GPU work on every compiled bucket.