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

Hyoseok Park; Yeonsang Park

physics.opticsMar 23, 2026

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

Authors: Hyoseok Park, Yeonsang Park

Organizations: Department of Physics, Chungnam National University, Daejeon 34134, Republic of Korea

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).

Figures & tables

Figure 1: Conceptual comparison of KV cache access strategies. Left: Electronic GPU full scan—the processor sequentially reads all

N

KV blocks from HBM to compute attention, bottlenecked by memory bandwidth. Right: Prism photonic block selection—the query is broadcast optically to all

N

signature channels in parallel; only the top-

k

highest-scoring blocks are fetched from memory, reducing traffic by

N/k

times.

Figure 2: Prism system architecture (five-stage pipeline). Stage 1 (Query Encoding): The GPU/ASIC computes the query sketch

\mathbf{q}=[q_{1},\ldots,q_{d}]

and encodes each component onto a WDM wavelength via DAC-driven modulators, producing a WDM query signal where

P(\lambda_{j})=q_{j}

. Stage 2 (Broadcast): A

1\times N

optical splitter distributes identical copies of the

d

-wavelength signal to all

N

signature channels (splitting loss:

-10\log_{10}N

dB). Stage 3 (Signature Weighting): Each channel passes through a row of

d

MRRs on the TFLN photonic chip; the transmission

t_{ij}=s_{ij}

of each MRR is electro-optically programmed via DC bias electrodes to encode the block signature weight, performing wavelength-selective multiplication

P_{\text{out}}(\lambda_{j})=q_{j}\times s_{ij}

. Stage 4 (Summation): Broadband photodetectors integrate all wavelengths, yielding photocurrents

I_{i}=\mathcal{R}\cdot\sum_{j}(q_{j}\cdot s_{ij})

that are proportional to the inner product

\mathbf{q}\cdot\mathbf{s}_{i}

. Stage 5 (Top-

k

Selection): ADCs digitize the

N

photocurrents, a digital top-

k

selector identifies the

k

highest-scoring block indices, and a memory controller fetches only those KV blocks from HBM/flash storage.

Figure 3: Prism photonic chip layout for an

8\times 8

configuration (

d=8

WDM channels,

N=8

signature rows). Left: the WDM query input (

\lambda_{1}

–

\lambda_{8}

) enters and is split by cascaded

1\times 2

Y-junctions. Center: each row contains

d

MRRs coupled to a bus waveguide with coupling gap of

{\sim}

200–300 nm; EO DC bias electrodes program the MRR resonances to encode signature weights via the Pockels effect. Right: through-port and drop-port outputs route to balanced Ge-on-Si PD pairs (or optionally on-chip integrated photodetectors). Scale bar: 100

\mathrm{\SIUnitSymbolMicro m}

. The layout scales to

d=32

N=256

by increasing the splitter tree depth and the number of rows.

Figure 4: X-cut TFLN rib waveguide cross-section. The rib is etched 500 nm into a 600 nm LN film on SiO 2 , leaving a 100 nm slab. Lateral Au electrodes apply DC bias for electro-optic (Pockels) tuning of the MRR resonance wavelength. Waveguide width: 1.4

\mathrm{\SIUnitSymbolMicro m}

Parameter	Value	Notes
Platform	X-cut TFLN	600 nm LN/SiO 2
Waveguide	Rib, $1.4\!\times\!0.6\,\mu$ m	500 nm etch
MRR radius	$20\,\mu$ m	FSR $\,\approx\,$ 8.3 nm
$Q_{L}$	$\sim\!10^{4}$	FDTD: 12,500
Extinction	$>$ 15 dB	Add-drop
Wt. precision	5 bit	EO resolution

Table 1: Prism device parameters (TFLN platform).

Element	Loss (dB)	Cumulative (dBm)
Laser output	—	$+$ 20.0
Fiber–chip coupling	$-$ 2.0	$+$ 18.0
MZM modulator (avg.)	$-$ 3.0	$+$ 15.0
$1\times 256$ splitter	$-$ 25.7 a	$-$ 10.7
Waveguide (2 cm)	$-$ 1.0	$-$ 11.7
$d{=}32$ MRR (balanced, worst-case drop)	$-$ 3.2 b	$-$ 14.9

Table 2: Optical link budget for

d=32

N=256

(

P_{\text{laser}}=

20\text{,}\mathrm{d}\mathrm{B}\mathrm{m}$$ ).

Figure 5: Optical power budget analysis. (a) Per-detector received power vs. bank size

N

for three laser powers. The horizontal dashed line indicates the minimum detectable power (

-20\text{\,}\mathrm{d}\mathrm{B}\mathrm{m}

). (b) Electrical SNR at the photodetector vs. signature dimension

d

for

N=256

and

N=1024

. The shaded region marks SNR

>

