Echo-minn/MSpecKV

1

stars

5

commits

Python

primary language

Jan 22, 2026

updated

README

MSpecKV

Long-context inference in large language models is increasingly bottlenecked by key–value (KV) cache traffic, including offloading, transfer, and memory bandwidth costs, which scale linearly with context length and can dominate computation. While speculative decoding techniques reduce redundant computation, they retain full- precision KV caches and therefore fail to address this memory bottleneck. We present MSpecKV, a plug-and-play multilevel speculative decoding framework that augmented with KV-cache quantization and GPU-resident KV reuse to reduce KV movement without modifying attention kernels or retraining models. MSpecKV stores offloaded KV in quantized form and reconstructs only the necessary subsets on demand, while maintaining high speculative acceptance. Across prefill lengths from 8K to 32K, MSpecKV improves throughput by up to 14.1× over naive autoregressive decoding and up to 3.8× over non-quantized variant, with comparable acceptance rates. These results demonstrate that reducing KV-cache traffic is critical for scalable long-context serving, and that quantization combined with structured speculative decoding provides a practical and effective solution.

Recorded Video

The recorded video is available at: https://youtu.be/fJIsocqpJF4 model

Evaluation Results(A1)

================================================================================
SPEEDUP SUMMARY (latency-based)
================================================================================

 prefill | vanilla_vs_base | no_kv_quant_vs_base | kv_quant_vs_base | quant_vs_no_quant
--------------------------------------------------------------------------------------
    8192 |           1.17x |           3.25x |           4.75x |              1.46x
   16384 |           1.06x |           3.76x |           8.01x |              2.13x
   24576 |           1.23x |           4.50x |          12.94x |              2.87x
   32768 |           1.21x |           3.92x |          14.45x |              3.69x

Resident Layer Abalation(A2)

Run (final):

cd /opt/dlami/nvme/mira/.www/MSpecKV
for s in 1 2 3; do
  CUDA_VISIBLE_DEVICES=7 OMP_NUM_THREADS=48 torchrun --nproc_per_node=1 evaluation/resident_layer_ablation.py \
    --dataset demo --target llama-7B-128K \
    --budget 4096 --on_chip 16 --gamma 8 \
    --prefill 16384 --gen_len 256 \
    --resident_layers_values "0,8,16" \
    --seed $s --output_json auto --output_csv /dev/null
done

Throughput

resident_layersbaselinequantdeltaspeedup
012.42 ± 4.4013.21 ± 3.580.78 ± 1.171.09 ± 0.10
812.18 ± 4.5914.89 ± 4.312.71 ± 1.171.26 ± 0.17
168.50 ± 1.3124.12 ± 1.9415.62 ± 2.532.89 ± 0.56

On-Chip Layer Abalation(B1)

Run:

cd /opt/dlami/nvme/mira/.www/MSpecKV
for s in 1 2 3; do
  CUDA_VISIBLE_DEVICES=7 OMP_NUM_THREADS=48 torchrun --nproc_per_node=1 evaluation/on_chip_ablation.py \
    --dataset demo --target llama-7B-128K \
    --budget 4096 --gamma 8 \
    --prefill 16384 --gen_len 256 \
    --resident_layers 16 --on_chip_layers "4,8,12,16,18" \
    --seed $s --output_json auto --output_csv /dev/null
done

Throughput

on_chip_layersbaselinequantdeltaspeedup
46.72 ± 0.529.74 ± 0.853.02 ± 1.361.46 ± 0.25
87.52 ± 1.589.69 ± 1.552.18 ± 2.131.32 ± 0.35
128.19 ± 0.6810.89 ± 0.612.70 ± 0.741.33 ± 0.11
169.23 ± 1.0520.01 ± 2.4710.78 ± 1.962.17 ± 0.20
189.53 ± 0.7322.08 ± 1.9212.55 ± 1.842.32 ± 0.23

Output Length Sensitivity(B2)

Run:

cd /opt/dlami/nvme/mira/.www/MSpecKV
for s in 1 2 3; do
  CUDA_VISIBLE_DEVICES=7 OMP_NUM_THREADS=48 torchrun --nproc_per_node=1 evaluation/output_len_ablation.py \
    --dataset demo --target llama-7B-128K \
    --budget 4096 --on_chip 16 --gamma 8 --resident_layers 16 \
    --prefill 8192 --gen_lens "256,512,768,1024,1536,2048" \
    --seed $s --output_json auto --output_csv /dev/null
done

Throughput

xbaselinequantdeltaspeedup
25618.14 ± 1.0724.74 ± 1.736.61 ± 2.691.37 ± 0.17
51218.34 ± 1.5428.81 ± 4.9110.47 ± 3.731.57 ± 0.17
76819.94 ± 1.7225.71 ± 1.875.77 ± 3.071.30 ± 0.17
102419.86 ± 0.2427.75 ± 1.617.89 ± 1.401.40 ± 0.07
153619.87 ± 0.3128.03 ± 3.868.16 ± 3.791.41 ± 0.19
204819.42 ± 1.2631.12 ± 2.0811.70 ± 2.961.61 ± 0.19

Contributors

ApiaoSamaa

3 commits

Echo-minn

2 commits

