NPU KV Cache by Codex

engine/README.md

@@ -0,0 +1,16 @@
+# Refactored Project Copy
+
+This directory contains a cleaned-up copy of the original demo project.
+
+## Structure
+
+- `src/engine/`: application source code
+- `tests/`: test and demo scripts
+- `examples/`: runnable training example
+
+## Main improvements
+
+- Reorganized from flat files into a Python package layout
+- Replaced mixed naming styles with `snake_case` modules and clearer class names
+- Added type hints and docstrings to improve readability and maintenance
+- Kept the original source files in the project root untouched

engine/examples/train_transformer.py

@@ -0,0 +1,35 @@
+"""Example script for training and running the decoder-only transformer."""
+
+from torch import nn, optim
+
+from engine.demo_data import INDEX_TO_TOKEN, TARGET_VOCAB, make_inference_batch, make_training_batch
+from engine.transformer_decoder_only import Transformer
+
+
+def main() -> None:
+    sentences = [["S i want a beer", "i want a beer E"], ["S i want a beer", "i want a beer E"]]
+    model = Transformer(len(TARGET_VOCAB))
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+    decoder_inputs, target_batch = make_training_batch(sentences)
+
+    for epoch in range(20):
+        optimizer.zero_grad()
+        outputs = model(decoder_inputs)
+        loss = criterion(outputs.view(-1, outputs.size(-1)), target_batch.contiguous().view(-1))
+        print(f"Epoch: {epoch + 1:04d} cost = {loss:.6f}")
+        loss.backward()
+        optimizer.step()
+
+    inference_sentences = ["S P P P P"] * 10
+    decoder_inputs = make_inference_batch(inference_sentences)
+    predictions = model(decoder_inputs, 16)
+    predictions = predictions[:, :, : len(TARGET_VOCAB) - 1].data.max(2, keepdim=True)[1]
+
+    for sentence, output in zip(inference_sentences, predictions):
+        print(sentence, "->", [INDEX_TO_TOKEN[index.item()] for index in output.squeeze()])
+
+
+if __name__ == "__main__":
+    main()

engine/pyproject.toml

@@ -0,0 +1,15 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "engine-refactored"
+version = "0.1.0"
+description = "Refactored copy of the KV cache demo project."
+requires-python = ">=3.9"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

engine/src/__init__.py

@@ -0,0 +1,21 @@
+"""Core package for the refactored KV cache demo project."""
+
+from .block_manager import BlockAllocator, MetadataEngine, MindsporeAPI
+from .demo_data import INDEX_TO_TOKEN, TARGET_VOCAB, make_inference_batch, make_training_batch
+from .kv_cache import Cache, CacheBackend, MindsporeCacheBackend, TorchCacheBackend
+from .transformer_decoder_only import Transformer
+
+__all__ = [
+    "BlockAllocator",
+    "MetadataEngine",
+    "MindsporeAPI",
+    "TARGET_VOCAB",
+    "INDEX_TO_TOKEN",
+    "make_inference_batch",
+    "make_training_batch",
+    "Cache",
+    "CacheBackend",
+    "MindsporeCacheBackend",
+    "TorchCacheBackend",
+    "Transformer",
+]

engine/src/block_manager.py

@@ -0,0 +1,90 @@
+"""Memory block helpers used by the KV cache implementation."""
+
+from __future__ import annotations
+
+from typing import Iterable, List, Sequence
+
+import torch
+from mindspore import Parameter, float32, tensor
+from mindspore.common.initializer import Zero, initializer
+from mindspore.ops import ScatterNdUpdate, gather, stack
+
+
+class BlockAllocator:
+    """Stores cache data in fixed-size MindSpore blocks."""
+
+    def __init__(self, num_blocks: int, block_size: int, head_dim: int) -> None:
+        self.num_blocks = num_blocks
+        self.pool = [
+            Parameter(initializer(Zero(), [block_size, head_dim], float32))
+            for _ in range(num_blocks)
+        ]
+
+    def alloc(self, indices: Sequence[Sequence[int]], values: Iterable[torch.Tensor]) -> None:
+        """Write one token per head into the allocated block positions."""
+        for index, value in zip(indices, values):
+            block_index, offset = index
+            self.pool[block_index] = ScatterNdUpdate()(
+                self.pool[block_index],
+                tensor([[offset]]),
+                tensor(value.detach().numpy()).unsqueeze(0),
+            )
+
+    def get(self, indices: Sequence[Sequence[int]]):
+        """Read back a stacked tensor for the provided block coordinates."""
+        values = []
+        for block_index, offset in indices:
+            values.append(gather(self.pool[block_index], tensor(offset), 0))
+        return stack(values)
+
+
+class MetadataEngine:
+    """Tracks token-to-block mappings for all attention heads."""
+
+    def __init__(self, num_blocks: int, num_heads: int, block_size: int) -> None:
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads
+        self.block_size = block_size
+        self.block_table: List[List[List[int]]] = []
+        self.num_token = 0
+        self.block_cursor = 0
+        self.offset_cursor = 0
+
+    def alloc(self) -> List[List[int]]:
+        """Allocate one position for each head for the next token."""
+        token_indices: List[List[int]] = []
+        self.num_token += 1
+
+        for _ in range(self.num_heads):
+            if self.block_cursor == self.num_blocks:
+                raise ValueError("Out of memory in block allocator.")
+
+            token_indices.append([self.block_cursor, self.offset_cursor])
+            self.offset_cursor += 1
+
+            if self.offset_cursor == self.block_size:
+                self.block_cursor += 1
+                self.offset_cursor = 0
+
+        self.block_table.append(token_indices)
+        return token_indices
+
+    def get(self, head_index: int) -> List[List[int]]:
+        """Return all stored positions for a single attention head."""
+        return [self.block_table[token_index][head_index] for token_index in range(self.num_token)]
+
+
+class MindsporeAPI:
+    """Thin wrappers that keep tensor conversion logic in one place."""
+
+    @staticmethod
+    def union(values):
+        return stack(values)
+
+    @staticmethod
+    def transpose(values):
+        return values.transpose(-1, -2)
+
+    @staticmethod
+    def matmul(x: torch.Tensor, values) -> torch.Tensor:
+        return torch.from_numpy(tensor(x.detach().numpy()).matmul(values).asnumpy())

