[SSM] Follow-up fix for Mamba P/D KV transfer

[ZhanqiuHu]

Connector-only N-1 prefill fix for Mamba P/D disaggregation. Prevents state corruption where D-side recomputes last prompt token from already-complete state. Single file change: nixl_connector.py (+16 lines), no scheduler changes.

Quick sanity (2p2d, prompt: The capital of France is):

Without fix:

Paris. Thats a weird sentence: The capital of France is is Paris. It repeats is. The user wants exactly that sentence, no

With fix:

Paris. No extra punctuation? The sentence includes a period at the end. The user wrote: The capital of France is Paris. So we

[gemini-code-assist]

Code Review

This pull request introduces a fix for Mamba/SSM state corruption during P/D KV transfer by ensuring the P-side only prefills N-1 tokens. While the logic is sound, there's a critical issue with how the request object is modified. The function get_num_new_matched_tokens is not guaranteed to be called only once for a request, but the current implementation repeatedly truncates the prompt on subsequent calls, which is a bug. My review includes a suggestion to make this modification idempotent to prevent this issue.

[gemini-code-assist]

The modification of the request object here introduces a side effect in get_num_new_matched_tokens. This function can be called multiple times for the same request if it remains in the scheduler's waiting queue across several scheduling cycles. The current implementation is not idempotent and will repeatedly pop tokens from prompt_token_ids on each call, leading to an incorrect prompt length.

To fix this, the modification should only happen once. I suggest adding a flag to the kv_transfer_params to track whether the truncation has already been performed.

Suggested change

	if (self._is_hma_required
	and params is not None
	and params.get("do_remote_decode")):
	if request.prompt_token_ids and len(request.prompt_token_ids) > 1:
	request.prompt_token_ids.pop()
	request._all_token_ids.pop()
	request.num_prompt_tokens -= 1
	request.max_tokens = 1
	if (self._is_hma_required
	and params is not None
	and params.get("do_remote_decode")
	and not params.get("_p_side_truncated")):
	if request.prompt_token_ids and len(request.prompt_token_ids) > 1:
	request.prompt_token_ids.pop()
	request._all_token_ids.pop()
	request.num_prompt_tokens -= 1
	request.max_tokens = 1
	params["_p_side_truncated"] = True