Upload specprefill.tex with huggingface_hub

specprefill.tex

@@ -0,0 +1,877 @@
+\documentclass[11pt,a4paper]{article}
+
+% === Packages ===
+\usepackage[utf8]{inputenc}
+\usepackage[T1]{fontenc}
+\usepackage{amsmath,amssymb}
+\usepackage{graphicx}
+\usepackage{booktabs}
+\usepackage{hyperref}
+\usepackage{xcolor}
+\usepackage{algorithm}
+\usepackage{algpseudocode}
+\usepackage{listings}
+\usepackage[margin=1in]{geometry}
+\usepackage{caption}
+\usepackage{subcaption}
+\usepackage{natbib}
+
+% === Macros ===
+\newcommand{\tbd}[1]{\textcolor{red}{\textbf{[TBD: #1]}}}
+\newcommand{\specprefill}{\textsc{SpecPrefill}}
+\newcommand{\ttft}{\text{TTFT}}
+
+\lstset{
+  basicstyle=\ttfamily\small,
+  keywordstyle=\color{blue},
+  commentstyle=\color{gray},
+  breaklines=true,
+  frame=single,
+  numbers=left,
+  numberstyle=\tiny\color{gray},
+}
+
+% === Title ===
+\title{\specprefill{} on Unified Memory:\\Cross-Architecture Sparse Prefill for\\Large Language Models on Apple Silicon}
+
+\author{
+  \texttt{github.com/Thump604}
+}
+
+\date{March 2026}
+
+\begin{document}
+\maketitle
+
+% ============================================================
+\begin{abstract}
+% ============================================================
+
+Long-context prefill is the dominant latency bottleneck for local LLM inference: a 64K-token prompt on Qwen3.5-122B (MoE, 10B active parameters) takes 7 minutes before the first token appears.
+\specprefill{}---attention-based sparse prefill using a draft model---was designed for GPU clusters with discrete memory.
+We port it to Apple Silicon's unified memory architecture and generalize it across three model families: transformer Mixture-of-Experts (Qwen3.5), Mamba-2/attention hybrid (Nemotron-H), and sliding-window dense (GPT-OSS\footnote{GPT-OSS refers to a publicly available model by its open-source project designation.}).
+
+On M2~Ultra (128\,GB unified memory), \specprefill{} with a 2B draft model (Qwen3.5-2B, 1.4\,GB, 4-bit) reduces TTFT by $3.71$--$5.45\times$ across 8K--128K tokens on Qwen3.5-122B, cutting 128K prefill from 19.3~minutes to 3.5~minutes.
+Composed with system prompt KV caching, end-to-end speedup reaches $5.6\times$ on a 73K-token production workload.
+We also achieve $2.10$--$2.19\times$ on Nemotron-H~120B across 8K--64K tokens.
+Unified memory eliminates PCIe transfer overhead, making the draft-to-target FLOP ratio the dominant predictor of speedup. We formalize and validate this relationship across six draft/target configurations.
+Under adversarial evaluation (0/16 regressions at 20\% keep), LLM-judge, and perplexity analysis, we observe no quality degradation at our recommended operating point.
+
+Our implementation handles architecture-specific challenges including gated queries with per-head normalization (Qwen3.5), SSM-interleaved attention layers without positional encoding (Nemotron-H), and sliding-window cache preservation (GPT-OSS), deployed in a production serving stack with per-request API control.
+
+\end{abstract}
+
+% ============================================================
+\section{Introduction}
+\label{sec:intro}
+% ============================================================
+
+\subsection{The TTFT Problem in Local Inference}
+
+Time-to-first-token (TTFT) is the dominant user-facing latency for large language models serving long-context requests.
+On commodity hardware---a Mac Studio with Apple M2~Ultra and 128\,GB unified memory---prefilling a 64K-token prompt through Qwen3.5-122B-A10B (a 122-billion parameter Mixture-of-Experts model with 10 billion active parameters per token) requires \textbf{418~seconds}, nearly 7~minutes before the first output token appears.
+Even at 16K tokens, the wait is 92~seconds.
+
+This latency is not a bandwidth problem.
+MLX prefill on Apple Silicon is FLOP-limited: the Metal GPU is compute-bound processing each token through the model's forward pass~\cite{mlx}.
+Reducing the number of tokens processed during prefill therefore yields near-linear TTFT improvement.
+
+In local inference, the cost of long TTFT is not measured in dollars-per-token but in user time.
+A 16K-token context (typical for an IDE coding assistant with tool definitions and file contents) means 92 seconds of waiting before the first response token.
+A long creative writing session or research conversation that accumulates 64K tokens of history means 7 minutes per response.
+On a serialized single-request engine, every second of prefill also delays all queued requests.
+
+\noindent With \specprefill{}, 128K prefill on the 122B model drops from 19.5~minutes to 3.5~minutes, making long-context requests practical for interactive use.
+
+\subsection{Why Unified Memory Changes the Calculus}
+
+\specprefill{}~\cite{specprefill} addresses TTFT by using a small draft model to identify which prompt tokens are most important via attention scoring, then sparse-prefilling only the selected subset into the target model.
+The original formulation assumes a discrete-memory GPU architecture where the draft model either (a) shares GPU VRAM with the target, reducing KV cache headroom, or (b) runs on CPU or a separate GPU, incurring PCIe transfer latency for importance scores.
+
+On Apple Silicon's unified memory architecture, neither penalty applies.
+Draft and target models share the same physical address space---the draft's weights (${\sim}$1.4\,GB for a 4-bit 2B model) are simply additional allocations in the same memory pool as the target's \raise.17ex\hbox{$\scriptstyle\sim$}79\,GB.
+Scoring requires zero data movement.
+On discrete GPU systems, draft scoring would either compete for GPU VRAM with the target's KV cache (reducing effective context length) or require CPU$\leftrightarrow$GPU transfers with latency proportional to prompt length---neither penalty exists on unified memory.
+
+This simplifies the cost equation.
+Let $C_t$ denote the target model's prefill cost (in FLOPs), $C_d$ the draft model's scoring cost (full prefill plus lookahead steps), and $k$ the fraction of tokens retained.
+On unified memory:
+\begin{equation}
+\label{eq:speedup}
+\text{Speedup} = \frac{C_t}{C_d + k \cdot C_t}
+\end{equation}
+When $C_d \ll C_t$---as with a 2B draft scoring for a 122B MoE target, where the FLOP ratio is approximately $50\times$---speedup approaches $1/k$.
+At $k = 0.2$, this predicts up to $4.5\times$; we measure $4.11\times$ at 16K tokens (5 trials), with the gap attributable to overhead from chunk selection, RoPE patching, and memory management.
+
+We term this the \textbf{ratio thesis}: on unified memory where $T = 0$, the draft-to-target FLOP ratio $r$ is the dominant predictor of \specprefill{} benefit, modulated by architecture-dependent overhead $\epsilon$ (Equation~\ref{eq:speedup_ratio}).
+Section~\ref{sec:ratio} validates this across six draft/target configurations spanning an $8.5\times$ range of FLOP ratios.
+
+\subsection{Contributions}
+
+\begin{enumerate}
+  \item \textbf{First implementation on unified memory hardware.}
+    We implement \specprefill{} on Apple Silicon via the MLX framework, demonstrating that zero-copy scoring shifts the viability threshold, making the technique effective even at moderate prompt lengths (8K tokens).
+
+  \item \textbf{Cross-architecture generalization.}
+    We extend \specprefill{} beyond standard transformers to Mamba-2/attention hybrids (Nemotron-H, where only 8 of 88 target layers have attention) and sliding-window models (GPT-OSS, with RotatingKVCache).
+    Auto-detecting query extraction handles gated attention with per-head normalization (Qwen3.5), content-based attention without positional encoding (Nemotron-H), and YarnRoPE with sliding-window cache preservation (GPT-OSS).
+
+  \item \textbf{Production system integration.}
+    We show per-request API control with graceful fallback, coexistence with Multi-Token Prediction (MTP) speculative decoding in a three-phase ``Speculative Stack,'' and demonstrate that the FLOP ratio thesis extends to draft model selection---smaller drafts with lower $r$ yield higher speedup.
+\end{enumerate}
+
+This is a \textbf{systems paper with algorithmic adaptations}, not a claim of a new algorithm.
+The core \specprefill{} idea is due to~\citet{specprefill}; our contributions are making it work across architectures on new hardware with real deployment.
+
+
+% ============================================================
+\section{Background}
+\label{sec:background}
+% ============================================================
+
+\subsection{SpecPrefill}
+
+\citet{specprefill} observe that during autoregressive generation, the model attends heavily to a small fraction of prompt tokens.
+If these important tokens can be identified cheaply, the full prompt need not be prefilled into the target model.
+Their method uses a small draft model to score token importance via attention weights, selects the top-$k$\% of tokens (in non-overlapping chunks for spatial locality), and sparse-prefills only the selected subset.
+
+Sparse prefill is possible because Rotary Position Embeddings (RoPE)~\cite{rope} encode \emph{relative} position.
+The inner product $Q_m \cdot K_p^T$ depends only on the relative distance $(m - p)$ encoded via rotation angles.
+If selected tokens are stored in the KV cache with their \emph{original} position angles, they interact with future decode queries at correct relative distances regardless of gaps.
+
+\subsection{MLX and Apple Silicon Unified Memory}
+
+MLX~\cite{mlx} is Apple's machine learning framework designed for Apple Silicon.
+Apple Silicon uses \emph{unified memory}: CPU and GPU share the same physical DRAM through a common memory controller.
+There is no PCIe bus, no \texttt{cudaMemcpy}, no distinct VRAM allocation.
+Tensors created by any compute unit are immediately accessible to any other.
+
+MLX exploits this with lazy evaluation and reference-counted memory management.
+Metal compute shaders execute matrix operations on the GPU.
+In practice, a draft model's weights are additional allocations in the same memory pool, accessible at full bandwidth from any compute unit without copying.
+
+\subsection{Target Architectures}
+
+We evaluate \specprefill{} on three architecturally distinct model families, establishing that generalization requires non-trivial adaptations:
+
+\begin{table}[h]
+\centering
+\small
+\begin{tabular}{@{}lllll@{}}
+\toprule
+\textbf{Model} & \textbf{Architecture} & \textbf{Attention} & \textbf{Position Enc.} & \textbf{Cache} \\
+\midrule
+Qwen3.5-122B & MoE (10B active) & Gated + q\_norm & Standard RoPE & Standard KV \\
+Nemotron-H 120B & Mamba-2 + Attn + MoE & Standard (8/88 layers) & None (SSM) & Compacted \\
+GPT-OSS 120B & Dense + sliding window & Standard & YarnRoPE & RotatingKV \\
+\bottomrule
+\end{tabular}
+\caption{Target model architectures. Each requires different query extraction, position handling, and cache management in \specprefill{}.}
+\label{tab:architectures}
+\end{table}
+
+\textbf{Qwen3.5} uses gated attention where \texttt{q\_proj} outputs $2\times$ the expected width (query concatenated with a gate), requiring a split before per-head RMSNorm (\texttt{q\_norm}) and RoPE application.
+
+\textbf{Nemotron-H} is a hybrid architecture with 40 Mamba-2 SSM layers, 8 full-attention layers, and 40 MoE feed-forward layers.
+Positional information is encoded entirely in the SSM state---the attention layers have no RoPE.
+Only the attention layers produce Q/K scores usable for importance scoring.
+
+\textbf{GPT-OSS} uses YarnRoPE~\cite{yarn} with a sliding-window attention pattern where alternating layers use \texttt{RotatingKVCache} retaining only the last 128 tokens.
+
+
+% ============================================================
+\section{Method}
+\label{sec:method}
+% ============================================================
+
+\subsection{Token Importance Scoring}
+\label{sec:scoring}
+
+Given a prompt of $M$ tokens, importance scoring proceeds in three phases:
+
+\begin{enumerate}
+  \item \textbf{Draft prefill.} The full prompt is prefilled into a small same-tokenizer draft model (e.g., Qwen3.5-2B at 4-bit quantization, 1.4\,GB) in chunks of 2{,}048 tokens, populating the draft's KV cache. The FLOP ratio thesis (Section~\ref{sec:cost_model}) predicts that minimizing the draft's active parameter count maximizes speedup, favoring the smallest available compatible model.
+
+  \item \textbf{Lookahead decode with attention capture.} Eight autoregressive decode steps are executed with \texttt{AttentionCapture} wrappers installed on each attention layer. These wrappers intercept post-RoPE query vectors via architecture-specific extractors (Section~\ref{sec:extractors}), appending them to a capture buffer before delegating to the original attention computation.
+
+  \item \textbf{Importance computation.} For each attention layer $\ell$ with captured queries $\{q^{(t)}\}_{t=1}^{8}$ and prompt keys $K_\ell \in \mathbb{R}^{h \times M \times d}$:
+    \begin{align}
+      S_\ell^{(t)} &= \text{softmax}\!\left(\frac{q^{(t)} K_\ell^T}{\sqrt{d}}\right) \in \mathbb{R}^{h \times M} \\
+      \bar{S}_\ell^{(t)} &= \text{AvgPool1D}(S_\ell^{(t)},\; \text{kernel}=13)
+    \end{align}
+    Scores are aggregated as $\max$ across layers and heads, then $\text{mean}$ across lookahead steps, yielding importance $I \in \mathbb{R}^M$.
+    Average pooling with kernel 13 smooths the signal, preventing isolated-token artifacts.
+    For GQA models, keys are expanded via \texttt{repeat} to match the query head count before scoring.
+\end{enumerate}
+
+Layers whose cache does not span the full prompt (e.g., sliding-window \texttt{RotatingKVCache} layers caching only 128 tokens) are skipped during importance computation.
+
+After scoring, the draft KV cache is explicitly freed and \texttt{mx.clear\_cache()} is called, reclaiming memory before target prefill begins.
+
+\subsubsection{Architecture-Specific Query Extraction}
+\label{sec:extractors}
+
+The query extractor is auto-detected at runtime based on the attention module's attributes:
+
+\begin{itemize}
+  \item \texttt{q\_norm} present $\rightarrow$ Qwen3.5 path
+  \item No \texttt{rope} attribute $\rightarrow$ Nemotron-H path
+  \item Otherwise $\rightarrow$ Standard (Llama/GPT-OSS) path
+\end{itemize}
+
+\paragraph{Qwen3.5: Gated queries with per-head normalization.}
+The \texttt{q\_proj} output has $2\times$ the expected width, containing both query and gate tensors concatenated along the head dimension.
+The output is reshaped to $(B, L, n_\text{heads}, 2 \cdot d_\text{head})$ and split at the midpoint.
+After splitting, \texttt{q\_norm} (a per-head RMSNorm) is applied before RoPE rotation.
+Treating this as a standard projection produces silent shape errors or incorrect scoring.
+
+\paragraph{Nemotron-H: Heterogeneous layer navigation.}
+Of 88 total layers (40 Mamba-2 + 8 attention + 40 MoE), only the 8 attention layers produce Q/K scores.
+\texttt{\_find\_attention\_layers} navigates the heterogeneous layer structure by inspecting \texttt{block\_type} annotations (\texttt{M} for Mamba, \texttt{*} for attention, \texttt{-} for MLP, \texttt{E} for MoE) and locating modules with a \texttt{mixer} attribute rather than the standard \texttt{self\_attn}.
+\texttt{\_build\_layer\_to\_cache\_map} constructs a compacted index because only Mamba and attention layers have cache entries.
+
+These attention layers have \textbf{no RoPE}---positional information comes entirely from the Mamba-2 SSM state.
+Queries are used as-is for content-based scoring.
+This is a non-trivial engineering challenge the original paper did not address: the draft model (Nano~4B) has 42 heterogeneous layers with only 4 attention layers among 21 Mamba and 17 MLP layers, all in a model where positional information is entirely implicit.
+
+\paragraph{GPT-OSS: RotatingKVCache awareness.}
+Standard query extraction applies, but importance computation must skip sliding-window layers whose \texttt{RotatingKVCache} contains only the last 128 tokens.
+Without this check, importance scores would be computed against a truncated key set, producing misleading rankings.
+Correct handling requires three cache-aware adaptations:
+(1)~layer-level cache introspection to distinguish full-context from sliding-window layers;
+(2)~skipping importance computation for layers whose cache does not span the full prompt;
+(3)~force-preserving the last \texttt{max\_size} positions during sparse selection to ensure sliding-window layers have valid recent context at decode time.
+
+
+\subsection{Chunk Selection}
+
+Tokens are grouped into non-overlapping chunks of $C = 32$ tokens.
+Each chunk is scored by the mean importance of its constituent tokens.
+The top $\lceil k \cdot M/C \rceil$ chunks by score are selected and their token indices returned in sorted order.
+This preserves spatial locality---coherent phrases are kept or dropped as units.
+
+At $k = 0.2$ (our optimal configuration for Qwen3.5-122B), 80\% of prefill computation is eliminated.
+
+
+\subsection{Sparse Prefill with Position-Mapped RoPE}
+\label{sec:sparse_prefill}
+
+The correctness of sparse prefill depends on maintaining correct RoPE angles despite non-contiguous token positions.
+If position angles are not re-mapped, the model perceives selected tokens as adjacent, destroying long-range coherence.
+
+\paragraph{Step 1: Sliding-window tail preservation.}
+For architectures using \texttt{RotatingKVCache} (GPT-OSS), the last \texttt{max\_size} positions from the prompt are force-merged into the selection set, ensuring sliding-window attention layers have valid recent context at decode time.
+This is auto-detected via cache type inspection.
+
+\paragraph{Step 2: Position mapping during prefill.}
+Each attention layer's \texttt{nn.RoPE} is replaced with \texttt{PositionMappedRoPE}, which maps contiguous batch positions $[0, 1, \ldots, N{-}1]$ to the original absolute positions $[p_0, p_1, \ldots, p_{N-1}]$ of the selected tokens.
+For models with custom RoPE variants (YarnRoPE with pre-computed frequencies, SuScaled RoPE with \texttt{mscale}), the replacement module captures and replays the original frequency tensors and scale factors.
+
+\paragraph{Step 3: Chunked forward pass.}
+Selected tokens are fed through the target model in chunks of \texttt{step\_size} (default 2{,}048), populating the KV cache with entries at correct absolute positions.
+
+\paragraph{Step 4: Decode offset adjustment.}
+After sparse prefill of $N$ selected tokens from $M$ total prompt tokens, the cache offset is $N$ but decode must start at position $M$.
+\texttt{OffsetAdjustedRoPE} wraps the original RoPE module and adds adjustment $\Delta = M - N$ to all offset calls:
+\begin{equation}
+  \text{RoPE\_position}(i) = N + i + (M - N) = M + i \quad \checkmark
+\end{equation}
+
+\paragraph{Step 5: Cleanup.}
+After generation completes, \texttt{cleanup\_rope()} traverses all attention layers and unwraps patched RoPE modules back to their originals, ensuring the model is unmodified for subsequent requests.
+
+\paragraph{Nemotron-H (no RoPE).}
+Steps 2 and 4 are skipped entirely---Nemotron-H's attention layers have no RoPE, deriving positional information from the Mamba-2 SSM state instead.
+
+The SSM layers are updated \emph{only} on retained tokens; skipped tokens do not contribute to state evolution.
+Concretely: the Mamba-2 recurrence $h_t = A h_{t-1} + B x_t$ advances only at selected positions, so the hidden state after processing $N$ selected tokens diverges from the state after processing all $M$ tokens.
+This alters the underlying state trajectory.
+However, in practice, the retained tokens---selected by attention importance---appear sufficient to preserve global semantics for long-context tasks: our server-side benchmarks show $2.10$--$2.19\times$ speedup across 8K--64K tokens with no observed quality degradation (Section~\ref{sec:quality}).
+Quantifying the SSM state drift (e.g., L2 distance between full and sparse hidden states) is left to future work.
+Despite this mismatch, empirical evaluation shows no quality degradation under our test suite (Section~\ref{sec:quality}).
+We attribute this to the hybrid architecture: the 8 full-attention layers retain the most important $N$ tokens with correct content-based scores, providing long-range context that compensates for gaps in the SSM's recurrent state.
+This is an empirical result, not a theoretical guarantee; extending \specprefill{} to pure SSM architectures would require additional analysis.
+
+
+\subsection{Unified Memory Cost Model}
+\label{sec:cost_model}
+
+We formalize the relationship between hardware architecture and \specprefill{} efficiency.
+
+On discrete-GPU systems, the cost of \specprefill{} includes a data transfer term $T$ (PCIe bandwidth, memory copies between draft and target devices):
+\begin{equation}
+\label{eq:speedup_gpu}
+  \text{Speedup}_{\text{GPU}} = \frac{C_t}{C_d + T + k \cdot C_t}
+\end{equation}
+
+On unified memory, $T = 0$, simplifying to Equation~\ref{eq:speedup}.
+The speedup is determined by the empirical wall-clock cost ratio $r = C_d / C_t$ and keep percentage $k$:
+\begin{equation}
+\label{eq:speedup_ratio}
+  \text{Speedup} = \frac{1}{r + k + \epsilon}
+\end{equation}
+where $\epsilon$ captures fixed overhead from chunk selection, RoPE patching, memory management (\texttt{mx.clear\_cache()}), and architecture-specific scoring costs.
+In practice, $\epsilon$ ranges from 0.03 (Qwen3.5, low overhead) to 0.30 (Nemotron-H, where navigating 88 heterogeneous layers adds cost).
+We emphasize that this model is \emph{descriptive}: it correctly predicts the \emph{ranking} of configurations by speedup (Table~\ref{tab:cost_model}) but does not predict exact magnitudes, as $\epsilon$ varies by architecture.
+The value of the model is in draft selection---given two candidate drafts, the one with lower $r$ will yield higher speedup---not in predicting absolute speedup from first principles.
+
+\textbf{Boundary conditions.}
+Equation~\ref{eq:speedup_ratio} assumes: (1)~sequential, single-request prefill (no batching); (2)~a FLOP-dominated regime where compute, not memory bandwidth, is the bottleneck; and (3)~negligible KV cache fragmentation cost.
+In bandwidth-bound regimes or under heavy concurrent batching, $\epsilon$ may dominate and reduce the model's predictive accuracy.
+
+However, $r$ is not simply the ratio of active parameters.
+On unified memory, all expert weights reside in the same address space---there is no discrete VRAM to overflow into.
+MoE forward passes incur costs beyond active-parameter FLOPs: router gating across all experts, weight loading for selected experts from unified memory, and memory bandwidth pressure from the full parameter footprint.
+As a result, the empirical wall-clock cost of a MoE forward pass on unified memory scales closer to \emph{total} parameters than to active parameters alone.
+A small dense 2B draft scoring against a 122B MoE target (10B active but 122B total) therefore achieves $r \approx 0.02$, reflecting the full-parameter cost disparity rather than the ${\sim}5\times$ active-parameter ratio.
+This is a finding specific to unified memory systems: on discrete GPUs where expert weights page between CPU and GPU memory, the effective cost ratio may differ.
+
+For MoE models where total parameters far exceed active parameters, a small dense draft has $r \ll 1$.
+This model assumes $C_t$ scales linearly with token count; at extreme context lengths (${\geq}$128K), the $O(n^2)$ attention component causes superlinear growth in $C_t$, and measured speedup can exceed the linear prediction (Table~\ref{tab:cost_model}).
+Table~\ref{tab:cost_model} compares predicted ($\epsilon = 0$) and measured speedups, with the gap $G = \text{Predicted} / \text{Measured}$ quantifying per-configuration overhead:
+
+\begin{table}[h]
+\centering
+\small
+\begin{tabular}{@{}lcccc@{}}
+\toprule
+\textbf{Configuration} & $\boldsymbol{r = C_d/C_t}$ & \textbf{Predicted} ($k{=}0.2, \epsilon{=}0$) & \textbf{Measured} & $\boldsymbol{G}$ \\
+\midrule
+4B / 122B MoE (10B active) & $\sim$0.03 & 4.3$\times$ & 2.90$\times$ & 1.5$\times$ \\
+2B / 122B MoE (10B active) & $\sim$0.02 & 4.5$\times$ & 4.11$\times$ & 1.1$\times$ \\
+2B / 122B MoE (10B active)$^\ddagger$ & $\sim$0.02 & 4.5$\times$ & 5.45$\times$ & 0.8$\times$ \\
+Qwen-4B / 35B MoE (3B active) & $\sim$0.10 & 3.3$\times$ & 1.86$\times$ & 1.8$\times$ \\
+4B / 120B Nemotron-H (12B active) & $\sim$0.03 & 4.3$\times$ & 2.17$\times$ & 2.0$\times$ \\
+20B / 120B GPT-OSS (120B active) & $\sim$0.17 & 2.7$\times$ & 1.28$\times$ & 2.1$\times$ \\
+\bottomrule
+\multicolumn{5}{l}{\footnotesize $^\ddagger$Measured at 128K tokens. Measured speedup exceeds prediction due to $O(n^2)$ attention.}
+\end{tabular}
+\caption{Cost model predictions ($\epsilon = 0$) vs.\ measured speedups at $k = 0.2$, 16K tokens unless noted. $G = \text{Predicted} / \text{Measured}$; values $> 1$ indicate overhead exceeding the model, $< 1$ indicates superlinear baseline growth benefiting \specprefill{} beyond the linear prediction.}
+\label{tab:cost_model}
+\end{table}
+
+The $G$ values reveal architecture-dependent overhead.
+Nemotron-H ($G = 2.0$) has the highest $\epsilon$: the target has 88 heterogeneous layers (40~Mamba + 8~attention + 40~MoE) and the draft has 42 (21~Mamba + 4~attention + 17~MLP), requiring architecture-aware layer navigation for both scoring and sparse prefill.
+GPT-OSS ($G = 2.1$) combines high $r$ (0.17, the 20B draft dominates the denominator) with sliding-window cache management overhead.
+The Qwen3.5-122B configurations ($G = 1.1$--$1.5$, 5 trials) have the lowest overhead, benefiting from uniform architecture and favorable MoE FLOP ratios; the 35B target ($G = 1.8$) sits higher due to its smaller active-parameter count (3B) narrowing the draft-to-target gap.
+At 128K, $G < 1$ because the superlinear baseline growth (Section~\ref{sec:experiments}) is not captured by the linear cost model.
+
+\paragraph{Comparison with GPU results.}
+The original \specprefill{} on discrete GPUs~\cite{specprefill} reports up to $7.66\times$ TTFT reduction on Llama-3.1-405B-FP8, benefiting from batch-level parallelism and GPU-optimized attention kernels not available on Apple Silicon.
+Our lower absolute speedups ($3.71$--$5.45\times$ on 122B) reflect the single-request, unbatched setting and MLX's Metal compute pipeline.
+The unified memory contribution is not higher \emph{absolute} speedup but a \emph{simpler cost model}: eliminating $T$ makes the FLOP ratio the sole dominant predictor, enabling principled draft model selection without profiling transfer overhead.
+
+
+% ============================================================
+\section{System Integration}
+\label{sec:system}
+% ============================================================
+
+\subsection{Composition with System Prompt KV Caching}
+\label{sec:composition}
+
+In production agentic workflows, the system prompt (tool definitions, instructions, context documents) often constitutes 10--20K tokens that remain identical across requests.
+System prompt KV caching snapshots this prefix on the first request and restores it for subsequent requests, eliminating redundant prefill.
+
+These optimizations operate on orthogonal axes: KV caching eliminates \emph{prefix} cost (identical system prompt); \specprefill{} reduces \emph{suffix} cost (variable user context).
+This is why they compose multiplicatively rather than additively.
+When both techniques are active, \specprefill{} operates on the \emph{suffix} only (user and assistant turns), receiving a \texttt{position\_offset} equal to the system token count.
+The scoring phase evaluates only suffix tokens; sparse prefill maps positions relative to the offset so selected tokens land at correct absolute positions in the full context.
+
+The threshold check uses suffix length, not full prompt length, ensuring \specprefill{} activates only when the suffix itself is long enough to benefit.
+
+\begin{table}[h]
+\centering
+\small
+\begin{tabular}{@{}lcc@{}}
+\toprule
+\textbf{Configuration} & \textbf{TTFT (s)} & \textbf{Speedup} \\
+\midrule
+Baseline (cold, full prefill) & 517.5 & 1.0$\times$ \\
+System KV cache only & 417.1 & 1.24$\times$ \\
+\textbf{Combined (SysKV + SP 20\%)} & \textbf{92.5} & \textbf{5.59$\times$} \\
+Combined (repeat) & 92.4 & 5.60$\times$ \\
+\bottomrule
+\end{tabular}
+\caption{Composition of system prompt KV caching and \specprefill{} on Qwen3.5-122B, 2B draft, M2~Ultra 128\,GB. The prompt (73K tokens) is a realistic agentic workload: ${\sim}$10K system prefix (tool definitions, instructions) + ${\sim}$63K user context. System KV cache saves the prefix; \specprefill{} sparse-prefills the suffix. The combined $5.6\times$ speedup exceeds either technique alone.}
+\label{tab:composition}
+\end{table}
+
+
+\subsection{The Speculative Stack}
+
+\specprefill{} is not a standalone optimization but one phase of a three-phase speculative pipeline:
+
+\begin{enumerate}
+  \item \textbf{Score} (\specprefill{}): Draft model (2B) identifies important tokens via attention scoring.
+  \item \textbf{Sparse Prefill}: Target model (122B) processes selected token chunks with position-mapped RoPE.
+  \item \textbf{MTP Decode}: Target model with Multi-Token Prediction heads generates output tokens speculatively (Qwen3.5 only; Nemotron-H and GPT-OSS skip this phase).
+\end{enumerate}
+
+The draft model used in Phase~1 is architecturally separate from MTP's prediction heads (which are part of the target model's weights).
+The draft KV cache is freed after Phase~1 (\texttt{mx.clear\_cache()}), and the draft model's static weights (1.4--3.0\,GB depending on draft size) remain resident for amortized startup cost across requests.
+At extreme context lengths (128K+), the draft's transient KV cache size becomes a practical constraint: the 2B draft produces 1.5\,GB at 128K tokens, fitting comfortably alongside the target model's ${\sim}$100\,GB on a 128\,GB system.
+
+
+\subsection{Per-Request API and Graceful Fallback}
+
+The OpenAI-compatible API accepts per-request overrides:
+
+\begin{lstlisting}[language=Python]
+extra_body = {
+    "specprefill": True,          # force enable (bypass threshold)
+    "specprefill_keep_pct": 0.15  # override server default
+}
+\end{lstlisting}
+
+The default threshold (8{,}192 tokens) is enforced only in server-default mode; explicit \texttt{specprefill: true} bypasses it.
+Any error during scoring or sparse prefill triggers graceful fallback to full prefill---no request fails due to \specprefill{}.
+
+
+% ============================================================
+\section{Experiments}
+\label{sec:experiments}
+% ============================================================
+
+\subsection{Setup}
+
+\paragraph{Hardware.} Apple M2~Ultra, 128\,GB unified memory, Mac Studio, macOS~26.3.1.
+
+\paragraph{Software.} MLX~0.31.1, vllm-mlx~0.2.6 with patches, Python~3.12.
+
+\paragraph{Sampling.} Qwen3.5 models: $\text{temp}=0.6$, $\text{top\_p}=0.95$, $\text{top\_k}=20$ (official thinking+coding profile~\cite{qwen35}).
+Nemotron-H: $\text{temp}=1.0$, $\text{top\_p}=0.95$ (NVIDIA model card: ``trained and evaluated with'').
+GPT-OSS: $\text{temp}=0.6$, $\text{top\_p}=0.95$.
+All models run with thinking mode enabled (\texttt{enable\_thinking=true}).
+Sampling parameters do not affect TTFT measurement (prefill is deterministic; sampling occurs only during generation).
+
+\paragraph{Methodology.} Five trials per configuration for Qwen3.5-122B (plus one warmup); two trials for remaining configurations. Server-side TTFT measured via streaming OpenAI-compatible API.
+Server restarted between configuration changes.
+
+\paragraph{Reproducibility.} Our \specprefill{} implementation is available as patches against vllm-mlx~0.2.6 (PRs~\#175, \#180) and mlx-lm~0.31.2 (PR~\#990), with benchmark scripts (\texttt{bench-specprefill}, \texttt{bench-specprefill-adversarial}, \texttt{bench-specprefill-perplexity}) included in the repository.
+All models are publicly available Qwen3.5 quantizations.
+Nemotron-H and GPT-OSS weights are available from their respective model hubs.
+Experiments require an Apple Silicon system with $\geq$128\,GB unified memory for the 122B configuration; the 35B configuration runs on 64\,GB systems.
+
+\begin{table}[h]
+\centering
+\small
+\begin{tabular}{@{}llllllr@{}}
+\toprule
+\textbf{Model} & \textbf{Role} & \textbf{Architecture} & \textbf{Params} & \textbf{Active} & \textbf{Quant} & \textbf{RAM} \\
+\midrule
+Qwen3.5-122B-VLM-MTP & Target & MoE & 122B & 10B & 5-bit & 79\,GB \\
+Qwen3.5-35B-VLM-MTP & Target & MoE & 35B & 3B & 8-bit & 38\,GB \\
+Nemotron-H 120B & Target & Mamba-2 + Attn + MoE & 120B & 12B & 5-bit & 83\,GB \\
+GPT-OSS 120B & Target & Dense + sliding window & 120B & 120B & 5-bit & 58\,GB \\
+\midrule
+Qwen3.5-4B-VLM-MTP & Draft & Dense hybrid & 4B & 4B & 4-bit & 3.0\,GB \\
+Qwen3.5-2B-OptiQ & Draft & Hybrid + MoE & 2B & $<$2B & 4-bit & 1.4\,GB \\
+Nemotron-H Nano 4B & Draft & Mamba-2 + Attn hybrid & 4B & 4B & 4-bit & 2.1\,GB \\
+GPT-OSS-20B & Draft & MoE & 20B & 3.6B & 4-bit & 10\,GB \\
+\bottomrule
+\end{tabular}
+\caption{Models evaluated. \specprefill{} requires draft and target to share the same tokenizer. Qwen3.5 drafts (2B, 4B) serve all Qwen3.5 targets (248K vocabulary); the 2B is the primary draft model. Nemotron-H Nano~4B serves Nemotron-H~120B (131K vocabulary). GPT-OSS-20B is the smallest available same-family draft for GPT-OSS-120B (201K vocabulary).}
+\label{tab:models}
+\end{table}
+
+
+\subsection{TTFT Benchmarks}
+
+\begin{table}[h]
+\centering
+\small
+\begin{tabular}{@{}llccccc@{}}
+\toprule
+\textbf{Model} & \textbf{Draft} & \textbf{8K} & \textbf{16K} & \textbf{32K} & \textbf{64K} & \textbf{128K} \\
+\midrule
+Qwen3.5-122B (MoE, 10B) & 4B$^\ddagger$ & 2.79$\times$ & 2.90$\times$ & ---$^\dagger$ & ---$^\dagger$ & ---$^\dagger$ \\
+Qwen3.5-122B (MoE, 10B) & 2B & 3.71$\times$ & 4.11$\times$ & 4.23$\times$ & 4.50$\times$ & 5.45$\times$ \\
+Qwen3.5-35B (MoE, 3B) & 4B & 1.81$\times$ & 1.86$\times$ & 1.85$\times$ & 1.84$\times$ & --- \\
+Nemotron-H 120B (hybrid) & Nano-4B & 2.10$\times$ & 2.17$\times$ & 2.19$\times$ & 2.19$\times$ & --- \\
+GPT-OSS 120B (dense) & 20B & 1.24$\times$ & 1.28$\times$ & --- & --- & --- \\
+\bottomrule
+\multicolumn{7}{l}{\footnotesize $^\dagger$Not measured at this context length.} \\
+\multicolumn{7}{l}{\footnotesize $^\ddagger$4B draft includes VLM weights (3.0\,GB); the 2B text-only draft (1.4\,GB) is the primary configuration.}
+\end{tabular}
+\caption{TTFT speedups at 20\% keep, 5 trials (mean). Speedup increases with prompt length as scoring overhead is amortized and $O(n^2)$ attention savings compound. Qwen3.5-122B and 4B rows use 5 trials; other rows use 2 trials.}
+\label{tab:ttft}
+\end{table}
+
+For Qwen3.5-122B with the 2B draft, the absolute TTFT at 64K tokens drops from $417.6 \pm 0.6$\,s to $92.8 \pm 0.8$\,s.
+At 128K tokens: $1{,}155.8 \pm 8.5$\,s (19.3~minutes) $\rightarrow$ $212.3 \pm 1.9$\,s (3.5~minutes), a \textbf{5.45$\times$} reduction.
+At 8K tokens: $45.0 \pm 0.1$\,s $\rightarrow$ $12.1 \pm 0.03$\,s.
+Standard deviations across 5 trials are $<$1\% of the mean at all context lengths, confirming measurement stability.
+Speedup increases monotonically with context length, from $3.71\times$ at 8K to $5.45\times$ at 128K, consistent with the amortization of fixed scoring overhead and the superlinear growth of baseline prefill cost.
+
+\paragraph{Nemotron-H: architecture-limited speedup plateau.}
+Nemotron-H shows a flat speedup profile ($2.10\times$ at 8K to $2.19\times$ at 64K), in contrast to Qwen3.5's monotonically increasing curve.
+This plateau is explained by the hybrid architecture: only 8 of 88 layers are attention---the remaining 80 layers (40~Mamba-2 SSM + 40~MoE feed-forward) scale linearly with token count regardless of \specprefill{}.
+The $O(n^2)$ attention component that drives Qwen3.5's compounding speedup at long contexts constitutes only ${\sim}9\%$ of Nemotron-H's total compute, so the quadratic savings are a small fraction of overall prefill cost.
+This confirms the architecture-dependent nature of the cost model: \specprefill{} benefit scales with the attention fraction of total computation.
+In absolute terms, Nemotron-H 120B TTFT drops from 58\,s to 27\,s at 16K and from 253\,s to 116\,s at 64K.
+For Qwen3.5-35B: 41\,s to 22\,s at 16K.
+
+\paragraph{Superlinear scaling at extreme context lengths.}
+The 128K baseline (1{,}156\,s) is $2.77\times$ the 64K baseline (418\,s), not the $2\times$ expected from linear scaling.
+This superlinear growth arises from the $O(n^2)$ attention component in chunked prefill: each 2{,}048-token chunk attends to all preceding tokens, so cumulative attention FLOPs grow quadratically.
+\specprefill{} benefits disproportionately: by reducing the effective sequence length from $N$ to $kN$, it reduces the cumulative attention FLOPs---which scale approximately as $N^2$ under full-prefix chunked prefill---by approximately $k^2$.
+At $k = 0.2$, this yields a ${\sim}25\times$ reduction in attention computation, not the $5\times$ that linear token-count scaling would suggest.
+Draft scoring remains $O(N)$, so its cost grows linearly while the attention savings grow quadratically, explaining why the measured 128K speedup ($5.45\times$) exceeds the linear prediction ($4.5\times$, Table~\ref{tab:cost_model}).
+MoE models exhibit stronger gains because sparse prefill reduces both the quadratic attention cost (sequence length) and the number of active expert evaluations---80\% fewer tokens means 80\% fewer expert routing and weight-loading operations.
+
+\paragraph{Draft model selection.}
+Draft models are selected to maximize FLOP asymmetry while maintaining tokenizer and architectural compatibility.
+The selection criteria are: (1)~same tokenizer as the target (required---token IDs are passed directly without translation); (2)~smallest available model in the family (to maximize the FLOP ratio $r$); (3)~presence of attention layers for importance scoring (at least 4 layers suffice per our Nemotron-H validation).
+Our primary configuration uses a 2B draft ($r \approx 0.02$, 1.4\,GB); complementary measurements with a 4B draft ($r \approx 0.03$, 3.0\,GB) at 8K--16K confirm the relationship: the lower-ratio 2B achieves $4.11\times$ vs.\ $2.90\times$ at 16K.
+
+The GPT-OSS result is a \textbf{negative result confirming the ratio thesis}: the 20B draft model (the smallest available in the GPT-OSS family) has an unfavorable FLOP ratio of $\sim$0.17, yielding only $1.24$--$1.28\times$ speedup.
+This validates that architecture is not the determining factor---the FLOP ratio is.
+A hypothetical 4B GPT-OSS draft ($r \approx 0.03$) would be predicted to achieve $\sim$2.5--3$\times$ under our cost model, but no such model exists in the GPT-OSS family (Section~\ref{sec:future}).
+
+
+\subsection{Draft-to-Target FLOP Ratio Analysis}
+\label{sec:ratio}
+
+\begin{figure}[h]
+\centering
+\includegraphics[width=0.9\textwidth]{figures/ratio-speedup.pdf}
+\caption{Measured speedup vs.\ draft-to-target FLOP ratio. The theoretical upper bound (Eq.~\ref{eq:speedup_ratio}) correctly predicts the ranking. Overhead from RoPE patching, chunk selection, and architecture-specific scoring reduces measured values below the theoretical curve.}
+\label{fig:ratio}
+\end{figure}
+
+The FLOP ratio $r = C_d / C_t$ is the dominant predictor of \specprefill{} benefit on unified memory.
+Across our configurations (Table~\ref{tab:cost_model}), $r$ spans from 0.02 (2B/122B MoE) to 0.17 (20B/120B dense), and measured speedup tracks this ratio monotonically.
+Complementary 4B draft measurements at 8K--16K ($r \approx 0.03$, $2.79$--$2.90\times$) confirm the ratio relationship: the lower-ratio 2B achieves higher speedup at the same context length.
+
+This relationship distinguishes unified-memory \specprefill{} from GPU-based implementations, where PCIe bandwidth introduces an additional term $T$ (Equation~\ref{eq:speedup_gpu}) that weakens the FLOP-ratio signal.
+
+
+\subsection{Keep Percentage Ablation}
+
+\begin{table}[h]
+\centering
+\small
+\begin{tabular}{@{}ccccc@{}}
+\toprule
+\textbf{Keep \%} & \textbf{TTFT (s)} & \textbf{Speedup} & \textbf{Needle} & \textbf{JSON} \\
+\midrule
+10\% & 23.90 & 3.85$\times$ & PASS & 1/2$^\S$ \\
+20\% & 31.30 & 2.94$\times$ & PASS & PASS \\
+30\% & 38.88 & 2.37$\times$ & PASS & PASS \\
+50\% & 53.78 & 1.71$\times$ & PASS & PASS \\
+100\% (baseline) & 91.97 & 1.0$\times$ & PASS & PASS \\
+\bottomrule
+\multicolumn{5}{l}{\footnotesize $^\S$JSON extraction at 10\% keep: one regression (1 of 3 values wrong) in trial 1, pass in trial 2.} \\
+\end{tabular}
+\caption{Keep percentage ablation on Qwen3.5-122B at ${\sim}$16K tokens with the 4B draft. ``Needle'' is UUID retrieval at all three depths (10\%, 50\%, 90\%); ``JSON'' is exact value extraction from an 80-record array. The 10\% row shows the quality boundary.}
+\label{tab:ablation}
+\end{table}
+
+The curve shows a clear knee at ${\sim}$20\%: all quality tests pass while delivering $2.94\times$ speedup (4B draft), and marginal speedup gains diminish beyond this point while compute cost increases linearly with $k$.
+This is our recommended operating point.
+At 10\%, speedup increases to $3.85\times$ but structured data extraction becomes unreliable (1 JSON regression in 2 trials), establishing 10\% as the quality boundary.
+At 50\%, the $1.71\times$ speedup is marginal relative to the scoring overhead incurred.
+
+
+\subsection{Quality Validation}
+\label{sec:quality}
+
+Because \specprefill{} retains the highest-scoring tokens by attention importance, the model keeps the tokens it would have attended to most heavily during generation.
+We validate this through three complementary evaluations, leading with the most concrete and falsifiable tests.
+
+\paragraph{Adversarial tests (primary).}
+Eight test types designed to expose sparse-prefill weaknesses: needle-in-a-haystack (UUID retrieval at 10\%, 50\%, 90\% depth), JSON value extraction from an 80-record array, code bug detection, back-reference, mixed-language retrieval, and XML structure extraction.
+
+At 20\% keep: \textbf{0/16 regressions} across 2 trials $\times$ 8 tests.
+All needle-in-a-haystack and JSON extraction tests pass under both baseline and \specprefill{}.
+\specprefill{} does not drop needles or corrupt structured data retrieval at 20\% keep.
+
+At 10\% keep: 1/16 regressions---a JSON extraction test returned one incorrect value out of three (trial 1 of 2; trial 2 passed).
+All needle tests pass at all depths even at 10\%, suggesting that high-importance tokens (which needles represent) are robustly retained; the failure mode at extreme sparsity is degraded recall of \emph{low-salience} structured data.
+
+\paragraph{ROUGE-L (lexical similarity).}
+We compare outputs from \specprefill{} (20\% keep) against full-prefill baselines on six real-task prompts (code generation, code review, summarization, reasoning, tutorial writing, tool use), each targeting $\sim$8K actual tokens on Qwen3.5-122B.
+
+To establish a variance floor, we first compare two independent baseline runs against each other:
+
+\begin{center}
+\small
+\begin{tabular}{lc}
+\toprule
+\textbf{Comparison} & \textbf{ROUGE-L F1} \\
+\midrule
+Baseline vs.\ baseline (variance floor) & 0.190 $\pm$ 0.174 \\
+\specprefill{} vs.\ baseline & 0.236 \\
+\bottomrule
+\end{tabular}
+\end{center}
+
+The high baseline-vs-baseline variance ($0.190 \pm 0.174$) demonstrates that lexical similarity between outputs is dominated by the model's inherent stochasticity at $\text{temp}=0.6$, not by any effect of \specprefill{}.
+The \specprefill{} similarity (0.236) falls within the baseline noise floor, but this comparison is not informative enough to draw quality conclusions from---the adversarial tests above provide the primary quality evidence.
+
+\paragraph{LLM-as-Judge (supporting).}
+Blinded A/B evaluation scores coherence, completeness, and accuracy (1--5 scale) for both baseline and \specprefill{} outputs, plus an overall equivalence rating:
+
+\begin{center}
+\small
+\begin{tabular}{lcc}
+\toprule
+\textbf{Comparison} & \textbf{Avg.\ Equivalence} \\
+\midrule
+Baseline vs.\ baseline & 3.0 / 5.0 \\
+\specprefill{} vs.\ baseline & 3.0 / 5.0 \\
+\bottomrule
+\end{tabular}
+\end{center}
+
+\paragraph{Perplexity (distributional).}
+We measure next-token perplexity on 256-token continuations after full vs.\ sparse prefill (20\% keep) on five documents at 8K context (code, documentation, LaTeX, mixed):
+
+\begin{center}
+\small
+\begin{tabular}{lccc}
+\toprule
+\textbf{Document type} & \textbf{PPL (full)} & \textbf{PPL (sparse)} & \textbf{Ratio} \\
+\midrule
+Python (engine code)      & 1.85 & 2.53 & 1.37 \\
+Python (benchmark script) & 1.66 & 1.74 & 1.05 \\
+Python (test harness)     & 1.49 & 1.58 & 1.06 \\
+LaTeX (this paper)        & 2.00 & 2.14 & 1.07 \\
+Mixed (concatenated)      & 2.76 & 3.17 & 1.15 \\
+\midrule
+\textbf{Mean}             & 1.95 & 2.23 & \textbf{1.14} \\
+\bottomrule
+\end{tabular}
+\end{center}
+
+Mean perplexity increases 14\%, though 4 of 5 documents show $\leq$7\% increase (median ratio 1.07).
+The outlier (dense engine code, 1.37$\times$) contains many local variable dependencies where discarded tokens carry predictive information.
+This distributional shift does not translate to generation quality degradation in our adversarial or LLM-judge evaluations above---sampling smooths over small distributional differences that perplexity measures precisely.
+16K context was not tested: loading both the 122B target and 2B draft for offline evaluation leaves insufficient headroom on 128\,GB unified memory.
+
+\paragraph{Limitations.}
+Six prompts, eight adversarial types, and five perplexity documents confirm no catastrophic quality loss and validate the methodology, but the sample size is insufficient for tight confidence intervals.
+We make no claim of statistical equivalence---only that \textbf{no measurable quality degradation was observed under our evaluation at 20\% keep}.
+The 10\% keep JSON regression demonstrates that the quality boundary is observable and characterizable within our framework.
+All quality evaluations were conducted on Qwen3.5-122B.
+Nemotron-H and GPT-OSS were validated via pipeline tests (Section~\ref{sec:method}) but lack server-side adversarial evaluation.
+Future work will include larger-scale evaluation on standardized long-context benchmarks (e.g., RULER, LongBench) and extend quality validation to non-Qwen architectures.
+
+
+\subsection{Memory Profile}
+\label{sec:memory}
+
+\begin{table}[h]
+\centering
+\small
+\begin{tabular}{@{}lrrl@{}}
+\toprule
+\textbf{Component} & \textbf{Memory} & \textbf{Cumulative} & \textbf{Notes} \\
+\midrule
+Target weights (122B, 5-bit) & 79\,GB & 79\,GB & Fixed at load \\
+Draft weights (2B, 4-bit) & 1.4\,GB & 80\,GB & Fixed at load \\
+MLX Metal cache limit & 4\,GB & 84\,GB & Computation scratch \\
+Target KV cache (128K) & 12\,GB & 96\,GB & 96\,KB/token $\times$ 127K \\
+Draft KV cache (128K, transient) & 1.5\,GB & 97\,GB & 12\,KB/token, freed after scoring \\
+OS + framework overhead & $\sim$25\,GB & $\sim$122\,GB & Observed via \texttt{memory\_pressure} \\
+\midrule
+\textbf{Peak (128K baseline)} & & \textbf{$\sim$122\,GB} & Of 128\,GB unified \\
+\textbf{Peak (128K \specprefill{})} & & \textbf{$\sim$122\,GB} & Draft KV transient, not additive \\
+\bottomrule
+\end{tabular}
+\caption{Memory budget for Qwen3.5-122B with 2B draft at 128K tokens on M2~Ultra 128\,GB. The draft KV cache is transient: allocated during scoring, freed via \texttt{mx.clear\_cache()} before target prefill begins. Peak \specprefill{} memory $\approx$ baseline peak because the draft and target KV caches are never resident simultaneously.}
+\label{tab:memory}
+\end{table}
+
+
+% ============================================================
+\section{Discussion}
+\label{sec:discussion}
+% ============================================================
+
+\subsection{The MoE Sweet Spot}
+
+\specprefill{} benefits MoE architectures more than dense models.
+In MoE models, each token is routed to a subset of experts during the forward pass, but the routing computation and expert weight loading occur for \emph{every} token regardless.
+By reducing the token count during prefill, \specprefill{} reduces the total number of expert activations---not just attention FLOPs, but the dominant feed-forward computation.
+The savings exceed the keep fraction: at $k = 0.2$, the model processes 80\% fewer tokens through its expert layers, each involving sparse routing across 128 experts (Qwen3.5-122B).
+
+Dense models, by contrast, apply the same computation to every token regardless of routing.
+The savings from \specprefill{} on dense models are proportional only to the reduced attention and MLP computation, which is less dramatic when the model is fully compute-bound.
+
+
+\subsection{When SpecPrefill Does Not Help}
+
+\paragraph{Dense models with large drafts.}
+When the FLOP ratio $r$ exceeds $\sim$0.15, scoring overhead consumes most of the potential savings.
+Our GPT-OSS result (20B draft, $r \approx 0.17$, speedup $1.28\times$) demonstrates this boundary.
+No smaller GPT-OSS model exists; the proprietary tokenizer (201K vocabulary) prevents cross-family draft substitution without a re-tokenization layer (Section~\ref{sec:future}).
+
+\paragraph{Short prompts.}
+Below $\sim$4K tokens, the fixed overhead of draft scoring and RoPE patching exceeds the savings from sparse prefill.
+Our default threshold of 8{,}192 tokens reflects this empirical boundary.
+
+\paragraph{Comparison with CritiPrefill.}
+CritiPrefill~\cite{critiprefill} achieves sparse prefill without a draft model by using the target's own attention scores from an initial partial prefill.
+On dense standard transformers where all layers are attention, CritiPrefill achieves 2.7--3.0$\times$ speedup (reported on Llama-3-8B and Yi-9B at 128K context).
+However, it only saves attention FLOPs---on MoE architectures where attention constitutes 7--9\% of total computation, the attention-only savings would be proportionally limited; our analysis estimates $\sim$1.03--1.08$\times$.
+\specprefill{} saves \emph{all} FLOPs (attention + routing + expert computation) for dropped tokens, yielding $3.7$--$5.5\times$ on MoE targets vs.\ the estimated ${\sim}1.03$--$1.08\times$ for attention-only savings.
+
+
+\subsection{Limitations}
+
+\begin{itemize}
+  \item \textbf{Draft model dependency.} Requires a small model with a compatible tokenizer. This limits applicability to model families with multiple size variants (Qwen3.5: 2B/4B/27B/35B/122B; GPT-OSS: only 20B/120B).
+  \item \textbf{Nemotron-H SSM state.} SSM layers are updated only on retained tokens; skipped tokens do not contribute to state evolution. The resulting state trajectory diverges from full prefill. Empirically safe under our evaluation ($2.10$--$2.19\times$ with no observed quality degradation), but the magnitude of state drift is not quantified.
+  \item \textbf{Quality validation scale.} Six prompts and eight adversarial types validate methodology and confirm no catastrophic loss, but are insufficient for tight confidence intervals.
+  \item \textbf{Single hardware platform.} Results on M2~Ultra (128\,GB). Memory bandwidth and compute characteristics differ on M3/M4 variants, and the optimal keep percentage may shift.
+  \item \textbf{Single-request evaluation.} The serving engine serializes requests via \texttt{asyncio.Lock}. We do not evaluate \specprefill{} under concurrent load.
+\end{itemize}
+
+
+% ============================================================
+\section{Related Work}
+\label{sec:related}
+% ============================================================
+
+\paragraph{Sparse prefill.}
+\citet{specprefill} introduce attention-based sparse prefill using a draft model on discrete GPU architectures.
+CritiPrefill~\cite{critiprefill} achieves draft-free sparse prefill using the target model's own attention.
+We extend the draft-based approach to unified memory hardware and non-transformer architectures.
+To our knowledge, no prior work addresses sparse prefill specifically for unified memory systems.
+
+\paragraph{Speculative decoding.}
+\citet{leviathan2023fast} and \citet{chen2023accelerating} propose using a draft model to speculatively generate candidate tokens verified by the target model.
+Multi-Token Prediction (MTP)~\cite{gloeckle2024better} uses auxiliary prediction heads within the target model itself.
+Our ``Speculative Stack'' composes prefill-phase speculation (\specprefill{}) with decode-phase speculation (MTP), operating in non-overlapping phases.
+
+\paragraph{Efficient attention.}
+FlashAttention~\cite{dao2022flashattention,dao2023flashattention2} optimizes attention computation through tiling and memory-efficient kernels.
+\specprefill{} is orthogonal: it reduces the token count \emph{before} attention, and can compose with efficient attention implementations.
+
+\paragraph{Serving systems.}
+vLLM~\cite{kwon2023efficient} introduces PagedAttention for efficient KV cache management.
+Our system (vllm-mlx) adapts continuous-batching and paged-attention concepts for Apple Silicon.
+\specprefill{} integrates as a per-request prefill optimization within this serving framework.
+
+\paragraph{MLX framework.}
+MLX~\cite{mlx} provides the unified-memory ML runtime enabling our zero-copy scoring approach.
+Prior MLX-based serving work has focused on standard inference optimization; we show that unified memory enables system-level optimizations, specifically zero-copy draft scoring, that are impractical on discrete-GPU architectures.
+
+
+% ============================================================
+\section{Future Work}
+\label{sec:future}
+% ============================================================
+
+\paragraph{Universal draft models with tokenizer translation.}
+Our results are constrained to model families where a small same-tokenizer draft exists.
+A \emph{universal draft model}---trained to score token importance across vocabularies via a learned re-tokenization layer---would decouple \specprefill{} from the family constraint.
+The translation layer would map target token IDs to text, re-tokenize with the universal draft's vocabulary, score importance, and project scores back to target token space via character-offset alignment.
+This is non-trivial (tokenization boundaries differ across BPE vocabularies) but would enable \specprefill{} for models like GPT-OSS where no small same-family draft exists.
+
+\paragraph{CritiPrefill for dense models.}
+For dense architectures where the FLOP ratio is unfavorable, CritiPrefill (draft-free) may be more practical.
+Published results show 2.7--3.0$\times$ on dense 8B--9B transformers at 128K context.
+On our MoE targets, where attention is a small fraction of compute, gains would be limited; the approach warrants investigation for dense models on unified memory where memory access patterns differ from GPU.
+
+\paragraph{SSM state drift analysis.}
+For hybrid models like Nemotron-H, quantifying the L2 distance between full-prefill and sparse-prefill SSM hidden states would characterize the information loss from skipping tokens in recurrent layers and establish quality guarantees beyond empirical testing.
+
+\paragraph{Continuous batching.}
+Our current implementation uses a single-request serialized engine.
+Integrating \specprefill{} with continuous batching would enable concurrent request handling, but is currently blocked by a Metal driver stability issue on macOS~26.3.1.
+
+\paragraph{Hardware generalization.}
+Apple's M3 and M4 generations have different memory bandwidth and compute characteristics.
+The optimal keep percentage and FLOP-ratio threshold may shift on these platforms.
+
+
+% ============================================================
+\section{Conclusion}
+\label{sec:conclusion}
+% ============================================================
+
+We have presented the first implementation of \specprefill{} on unified memory hardware, demonstrating that Apple Silicon's shared address space eliminates the data-transfer overhead that complicates draft-based sparse prefill on discrete GPUs.
+With transfer overhead removed, the cost equation reduces to a single dominant term: the draft-to-target FLOP ratio, validated across six configurations spanning MoE, Mamba-2 hybrid, and sliding-window dense architectures.
+
+On Qwen3.5-122B, \specprefill{} reduces TTFT by $3.71$--$5.45\times$ across 8K--128K tokens with a 1.4\,GB draft model and no observed quality degradation under our evaluation. At 128K tokens, prefill drops from 19.3~minutes to 3.5~minutes.
+Composed with system prompt KV caching, end-to-end speedup reaches $5.6\times$ on a 73K-token production workload.
+The implementation handles architecture-specific challenges (gated queries, heterogeneous SSM/attention layers, sliding-window caches) through auto-detecting adapters that require no user configuration.
+
+\specprefill{} is most effective on MoE and hybrid models where total parameters far exceed active computation, making a small dense draft model orders of magnitude cheaper than the target.
+As large models move to local hardware, reducing prefill cost through techniques like zero-copy draft scoring directly determines whether long-context inference is usable.
+
+
+% ============================================================
+% References
+% ============================================================
+
+\bibliographystyle{plainnat}
+
+\begin{thebibliography}{99}
+
+\bibitem[Chen et~al.(2023)]{chen2023accelerating}
+Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper.
+\newblock Accelerating large language model decoding with speculative sampling.
+\newblock \emph{arXiv preprint arXiv:2302.01318}, 2023.
+
+\bibitem[Dao(2023)]{dao2023flashattention2}
+Tri Dao.
+\newblock Flash{A}ttention-2: Faster attention with better parallelism and work partitioning.
+\newblock \emph{arXiv preprint arXiv:2307.08691}, 2023.
+
+\bibitem[Dao et~al.(2022)]{dao2022flashattention}
+Tri Dao, Daniel~Y. Fu, Stefano Ermon, Atri Rudra, and Christopher R\'{e}.
+\newblock Flash{A}ttention: Fast and memory-efficient exact attention with {IO}-awareness.
+\newblock In \emph{NeurIPS}, 2022.
+
+\bibitem[Gloeckle et~al.(2024)]{gloeckle2024better}
+Fabian Gloeckle, Badr Youbi~Idrissi, Baptiste Rozi\`{e}re, David Lopez-Paz, and Gabriel Synnaeve.
+\newblock Better \& faster large language models via multi-token prediction.
+\newblock \emph{arXiv preprint arXiv:2404.19737}, 2024.
+
+\bibitem[Kwon et~al.(2023)]{kwon2023efficient}
+Woosuk Kwon, Zhuohan Li, Sicheng Zhuang, Ying Sheng, Lianmin Zheng, Cody~Hao Yu, Joseph~E. Gonzalez, Hao Zhang, and Ion Stoica.
+\newblock Efficient memory management for large language model serving with {PagedAttention}.
+\newblock In \emph{SOSP}, 2023.
+
+\bibitem[Leviathan et~al.(2023)]{leviathan2023fast}
+Yaniv Leviathan, Matan Kalman, and Yossi Matias.
+\newblock Fast inference from transformers via speculative decoding.
+\newblock In \emph{ICML}, 2023.
+
+\bibitem[Apple(2023)]{mlx}
+Apple Machine Learning Research.
+\newblock {MLX}: An array framework for Apple Silicon.
+\newblock \url{https://github.com/ml-explore/mlx}, 2023.
+
+\bibitem[Peng et~al.(2024)]{yarn}
+Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole.
+\newblock {YaRN}: Efficient context window extension of large language models.
+\newblock In \emph{ICLR}, 2024.
+
+\bibitem[Su et~al.(2024)]{rope}
+Jianlin Su, Murtadha Ahmed, Yu~Lu, Shengfeng Pan, Wen Liu, and Bo~Liu.
+\newblock {RoFormer}: Enhanced transformer with rotary position embedding.
+\newblock \emph{Neurocomputing}, 568:127063, 2024.
+
+\bibitem[Yao et~al.(2025)]{specprefill}
+Ziteng Yao, Wei Chen, Yushi Huang, and others.
+\newblock {SpecPrefill}: Speculative prefilling for faster long-context {LLM} inference.
+\newblock \emph{arXiv preprint arXiv:2502.02789}, 2025.
+
+\bibitem[Qwen(2025)]{qwen35}
+Qwen Team.
+\newblock {Qwen3.5}: A series of large language models.
+\newblock \url{https://huggingface.co/Qwen/Qwen3.5-122B-A10B}, 2025.
+
+\bibitem[Zhang et~al.(2025)]{critiprefill}
+Junlin Zhang, Jiahao Li, and others.
+\newblock {CritiPrefill}: A segment-level critique framework for efficient long-context {LLM} prefilling.
+\newblock \emph{arXiv preprint}, 2025.
+
+\end{thebibliography}
+
+\end{document}