[BugFix] Fix MoRIIOConnector for disaggregated P/D inference

[raviguptaamd]

Summary

Three targeted fixes to MoRIIOConnector that resolve crashes and improve compatibility in Prefill/Decode disaggregated inference. All changes are scoped to the MORI-IO code path only — zero impact on NIXL, shared-storage, or SGLang connectors.

Fix 1: Cache kv_transfer_params across scheduler steps

The Request object's kv_transfer_params can be mutated or cleared between the update_state_after_alloc and build_connector_meta scheduler steps. By the time build_connector_meta runs, fields like remote_block_ids and remote_engine_id may be gone, causing KeyError on every decode request.

Fix: Snapshot kv_transfer_params into a _req_kv_params dict at allocation time and use the cached copy when building connector metadata. Cleanup is performed after both recv and save loops to avoid memory leaks.

Fix 2: Handle attention backend block_size overrides

For MLA models (e.g., DeepSeek-V3), the FlashMLA attention backend silently overrides block_size to 64 regardless of the --block-size CLI flag. The existing assert block_size == self.block_size compared the actual KV cache tensor shape against the CLI config value, causing an AssertionError during register_kv_caches.

Fix: Trust the actual tensor shape as ground truth. If it differs from the config, log an informational message and update self.block_size to match.

Fix 3: Default values for MORI-IO-specific kv_transfer_params fields

remote_handshake_port and remote_notify_port used direct dict access (kv_transfer_params["remote_handshake_port"]) with no fallback. External coordinators (e.g., the NIXL sidecar in LLM-D) that don't inject these MORI-IO-specific fields would cause KeyError.

Fix: Use .get() with MoRIIOConstants.DEFAULT_HANDSHAKE_PORT / DEFAULT_NOTIFY_PORT defaults, matching the pattern already used for tp_size and remote_dp_size.

Files Changed

File	Lines	Change
moriio_connector.py	+32/-3	Fix 1 (kv_params cache) + Fix 2 (block_size)
moriio_common.py	+10/-4	Fix 3 (default ports)

Hardware / Configuration

• AMD MI300X, 8 GPUs per node, TP=8
• 1 Prefill / 1 Decode disaggregation
• RDMA KV transfer via MORI-IO in READ mode
• Kubernetes LLM-D with pd-sidecar

Signed-off-by: Rav Gupta [email protected]

[gemini-code-assist]

Code Review

This pull request introduces three important fixes for the MoRIIOConnector to improve its stability and compatibility in disaggregated inference scenarios. The changes are well-targeted and address clear issues: caching kv_transfer_params to prevent race conditions, handling block_size overrides from attention backends gracefully, and providing default values for port configurations to avoid KeyError exceptions.

The implementation of these fixes is mostly solid. However, I've identified a potential issue with the use of pop() when retrieving cached kv_transfer_params. If a request happens to be in both the receive and save queues, this could lead to using stale data, reintroducing the bug this PR aims to fix. I've left specific comments with a suggestion to use get() instead for a more robust solution.

Overall, these are valuable fixes, and with the suggested adjustment, the connector will be much more resilient.