Skip to content

FUTURE_WORK.md

Latest commit

 

History

History
85 lines (63 loc) · 2.34 KB

FUTURE_WORK.md

File metadata and controls

85 lines (63 loc) · 2.34 KB
Raw

Future Work: Dynamic Expert Loading for Kimi-K2

Proposed Architecture

Overview

Instead of compressing the model, we keep all 1.07T parameters but dynamically load only what is needed for each inference.

Key Components

  1. Base Model (Always in GPU): ~10GB

    • Embeddings
    • Layer normalization
    • Attention mechanisms
    • Expert routing networks

  2. Expert Storage (CPU/Disk): ~1TB

    • All 23,424 experts (61 layers × 384 experts)
    • Each expert: ~45GB / 384 ≈ 117MB
    • Stored in quantized format (INT8/INT4)

  3. Dynamic Loading System

    • GPU cache: 32-64 experts (~4-8GB)
    • CPU cache: 256-512 experts (~30-60GB)
    • Disk storage: All remaining experts

Implementation Plan

Phase 1: Proof of Concept (1 week)

  • Implement basic expert loading mechanism
  • Test with single layer (384 experts)
  • Measure loading latency and throughput

Phase 2: Optimization (2 weeks)

  • Implement LRU caching strategy
  • Add predictive preloading
  • Optimize CPU-GPU transfer pipeline

Phase 3: Production Ready (2 weeks)

  • Full model integration
  • Performance tuning
  • API development

Expected Performance

Metric Full Model Dynamic Loading
GPU Memory 1.07TB 40-50GB
Inference Speed 100% 30-50%
Model Quality 100% 100%
First Token Latency 50ms 200-500ms

Technical Requirements

  1. Hardware

    • GPU: H100 80GB (via CloudExe)
    • CPU RAM: 2TB (available)
    • Fast SSD: 2TB+ (needed)

  2. Software

    • Custom PyTorch extensions
    • Efficient serialization (SafeTensors)
    • Async I/O handling

Research Questions

  1. Can we predict which experts will be needed based on context?
  2. What is the optimal cache size vs performance tradeoff?
  3. Can we cluster related experts for faster bulk loading?
  4. How does quantization affect expert routing decisions?

Alternative Approaches

  1. Hybrid Static-Dynamic

    • Keep top 10% most-used experts always in GPU
    • Dynamically load the remaining 90%

  2. Hierarchical Experts

    • Group experts into clusters
    • Load entire clusters instead of individual experts

  3. Speculative Loading

    • Predict next N tokens expert needs
    • Preload in parallel with current computation

This approach preserves the full capabilities of Kimi-K2 while making it deployable on reasonable hardware.