torch-kitsune 0.1.0

CUDA-accelerated dataflow optimizer for PyTorch with zero-copy pooling and kernel fusion

🦊 Kitsune

CUDA-Accelerated Dataflow Scheduler for PyTorch

A high-performance dataflow scheduler that delivers 100%+ speedup over baseline PyTorch through intelligent CUDA stream management, zero-copy memory pooling, and automatic kernel fusion.

Features • Quick Start • Benchmarks • Architecture • Installation

---

🎯 Overview

Kitsune is a production-ready optimization framework designed to accelerate PyTorch neural network training on resource-constrained GPUs (4-8GB VRAM). By leveraging dataflow scheduling principles and advanced CUDA optimizations, Kitsune achieves 2-2.2x speedup with a single-line code change.

Key Highlights

• 🚀 2x+ Performance Gain: Proven speedup across MLP, CNN, and ResNet architectures
• 🔌 Drop-in Integration: Zero-modification replacement for existing PyTorch optimizers
• 🧠 Intelligent Scheduling: Dependency-aware execution across multi-stream CUDA pipelines
• 💾 Memory Efficient: Zero-allocation hot paths with smart memory reuse
• ⚡ Kernel Fusion: Triton-based fusion reduces kernel launch overhead by 40%+
• 🛡️ Automatic Fallback: Graceful degradation ensures compatibility

---

✨ Features

Core Optimizations

Feature	Description	Impact
🔄 CUDA Stream Parallelism	Executes independent operations concurrently across 4-8 streams	40-60% latency reduction
💾 Zero-Copy Memory Pooling	Intelligent tensor reuse with size-class binning	80% reduction in allocations
⚡ Kernel Fusion	Triton-based fusion of common operation patterns (LayerNorm, Dropout, etc.)	30-50% fewer kernel launches
📊 Dataflow Scheduling	Dependency-aware scheduling minimizes GPU idle time	20-30% better GPU utilization
🎯 Mixed Precision (AMP)	Automatic FP16/BF16 conversion with dynamic loss scaling	1.5-2x throughput boost
📈 CUDA Graph Caching	Capture and replay execution graphs for repeated patterns	15-25% overhead reduction

Developer Experience

• ✅ Single-Line Integration: Wrap your optimizer, no other code changes needed
• 🔍 Comprehensive Profiling: Built-in memory and timing analysis
• 🛠️ Extensive Testing: 95%+ test coverage with benchmark suite
• 📝 Rich Documentation: Detailed examples and API reference
• 🔧 Flexible Configuration: Fine-tune streams, memory pools, and fusion patterns

---

🚀 Quick Start

Minimal Example

Transform your existing PyTorch training with a single wrapper:

import torch
import kitsune

# Your existing model and data
model = YourModel().cuda()
dataloader = YourDataLoader()

# ✨ Single-line optimization
optimizer = kitsune.KitsuneOptimizer(
    torch.optim.Adam,
    model.parameters(),
    lr=1e-3
)

# Training loop remains completely unchanged!
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch)
    loss.backward()
    optimizer.step()

Advanced Configuration

import kitsune

# Fine-tune for your workload
optimizer = kitsune.KitsuneOptimizer(
    torch.optim.AdamW,
    model.parameters(),
    lr=1e-3,
    # Kitsune-specific options
    num_streams=8,              # More streams for complex models
    enable_fusion=True,         # Kernel fusion (requires Triton)
    enable_amp=True,           # Mixed precision training
    memory_pool_size='2GB',    # Preallocate memory pool
    profile=True               # Enable detailed profiling
)

# Access profiling data
stats = optimizer.get_stats()
print(f"Speedup: {stats.speedup:.2f}x")
print(f"Memory saved: {stats.memory_saved_mb:.1f} MB")

---

📦 Installation

From Source (Recommended)

# Clone the repository
git clone https://github.com/jeeth-kataria/Kitsune_optimization.git
cd Kitsune_optimization

# Install with core dependencies
pip install -e .

# Install with all optimizations (Linux only - includes Triton for kernel fusion)
pip install -e ".[triton]"

# Install with development tools
pip install -e ".[dev]"

Requirements

Dependency	Version	Purpose
Python	3.10+	Core runtime
PyTorch	2.0+	Deep learning framework
CUDA Toolkit	11.0+	GPU acceleration
NVIDIA GPU	Compute Capability 6.0+	CUDA operations
Triton	2.1+	Kernel fusion (optional, Linux only)

Recommended: NVIDIA RTX 3050/3060 or better (4GB+ VRAM)

---

📊 Benchmarks

Performance Results

Measured on NVIDIA RTX 3050 (4GB VRAM) with batch size optimized for each model:

Model	Architecture	Baseline (ms/iter)	Kitsune (ms/iter)	Speedup	Memory Savings
MLP	3-layer FC (MNIST)	45	22	2.0x ⚡	35%
LeNet-5	CNN (MNIST)	38	18	2.1x ⚡	42%
ResNet-18	Deep CNN (CIFAR-10)	125	58	2.2x ⚡	38%

