Getting Started with Keys & Caches

Get up and running with Keys & Caches in under 5 minutes! This guide will walk you through installation, your first GPU job, and essential concepts.

Prerequisites

Python 3.8+ (Python 3.9+ recommended)
PyTorch installed in your environment
Basic familiarity with PyTorch and command line

Installation

Install via pip

pip install kandc

Verify Installation

kandc --version

You should see something like:

Keys & Caches CLI v0.1.0

Quick Example

Copy and paste this simple example:

#!/usr/bin/env python3
import torch
import torch.nn as nn
from kandc import capture_model_instance

def create_simple_model():
    model = nn.Sequential(
        nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 1), nn.Sigmoid()
    )
    return model

def main():
    print("Creating model instance...")
    model = create_simple_model()
    print(f"   Original model: {type(model).__name__}")

    model = capture_model_instance(
        model, model_name="SimpleSequential", record_shapes=True, profile_memory=True
    )
    print(f"   Wrapped model: {type(model).__name__}")

    print("\nRunning forward passes...")
    batch_sizes = [1, 4, 8]

    for i, batch_size in enumerate(batch_sizes, 1):
        print(f"   Pass {i}: batch_size={batch_size}")
        x = torch.randn(batch_size, 10)
        with torch.no_grad():
            output = model(x)
        print(f"   Output shape: {output.shape}")

    print("\nAll forward passes completed!")

if __name__ == "__main__":
    main()

Then run:

kandc python your_script.py

Interested in more examples? Check out the examples page.