Trait Backend 

pub trait Backend:
    Send
    + Sync
    + Sized
    + 'static {
    type Buffer: Send + Sync;
    type Context;
    type Timer: BackendTimer<Self>;

    // Required methods
    fn make_timer() -> Self::Timer;
    fn new_context() -> Self::Context;
    fn sync(ctx: &mut Self::Context);
    fn alloc_typed(dtype: Dtype, n: usize) -> Self::Buffer;
    fn from_slice_typed<T: HostDtype>(data: &[T]) -> Self::Buffer;
    fn write_typed<T: HostDtype>(
        ctx: &mut Self::Context,
        dst: &mut Self::Buffer,
        data: &[T],
    );
    fn gemm(
        ctx: &mut Self::Context,
        a: &Self::Buffer,
        b: &Self::Buffer,
        out: &mut Self::Buffer,
        m: usize,
        n: usize,
        k: usize,
    );
    fn rms_norm(
        ctx: &mut Self::Context,
        x: &Self::Buffer,
        w: &Self::Buffer,
        eps: f32,
        out: &mut Self::Buffer,
        tokens: usize,
        dim: usize,
    );
    fn fused_add_rms_norm(
        ctx: &mut Self::Context,
        residual: &mut Self::Buffer,
        x: &Self::Buffer,
        w: &Self::Buffer,
        eps: f32,
        out: &mut Self::Buffer,
        tokens: usize,
        dim: usize,
    );
    fn flash_attention(
        ctx: &mut Self::Context,
        q: &Self::Buffer,
        k: &Self::Buffer,
        v: &Self::Buffer,
        out: &mut Self::Buffer,
        batch: usize,
        q_len: usize,
        kv_len: usize,
        pos_offset: usize,
        cfg: &AttnConfig,
    );
    fn copy_slice(
        ctx: &mut Self::Context,
        src: &Self::Buffer,
        src_offset: usize,
        dst: &mut Self::Buffer,
        dst_offset: usize,
        len: usize,
    );
    fn embedding_lookup(
        ctx: &mut Self::Context,
        table: &Self::Buffer,
        ids: &[u32],
        out: &mut Self::Buffer,
        dim: usize,
    );
    fn split_qkv(
        ctx: &mut Self::Context,
        qkv: &Self::Buffer,
        q: &mut Self::Buffer,
        k: &mut Self::Buffer,
        v: &mut Self::Buffer,
        tokens: usize,
        q_dim: usize,
        kv_dim: usize,
    );
    fn fused_silu_mul_split(
        ctx: &mut Self::Context,
        gate_up: &Self::Buffer,
        out: &mut Self::Buffer,
        tokens: usize,
        im: usize,
    );
    fn qk_norm_rope(
        ctx: &mut Self::Context,
        input: &Self::Buffer,
        norm_w: &Self::Buffer,
        cos: &Self::Buffer,
        sin: &Self::Buffer,
        output: &mut Self::Buffer,
        tokens: usize,
        heads: usize,
        head_dim: usize,
        pos_offset: usize,
        eps: f32,
        mode: i32,
    );
    fn kv_cache_append_head_major(
        ctx: &mut Self::Context,
        cache_k: &mut Self::Buffer,
        cache_v: &mut Self::Buffer,
        cache_len: usize,
        cache_capacity: usize,
        new_k_head_major: &Self::Buffer,
        new_v_head_major: &Self::Buffer,
        new_tokens: usize,
        nkv: usize,
        hd: usize,
    );
    fn transpose_head_to_token(
        ctx: &mut Self::Context,
        src: &Self::Buffer,
        dst: &mut Self::Buffer,
        tokens: usize,
        heads: usize,
        dim: usize,
    );
    fn add_inplace(
        ctx: &mut Self::Context,
        residual: &mut Self::Buffer,
        x: &Self::Buffer,
        len: usize,
    );
    fn add_bias(
        ctx: &mut Self::Context,
        data: &mut Self::Buffer,
        bias: &Self::Buffer,
        rows: usize,
        cols: usize,
    );
    fn layer_norm(
        ctx: &mut Self::Context,
        x: &Self::Buffer,
        gamma: &Self::Buffer,
        beta: &Self::Buffer,
        eps: f32,
        out: &mut Self::Buffer,
        tokens: usize,
        dim: usize,
    );
    fn gelu(
        ctx: &mut Self::Context,
        x: &Self::Buffer,
        out: &mut Self::Buffer,
        len: usize,
    );
    fn alloc(len: usize) -> Self::Buffer;
    fn to_vec(buf: &Self::Buffer, len: usize) -> Vec<f32>;
    fn from_slice(data: &[f32]) -> Self::Buffer;

    // Provided methods
    fn with_device_ordinal<R>(
        _device_ordinal: Option<usize>,
        body: impl FnOnce() -> R,
    ) -> R { ... }
    fn supports_device_ordinal_scope() -> bool { ... }
    fn graph_capture_in_flight(_ctx: &Self::Context) -> bool { ... }
    fn sync_before_host_readback(_ctx: &mut Self::Context) { ... }
    fn activation_elem_size_bytes() -> usize { ... }
    fn supports_llama_family_batched_decode() -> bool { ... }
    fn supports_qk_norm_rope_batched_per_item() -> bool { ... }
    fn supports_kv_cache_append_batched_per_cache() -> bool { ... }
    fn supports_flash_attention_batched_per_cache() -> bool { ... }
    fn zero_buffer(
        _ctx: &mut Self::Context,
        _buf: &mut Self::Buffer,
        _len: usize,
    ) -> Result<()> { ... }
    fn mla_attention(
        _ctx: &mut Self::Context,
        _q: &Self::Buffer,
        _kv_compressed: &Self::Buffer,
        _kv_rope: &Self::Buffer,
        _out: &mut Self::Buffer,
        _batch: usize,
        _q_len: usize,
        _kv_len: usize,
        _pos_offset: usize,
        _cfg: &AttnConfig,
        _kv_lora_rank: usize,
        _qk_rope_head_dim: usize,
    ) -> Result<()> { ... }
    fn recurrent_gated_delta_rule_f32(
        _ctx: &mut Self::Context,
        _query: &Self::Buffer,
        _key: &Self::Buffer,
        _value: &Self::Buffer,
        _g: &Self::Buffer,
        _beta: &Self::Buffer,
        _initial_state: &Self::Buffer,
        _out: &mut Self::Buffer,
        _final_state: &mut Self::Buffer,
        _tokens: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _use_qk_l2norm: bool,
        _scale: f32,
    ) -> Result<()> { ... }
    fn recurrent_gated_delta_rule_batch_f32(
        _ctx: &mut Self::Context,
        _query: &Self::Buffer,
        _key: &Self::Buffer,
        _value: &Self::Buffer,
        _g: &Self::Buffer,
        _beta: &Self::Buffer,
        _initial_states: &Self::Buffer,
        _out: &mut Self::Buffer,
        _final_states: &mut Self::Buffer,
        _batch: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _use_qk_l2norm: bool,
        _scale: f32,
    ) -> Result<()> { ... }
    fn supports_qwen35_indexed_recurrent_state() -> bool { ... }
    fn qwen35_indexed_recurrent_state_dtype() -> Dtype { ... }
    fn supports_qwen35_packed_gdn_decode_prepare() -> bool { ... }
    fn supports_qwen35_packed_gdn_prefill_prepare() -> bool { ... }
    fn supports_qwen35_packed_gdn_recurrent_decode() -> bool { ... }
    fn recurrent_gated_delta_rule_batch_indexed_f32(
        _ctx: &mut Self::Context,
        _query: &Self::Buffer,
        _key: &Self::Buffer,
        _value: &Self::Buffer,
        _g: &Self::Buffer,
        _beta: &Self::Buffer,
        _state_slots: &mut Self::Buffer,
        _slot_indices: &Self::Buffer,
        _out: &mut Self::Buffer,
        _batch: usize,
        _max_slots: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _use_qk_l2norm: bool,
        _scale: f32,
    ) -> Result<()> { ... }
    fn recurrent_gated_delta_rule_batch_indexed_packed_f32(
        _ctx: &mut Self::Context,
        _mixed_qkv: &Self::Buffer,
        _ba_raw: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _state_slots: &mut Self::Buffer,
        _slot_indices: &Self::Buffer,
        _out: &mut Self::Buffer,
        _batch: usize,
        _max_slots: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _scale: f32,
    ) -> Result<()> { ... }
    fn recurrent_gated_delta_rule_varlen_f32(
        _ctx: &mut Self::Context,
        _query: &Self::Buffer,
        _key: &Self::Buffer,
        _value: &Self::Buffer,
        _g: &Self::Buffer,
        _beta: &Self::Buffer,
        _initial_states: &Self::Buffer,
        _cu_seqlens: &Self::Buffer,
        _out: &mut Self::Buffer,
        _final_states: &mut Self::Buffer,
        _batch: usize,
        _total_tokens: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _use_qk_l2norm: bool,
        _scale: f32,
    ) -> Result<()> { ... }
    fn linear_attention_prepare_f32(
        _ctx: &mut Self::Context,
        _mixed_qkv_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _a_raw: &Self::Buffer,
        _b_raw: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _query: &mut Self::Buffer,
        _key: &mut Self::Buffer,
        _value: &mut Self::Buffer,
        _g: &mut Self::Buffer,
        _beta: &mut Self::Buffer,
        _tokens: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
        _apply_qk_l2norm: bool,
    ) -> Result<()> { ... }
    fn linear_attention_prepare_varlen_f32(
        _ctx: &mut Self::Context,
        _mixed_qkv_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _initial_conv_states: &Self::Buffer,
        _a_raw: &Self::Buffer,
        _b_raw: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _cu_seqlens: &Self::Buffer,
        _token_seq_indices: &Self::Buffer,
        _query: &mut Self::Buffer,
        _key: &mut Self::Buffer,
        _value: &mut Self::Buffer,
        _g: &mut Self::Buffer,
        _beta: &mut Self::Buffer,
        _final_conv_states: &mut Self::Buffer,
        _batch: usize,
        _total_tokens: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
        _apply_qk_l2norm: bool,
    ) -> Result<()> { ... }
    fn linear_attention_prepare_varlen_packed_qkvz_ba_f32(
        _ctx: &mut Self::Context,
        _mixed_qkvz_raw: &Self::Buffer,
        _ba_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _initial_conv_states: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _cu_seqlens: &Self::Buffer,
        _token_seq_indices: &Self::Buffer,
        _query: &mut Self::Buffer,
        _key: &mut Self::Buffer,
        _value: &mut Self::Buffer,
        _z: &mut Self::Buffer,
        _g: &mut Self::Buffer,
        _beta: &mut Self::Buffer,
        _final_conv_states: &mut Self::Buffer,
        _batch: usize,
        _total_tokens: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
        _apply_qk_l2norm: bool,
    ) -> Result<()> { ... }
    fn linear_attention_decode_prepare_f32(
        _ctx: &mut Self::Context,
        _mixed_qkv_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _conv_state: &Self::Buffer,
        _a_raw: &Self::Buffer,
        _b_raw: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _query: &mut Self::Buffer,
        _key: &mut Self::Buffer,
        _value: &mut Self::Buffer,
        _g: &mut Self::Buffer,
        _beta: &mut Self::Buffer,
        _next_conv_state: &mut Self::Buffer,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
        _apply_qk_l2norm: bool,
    ) -> Result<()> { ... }
    fn linear_attention_decode_prepare_batch_f32(
        _ctx: &mut Self::Context,
        _mixed_qkv_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _conv_states: &Self::Buffer,
        _a_raw: &Self::Buffer,
        _b_raw: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _query: &mut Self::Buffer,
        _key: &mut Self::Buffer,
        _value: &mut Self::Buffer,
        _g: &mut Self::Buffer,
        _beta: &mut Self::Buffer,
        _next_conv_states: &mut Self::Buffer,
        _batch: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
        _apply_qk_l2norm: bool,
    ) -> Result<()> { ... }
    fn linear_attention_decode_prepare_batch_indexed_f32(
        _ctx: &mut Self::Context,
        _mixed_qkv_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _conv_state_slots: &mut Self::Buffer,
        _slot_indices: &Self::Buffer,
        _a_raw: &Self::Buffer,
        _b_raw: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _query: &mut Self::Buffer,
        _key: &mut Self::Buffer,
        _value: &mut Self::Buffer,
        _g: &mut Self::Buffer,
        _beta: &mut Self::Buffer,
        _batch: usize,
        _max_slots: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
        _apply_qk_l2norm: bool,
    ) -> Result<()> { ... }
    fn linear_attention_decode_prepare_batch_indexed_packed_qkvz_ba_f32(
        _ctx: &mut Self::Context,
        _mixed_qkvz_raw: &Self::Buffer,
        _ba_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _conv_state_slots: &mut Self::Buffer,
        _slot_indices: &Self::Buffer,
        _a_log: &Self::Buffer,
        _dt_bias: &Self::Buffer,
        _query: &mut Self::Buffer,
        _key: &mut Self::Buffer,
        _value: &mut Self::Buffer,
        _z: &mut Self::Buffer,
        _g: &mut Self::Buffer,
        _beta: &mut Self::Buffer,
        _batch: usize,
        _max_slots: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
        _apply_qk_l2norm: bool,
    ) -> Result<()> { ... }
    fn linear_attention_decode_prepare_batch_indexed_packed_qkvz_to_mixed_f32(
        _ctx: &mut Self::Context,
        _mixed_qkvz_raw: &Self::Buffer,
        _conv_weight: &Self::Buffer,
        _conv_state_slots: &mut Self::Buffer,
        _slot_indices: &Self::Buffer,
        _mixed_qkv: &mut Self::Buffer,
        _z: &mut Self::Buffer,
        _batch: usize,
        _max_slots: usize,
        _key_heads: usize,
        _value_heads: usize,
        _key_dim: usize,
        _value_dim: usize,
        _conv_kernel: usize,
    ) -> Result<()> { ... }
    fn gated_rms_norm_f32(
        _ctx: &mut Self::Context,
        _core: &Self::Buffer,
        _z: &Self::Buffer,
        _weight: &Self::Buffer,
        _out: &mut Self::Buffer,
        _tokens: usize,
        _heads: usize,
        _dim: usize,
        _eps: f32,
    ) -> Result<()> { ... }
    fn embedding_lookup_dev(
        ctx: &mut Self::Context,
        table: &Self::Buffer,
        ids: &Self::Buffer,
        out: &mut Self::Buffer,
        batch: usize,
        dim: usize,
    ) { ... }
    fn fused_gelu_tanh_mul_split(
        _ctx: &mut Self::Context,
        _gate_up: &Self::Buffer,
        _out: &mut Self::Buffer,
        _tokens: usize,
        _im: usize,
    ) { ... }
    fn scale_inplace(
        ctx: &mut Self::Context,
        buf: &mut Self::Buffer,
        scale: f32,
        len: usize,
    ) { ... }
    fn qk_norm_rope_partial(
        ctx: &mut Self::Context,
        input: &Self::Buffer,
        norm_w: &Self::Buffer,
        cos: &Self::Buffer,
        sin: &Self::Buffer,
        output: &mut Self::Buffer,
        tokens: usize,
        heads: usize,
        head_dim: usize,
        rope_dim: usize,
        input_stride: usize,
        input_offset: usize,
        input_head_stride: usize,
        pos_offset: usize,
        eps: f32,
        mode: i32,
    ) -> Result<()> { ... }
    fn qwen35_apply_attention_gate(
        _ctx: &mut Self::Context,
        _context: &mut Self::Buffer,
        _query_raw: &Self::Buffer,
        _tokens: usize,
        _q_total: usize,
        _q_proj_total: usize,
        _head_dim: usize,
    ) -> Result<()> { ... }
    fn qwen35_apply_token_gate(
        _ctx: &mut Self::Context,
        _values: &mut Self::Buffer,
        _gate: &Self::Buffer,
        _tokens: usize,
        _hidden_size: usize,
    ) -> Result<()> { ... }
    fn qwen35_apply_token_gate_and_add_inplace(
        ctx: &mut Self::Context,
        dst: &mut Self::Buffer,
        values: &mut Self::Buffer,
        gate: &Self::Buffer,
        tokens: usize,
        hidden_size: usize,
    ) -> Result<()> { ... }
    fn qwen35_interleave_gate_up(
        ctx: &mut Self::Context,
        gate: &Self::Buffer,
        up: &Self::Buffer,
        out: &mut Self::Buffer,
        tokens: usize,
        intermediate: usize,
    ) -> Result<()> { ... }
    fn kv_cache_append_batched_per_cache(
        _ctx: &mut Self::Context,
        _caches: &[&Self::Buffer],
        _new_data: &Self::Buffer,
        _cache_lens: &Self::Buffer,
        _capacity: usize,
        _m: usize,
        _nkv: usize,
        _hd: usize,
        _slot: usize,
    ) -> Result<()> { ... }
    fn flash_attention_batched_per_cache(
        _ctx: &mut Self::Context,
        _q: &Self::Buffer,
        _k_caches: &[&Self::Buffer],
        _v_caches: &[&Self::Buffer],
        _kv_lens: &Self::Buffer,
        _out: &mut Self::Buffer,
        _nq: usize,
        _nkv: usize,
        _hd: usize,
        _scale: f32,
        _max_valid_kv: usize,
        _capacity: usize,
        _sliding_window: usize,
        _slot: usize,
    ) -> Result<()> { ... }
    fn qk_norm_rope_batched_per_item(
        _ctx: &mut Self::Context,
        _input: &Self::Buffer,
        _norm_w: &Self::Buffer,
        _cos: &Self::Buffer,
        _sin: &Self::Buffer,
        _output: &mut Self::Buffer,
        _positions: &Self::Buffer,
        _m: usize,
        _heads: usize,
        _head_dim: usize,
        _eps: f32,
        _mode: i32,
    ) -> Result<()> { ... }
    fn split_qkv_norm_rope(
        _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,
        _k_out: &mut Self::Buffer,
        _v_out: &mut Self::Buffer,
        _tokens: usize,
        _q_heads: usize,
        _kv_heads: usize,
        _head_dim: usize,
        _pos_offset: usize,
        _eps: f32,
        _qk_mode: i32,
    ) -> Result<()> { ... }
    fn split_qkv_norm_rope_into_cache(
        _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,
        _tokens: usize,
        _q_heads: usize,
        _kv_heads: usize,
        _head_dim: usize,
        _pos_offset: usize,
        _eps: f32,
        _qk_mode: i32,
        _cache_len: usize,
        _cache_capacity: usize,
    ) -> Result<()> { ... }
    fn transpose_token_to_head(
        _ctx: &mut Self::Context,
        _src: &Self::Buffer,
        _dst: &mut Self::Buffer,
        _tokens: usize,
        _heads: usize,
        _dim: usize,
    ) { ... }
    fn scaled_add_inplace(
        _ctx: &mut Self::Context,
        dst: &mut Self::Buffer,
        src: &Self::Buffer,
        scale: f32,
        len: usize,
    ) { ... }
    fn fused_silu_mul_split_strided(
        _ctx: &mut Self::Context,
        _gate_up: &Self::Buffer,
        _in_row_offset: usize,
        _out: &mut Self::Buffer,
        _out_row_offset: usize,
        _tokens: usize,
        _intermediate: usize,
    ) { ... }
    fn write_f32_to_activation(
        ctx: &mut Self::Context,
        dst: &mut Self::Buffer,
        data: &[f32],
    ) { ... }
    fn f32_to_activation(
        ctx: &mut Self::Context,
        input_f32: &Self::Buffer,
        out: &mut Self::Buffer,
        len: usize,
    ) { ... }
    fn supports_device_f32_residual_shadow() -> bool { ... }
    fn activation_to_f32_shadow(
        ctx: &mut Self::Context,
        src: &Self::Buffer,
        dst_f32: &mut Self::Buffer,
        len: usize,
    ) { ... }
    fn activation_add_to_f32_shadow(
        ctx: &mut Self::Context,
        src: &Self::Buffer,
        residual_f32: &mut Self::Buffer,
        scratch_f32: &mut Self::Buffer,
        len: usize,
    ) { ... }
    fn rms_norm_activation_to_f32(
        ctx: &mut Self::Context,
        input: &Self::Buffer,
        weight: &Self::Buffer,
        eps: f32,
        out_f32: &mut Self::Buffer,
        tokens: usize,
        dim: usize,
    ) { ... }
    fn rms_norm_activation_add_to_f32(
        ctx: &mut Self::Context,
        input: &Self::Buffer,
        weight: &Self::Buffer,
        eps: f32,
        residual_f32: &mut Self::Buffer,
        scratch_f32: &mut Self::Buffer,
        tokens: usize,
        dim: usize,
    ) { ... }
    fn rms_norm_f32_to_activation(
        ctx: &mut Self::Context,
        input_f32: &Self::Buffer,
        weight: &Self::Buffer,
        eps: f32,
        out: &mut Self::Buffer,
        tokens: usize,
        dim: usize,
    ) { ... }
    fn argmax_rows_f16(
        _ctx: &mut Self::Context,
        logits: &Self::Buffer,
        m: usize,
        n: usize,
    ) -> Result<Vec<u32>> { ... }
    fn argmax_rows_f16_masked(
        _ctx: &mut Self::Context,
        _logits: &Self::Buffer,
        _valid_token_mask: &Self::Buffer,
        _mask_len: usize,
        _m: usize,
        _n: usize,
    ) -> Result<Vec<u32>> { ... }
    fn supports_argmax_rows_f16_sparse_repetition_penalty() -> bool { ... }
    fn argmax_rows_f16_sparse_repetition_penalty(
        _ctx: &mut Self::Context,
        _logits: &mut Self::Buffer,
        _valid_token_mask: Option<(&Self::Buffer, usize)>,
        _row_offsets: &Self::Buffer,
        _token_ids: &Self::Buffer,
        _repetition_penalties: &Self::Buffer,
        _total_token_ids: usize,
        _m: usize,
        _n: usize,
    ) -> Result<Vec<u32>> { ... }
    fn from_weight_bytes(raw: &[u8], src_dtype: SrcDtype) -> Self::Buffer { ... }
}

