TurboQuant TQ4 KV cache compression for Qwen 3.5 MoE (pytorch#18687)

TurboQuant TQ4 KV cache compression for Qwen 3.5 MoE

Add TurboQuant (arXiv 2504.19874) KV cache compression to the CUDA
backend, reducing KV cache memory 3.8x by storing nibble-packed uint8
indices + bf16 norms instead of bf16 tensors. A fused Triton SDPA
kernel decompresses K/V per-tile in the attention inner loop so the
full cache is never materialized.

Components:

- backends/cuda/triton/kernels/tq4_sdpa.py: Fused TQ4 Flash Attention
  kernel with Pack GQA optimization (adapted from sdpa.py structure),
  precomputed [256]-entry bf16 LUT for zero-bit-op centroid gather,
  and norm factoring (norms multiplied on [M,N] QK/P matrices instead
  of [N,D] K/V tiles). NaN-safe softmax guards for sparse masks.
  Registered as @triton_op for torch.export + CUDA backend lowering.

- extension/llm/modules/turboquant/: TurboQuantKVCache nn.Module with
  bf16 compression path and self-contained Lloyd-Max codebook solver
  (no external dependencies; scipy used lazily for codebook init only).

- examples/models/qwen3_5_moe/: --turboquant flag in export.py, branch
  in FullAttention.forward() between standard SDPA and tq4_sdpa.

- backends/aoti/: Added aoti_torch_dtype_uint8 shim and Byte ScalarType
  to slim headers (required for uint8 KV cache tensors in C++ runtime).

Performance (A100, Qwen 3.5 MoE, B=1, GQA 16:2, D=256, seq=4096):
  TQ4 SDPA kernel: 0.66ms (was 3.74ms before optimizations)
  Baseline bf16:   0.45ms (1.5x overhead for 3.8x memory savings)
  Full AOTI path:  0.79ms (inductor fuses compress ops)

The full Qwen 3.5 MoE model with TurboQuant KV cache compression runs at
75% (78->60)of baseline decode speed with 3.8x memory savings (relative
to normal KVcache). End-to-end validated through the C++ runner. But
full attention kvcache is small part of Qwen3.5 MoE because 3/4 of it is
recurrent states. In 200K context length, turboquant saves about 3GB
(4GB -> 1GB) of kvcache.

Author: mergennachin

.ci/scripts/test_model_e2e.sh

@@ -354,7 +354,7 @@ EOF
     fi
     ;;
   qwen3_5_moe)
-    RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 32"
+    RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt 'What is the capital of France?' --max_new_tokens 128 --temperature 0"
     ;;
   voxtral_realtime)
     RUNNER_ARGS="--model_path ${MODEL_DIR}/model.pte --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --preprocessor_path ${MODEL_DIR}/$PREPROCESSOR --audio_path ${MODEL_DIR}/$AUDIO_FILE --temperature 0"

.github/workflows/cuda.yml

@@ -145,8 +145,8 @@ jobs:
         # Run CUDA backend Python tests
         python -m pytest backends/cuda/tests backends/cuda/passes/tests -v -o "addopts="
 
-        # Run quantize roundtrip tests (Qwen 3.5 MoE save/load prequantized)
-        python -m pytest examples/models/qwen3_5_moe/test_quantize_roundtrip.py -v -o "addopts="
+        # Run Qwen 3.5 MoE tests (quantize roundtrip + TurboQuant KV cache)
+        python -m pytest examples/models/qwen3_5_moe/test_quantize_roundtrip.py examples/models/qwen3_5_moe/test_turboquant.py -v -o "addopts="
 
   export-model-cuda-artifact:
     name: export-model-cuda-artifact

backends/aoti/common_shims_slim.cpp

@@ -134,6 +134,10 @@ int32_t aoti_torch_dtype_int8() {
   return 1; // ScalarType::Char
 }
 
+int32_t aoti_torch_dtype_uint8() {
+  return 0; // ScalarType::Byte
+}
+
 int32_t aoti_torch_dtype_bool() {
   return 11; // ScalarType::Bool
 }

backends/aoti/common_shims_slim.h

@@ -76,6 +76,7 @@ AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int64();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int32();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int16();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int8();
+AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_uint8();
 AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bool();
 
 // ============================================================

backends/aoti/slim/c10/core/ScalarType.h

@@ -23,7 +23,7 @@ using BFloat16 = ::executorch::runtime::etensor::BFloat16;
 /// Enum representing the scalar type (dtype) of tensor elements.
 /// Note: Enum values must match PyTorch's c10::ScalarType for compatibility.
 enum class ScalarType : int8_t {
-  // Byte = 0,     // uint8_t - not currently needed
+  Byte = 0, // uint8_t
   Char = 1, // int8_t
   Short = 2, // int16_t
   Int = 3, // int32_t
@@ -43,6 +43,7 @@ enum class ScalarType : int8_t {
 };
 
 // Type alias constants for convenience
+constexpr ScalarType kByte = ScalarType::Byte;
 constexpr ScalarType kChar = ScalarType::Char;
 constexpr ScalarType kShort = ScalarType::Short;
 constexpr ScalarType kInt = ScalarType::Int;
@@ -56,6 +57,8 @@ constexpr ScalarType kBFloat16 = ScalarType::BFloat16;
 /// @return The size in bytes of a single element.
 inline size_t elementSize(ScalarType t) {
   switch (t) {
+    case ScalarType::Byte:
+      return sizeof(uint8_t);
     case ScalarType::Char:
       return sizeof(int8_t);
     case ScalarType::Short:
@@ -80,6 +83,8 @@ inline size_t elementSize(ScalarType t) {
 /// @return The name of the scalar type.
 inline const char* toString(ScalarType t) {
   switch (t) {
+    case ScalarType::Byte:
+      return "Byte";
     case ScalarType::Char:
       return "Char";
     case ScalarType::Short:
@@ -114,6 +119,7 @@ inline bool isFloatingType(ScalarType t) {
 /// @return true if the scalar type is integral, false otherwise.
 inline bool isIntegralType(ScalarType t, bool includeBool) {
   switch (t) {
+    case ScalarType::Byte:
     case ScalarType::Char:
     case ScalarType::Short:
     case ScalarType::Int:
@@ -138,6 +144,7 @@ inline bool isBoolType(ScalarType t) {
 /// @return true if the scalar type is valid, false otherwise.
 inline bool isValidScalarType(ScalarType t) {
   switch (t) {
+    case ScalarType::Byte:
     case ScalarType::Char:
     case ScalarType::Short:
     case ScalarType::Int: