9was9 0.1.0__py3-none-any.whl

9was9-0.1.0.dist-info/METADATA

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: 9was9
+Version: 0.1.0
+Summary: Универсальный ускоритель numpy-вычислений (KV-кэш, RoPE, fused linear).
+Author-email: CharmLxt <reception.mailspypi@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/CharmLxt/9was9
+Project-URL: Source, https://github.com/CharmLxt/9was9
+Keywords: numpy,acceleration,kv-cache,rope,llm,inference
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.21.0
+
+9was9. Версия 0.1.0. Автор CharmLxt (PyPI: CharmLxt, email: reception.mailspypi@gmail.com). Ссылка на GitHub: https://github.com/CharmLxt/9was9.
+
+
+
+Важно про имя пакета. Установка производится командой pip install 9was9 (так пакет называется на PyPI), но импортировать его нужно как ninewas9, потому что имена Python-модулей не могут начинаться с цифры. Пример: from ninewas9 import KVCache, RoPECache, FusedLinear, fast\_softmax.
+
+
+
+9was9 — это универсальная библиотека-ускоритель для numpy-вычислений, созданная в первую очередь для нейросетевого инференса, но применимая в любом проекте, где важна скорость работы с массивами. В основе лежит несколько идей: предвыделение памяти вместо повторяющихся аллокаций, кэширование часто используемых таблиц, объединение нескольких матричных умножений в одно и численно стабильные реализации классических операций.
+
+
+
+Первая версия содержит пять групп функций. KVCache — это менеджер памяти для ключей и значений в трансформерах, который заранее выделяет буфер максимального размера и записывает новые значения в срез, избегая np.concatenate на каждом токене. Это уменьшает сложность вставки с O(n) до O(1) и убирает до 20% накладных расходов на инференс. RoPECache предвычисляет sin/cos таблицы для Rotary Position Embeddings с поддержкой нескольких типов масштабирования (linear, dynamic, llama3) и возвращает view без копирования. Это экономит до 15% времени forward-pass.
+
+
+
+FusedLinear объединяет несколько линейных слоёв (например, Q, K и V проекции или gate и up MLP) в одно матричное умножение за счёт конкатенации весов. Это экономит один проход по памяти вместо трёх и лучше утилизирует BLAS-ядра. FusedMLP — это готовый MLP-блок с fused gate/up и SwiGLU/GELU/ReLU активацией. Fast-softmax содержит численно стабильные реализации softmax и log-softmax (защита от inf/nan на больших логитах), causal mask, scaled dot-product attention с поддержкой GQA и MQA через broadcast, а также top-k и top-p фильтры для вероятностного сэмплинга.
+
+
+
+Установка: pip install 9was9. Зависимость только одна — numpy. Библиотека не требует torch, tensorflow или других тяжёлых фреймворков. Все функции работают с обычными numpy-массивами и не имеют побочных эффектов. Связь с автором: reception.mailspypi@gmail.com.
+

9was9-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+ninewas9/__init__.py,sha256=uGhoAHEsEqW9J1suSexkqtOG4Zqtspq8gw6_O43W9pg,762
+ninewas9/fast_softmax.py,sha256=6uhdzt_IwiAyFDfWJZbHqPPbwALZ8Bet9Re8V1U1PNY,6407
+ninewas9/fused_linear.py,sha256=aRo_Xs5dIOEL_CwbzwWQNzdBCcLtK-b07jsjcETD6tA,3701
+ninewas9/kvcache.py,sha256=UFGDz3Qn5eI4ej1TnRXnfGkKNaOtLGcMzEatXXA0ICw,4825
+ninewas9/ropecache.py,sha256=RGFhRYDtZJQI3FehnGz7OYQcOkjbitFpSuSAZl46irY,6728
+ninewas9/version.py,sha256=Pru0BlFBASFCFo7McHdohtKkUtgMPDwbGfyUZlE2_Vw,21
+9was9-0.1.0.dist-info/METADATA,sha256=XXZ6jFr_nmPOb1xcDc1LcKJS43_J0psJD_jGdsPD-uQ,4353
+9was9-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+9was9-0.1.0.dist-info/top_level.txt,sha256=H9pMoxbJfc1_XbT25E3iJEjmKN9UquHui9pXjBuUU8g,9
+9was9-0.1.0.dist-info/RECORD,,