The core abstraction over CUDA / Metal / CPU.

Key design: operations take a &mut Self::Context which accumulates work.

• CPU: Context is () — ops execute immediately.
• Metal: Context is a CommandBuffer — ops encode into it, flushed on sync().
• CUDA: Context is a CudaStream — ops launch on the stream, synced on sync().

layer_forward passes the context through all ops in a layer. ModelRunner calls sync() only when it needs results (e.g., reading logits).

Required Associated Types

type Buffer: Send + Sync

type Context

Execution context that accumulates GPU work.

• CPU: () (no-op, ops execute inline)
• Metal: wraps a CommandBuffer
• CUDA: wraps a CudaStream

type Timer: BackendTimer<Self>

GPU-side timer scoped to this backend. See super::timer — CPU: Instant; Metal: sync-wrap; CUDA: cuEvent. PLAYBOOK § 1.1.

Required Methods

fn make_timer() -> Self::Timer

Factory for Self::Timer — exists so call sites that have a <B: Backend> parameter can spawn a timer without importing the concrete impl. PLAYBOOK § 1.2.

fn new_context() -> Self::Context

Opaque per-backend GPTQ weight representation.

• CPU: dequantized f32 weights (run as regular GEMM)
• Metal: () — unsupported; gemm_gptq errors Create a new execution context (begin accumulating work).

fn sync(ctx: &mut Self::Context)

Flush accumulated work and wait for completion. CPU: no-op. Metal: commit + waitUntilCompleted. CUDA: stream sync.

fn alloc_typed(dtype: Dtype, n: usize) -> Self::Buffer

Phase D step 2+3: unified typed allocator. Replaces per-dtype alloc_u32 / alloc_typed_i32 / etc. The buffer is dtype- tagged at the wrapper level (CudaBuf::U32, MetalBuf with Dtype::U32, CpuBuf::U32), so reads/writes through .as_<T>() accessors get the correct byte count automatically.

fn from_slice_typed<T: HostDtype>(data: &[T]) -> Self::Buffer

Upload typed host data — replaces from_slice_i32 / from_slice_u32 etc. The host element type T carries its Dtype via the HostDtype marker so dispatch in the impl is a one-line match T::DTYPE.

fn write_typed<T: HostDtype>( ctx: &mut Self::Context, dst: &mut Self::Buffer, data: &[T], )

In-place typed write — replaces write_u32 / write_i32_into / write_f32_into. The buffer must already be dtype-tagged matching T::DTYPE (typically alloc’d via alloc_typed or from_slice_typed).

fn gemm( ctx: &mut Self::Context, a: &Self::Buffer, b: &Self::Buffer, out: &mut Self::Buffer, m: usize, n: usize, k: usize, )

