Skip to main content

BackendPagedKv

Trait BackendPagedKv 

Source
pub trait BackendPagedKv: Backend {
Show 20 methods // Provided methods fn supports_paged_kv() -> bool { ... } fn populate_batched_pointers( _ctx: &mut Self::Context, _k_caches: &[&Self::Buffer], _v_caches: &[&Self::Buffer], _num_layers: usize, _m: usize, ) -> Result<()> { ... } fn split_qkv_norm_rope_into_paged_cache( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _qkv_byte_offset: u64, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _q_out_byte_offset: u64, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _block_table: &Self::Buffer, _tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _pos_offset: usize, _eps: f32, _qk_mode: i32, _cache_len: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()> { ... } fn paged_decode_attention( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _context_lens: &Self::Buffer, _num_seqs: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, _q_len: usize, ) -> Result<()> { ... } fn supports_varlen_qkv() -> bool { ... } fn split_qkv_norm_rope_into_paged_cache_varlen( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _m_total: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()> { ... } fn paged_varlen_attention( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _sliding_window: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()> { ... } fn paged_varlen_attention_fa2_ffi( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _lse: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _seq_lens: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _max_q_len: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()> { ... } fn paged_batched_decode_attention( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _valid_kv_lens: &Self::Buffer, _num_seqs: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()> { ... } fn supports_vllm_paged_attn() -> bool { ... } fn supports_qwen35_paged_qkv() -> bool { ... } fn qwen35_split_qkv_norm_rope_into_paged_cache_varlen( _ctx: &mut Self::Context, _query_raw: &Self::Buffer, _key_raw: &Self::Buffer, _value_raw: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _token_seq_indices: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _rope_dim: usize, _q_proj_stride: usize, _q_head_stride: usize, _kv_proj_stride: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()> { ... } fn supports_qwen35_paged_qkv_vllm() -> bool { ... } fn qwen35_split_qkv_norm_rope_into_paged_cache_varlen_vllm( _ctx: &mut Self::Context, _query_raw: &Self::Buffer, _key_raw: &Self::Buffer, _value_raw: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _token_seq_indices: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _rope_dim: usize, _q_proj_stride: usize, _q_head_stride: usize, _kv_proj_stride: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()> { ... } fn split_qkv_norm_rope_into_paged_cache_vllm( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _qkv_byte_offset: u64, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _q_out_byte_offset: u64, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _block_table: &Self::Buffer, _tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _pos_offset: usize, _eps: f32, _qk_mode: i32, _cache_len: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()> { ... } fn split_qkv_norm_rope_into_paged_cache_varlen_vllm( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _m_total: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()> { ... } fn paged_decode_attention_v2( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _context_lens: &Self::Buffer, _num_seqs: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, _max_seq_len: usize, ) -> Result<()> { ... } fn paged_varlen_attention_vllm_layout( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _context_lens: &Self::Buffer, _num_seqs: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, _q_len: usize, ) -> Result<()> { ... } fn paged_varlen_attention_vllm( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()> { ... } fn paged_varlen_attention_vllm_tiled_q4( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _tile_seqs: &Self::Buffer, _tile_starts: &Self::Buffer, _num_tiles: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()> { ... }
}
Expand description

Capability-trait for backends that support paged KV cache + paged attention.

Provided Methods§

Source

fn supports_paged_kv() -> bool

Whether this backend has a paged-KV decode path (paged_decode_attention etc.). Currently true for Metal, false for CPU. Used to decide the default of FERRUM_METAL_PAGED_KV — the serve path should opt in automatically when supported so users get the bench-quality concurrent-decode numbers without having to learn the flag.

Source

fn populate_batched_pointers( _ctx: &mut Self::Context, _k_caches: &[&Self::Buffer], _v_caches: &[&Self::Buffer], _num_layers: usize, _m: usize, ) -> Result<()>

Pre-populate the per-slot device-pointer scratch arrays used by the batched kernels (kv_cache_append_batched_per_cache and flash_attention_batched_per_cache). Required by the CUDA-graph capture path: the captured graph contains only kernel launches (no captured memcpy_htod), so the device scratch must be fresh when the graph replays.

Caller passes flat layer-major slices: k_caches[li * m + i] and v_caches[li * m + i]. Backend extracts each cache’s device pointer and writes into its corresponding slot in the device scratch via SYNCHRONOUS memcpy (not captured by stream capture).

CUDA-only; other backends fall through to the default unsupported and the caller skips the population call.

Source

