PRISM: Breaking the O(n) Memory Wall in Long-Context LLM Inference via O(1) Photonic Block Selection

Abstract: Long-context LLM inference is bottlenecked not by compute but by the O(n) memory bandwidth cost of scanning the KV cache at every decode step -- a wall that no amount of arithmetic scaling can break. Recent photonic accelerators have demonstrated impressive throughput for dense attention computation; however, these approaches inherit the same O(n) memory scaling as electronic attention when applied to long contexts. We observe that the real leverage point is the coarse block-selection step: a memory-bound similarity search that determines which KV blocks to fetch. We identify, for the first time, that this task is structurally matched to the photonic broadcast-and-weight paradigm -- the query fans out to all candidates via passive splitting, signatures are quasi-static (matching electro-optic MRR programming), and only rank order matters (relaxing precision to 4-6 bits). Crucially, the photonic advantage grows with context length: as N increases, the electronic scan cost rises linearly while the photonic evaluation remains O(1). We instantiate this insight in PRISM (Photonic Ranking via Inner-product Similarity with Microring weights), a thin-film lithium niobate (TFLN) similarity engine. Hardware-impaired needle-in-a-haystack evaluation on Qwen2.5-7B confirms 100% accuracy from 4K through 64K tokens at k=32, with 16x traffic reduction at 64K context. PRISM achieves a four-order-of-magnitude energy advantage over GPU baselines at practical context lengths (n >= 4K).

• The paper introduces a photonic block selection method that reduces KV cache access latency from O(n) to O(1) by leveraging a broadcast-and-weight architecture.
• The method achieves significant memory traffic and energy reductions, with evaluations showing a recall@8 of 77.3% and energy improvements of over four orders of magnitude.
• System-level tests confirm PRISM's scalability and robustness under realistic hardware impairments, offering a promising path for efficient long-context LLM inference.

Photonic Block Selection for O(1) KV Cache Access in LLMs: An Authoritative Review of "PRISM: Breaking the O(n) Memory Wall in Long-Context LLM Inference via O(1) Photonic Block Selection" (2603.21576)

Motivation and Memory Bottleneck in LLM Inference

Transformer-based LLMs, particularly at scale, are bottlenecked by memory bandwidth rather than arithmetic throughput. Each autoregressive decoding step demands full scans of the key-value (KV) cache, with memory traffic scaling linearly with context length $n$. As real-world context windows approach millions of tokens, the KV cache management becomes a first-class constraint; contemporary architectures like NVIDIA's Vera Rubin confirm this, dedicating hardware to the management and storage of ever-larger KV caches. Recent photonic accelerators for dense attention offer high arithmetic throughput (e.g., Mach-Zehnder mesh-based photonic transformers), but the KV cache scan remains $O(n)$ in memory cost and constitutes the dominant bottleneck for inference.

PRISM Architecture Overview and Photonic Broadcast Paradigm

PRISM leverages a structural correspondence between the block-selection operation in KV cache access and the broadcast-and-weight paradigm available in photonic hardware. In this paradigm, a wavelength-division multiplexed (WDM) laser encodes query sketches onto optical channels. Passive splitting enables the identical broadcast of this query to all signature block channels. Electro-optically programmable microring resonators (MRRs) encode block signatures as weights, and broadband photodetectors enable the computation of analog inner products for all blocks in parallel.

Figure 1: Comparison of GPU-based full scan (O(n) memory bandwidth) vs. PRISM's photonic selective fetch (O(1) access via broadcast channel and block-ranking with top-k retrieval).

The resulting system performs block similarity search in O(1) optical latency and eliminates sequential memory access patterns—crucially, as the context grows, the electronic scan cost rises linearly, while PRISM evaluation remains constant in latency.

Figure 2: PRISM's five-stage pipeline: query encoding, broadcast, MRR-based signature weighting, photodetector summation, and top-k selection/comparators.

Retrieval Head Profiling and Double-Scaling Advantage

Recent advances in attention analysis (DuoAttention, RazorAttention) reveal that only a fraction of attention heads behave as retrieval heads—those attending to distant tokens. The remainder ("streaming heads") operate over local windows and do not require full-cache retrieval. PRISM's system-level evaluation across Qwen2.5-7B and Qwen3-8B demonstrates that the retrieval head fraction exceeds 90% at realistic context lengths ($n \geq 8$K) and increases further with longer contexts.

Figure 3: Retrieval head fraction and retrieval ratio increase monotonically with context, compounding the benefit of O(1) photonic acceleration.

This double-scaling advantage underscores that both the fraction of compute dominated by retrieval heads and the benefit of O(1) selection increase with context—aligning PRISM's hardware leverage with future model scaling trends.

Block Signature Construction and Recall Analysis

Performance is critically linked to block-level signature fidelity. Approaches evaluated include mean-key projection, PCA-based projection, random projection (via JL lemma), and learned projection. Mean-key projection with $d=32$ achieves recall@8 of 77.3% in aggressive settings, with full recall at operational points ($k=32$) for contexts up to 8K tokens. Signed weight encoding (via balanced photodetection) is pivotal, with significant improvements in recall compared to unsigned or split encoding.