fn rms_norm( ctx: &mut Self::Context, x: &Self::Buffer, w: &Self::Buffer, eps: f32, out: &mut Self::Buffer, tokens: usize, dim: usize, )

fn fused_add_rms_norm( ctx: &mut Self::Context, residual: &mut Self::Buffer, x: &Self::Buffer, w: &Self::Buffer, eps: f32, out: &mut Self::Buffer, tokens: usize, dim: usize, )

fn flash_attention( ctx: &mut Self::Context, q: &Self::Buffer, k: &Self::Buffer, v: &Self::Buffer, out: &mut Self::Buffer, batch: usize, q_len: usize, kv_len: usize, pos_offset: usize, cfg: &AttnConfig, )

fn copy_slice( ctx: &mut Self::Context, src: &Self::Buffer, src_offset: usize, dst: &mut Self::Buffer, dst_offset: usize, len: usize, )

Copy len floats from src[src_offset..] to dst[dst_offset..].

Needed for Qwen3Model::prefill to pluck the last token’s hidden state out of residual[seq_len, h] without round-tripping through host RAM. Backend::copy is the offset-free variant; copy_slice additionally supports non-zero source and destination offsets.

fn embedding_lookup( ctx: &mut Self::Context, table: &Self::Buffer, ids: &[u32], out: &mut Self::Buffer, dim: usize, )

fn split_qkv( ctx: &mut Self::Context, qkv: &Self::Buffer, q: &mut Self::Buffer, k: &mut Self::Buffer, v: &mut Self::Buffer, tokens: usize, q_dim: usize, kv_dim: usize, )

Split fused QKV [tokens, q_dim+2*kv_dim] into separate Q, K, V buffers. Q: [tokens, q_dim], K: [tokens, kv_dim], V: [tokens, kv_dim]

fn fused_silu_mul_split( ctx: &mut Self::Context, gate_up: &Self::Buffer, out: &mut Self::Buffer, tokens: usize, im: usize, )

Split fused gate_up [tokens, 2*im] into gate [tokens, im] and up [tokens, im], then compute SiLU(gate) * up → out [tokens, im].

fn qk_norm_rope( ctx: &mut Self::Context, input: &Self::Buffer, norm_w: &Self::Buffer, cos: &Self::Buffer, sin: &Self::Buffer, output: &mut Self::Buffer, tokens: usize, heads: usize, head_dim: usize, pos_offset: usize, eps: f32, mode: i32, )

Fused QK-norm + RoPE + transpose-to-head-major.

mode selects the operation: 0 = transpose only (typical for V, which needs no norm and no RoPE) 1 = per-head RMS norm + RoPE + transpose (Q/K with QK-norm, Qwen3) 2 = RoPE + transpose (Q/K without QK-norm, Llama/Mistral)

input: [tokens, heads, head_dim] (token-major, output of split_qkv) output: [heads, tokens, head_dim] (head-major, ready for flash_attn / kv_cache_append)

pos_offset is the position of token 0 (decode uses current seq len; prefill uses 0). Within the batch, positions are taken as pos_offset + i.

This is the primary attention-input preparation op. Backends that have a fused kernel (Metal’s qk_norm_rope_transpose_f32) will be dramatically faster than composing norm + rope + transpose separately; the CPU fallback lowers to the individual ops.

fn kv_cache_append_head_major( ctx: &mut Self::Context, cache_k: &mut Self::Buffer, cache_v: &mut Self::Buffer, cache_len: usize, cache_capacity: usize, new_k_head_major: &Self::Buffer, new_v_head_major: &Self::Buffer, new_tokens: usize, nkv: usize, hd: usize, )

Append new K/V into a pre-allocated head-major cache buffer.

cache_k / cache_v: [nkv, capacity, hd] (head-major, pre-allocated) new_k_head_major / new_v_head_major: [nkv, new_tokens, hd] — produced directly by qk_norm_rope, no extra transpose needed.

In-place append at slot [nkv, cache_len..cache_len+new_tokens, hd]. Caller owns cache_len bookkeeping.

fn transpose_head_to_token( ctx: &mut Self::Context, src: &Self::Buffer, dst: &mut Self::Buffer, tokens: usize, heads: usize, dim: usize, )

Transpose [heads, tokens, dim] → [tokens, heads, dim]. Called after flash_attention to restore token-major layout for O-proj.

fn add_inplace( ctx: &mut Self::Context, residual: &mut Self::Buffer, x: &Self::Buffer, len: usize, )

residual[i] += x[i] (in-place)

fn add_bias( ctx: &mut Self::Context, data: &mut Self::Buffer, bias: &Self::Buffer, rows: usize, cols: usize, )

Broadcast bias add: data[r, c] += bias[c] for every row. Required by Bert / Clip / Whisper whose linear projections carry a bias.

fn layer_norm( ctx: &mut Self::Context, x: &Self::Buffer, gamma: &Self::Buffer, beta: &Self::Buffer, eps: f32, out: &mut Self::Buffer, tokens: usize, dim: usize, )

Full LayerNorm (mean + variance normalisation + affine), distinct from the rms_norm used by Llama-family decoders. out[r, c] = ((x[r, c] - mean) / sqrt(var + eps)) * gamma[c] + beta[c] Where mean and var are reduced over the last dim (cols).

fn gelu( ctx: &mut Self::Context, x: &Self::Buffer, out: &mut Self::Buffer, len: usize, )

Element-wise GELU activation (erf-based, matches PyTorch default).

fn alloc(len: usize) -> Self::Buffer

fn to_vec(buf: &Self::Buffer, len: usize) -> Vec<f32>

fn from_slice(data: &[f32]) -> Self::Buffer

Provided Methods

fn with_device_ordinal<R>( _device_ordinal: Option<usize>, body: impl FnOnce() -> R, ) -> R

Run body while binding context-free backend operations to an explicit device ordinal when the backend supports multi-device scopes.

Most backends have no per-ordinal concept and use the default no-op implementation. CUDA overrides this once its stream/context caches are device-keyed, allowing layer-split stages to load and execute on their selected GPU instead of relying on process-global defaults.

fn supports_device_ordinal_scope() -> bool

Whether Self::with_device_ordinal actually switches backend execution to the requested ordinal.

fn graph_capture_in_flight(_ctx: &Self::Context) -> bool

Whether the backend context is currently inside a graph-capture window.

Synchronizing a CUDA stream while capture is active raises CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED; diagnostic probes that time sub-ops with explicit sync boundaries must skip those boundaries while this returns true. Backends without graph capture use the default.

fn sync_before_host_readback(_ctx: &mut Self::Context)

