[Model Runner V2] Add Support for XD-RoPE

[santiramos27]

Summary

• Adds XD-RoPE (variable-dimension RoPE) support to Model Runner V2, enabling models like HunyuanOCR that use 3 or 4-dim rotary embeddings
• Unifies M-RoPE and XD-RoPE into a single RopeState class with one parameterized Triton kernel (NUM_DIMS: tl.constexpr), replacing the previous M-RoPE-only MRopeState
• The kernel always loads and applies delta (non-zero for M-RoPE, zero for XD-RoPE), avoiding a separate code path per variant

Changes

• vllm/v1/worker/gpu/mm/rope.py (new): Unified RopeState class and get_rope_state() factory. The factory validates the model interface (SupportsMRoPE / SupportsXDRoPE) and dispatches based on
  ModelConfig
• vllm/v1/worker/gpu/mm/mrope_utils.py (deleted): Replaced by rope.py
• vllm/v1/worker/gpu/model_states/default.py: Simplified to use get_rope_state() — all M-RoPE/XD-RoPE branching now lives in the factory
• vllm/v1/worker/gpu/model_runner.py: Pass intermediate_tensors in model_inputs dict directly; non-first PP rank asserts it's not None
• vllm/v1/worker/gpu/cudagraph_utils.py: Add intermediate_tensors: None for CUDA graph capture with TODO for PP support ([Model Runner V2] Enable piecewise & full CUDA graphs for pipeline parallelism #35162)

Test Result

• Verified numerical equivalence between V1 and V2 model runners for both XD-RoPE (HunyuanOCR with image input) and M-RoPE (Qwen2-VL-2B-Instruct with image input)
• Identical OCRBench result on V1 vs V2 model runner for HunyuanOCR and Qwen2-VL-2B-Instruct

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request adds support for XD-RoPE. The changes include a new XDRopeState class with a Triton kernel to manage XD-RoPE positions, and integration into the DefaultModelState. The core logic seems correct, but I've identified an opportunity to improve the clarity and maintainability of the new Triton kernel by refactoring its indexing logic to be more conventional. The rest of the changes for integrating XD-RoPE support are well-structured.

[gemini-code-assist]

The current implementation of _prepare_xdrope_positions_kernel and its call site in prepare_xdrope_positions is functionally correct but hard to understand due to unconventional argument passing for tensor indexing. The arguments prefill_xdrope_positions_stride0 and prefill_xdrope_positions_stride1 are not standard tensor strides, making the address calculation logic within the kernel confusing.

For better readability and maintainability, I suggest refactoring to use standard stride-based indexing. This involves passing the actual stride of the prefill_xdrope_positions tensor and adjusting the kernel to calculate the row index explicitly. This change does not alter the logic but makes the code much clearer and aligned with common practices for writing Triton kernels.

    def prepare_xdrope_positions(
        self,
        idx_mapping: torch.Tensor,
        query_start_loc: torch.Tensor,
        prefill_lens: torch.Tensor,
        num_computed_tokens: torch.Tensor,
    ) -> None:
        num_reqs = idx_mapping.shape[0]
        prefill_positions = self.prefill_xdrope_positions.gpu
        _prepare_xdrope_positions_kernel[(num_reqs,)](
            self.xdrope_positions,
            self.xdrope_positions.stride(0),
            prefill_positions,
            prefill_positions.stride(0),
            idx_mapping,
            query_start_loc,
            prefill_lens,
            num_computed_tokens,
            BLOCK_SIZE=1024,
            USES_XDROPE_DIM=self.uses_xdrope_dim,
        )


@triton.jit
def _prepare_xdrope_positions_kernel(
    xdrope_positions_ptr,
    xdrope_positions_stride,
    prefill_xdrope_positions_ptr,
    prefill_xdrope_positions_stride0,
    idx_mapping_ptr,
    query_start_loc_ptr,
    prefill_lens_ptr,
    num_computed_tokens_ptr,
    BLOCK_SIZE: tl.constexpr,
    USES_XDROPE_DIM: tl.constexpr,
):
    batch_idx = tl.program_id(0)
    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)

    prefill_len = tl.load(prefill_lens_ptr + req_state_idx)
    num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
    is_prefill = num_computed < prefill_len

    query_start = tl.load(query_start_loc_ptr + batch_idx)
    query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
    query_len = query_end - query_start

    for i in range(0, query_len, BLOCK_SIZE):
        block = i + tl.arange(0, BLOCK_SIZE)
        mask = block < query_len
        orig_pos = num_computed + block

        for j in tl.static_range(USES_XDROPE_DIM):
            if is_prefill:
                # Read from pre-computed XD-RoPE positions.
                row_idx = req_state_idx * USES_XDROPE_DIM + j
                pos = tl.load(
                    prefill_xdrope_positions_ptr
                    + row_idx * prefill_xdrope_positions_stride0
                    + orig_pos,
                    mask=mask,
                )
            else:
                pos = orig_pos
            tl.store(
                xdrope_positions_ptr
                + j * xdrope_positions_stride
                + query_start
                + block,
                pos,
                mask=mask,
            )

