Module qwen3_5

Module qwen3_5 

Source
Expand description

Vendor-agnostic Qwen3.5 per-layer decoder forward.

Extracted verbatim from the original Metal end-to-end driver (crates/spark-runtime/examples/metal_qwen35_inference/). The two exported functions — forward_full_attention and forward_linear_attention — drive a single decoder layer end to end (norm → projections → attention/GDN → residual+post-norm → MLP → residual). Any end-to-end inference example calls these through &dyn GpuBackend + a QuantWeights impl, regardless of which hardware target the backend speaks.

What the module does not do (intentional):

  • Multi-token prefill / batched dispatch — single-token decode only (KV-append at cache_pos, attention at seq_len_attn = cache_pos+1).
  • CUDA-graph capture, NCCL, paged-KV — the production decode path (crate::model::trait_impl::decode_a) handles those; this is the simpler shape an example or smoke driver wants.
  • Tokenizer / sampler / weight loading — the caller owns these.

Performance: the fused kernels (gemv_silu_gate_resid, gemv_gate_up_with, add_rms_norm) all dispatch through trait methods that backends override with their fused launches. Atlas’s Metal backend keeps decode at ~20 tok/s through this path identically to the inlined version it replaces.

Structs§

FullAttentionLayer
Full-attention layer weights, parameterised over the backend’s quantised weight type.
FullAttentionScratch
Per-call scratch buffers for the full-attention forward.
LayerKvCache
Per-layer KV cache for a full-attention layer (single-batch).
LinearAttentionLayer
Linear-attention (GDN) layer weights.
LinearAttentionScratch
Per-call scratch buffers for the linear-attention forward.
LinearAttentionState
Per-layer SSM/conv state for a linear-attention layer. Persists across tokens. Caller owns alloc + zero-init.
Qwen35ForwardConfig
Compile-time-fixed dimensions for a Qwen3.5 checkpoint. Populate from the model’s config.json (text_config) at startup.
Qwen35Kernels
Pre-resolved kernel handles. Resolve once at startup; pass & to every per-layer call so name-lookup overhead doesn’t appear in the hot path.

Enums§

MetalKvDtype
KV cache storage format for the Metal contiguous cache.

Functions§

forward_full_attention
Single-token full-attention decoder forward. Returns the DevicePtr containing the layer’s output residual stream (caller-owned scratch.x_out).
forward_linear_attention
Single-token GDN (linear-attention) decoder forward. Returns the DevicePtr containing the layer’s output residual stream — that pointer is x_buf, into which scratch.x_final was copied so the caller’s residual-stream buffer stays stable across layers.