Figure 4: Recall@k vs. signature dimension; mean-key projection outperforms random projections in all tested settings.

PRISM's error-tolerance is established through "needle-in-a-haystack" (NIAH) benchmarks, confirming that even under hardware impairments (5-bit quantization, 30 pm thermal drift), block selection preserves full-attention accuracy. Recall degrades slowly with increased impairment; for 6-bit precision, the drop is $<$5% relative to floating-point ideal.

Figure 5: Impact of weight quantization on recall remains modest for 6-bit precision, with thermal drift and detector noise contributing additional but manageable degradation.

Hardware Impairment Modeling and Photonic Link Budget

An extensive device-level impairment study aggregates quantization precision, thermal drift, insertion loss, detector noise, MRR crosstalk, and DAC noise. The optical link budget confirms that $d=32$, $N=256$ configurations with standard laser powers ($P_\text{laser} = 20$ dBm) meet SNR requirements for robust top-k retrieval, with photodetector SNR well above 20 dB for reliable ranking.

Figure 6: Electrical SNR and recall@8 as a function of signature dimension and bank size, demonstrating the SNR threshold for 90%+ recall.

MRR count and chip area scale linearly with $d$ and $N$; emerging TFLN platforms overcome SOI's high heater power needs via capacitive EO tuning. Chip area and thermal power budget analyses demonstrate that multi-bank and time-multiplexed configurations (e.g., $d=64$, $N=1024$, $M=8$) are feasible with current integrated photonic technology.

Figure 7: Scaling projections of MRR count, heater power (for SOI benchmark), and chip area vs. configuration.

Energy and Latency Crossover Analysis

PRISM offers decisive energy benefits at context lengths as low as 4K tokens, with four orders of magnitude lower selection energy ($\sim$2.3 nJ) than GPU full scan ($\sim$16.3 μJ) per query. The photonic pipeline latency ($\sim$9 ns) is three orders of magnitude lower than GPU-based selection ($\sim$1–5 μs), although interface latencies (e.g., PCIe) offset this in practical deployments. Time-multiplexed operation allows trade-off in chip area against selection latency, with O(1) scaling maintained.

Figure 8: Energy crossover contour for PRISM vs. electronic baselines, highlighting the practical regime (context length > 4K) where PRISM is advantageous.

System-Level Evaluation and Benchmarking

End-to-end evaluation on NIAH and LongBench-v2 demonstrates that MRR-impaired block selection achieves identical accuracy to full dense attention for contexts up to 64K tokens; at longer contexts, model-level accuracy constraints predominate over selection method. Block selection reduces KV cache traffic by up to 16x at 64K, projecting to >244x at 1M context.

Figure 9: Linear scaling of memory traffic reduction factor with context length; stress-tested at $k=8$, traffic reduction exceeds 977x at one million tokens.

Photonic Hardware Integration, Practical Considerations, and Outlook

While the results rely on device-level simulation, all impairment parameters are grounded in FDTD and published device data. Integration at $d=64$, $N=1024$ requiring 65,536 MRRs is at the projected threshold for TFLN fabrication and would benefit from advanced packaging and fan-out. The add-drop MRR configuration with balanced photodetection enables true signed inner products and eliminates the need for ReLU or split encoding, doubling the photodetector count but with negligible area penalty.

Figure 10: PRISM chip layout concept (8x8 configuration); scalable to larger arrays by deeper splitter trees and increased MRR rows.

Integration form factors could include PCIe add-in cards, CXL devices, or chiplet modules. A block-index API would facilitate drop-in usage in modern LLM serving stacks.

Comparison with Related Work

Block-level top-k selection is validated by Quest, DuoAttention, and InfLLM. PRISM's contribution is the hardware mapping: it eliminates O(N) scan latency/traffic by storing signatures in the photonic weight bank, where selection is O(1), contrasting with GPU and electronic selection approaches that must scan all N signatures each decode step.

Figure 11: Time-multiplexed operation enables latency–area trade-off, keeping selection latency four orders below GPU baseline even with reduced parallel MRR count.

Conclusion

PRISM demonstrates that photonic broadcast search, particularly via broadcast-and-weight in TFLN MRR chip architectures, fundamentally breaks the O(n) memory scaling of KV cache access in LLM inference. Device-accurate modeling shows full accuracy retention for realistic hardware impairment regimes, with significant memory traffic reduction and practical energy/latency advantages projected for all context lengths of practical relevance. Photonic block selection not only addresses the memory-wall bottleneck in LLMs but generalizes to similarity-search workloads in data centers wherever single query–large stored vector tasks are pervasive.

Future research directions include fabrication of test-scale TFLN MRR arrays, end-to-end energy/latency benchmarking in GPU-integrated inference pipelines, and the application of non-volatile photonic weight storage for further energy reduction. The paradigm is poised for broad adoption as context windows and memory-bound workloads continue to grow.