[njhill]

Thanks @santiramos27!

[njhill]

May be good to keep common case at the top, i.e.

if not self.uses_mrope and not self.uses_xdrope:
    # Common case (1D positions).
    return {}

if self.uses_mrope:
    # ...
    return ...

# xdrope logic
return ...

[WoosukKwon]

+1. I intentionally put the early exit for the common case at the top. It'd be nice to keep it!

[santiramos27]

Early exit makes sense. updated.

[njhill]

This can be a bool

Suggested change

	self.uses_xdrope_dim = self.model_config.uses_xdrope_dim
	self.uses_xdrope = self.model_config.uses_xdrope_dim > 0

I actually like using self.xdrope_state = XDRopeState | None for this rather than separate bool, but I don't think @WoosukKwon likes that :)

Or, I feel it might be better/clearer to instead have something like a pe_type enum (MROPE, XDROPE or None)

[WoosukKwon]

+1. I'm also ok with if self.xdrope_state is not None.

[santiramos27]

I will change to using if self.xdrope_state is not None. From what I can tell, these two PE variants are the only ones getting special treatment in model runner. Should we wait to use enums once we have more PE's to support? It already feels like a good amount of duplication in mrope/xdrope that can probably get consolidated once we have more of them to worry about?

[njhill]

sure, I just thought it's maybe clearer to have a single enum than two bools which can't both be true

[njhill]

we could consider adding this for clarity

Suggested change

	model_inputs["inputs_embeds"] = None
	model_inputs["inputs_embeds"] = None
	assert intermediate_tensors is not None

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @santiramos27.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @santiramos27, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[njhill]

Looking at the code now ... it would simplify things to have a shared RopeState interface with init_prefill_positions, prepare_positions and get_positions(num_tokens) methods

[santiramos27]

Done. I followed the interface pattern under model_states/ and did a slight renaming for clarity. Let me know what you think.

[njhill]

Thanks @santiramos27 it looks good to me now.

I can see that there's a natural generalization of both of these, where the mrope uses dim 3 and xdrope uses bias 0. We could have a single kernel and other state logic used by both which would eliminate more of the code. But I have a feeling that @WoosukKwon won't like that idea :-)

[santiramos27]

Thanks for the feedback @njhill! I incorporated your edits.

@WoosukKwon wdyt about consolidating the /mm/rope triton kernels?

[santiramos27]

@njhill Updated with the unified position prep triton kernel -- this feels cleaner to me than the redundant util files + interface.

[njhill]

Thanks @santiramos27 this looks great. Just have one more suggestion.

[njhill]

Thanks @santiramos27 just couple of last minor things!

[njhill]

I guess this should be renamed to e.g. is_mrope now?

[santiramos27]

Since get_rope_state() already handles M-RoPE vs XD-RoPE dispatch, has_delta feels more natural here — it describes the behavioral difference rather than leaking model-type identity into the unified RopeState (ie. mrope and xdrope are instances of RopeState). Wdyt?

[njhill]

probably still worth has_delta (or is_mrope) check here?

[santiramos27]

You're right -- for xdrope the UVA tensor is zero-initialized and we never write to it, so this copy is unnecessary

[njhill]

nit

Suggested change

	elif model_config.uses_xdrope_dim > 0:
	if model_config.uses_xdrope_dim > 0:

[santiramos27]

mutually exclusive but i forgot i added early return here :')

[njhill]

yeah it's not incorrect maybe more a code style preference :)

[njhill]

Thanks @santiramos27 for the nice work!