Prepare pending GPU work for a following host readback.

Most backends either execute eagerly or synchronize as part of their device-to-host copy. Metal shared-buffer reads use the CPU pointer directly, so Metal must flush its command buffer before to_vec.

fn activation_elem_size_bytes() -> usize

Byte width of buffers returned by Self::alloc.

CUDA activation scratch is fp16, while Metal and CPU scratch are fp32. Generic model code uses this for byte offsets into batched scratch buffers without checking concrete backend types.

fn supports_llama_family_batched_decode() -> bool

Whether LlamaFamilyModel::decode_batch_internal may use its optimized batched decode path on this backend.

Backends that do not yet produce correct follow-up logits under concurrent dense decode should override this to force the per-item fallback until the optimized path is fixed.

fn supports_qk_norm_rope_batched_per_item() -> bool

Whether this backend implements the fused per-item batched Q/K/V normalization and RoPE kernel used by the Llama-family batched decode path.

fn supports_kv_cache_append_batched_per_cache() -> bool

Whether this backend implements batched KV-cache append across multiple independent per-request caches.

fn supports_flash_attention_batched_per_cache() -> bool

Whether this backend implements batched decode attention across multiple independent per-request caches.

fn zero_buffer( _ctx: &mut Self::Context, _buf: &mut Self::Buffer, _len: usize, ) -> Result<()>

Zero the first len elements of a Self::Buffer. CUDA path uses cuMemsetD16Async; default returns unsupported.

fn mla_attention( _ctx: &mut Self::Context, _q: &Self::Buffer, _kv_compressed: &Self::Buffer, _kv_rope: &Self::Buffer, _out: &mut Self::Buffer, _batch: usize, _q_len: usize, _kv_len: usize, _pos_offset: usize, _cfg: &AttnConfig, _kv_lora_rank: usize, _qk_rope_head_dim: usize, ) -> Result<()>

Multi-Head Latent Attention — DeepSeek V2 / V3’s compressed-KV attention variant. Extension point only; no backend implements it yet. DeepSeek V3 landing in Phase D/E will fill this in.

q: full Q [batch, num_heads, q_len, head_dim] kv_compressed: latent KV [batch, kv_len, kv_lora_rank] kv_rope: per-position rope-applied key heads [batch, kv_len, qk_rope_head_dim] out: [batch, num_heads, q_len, head_dim]

fn recurrent_gated_delta_rule_f32( _ctx: &mut Self::Context, _query: &Self::Buffer, _key: &Self::Buffer, _value: &Self::Buffer, _g: &Self::Buffer, _beta: &Self::Buffer, _initial_state: &Self::Buffer, _out: &mut Self::Buffer, _final_state: &mut Self::Buffer, _tokens: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _use_qk_l2norm: bool, _scale: f32, ) -> Result<()>

Recurrent gated DeltaNet update used by linear-attention layers.

Layouts are token-major:

• query / key: [tokens, key_heads, key_dim]
• value / out: [tokens, value_heads, value_dim]
• g / beta: [tokens, value_heads]
• initial_state / final_state: [value_heads, value_dim, key_dim]

Backends may require these buffers to be F32. CUDA currently provides the native W3 path; unsupported backends should use the model-level reference path instead of silently round-tripping through the host.

fn recurrent_gated_delta_rule_batch_f32( _ctx: &mut Self::Context, _query: &Self::Buffer, _key: &Self::Buffer, _value: &Self::Buffer, _g: &Self::Buffer, _beta: &Self::Buffer, _initial_states: &Self::Buffer, _out: &mut Self::Buffer, _final_states: &mut Self::Buffer, _batch: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _use_qk_l2norm: bool, _scale: f32, ) -> Result<()>

Batched one-token recurrent gated DeltaNet update.

Layouts are independent-sequence token-major:

• query / key: [batch, key_heads, key_dim]
• value / out: [batch, value_heads, value_dim]
• g / beta: [batch, value_heads]
• initial_states / final_states: [batch, value_heads, value_dim, key_dim]

This is the decode-time counterpart of Self::recurrent_gated_delta_rule_f32 for continuous batching. Each batch row has its own recurrent state; there is no temporal dependency across rows.

fn supports_qwen35_indexed_recurrent_state() -> bool

Whether this backend can update Qwen3.5 decode-time recurrent state directly from a persistent slot-indexed state slab.

fn qwen35_indexed_recurrent_state_dtype() -> Dtype

Persistent state-slab dtype supported by the fast indexed Qwen3.5 GDN kernels. Activation/cache dtype alone is not sufficient: each backend must report the dtype its indexed conv and DeltaNet kernels can update directly.

fn supports_qwen35_packed_gdn_decode_prepare() -> bool

Whether this backend can consume Qwen3.5 GDN decode projections in the vLLM-packed layout:

• in_proj_qkvz: [q, k, v, z]
• in_proj_ba: [b, a]

This avoids two small decode projection launches and lets the prepare kernel split the packed outputs while updating indexed recurrent state.

fn supports_qwen35_packed_gdn_prefill_prepare() -> bool

Whether this backend can consume Qwen3.5 GDN prefill projections in the vLLM-packed layout:

• in_proj_qkvz: [q, k, v, z]
• in_proj_ba: [b, a]

This avoids two projection launches on chunked/varlen prefill and lets the prepare kernel split the packed outputs while doing causal conv.

fn supports_qwen35_packed_gdn_recurrent_decode() -> bool

Whether this backend can keep Qwen3.5 packed GDN decode projections packed through the recurrent update, without splitting q/k/v/g/beta into intermediate buffers.

fn recurrent_gated_delta_rule_batch_indexed_f32( _ctx: &mut Self::Context, _query: &Self::Buffer, _key: &Self::Buffer, _value: &Self::Buffer, _g: &Self::Buffer, _beta: &Self::Buffer, _state_slots: &mut Self::Buffer, _slot_indices: &Self::Buffer, _out: &mut Self::Buffer, _batch: usize, _max_slots: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _use_qk_l2norm: bool, _scale: f32, ) -> Result<()>

Batched one-token recurrent gated DeltaNet update over a persistent slot-indexed state slab.

Layouts:

• query / key: [batch, key_heads, key_dim]
• value / out: [batch, value_heads, value_dim]
• g / beta: [batch, value_heads]
• state_slots: [max_slots, value_heads, value_dim, key_dim]
• slot_indices: [batch] u32 indices into state_slots

Each row reads and updates the state slot selected by slot_indices[row].

fn recurrent_gated_delta_rule_batch_indexed_packed_f32( _ctx: &mut Self::Context, _mixed_qkv: &Self::Buffer, _ba_raw: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _state_slots: &mut Self::Buffer, _slot_indices: &Self::Buffer, _out: &mut Self::Buffer, _batch: usize, _max_slots: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _scale: f32, ) -> Result<()>

Batched one-token recurrent gated DeltaNet update directly from packed decode-time q/k/v and raw b/a projections.

Layouts:

• mixed_qkv: [batch, q, k, v] where q/k are [key_heads, key_dim] and v is [value_heads, value_dim]
• ba_raw: [batch, b, a] with each half [value_heads]
• state_slots: [max_slots, value_heads, value_dim, key_dim]
• slot_indices: [batch] u32 indices into state_slots

fn recurrent_gated_delta_rule_varlen_f32( _ctx: &mut Self::Context, _query: &Self::Buffer, _key: &Self::Buffer, _value: &Self::Buffer, _g: &Self::Buffer, _beta: &Self::Buffer, _initial_states: &Self::Buffer, _cu_seqlens: &Self::Buffer, _out: &mut Self::Buffer, _final_states: &mut Self::Buffer, _batch: usize, _total_tokens: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _use_qk_l2norm: bool, _scale: f32, ) -> Result<()>

Variable-length batched recurrent gated DeltaNet prefill update.

Layouts are token-major over all concatenated prefill chunks:

