| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Structured key for identifying kernel config/autotune/benchmark cases.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import json |
| 8 | from typing import Any |
| 9 | |
| 10 | |
| 11 | class CaseKey(dict[str, Any]): |
| 12 | """Immutable, hashable dict for identifying kernel cases. |
| 13 | |
| 14 | Used as the key for config lookup, autotuning, benchmarking, and |
| 15 | input generation. Behaves like a read-only dict and can be used |
| 16 | as a dict key or in sets. |
| 17 | |
| 18 | The canonical string form (``__str__``) is stable JSON with sorted |
| 19 | keys. Use ``CaseKey.default()`` for the default/fallback key. |
| 20 | The regular constructor requires at least one key-value pair:: |
| 21 | |
| 22 | CaseKey({"intermediate": 2048, "numtokens": 256}) |
| 23 | CaseKey.default() # default/fallback |
| 24 | """ |
| 25 | |
| 26 | def __init__(self, *args: Any, _allow_empty: bool = False, **kwargs: Any): |
| 27 | super().__init__(*args, **kwargs) |
| 28 | if not self and not _allow_empty: |
| 29 | raise TypeError( |
| 30 | "CaseKey requires at least one key-value pair. " |
| 31 | "Use CaseKey.default() for the default config key." |
| 32 | ) |
| 33 | self._str: str | None = None |
| 34 | self._hash: int | None = None |
| 35 | |
| 36 | @classmethod |
| 37 | def default(cls) -> CaseKey: |
| 38 | """Create a default case key (empty).""" |
| 39 | return cls(_allow_empty=True) |
| 40 | |
| 41 | def __hash__(self) -> int: # type: ignore[override] |
| 42 | if self._hash is None: |
| 43 | self._hash = hash(str(self)) |
| 44 | return self._hash |
| 45 | |
| 46 | def __str__(self) -> str: |
| 47 | if self._str is None: |
| 48 | self._str = json.dumps(dict(self), sort_keys=True, separators=(",", ":")) |
| 49 | return self._str |
| 50 | |
| 51 | def __repr__(self) -> str: |
| 52 | if not self: |
| 53 | return "CaseKey.default()" |
| 54 | return f"CaseKey({dict(self)})" |
| 55 | |
| 56 | def is_default(self) -> bool: |
| 57 | """Return True if this is the default case key (empty).""" |
| 58 | return not self |
| 59 | |
| 60 | def _readonly(self, *args: Any, **kwargs: Any) -> Any: |
| 61 | raise TypeError("CaseKey is immutable") |
| 62 | |
| 63 | __setitem__ = _readonly # type: ignore[assignment] |
| 64 | __delitem__ = _readonly # type: ignore[assignment] |
| 65 | __ior__ = _readonly # type: ignore[assignment] |
| 66 | update = _readonly # type: ignore[assignment] |
| 67 | pop = _readonly # type: ignore[assignment] |
| 68 | popitem = _readonly # type: ignore[assignment] |
| 69 | setdefault = _readonly # type: ignore[assignment] |
| 70 | clear = _readonly # type: ignore[assignment] |