cs.LGAug 3, 2026

Designing a Good Virtual Node: Addressable and Cardinality-Preserving Global Memory for Message Passing Architectures

Authors: Félix Marcoccia

Abstract

Virtual nodes give message-passing neural networks a simple global communication route, but the standard node--VN--node pipeline compresses the graph into one homogeneous state and broadcasts it identically to every node. Building on the Two-Radius analysis of Mishayev et al., we ask how auxiliary virtual memory can relieve this finite-capacity bottleneck without self-attention. We identify two requirements. First, the global memory should be factorized into independently writable and readable states: this can be achieved using addressable cross-attention slots. Second, addressability alone does not preserve multiplicity, because softmax attention is invariant to uniform replication. Inserting each slot query as a private key/value anchor recovers the discarded normalization mass and yields, on bounded color domains, an injective multiset representation able to implement a 1-WL refinement. Experiments on multiplicity-aware Two-Radius, motif counting, and constrained link-set prediction support this addressable and cardinality-preserving virtual memory at (O(nMd)) arithmetic cost.

Explore similar work

May 11, 2026cs.LG

Variational Linear Attention: Stable Associative Memory for Long-Context Transformers

Linear attention reduces the quadratic cost of softmax attention to O(T)\mathcal{O}(T), but its memory state grows as O(T)\mathcal{O}(T) in Frobenius norm, causing progressive interference between stored associations. We introduce \textbf{Variational Linear Attention} (VLA), which reframes the memory update as an online regularised least-squares problem with an adaptive penalty matrix maintained via the Sherman-Morrison rank-1 formula. We prove that normalising the write direction to unit length gives the recurrence Jacobian spectral norm exactly 11 for all sequence lengths and head dimensions (Proposition 2), and that the state norm is self-limiting under bounded inputs (Proposition 1). Empirically, VLA reduces StF\|S_t\|_F by 109×109\times relative to standard linear attention at T=1,000T{=}1{,}000, achieves near-perfect exact-match accuracy on multi-query associative recall within the effective per-head memory regime (npairs<dhn_\text{pairs} < d_h), maintaining substantially higher retrieval performance than DeltaNet and standard linear attention under increasing memory load, and maintains 62% accuracy at the per-head capacity boundary. A Triton-fused kernel achieves 14×14\times speedup over sequential Python and O(T)\mathcal{O}(T) scaling, crossing below softmax attention latency at approximately 43,000 tokens.
Vishal Pandey, Gopal Singh
Aug 13, 2026cs.AI

vToken: Token-Level Virtualization for Reclaimable KV Caches

Large language model serving faces a critical memory bottleneck: the KV cache grows with sequence length and batch size. PagedAttention uses fixed-size memory blocks to reduce allocator-level fragmentation, but recent KV eviction algorithms operate at a token granularity finer than block-level management. This mismatch causes intra-block fragmentation, leaving a large fraction of allocated KV memory unreclaimable. We present vToken, a lightweight token-level virtualization layer that decouples logical token liveness from physical block placement. vToken maintains a stable logical token view through token-table indirection and realizes physical reclamation by repacking live tokens asynchronously. The design preserves PagedAttention kernels and CUDA Graph compatibility. We implement vToken in vLLM and evaluate it with H2O, Random, and Scissorhands across models. Compared with a paired Naive-Evict baseline, vToken reduces retained KV blocks per request by 27.2%--72.3% and improves SLA-constrained throughput by up to 1.37×\times. Under a constrained active-KV budget, it extends the maximum feasible concurrency by up to 2×\times, while reducing the per-policy integration footprint from 500+ lines to under 50.
Yuanhang Gao, Xiangrui Yang, Yuanfeng Chen +4
May 21, 2026cs.LG

Asymmetric Virtual Memory Paging for Hybrid Mamba-Transformer Inference

Hybrid language models like Jamba mix attention layers with State Space Models (SSMs), creating two memory cache types with opposite profiles: Key-Value (KV) caches grow linearly with sequence length, while SSM states stay fixed per layer. Current inference engines handle this poorly. Unified pools pad SSM states to attention page sizes, wasting up to 7.3x capacity. Static dual pools cannot adapt when prompt distributions shift between requests. We present Asymmetric Virtual Memory Paging (AVMP). The allocator separates the two cache types into physically distinct pools behind a unified virtual address space, and migrates capacity between pools when one runs out. Migration triggers only on allocation failure, keeping behavior deterministic. We evaluate AVMP across 270 synthetic cells plus 60 cells of ShareGPT trace replay on an RTX 3060 12GB. Out-of-Memory events drop 7.6% and request throughput improves 1.83x to 13.3x across synthetic workloads and 2.36x on ShareGPT. All gains hold under paired-bootstrap 95% confidence intervals. A phase-time breakdown reveals two distinct mechanisms: shorter OOM recovery on capacity-pressured workloads, and faster allocation calls on KV-heavy workloads. Implementation is pure Python; Triton integration is future work.
An Xuan Nguyen