Echo-minn/MSpecKV

1

stars

5

commits

Python

primary language

Jan 22, 2026

updated

README

MSpecKV

Long-context inference in large language models is increasingly bottlenecked by key–value (KV) cache traffic, including offloading, transfer, and memory bandwidth costs, which scale linearly with context length and can dominate computation. While speculative decoding techniques reduce redundant computation, they retain full- precision KV caches and therefore fail to address this memory bottleneck. We present MSpecKV, a plug-and-play multilevel speculative decoding framework that augmented with KV-cache quantization and GPU-resident KV reuse to reduce KV movement without modifying attention kernels or retraining models. MSpecKV stores offloaded KV in quantized form and reconstructs only the necessary subsets on demand, while maintaining high speculative acceptance. Across prefill lengths from 8K to 32K, MSpecKV improves throughput by up to 14.1× over naive autoregressive decoding and up to 3.8× over non-quantized variant, with comparable acceptance rates. These results demonstrate that reducing KV-cache traffic is critical for scalable long-context serving, and that quantization combined with structured speculative decoding provides a practical and effective solution.

Recorded Video

The recorded video is available at: https://youtu.be/fJIsocqpJF4 model

Evaluation Results(A1)

================================================================================
SPEEDUP SUMMARY (latency-based)
================================================================================

 prefill | vanilla_vs_base | no_kv_quant_vs_base | kv_quant_vs_base | quant_vs_no_quant
--------------------------------------------------------------------------------------
    8192 |           1.17x |           3.25x |           4.75x |              1.46x
   16384 |           1.06x |           3.76x |           8.01x |              2.13x
   24576 |           1.23x |           4.50x |          12.94x |              2.87x
   32768 |           1.21x |           3.92x |          14.45x |              3.69x

Resident Layer Abalation(A2)

Run (final):

cd /opt/dlami/nvme/mira/.www/MSpecKV
for s in 1 2 3; do
  CUDA_VISIBLE_DEVICES=7 OMP_NUM_THREADS=48 torchrun --nproc_per_node=1 evaluation/resident_layer_ablation.py \
    --dataset demo --target llama-7B-128K \
    --budget 4096 --on_chip 16 --gamma 8 \
    --prefill 16384 --gen_len 256 \
    --resident_layers_values "0,8,16" \
    --seed $s --output_json auto --output_csv /dev/null
done

Throughput

resident_layersbaselinequantdeltaspeedup
012.42 ± 4.4013.21 ± 3.580.78 ± 1.171.09 ± 0.10
812.18 ± 4.5914.89 ± 4.312.71 ± 1.171.26 ± 0.17
168.50 ± 1.3124.12 ± 1.9415.62 ± 2.532.89 ± 0.56

On-Chip Layer Abalation(B1)

Run:

cd /opt/dlami/nvme/mira/.www/MSpecKV
for s in 1 2 3; do
  CUDA_VISIBLE_DEVICES=7 OMP_NUM_THREADS=48 torchrun --nproc_per_node=1 evaluation/on_chip_ablation.py \
    --dataset demo --target llama-7B-128K \
    --budget 4096 --gamma 8 \
    --prefill 16384 --gen_len 256 \
    --resident_layers 16 --on_chip_layers "4,8,12,16,18" \
    --seed $s --output_json auto --output_csv /dev/null
done

Throughput

on_chip_layersbaselinequantdeltaspeedup
46.72 ± 0.529.74 ± 0.853.02 ± 1.361.46 ± 0.25
87.52 ± 1.589.69 ± 1.552.18 ± 2.131.32 ± 0.35
128.19 ± 0.6810.89 ± 0.612.70 ± 0.741.33 ± 0.11
169.23 ± 1.0520.01 ± 2.4710.78 ± 1.962.17 ± 0.20
189.53 ± 0.7322.08 ± 1.9212.55 ± 1.842.32 ± 0.23

Output Length Sensitivity(B2)

Run:

cd /opt/dlami/nvme/mira/.www/MSpecKV
for s in 1 2 3; do
  CUDA_VISIBLE_DEVICES=7 OMP_NUM_THREADS=48 torchrun --nproc_per_node=1 evaluation/output_len_ablation.py \
    --dataset demo --target llama-7B-128K \
    --budget 4096 --on_chip 16 --gamma 8 --resident_layers 16 \
    --prefill 8192 --gen_lens "256,512,768,1024,1536,2048" \
    --seed $s --output_json auto --output_csv /dev/null
done

Throughput

xbaselinequantdeltaspeedup
25618.14 ± 1.0724.74 ± 1.736.61 ± 2.691.37 ± 0.17
51218.34 ± 1.5428.81 ± 4.9110.47 ± 3.731.57 ± 0.17
76819.94 ± 1.7225.71 ± 1.875.77 ± 3.071.30 ± 0.17
102419.86 ± 0.2427.75 ± 1.617.89 ± 1.401.40 ± 0.07
153619.87 ± 0.3128.03 ± 3.868.16 ± 3.791.41 ± 0.19
204819.42 ± 1.2631.12 ± 2.0811.70 ± 2.961.61 ± 0.19

Contributors

ApiaoSamaa

3 commits

Echo-minn

2 commits

Languages

Python

100.0%