engine/src/demo_data.py

@@ -0,0 +1,25 @@
+"""Shared helper functions for the demo scripts."""
+
+from __future__ import annotations
+
+from torch import LongTensor
+
+
+TARGET_VOCAB = {"P": 0, "i": 1, "want": 2, "a": 3, "beer": 4, "S": 5, "E": 6, "U": 7}
+INDEX_TO_TOKEN = {index: token for token, index in TARGET_VOCAB.items()}
+
+
+def make_training_batch(sentences):
+    output_batch = []
+    target_batch = []
+    for source_text, target_text in sentences:
+        output_batch.append([TARGET_VOCAB[token] for token in source_text.split()])
+        target_batch.append([TARGET_VOCAB[token] for token in target_text.split()])
+    return LongTensor(output_batch), LongTensor(target_batch)
+
+
+def make_inference_batch(sentences):
+    output_batch = []
+    for sentence in sentences:
+        output_batch.append([TARGET_VOCAB[token] for token in sentence.split()])
+    return LongTensor(output_batch)

engine/src/kv_cache.py

@@ -0,0 +1,101 @@
+"""KV cache backends used by the decoder-only transformer demo."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+import torch
+
+from .block_manager import BlockAllocator, MetadataEngine, MindsporeAPI
+
+
+class CacheBackend(Protocol):
+    """Protocol implemented by all cache backends."""
+
+    def write(self, values: torch.Tensor) -> None:
+        ...
+
+    def transposed_matmul(self, x: torch.Tensor) -> torch.Tensor:
+        ...
+
+    def matmul(self, x: torch.Tensor) -> torch.Tensor:
+        ...
+
+
+@dataclass
+class TorchCacheBackend:
+    """Simple PyTorch-backed cache used as a reference implementation."""
+
+    num_blocks: int
+    num_heads: int
+    block_size: int
+    head_dim: int
+
+    def __post_init__(self) -> None:
+        self.cache = torch.empty(self.num_heads, 0, self.head_dim)
+
+    def write(self, values: torch.Tensor) -> None:
+        self.cache = torch.cat([self.cache, values.unsqueeze(1)], dim=1)
+
+    def transposed_matmul(self, x: torch.Tensor) -> torch.Tensor:
+        return x.unsqueeze(1).matmul(self.cache.transpose(-1, -2)).squeeze(1)
+
+    def matmul(self, x: torch.Tensor) -> torch.Tensor:
+        return x.unsqueeze(1).matmul(self.cache).squeeze(1)
+
+
+class MindsporeCacheBackend:
+    """MindSpore-backed cache that stores values in fixed-size blocks."""
+
+    def __init__(self, num_blocks: int, num_heads: int, block_size: int, head_dim: int) -> None:
+        self.block_allocator = BlockAllocator(num_blocks, block_size, head_dim)
+        self.metadata_engine = MetadataEngine(num_blocks, num_heads, block_size)
+        self.num_heads = num_heads
+
+    def write(self, values: torch.Tensor) -> None:
+        self.block_allocator.alloc(self.metadata_engine.alloc(), values)
+
+    def _stack_cached_values(self):
+        return MindsporeAPI.union(
+            [self.block_allocator.get(self.metadata_engine.get(head_index)) for head_index in range(self.num_heads)]
+        )
+
+    def transposed_matmul(self, x: torch.Tensor) -> torch.Tensor:
+        cached_values = self._stack_cached_values()
+        return MindsporeAPI.matmul(
+            x.unsqueeze(1),
+            MindsporeAPI.transpose(cached_values),
+        ).squeeze(1)
+
+    def matmul(self, x: torch.Tensor) -> torch.Tensor:
+        return MindsporeAPI.matmul(x.unsqueeze(1), self._stack_cached_values()).squeeze(1)
+
+
+class Cache:
+    """Facade that exposes a consistent cache API to the transformer."""
+
+    def __init__(
+        self,
+        num_blocks: int,
+        num_heads: int,
+        block_size: int,
+        head_dim: int,
+        backend: str = "mindspore",
+    ) -> None:
+        if backend == "torch":
+            self.backend: CacheBackend = TorchCacheBackend(num_blocks, num_heads, block_size, head_dim)
+        else:
+            self.backend = MindsporeCacheBackend(num_blocks, num_heads, block_size, head_dim)
+
+    def write(self, values: torch.Tensor) -> None:
+        self.backend.write(values)
+
+    def transposed_matmul(self, x: torch.Tensor) -> torch.Tensor:
+        return self.backend.transposed_matmul(x)
+
+    def matmul(self, x: torch.Tensor) -> torch.Tensor:
+        return self.backend.matmul(x)
+
+    def delete(self) -> None:
+        del self.backend