Skip to content

[Feature Request] Support for TurboQuant (PolarQuant + QJL) for 6x KV Cache Compression #799

Description

@rbf22

Is your feature request related to a problem? Please describe. WebLLM is a powerful tool for local-first AI, but it is often constrained by browser VRAM limits, especially when dealing with long-context windows. Maintaining a large Key-Value (KV) cache for long conversations or document processing significantly limits the maximum context length and slows down inference on consumer-grade hardware.

Describe the solution you'd like. I would like to see support for TurboQuant, a new quantization framework recently introduced by Google Research. TurboQuant uses two key techniques to achieve extreme compression:

  1. PolarQuant: Converts Cartesian coordinates to polar coordinates (radius and angles) to eliminate memory overhead from data normalization.
  2. QJL (Quantized Johnson-Lindenstrauss): A "1-bit trick" that uses a mathematical transform to eliminate residual errors from the first stage.

Integrating these kernels into the WebLLM/TVM Unity pipeline would allow for 3-bit or 4-bit KV cache quantization with zero to negligible accuracy loss.

Describe alternatives you've considered. Currently, WebLLM supports standard weight quantization (4-bit/3-bit), but extreme KV cache compression is still a major bottleneck for long-context performance.

Additional context According to Google Research, TurboQuant provides:

  • 6x Reduction in KV cache memory footprint.
  • Up to 8x Performance Speedup in computing attention logits on modern hardware (like H100s, with significant benefits expected for WebGPU).
  • Data-Oblivious: It does not require dataset-specific tuning or model retraining.

References:

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions