@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ef5a226819

Fork 0

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 5, 09:45 UTC

[PD][HeteroArch]Fix accuracy issue with CPU_ATTN as Decoder and Flash_ATTN as prefiller (#38935)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

Chendi.Xue committed Apr 8, 2026 at 22:19 UTC ef5a226819e02d0d4a931bcf42b80314076e3629

3 files changed +93

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -523,6 +523,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

523 kv_cache_layout="HND",

524 block_size=self.block_size,

525 ssm_sizes=(0, 0),

526 + attn_backend_name=self.backend_name,

527 ),

528 remote_tp_rank=remote_tp_rank,

529 remote_tp_size=remote_tp_size,

@@ -972,6 +973,7 @@ class TestNixlHandshake:

973 kv_cache_layout=mismatched_layout,

974 block_size=worker.block_size,

975 ssm_sizes=(0, 0),

976 + attn_backend_name=worker.backend_name,

977 )

978

979 with pytest.raises(RuntimeError):

@@ -1028,6 +1030,7 @@ class TestNixlHandshake:

1030 kv_cache_layout="HND",

1031 block_size=worker.block_size,

1032 ssm_sizes=(0, 0),

1033 + attn_backend_name=worker.backend_name,

1034 )

1035

1036 # We don't check layout for homogeneous TP and MLA for now, as the

@@ -2347,6 +2350,7 @@ def test_compatibility_hash_validation(

2350 kv_cache_layout="HND",

2351 block_size=prefill_block_size,

2352 ssm_sizes=(0, 0),

2353 + attn_backend_name=decode_worker.backend_name,

2354 )

2355 handshake_payload = NixlHandshakePayload(

2356 compatibility_hash=remote_hash,

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+49

@@ -173,6 +173,7 @@ class NixlAgentMetadata:

173 kv_cache_layout: str

174 block_size: int

175 ssm_sizes: tuple[int, int]

176 + attn_backend_name: str

177

178

179 @dataclass

@@ -1116,6 +1117,7 @@ class NixlConnectorWorker:

1117

1118 self.num_blocks = kv_cache_config.num_blocks

1119 self.enable_permute_local_kv = False

1120 + self.enable_heterogeneous_attn_post_process = False

1121

1122 # KV Caches and nixl tracking data.

1123 self.device_type = current_platform.device_type

@@ -1776,6 +1778,7 @@ class NixlConnectorWorker:

1778 else self.host_buffer_kv_cache_layout,

1779 block_size=self.block_size,

1780 ssm_sizes=self._mamba_ssm_size,

1781 + attn_backend_name=self.backend_name,

1782 )

1783 # Wrap metadata in payload with hash for defensive decoding

1784 assert self.compat_hash is not None

@@ -2369,6 +2372,21 @@ class NixlConnectorWorker:

2372 "Or enable experimental feature to use HND to NHD support by "

2373 "setting 'enable_permute_local_kv'=True in --kv-transfer-config."

2374 )

2375 + # if remote_agent used attn is not same as local,

2376 + # hint heterogenuous attn post process

2377 + if (

2378 + nixl_agent_meta.attn_backend_name != self.backend_name

2379 + and self.backend_name in ["CPU_ATTN"]

2380 + ):

2381 + if self._is_hma_required:

2382 + raise RuntimeError(

2383 + "heterogeneous attn post process is not supported with HMA"

2384 + )

2385 + logger.info(

2386 + "[Experimental] CPU_ATTN backend is used, "

2387 + "hint heterogeneous attn post process"

2388 + )

2389 + self.enable_heterogeneous_attn_post_process = True

2390

2391 # Heterogeneous TP requires head-splitting, which only works with

2392 # HND layout. MLA and replicated-KV cases don't split on heads.

@@ -2542,6 +2560,28 @@ class NixlConnectorWorker:

2560 cache, indices, block_size_ratio

2561 )

2562

2563 + def post_process_device_kv_on_receive_heterogeneous_attn(

2564 + self, block_ids: list[int]

2565 + ):