• query / key: [total_tokens, key_heads, key_dim]
• value / out: [total_tokens, value_heads, value_dim]
• g / beta: [total_tokens, value_heads]
• cu_seqlens: [batch + 1] u32 prefix sum into the flat token axis
• initial_states / final_states: [batch, value_heads, value_dim, key_dim]

Each sequence advances independently from its own initial recurrent state and writes one final state. This is the prefill counterpart of Self::recurrent_gated_delta_rule_batch_f32 and matches the cu_seqlens shape used by vLLM-style chunked GDN prefill.

fn linear_attention_prepare_f32( _ctx: &mut Self::Context, _mixed_qkv_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _a_raw: &Self::Buffer, _b_raw: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _query: &mut Self::Buffer, _key: &mut Self::Buffer, _value: &mut Self::Buffer, _g: &mut Self::Buffer, _beta: &mut Self::Buffer, _tokens: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, _apply_qk_l2norm: bool, ) -> Result<()>

Prepare a gated-Delta linear-attention block: depthwise causal conv + SiLU over mixed_qkv_raw, split into Q/K/V, and compute GDN gates g and beta.

fn linear_attention_prepare_varlen_f32( _ctx: &mut Self::Context, _mixed_qkv_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _initial_conv_states: &Self::Buffer, _a_raw: &Self::Buffer, _b_raw: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _cu_seqlens: &Self::Buffer, _token_seq_indices: &Self::Buffer, _query: &mut Self::Buffer, _key: &mut Self::Buffer, _value: &mut Self::Buffer, _g: &mut Self::Buffer, _beta: &mut Self::Buffer, _final_conv_states: &mut Self::Buffer, _batch: usize, _total_tokens: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, _apply_qk_l2norm: bool, ) -> Result<()>

Varlen prefill-time gated-Delta linear-attention preparation.

This is the batched/stateful counterpart of Self::linear_attention_prepare_f32:

• mixed_qkv_raw: [total_tokens, conv_channels]
• a_raw / b_raw / g / beta: [total_tokens, value_heads]
• query / key: [total_tokens, key_heads, key_dim]
• value: [total_tokens, value_heads, value_dim]
• cu_seqlens: [batch + 1] u32 prefix sum into the flat token axis
• token_seq_indices: [total_tokens] u32 sequence row per flat token
• initial_conv_states / final_conv_states: [batch, conv_channels, conv_kernel - 1]

Each sequence’s depthwise causal conv reads only that sequence plus its own initial conv state and writes one final conv state. That boundary handling is required before a varlen recurrent GDN pass can be used for product prefill batching.

fn linear_attention_prepare_varlen_packed_qkvz_ba_f32( _ctx: &mut Self::Context, _mixed_qkvz_raw: &Self::Buffer, _ba_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _initial_conv_states: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _cu_seqlens: &Self::Buffer, _token_seq_indices: &Self::Buffer, _query: &mut Self::Buffer, _key: &mut Self::Buffer, _value: &mut Self::Buffer, _z: &mut Self::Buffer, _g: &mut Self::Buffer, _beta: &mut Self::Buffer, _final_conv_states: &mut Self::Buffer, _batch: usize, _total_tokens: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, _apply_qk_l2norm: bool, ) -> Result<()>

Varlen prefill-time gated-Delta linear-attention preparation from vLLM-packed Qwen3.5 projections.

Layouts:

• mixed_qkvz_raw: [total_tokens, q, k, v, z]
• ba_raw: [total_tokens, b, a]
• z: [total_tokens, value_heads, value_dim] Other outputs and state layouts match Self::linear_attention_prepare_varlen_f32.

fn linear_attention_decode_prepare_f32( _ctx: &mut Self::Context, _mixed_qkv_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _conv_state: &Self::Buffer, _a_raw: &Self::Buffer, _b_raw: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _query: &mut Self::Buffer, _key: &mut Self::Buffer, _value: &mut Self::Buffer, _g: &mut Self::Buffer, _beta: &mut Self::Buffer, _next_conv_state: &mut Self::Buffer, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, _apply_qk_l2norm: bool, ) -> Result<()>

Decode-time gated-Delta linear-attention preparation for one token.

This is the stateful counterpart of Self::linear_attention_prepare_f32: it reads [conv_channels, conv_kernel - 1] causal-conv state, appends the current raw QKV token, writes the next conv state, then emits Q/K/V and GDN gates for the current token. The layout mirrors vLLM’s Qwen GDN conv_state + temporal-state split.

fn linear_attention_decode_prepare_batch_f32( _ctx: &mut Self::Context, _mixed_qkv_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _conv_states: &Self::Buffer, _a_raw: &Self::Buffer, _b_raw: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _query: &mut Self::Buffer, _key: &mut Self::Buffer, _value: &mut Self::Buffer, _g: &mut Self::Buffer, _beta: &mut Self::Buffer, _next_conv_states: &mut Self::Buffer, _batch: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, _apply_qk_l2norm: bool, ) -> Result<()>

Batched stateful one-token linear-attention preparation.

This processes batch independent decode rows:

• mixed_qkv_raw: [batch, conv_channels]
• conv_states / next_conv_states: [batch, conv_channels, conv_kernel - 1]
• a_raw / b_raw / g / beta: [batch, value_heads]
• query / key: [batch, key_heads, key_dim]
• value: [batch, value_heads, value_dim]

fn linear_attention_decode_prepare_batch_indexed_f32( _ctx: &mut Self::Context, _mixed_qkv_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _conv_state_slots: &mut Self::Buffer, _slot_indices: &Self::Buffer, _a_raw: &Self::Buffer, _b_raw: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _query: &mut Self::Buffer, _key: &mut Self::Buffer, _value: &mut Self::Buffer, _g: &mut Self::Buffer, _beta: &mut Self::Buffer, _batch: usize, _max_slots: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, _apply_qk_l2norm: bool, ) -> Result<()>

Batched stateful one-token linear-attention preparation over a persistent slot-indexed conv-state slab.

conv_state_slots has layout [max_slots, conv_channels, conv_kernel-1] and is updated in place at slot_indices[row].

fn linear_attention_decode_prepare_batch_indexed_packed_qkvz_ba_f32( _ctx: &mut Self::Context, _mixed_qkvz_raw: &Self::Buffer, _ba_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _conv_state_slots: &mut Self::Buffer, _slot_indices: &Self::Buffer, _a_log: &Self::Buffer, _dt_bias: &Self::Buffer, _query: &mut Self::Buffer, _key: &mut Self::Buffer, _value: &mut Self::Buffer, _z: &mut Self::Buffer, _g: &mut Self::Buffer, _beta: &mut Self::Buffer, _batch: usize, _max_slots: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, _apply_qk_l2norm: bool, ) -> Result<()>

fn linear_attention_decode_prepare_batch_indexed_packed_qkvz_to_mixed_f32( _ctx: &mut Self::Context, _mixed_qkvz_raw: &Self::Buffer, _conv_weight: &Self::Buffer, _conv_state_slots: &mut Self::Buffer, _slot_indices: &Self::Buffer, _mixed_qkv: &mut Self::Buffer, _z: &mut Self::Buffer, _batch: usize, _max_slots: usize, _key_heads: usize, _value_heads: usize, _key_dim: usize, _value_dim: usize, _conv_kernel: usize, ) -> Result<()>

fn gated_rms_norm_f32( _ctx: &mut Self::Context, _core: &Self::Buffer, _z: &Self::Buffer, _weight: &Self::Buffer, _out: &mut Self::Buffer, _tokens: usize, _heads: usize, _dim: usize, _eps: f32, ) -> Result<()>

Gated RMSNorm used after recurrent DeltaNet core: out = rms_norm(core) * weight * silu(z).

fn embedding_lookup_dev( ctx: &mut Self::Context, table: &Self::Buffer, ids: &Self::Buffer, out: &mut Self::Buffer, batch: usize, dim: usize, )