9was9-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

9was9-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+ninewas9

ninewas9/__init__.py

@@ -0,0 +1,26 @@
+"""9was9 — универсальный ускоритель numpy-вычислений."""
+from .version import __version__
+from .kvcache import KVCache
+from .ropecache import RoPECache
+from .fused_linear import (
+    FusedLinear, FusedMLP, fuse_qkv, fuse_gate_up,
+)
+from .fast_softmax import (
+    fast_softmax, log_softmax,
+    apply_causal_mask, scaled_dot_product_attention,
+    top_k_filter, top_p_filter,
+)
+
+__all__ = [
+    "__version__",
+    # Memory
+    "KVCache",
+    # Rotary
+    "RoPECache",
+    # Fused linear
+    "FusedLinear", "FusedMLP", "fuse_qkv", "fuse_gate_up",
+    # Softmax / attention
+    "fast_softmax", "log_softmax",
+    "apply_causal_mask", "scaled_dot_product_attention",
+    "top_k_filter", "top_p_filter",
+]

ninewas9/fast_softmax.py

@@ -0,0 +1,156 @@
+"""fast_softmax — numerically stable softmax и scaled dot-product attention.
+
+Проблемы наивной реализации:
+    exp(x) / sum(exp(x))
+на больших x даёт inf/inf = nan. Поэтому используем трюк
+"вычитание максимума": exp(x - max(x)) / sum(exp(x - max(x))).
+
+Публичный API:
+    fast_softmax(x, axis=-1)
+    log_softmax(x, axis=-1)
+    scaled_dot_product_attention(q, k, v, mask=None, scale=None)
+    apply_causal_mask(scores, seq_len)
+"""
+import numpy as np
+
+
+def fast_softmax(x, axis=-1):
+    """Numerically stable softmax.
+
+    Параметры:
+        x: np.ndarray
+        axis: ось, по которой нормализуем (по умолчанию последняя)
+
+    Возвращает:
+        np.ndarray с той же формой, сумма по axis равна 1.
+    """
+    x = np.asarray(x)
+    x_max = np.max(x, axis=axis, keepdims=True)
+    # Обрабатываем случай, когда все значения -inf (например, полностью замьюченный ряд)
+    x_max = np.where(np.isfinite(x_max), x_max, 0.0)
+    exp_x = np.exp(x - x_max)
+    denom = np.sum(exp_x, axis=axis, keepdims=True)
+    # Защита от деления на ноль (если весь ряд был -inf)
+    denom = np.where(denom == 0, 1.0, denom)
+    return exp_x / denom
+
+
+def log_softmax(x, axis=-1):
+    """Numerically stable log-softmax. Полезен для cross-entropy и sampling."""
+    x = np.asarray(x)
+    x_max = np.max(x, axis=axis, keepdims=True)
+    x_max = np.where(np.isfinite(x_max), x_max, 0.0)
+    x_shifted = x - x_max
+    log_sum_exp = np.log(np.sum(np.exp(x_shifted), axis=axis, keepdims=True) + 1e-12)
+    return x_shifted - log_sum_exp
+
+
+def apply_causal_mask(scores, seq_len=None):
+    """Применяет causal mask: score[i, j] = -inf для j > i.
+
+    Параметры:
+        scores: (..., seq_len, seq_len)
+        seq_len: длина (если None — берётся из scores)
+
+    Возвращает:
+        masked scores (новый массив)
+    """
+    if seq_len is None:
+        seq_len = scores.shape[-1]
+
+    # Нижний треугольник = 0, верхний = -inf
+    mask = np.triu(np.full((seq_len, seq_len), -np.inf, dtype=np.float32), k=1)
+    return scores + mask
+
+
+def scaled_dot_product_attention(q, k, v, mask=None, scale=None, dropout_p=0.0):
+    """Стабильный scaled dot-product attention.
+
+    Параметры:
+        q: (batch, num_heads_q, seq_len_q, head_dim)
+        k: (batch, num_kv_heads, seq_len_kv, head_dim)
+        v: (batch, num_kv_heads, seq_len_kv, head_dim)
+        mask: (seq_len_q, seq_len_kv) или None. 0 = можно смотреть, -inf = нельзя.
+        scale: 1/sqrt(head_dim) по умолчанию
+        dropout_p: вероятность dropout (по умолчанию 0)
+
+    Возвращает:
+        (batch, num_heads_q, seq_len_q, head_dim)
+    """
+    if scale is None:
+        scale = 1.0 / np.sqrt(q.shape[-1])
+
+    num_heads_q = q.shape[1]
+    num_kv_heads = k.shape[1]
+
+    # GQA / MQA: если голов меньше — размножаем через broadcast (без копии)
+    if num_kv_heads != num_heads_q:
+        if num_heads_q % num_kv_heads != 0:
+            raise ValueError(
+                f"num_heads_q={num_heads_q} не делится на num_kv_heads={num_kv_heads}"
+            )
+        repeat_factor = num_heads_q // num_kv_heads
+        # Вместо np.repeat используем расширение размерности и broadcast
+        # k: (batch, num_kv_heads, seq, dim) → (batch, num_kv_heads, 1, seq, dim)
+        #    → broadcast с q (batch, num_heads_q, seq_q, dim) при reshape
+        # Проще: reshape q в (batch, num_kv_heads, repeat_factor, seq_q, dim)
+        k = np.repeat(k, repeat_factor, axis=1)
+        v = np.repeat(v, repeat_factor, axis=1)
+
+    # scores: (batch, num_heads_q, seq_len_q, seq_len_kv)
+    scores = np.matmul(q, np.swapaxes(k, -1, -2)) * scale
+
+    if mask is not None:
+        # mask broadcast: (seq_len_q, seq_len_kv) → (1, 1, seq_len_q, seq_len_kv)
+        scores = scores + mask
+
+    # Softmax по последней оси (по seq_len_kv)
+    attn_weights = fast_softmax(scores, axis=-1)
+
+    if dropout_p > 0:
+        # Простой dropout без random seed control (только для инференса не нужен)
+        keep_mask = np.random.rand(*attn_weights.shape) > dropout_p
+        attn_weights = attn_weights * keep_mask / (1.0 - dropout_p)
+
+    # output: (batch, num_heads_q, seq_len_q, head_dim)
+    output = np.matmul(attn_weights, v)
+    return output
+
+
+def top_k_filter(logits, k):
+    """Оставляет только k самых вероятных токенов, остальные = -inf."""
+    if k <= 0 or k >= logits.shape[-1]:
+        return logits
+    # Порог = k-й по величине
+    thresh = np.partition(logits, -k, axis=-1)[..., -k][..., np.newaxis]
+    return np.where(logits >= thresh, logits, -np.inf)
+
+
+def top_p_filter(logits, p):
+    """Nucleus sampling: оставляет минимальное множество, сумма вероятностей ≥ p."""
+    if p >= 1.0:
+        return logits
+
+    # Сортируем по убыванию
+    sorted_indices = np.argsort(-logits, axis=-1)
+    sorted_logits = np.take_along_axis(logits, sorted_indices, axis=-1)
+    sorted_probs = fast_softmax(sorted_logits, axis=-1)
+    cumulative_probs = np.cumsum(sorted_probs, axis=-1)
+
+    # Маска: элементы, где cumulative_probs > p — отбрасываем
+    # Но оставляем первый, где cumulative_probs > p
+    sorted_mask = cumulative_probs - sorted_probs >= p
+    sorted_logits = np.where(sorted_mask, -np.inf, sorted_logits)
+
+    # Восстанавливаем исходный порядок
+    result = np.empty_like(logits)
+    np.put_along_axis(result, sorted_indices, sorted_logits, axis=-1)
+    return result
+
+
+__all__ = [
+    "fast_softmax", "log_softmax",
+    "apply_causal_mask",
+    "scaled_dot_product_attention",
+    "top_k_filter", "top_p_filter",
+]

ninewas9/fused_linear.py

@@ -0,0 +1,104 @@
+"""FusedLinear — объединение нескольких линейных слоёв в один matmul."""
+import numpy as np
+
+
+class FusedLinear:
+    """Объединяет несколько матриц весов в одну и делает один matmul.
+
+    Параметры:
+        weights: список кортежей (name, W)
+                 где W — 2D массив формы (out_features, in_features)
+    """
+
+    def __init__(self, weights):
+        if not weights:
+            raise ValueError("FusedLinear: пустой список весов")
+
+        self.names = []
+        self.shapes = []
+        matrices = []
+
+        for name, W in weights:
+            if W.ndim != 2:
+                raise ValueError(f"FusedLinear: {name}.ndim={W.ndim}, ожидалось 2")
+            self.names.append(name)
+            self.shapes.append(W.shape)
+            matrices.append(W)
+
+        self.fused_weight = np.concatenate(matrices, axis=0)
+        self.fused_weight_T = self.fused_weight.T
+        self.total_out = self.fused_weight.shape[0]
+
+    def __call__(self, x):
+        if x.shape[-1] != self.shapes[0][1]:
+            raise ValueError(
+                f"FusedLinear: ожидалось in_features={self.shapes[0][1]}, "
+                f"получено {x.shape[-1]}"
+            )
+        out = x @ self.fused_weight_T
+        result = {}
+        offset = 0
+        for name, shape in zip(self.names, self.shapes):
+            n = shape[0]
+            result[name] = out[..., offset:offset + n]
+            offset += n
+        return result
+
+    def __repr__(self):
+        parts = [f"{name}{shape}" for name, shape in zip(self.names, self.shapes)]
+        return f"<FusedLinear [{', '.join(parts)}] fused={self.fused_weight.shape}>"
+
+
+def fuse_qkv(w_q, w_k, w_v):
+    """Объединяет Q, K, V проекции в один FusedLinear."""
+    return FusedLinear([("q", w_q), ("k", w_k), ("v", w_v)])
+
+
+def fuse_gate_up(w_gate, w_up):
+    """Объединяет gate и up проекции MLP."""
+    return FusedLinear([("gate", w_gate), ("up", w_up)])
+
+
+class FusedMLP:
+    """Объединённый MLP: [gate|up] @ x, потом down @ (act(gate) * up).
+
+    Параметры:
+        w_gate: (intermediate, hidden)
+        w_up:   (intermediate, hidden)
+        w_down: (hidden, intermediate)
+        activation: "silu" | "gelu" | "relu" | "identity" | callable
+    """
+
+    def __init__(self, w_gate, w_up, w_down, activation="silu"):
+        self.fused_gate_up = fuse_gate_up(w_gate, w_up)
+        self.w_down = w_down
+        self.w_down_T = w_down.T
+        self.activation = activation
+
+    def _apply_activation(self, x):
+        if callable(self.activation):
+            return self.activation(x)
+        if self.activation in ("silu", "swish"):
+            return x / (1.0 + np.exp(-x))
+        if self.activation == "gelu":
+            # Точный GELU через tanh-приближение
+            return 0.5 * x * (1.0 + np.tanh(
+                0.7978845608 * (x + 0.044715 * x ** 3)
+            ))
+        if self.activation == "relu":
+            return np.maximum(x, 0.0)
+        if self.activation == "identity":
+            return x
+        raise ValueError(f"FusedMLP: неизвестная активация '{self.activation}'")
+
+    def __call__(self, x):
+        gate_up = self.fused_gate_up(x)
+        gate = gate_up["gate"]
+        up = gate_up["up"]
+        return self._apply_activation(gate) * up @ self.w_down_T
+
+
+__all__ = [
+    "FusedLinear", "FusedMLP",
+    "fuse_qkv", "fuse_gate_up",
+]

