refactor: extract shared virtual memory primitives.

AGENTS.md

@@ -1,36 +1,5 @@
 # xLLM Coding Agent Instructions
 
-## Directory Structure
-
-```
-├── xllm/
-|   : main source folder
-│   ├── api_service/               # code for api services
-│   ├── c_api/                     # code for c api
-│   ├── cc_api/                    # code for cc api 
-│   ├── core/  
-│   │   : xllm core features folder
-│   │   ├── common/                
-│   │   ├── distributed_runtime/   # code for distributed and pd serving
-│   │   ├── framework/             # code for execution orchestration
-│   │   ├── kernels/               # adaption for npu kernels adaption
-│   │   ├── layers/                # model layers impl
-│   │   ├── platform/              # adaption for various platform
-│   │   ├── runtime/               # code for worker and executor
-│   │   ├── scheduler/             # code for batch and pd scheduler
-│   │   └── util/
-│   ├── function_call              # code for tool call parser
-│   ├── models/                    # models impl
-│   ├── parser/                    # parser reasoning
-│   ├── processors/                # code for vlm pre-processing
-│   ├── proto/                     # communication protocol
-│   ├── pybind/                    # code for python bind
-|   └── server/                    # xLLM server
-├── examples/                      # examples of calling xLLM
-├── tools/                         # code for npu time generations
-└── xllm.cpp                       # entrypoint of xLLM
-```
-
 ## Code Style Guide
 
 * Before editing, creating, refactoring, or reviewing any file under `xllm/`, you **MUST** read [custom-code-style.md](.agents/skills/code-review/references/custom-code-style.md).

docs/src/content/docs/en/cli_reference.md

@@ -54,7 +54,7 @@ xLLM uses gflags to manage service startup parameters. `--model <PATH>` is the o
 
 | Parameter | Type | Default | Description |
 |:----------|:-----|:--------|:------------|
-| `enable_manual_loader` | `bool` | `false` | Pin decoder layer weights to host memory and use async H2D transfer. Required by `enable_rolling_load`; also implied by `enable_xtensor`. |
+| `enable_manual_loader` | `bool` | `false` | Pin decoder layer weights to host memory and use async H2D transfer. Required by `enable_rolling_load`; also implied by `enable_virtual_memory`. |
 | `enable_rolling_load` | `bool` | `false` | Enable rolling weight load: keep only N decoder layer weight slots in HBM and stream-load each layer just in time. Requires `enable_manual_loader=true`. NPU only. |
 | `rolling_load_num_cached_layers` | `int32` | `2` | Number of decoder layer weight slots to keep in HBM when `enable_rolling_load=true`. |
 | `rolling_load_num_rolling_slots` | `int32` | `-1` | Number of rolling slots used by decoder rolling load. Fixed slots are `rolling_load_num_cached_layers - rolling_load_num_rolling_slots`. `-1` means auto, `min(2, preload_count)`. Must be in `[-1, rolling_load_num_cached_layers]`. |
@@ -73,7 +73,7 @@ xLLM uses gflags to manage service startup parameters. `--model <PATH>` is the o
 | `enable_in_batch_prefix_cache` | `bool` | `false` | Whether to cache admitted prefill full blocks into the prefix cache so that later requests in the same batch can share them. |
 | `max_linear_state_cache_slots` | `int64` | `0` | Maximum number of active linear-attention state cache slots. `0` derives an automatic capacity from the available KV Cache budget. |
 | `xxh3_128bits_seed` | `uint32` | `1024` | Default XXH3 128-bit hash seed. |
-| `enable_xtensor` | `bool` | `false` | Whether to enable XTensor for model weights with the physical page pool. |
+| `enable_virtual_memory` | `bool` | `false` | Enable virtual memory for model weights and KV Cache with the physical page pool. |
 | `phy_page_granularity_size` | `int64` | `2097152` | Granularity size of one physical page in bytes, default 2 MiB, for continuous KV Cache. |
 
 ## KVCacheStoreConfig
@@ -155,7 +155,7 @@ xLLM uses gflags to manage service startup parameters. `--model <PATH>` is the o
 | Parameter | Type | Default | Description |
 |:----------|:-----|:--------|:------------|
 | `master_node_addr` | `string` | `"127.0.0.1:19888"` | Master address for multi-node distributed serving, for example `10.18.1.1:9999`. |
-| `xtensor_master_node_addr` | `string` | `"127.0.0.1:19889"` | Master address for the XTensor distributed service, for example `10.18.1.1:9999`. |
+| `worker_memory_master_node_addr` | `string` | `"127.0.0.1:19889"` | Rank 0 address where worker memory RPC servers register their endpoints, for example `10.18.1.1:9999`. Used only when virtual memory is enabled and `nnodes > 1`. |
 | `nnodes` | `int32` | `1` | Number of multi-node nodes. |
 | `node_rank` | `int32` | `0` | Rank of the current node. |
 | `etcd_addr` | `string` | `""` | etcd address used to save instance metadata. |

