[Bugfix][XPU] Fix KDA RecoverSSM device-pointer overflow

[pmanczak]

Purpose

KDARecoverSSMCommitContext stores raw device pointers in a torch.int64
tensor. XPU USM addresses exceed 2**63, so building the metadata raises before
any kernel launches:

XPU data_ptr() = 0xffffd556b6400000 = 18446697167134392320 > int64 max 9223372036854775807

vllm/models/kimi_k3/nvidia/ops/recoverssm.py:856: ValueError
E   ValueError: Overflow when unpacking long long

Not a regression: recoverssm.py has one commit ever (#51855), which added the
code and the failing tests together. CUDA/ROCm pointers are always below
2**63.

Fix

Wrap data_ptr() in the two's-complement rewrite the tree already uses twice for
this exact problem — _reinterpret_u64_as_i64 in vllm/v1/worker/mamba_utils.py
(#48109) and _to_i64_ptr in vllm/device_allocator/xpumem.py. The kernels
recover the address via tl.pointer_type, which reinterprets rather than
converts, so the wrapped value round-trips. Below 2**63 the helper is the
identity, leaving CUDA/ROCm provably unchanged.

_base_addrs is the single funnel for all four pointer buffers. Strides stay
int64 — they are element counts, not addresses. Deduplicating the three copies
is RFC #50834's scope (open, unimplemented); this PR stays on the bug.

Not a duplicate

is:pr is:open recoverssm (#55688, #55215, #54495, #54255, #54103, #53298) and
is:pr is:open pointer overflow xpu — none touches device pointers in this file.
RFC #50834 has no implementation and does not mention K3.

Test plan and result

Arc Pro B70, torch 2.13.0+xpu, triton 3.7.2. Tested file sha256-matched to the
committed blob.

pytest tests/models/kimi_k3/test_kda.py -k recoverssm -q  # 4F/1P    -> 5 passed
pytest tests/models/kimi_k3/test_kda.py -q                # 47P/4F/11S -> 51P/0F/11S

The 11 skips are pre-existing platform gates, unchanged. No CUDA run: the change
is an identity below 2**63.

The reachable XPU config is Kimi-Linear-48B --use-replayssm --num-speculative-tokens N --mamba-backend triton — the validator accepts
KimiLinearForCausalLM and requires the Triton mamba backend. I have not run
that end-to-end; it needs TP=4 and my box is single-card.

Model evaluation

N/A: a provable no-op on CUDA/ROCm, and on XPU the path raised before any kernel
ran, so there is no prior behaviour to compare against. The four existing tests
check RecoverSSM against a reference, not just the absence of a crash.

AI assistance

AI assistance (Claude Code) was used for the investigation, the XPU measurements
and this description. As the submitting human I reviewed every changed line.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}