@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 257 lines 12.7 KB
Rendered Raw
1 # HiSparse local KV offload architecture
2
3 Status: experimental
4
5 ## The short version
6
7 There are three different jobs:
8
9 1. The normal KV cache system manages GPU block pools and tables.
10 2. `HiSparseCoordinator` manages logical host blocks, source-prefix identity, and
11 host/GPU residency transitions.
12 3. `HiSparseConnector` carries residency work between scheduler and worker;
13 `HiSparseWorker` coordinates transfers, while per-cache
14 `HiSparseRuntime` objects own host/hot views and GPU replacement state.
15
16 Neither worker-side object allocates or frees logical blocks. HMA provides the
17 GPU allocation shared by the resident and hot groups; it does not manage CPU
18 memory or KV identity.
19
20 ```text
21 request ────► HiSparseCoordinator ── source blocks + residency policy
22 │
23 ├──► KV cache manager ── resident/hot GPU leases (HMA)
24 │
25 └──► HiSparseConnector ── host bytes + copies + GPU LRU
26 ```
27
28 This is a local KV connector. When P/D or another offload connector is also
29 configured, `MultiConnector` composes it with `HiSparseConnector`.
30
31 `host_pool_gib` is configured on `HiSparseConnector` and is the usable
32 host-cache capacity per data-parallel replica, not a node-wide memory budget.
33 Tensor-parallel ranks hold replicated views of that logical cache. Those views
34 may use private per-rank backing or one shared
35 physical allocation without changing the configured capacity. Physical host
36 memory consumption is therefore topology- and implementation-dependent. The
37 realized capacity may be slightly smaller because the budget is rounded down
38 to complete host blocks.
39
40 ## Ownership
41
42 | Thing | Owner | What “owner” means |
43 | --- | --- | --- |
44 | HiSparse source and prefix identity | `HiSparseCoordinator` | maps tokens to logical host blocks |
45 | Resident GPU block leases | normal KV cache manager | allocates and frees HMA blocks |
46 | Resident block tables | normal KV cache manager | tells attention where resident pages are |
47 | Residency transitions | `HiSparseCoordinator` | plans spill-before-free transactions |
48 | Logical host block allocation | `HiSparseCoordinator` | owns the separate CPU block pool and its lifecycle |
49 | Pinned host-pool lifecycle | `HiSparseWorker` | worker-wide backing and teardown |
50 | Per-cache host view and hot contents | `HiSparseRuntime` | binds host/hot storage and fills cache-manager-provided hot leases |
51 | Hot row map and LRU | `HiSparseRuntime` | resolves hits and chooses victims on GPU |
52 | Resident-cache route | `HiSparseCacheHandle` | exposes resident or host/hot resolution to attention |
53 | Sparse attention | attention backend | consumes a device cache and physical row IDs |
54 | HMA | allocator | provides GPU capacity; owns no KV meaning |
55
56 The key distinction is logical allocation versus contents. `HiSparseCoordinator`
57 owns host block IDs and their request/prefix associations. `HiSparseWorker` and
58 its per-cache runtimes own the corresponding bytes. The normal cache manager
59 sees only device pools.
60 The source group has `block_pool_id=None`; device-pool consumers must narrow it
61 before indexing, so host ownership cannot masquerade as a numeric GPU pool.
62
63 For single-node MP tensor parallelism, every TP worker maps the same pinned host
64 pool and uses the same block and layer offsets. MLA source KV is replicated
65 across TP ranks, so this stores one physical copy instead of one copy per rank.
66 TP rank 0 writes the shared host pool; peers wait on its IPC events before
67 reading it. Other executor and parallel layouts retain private per-rank pools.
68
69 The shared layout backs the per-replica logical capacity with one physical pool;
70 the private layout allocates one physical pool per rank. Physical pool size
71 includes block-stride alignment.
72
73 ## Code boundary
74
75 ```text
76 scheduler process worker process
77
78 HiSparseConnector HiSparseConnector
79 └─ HiSparseCoordinator └─ HiSparseWorker
80 │ │
81 │ connector metadata ├─ host bytes
82 │ - page transfers ├─ copy scheduling
83 │ - block-table replacements └─ per-layer hot state
84 └───────────────────────────────────────────────►│
85 ◄──────── connector worker metadata ─────────────┘
86 enqueued and completed transfer IDs
87 ```
88
89 The command travels in `kv_connector_metadata`; transfer updates return in
90 `KVConnectorOutput.kv_connector_worker_meta`. The model runner does not
91 interpret page transfers. Enqueue acknowledgements let the scheduler release
92 source leases in stream order; completion acknowledgements publish the copied
93 host pages.
94
95 ## Resident device pages
96
97 Resident pages are intentionally outside `HiSparseRuntime`.
98
99 KV-cache initialization binds cache-manager allocations to the attention-facing
100 `HiSparseCacheHandle` before constructing `HiSparseWorker`. That same
101 handle's runtime retains the resident source index needed by a transfer plan.
102 There is no second resident object or registration wrapper.
103
104 ```text
105 KV cache setup
106 │
107 ├─ bind resident allocation ──► HiSparseCacheHandle
108 │ cache + block table + slot mapping
109 │
110 ├─ bind host/hot allocation ──► HiSparseRuntime
111 │ host + hot + GPU LRU
112 │
113 └─ register cache handles ────► HiSparseWorker
114 step-level transfers
115
116 HiSparseWorker registers the same HiSparseCacheHandle objects directly
117 ```
118
119 Attention construction links each layer to the most recent layer that actually
120 owns an indexer. This releases a follower's duplicate LRU tensors before GPU
121 memory profiling. Cache binding only attaches storage; it does not infer
122 semantic groups from the physical packed-tensor order. The construction cursor
123 is discarded with the worker's pinned state.
124
125 Every HiSparse decode batch uses the same fused resolver. It checks resident
126 pages first, then hot rows, then pinned host memory. A resident hit exits inside
127 the kernel before hot-LRU lookup or host copying; there is no framework-level
128 residency route or separate CUDA graph. No CPU decision is added to the decode
129 path. The resolver consumes the existing graph-stable request mapping from
130 attention metadata; neither the worker nor individual cache handles keep a
131 duplicate mapping.
132
133 Speculative decoding resolves all verification rows of a request in one pass:
134 one block resolves the union of the rows' top-k against the request's hot-cache
135 state, so rows that select the same host row share its hot row and no row
136 evicts a hot row another row of the step still reads. Draft layers write their
137 rows after the target forward, so their host mirror runs at the start of the
138 next step, after the drafter, and the step's page transfers are submitted
139 behind it.
140
141 ## P/D import target
142
143 The decoder chooses the landing target once per request from the normal cache
144 admission calculation. If the complete imported prefix fits the device pools,
145 NIXL transfers it directly into resident GPU pages. Otherwise, if the fixed
146 host-backed GPU footprint and host source blocks fit, the request imports into
147 the host tier. There is no context-length threshold or other heuristic, and a
148 request waiting for capacity retains its choice across admission retries.
149
150 A host import reads through a bounded decoder-GPU staging pool before copying
151 into registered host memory. Pages needed immediately are mirrored into their
152 resident destinations during that copy. Both landing targets then use the same
153 fused decode resolver described above.
154
155 ## Indexer KV offloading
156
157 HiSparse does not keep a private CPU copy of indexer KV. The indexer remains a
158 normal prefix-cacheable GPU cache group. If `OffloadingConnector` is configured
159 with HiSparse, it stores and restores that group through the generic KV
160 offloading path; HiSparse continues to own only the sparse MLA host tier.
161
162 The two prefix sources can have different hit lengths. When the HiSparse host
163 prefix extends beyond the GPU-resident indexer prefix, the scheduler asks
164 `OffloadingConnector` to restore only the missing indexer suffix, capped at the
165 host prefix boundary. If that suffix is unavailable, all groups fall back to
166 the shorter prefix they share. NIXL P/D transfers continue to place indexer KV
167 directly in its GPU group.
168
169 ## Spill transaction
170
171 A resident block cannot be reused until its contents have been handed to the
172 worker.
173
174 ```text
175 HiSparseCoordinator HiSparseWorker
176 │ │
177 │ pin source and destination leases │
178 │── SparseKVPageTransfer ─────────────►│
179 │ │ enqueue GPU-to-host copy
180 │◄── enqueued transfer ID ────────────│
181 │ replace resident table entry │
182 │ release resident lease to HMA │
183 │ │ copy reaches its event
184 │◄── completed transfer ID ───────────│
185 │ mark host page valid │
186 │ release destination host lease │
187 ```
188
189 “Enqueued” means the copy has entered the worker stream. Stream ordering makes
190 it safe to reuse the resident GPU block for later work, but the host page is
191 not yet published. “Completed” means the worker has observed the copy's event;
192 only then does the coordinator publish the host page for prefix reuse and
193 release its destination lease. A host-write event separately protects direct
194 CPU readers from writes already queued on the accelerator.
195
196 The worker transfer contains only its transfer ID and physical copy
197 coordinates. Request identity and logical page state remain in the scheduler.
198
199 ## Hot lookup and LRU
200
201 The NVIDIA CUDA path keeps replacement entirely on the accelerator:
202
203 ```text
204 top-K logical positions
205 │
206 ▼
207 resident page? ── yes ──► resident physical row
208 │ no
209 ▼
210 hot row? ──────── yes ──► existing hot physical row + update GPU LRU
211 │ no
212 ▼
213 choose GPU LRU victim ──► copy pinned host row ──► hot physical row
214 ```
215
216 ROCm is not currently supported because the fused HiSparse cache operations are
217 implemented only by CUDA kernels. A future platform-specific worker may provide
218 the same command, output, and cache-resolution boundaries.
219
220 ## Main classes
221
222 | Class | Inherits / implements | Responsibility |
223 | --- | --- | --- |
224 | `HiSparseCoordinator` | plain scheduler component | host allocation, source prefixes, resident leases, and spill state machine |
225 | `HiSparseConnector` | `KVConnectorBase_V1`, `SupportsHMA` | scheduler/worker metadata and lifecycle boundary |
226 | `HiSparseResidentManager` | `SingleTypeKVCacheManager` | normal block-pool bookkeeping with host-backed holes |
227 | `PagedCacheView` | immutable data object | shared resident/hot HMA tensor binding |
228 | `HiSparseWorker` | connector-owned worker component | worker-wide transfer scheduling and host-pool lifecycle |
229 | `HiSparseRuntime` | plain worker-owned component | per-cache host/hot tensors, GPU LRU, and fused resolution |
230 | `HiSparseCacheHandle` | plain attention component | resident view and fused cache resolution |
231 | `SparseKVOffloadCommand` | dataclass | opaque scheduler-to-worker work |
232
233 ## Performance invariants
234
235 - Resident hits bypass hot-LRU lookup and host copies inside the fused resolver.
236 - Hot lookup, victim selection, and LRU updates stay on the GPU.
237 - A hot miss still copies directly from registered pinned host memory.
238 - Top-K resolution stays inside the attention invocation and remains graph
239 capturable.
240 - Compatible layers still share one miss plan.
241 - Index-sharing followers release their private LRU state before memory sizing.
242 - Indexer KV is untouched by HiSparse unless a generic KV offloader is configured.
243 - Resident and hot leases can still share one packed HMA allocation.
244 - No device scalar readback or CPU/device metadata round trip is added.
245 - The abstraction wraps the fused kernel; it does not add another kernel
246 launch.
247 - When HiSparse is disabled, the scheduler does not construct an offload
248 command or empty update table.
249
250 ## What remains platform-specific
251
252 The command/result and attention-layer boundaries can be shared. The host
253 allocator, copy implementation, hot layout, and replacement policy should stay
254 platform-specific. NVIDIA uses the current accelerator LRU and fused host/hot
255 kernel. ROCm is not currently supported; AMD or other accelerator backends can
256 implement their own worker without forcing NVIDIA's policy into the shared
257 boundary.