docs/src/content/docs/en/dev_guide/code_arch.md

@@ -32,7 +32,7 @@ same basic flow; you do not need to understand them to locate the main path.
 | `Batch`, `BatchGroup` | A `Batch` contains work selected for one execution step. A `BatchGroup` holds the batches for the data-parallel (DP) ranks in that step. | [`batch.h`](https://github.com/xLLM-AI/xllm/blob/main/xllm/core/framework/batch/batch.h), [`batch_group.h`](https://github.com/xLLM-AI/xllm/blob/main/xllm/core/framework/batch/batch_group.h) |
 | Rank and parallelism | A rank identifies a participant in parallel execution. DP distributes request work across groups; tensor parallelism (TP) splits model tensor computation across ranks; expert parallelism (EP) distributes mixture-of-experts model experts. | [`parallel_state/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/framework/parallel_state) |
 | Continuous batching | The scheduler selects work again on each step, allowing new requests to join while completed requests leave. A batch is not fixed for the full lifetime of its requests. | [`continuous_scheduler.h`](https://github.com/xLLM-AI/xllm/blob/main/xllm/core/scheduler/continuous_scheduler.h) |
-| KV cache, block, prefix cache | The KV cache stores attention state for reuse. Block managers allocate and reclaim cache capacity; prefix caching allows compatible requests to reuse cached prompt prefixes. | [`kv_cache/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/framework/kv_cache), [`block/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/framework/block), [`prefix_cache/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/framework/prefix_cache) |
+| KV cache, block, prefix cache | The KV cache stores attention state for reuse. Block managers allocate and reclaim cache capacity; prefix caching allows compatible requests to reuse cached prompt prefixes. | [`kv_cache/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/kv_cache), [`block/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/framework/block), [`prefix_cache/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/kv_cache/prefix) |
 | Logits and sampling | Logits are the model's scores for vocabulary tokens. Sampling applies generation settings to choose the next token from those scores. | [`sampling/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/framework/sampling) |
 | Master, engine, worker, executor | The master connects request handling and scheduling; the engine coordinates batch execution across workers; a worker manages device-side execution; its executor selects how the model runs. | [`distributed_runtime/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/distributed_runtime), [`runtime/`](https://github.com/xLLM-AI/xllm/tree/main/xllm/core/runtime) |
 
@@ -56,19 +56,22 @@ entry is **`xllm/xllm.cpp`**, inside the source directory.
 │   │   │   ├── config/                    # Configuration categories and CLI options
 │   │   │   ├── request/                   # Request, sequence, and stopping state
 │   │   │   ├── batch/                     # Batches and model input builders
-│   │   │   ├── block/                     # Cache allocation and block management
-│   │   │   ├── kv_cache/                  # Cache tensors, layouts, and capacity estimation
-│   │   │   ├── prefix_cache/              # Prompt-prefix reuse
-│   │   │   ├── kv_cache_transfer/         # Cache transfer and external cache storage
+│   │   │   ├── block/                     # Composite cache and runtime allocation adapters
+│   │   │   ├── kv_cache/                  # Model-aware linear-state restore adapter
 │   │   │   ├── model/                     # Model interfaces and input/output types
 │   │   │   ├── model_loader/              # Model configuration and checkpoint loading
+│   │   │   │   └── weight/                # Weight resources, allocation policy and transfer
 │   │   │   ├── tokenizer/                 # Text and token conversion
 │   │   │   ├── chat_template/             # Chat messages to model prompts
 │   │   │   ├── sampling/                  # Token sampling and constrained decoding
 │   │   │   ├── parallel_state/            # Parallel groups and communication state
 │   │   │   ├── eplb/                      # Expert parallel load balancing
-│   │   │   └── xtensor/                   # Virtual memory and model memory management
-│   │   ├── runtime/                       # Workers and model executors
+│   │   │   ├── transfer/                  # Shared byte transport, sessions and registration
+│   │   │   └── allocator/                 # Generic memory facilities and page budget accounting
+│   │   │       ├── virtual_memory/        # Generic physical pages and virtual mappings
+│   │   │       └── torch/                 # Tensor views over mapped byte regions
+│   │   ├── kv_cache/                      # Cache domain: layout, storage, blocks, prefixes, transfer
+│   │   ├── runtime/                       # Workers, local resources and model executors
 │   │   ├── layers/                        # C++ neural-network layers
 │   │   ├── kernels/                       # Device operator implementations and wrappers
 │   │   ├── platform/                      # Devices, streams, memory, and communication APIs
@@ -106,6 +109,172 @@ entry is **`xllm/xllm.cpp`**, inside the source directory.
 └── setup.py                               # Python packaging and build entry
 ```
 
+## KV cache and virtual memory boundaries
+
+The memory facilities serve both KV cache and weights. Resource-specific
+mapping and allocation policy stay with each resource. Runtime owns local
+resources; distributed coordination manages shared budgets, worker operations,
+and model sleep/wakeup. In the diagram, a solid arrow means "uses or owns";
+a dashed arrow marks interface implementation or a coupling that remains to
+be separated.
+
+```mermaid
+flowchart TB
+  subgraph runtime["runtime"]
+    worker["WorkerImpl<br/>composition and injection"]
+    resources["WorkerMemoryResources<br/>local device and resource lifecycle"]
+  end
+
+  blocks["framework/block<br/>composite managers and paged adapter"]
+
+  subgraph distributed["distributed_runtime"]
+    residency["ModelResidencyCoordinator<br/>model initialization, sleep and wakeup"]
+    coordination["DistributedVirtualMemoryCoordinator<br/>shared KV/weight budget, worker selection and dispatch"]
+    rpc["WorkerMemoryRpc<br/>client, service and server"]
+  end
+
+  subgraph cache["core/kv_cache"]
+    policy["layout, block and prefix<br/>KV policy and state"]
+    pages["storage / VirtualKVCachePageAllocatorImpl<br/>page queues, preallocation and block-to-page conversion"]
+    backend["storage / KVCachePageBackend<br/>capacity reservation and mapping interface"]
+    ports["storage and transfer interfaces<br/>page allocator, tensor allocator, memory provider"]
+    storage["storage / VirtualKVCacheMemoryRegistry<br/>local KV mappings and paged tensor allocation"]
+    kvtransfer["transfer<br/>memory provider, KV peer/layout state<br/>resharding and cache transfer"]
+  end
+
+  subgraph framework["framework"]
+    weight["model_loader/weight<br/>WeightVirtualMemoryManager, store<br/>allocation policy and weight transfer"]
+    tensor["allocator/torch<br/>mapped-memory tensor views"]
+    shared["allocator<br/>WorkerPhysicalPageBudget / GlobalVirtualMemoryRegion"]
+    memory["allocator/virtual_memory<br/>physical pages and byte mappings"]
+    transport["transfer<br/>Mooncake byte transport<br/>sessions and memory registration"]
+  end
+
+  remaining["framework/request / Sequence<br/>framework/config and runtime options"]
+  worker --> resources
+  worker --> blocks
+  worker --> storage
+  worker --> kvtransfer
+  residency --> coordination
+  residency -->|worker lifecycle operations| worker
+  coordination --> pages
+  coordination --> rpc
+  coordination -->|local execution| resources
+  coordination --> shared
+  rpc --> resources
+  resources --> storage
+  resources --> weight
+  resources --> shared
+  blocks --> policy
+  blocks --> ports
+  pages -.->|implements page allocator| ports
+  pages --> backend
+  coordination -.->|implements mapping backend| backend
+  storage -.->|implements tensor allocator| ports
+  kvtransfer -.->|implements memory provider| ports
+  storage --> tensor
+  tensor --> memory
+  shared --> memory
+  weight --> memory
+  weight --> shared
+  weight --> transport
+  kvtransfer --> storage
+  kvtransfer --> shared
+  kvtransfer --> transport
+  policy -.-> remaining
+```
+
+`framework/allocator/virtual_memory` provides `PhysicalPage`, `PhysicalPagePool`,
+`MappedVirtualMemoryRegion` and `SharedPhysicalPageMapping`. A mapped region
+manages bytes, physical pages and virtual addresses. Tensor shape, dtype and view
+creation live in `framework/allocator/torch`; the returned tensor view does not own the
+region. Generic memory code has no KV layout, weight policy or model lifecycle
+responsibilities. `GlobalVirtualMemoryRegion` is a shared mapping facade with lifetime
+leases. Transport registration and its retained mapping lease belong to
+`framework/transfer`, which has no KV or distributed-runtime target dependency.
+
+`core/kv_cache/storage` owns `VirtualKVCacheMemoryRegions` and
+`VirtualKVCacheMemoryRegistry`, including model-local cache mappings and physical
+offset lookup. `VirtualKVCacheTensorAllocator` uses an injected registry to
+implement `KVCacheTensorAllocator`. In `core/kv_cache/transfer`,
+`VirtualKVCacheTransferMemoryProviderImpl` uses that registry and an injected
+`GlobalVirtualMemoryRegion` to implement `VirtualKVCacheTransferMemoryProvider`.
+`WorkerImpl` supplies these dependencies. Concrete cache backends consume
+generic allocation facilities without depending on distributed coordination.
+
+`framework/model_loader/weight` owns `WeightVirtualMemoryManager`, `ModelWeightStore`,
+`WeightAllocation` and the weight transfer adapter. It manages weight
+reservation ownership, contiguous or fragmented mappings, suballocation
+and cleanup using a physical-page pool and shared mapping supplied at
+construction. Both weight transfer and KV transfer use the shared Mooncake
+byte transport. KV manifests, peer state and resharding plans remain in
+`core/kv_cache/transfer`, including `MooncakeKVCacheTransferEngine`.
+
+`runtime/WorkerMemoryResources` owns the local KV registry and
+`WeightVirtualMemoryManager`, managing device initialization, the physical page pool,
+and local resource creation and cleanup. Its separate build target does not
+depend on the `runtime` aggregate, RPC or distributed-runtime targets.
+`distributed_runtime/WorkerMemoryRpcClient`, `WorkerMemoryRpcService` and
+`WorkerMemoryRpcServer` handle requests between workers. The service invokes
+local resources without calling back into distributed coordination.
+
+`core/kv_cache/storage/VirtualKVCachePageAllocatorImpl` owns KV page queues,
+preallocation and block-to-virtual-page conversion, implementing
+`VirtualKVCachePageAllocator`. An injected `KVCachePageBackend` supplies capacity
+reservation and map/unmap operations; the allocator has no weight policy or
+RPC code. Map/unmap receive byte offsets within each layer's K/V region,
+computed by the allocator as virtual page ID times page size.
+`KVCacheManagerFactory` injects the page allocator interface through
+pools and composite managers into `VirtualMemoryKVCacheBlockManager`, whose block
+adapter does not access a distributed singleton.
+
+`DistributedVirtualMemoryCoordinator` implements that backend, owns the paged
+allocator, accounts for capacity through generic `WorkerPhysicalPageBudget`, and
+selects workers to dispatch local or RPC operations. `ModelResidencyCoordinator`
+organizes model initialization, sleep and wakeup. Distributed coordination
+atomically reserves the combined KV and weight budget in one transaction
+during wakeup. On initial awake startup, weights reserve capacity first;
+later KV page allocations debit the same shared budget. Suspension
+first freezes KV allocation and recycling and drains preallocation and
+in-flight mappings before changing mappings and the weight lifecycle. This
+prevents capacity release from racing with background mapping work. KV cache
+owns page policy; distributed-runtime owns cross-resource transactions and
+execution across workers.
+
+Block value types form `kv_cache_block_types`, used by request and prefix
+cache. Cache domain leaves are compiled by `kv_cache_block`; framework
+composition and allocation adapters form the `block` aggregate target.
+This is an incremental boundary change: cache block policies still operate
+on `Sequence`, capacity estimation still reads framework configuration and
+runtime options, and model-specific cache policies remain. Those policy and
+input boundaries need further slices before `core/kv_cache` can be treated as
+a fully independent domain module.
+
+The protobuf schema `proto/mooncake_transfer_engine.proto` remains a
+compatibility bridge: it retains the existing service and method names on
+one listener. The generic transport serves session RPCs and forwards cache
+RPCs to a KV extension service without interpreting their messages. The
+schema still contains KV messages, so this boundary removes C++ domain and
+target dependencies rather than completing protocol separation.
+
+The options are `enable_virtual_memory` and `worker_memory_master_node_addr`.
+When virtual memory is enabled with `nnodes > 1`, rank 0 listens on
+`worker_memory_master_node_addr` to collect worker memory RPC server endpoints.
+Each worker registers its endpoint there, and rank 0 connects to those endpoints
+to dispatch memory operations. This registration service is separate from the
+execution worker registration service configured by `master_node_addr`.
+Explicit CLI settings override JSON configuration.
+
+Splitting the local resource, distributed coordinator and RPC C++ classes
+preserves protobuf field numbers. `proto/model_memory_dist.proto` uses the
+`ModelMemoryDist` service and `GetVirtualKVCacheOffsets` method. The method serves
+the virtual-memory KV backend and returns byte offsets in the shared physical-page
+pool mapping, rather than offsets in model-local tensor views. The heartbeat JSON
+field is `virtual_memory_info`.
+These protocol names differ from the former xtensor version. Upgrade masters,
+workers and external heartbeat consumers together; RPC routes from before
+and after that rename cannot be mixed.
+
 The most useful distinction is between **state**, **scheduling**, and
 **execution**: `framework/` defines requests, batches, caches, and model
 interfaces; `scheduler/` decides what runs next; `distributed_runtime/` and