arctic_inference/vllm/ulysses.py

# Copyright 2025 Snowflake Inc.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import threading
import weakref
from contextlib import contextmanager
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Any

import torch
import vllm.distributed.parallel_state as parallel_state
import vllm.envs as envs
from vllm.attention.layer import Attention
from vllm.config import ModelConfig, ParallelConfig
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
from vllm.distributed.parallel_state import (init_model_parallel_group,
                                             get_world_group,
                                             destroy_model_parallel,
                                             destroy_distributed_environment)
from vllm.executor.multiproc_worker_utils import (
    set_multiprocessing_worker_envs)
from vllm.utils import get_distributed_init_method, get_open_port
from vllm.v1.executor.abstract import FailureCallback
from vllm.v1.executor.multiproc_executor import (MultiprocExecutor, WorkerProc,
                                                 UnreadyWorkerProcHandle)
from vllm.platforms import current_platform
from vllm.utils import resolve_obj_by_qualname
from vllm.compilation.backends import PiecewiseCompileInterpreter
from vllm.model_executor.layers.fused_moe import FusedMoE

from arctic_inference.patching import ArcticPatch


def apply_shift_parallel_patches():
    UlyssesModelConfigPatch.apply_patch()
    UlyssesParallelStatePatch.apply_patch()
    UlyssesWorkerProcPatch.apply_patch()
    UlyssesMultiprocExecutorPatch.apply_patch()
    UlyssesAttentionPatch.apply_patch()
    PiecewiseCompileInterpreterPatch.apply_patch()
    UlyssesFusedMoEPatch.apply_patch()