ninewas9/kvcache.py

@@ -0,0 +1,108 @@
+"""KVCache — эффективный менеджер KV-кэша для трансформеров.
+
+Ключевая идея: вместо np.concatenate на каждом токене (что даёт
+квадратичную сложность), мы заранее выделяем буфер максимального размера
+и записываем новые K/V в срез. Чтение возвращает view, а не копию.
+
+Публичный API:
+    cache = KVCache(num_layers, num_kv_heads, head_dim, max_seq_len)
+    cache.append(layer_idx, k_new, v_new)   # k_new: (num_kv_heads, 1, head_dim)
+    k, v = cache.get(layer_idx)             # (num_kv_heads, cur_len, head_dim)
+    cache.reset()
+"""
+import numpy as np
+
+
+class KVCache:
+    """Менеджер KV-кэша с предвыделенной памятью.
+
+    Параметры:
+        num_layers: количество слоёв
+        num_kv_heads: количество KV-голов (для GQA/MQA меньше, чем Q-головы)
+        head_dim: размер одной головы
+        max_seq_len: максимальная длина последовательности (контекст)
+        batch_size: размер батча (пока 1)
+        dtype: dtype для хранения (внутри всегда float32 для скорости)
+
+    Атрибуты:
+        k: np.ndarray (num_layers, batch, num_kv_heads, max_seq_len, head_dim)
+        v: np.ndarray (num_layers, batch, num_kv_heads, max_seq_len, head_dim)
+        cur_len: текущая длина (число записанных токенов)
+    """
+
+    def __init__(self, num_layers, num_kv_heads, head_dim,
+                 max_seq_len, batch_size=1, dtype=np.float32):
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+        self.batch_size = batch_size
+        self.dtype = dtype
+
+        # Предвыделяем всё сразу. Это ~N_layers * B * H_kv * L * D * 4 байт.
+        # Для Llama-3-8B (32 слоя, 8 kv_heads, 128 head_dim, 8192 контекст):
+        # 32 * 1 * 8 * 8192 * 128 * 4 ≈ 1.07 GB на k + столько же на v.
+        # Это цена за O(1) запись вместо O(n) concatenate.
+        shape = (num_layers, batch_size, num_kv_heads, max_seq_len, head_dim)
+        self.k = np.zeros(shape, dtype=dtype)
+        self.v = np.zeros(shape, dtype=dtype)
+
+        self.cur_len = 0
+
+    def append(self, layer_idx, k_new, v_new):
+        """Записывает новые K/V для одного слоя.
+
+        Параметры:
+            layer_idx: индекс слоя (0 .. num_layers-1)
+            k_new: (batch_size, num_kv_heads, num_new_tokens, head_dim)
+            v_new: то же для V
+        """
+        n_new = k_new.shape[2]
+        pos = self.cur_len
+        if pos + n_new > self.max_seq_len:
+            raise ValueError(
+                f"KV-кэш переполнен: {pos + n_new} > {self.max_seq_len}. "
+                f"Увеличьте max_seq_len или используйте sliding window."
+            )
+        self.k[layer_idx, :, :, pos:pos + n_new, :] = k_new
+        self.v[layer_idx, :, :, pos:pos + n_new, :] = v_new
+
+    def get(self, layer_idx):
+        """Возвращает (k, v) для слоя. Возвращает view, а не копию.
+
+        Возвращает:
+            k: (batch_size, num_kv_heads, cur_len, head_dim)
+            v: то же для V
+        """
+        return (
+            self.k[layer_idx, :, :, :self.cur_len, :],
+            self.v[layer_idx, :, :, :self.cur_len, :],
+        )
+
+    def advance(self, n=1):
+        """Продвигает счётчик длины после записи всех слоёв."""
+        self.cur_len += n
+
+    def reset(self):
+        """Полностью сбрасывает кэш (новый запрос)."""
+        self.cur_len = 0
+
+    def memory_usage_bytes(self):
+        """Сколько байт занимает кэш."""
+        return self.k.nbytes + self.v.nbytes
+
+    def memory_usage_human(self):
+        """Читаемый размер."""
+        n = self.memory_usage_bytes()
+        for unit in ("B", "KB", "MB", "GB"):
+            if n < 1024:
+                return f"{n:.2f} {unit}"
+            n /= 1024
+        return f"{n:.2f} TB"
+
+    def __repr__(self):
+        return (
+            f"<KVCache layers={self.num_layers} heads={self.num_kv_heads} "
+            f"head_dim={self.head_dim} len={self.cur_len}/{self.max_seq_len} "
+            f"mem={self.memory_usage_human()}>"
+        )