fn split_qkv_norm_rope_into_paged_cache( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _qkv_byte_offset: u64, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _q_out_byte_offset: u64, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _block_table: &Self::Buffer, _tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _pos_offset: usize, _eps: f32, _qk_mode: i32, _cache_len: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()>

Paged-KV variant of [Self::split_qkv_norm_rope_into_cache].

Same fused split + qk-norm + RoPE, but K/V are written into a paged pool [num_blocks, kv_heads, block_size, head_dim] indexed via block_table[logical_block] → physical_block. Q still goes to head-major scratch.

Default returns Unsupported. Backends that lack a paged kernel keep using the contiguous variant. qkv_byte_offset / q_out_byte_offset let the caller pass a slice of a larger batched buffer (used by the multi-seq paged path in decode_batch_internal). For single-seq dispatch they should be 0.

Source

fn paged_decode_attention( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _context_lens: &Self::Buffer, _num_seqs: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, _q_len: usize, ) -> Result<()>

Paged-KV variant of [Self::flash_attention].

Decode (q_len == 1): q/out: [num_seqs, num_heads, head_dim] (token-major)

Causal prefill (q_len > 1, single seq): q/out: [num_heads, q_len, head_dim] (head-major — the layout produced by split_qkv_norm_rope_into_paged_cache) The kernel applies a per-q-token causal mask using context_lens[seq] as the FINAL kv_len (= pos_offset + q_len): token i sees positions [0, context_lens - q_len + 1 + i).

Common to both: k_pool/v_pool: [num_blocks, num_kv_heads, block_size, head_dim] block_tables: [num_seqs, max_num_blocks_per_seq] u32 context_lens: [num_seqs] u32

Backends without a paged kernel return Unsupported; callers are expected to fall back to contiguous KV.

Source

fn supports_varlen_qkv() -> bool

Capability: does this backend implement split_qkv_norm_rope_into_paged_cache_varlen and paged_varlen_attention? Required by the unified mixed-batch forward path used by LlamaFamilyModel::unified_forward. Default false; backends that ship the varlen kernels override.

Source

fn split_qkv_norm_rope_into_paged_cache_varlen( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _m_total: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()>

Varlen variant of Self::split_qkv_norm_rope_into_paged_cache.

Single launch covering ALL sequences in the batch. Reads pos_offsets[seq], cu_seqlens_q[seq], and the per-seq block_table from device buffers — graph-capturable (the per-iter state is in buffers, not kernel scalars). Replaces the per-item dispatch loop in unified_forward_layer with one call.

Layouts:

  • qkv: [m_total, q_dim + 2 * kv_dim] token-major
  • q_out: [m_total, q_heads, head_dim] token-major (matches what paged_varlen_attention reads)
  • cache_k / cache_v: paged pool same as paged_varlen_attention
  • cu_seqlens_q: [num_seqs + 1] u32 prefix sum
  • pos_offsets: [num_seqs] u32, starting kv_pos per seq
  • block_tables: [num_seqs, max_blocks_per_seq] i32 stacked
Source

fn paged_varlen_attention( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _sliding_window: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()>

Variable-length paged attention with GQA + causal mask.

Supports a unified mixed batch where each sequence contributes 1 (decode) or N (prefill chunk) query tokens — the workhorse for chunked-prefill. See kernels/paged_varlen_attention.cu for the kernel itself.

Layouts:

  • q / out: [total_q_tokens, num_heads, head_dim] (token- major, FP16). total_q_tokens = cu_seqlens_q[num_seqs].
  • k_pool / v_pool: paged block pool, layout matches paged_decode_attention.
  • cu_seqlens_q: [num_seqs + 1] u32 prefix sum, with cu_seqlens_q[0] = 0 and cu_seqlens_q[num_seqs] = total_q_tokens.
  • pos_offsets: [num_seqs] u32, the starting absolute KV position of each seq’s first q token (= prior kv_len).
  • block_tables: [num_seqs, max_num_blocks_per_seq] i32 grid.

Each query token attends causally to KV positions [0, pos_offsets[s] + local_idx] when sliding_window == 0, or only the most recent sliding_window positions when non-zero.

Source

fn paged_varlen_attention_fa2_ffi( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _lse: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _seq_lens: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _max_q_len: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()>

Opt-in vLLM FlashAttention-2 FFI path for FA-layout paged KV.

This is intentionally separate from Self::paged_varlen_attention: it needs the final per-sequence KV lengths (seq_lens) and an explicit LSE scratch buffer because the external FA2 runner writes softmax LSE. Default returns Err(unsupported); CUDA overrides when a runtime shim is provided via FERRUM_FA2_DIRECT_FFI_SHIM.

Source

fn paged_batched_decode_attention( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _valid_kv_lens: &Self::Buffer, _num_seqs: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()>

Batched paged decode attention — multi-seq, single token per seq. Faster path for the unified_forward layer when m_total == num_seqs (every item is a single-token decode). Skips the cu_seqlens_q linear scan that paged_varlen_attention does in the fully-mixed case.

Layouts: q : [num_seqs, num_q_heads, head_dim] k_pool/v_pool : paged pool (same as paged_varlen) block_tables : [num_seqs, max_num_blocks_per_seq] valid_kv_lens : [num_seqs] — current kv_len per seq out : [num_seqs, num_q_heads, head_dim]

Default returns Err(unsupported); CUDA backend overrides.

Source

fn supports_vllm_paged_attn() -> bool

Capability: backend has vLLM-layout paged KV write kernels and the paged_attention_v2 decode kernel. Models that opt into this layout at construction time (via FERRUM_USE_VLLM_PAGED_ATTN=1) must dispatch ALL paged writes and reads through the _vllm variants — the layouts are not compatible. Default false.

Source

fn supports_qwen35_paged_qkv() -> bool

Qwen3.5 full-attention uses separate q/k/v projections and partial RoPE/gated-Q layout, so it cannot use the fused-QKV paged writer directly. Backends that implement this method can write those separate projections into Ferrum’s legacy paged pool consumed by Self::paged_batched_decode_attention.

Source

fn qwen35_split_qkv_norm_rope_into_paged_cache_varlen( _ctx: &mut Self::Context, _query_raw: &Self::Buffer, _key_raw: &Self::Buffer, _value_raw: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _token_seq_indices: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _rope_dim: usize, _q_proj_stride: usize, _q_head_stride: usize, _kv_proj_stride: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()>

Source

fn supports_qwen35_paged_qkv_vllm() -> bool

vLLM-layout variant of the Qwen3.5 separate q/k/v writer. K/V are written in the layout consumed by Self::paged_decode_attention_v2, while Q remains token-major [total_q_tokens, q_heads, head_dim].

Source

fn qwen35_split_qkv_norm_rope_into_paged_cache_varlen_vllm( _ctx: &mut Self::Context, _query_raw: &Self::Buffer, _key_raw: &Self::Buffer, _value_raw: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _token_seq_indices: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _rope_dim: usize, _q_proj_stride: usize, _q_head_stride: usize, _kv_proj_stride: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()>

Source

fn split_qkv_norm_rope_into_paged_cache_vllm( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _qkv_byte_offset: u64, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _q_out_byte_offset: u64, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _block_table: &Self::Buffer, _tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _pos_offset: usize, _eps: f32, _qk_mode: i32, _cache_len: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()>

vLLM-layout variant of Self::split_qkv_norm_rope_into_paged_cache. K/V are written in vLLM’s paged_attention_v2 layout: K is [num_blocks, kv_heads, head_dim/x, block_size, x] (x = 16/sizeof(elem)), V is [num_blocks, kv_heads, head_dim, block_size]. Q output and every other argument matches the non-vllm variant exactly so the model layer can swap dispatchers based on a single flag.

Source

fn split_qkv_norm_rope_into_paged_cache_varlen_vllm( _ctx: &mut Self::Context, _qkv: &Self::Buffer, _q_norm_w: &Self::Buffer, _k_norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _q_out: &mut Self::Buffer, _cache_k: &mut Self::Buffer, _cache_v: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _m_total: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _eps: f32, _qk_mode: i32, _block_size: usize, _max_blocks_per_seq: usize, ) -> Result<()>

vLLM-layout variant of Self::split_qkv_norm_rope_into_paged_cache_varlen. Same signature — only the K/V cache layout changes.

Source

fn paged_decode_attention_v2( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _context_lens: &Self::Buffer, _num_seqs: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, _max_seq_len: usize, ) -> Result<()>

vLLM paged_attention_v2 — multi-partition split-K decode attention reading the vLLM K/V layout. q_len is implicitly 1 (decode only; vLLM’s v2 kernel does not support q_len > 1). max_seq_len is the max kv_len across the batch — used to size the partition reduction.

Source

fn paged_varlen_attention_vllm_layout( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _block_tables: &Self::Buffer, _context_lens: &Self::Buffer, _num_seqs: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, _q_len: usize, ) -> Result<()>

q_len>1 prefill/chunk-prefill attention over vLLM-layout paged KV. This keeps cache layout consistent when FERRUM_USE_VLLM_PAGED_ATTN=1 and the prompt path writes K/V in the layout consumed later by paged_decode_attention_v2.

Source

fn paged_varlen_attention_vllm( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _num_seqs: usize, _total_q_tokens: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()>

Variable-length paged attention over vLLM-layout paged KV.

Unlike Self::paged_varlen_attention_vllm_layout, this accepts the same varlen index tensors as Self::paged_varlen_attention and writes token-major output directly. It is the unified mixed-batch companion for split_qkv_norm_rope_into_paged_cache_varlen_vllm.

Source

fn paged_varlen_attention_vllm_tiled_q4( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_pool: &Self::Buffer, _v_pool: &Self::Buffer, _out: &mut Self::Buffer, _cu_seqlens_q: &Self::Buffer, _pos_offsets: &Self::Buffer, _block_tables: &Self::Buffer, _tile_seqs: &Self::Buffer, _tile_starts: &Self::Buffer, _num_tiles: usize, _max_kv_len: usize, _num_heads: usize, _num_kv_heads: usize, _head_dim: usize, _block_size: usize, _max_num_blocks_per_seq: usize, ) -> Result<()>

Q-tiled vLLM-layout varlen attention. tile_seqs and tile_starts describe a compact list of q-token tiles, avoiding empty grid blocks for mixed batches that contain both long prefill items and q_len=1 decode items. Semantics match Self::paged_varlen_attention_vllm.

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§