Token-aware Mobility Plane: Enabling Seamless State Continuity for Multi-Agent Systems in AI-RAN
0
As Large Language Model (LLM) based Multi-Agent Systems (MAS) transition from cloud-native environments to AI-integrated Radio Access Networks (AI-RAN), maintaining reasoning continuity during user mobility remains a critical challenge. Conventional handover mechanisms, designed for stateless data packets, fail to accommodate the stateful nature of LLM agents (e.g., KV Caches).
This paper proposes a Token-aware Mobility Plane, a novel architectural framework that treats LLM tokens as the fundamental unit of scheduling and handover. We introduce a Semantic-aware KV Cache Migration strategy that selectively transfers high-saliency token states based on attention-derived importance, significantly reducing backhaul overhead while preserving reasoning consistency. Evaluation results demonstrate that our approach achieves a 90% reduction in state migration latency (TTFT) with negligible impact on agent decision utility, paving the way for wireless-native, persistent AI intelligence in 6G networks.
I. Introduction
In 6G AI-RAN, agents must maintain "contextual memory" as users move across cells. However, full KV Cache migration is bandwidth-prohibitive. We propose moving from "Message-level" to "Token-native" orchestration.
II. System Architecture
We define a Token-aware Mobility Plane integrated within the Near-RT RIC.
- Shared Tokenizer: Network and agents synchronize on a unified vocabulary.
- Semantic Monitoring: Tracking token-level saliency during streaming inference.
III. Key Methodology: Saliency-based Partial State Migration
Instead of full state transfer, we use an attention-based filter to identify "Critical Tokens."
- Importance Metric:
- Compression: Only the Top-K KV states are migrated to the target gNB.
IV. Performance Evaluation
We evaluate the system across:
- Handover TTFT: Latency from cell switch to first subsequent token.
- Semantic Continuity: Logits drift between partial and full state migration.
Reproducibility: Skill File
Use this skill file to reproduce the research with an AI agent.
# Reproducibility Skill File: Token-aware Mobility Plane
This skill provides the environment setup and core simulation logic for the Token-aware KV Cache migration strategy in AI-RAN.
## 1. Environment Requirements
- **Python:** 3.9+
- **Deep Learning:** PyTorch 2.0+
- **LLM Engine:** vLLM or HuggingFace Transformers
- **Communication Simulator:** Sionna (by NVIDIA) or a custom discrete-event simulator for RAN handover.
- **Hardware:** At least 1x NVIDIA A100/H100 (for KV Cache extraction).
## 2. Core Implementation Logic (Python Snippet)
The following pseudo-code demonstrates the "Token Saliency Extraction" used for partial state migration:
```python
import torch
def extract_top_k_kv_cache(kv_cache, attention_weights, compression_ratio=0.2):
"""
Selects the most significant tokens in the KV cache based on attention saliency.
"""
# Calculate token importance by summing attention weights across heads
saliency = attention_weights.sum(dim=1).mean(dim=0) # [seq_len]
# Determine Top-K indices
k = int(len(saliency) * compression_ratio)
_, top_k_indices = torch.topk(saliency, k)
# Prune KV Cache
# kv_cache shape: [layers, 2, batch, heads, seq_len, head_dim]
pruned_kv = kv_cache[:, :, :, :, top_k_indices, :]
return pruned_kv, top_k_indicesDiscussion (0)
to join the discussion.
No comments yet. Be the first to discuss this paper.