A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 20:15 UTC
[Bugfix] Fix mismatched kernel-per-logical blocks in NIXL HMA transfer (#42097)
Signed-off-by: ZhanqiuHu <zhu@redhat.com> Signed-off-by: Zhanqiu Hu <zhu@redhat.com> Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: NickLucche <nlucches@redhat.com>
zhanqiuhu committed
May 12, 2026 at 09:53 UTC
593d5a40339ab0c7b0e6b3efb7540bdeab82161c
2 files changed
+395
-29
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+300
-8
@@ -93,27 +93,39 @@ def test_logical_to_kernel_block_ids_with_hma():
93
94
@pytest.mark.cpu_test
95
@pytest.mark.parametrize(
96
- "group_spec_types,expansion_stride,remote_block_ids,expected_remote_block_ids",
96
+ "group_spec_types,remote_physical_per_logical,"
97
+ "local_physical_per_logical,tp_ratio,remote_block_ids,"
98
+ "expected_remote_block_ids",
99
[
100
pytest.param(
101
("FullAttentionSpec", "SlidingWindowSpec"),
102
2,
103
+ 2,
104
+ 1,
105
([0, 1, 2], [3, 4]),
106
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]],
107
id="dense_fa_swa",
108
),
109
+ # Nemotron-3-Nano-30B-A3B 4p1d (P_TP=4, D_TP=1):
110
+ # remote_physical_per_logical=34, local_physical_per_logical=66.
111
+ # FA logical block 5 → kernel [170..203], block 6 → [204..237].
112
+ # Mamba block unchanged.
113
pytest.param(
114
("FullAttentionSpec", "MambaSpec"),
107
- 261,
108
- ([0, 1, 2], [10, 11]),
109
- [[0, 1, 261, 262, 522, 523], [10, 11]],
115
+ 34,
116
+ 66,
117
+ -4,
118
+ ([5, 6], [2]),
119
+ [list(range(170, 238)), [2]],
120
id="mamba_fa_ssm",
121
),
122
],
123
)
124
def test_read_blocks_for_req_expands_remote_ids(
125
group_spec_types,
116
- expansion_stride,
126
+ remote_physical_per_logical,
127
+ local_physical_per_logical,
128
+ tp_ratio,
129
remote_block_ids,
130
expected_remote_block_ids,
131
):
@@ -148,7 +160,7 @@ def test_read_blocks_for_req_expands_remote_ids(
160
resolved_types = tuple(spec_name_to_type[n] for n in group_spec_types)
161
162
worker = object.__new__(NixlConnectorWorker)
151
- worker._physical_blocks_per_logical_kv_block = 2
163
+ worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
164
165
has_mamba = any(t is MambaSpec for t in resolved_types)
166
has_swa = any(t is SlidingWindowSpec for t in resolved_types)
@@ -159,9 +171,11 @@ def test_read_blocks_for_req_expands_remote_ids(
171
remote_engine_id = "remote-engine"
172
173
worker.transfer_topo = MagicMock()
162
- worker.transfer_topo.tp_ratio.return_value = 1
174
+ # tp_ratio not exercised (all_source_ranks is empty so no reads run),
175
+ # but set for realism.
176
+ worker.transfer_topo.tp_ratio.return_value = tp_ratio
177
remote_info = MagicMock()
164
- remote_info.remote_physical_blocks_per_logical = expansion_stride
178
+ remote_info.remote_physical_blocks_per_logical = remote_physical_per_logical
179
worker.transfer_topo.get_engine_info.return_value = remote_info
180
worker.use_mla = False
181
@@ -192,6 +206,168 @@ def test_read_blocks_for_req_expands_remote_ids(
206
)
207
208
209
+@pytest.mark.cpu_test
210
+@pytest.mark.parametrize(
211
+ "local_physical_per_logical,remote_physical_per_logical,"
212
+ "local_block_ids,remote_block_ids,"
213
+ "expected_local,expected_remote",
214
+ [
215
+ # 10 kernel blocks of data, local has more logical blocks.
216
+ # remote physical_per_logical=10 → 1 logical → 10 kernel blocks
217
+ # local physical_per_logical=6 → 2 logical → 12 kernel blocks
218
+ # Trim local from 12 to 10.
219
+ pytest.param(
220
+ 6,
221
+ 10,
222
+ [list(range(12)), [42]],
223
+ [list(range(10)), [42]],
224
+ [list(range(10)), [42]],
225
+ [list(range(10)), [42]],
226
+ id="align_local6_remote10",
227
+ ),
228
+ # 10 kernel blocks of data, remote has more logical blocks.
229
+ # remote physical_per_logical=6 → 2 logical → 12 kernel blocks
230
+ # local physical_per_logical=10 → 1 logical → 10 kernel blocks
231
+ # Trim remote from 12 to 10.
232
+ pytest.param(
233
+ 10,
234
+ 6,
235
+ [list(range(10)), [42]],
236
+ [list(range(12)), [42]],
237
+ [list(range(10)), [42]],
238
+ [list(range(10)), [42]],
239
+ id="align_local10_remote6",
240
+ ),
241
+ ],
242
+)
243
+def test_apply_prefix_caching_mamba_hybrid(
244
+ local_physical_per_logical,
245
+ remote_physical_per_logical,
246
+ local_block_ids,
247
+ remote_block_ids,
248
+ expected_local,
249
+ expected_remote,
250
+):
251
+ """_apply_prefix_caching front-trims FA groups to
252
+ min(local, remote) for Mamba hybrid models with heterogeneous TP.
253
+ """
254
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
255
+ NixlConnectorWorker,
256
+ )
257
+ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
258
+
259
+ worker = object.__new__(NixlConnectorWorker)
260
+ worker._has_mamba = True
261
+ worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
262
+ worker._group_spec_types = (FullAttentionSpec, MambaSpec)
263
+ worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
264
+
265
+ aligned_local, aligned_remote = worker._apply_prefix_caching(
266
+ local_block_ids, remote_block_ids, remote_physical_per_logical
267
+ )
268
+
269
+ assert aligned_local == expected_local, (
270
+ f"Expected local {expected_local}, got {aligned_local}"
271
+ )
272
+ assert aligned_remote == expected_remote, (
273
+ f"Expected remote {expected_remote}, got {aligned_remote}"
274
+ )
275
+
276
+
277
+@pytest.mark.cpu_test
278
+@pytest.mark.parametrize(
279
+ "local_physical_per_logical,remote_physical_per_logical,"
280
+ "remote_fa_blocks,local_fa_blocks,ssm_blocks,"
281
+ "correct_remote_fa,correct_local_fa",
282
+ [
283
+ # 10 kernel blocks of data (640 tokens).
284
+ # remote physical_per_logical=10 → 1 logical → 10 kernel [0..9]
285
+ # local physical_per_logical=6 → 2 logical → 12 kernel [0..11]
286
+ # 1st local logical block cached → suffix [6..11]
287
+ # Correct: transfer only uncached suffix tokens (384-639)
288
+ # = remote [6,7,8,9] → local [6,7,8,9].
289
+ # Actual (front-trim): remote[:6]=[0..5] → local [6..11]. Wrong.
290
+ pytest.param(
291
+ 6,
292
+ 10,
293
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
294
+ [6, 7, 8, 9, 10, 11],
295
+ [42],
296
+ [6, 7, 8, 9],
297
+ [6, 7, 8, 9],
298
+ id="local6_remote10_fail",
299
+ ),
300
+ # 15 kernel blocks of data (960 tokens).
301
+ # remote physical_per_logical=6 → 3 logical → 18 kernel [0..17]
302
+ # local physical_per_logical=10 → 2 logical → 20 kernel [0..19]
303
+ # 1st local logical block cached → suffix [10..19]
304
+ # Correct: transfer only uncached suffix tokens (640-959)
305
+ # = remote [10,11,12,13,14] → local [10,11,12,13,14].
306
+ # Actual (front-trim): remote[:10]=[0..9] → local [10..19]. Wrong.
307
+ pytest.param(
308
+ 10,
309
+ 6,
310
+ list(range(18)),
311
+ list(range(10, 20)),
312
+ [42],
313
+ [10, 11, 12, 13, 14],
314
+ [10, 11, 12, 13, 14],
315
+ id="local10_remote6_fail",
316
+ ),
317
+ ],
318
+)
319
+def test_mismatched_physical_per_logical_fails_with_prefix_caching(
320
+ local_physical_per_logical,
321
+ remote_physical_per_logical,
322
+ remote_fa_blocks,
323
+ local_fa_blocks,
324
+ ssm_blocks,
325
+ correct_remote_fa,
326
+ correct_local_fa,
327
+):
328
+ """Demonstrate that _apply_prefix_caching front-trims ([:N])
329
+ in the Mamba hybrid path, which fails when prefix caching produces
330
+ suffix-only local blocks.
331
+
332
+ Prefix caching operates at logical block granularity. When a logical
333
+ block is cached locally, the decode side only allocates kernel blocks
334
+ for the uncached suffix. The front-trim pairs remote prefix blocks
335
+ with local suffix slots — a silent data corruption.
336
+ """
337
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
338
+ NixlConnectorWorker,
339
+ )
340
+
341
+ worker = object.__new__(NixlConnectorWorker)
342
+ worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
343
+ worker.kv_cache_config = make_kv_cache_config(
344
+ block_size=16,
345
+ mamba_enabled=True,
346
+ )
347
+ worker._has_mamba = True
348
+ worker._group_spec_types = tuple(
349
+ type(g.kv_cache_spec) for g in worker.kv_cache_config.kv_cache_groups
350
+ )
351
+
352
+ local_block_ids = (local_fa_blocks, ssm_blocks)
353
+ remote_block_ids = (remote_fa_blocks, ssm_blocks)
354
+
355
+ aligned_local, aligned_remote = worker._apply_prefix_caching(
356
+ local_block_ids,
357
+ remote_block_ids,
358
+ remote_physical_per_logical,
359
+ )
360
+
361
+ assert (
362
+ aligned_remote[0] != correct_remote_fa or aligned_local[0] != correct_local_fa
363
+ ), (
364
+ f"Prefix caching with mismatched physical_per_logical should not "
365
+ f"produce correct transfer ids: "
366
+ f"remote={aligned_remote[0]}, local={aligned_local[0]}, "
367
+ f"correct_remote={correct_remote_fa}, correct_local={correct_local_fa}"
368
+ )
369
+
370
+
371
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
372
def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
373
"""Test that a prefill instance returns fewer "remote blocks" for the SWA groups
@@ -564,3 +740,119 @@ def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_rati
740
)
741
742
assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio
743
+
744
+
745
+@pytest.mark.cpu_test
746
+@pytest.mark.parametrize(
747
+ "mamba_enabled,swa_enabled,"
748
+ "local_physical_per_logical,remote_physical_per_logical,"
749
+ "logical_block_ids,expected_kernel_block_ids",
750
+ [
751
+ # Qwen3.5-0.8B 4P2D (kernel_block_size=64):
752
+ # prefill TP=4: logical_block_size=384 → physical_per_logical=6
753
+ # decode TP=2: logical_block_size=640 → physical_per_logical=10
754
+ # FA logical [0] → remote kernel [0..9] (1 * 10)
755
+ # SSM logical [10] → unchanged [10]
756
+ pytest.param(
757
+ True,
758
+ False,
759
+ 6,
760
+ 10,
761
+ ([0], [10]),
762
+ [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10]],
763
+ id="qwen35_4p2d",
764
+ ),
765
+ # Qwen3.5-0.8B 2P4D (kernel_block_size=64):
766
+ # prefill TP=2: logical_block_size=640 → physical_per_logical=10
767
+ # decode TP=4: logical_block_size=384 → physical_per_logical=6
768
+ # FA logical [0, 1] → remote kernel [0..5, 6..11] (2 * 6)
769
+ # SSM logical [10] → unchanged [10]
770
+ pytest.param(
771
+ True,
772
+ False,
773
+ 10,
774
+ 6,
775
+ ([0, 1], [10]),
776
+ [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], [10]],
777
+ id="qwen35_2p4d",
778
+ ),
779
+ # Homogeneous TP (kernel_block_size=64):
780
+ # both sides: logical_block_size=640 → physical_per_logical=10
781
+ # FA logical [0] → kernel [0..9], SSM unchanged
782
+ pytest.param(
783
+ True,
784
+ False,
785
+ 10,
786
+ 10,
787
+ ([0], [10]),
788
+ [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10]],
789
+ id="homo_tp",
790
+ ),
791
+ # remote physical_per_logical=1: early return, no expansion
792
+ pytest.param(
793
+ True,
794
+ False,
795
+ 10,
796
+ 1,
797
+ ([0, 1, 2], [5]),
798
+ [[0, 1, 2], [5]],
799
+ id="mamba_remote_physical_per_logical_1",
800
+ ),
801
+ # Pure FA (no mamba): single group expanded with remote stride
802
+ pytest.param(
803
+ False,
804
+ False,
805
+ 2,
806
+ 4,
807
+ ([0, 1],),
808
+ [[0, 1, 2, 3, 4, 5, 6, 7]],
809
+ id="pure_fa",
810
+ ),
811
+ # FA + SWA (no mamba): both groups expanded
812
+ pytest.param(
813
+ False,
814
+ True,
815
+ 2,
816
+ 3,
817
+ ([0, 1], [2, 3]),
818
+ [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]],
819
+ id="fa_swa",
820
+ ),
821
+ ],
822
+)
823
+def test_logical_to_remote_kernel_block_ids(
824
+ mamba_enabled,
825
+ swa_enabled,
826
+ local_physical_per_logical,
827
+ remote_physical_per_logical,
828
+ logical_block_ids,
829
+ expected_kernel_block_ids,
830
+):
831
+ """Verify _logical_to_remote_kernel_block_ids uses the remote
832
+ physical_per_logical for FA expansion, not the local one.
833
+
834
+ This was the root cause of silent accuracy corruption in Qwen3.5
835
+ heterogeneous TP (e.g. 4P2D): the old code used local physical_per_logical
836
+ for the expansion arange, producing wrong kernel block indices.
837
+
838
+ Qwen3.5-0.8B values verified by verify_conv_split.py (issue #13).
839
+ """
840
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
841
+ NixlConnectorWorker,
842
+ )
843
+
844
+ worker = object.__new__(NixlConnectorWorker)
845
+ worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
846
+ worker.kv_cache_config = make_kv_cache_config(
847
+ block_size=16,
848
+ mamba_enabled=mamba_enabled,
849
+ swa_enabled=swa_enabled,
850
+ )
851
+
852
+ result = worker._logical_to_remote_kernel_block_ids(
853
+ logical_block_ids,
854
+ remote_physical_per_logical,
855
+ )
856
+ assert list(result) == expected_kernel_block_ids, (
857
+ f"Expected {expected_kernel_block_ids}, got {result}"
858
+ )
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+95
-21
@@ -1446,6 +1446,23 @@ class NixlConnectorWorker:
1446
tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)
1447
)
1448
1449
+ remote_physical_per_logical = (
1450
+ nixl_agent_meta.physical_blocks_per_logical_kv_block
1451
+ )
1452
+ if (
1453
+ self._has_mamba
1454
+ and remote_physical_per_logical
1455
+ != self._physical_blocks_per_logical_kv_block
1456
+ and self.vllm_config.cache_config.enable_prefix_caching
1457
+ ):
1458
+ raise RuntimeError(
1459
+ "Prefix caching with heterogeneous physical_blocks_per_logical "
1460
+ "is not supported for Mamba hybrid models. "
1461
+ f"Local: {self._physical_blocks_per_logical_kv_block}, "
1462
+ f"Remote: {remote_physical_per_logical}. "
1463
+ "Disable prefix caching with --no-enable-prefix-caching."
1464
+ )
1465
+
1466
if self._is_hma_required:
1467
assert block_size_ratio == 1, (
1468
"HMA does not support different remote block size yet"
@@ -2182,15 +2199,10 @@ class NixlConnectorWorker:
2199
== len(local_block_ids)
2200
== len(self.kv_cache_config.kv_cache_groups)
2201
)
2185
- # Partial prefix cache hit: just read uncomputed blocks.
2186
- # Skip mamba groups — their blocks represent full state (conv+ssm),
2187
- # not per-token data, so trimming would corrupt the transfer.
2188
- remote_block_ids = list(remote_block_ids)
2189
- for i, remote_group in enumerate(remote_block_ids):
2190
- num_local_blocks = len(local_block_ids[i])
2191
- assert num_local_blocks <= len(remote_group)
2192
- if num_local_blocks < len(remote_group):
2193
- remote_block_ids[i] = remote_group[-num_local_blocks:]
2202
+ remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical
2203
+ local_block_ids, remote_block_ids = self._apply_prefix_caching(
2204
+ local_block_ids, remote_block_ids, remote_physical_per_logical
2205
+ )
2206
2207
# NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
2208
# corresponding rank. With heterogeneous TP, fixing D>P, the D tp
@@ -2287,6 +2299,66 @@ class NixlConnectorWorker:
2299
for i, group in enumerate(block_ids)
2300
]
2301
2302
+ def _apply_prefix_caching(
2303
+ self,
2304
+ local_block_ids: BlockIds,
2305
+ remote_block_ids: BlockIds,
2306
+ remote_physical_per_logical: int,
2307
+ ) -> tuple[BlockIds, list]:
2308
+ """Apply prefix caching by trimming local/remote block ID lists.
2309
+
2310
+ For non-Mamba models: end-trim remote to match local count, so that
2311
+ already-cached prefix blocks are skipped in the transfer.
2312
+
2313
+ For Mamba hybrid (prefix caching not yet supported): front-trim both
2314
+ to the minimum count to handle kernel block count discrepancies from
2315
+ logical block rounding in heterogeneous TP.
2316
+ """
2317
+ # Partial prefix cache hit: just read uncomputed blocks.
2318
+ # Skip mamba groups — their blocks represent full state (conv+ssm),
2319
+ # not per-token data, so trimming would corrupt the transfer.
2320
+ remote_block_ids = list(remote_block_ids)
2321
+ if not self._has_mamba:
2322
+ for i, remote_group in enumerate(remote_block_ids):
2323
+ num_local_blocks = len(local_block_ids[i])
2324
+ assert num_local_blocks <= len(remote_group)
2325
+ if num_local_blocks < len(remote_group):
2326
+ remote_block_ids[i] = remote_group[-num_local_blocks:]
2327
+ else:
2328
+ # (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP
2329
+ # can cause different kernel block counts due to logical block rounding.
2330
+ # Example: 640 prompt tokens, kernel_block_size=64
2331
+ # remote physical_per_logical=10, local physical_per_logical=6
2332
+ # remote logical ids from kv_transfer_params = [0]
2333
+ # local logical ids allocated = [0, 1]
2334
+ # remote kernel blocks: [0..9] (1*10=10)
2335
+ # local kernel blocks: [0..11] (2*6=12)
2336
+ # actual data blocks = ceil(640/64) = 10, trim both to 10
2337
+ # Vice versa (remote physical_per_logical=6, local=10):
2338
+ # remote logical ids = [0, 1], local logical ids = [0]
2339
+ # remote kernel blocks: [0..11] (2*6=12)
2340
+ # local kernel blocks: [0..9] (1*10=10)
2341
+ # actual data blocks = ceil(640/64) = 10, trim both to 10
2342
+ local_block_ids = list(local_block_ids)
2343
+ for i, remote_group in enumerate(remote_block_ids):
2344
+ num_local_blocks = len(local_block_ids[i])
2345
+ num_remote_blocks = len(remote_group)
2346
+ if _is_ssm_spec(self._group_spec_types[i]):
2347
+ assert num_local_blocks == num_remote_blocks
2348
+ else:
2349
+ max_padding = max(
2350
+ self._physical_blocks_per_logical_kv_block,
2351
+ remote_physical_per_logical,
2352
+ )
2353
+ assert abs(num_local_blocks - num_remote_blocks) < max_padding, (
2354
+ f"Group {i}: |{num_local_blocks} - "
2355
+ f"{num_remote_blocks}| >= {max_padding}"
2356
+ )
2357
+ num_blocks = min(num_local_blocks, num_remote_blocks)
2358
+ local_block_ids[i] = local_block_ids[i][:num_blocks]
2359
+ remote_block_ids[i] = remote_group[:num_blocks]
2360
+ return local_block_ids, remote_block_ids
2361
+
2362
def _logical_to_remote_kernel_block_ids(
2363
self, block_ids: BlockIds, remote_physical_per_logical: int
2364
) -> BlockIds:
@@ -2299,23 +2371,25 @@ class NixlConnectorWorker:
2371
2372
Returns:
2373
Same structure with FA groups expanded (each logical block L
2302
- becomes kernel blocks [L*ratio .. L*ratio + local_ratio - 1]).
2374
+ becomes kernel blocks [L*remote_physical_per_logical, ..
2375
+ L*remote_physical_per_logical +
2376
+ remote_physical_per_logical - 1]).
2377
Mamba groups are passed through unchanged.
2378
"""
2305
- local_ratio = self._physical_blocks_per_logical_kv_block
2379
if remote_physical_per_logical == 1:
2380
return block_ids
2308
- local_arange = np.arange(local_ratio).reshape(1, -1)
2381
+ remote_arange = np.arange(remote_physical_per_logical).reshape(1, -1)
2382
group_specs = self.kv_cache_config.kv_cache_groups
2310
- result: list[list[int]] = []
2311
- for i, group in enumerate(block_ids):
2312
- if not isinstance(group_specs[i].kv_cache_spec, MambaSpec):
2313
- arr = np.array(group).reshape(-1, 1)
2314
- expanded = (arr * remote_physical_per_logical + local_arange).flatten()
2315
- result.append(expanded.tolist())
2316
- else:
2317
- # Mamba blocks are 1:1 logical-to-physical (no expansion).
2318
- result.append(group)
2383
+ result = [
2384
+ BlockTable.map_to_kernel_blocks(
2385
+ np.array(group),
2386
+ remote_physical_per_logical,
2387
+ remote_arange,
2388
+ ).tolist()
2389
+ if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
2390
+ else group
2391
+ for i, group in enumerate(block_ids)
2392
+ ]
2393
return result
2394
2395
def get_backend_aware_kv_block_len(