20\text{,}\mathrm{d}\mathrm{B}$$ , sufficient for reliable top-

k

ranking.

Figure 6: MRR weight encoding principle. (a) Through-port and drop-port transmission of a single add-drop MRR (

Q_{L}=10{,}000

, ER

=20

dB). The balanced weight

w=T_{\text{through}}-T_{\text{drop}}

maps from

-1

(on-resonance) to

+1

(fully detuned). (b) Weight-to-balanced-transmission mapping for different DAC precisions.

Impairment	Parameter	Range
Weight quantization	$b$ (bits)	4–8
Thermal drift	$\sigma_{\text{drift}}$ ( $\mathrm{p}\mathrm{m}$ )	10–100
MRR insertion loss	IL ${}_{\text{MRR}}$ ( $\mathrm{d}\mathrm{B}$ )	0.02–0.05 c
Splitter excess loss	$\alpha_{\text{excess}}$ ( $\mathrm{d}\mathrm{B}\mathrm{/}\mathrm{s}\mathrm{t}\mathrm{a}\mathrm{g}\mathrm{e}$ )	0.1–0.3
Detector NEP	( $\mathrm{p}\mathrm{W}\mathrm{/}\sqrt{Hz}$ )	1–20
MRR crosstalk	Isolation ( $\mathrm{d}\mathrm{B}$ )	$-$ 15 to $-$ 30

Table 3: Device impairment parameter ranges used in hardware simulation.

Figure 7: Digital vs. MRR photonic inner-product comparison ( $d=32$ , $N=256$ , $K=8$ ). (a) Score correlation between exact (FP64) and MRR-computed similarity for three hardware configurations. Pearson correlation $\rho>0.98$ for all configs. (b) Rank agreement for the nominal config (5-bit, $20\text{\,}\mathrm{p}\mathrm{m}$ ): green triangles indicate correctly identified top- $K$ blocks (7/8 match, Recall@8 $=88$ %). (c) Normalised score error distributions; pessimistic config (4-bit, $30\text{\,}\mathrm{p}\mathrm{m}$ ) shows wider tails but remains zero-centred.

Figure 8: Combined impairment sensitivity: Recall@8 as a function of weight precision (bits) and thermal drift $\sigma$ ( $d=32$ , $N=500$ , 50 Monte Carlo trials per cell). Markers indicate the three operating points studied in this work: nominal (5-bit, $\sigma=0.01$ ), optimistic (6-bit, $\sigma=0.005$ ), and pessimistic (4-bit, $\sigma=0.02$ ). Recall exceeds 80% for ${\geq}5$ -bit precision and $\sigma\leq 0.02$ .

Component	Power ( $\mathrm{m}\mathrm{W}$ )	Energy/query ( $\mathrm{p}\mathrm{J}$ )
Laser source	20.0	180
TEC (thermal stab.)	1000	9000 †
Voltage driver array	5.0	45
DACs ( $d$ channels)	32.0	288
MZM modulators	6.4	58
EO bias (static)	${\sim}0$	${\sim}0$ ∗

Table 4: Prism energy breakdown per query (

d=64

N=1024

k=32

, TFLN platform).

Stage	Latency	Notes
DAC	${\sim}$ 1 ns	4-bit
MZM	${\sim}$ 0.1 ns	Si depl.
Opt. prop.	${\sim}$ 0.5 ns	5 cm
MRR decay	${\sim}$ 0.1 ns	$Q{=}10^{4}$
PD	${\sim}$ 0.2 ns	Ge
TIA+ADC	${\sim}$ 2 ns	6-bit flash

Table 5: Prism latency breakdown.

Figure 9: SNR and recall analysis. (a) Electrical SNR at the photodetector as a function of signature dimension

d

for three bank sizes. (b) Recall@8 vs. SNR showing that reliable top-

k

selection (

>90

% recall) requires SNR

\gtrsim

15\text{,}\mathrm{d}\mathrm{B}$$ .

Figure 10: Retrieval ratio $R_{h}^{(\ell,h)}$ for each KV head across all layers. Heads with $R_{h}>0.3$ (dashed line) are classified as retrieval heads. (a) Qwen2.5-7B: 102/112 heads are retrieval heads (91.1%) at 8K context. (b) Qwen3-8B: 258/288 heads are retrieval heads (89.6%) at 8K context.

Model	Context	$R_{h}(\tau{=}0.3)$ (%)	Heads	Mean $\bar{R}_{h}$
Qwen2.5-7B	2K	83.9	94/112	0.574
	4K	83.0	93/112	0.560
	8K	91.1	102/112	0.627
	16K	92.9	104/112	0.639
	32K	95.5	107/112	0.656
	65K	92.0	103/112	0.633

Table 6: Retrieval head fraction at threshold

\tau=0.3

across context lengths.