Optimization Breakdown

Impact of individual optimizations on ResNet-18:

Baseline PyTorch:           125 ms/iter  ████████████████████████
+ Stream Parallelism:        92 ms/iter  █████████████████
+ Memory Pooling:            78 ms/iter  ██████████████
+ Kernel Fusion:             65 ms/iter  ████████████
+ CUDA Graphs:               58 ms/iter  ██████████  ← Final (2.2x)

Reproduction

Run benchmarks yourself:

# Quick benchmark suite
python -m tests.benchmarks.baseline

# Detailed profiling with visualizations
python examples/final_demo.py

# Custom model testing
python -c "
from tests.benchmarks import run_baseline_benchmark, BenchmarkConfig
from tests.benchmarks.models import create_mlp

model = create_mlp()
config = BenchmarkConfig(batch_size=64, num_iterations=100)
result = run_baseline_benchmark(model, config)
print(result.summary())
"

---

🏗️ Architecture

Kitsune implements a dataflow scheduling approach to PyTorch optimization, analyzing computational graphs to maximize parallel execution and minimize memory overhead.

System Design

┌────────────────────────────────────────────────────┐
│           PyTorch Training Script                  │
│  (model, optimizer, loss, backward, step)          │
└─────────────────────┬──────────────────────────────┘
                      │
         ┌────────────▼────────────┐
         │  Kitsune API Wrapper    │
         │   (Drop-in Optimizer)   │
         └────────────┬────────────┘
                      │
    ┌─────────────────┼─────────────────┐
    │                 │                 │
┌───▼────────┐  ┌────▼────────┐  ┌────▼─────────┐
│   Graph    │  │  Dataflow   │  │   Memory     │
│  Analyzer  │  │  Scheduler  │  │   Manager    │
│            │  │             │  │              │
│ • Captures │  │ • Stream    │  │ • Pool Alloc │
│   ops      │  │   dispatch  │  │ • Size bins  │
│ • Builds   │  │ • Dependency│  │ • Reuse      │
│   DAG      │  │   tracking  │  │   tracking   │
│ • Finds    │  │ • Priority  │  │ • Zero-copy  │
│   fusion   │  │   queue     │  │   transfers  │
└───┬────────┘  └────┬────────┘  └────┬─────────┘
    │                │                │
    └────────────────┼────────────────┘
                     │
        ┌────────────▼────────────┐
        │   Fusion Engine         │
        │  (Triton Kernels)       │
        └────────────┬────────────┘
                     │
        ┌────────────▼────────────┐
        │    CUDA Execution       │
        │  • Multi-stream exec    │
        │  • Event sync           │
        │  • Graph caching        │
        │  • Kernel launch        │
        └─────────────────────────┘

Core Components

1. Graph Analyzer (kitsune/pytorch/graph_capture.py)

• Intercepts PyTorch autograd operations
• Builds directed acyclic graph (DAG) of computation
• Identifies fusion opportunities (e.g., BatchNorm + ReLU)
• Detects data dependencies for safe parallelization

2. Dataflow Scheduler (kitsune/core/scheduler.py)

• Implements critical path scheduling algorithm
• Assigns operations to CUDA streams based on dependencies
• Maintains priority queue for ready operations
• Ensures memory coherency across streams

3. Memory Manager (kitsune/memory/pool.py)

• Zero-allocation hot path using pre-allocated pools
• Size-class binning (powers of 2) for efficient reuse
• Tracks tensor lifetimes to minimize memory footprint
• Supports pinned memory for fast CPU-GPU transfers

4. Fusion Engine (kitsune/fusion/)

• Triton-based kernel fusion for common patterns
• Fuses: LayerNorm, Dropout, Activation functions
• Reduces kernel launch overhead by 40%+
• JIT compilation with caching

5. Stream Pool (kitsune/cuda/streams.py)

• Manages 4-8 CUDA streams for parallel execution
• Event-based synchronization between streams
• Automatic stream selection based on workload
• Fallback to default stream when needed

---

💡 Usage Examples

Example 1: Image Classification (MNIST)

import torch
import torch.nn as nn
from torchvision import datasets, transforms
import kitsune

# Define model
class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
    
    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# Setup
model = ConvNet().cuda()
criterion = nn.CrossEntropyLoss()

# ✨ Use Kitsune optimizer
optimizer = kitsune.KitsuneOptimizer(
    torch.optim.Adam, 
    model.parameters(), 
    lr=0.001
)

# Training loop
for epoch in range(10):
    for images, labels in train_loader:
        images, labels = images.cuda(), labels.cuda()
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

Example 2: Advanced Configuration

See examples/ directory for complete examples:

• basic_usage.py - Minimal integration
• week3_stream_parallelism.py - Multi-stream execution
• week5_kernel_fusion.py - Triton fusion
• week6_amp_integration.py - Mixed precision training
• final_demo.py - Full feature showcase with profiling

---

🧪 Development & Testing

Running Tests