ninewas9/ropecache.py

@@ -0,0 +1,151 @@
+"""RoPECache — предвычисленные sin/cos таблицы для Rotary Position Embeddings.
+
+RoPE применяется в каждой голове каждого слоя к каждому токену. Наивный
+подход — считать sin/cos на лету — тратит до 15% времени. Мы предвычисляем
+их один раз для максимальной длины контекста и кэшируем.
+
+Поддерживаемые варианты:
+- linear  (rope_scaling: type="linear", factor=F)
+- dynamic (rope_scaling: type="dynamic", factor=F)
+- llama3  (rope_scaling: type="llama3", low_freq_factor, high_freq_factor,
+           original_max_position_embeddings)
+
+Публичный API:
+    cache = RoPECache(head_dim, max_seq_len, rope_theta=10000.0,
+                      rope_scaling=None)
+    cos, sin = cache.get(seq_len)  # (seq_len, head_dim)
+    q_rot, k_rot = RoPECache.apply(q, k, cos, sin)
+"""
+import numpy as np
+
+
+class RoPECache:
+    """Предвычисленные cos/sin для RoPE.
+
+    Параметры:
+        head_dim: размер одной головы (должен быть чётным)
+        max_seq_len: максимальная длина последовательности
+        rope_theta: базовая частота (обычно 10000.0)
+        rope_scaling: dict с настройками или None
+    """
+
+    def __init__(self, head_dim, max_seq_len, rope_theta=10000.0,
+                 rope_scaling=None, dtype=np.float32):
+        if head_dim % 2 != 0:
+            raise ValueError(f"head_dim должен быть чётным, получено {head_dim}")
+
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling or {}
+        self.dtype = dtype
+
+        # Частоты: 1 / theta^(2i/head_dim), i = 0..head_dim/2-1
+        half_dim = head_dim // 2
+        freqs = 1.0 / (rope_theta ** (np.arange(0, half_dim, dtype=np.float64) / half_dim))
+
+        # Применяем scaling
+        scaling_type = self.rope_scaling.get("type", "linear")
+        if scaling_type == "linear":
+            factor = self.rope_scaling.get("factor", 1.0)
+            freqs = freqs / factor
+        elif scaling_type == "dynamic":
+            # Динамическое масштабирование (NTK-aware). Упрощённая версия:
+            # тот же линейный скейл, но по факту применяется только к длинным
+            # последовательностям. Для простоты делаем как linear.
+            factor = self.rope_scaling.get("factor", 1.0)
+            freqs = freqs / factor
+        elif scaling_type == "llama3":
+            # Llama 3 style scaling
+            factor = self.rope_scaling.get("factor", 8.0)
+            low_freq_factor = self.rope_scaling.get("low_freq_factor", 1.0)
+            high_freq_factor = self.rope_scaling.get("high_freq_factor", 4.0)
+            orig_max_pos = self.rope_scaling.get("original_max_position_embeddings", 8192)
+
+            low_freq_wavelen = orig_max_pos / low_freq_factor
+            high_freq_wavelen = orig_max_pos / high_freq_factor
+
+            wavelens = 2 * np.pi / freqs
+
+            # Низкие частоты (длинные волны) — делим на фактор
+            freqs = np.where(
+                wavelens > low_freq_wavelen,
+                freqs / factor,
+                freqs,
+            )
+            # Высокие частоты (короткие волны) — оставляем
+            # Промежуточные — плавный переход
+            medium_mask = (wavelens <= low_freq_wavelen) & (wavelens >= high_freq_wavelen)
+            smooth = (wavelens - high_freq_wavelen) / (low_freq_wavelen - high_freq_wavelen)
+            freqs = np.where(
+                medium_mask,
+                freqs / (factor * (1 - smooth) + smooth),
+                freqs,
+            )
+        # иначе (тип неизвестен) — игнорируем scaling
+
+        # Позиции: 0..max_seq_len-1
+        positions = np.arange(max_seq_len, dtype=np.float64)
+        # Внешнее произведение: (max_seq_len, half_dim)
+        angles = np.outer(positions, freqs)
+
+        cos_half = np.cos(angles).astype(dtype)
+        sin_half = np.sin(angles).astype(dtype)
+
+        # Дублируем на весь head_dim: [cos, cos] и [sin, sin]
+        self.cos = np.concatenate([cos_half, cos_half], axis=-1)  # (L, D)
+        self.sin = np.concatenate([sin_half, sin_half], axis=-1)  # (L, D)
+
+        self.cur_len = 0
+
+    def get(self, seq_len):
+        """Возвращает (cos, sin) для первых seq_len позиций.
+
+        Возвращает:
+            cos: (seq_len, head_dim)
+            sin: (seq_len, head_dim)
+        """
+        if seq_len > self.max_seq_len:
+            raise ValueError(
+                f"seq_len={seq_len} > max_seq_len={self.max_seq_len}"
+            )
+        return self.cos[:seq_len], self.sin[:seq_len]
+
+    def slice_for_positions(self, start, end):
+        """Возвращает cos/sin для позиций [start, end)."""
+        return self.cos[start:end], self.sin[start:end]
+
+    @staticmethod
+    def rotate_half(x):
+        """Поворот половины: [-x2, x1] для x = [x1, x2]."""
+        half = x.shape[-1] // 2
+        x1 = x[..., :half]
+        x2 = x[..., half:]
+        return np.concatenate([-x2, x1], axis=-1)
+
+    @staticmethod
+    def apply(q, k, cos, sin):
+        """Применяет RoPE к Q и K.
+
+        Параметры:
+            q: (batch, num_heads, seq_len, head_dim)
+            k: (batch, num_kv_heads, seq_len, head_dim)
+            cos: (seq_len, head_dim)
+            sin: (seq_len, head_dim)
+
+        Возвращает:
+            q_rot, k_rot с той же формой.
+        """
+        # cos/sin нужно broadcast'ить: (1, 1, seq_len, head_dim)
+        cos = cos[np.newaxis, np.newaxis, :, :]
+        sin = sin[np.newaxis, np.newaxis, :, :]
+
+        q_rot = (q * cos) + (RoPECache.rotate_half(q) * sin)
+        k_rot = (k * cos) + (RoPECache.rotate_half(k) * sin)
+        return q_rot, k_rot
+
+    def __repr__(self):
+        return (
+            f"<RoPECache head_dim={self.head_dim} max_seq={self.max_seq_len} "
+            f"theta={self.rope_theta} scaling={self.rope_scaling.get('type', 'none')}>"
+        )

ninewas9/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"