class UlyssesModelConfigPatch(ArcticPatch[ModelConfig]):

    _orig_get_num_kv_heads = ModelConfig.get_num_kv_heads
    _orig_get_num_attention_heads = ModelConfig.get_num_attention_heads

    def get_num_kv_heads(self: ModelConfig,
                         parallel_config: ParallelConfig) -> int:
        num_kv_heads = self._orig_get_num_kv_heads(parallel_config)
        sp_size = parallel_config.ulysses_sequence_parallel_size
        return max(1, num_kv_heads // sp_size)

    def get_num_attention_heads(self: ModelConfig,
                                parallel_config: ParallelConfig) -> int:
        num_heads = self._orig_get_num_attention_heads(parallel_config)
        sp_size = parallel_config.ulysses_sequence_parallel_size
        return max(1, num_heads // sp_size)

    def get_layers_start_end_indices(
            self, parallel_config: "ParallelConfig") -> tuple[int, int]:
        from vllm.distributed.utils import get_pp_indices
        if (self.hf_text_config.model_type == "deepseek_mtp"
                or self.hf_config.model_type == "mimo_mtp"):
            total_num_hidden_layers = getattr(self.hf_text_config,
                                              "num_nextn_predict_layers", 0)
        else:
            total_num_hidden_layers = getattr(self.hf_text_config,
                                              "num_hidden_layers", 0)
        # the layout order is: DP x PP x SP x TP
        pp_rank = (parallel_config.rank //
                   (parallel_config.tensor_parallel_size *
                    parallel_config.ulysses_sequence_parallel_size)
                   ) % parallel_config.pipeline_parallel_size
        pp_size = parallel_config.pipeline_parallel_size
        start, end = get_pp_indices(total_num_hidden_layers, pp_rank, pp_size)
        return start, end


class UlyssesParallelStatePatch(ArcticPatch[parallel_state]):

    _SP = None
    _SP_TP = None
    _SP_AA = None
    _SP_AG = None
    # Rationale for SP_AA and SP_AG groups:
    # When num_kv_heads > SP, the kv heads are distributed and replicated as in TP.
    # To implement the logic, the distributed kv heads are exchanged with a local
    # all-to-all within SP_AA group followed by an local all-gather within SP_AG
    # group. The SP_AA and SP_AG groups partitions the SP group into two orthogonal
    # sub-groups and will not be initialized if max(1, num_kv_heads / TP) < SP.
    # See the figure in PR #126 https://github.com/snowflakedb/ArcticInference/pull/126

    def initialize_model_parallel(
        tensor_model_parallel_size: int = 1,
        pipeline_model_parallel_size: int = 1,
        backend: Optional[str] = None,
    ) -> None:
        """
        Initialize model parallel groups.

        Arguments:
            tensor_model_parallel_size: number of GPUs used for tensor model
                parallelism.
            pipeline_model_parallel_size: number of GPUs used for pipeline model
                parallelism.

        Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
        use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
        the model pipeline. The present function will
        create 4 tensor model-parallel groups and 2 pipeline model-parallel groups:
            4 tensor model-parallel groups:
                [g0, g1], [g2, g3], [g4, g5], [g6, g7]
            2 pipeline model-parallel groups:
                [g0, g2, g4, g6], [g1, g3, g5, g7]
        Note that for efficiency, the caller should make sure adjacent ranks
        are on the same DGX box. For example if we are using 2 DGX-1 boxes
        with a total of 16 GPUs, rank 0 to 7 belong to the first box and
        ranks 8 to 15 belong to the second box.
        """
        from vllm.distributed.parallel_state import _DP, _EP, _PP, _TP
        # Get world size and rank. Ensure some consistencies.
        assert torch.distributed.is_initialized()
        world_size: int = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        backend = backend or torch.distributed.get_backend(
            get_world_group().device_group)

        data_parallel_size = 1
        from vllm.config import get_current_vllm_config
        config = get_current_vllm_config()
        if config is not None:
            data_parallel_size = config.parallel_config.data_parallel_size

        sequence_parallel_size = \
            config.parallel_config.ulysses_sequence_parallel_size

        # the layout order is: ExternalDP x DP x PP x SP x TP
        # ExternalDP is the data parallel group that is not part of the model,
        # every dp rank can generate independently (in verl integration).
        # DP is the data parallel group that is part of the model,
        # all the ranks in the same DP group should generate simultaneously,
        # i.e. the `generate` call in the same DP group should be called together,
        # otherwise it will cause deadlock.
        # to get group_ranks for each dimension, transpose that dimension to the
        # last dimension, then reshape to 2D, then unbind the last dimension
        all_ranks = torch.arange(world_size).reshape(
            -1, data_parallel_size, pipeline_model_parallel_size,
            sequence_parallel_size, tensor_model_parallel_size)  # noqa

        # Build the tensor model-parallel groups.
        assert _TP is None, ("tensor model parallel group is already initialized")
        group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
        group_ranks = [x.tolist() for x in group_ranks]
        TP_group_ranks = group_ranks
        # message queue broadcaster is only used in tensor model parallel group
        _TP = init_model_parallel_group(group_ranks,
                                        get_world_group().local_rank,
                                        backend,
                                        use_message_queue_broadcaster=True,
                                        group_name="tp")

        # Build the pipeline model-parallel groups.
        assert _PP is None, (
            "pipeline model parallel group is already initialized")
        group_ranks = all_ranks.transpose(2, 4).reshape(
            -1, pipeline_model_parallel_size).unbind(0)
        group_ranks = [x.tolist() for x in group_ranks]
        PP_group_ranks = group_ranks
        _PP = init_model_parallel_group(group_ranks,
                                        get_world_group().local_rank,
                                        backend,
                                        group_name="pp")

        assert _DP is None, ("data parallel group is already initialized")
        group_ranks = all_ranks.transpose(1,
                                          4).reshape(-1,
                                                     data_parallel_size).unbind(0)
        group_ranks = [x.tolist() for x in group_ranks]
        DP_group_ranks = group_ranks
        _DP = init_model_parallel_group(group_ranks,
                                        get_world_group().local_rank,
                                        backend,
                                        group_name="dp")

        assert _EP is None, ("expert parallel group is already initialized")
        group_ranks = all_ranks.transpose(1, 3).reshape(
            -1, data_parallel_size * tensor_model_parallel_size).unbind(0)
        group_ranks = [x.tolist() for x in group_ranks]
        EP_group_ranks = group_ranks
        _EP = init_model_parallel_group(group_ranks,
                                        get_world_group().local_rank,
                                        backend,
                                        group_name="ep")

        # Build the sequence parallel groups.
        assert parallel_state._SP is None, (
            "sequence parallel group is already initialized")
        group_ranks = all_ranks.transpose(3, 4).reshape(
            -1, sequence_parallel_size).unbind(0)
        group_ranks = [x.tolist() for x in group_ranks]
        SP_group_ranks = group_ranks
        _SP = init_model_parallel_group(group_ranks,
                                        get_world_group().local_rank,
                                        backend,
                                        group_name="sp")

        # Build full-TP groups for ShiftParallel
        shift_parallel_size = (tensor_model_parallel_size *
                               sequence_parallel_size)
        assert parallel_state._SP_TP is None, (
            "full-TP group is already initialized")
        # transpose(3, 4) for obtaining the correct attn head order
        group_ranks = all_ranks.transpose(3, 4).reshape(
            -1, shift_parallel_size).unbind(0)
        group_ranks = [x.tolist() for x in group_ranks]
        SP_TP_group_ranks = group_ranks
        _SP_TP = init_model_parallel_group(group_ranks,
                                           get_world_group().local_rank,
                                           backend,
                                           group_name="sp_tp")

        parallel_state.logger.info(
            "rank %s in world size %s is assigned as DP rank %s, PP rank %s, "
            "TP rank %s, EP rank %s, SP rank %s, SP_TP rank %s", rank,
            world_size, _DP.rank_in_group, _PP.rank_in_group,
            _TP.rank_in_group, _EP.rank_in_group, _SP.rank_in_group,
            _SP_TP.rank_in_group)

        parallel_state._TP = _TP
        parallel_state._PP = _PP
        parallel_state._SP = _SP
        parallel_state._SP_TP = _SP_TP
        parallel_state._DP = _DP

        # check if SP requires kv replication
        num_kv_heads = config.model_config._orig_get_num_kv_heads(config.parallel_config)
        if num_kv_heads < sequence_parallel_size:

            # divide SP group into two orthogonal sub-groups:
            sp_aa_size = num_kv_heads
            sp_ag_size = sequence_parallel_size // num_kv_heads
            all_ranks_ = torch.arange(world_size).reshape(
            -1, data_parallel_size, pipeline_model_parallel_size,
            sp_aa_size, sp_ag_size, tensor_model_parallel_size)

            group_ranks = all_ranks_.transpose(3, 5).reshape(
                -1, sp_aa_size).unbind(0)
            group_ranks = [x.tolist() for x in group_ranks]
            SP_AA_group_ranks = group_ranks
            # SP_AA group is used for all-to-all communication of kv heads
            _SP_AA = init_model_parallel_group(group_ranks,
                                            get_world_group().local_rank,
                                            backend,
                                            group_name="sp_aa")
            
            group_ranks = all_ranks_.transpose(4, 5).reshape(
                -1, sp_ag_size).unbind(0)
            group_ranks = [x.tolist() for x in group_ranks]
            SP_AG_group_ranks = group_ranks
            # SP_AG group is used for all-gather communication of kv heads
            _SP_AG = init_model_parallel_group(group_ranks,
                                            get_world_group().local_rank,
                                            backend,
                                            group_name="sp_ag")

            parallel_state._SP_AA = _SP_AA
            parallel_state._SP_AG = _SP_AG

        if get_world_group().local_rank == 0:
            parallel_state.logger.info(
                    f"UlyssesParallelStatePatch initialized:\n"
                    f"  PP {_PP.world_size} ranks {PP_group_ranks}\n"
                    f"  TP {_TP.world_size} ranks {TP_group_ranks}\n"
                    f"  SP {_SP.world_size} ranks {SP_group_ranks}\n"
                    f"  DP {_DP.world_size} ranks {DP_group_ranks}\n"
                    f"  EP {_EP.world_size} ranks {EP_group_ranks}\n"
                    f"  SP_TP {_SP_TP.world_size} ranks {SP_TP_group_ranks}")
            if num_kv_heads < sequence_parallel_size:
                parallel_state.logger.info(
                    f"  SP_AA {parallel_state._SP_AA.world_size} ranks {SP_AA_group_ranks}\n"
                    f"  SP_AG {parallel_state._SP_AG.world_size} ranks {SP_AG_group_ranks}\n")

    @contextmanager
    def graph_capture(device: torch.device):
        """
        `graph_capture` is a context manager which should surround the code that
        is capturing the CUDA graph. Its main purpose is to ensure that the
        some operations will be run after the graph is captured, before the graph
        is replayed. It returns a `GraphCaptureContext` object which contains the
        necessary data for the graph capture. Currently, it only contains the
        stream that the graph capture is running on. This stream is set to the
        current CUDA stream when the context manager is entered and reset to the
        default stream when the context manager is exited. This is to ensure that
        the graph capture is running on a separate stream from the default stream,
        in order to explicitly distinguish the kernels to capture
        from other kernels possibly launched on background in the default stream.
        """
        from vllm.distributed.parallel_state import GraphCaptureContext
        context = GraphCaptureContext(torch.cuda.Stream(device=device))
        with parallel_state._TP.graph_capture(context), parallel_state._PP.graph_capture(
                context), parallel_state._SP_TP.graph_capture(context):
            yield context


class UlyssesWorkerProcPatch(ArcticPatch[WorkerProc]):

    def destroy_model_parallel(self):
        from vllm.distributed.parallel_state import _SP, _SP_TP, _SP_AA, _SP_AG
        if _SP:
            _SP.destroy()
        _SP = None
        if _SP_TP:
            _SP_TP.destroy()
        _SP_TP = None
        if _SP_AA:
            _SP_AA.destroy()
        _SP_AA = None
        if _SP_AG:
            _SP_AG.destroy()
        _SP_AG = None

    def shutdown(self):
        self.rpc_broadcast_mq = None
        self.worker_response_mq = None
        destroy_model_parallel()
        # destroy Ulysses communicators here
        self.destroy_model_parallel()
        destroy_distributed_environment()


class UlyssesMultiprocExecutorPatch(ArcticPatch[MultiprocExecutor]):

    def _init_executor(self) -> None:
        # Call self.shutdown at exit to clean up
        # and ensure workers will be terminated.
        self._finalizer = weakref.finalize(self, self.shutdown)
        self.is_failed = False
        self.shutdown_event = threading.Event()
        self.failure_callback: Optional[FailureCallback] = None
        self.io_thread_pool: Optional[ThreadPoolExecutor] = None

        self.world_size = self.parallel_config.world_size
        tensor_parallel_size = self.parallel_config.tensor_parallel_size
        pp_parallel_size = self.parallel_config.pipeline_parallel_size
        sp_parallel_size = self.parallel_config.ulysses_sequence_parallel_size
        assert (self.world_size ==
                tensor_parallel_size * pp_parallel_size * sp_parallel_size), (
            f"world_size ({self.world_size}) must be equal to the "
            f"tensor_parallel_size ({tensor_parallel_size}) x pipeline"
            f"_parallel_size ({pp_parallel_size}) x ulysses_sequence_parallel"
            f"_size ({sp_parallel_size}).")

        # Set multiprocessing envs that are common to V0 and V1
        set_multiprocessing_worker_envs(self.parallel_config)

        # Multiprocessing-based executor does not support multi-node setting.
        # Since it only works for single node, we can use the loopback address
        # 127.0.0.1 for communication.
        distributed_init_method = get_distributed_init_method(
            "127.0.0.1", get_open_port())

        # Initialize worker and set up message queues for SchedulerOutputs
        # and ModelRunnerOutputs
        max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024
        self.rpc_broadcast_mq = MessageQueue(self.world_size,
                                             self.world_size,
                                             max_chunk_bytes=max_chunk_bytes)
        scheduler_output_handle = self.rpc_broadcast_mq.export_handle()

        # Create workers
        unready_workers: list[UnreadyWorkerProcHandle] = []
        success = False
        try:
            for rank in range(self.world_size):
                unready_workers.append(
                    WorkerProc.make_worker_process(
                        vllm_config=self.vllm_config,
                        local_rank=rank,
                        rank=rank,
                        distributed_init_method=distributed_init_method,
                        input_shm_handle=scheduler_output_handle,
                    ))

            # Workers must be created before wait_for_ready to avoid
            # deadlock, since worker.init_device() does a device sync.
            self.workers = WorkerProc.wait_for_ready(unready_workers)

            # Ensure message queues are ready. Will deadlock if re-ordered
            # Must be kept consistent with the WorkerProc.
            self.rpc_broadcast_mq.wait_until_ready()
            for w in self.workers:
                w.worker_response_mq.wait_until_ready()

            self.start_worker_monitor()
            success = True
        finally:
            if not success:
                # Clean up the worker procs if there was a failure.
                self._ensure_worker_termination(
                    [w.proc for w in unready_workers])

        # For pipeline parallel, we use a thread pool for asynchronous
        # execute_model.
        if self.max_concurrent_batches > 1:
            # Note: must use only 1 IO thread to keep dequeue sequence
            # from the response queue
            self.io_thread_pool = ThreadPoolExecutor(
                max_workers=1, thread_name_prefix="mp_exec_io")

        self.output_rank = self._get_output_rank()


class UlyssesAttentionPatch(ArcticPatch[Attention]):

    _orig_init = Attention.__init__
    _orig_forward = Attention.forward

    def __init__(self, num_heads, *args, **kwargs):
        from .model_runner import is_shift_parallel_mode
        self.sp_size = parallel_state._SP.world_size
        self.sp_device_group = parallel_state._SP.device_group
        if not is_shift_parallel_mode():
            num_heads //= self.sp_size
            num_kv_heads = kwargs["num_kv_heads"]
            self.is_kv_replicated = True if num_kv_heads < self.sp_size else False
            if self.is_kv_replicated:
                num_kv_heads = 1
                assert parallel_state._SP_AA is not None and parallel_state._SP_AG is not None, (
                    "UlyssesAttentionPatch requires SP_AA and SP_AG groups to be initialized.")
                self.sp_aa_device_group = parallel_state._SP_AA.device_group
                self.sp_ag_device_group = parallel_state._SP_AG.device_group
                self.sp_aa_size = parallel_state._SP_AA.world_size
                self.sp_ag_size = parallel_state._SP_AG.world_size
                # this reorders the all-gathered sequence
                self.order = [j * self.sp_aa_size + i 
                              for i in range(self.sp_aa_size) 
                              for j in range(self.sp_ag_size)]
            else:
                num_kv_heads //= self.sp_size
            kwargs["num_kv_heads"] = num_kv_heads
        return self._orig_init(num_heads, *args, **kwargs)

    def forward(self, query, key, value, **kwargs):
        from .model_runner import is_shift_parallel_mode
        if self.sp_size == 1 or is_shift_parallel_mode():
            return self._orig_forward(query, key, value, **kwargs)

        if self.is_kv_replicated:
            # Ulysses all-to-all 1/2 (query)
            q = query.view(-1,
                           self.sp_size, self.num_heads * self.head_size).transpose(
                               0, 1).reshape(-1,
                                             self.num_heads * self.head_size)
            q_ = torch.empty_like(q)
            torch.distributed.all_to_all_single(q_, q, group=self.sp_device_group)
            # Ulysses pack (key, value)
            kv = torch.cat((key.view(-1, self.sp_aa_size, self.num_kv_heads * self.head_size),
                            value.view(-1, self.sp_aa_size, self.num_kv_heads * self.head_size)),
                           dim=-1).transpose(0, 1).reshape(
                               -1, 2 * self.num_kv_heads * self.head_size)
            # Ulysses all-to-all (key, value)
            kv_part = torch.empty_like(kv)
            torch.distributed.all_to_all_single(kv_part, kv, group=self.sp_aa_device_group)
            # Ulysses all-gather (key, value)
            kv_ = torch.empty(q_.shape[0],
                              2 * self.num_kv_heads * self.head_size,
                              dtype=query.dtype,
                              device=query.device)
            torch.distributed.all_gather_into_tensor(kv_,
                                                     kv_part,
                                                     group=self.sp_ag_device_group)
            # reorder
            kv_chunk = kv_.chunk(self.sp_size)
            kv_ordered = torch.cat([kv_chunk[i] for i in self.order])
            # unpack (key, value)
            k_, v_ = kv_ordered.split([self.num_kv_heads * self.head_size] * 2, dim=-1)
        else:
            # pack
            qkv = (torch.cat(
                (query.view(-1, self.sp_size, self.num_heads * self.head_size),
                key.view(-1, self.sp_size, self.num_kv_heads * self.head_size),
                value.view(-1, self.sp_size, self.num_kv_heads * self.head_size)),
                dim=-1)
                .transpose(0, 1)
                .reshape(-1, (self.num_heads + 2 * self.num_kv_heads) * self.head_size))
            # Ulysses all-to-all 1/2
            qkv_ = torch.empty_like(qkv)
            torch.distributed.all_to_all_single(qkv_, qkv, group=self.sp_device_group)
            # unpack
            q_, k_, v_ = qkv_.split([
                self.num_heads * self.head_size, self.num_kv_heads *
                self.head_size, self.num_kv_heads * self.head_size
            ], dim=-1)

        # original attention
        c_ = self._orig_forward(q_, k_, v_, **kwargs)

        # Ulysses all-to-all 2/2
        c = torch.empty_like(c_)
        torch.distributed.all_to_all_single(c, c_, group=self.sp_device_group)
        output = (c.view(self.sp_size, -1, self.num_heads * self.head_size)
                  .transpose(0, 1)
                  .reshape(-1, self.num_heads * self.sp_size * self.head_size))
        
        return output


class PiecewiseCompileInterpreterPatch(ArcticPatch[PiecewiseCompileInterpreter]):

    # find the symbolic shape of the subgraph
    def find_symbolic_shape(self, args: tuple[torch.fx.node.Argument,
                                ...]) -> torch.SymInt:
        symbols = set()
        for x in args:
            if isinstance(x, torch._subclasses.fake_tensor.FakeTensor):
                for dim in x.shape:
                    if isinstance(dim, torch.SymInt):
                        symbols.update(dim.node.expr.free_symbols)
        assert len(symbols) == 1, (
            f"Expected exactly one symbolic shape, but found {len(symbols)}: {symbols}")
        return list(symbols)[0]
  
    def call_module(self, target: torch.fx.node.Target,
                    args: tuple[torch.fx.node.Argument,
                                ...], kwargs: dict[str, Any]) -> Any:
        assert isinstance(target, str)
        # [Arctic Inference]
        # Since monkeypatching inherits the original class
        # through ArcticPatch class, we lose the access to the original class'
        # super() function. Instead of using super(), we directly invoke call_module
        # from the super class torch.fx.Interpreter of PiecewiseCompileInterpreter.
        # see - v0.9.0.1/compilation/backends.py#L241
        output = torch.fx.Interpreter.call_module(self, target, args, kwargs)

        if target in self.compile_submod_names:
            index = self.compile_submod_names.index(target)
            submod = self.fetch_attr(target)
            # [Arctic Inference]
            # Compiler may create subgraphs with certain symbolic
            # integer values that violates vllm's assumption here:
            # - v0.9.0.1/compilation/base_piecewise_backend.py#L64
            # The index of the significant symbol determines the runtime shape here:
            # - v0.9.0.1/compilation/cuda_piecewise_backend.py#L112
            # The fix is relaxing vllm's original assumption that there is only a
            # single symbolic that determines the shape.We then find the matching 
            # symbol indices.
            sym_shape = self.find_symbolic_shape(args)
            sym_shape_indices = []
            for i, x in enumerate(args):
                if isinstance(x, torch.SymInt):
                    if sym_shape == x:
                        sym_shape_indices.append(i)

            global compilation_start_time
            compiled_graph_for_general_shape = self.vllm_backend.\
                compiler_manager.compile(
                submod,
                args,
                self.compilation_config.inductor_compile_config,
                self.compilation_config,
                graph_index=index,
                num_graphs=len(self.compile_submod_names),
                runtime_shape=None)

            piecewise_backend = resolve_obj_by_qualname(
                current_platform.get_piecewise_backend_cls())
            self.module.__dict__[target] = piecewise_backend(
                submod, self.vllm_config, self.graph_pool, index,
                len(self.compile_submod_names), sym_shape_indices,
                compiled_graph_for_general_shape, self.vllm_backend)

            from vllm.compilation.counter import compilation_counter
            compilation_counter.num_piecewise_capturable_graphs_seen += 1

        return output


class UlyssesFusedMoEPatch(ArcticPatch[FusedMoE]):

    def forward(self, hidden_states: torch.Tensor,
                router_logits: torch.Tensor):
        # directly call forward_impl to bypass custom opt
        # custom opt prevents using the shift model
        # we will expand this function to fuse SP with EP
        return self.forward_impl(hidden_states, router_logits)