intel review

Signed-off-by: NickLucche <[email protected]>

Author: NickLucche

tests/out_prefill

@@ -419,6 +419,7 @@ def _nixl_handshake(
                     engine_id=self.REMOTE_ENGINE_ID,
                     agent_metadata=FakeNixlWrapper.AGENT_METADATA,
                     kv_caches_base_addr=[0],
+                    device_id=remote_tp_rank,
                     num_blocks=1,
                     block_lens=remote_block_lens,
                     attn_backend_name=self.backend_name,

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -956,7 +956,8 @@ def _nixl_handshake(
                 metadata = decoder.decode(metadata_bytes)
                 got_metadata_time = time.perf_counter()
                 logger.debug(
-                    "NIXL handshake: get metadata took: %s", got_metadata_time - start_time
+                    "NIXL handshake: get metadata took: %s",
+                    got_metadata_time - start_time,
                 )
 
                 # Ensure engine id matches.
@@ -1474,9 +1475,6 @@ def _validate_remote_agent_handshake(
             )
             remote_block_size = remote_block_len // (self.slot_size_per_layer[0])
         else:
-            if tp_ratio != 1 and self.device_type == "xpu":
-                # XPU uses NHD, hence it does not support splitting on H
-                raise ValueError("Heterogeneous TP is not supported on XPU")
             # When MLA is not used, this is a list of the same block length
             for block_len in nixl_agent_meta.block_lens:
                 assert block_len == remote_block_len, (