Device-buffer variant of embedding_lookup for graph-capturable MoE routing — the gather step before phase-1 GEMM in moe_forward_bucketed. The host-slice embedding_lookup does clone_htod(ids) internally, which records stale host pointers under CUDA Graph capture replay.

ids: &Self::Buffer must be a device I32 buffer of batch elements (e.g. Qwen3MoeScratch::route_packed_idx_dev). batch is passed explicitly since a typed CudaBuf carries its element count but the caller often wants a partial gather.

Default impl: round-trip via to_vec + dispatch the host-slice variant. CUDA overrides.

fn fused_gelu_tanh_mul_split( _ctx: &mut Self::Context, _gate_up: &Self::Buffer, _out: &mut Self::Buffer, _tokens: usize, _im: usize, )

GeGLU variant of Backend::fused_silu_mul_split: gelu_tanh(gate) * up → out. Matches HF gelu_pytorch_tanh (Gemma family MLP). Panics by default: wire a real kernel on any backend that loads a GeGLU model.

fn scale_inplace( ctx: &mut Self::Context, buf: &mut Self::Buffer, scale: f32, len: usize, )

buf[i] *= scale over the first len elements. Gemma-family embedding scaling (×√hidden_size on residual-stream entry). Default round-trips through host memory — correct but slow; override on backends that serve Gemma models.

fn qk_norm_rope_partial( ctx: &mut Self::Context, input: &Self::Buffer, norm_w: &Self::Buffer, cos: &Self::Buffer, sin: &Self::Buffer, output: &mut Self::Buffer, tokens: usize, heads: usize, head_dim: usize, rope_dim: usize, input_stride: usize, input_offset: usize, input_head_stride: usize, pos_offset: usize, eps: f32, mode: i32, ) -> Result<()>

Q/K preparation variant for Qwen3.5 full attention.

input_stride is the per-token feature width in input, and input_offset is the first feature for this projection inside each token row. This lets Q read the first num_heads * head_dim slice from Qwen3.5’s gated q_proj output while leaving the attention gate slice in place for a later device-side post-op.

rope_dim may be smaller than head_dim; dimensions outside rope_dim are normalized and copied but not rotated. mode follows Backend::qk_norm_rope: 0 transpose only, 1 RMSNorm+RoPE, 2 RoPE only, 3 RMSNorm+interleaved RoPE for Qwen3.5’s mrope layout.

fn qwen35_apply_attention_gate( _ctx: &mut Self::Context, _context: &mut Self::Buffer, _query_raw: &Self::Buffer, _tokens: usize, _q_total: usize, _q_proj_total: usize, _head_dim: usize, ) -> Result<()>

Apply Qwen3.5 attention output gate in place. Gated Qwen3.5 full attention stores q_proj rows as per-head [query, gate] slices, so context[token, head, dim] *= sigmoid(q_proj[token, head, head_dim + dim]).

The CUDA implementation is a single device kernel. Backends that do not implement it must fail rather than silently copying data through host memory on product paths.

fn qwen35_apply_token_gate( _ctx: &mut Self::Context, _values: &mut Self::Buffer, _gate: &Self::Buffer, _tokens: usize, _hidden_size: usize, ) -> Result<()>

Apply one scalar gate per token to a token-major hidden buffer: values[token, dim] *= sigmoid(gate[token]).

fn qwen35_apply_token_gate_and_add_inplace( ctx: &mut Self::Context, dst: &mut Self::Buffer, values: &mut Self::Buffer, gate: &Self::Buffer, tokens: usize, hidden_size: usize, ) -> Result<()>

Apply one scalar gate per token to values, then add the gated values into dst in place: values[token, dim] *= sigmoid(gate[token]); dst[token, dim] += values[token, dim].

The default preserves the old two-dispatch behavior. CUDA overrides this for Qwen3.5 sparse-MoE shared-expert decode, where it removes the separate token-gate and merge launches while keeping the gated values buffer available for tracing/debugging.

fn qwen35_interleave_gate_up( ctx: &mut Self::Context, gate: &Self::Buffer, up: &Self::Buffer, out: &mut Self::Buffer, tokens: usize, intermediate: usize, ) -> Result<()>

Interleave shared-expert gate/up projections from token-major [tokens, intermediate] buffers into [tokens, 2 * intermediate].

The default keeps backend behavior unchanged. CUDA overrides this to avoid 2 * tokens tiny device-to-device copies per MoE layer.

fn kv_cache_append_batched_per_cache( _ctx: &mut Self::Context, _caches: &[&Self::Buffer], _new_data: &Self::Buffer, _cache_lens: &Self::Buffer, _capacity: usize, _m: usize, _nkv: usize, _hd: usize, _slot: usize, ) -> Result<()>

Batched kv_cache_append across M caches in one launch. Each item writes its (head-major) K-or-V row into its own cache at offset read from cache_lens[i]. Replaces M sequential kv_cache_append_head_major calls with a single dispatch.

new_data layout: [m, nkv, hd] item-major (each item’s slice is contiguous, identical to the k/v_normed_batched produced by qk_norm_rope_batched_per_item). caches: per-cache [nkv, capacity, hd] head-major. cache_lens: device buffer (u32 storage, length ≥ m). Caller fills via B::write_u32_into BEFORE the call. Required for CUDA-graph capture: the kernel reads from this stable device buffer, so a captured graph can be replayed with new lens by just rewriting the buffer between launches.

fn flash_attention_batched_per_cache( _ctx: &mut Self::Context, _q: &Self::Buffer, _k_caches: &[&Self::Buffer], _v_caches: &[&Self::Buffer], _kv_lens: &Self::Buffer, _out: &mut Self::Buffer, _nq: usize, _nkv: usize, _hd: usize, _scale: f32, _max_valid_kv: usize, _capacity: usize, _sliding_window: usize, _slot: usize, ) -> Result<()>

Batched flash_attention across M decode caches in one launch. Replaces the per-item flash_attention(q_len=1, ...) × M loop in the non-paged batched-decode path.

API takes Vec<&Buffer> for the per-cache K/V buffers (each [nkv, capacity, hd] head-major) plus host-side kv_lens. Backends that implement it must extract per-cache device pointers, build the device arrays the kernel needs, and launch one kernel covering all M items.

q layout: [m, nq, hd] item-major (matches the qk_norm_rope_batched_per_item output for q_len=1). out layout: [m, nq, hd] item-major — written directly into the caller’s batched attn_out buffer, no per-item copy needed.

CUDA-only for now (kernel batched_decode_attention exists in kernels/batched_decode_attention.cu). kv_lens: device buffer (u32 storage, length ≥ m) — same design as kv_cache_append_batched_per_cache::cache_lens. sliding_window: common decode window for every item; 0 means full causal attention, w > 0 means each item attends only to the last w valid KV positions.

fn qk_norm_rope_batched_per_item( _ctx: &mut Self::Context, _input: &Self::Buffer, _norm_w: &Self::Buffer, _cos: &Self::Buffer, _sin: &Self::Buffer, _output: &mut Self::Buffer, _positions: &Self::Buffer, _m: usize, _heads: usize, _head_dim: usize, _eps: f32, _mode: i32, ) -> Result<()>

Batched per-item-position variant of qk_norm_rope for the non-paged batched-decode path. Each of the m items has its own absolute RoPE position (read from a device i32 buffer of length m). Layout is item-major in both input and output:

input [m, heads, head_dim] output [m, heads, head_dim] (no head-major transpose)

Item-major output keeps the per-item flash_attention slice contiguous (output[i * heads * head_dim ..] is item i’s whole Q tensor in head-major-equivalent layout for q_len=1).

Replaces the M sequential single-item launches in the existing forward_layer_batched_decode path with one batched dispatch. CUDA-only for now; other backends fall through to the default unsupported and the caller falls back to the per-item loop.

fn split_qkv_norm_rope( _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, _k_out: &mut Self::Buffer, _v_out: &mut Self::Buffer, _tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _pos_offset: usize, _eps: f32, _qk_mode: i32, ) -> Result<()>

Fused split-QKV + QK-norm + RoPE + head-major transpose.

