Technology · Framework

PyTorch

ActiveKV extends PyTorch's XPU backend and ATen from native code; the capabilities model is a PyTorch fine-tune.

ActiveKV · Capabilities modelNative ATen operatorsTransformers training
Below the framework · on top of it
  • 22source files registering native operators in PyTorch's own registry
  • ~596 Mparameters in the open-weight base model
  • ≤ 14 minper full fine-tune, the shortest about two
  • 17.6 GiBpeak memory on one NVIDIA Blackwell GPU

I work with PyTorch from both ends: as a native extension that must be indistinguishable from the code it replaced, and as a training tool where the recipe, the size and the rebuild cost are measured choices. I work below the framework when the framework is in the way.

Native operators inside vLLM

In ActiveKV a native engine manages a model's key-value cache inside vLLM, which is built on PyTorch, so the engine has to speak PyTorch's tensor model exactly. I built it in C++ against ATen, PyTorch's tensor library, and registered its operators with PyTorch's own operator registry, so vLLM's attention path calls them as first-class operators. When the runtime moved from Python to C++ it had to stay indistinguishable from the Python version, and bitwise comparison proved it: Rewriting an engine extension without changing a token.

Python extension
=
C++ engine on ATen
Compared: shapes, data types, strides, lifetimes, the device stream vLLM owns, even exception types and their order.

A profile that changed what it measured

PyTorch's own timing events taught me a lesson. Timing-enabled events serialised the copy engine they were meant to observe, and I retracted a conclusion built on them: The observer that changed the answer.

With timing-enabled events

copies run one after another

Uninstrumented

the same copies overlap

Profiles built on framework events are checked against an uninstrumented run before I believe them.

A full fine-tune that rebuilds in minutes

In the Capabilities model I needed a model that knows a tool catalog by heart and can be rebuilt whenever the catalog changes, so I trained in PyTorch with Hugging Face Transformers and fine-tuned every parameter of a small open-weight base model. At that size a full fine-tune finishes in minutes, so I skipped adapters and their extra layer. The recipe is AdamW with a cosine schedule, full-precision master weights with BF16 autocast and gradient checkpointing.

A measured ladder chose the size

Three model sizes by several epoch counts, each cell trained and scored.

Training is cheap and scoring is the slow step, which makes a rebuild routine. Quantization afterwards takes about fifteen seconds and reproduces byte for byte. See A hundred tokens instead of forty-two thousand. I also used PyTorch on Intel's XPU stack to produce a published 4-bit quantization of a 27B model: A 27B model on one 32 GB Intel card. PyTorch, Transformers and the open-weight base model are third-party; the extension, recipe and measurements are mine.