class KVCacheSpecRegistry:
    """Global registry mapping KVCacheSpec classes to SingleTypeManager and the rules to define as an uniform type."""

    @classmethod
    def register(cls, spec_class, manager_class, grouping_base_class):
        """
        allows all the platforms, both in-tree and out-of-tree, to register the new KVCacheSpec custom their implements of different KVCacheSpecs
        """

    @classmethod
    def override(cls, spec_class, manager_class, grouping_base_class)
        """
        allows all the platforms to custom their implements of an existed KVCacheSpec
        """
    @classmethod
    def get_manager_class(cls, spec: KVCacheSpec) -> Type[SingleTypeKVCacheManager]
        """
        define the SingleTypeManager to manage the blocks of current KVCacheSpec
        """
    @classmethod
    def get_uniform_type_base_spec(cls, spec: KVCacheSpec) -> Type[KVCacheSpec]
        """
        define the uniform type of KVCacheSpec, this is used for UniformTypeKVCacheSpecs to determine which specs could be treated as the same type, normally it is the KVCacheSpec itself.
        """

---

def register_kv_cache_spec(
    manager_class: Type["SingleTypeKVCacheManager"] | None = None,
    uniform_type_base_spec: Type["KVCacheSpec"] | None = None,
    override: bool = False,
):
    """
    Decorator to register a custom KVCacheSpec class.

    Usage for new specs:
        @register_kv_cache_spec(
            manager_class=FullAttentionManager,
            uniform_type_base_spec=FullAttentionSpec
        )
        @dataclass(frozen=True, kw_only=True)
        class MyCustomFullAttentionSpec(FullAttentionSpec):
            custom_alignment: int = 64

            @property
            def page_size_bytes(self) -> int:
                base_size = super().page_size_bytes
                return ((base_size + self.custom_alignment - 1)
                        // self.custom_alignment * self.custom_alignment)

    Usage for overriding existing specs:
        @register_kv_cache_spec(
            manager_class=IntelOptimizedFullAttentionManager,
            override=True
        )
        class FullAttentionSpec:
            pass  # Just to trigger the decorator

    Args:
        manager_class: The SingleTypeKVCacheManager to use for this spec.
            Required when override=False, optional when override=True.
        uniform_type_base_spec: The base spec class for uniform type kv cache specs compatibility.
            If None, the spec is treated as a new base type (when override=False)
            or keeps existing value (when override=True).
        override: If True, calls override() instead of register(). Use this when
            you want to change the manager for an existing spec without creating
            a new subclass.
    """

    def decorator(spec_class: Type["KVCacheSpec"]) -> Type["KVCacheSpec"]:
        if override:
            # Use override() method for existing specs
            KVCacheSpecRegistry.override(
                spec_class=spec_class,
                manager_class=manager_class,
                uniform_type_base_spec=uniform_type_base_spec,
            )
        else:
            # Use register() method for new specs
            if manager_class is None:
                raise ValueError(
                    "manager_class is required when override=False"
                )
            KVCacheSpecRegistry.register(
                spec_class=spec_class,
                manager_class=manager_class,
                uniform_type_base_spec=uniform_type_base_spec,
            )
        return spec_class

    return decorator

---

@register_kv_cache_spec(
    spec_class=FullAttentionSpec
    manager_class=FullAttentionManager,
    override=True
)
class CustomFullAttentionSpec:
    ...