2566 + """

2567 + Post process device kv cache after receiving from remote

2568 + for heterogeneous attention.

2569 + """

2570 + assert self.enable_heterogeneous_attn_post_process

2571 +

2572 + indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)

2573 +

2574 + for _, cache_or_caches in self.device_kv_caches.items():

2575 + blocks_to_update = cache_or_caches.index_select(1, indices)

2576 + current_platform.pack_kv_cache(

2577 + key=blocks_to_update[0],

2578 + value=blocks_to_update[1],

2579 + key_cache=cache_or_caches[0],

2580 + value_cache=cache_or_caches[1],

2581 + block_ids=block_ids,

2582 + indices=indices,

2583 + )

2584 +

2585 def get_finished(self) -> tuple[set[str], set[str]]:

2586 """

2587 Get requests that are done sending or recving on this specific worker.

@@ -2566,6 +2606,7 @@ class NixlConnectorWorker:

2606 )

2607

2608 block_ids_for_blocksize_post_process = defaultdict(list)

2609 + block_ids_for_heterogeneous_attn_post_process = list[list[int]]()

2610 for req_id in done_recving:

2611 # clean up metadata for completed requests

2612 meta = self._recving_metadata.pop(req_id, None)

@@ -2585,12 +2626,20 @@ class NixlConnectorWorker:

2626 block_ids_for_blocksize_post_process[block_size_ratio].append(

2627 meta.local_physical_block_ids[0]

2628 )

2629 + # post processing for heterogeneous attention

2630 + if self.enable_heterogeneous_attn_post_process:

2631 + block_ids_for_heterogeneous_attn_post_process.append(

2632 + meta.local_physical_block_ids[0]

2633 + )

2634 for (

2635 block_size_ratio,

2636 block_ids_list,

2637 ) in block_ids_for_blocksize_post_process.items():

2638 self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list)

2639

2640 + for block_ids in block_ids_for_heterogeneous_attn_post_process:

2641 + self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)

2642 +

2643 # Handle timeout to avoid stranding blocks on remote.

2644 now = time.perf_counter()

2645 while self._reqs_to_send:

vllm/platforms/cpu.py

+40

@@ -520,3 +520,43 @@ class CpuPlatform(Platform):

520 import vllm._C # noqa: F401

521 except ImportError as e:

522 logger.warning("Failed to import from vllm._C: %r", e)

523 +

524 + @classmethod

525 + def pack_kv_cache(

526 + cls,

527 + key: torch.Tensor,

528 + value: torch.Tensor,

529 + key_cache: torch.Tensor,

530 + value_cache: torch.Tensor,

531 + block_ids: list[int],

532 + indices: torch.Tensor,

533 + ) -> None:

534 + """

535 + Rewrite the kv cache shape for the current platform.

536 + """

537 + # Import lazily: cpu_attn pulls in _custom_ops, which needs a fully

538 + # initialized vllm.platforms (avoid circular import while CpuPlatform loads).

539 + from vllm._custom_ops import cpu_attn_reshape_and_cache

540 + from vllm.v1.attention.backends.cpu_attn import _get_attn_isa

541 +

542 + dtype = key.dtype

543 + # For CPU_ATTN, the shape is [N, num_kv_heads, block_size, head_size]

544 + _, _, block_size, head_size = key_cache.shape

545 + key = key.permute(0, 2, 1, 3).flatten(0, 1)

546 + value = value.permute(0, 2, 1, 3).flatten(0, 1)

547 +

548 + isa = _get_attn_isa(dtype, block_size, head_size)

549 + block_offsets = torch.arange(block_size, device="cpu", dtype=torch.long)

550 + num_blocks = len(block_ids)

551 + slot_mapping = (

552 + block_offsets.reshape(1, block_size)

553 + + indices.reshape(num_blocks, 1) * block_size

554 + ).flatten()

555 + cpu_attn_reshape_and_cache(

556 + key,

557 + value,

558 + key_cache,

559 + value_cache,

560 + slot_mapping,

561 + isa,

562 + )