R_{h}(\tau)

: percentage of KV heads with

R_{h}>\tau

. Mean

\bar{R}_{h}

: average retrieval ratio across all heads.

Figure 11: Retrieval head statistics vs. context length. (a) Retrieval head fraction

R_{h}(\tau{=}0.3)

for Qwen2.5-7B (bf16 and 4-bit) and Qwen3-8B (bf16). The fraction exceeds 90% for

n\geq 8

K and approaches 99% at 256K context. (b) Mean retrieval ratio

\bar{R}_{h}

for Qwen2.5-7B (bf16), showing that individual-head retrieval strength also increases with context length.

Figure 12: Recall@

k

as a function of signature dimension

d

for different signature methods. Block size

B=128

k=8

(stress-test setting). Mean-key projection consistently outperforms random projection, achieving 77.3% recall@8 at

d=32

. At the primary operating point

k{=}32

, recall reaches 100% at 8K context ( table 7 ).

$n$	Blocks	R@8 (%)	R@16 (%)	R@32 (%)	NIAH (%)
4K	16	46.7	100	— ∗	100
8K	32	29.2	55.8	100	100
16K	64	26.7	41.7	57.5	100
32K	128	(OOM † )			100
64K	256	—			100

Table 7: Recall@

k

for Prism block selection across context lengths. Qwen2.5-7B,

B{=}128

d{=}32

, mean-key projection. Values averaged over 15 (layer, head) pairs.

Configuration	2K	4K	8K
Full attention	90.9	100.0	100.0
Ideal select	90.9	100.0	100.0
5-bit, 20 pm	90.9	100.0	100.0
4-bit, 30 pm	90.9	100.0	100.0
5-bit, 10 pm	90.9	100.0	100.0

Table 8: NIAH accuracy (%) with MRR-integrated block selection (Qwen2.5-7B, 11 positions,

k=8

stress-test setting).

Configuration	4K	8K	16K	32K	64K	128K †
Full attention	100	100	100	100	100	45.5
Ideal ( $k{=}32$ )	100	100	100	100	100	18.2 ‡
5-bit, 20 pm	100	100	100	100	100	27.3 ‡
4-bit, 30 pm	100	100	100	100	100	27.3 ‡
5-bit, 10 pm	100	100	100	100	100	27.3 ‡