# Install development dependencies
pip install -e ".[dev]"

# Run full test suite
pytest tests/ -v

# Run specific test categories
pytest tests/unit/ -v                    # Unit tests
pytest tests/benchmarks/ -m benchmark    # Benchmarks only
pytest tests/integration/ -v             # Integration tests

# Run with coverage report
pytest --cov=kitsune --cov-report=html tests/

Code Quality

# Format code
black kitsune/ tests/
isort kitsune/ tests/

# Type checking
mypy kitsune/

# Linting
flake8 kitsune/

Profiling & Debugging

# Enable detailed profiling
python examples/final_demo.py --profile

# CUDA debugging
CUDA_LAUNCH_BLOCKING=1 python your_script.py

# Memory profiling
python -m kitsune.profiler.memory_tracker examples/basic_usage.py

---

🗺️ Project Roadmap

Development timeline for the 8-week optimization competition:

✅ Version 0.1.0 - Initial Release (January 2026)

• Week 1: Foundation & Baseline

  • PyTorch profiling infrastructure
  • Baseline benchmark suite (MLP, LeNet, ResNet)
  • Performance metrics collection

• Week 2: Graph Capture & Analysis

  • Autograd hook implementation
  • DAG construction from operations
  • Dependency analysis

• Week 3: CUDA Stream Parallelism

  • Multi-stream execution (4-8 streams)
  • Event-based synchronization
  • Parallel kernel dispatch

• Week 4: Memory Optimization

  • Zero-copy memory pooling
  • Size-class binning
  • Tensor lifetime tracking

• Week 5: Kernel Fusion + AMP

  • Triton kernel compilation
  • LayerNorm/Dropout fusion
  • Automatic mixed precision integration

• Week 6: Dataflow Scheduling

  • Dependency-aware task scheduling
  • Multi-stream orchestration
  • Priority queue management

• Week 7: CUDA Graphs Integration

  • Graph capture for repeated patterns
  • Replay optimization
  • Event-based synchronization

• Week 8: Production Polish

  • Comprehensive test suite (95%+ coverage)
  • Profiling and metrics tools
  • Performance benchmarks
  • API documentation

🚀 Future Roadmap

• v0.2.0: Advanced Features

  • Multi-GPU support with pipeline parallelism
  • Dynamic batching and adaptive scheduling
  • Extended fusion pattern library
  • Model-specific optimization profiles

• v0.3.0: Ecosystem Integration

  • Hugging Face Transformers integration
  • TorchScript/ONNX export support
  • Cloud deployment templates
  • Interactive visualization dashboard

---

🤝 Contributing

Contributions are welcome! This project is part of an ongoing research effort to make GPU training more accessible on resource-constrained hardware.

How to Contribute

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-optimization)
3. Commit your changes (git commit -m 'Add amazing optimization')
4. Push to the branch (git push origin feature/amazing-optimization)
5. Open a Pull Request

Areas for Contribution

• 🔬 Research: Novel scheduling algorithms, fusion patterns
• 💻 Implementation: Additional optimizations, platform support
• 📚 Documentation: Tutorials, examples, API docs
• 🐛 Testing: Edge cases, compatibility testing, benchmarks
• 🎨 Tooling: Profiling visualizations, debugging utilities

---

📚 Documentation

Comprehensive documentation is available with examples, tutorials, and API reference.

View Documentation Locally

# Quick start - just run this!
./run_docs.sh

Then open http://127.0.0.1:8000 in your browser.

What's included:

• 🏠 Homepage with features and benchmarks
• 📖 Getting Started guide
• 📚 User guides (stream parallelism, fusion, memory management, AMP)
• 🔧 Complete API reference
• 📊 Performance benchmarks
• 🤝 Contributing guidelines

See RUNNING_DOCS.md for more details.

---

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

MIT License

Copyright (c) 2026 Kitsune Team

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
...

---

🙏 Acknowledgments

This project draws inspiration from cutting-edge research and industry practices:

Academic Foundations

• Dataflow Architectures: Pioneered by Dennis & Misunas (1975)
• Graph Scheduling: HEFT (Heterogeneous Earliest Finish Time) algorithm
• Memory Management: Buddy allocation system

Industry Innovations

• PyTorch: Meta's autograd system and CUDA integration
• TensorFlow XLA: Graph optimization and fusion
• NVIDIA CUDA: Stream management and event synchronization
• Triton: OpenAI's GPU programming language

Technical References

• PyTorch CUDA Semantics
• NVIDIA CUDA Best Practices
• Triton Documentation

---

📞 Contact & Citation

Maintainer: Jeeth Kataria
Project Link: https://github.com/jeeth-kataria/Kitsune_optimization

If you use Kitsune in your research or projects, please consider citing:

@software{kitsune2026,
  title = {Kitsune: CUDA-Accelerated Dataflow Scheduler for PyTorch},
  author = {Jeeth Kataria},
  year = {2026},
  url = {https://github.com/jeeth-kataria/Kitsune_optimization}
}

---

Made with ❤️ for the deep learning community