Single-dispatch replacement for the (split_qkv → 3× qk_norm_rope) chain on the decode-attention prelude. Reads the linear-layer fused-QKV output once and writes head-major Q/K/V directly into attention scratch.

qkv layout: [tokens, q_heads*hd + 2*kv_heads*hd]. q_out: [q_heads, tokens, hd]. k_out/v_out: [kv_heads, tokens, hd]. qk_mode: 1 = norm + half-split RoPE for Q/K (Qwen3 with QK-norm), 2 = half-split RoPE only for Q/K, 3 = interleaved RoPE only for Q/K (GGUF LLaMA / llama.cpp layout). V always falls through to transpose-only.

Default returns Unsupported. Backends that implement it are expected to be dramatically faster than the four-dispatch chain.

fn split_qkv_norm_rope_into_cache( _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, _tokens: usize, _q_heads: usize, _kv_heads: usize, _head_dim: usize, _pos_offset: usize, _eps: f32, _qk_mode: i32, _cache_len: usize, _cache_capacity: usize, ) -> Result<()>

Variant of Backend::split_qkv_norm_rope that writes the new K and V directly into pre-allocated head-major KV cache buffers at slot [kv_heads, cache_len .. cache_len + tokens, hd]. Eliminates the trailing kv_cache_append_head_major dispatch on the decode hot path. Q still lands in per-token head-major scratch (flash-attention reads it as the query).

Default returns Unsupported. Backends without the fused kernel can keep using split_qkv_norm_rope + kv_cache_append_head_major.

fn transpose_token_to_head( _ctx: &mut Self::Context, _src: &Self::Buffer, _dst: &mut Self::Buffer, _tokens: usize, _heads: usize, _dim: usize, )

Inverse of transpose_head_to_token: [tokens, heads, dim] → [heads, tokens, dim]. Used by the CUDA paged_decode_attention wrapper to convert paged_varlen_attention’s token-major output back to the head-major layout that Qwen3MoeModel expects. Default panics — backends without a paged-KV CUDA path don’t hit this code.

fn scaled_add_inplace( _ctx: &mut Self::Context, dst: &mut Self::Buffer, src: &Self::Buffer, scale: f32, len: usize, )

dst[i] += scale * src[i] — scalar-broadcast scaled add, in place.

MoE per-token combine writes out[b] += weight_k * expert_k(x[b]) for each top-K expert; this primitive is the per-call accumulate. Backends without a dedicated kernel can fall back to the default implementation, which round-trips through host memory — correct, but slow on a hot path. Override on any backend you actually dispatch MoE on.

fn fused_silu_mul_split_strided( _ctx: &mut Self::Context, _gate_up: &Self::Buffer, _in_row_offset: usize, _out: &mut Self::Buffer, _out_row_offset: usize, _tokens: usize, _intermediate: usize, )

Strided variant of Backend::fused_silu_mul_split for the bucketed MoE path: reads gate_up rows starting at in_row_offset, writes out rows starting at out_row_offset.

fn write_f32_to_activation( ctx: &mut Self::Context, dst: &mut Self::Buffer, data: &[f32], )

fn f32_to_activation( ctx: &mut Self::Context, input_f32: &Self::Buffer, out: &mut Self::Buffer, len: usize, )

Convert a typed F32 device buffer into the backend activation dtype.

CUDA activations are FP16 for tensor-core/Marlin kernels, while the Qwen3.5 gated-Delta core keeps recurrent math in F32. Backends with non-F32 activations should override this with a device-side conversion.

fn supports_device_f32_residual_shadow() -> bool

Whether this backend can keep Gemma-style sandwich residuals in a device-side F32 shadow while continuing to feed FP16 activations into projection kernels. The default is false so existing CPU/Metal paths keep their current host-side fallback behavior.

fn activation_to_f32_shadow( ctx: &mut Self::Context, src: &Self::Buffer, dst_f32: &mut Self::Buffer, len: usize, )

Copy an activation buffer into a typed F32 shadow buffer.

fn activation_add_to_f32_shadow( ctx: &mut Self::Context, src: &Self::Buffer, residual_f32: &mut Self::Buffer, scratch_f32: &mut Self::Buffer, len: usize, )

Add an activation buffer directly into an existing F32 residual shadow.

scratch_f32 is provided for portable fallback implementations. CUDA can fuse the activation-to-F32 conversion and residual add into one kernel.

fn rms_norm_activation_to_f32( ctx: &mut Self::Context, input: &Self::Buffer, weight: &Self::Buffer, eps: f32, out_f32: &mut Self::Buffer, tokens: usize, dim: usize, )

RMSNorm an activation buffer and write the result into a typed F32 scratch buffer. Used for Gemma post-attn/post-ffn branch norms.

fn rms_norm_activation_add_to_f32( ctx: &mut Self::Context, input: &Self::Buffer, weight: &Self::Buffer, eps: f32, residual_f32: &mut Self::Buffer, scratch_f32: &mut Self::Buffer, tokens: usize, dim: usize, )

RMSNorm an activation buffer and add the F32 result directly into an existing F32 residual shadow. scratch_f32 is provided for backend fallbacks that need to materialize the normalized branch.

fn rms_norm_f32_to_activation( ctx: &mut Self::Context, input_f32: &Self::Buffer, weight: &Self::Buffer, eps: f32, out: &mut Self::Buffer, tokens: usize, dim: usize, )

RMSNorm a typed F32 shadow buffer and write the normalized result back to the backend’s regular activation dtype.

fn argmax_rows_f16( _ctx: &mut Self::Context, logits: &Self::Buffer, m: usize, n: usize, ) -> Result<Vec<u32>>

Greedy-decode fast path: GPU argmax over each row of a [m, n] FP16 logits buffer, returning the m token indices on the host. Saves m × n × 2 bytes of D2H per call (e.g. 19.5 MB at c=32, vocab=152064) and the host-side argmax scan (~150 µs × m).

Default impl falls back to the slow path: full to_vec + host argmax. CUDA overrides with a native kernel + tiny D2H (m × 4 B). Backends that don’t override pay the same cost as to_vec + host argmax, so callers can call this unconditionally.

fn argmax_rows_f16_masked( _ctx: &mut Self::Context, _logits: &Self::Buffer, _valid_token_mask: &Self::Buffer, _mask_len: usize, _m: usize, _n: usize, ) -> Result<Vec<u32>>

fn supports_argmax_rows_f16_sparse_repetition_penalty() -> bool

Whether this backend can apply a per-row sparse repetition penalty and select the greedy token without reading full logits back to the host.

fn argmax_rows_f16_sparse_repetition_penalty( _ctx: &mut Self::Context, _logits: &mut Self::Buffer, _valid_token_mask: Option<(&Self::Buffer, usize)>, _row_offsets: &Self::Buffer, _token_ids: &Self::Buffer, _repetition_penalties: &Self::Buffer, _total_token_ids: usize, _m: usize, _n: usize, ) -> Result<Vec<u32>>

Greedy-decode fast path with sparse repetition penalty.

row_offsets has length m + 1 and indexes into token_ids; row r owns token_ids[row_offsets[r]..row_offsets[r + 1]]. Backends should apply each row’s repetition_penalties[r] to those logits in-place, then run raw or masked argmax and return one token id per row.

fn from_weight_bytes(raw: &[u8], src_dtype: SrcDtype) -> Self::Buffer

Load a weight tensor straight from its on-disk byte representation, letting the backend pick its preferred storage dtype.

Default impl upcasts bf16/f16 to f32 via an intermediate Vec, matching pre-existing loader behaviour. Backends override this to go straight from raw bytes into a native half-precision buffer (e.g. Metal with FERRUM_METAL_DTYPE=f16), avoiding the transient 2× RAM spike.

Dyn Compatibility

This trait is not dyn compatible.

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

Implementors

impl Backend for CpuBackend

type Buffer = Vec<f32>

type Context = ()

type Timer = CpuTimer