Table 9: Extended NIAH accuracy (%) with MRR block selection (Qwen2.5-7B,

B{=}128

d{=}32

k{=}32

, 10 positions per context). At 128K the base model degrades (

\dagger

Figure 13: NIAH accuracy heatmap across context lengths (4K–128K) and MRR configurations. At $k{=}32$ , all MRR variants match full attention perfectly up to 64K. Under the $k{=}8$ stress test, accuracy degrades gracefully with context length but remains above 90% at 16K. The 128K column shows model-intrinsic degradation (full attention itself drops to 45.5%).

$d$	$N$	$N_{\text{MRR}}$	Context	Feasibility
32	256	$8192$	32K	Current TFLN
64	1024	$65\,536$	128K	Near-term
128	4096	$524\,288$	512K	Multi-chip

Table 10: MRR count for representative Prism configurations. The rightmost column indicates the approximate context length supported at block size

B=128

Figure 14: Prism photonic scaling projections. MRR count, aggregate heater power (SOI), and estimated chip area as functions of the configuration parameters

d

and

N

. The dashed horizontal lines indicate practical limits:

200\text{\,}\mathrm{W}

thermal dissipation (active cooling) and

858\text{\,}{\mathrm{mm}}^{2}

single-reticle area. Configurations below both limits (shaded region) are realizable on a single photonic chip.

Figure 15: Time-multiplexed Prism operation. (a) Physical MRR count vs. multiplexing factor

M

(

d=64

N=1024

). (b) Query evaluation latency vs.

M

, compared to electronic baselines. The shaded region indicates the sweet spot (

M=4

–

8

Figure 16: Energy crossover map for Prism vs. electronic baselines. The crossover contour (

C_{\text{PRISM}}/C_{\text{baseline}}=1

) shifts to shorter context lengths as

d

decreases. (a) vs. GPU full scan: practical benefit at

n^{*}\approx 4

K. (b) vs. GPU ANN:

n^{*}\approx 2

K. (c) vs. NVIDIA ICMS:

n^{*}\approx 4

Figure 17: Energy advantage vs. context length. The ratio decreases as

{\sim}1/n

; at

1\,000\,000

tokens (

B=128

k{=}32

), the traffic reduction reaches

244\times

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure S1: Impact of weight quantization on recall. At

b=6

bits, recall degrades by less than 5% from the floating-point ideal (Recall@8 = 0.960 at 6-bit). Adding thermal drift (

\sigma_{\text{th}}=0.01

) and detector noise (

\sigma_{\text{det}}=0.01

) degrades recall by an additional 5%.

Figure S2: Recall degradation as a function of thermal drift

\sigma

. Recall remains above 95% for

\sigma\leq 0.005

(corresponding to

\sim

150\text{\,}\mathrm{p}\mathrm{m}

drift), achievable with standard thermal stabilization. At

\sigma=0.01

(

\sim

300\text{\,}\mathrm{p}\mathrm{m}

), recall is still 94.8%.

Figure S3: Weight encoding fidelity under MRR impairments ( $d=32$ , $N=128$ ). (a) Ideal vs. 5-bit quantised weights in $[0,1]$ : the staircase pattern shows 32 discrete levels with RMSE $=0.009$ . (b) Adding $20\text{\,}\mathrm{p}\mathrm{m}$ thermal drift and fabrication variation broadens the scatter (RMSE $=0.014$ ). (c) Error histograms for three configurations: 5-bit only, 5-bit with $20\text{\,}\mathrm{p}\mathrm{m}$ drift, and 4-bit with $30\text{\,}\mathrm{p}\mathrm{m}$ drift. Even the pessimistic case concentrates errors within $\pm 5$ % of the full weight range.

Figure S4: Recall@8 degradation as a function of photodetector noise

\sigma_{\text{det}}

(

d=32

N=500

, 100 trials). Recall remains above the 95% threshold for

\sigma_{\text{det}}\leq 0.01

, achievable with standard Ge-on-Si photodetectors (NEP

\sim

1\text{\,}\mathrm{p}\mathrm{W}\mathrm{/}\sqrt{Hz}

Figure S5: Signed vs. unsigned recall comparison. Balanced photodetection (signed

[-1,+1]

) consistently outperforms ReLU projection (unsigned

[0,1]

) and split encoding across all signature dimensions

d

Figure S6: KV cache traffic reduction factor

N/k

as a function of context length for different

k

values.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving
Aug 20, 2026Qihang Fan, Huaibo Huang, Zhiying Wu +2LLM Inference Acceleration Long-Context Language Model Inference
MAIS&NLPR, CASIA · UCAS · WeChat, Tencent
FlashMemory-DeepSeek-V4: Lightning Index Ultra-Long Context via Lookahead Sparse Attention
Jun 8, 2026Yan Wang, Qifan Zhang, Jiachen Yu +12KV Caching Long-Context Language Model Inference
Independent Researchers · Tencent · The Hong Kong University of Science and Technology (Guangzhou) +1
Faster Than Flash: Exploiting Attention Sparsity for Efficient Long-Context Decoding
Aug 31, 2026Zhigeng Liu, Zhiyuan Ning, Ruixiao Li +5Long-Context Language Model Inference LLM Inference Acceleration
Fudan University, Shanghai, China · Shanghai Innovation Institute, Shanghai, China · Harbin Institute of Technology, Harbin, China

Aug 20, 2026J/K move · Enter open · S save

FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving

Qihang Fan, Huaibo Huang, Zhiying Wu, Bingning Wang+1

MAIS&NLPR, CASIA · UCAS · WeChat, Tencent

LLM Inference Acceleration Long-Context Language Model Inference Block-Sparse Attention

Open paper AskPDF arXiv

Long-context modeling is a pivotal capability for Large Language Models, yet the quadratic complexity of attention remains a critical bottleneck, particularly during the compute-intensive prefilling phase. Our previous work, FlashPrefill, mitigates this cost through instantaneous pattern discovery and max-based dynamic thresholding; however, it remains an algorithmic prototype that is still distant from production deployment. In this paper, we present FlashPrefill V2, which evolves FlashPrefill from a prototype toward practical long-context serving along three dimensions. First, we introduce a mean correction term that effectively suppresses the approximation error, keeping performance degradation manageable even at extreme sparsity levels. Second, we redesign the sparse attention operator with PackGQA memory access, warp specialization, and pingpong pipelining, fully aligning with the latest FlashAttention-3/4 implementations and supporting FP8 inference to meet practical quantization requirements. Third, FlashPrefill V2 natively supports paged KV cache and continuous batching, allowing integration as an attention backend in modern inference frameworks such as SGLang. Extensive evaluations on NVIDIA H20 GPUs---among the most widely deployed inference accelerators---demonstrate that FlashPrefill V2 delivers up to 47.26x and 27.19x speedups over FlashAttention-2 at 128K context length under FP8 and BF16 precision, respectively, and, in FP8, still achieves a 30.49x speedup against an FA3/4-aligned dense baseline.

Abstract

LLM Inference Acceleration

Explore similar work

FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving

FlashMemory-DeepSeek-V4: Lightning Index Ultra-Long Context via Lookahead Sparse Attention

Faster Than Flash: Exploiting Attention Sparsity for Efficient Long-Context Decoding