Repository: vllm-project/vllm
File: vllm/utils/torch_utils.py (Python)

Code before:

class ModuleName:
    """Wraps a module name string for use as a torch opaque type.

    When torch >= 2.11, this is registered as a hoisted value-type opaque
    object so that torch.compile lifts it as a graph input instead of baking
    it as a constant.  This avoids per-layer recompilation for MOE ops.
    """

    def __init__(self, value: str):
        self.value = value

    def __eq__(self, other):
        return isinstance(other, ModuleName) and self.value == other.value

    def __hash__(self):
        return hash(self.value)

    def __fx_repr__(self):
        return (f"ModuleName({self.value!r})", {ModuleName})


if HAS_OPAQUE_TYPE:
    # OpaqueBase registration requires the metaclass; apply it dynamically
    # so the class definition itself doesn't import from torch internals
    # on older torch versions.
    from torch._opaque_base import OpaqueBase, OpaqueBaseMeta

    ModuleName = OpaqueBaseMeta(  # type: ignore[misc]
        "ModuleName",
        (OpaqueBase,),
        {
            "__init__": ModuleName.__init__,
            "__eq__": ModuleName.__eq__,
            "__hash__": ModuleName.__hash__,
            "__fx_repr__": ModuleName.__fx_repr__,
            "__doc__": ModuleName.__doc__,
            "__module__": __name__,
        },
    )

    from torch._library.opaque_object import register_opaque_type

    register_opaque_type(ModuleName, typ="value", hoist=True)


# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
def supports_xccl() -> bool:
    return torch.distributed.is_xccl_available()



Review comment (suggestion):
Can we not do something like:

```suggestion
if HAS_OPAQUE_TYPE:
    from torch._opaque_base import OpaqueBase
else:
    OpaqueBase = object
```

Code after:
if HAS_OPAQUE_TYPE:
    from torch._opaque_base import OpaqueBase
else:
    OpaqueBase = object  # type: ignore[misc, assignment]


class ModuleName(OpaqueBase):  # type: ignore[misc]
    """Wraps a module name string for use as a torch opaque type.

    When torch >= 2.11, this is registered as a hoisted value-type opaque
    object so that torch.compile lifts it as a graph input instead of baking
    it as a constant.  This avoids per-layer recompilation for MOE ops.
    """

    def __init__(self, value: str):
        self.value = value

    def __eq__(self, other):
        return isinstance(other, ModuleName) and self.value == other.value

    def __hash__(self):
        return hash(self.value)

    def __fx_repr__(self):
        return (f"ModuleName({self.value!r})", {ModuleName})


if HAS_OPAQUE_TYPE:
    from torch._library.opaque_object import register_opaque_type

    register_opaque_type(ModuleName, typ="value", hoist=True)


# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
def supports_xccl() -> bool:
    return torch.distributed.is_xccl_available()


# Supports XPU Graph with PyTorch versions >= 2.11.0.dev for XPU platform
def supports_xpu_graph() -> bool:
    return is_torch_equal_or_newer("2.11.0.dev")


# create a library to hold the custom op
vllm_lib = Library("vllm", "FRAGMENT")  # noqa


def direct_register_custom_op(
    op_name: str,
    op_func: Callable,
    mutates_args: list[str] | None = None,