Upload folder using huggingface_hub

README.md

@@ -8,17 +8,16 @@ tags:
 library_name: sparsevlm
 ---
 
-# SparseVLM — Production Inference Acceleration for Vision-Language Models
+# SparseVLM
 
+[![PyPI](https://img.shields.io/pypi/v/sparsevlm)](https://pypi.org/project/sparsevlm/)
 [![Paper](https://img.shields.io/badge/ICML_2025-Paper-blue)](https://arxiv.org/abs/2410.04417)
 [![License](https://img.shields.io/badge/License-Apache_2.0-green)](LICENSE)
 [![Tests](https://github.com/aryanchauhan31/SparseVLM/actions/workflows/tests.yml/badge.svg)](https://github.com/aryanchauhan31/SparseVLM/actions)
 
-Training-free visual token sparsification for Qwen2.5-VL.
-**2–4× faster inference. <3% accuracy drop. One function call.**
+Training-free visual token pruning for Qwen2.5-VL. Scores visual tokens by how much text attends to them, prunes the unimportant ones from the KV cache, and decodes with the smaller cache.
 
-Based on the ICML 2025 paper by Zhang et al.:
-[SparseVLM: Visual Token Sparsification for Efficient VLM Inference](https://arxiv.org/abs/2410.04417)
+Based on [SparseVLM: Visual Token Sparsification for Efficient VLM Inference](https://arxiv.org/abs/2410.04417) (ICML 2025).
 
 ---
 
@@ -28,7 +27,7 @@ Based on the ICML 2025 paper by Zhang et al.:
 pip install sparsevlm
 ```
 
-**Requirements:** Python 3.10+, PyTorch 2.1+, Triton 2.1+
+Requirements: Python 3.10+, PyTorch 2.1+, transformers 4.49+
 
 ---
 
@@ -38,28 +37,31 @@ pip install sparsevlm
 import torch
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
 from sparsevlm import sparsevlm_generate
+from PIL import Image
 
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     "Qwen/Qwen2.5-VL-7B-Instruct",
     torch_dtype=torch.bfloat16,
     device_map="auto",
-    attn_implementation="eager",   # required for attention-weight scoring
+    attn_implementation="eager",
 )
 processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
 
-# Prepare inputs normally
+image = Image.open("your_image.jpg")
 messages = [{"role": "user", "content": [
     {"type": "image", "image": image},
-    {"type": "text",  "text": "Describe this image."}
+    {"type": "text",  "text": "Describe this image in detail."}
 ]}]
 text   = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 inputs = processor(text=[text], images=[image], return_tensors="pt").to("cuda")
 
-# Run SparseVLM — keeps top-64 visual tokens out of 256 (25%)
+# count visual tokens
+n_vis = int((inputs["image_grid_thw"][0].prod() / 4).item())
+
 output = sparsevlm_generate(
     model, processor, inputs,
-    n_vis=256,          # visual tokens in your sequence
-    keep_n_vis=64,      # keep 25% — tune this
+    n_vis=n_vis,
+    keep_n_vis=n_vis // 4,   # keep 25% of visual tokens
     max_new_tokens=256,
 )
 print(processor.decode(output[0][1:], skip_special_tokens=True))
@@ -67,53 +69,126 @@ print(processor.decode(output[0][1:], skip_special_tokens=True))
 
 ---
 
-## Benchmark
+## Benchmark results
+
+Measured on **NVIDIA A100-SXM4-40GB**, Qwen2.5-VL-7B-Instruct, bfloat16, SDPA attention.
 
-A100 40GB, Qwen2.5-VL-7B-Instruct, batch size 1.
-**Replace these with your numbers from `python benchmark/bench_layer1.py`.**
+### Real photo — Fuji mountain + Milky Way (4928×2773px, 16320 visual tokens)
 
-| Tokens retained | Latency | Speedup | MME | TextVQA |
+| Config | Tokens kept | Time | Speedup | Output quality |
 |---|---|---|---|---|
-| 256 (100%) | 48ms | 1.0× | 100% | 100% |
-| 128 (50%)  | 22ms | 2.2× | 98.2% | 97.6% |
-| 96  (37%)  | 18ms | 2.7× | 97.1% | 96.4% |
-| 64  (25%)  | 14ms | 3.4× | 95.3% | 94.1% |
+| Baseline | 16320 (100%) | 9738ms | 1.00× | ✅ Identifies Fuji, Milky Way, snow cap, star colors |
+| SparseVLM 50% | 8192 | 9441ms | 1.03× | ✅ Same quality |
+| SparseVLM 25% | 4080 | 9297ms | 1.05× | ✅ All key details preserved |
+| SparseVLM 10% | 1632 | 9425ms | 1.03× | ✅ Still correctly describes scene |
+
+> **Key result:** Full 4K image (16K tokens) runs without OOM. Without SparseVLM's hook-based scoring, the 16K-token image requires materialising a 15GB attention matrix and crashes. The scorer computes only the text→visual submatrix (35 × 16320 = 32MB instead of 15GB).
+
+### Resized photo (896×504px, 576 visual tokens), batch=1
+
+| Tokens kept | Time | Speedup |
+|---|---|---|
+| 576 (100%) | 2167ms | 1.00× |
+| 288 (50%)  | 1685ms | 1.29× |
+| **144 (25%)** | **1565ms** | **1.39×** |
+| 72  (12%)  | 1620ms | 1.34× |
+
+### When to expect larger speedup
+
+Speedup grows when the KV cache is large relative to model weights:
+
+| Scenario | Expected speedup |
+|---|---|
+| Single image, short generation | ~1.1–1.4× |
+| Single image, 256+ output tokens | ~1.5–2.5× |
+| Batch=32, high-res images | ~2–4× |
+| Very long visual context (10K+ tokens) | ~2–4× |
 
 ---
 
 ## How it works
 
-SparseVLM hooks into the LLM decoder's attention layers and reuses
-attention weights the model already computes — zero extra parameters.
+### Token scoring (no extra parameters)
 
-At each target layer:
-1. **Rater selection** — text tokens with above-average visual attention
-2. **Visual token scoring** — sum of rater attention per visual token
-3. **Rank-adaptive pruning** — rank(A_rater) sets the pruning ratio
-4. **Token recycling** — pruned tokens clustered into compact representations
+At decoder layer 2, a lightweight hook intercepts the attention projection and computes:
 
-Three-layer optimisation stack:
-- **Layer 1** — Triton sparse attention kernel + sketch rank (15-50× faster than SVD)
-- **Layer 2** — FlashAttention varlen, variable-length packing (no padding waste)
-- **Layer 3** — CUDA graph bucketing (zero kernel-launch overhead)
+```
+A_tv = Q_text @ K_visual^T   # only the text→visual submatrix
+                              # 35 × 16320 instead of 16320 × 16320
+score_i = sum over text tokens of attention to visual token i
+```
+
+Visual tokens with high scores are important to the text query. Low-score tokens are pruned from the KV cache before decoding starts.
+
+### KV cache pruning
+
+After scoring, the KV cache is sliced to keep only the top-K visual entries plus all text entries. The model then decodes with a smaller cache — fewer keys to attend over per decode step.
+
+```
+Prefill:  build KV cache for all 16320 visual tokens
+Score:    rank each visual token by text attention (32MB op)
+Prune:    keep top-K, drop the rest
+Decode:   attend over K + N_text keys instead of 16320 + N_text
+```
+
+### Position fix (`rope_deltas`)
+
+After pruning, Qwen2.5-VL's internal position counter (`rope_deltas`) is adjusted so decode tokens get correct positional embeddings despite the shorter cache.
 
 ---
 
-## Configuration
+## API
+
+### `sparsevlm_generate`
 
 ```python
-state = apply_sparsevlm(
-    model,
-    n_vis=256,          # visual tokens per image
-    target_layers=None, # default: every 4th layer from layer 2
-    min_keep=32,        # never prune below this
-    tau=0.5,            # recycling fraction
-    theta=0.5,          # cluster ratio
+from sparsevlm import sparsevlm_generate
+
+output = sparsevlm_generate(
+    model,                  # Qwen2_5_VLForConditionalGeneration
+    processor,              # AutoProcessor
+    inputs,                 # dict from processor(...)
+    n_vis,                  # total visual tokens in the sequence
+    keep_n_vis,             # how many to keep (e.g. n_vis // 4 for 25%)
+    max_new_tokens=256,     # generation length
+    target_layer=2,         # which layer to score from (default 2)
+    device="cuda",          # primary device
 )
+# returns: token ids [B, max_new_tokens]
+```
+
+### `apply_sparsevlm` / `remove_hooks` (hook-based API)
+
+```python
+from sparsevlm import apply_sparsevlm, reset_n_vis, remove_hooks
+
+state = apply_sparsevlm(model, n_vis=256)
+reset_n_vis(state, n_vis=256)   # call before each generate
+output = model.generate(...)
+remove_hooks(state)
 ```
 
 ---
 
+## Model support
+
+| Model | Status |
+|---|---|
+| Qwen/Qwen2.5-VL-7B-Instruct | ✅ Tested |
+| Qwen/Qwen2.5-VL-3B-Instruct | ✅ Should work |
+| Qwen/Qwen2.5-VL-72B-Instruct | ✅ Should work |
+| Qwen/Qwen2-VL-* | ✅ Legacy support |
+
+---
+
+## Limitations
+
+- Requires `attn_implementation="eager"` or `"sdpa"`. Flash Attention 2 (separate package) is not required.
+- Speedup is modest (~1.1–1.4×) for single-image, short-generation use cases. The gain comes from long generations, high-resolution images, or batched serving.
+- Currently tested with Qwen2.5-VL. Other VLM families would need architecture-specific adaptation.
+
+---
+
 ## Citation
 
 ```bibtex
@@ -127,8 +202,4 @@ state = apply_sparsevlm(
 }
 ```
 
----
-
-## License
-
-Apache 2.0
+Apache 2.0 license.

dist/sparsevlm-0.1.2.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a38ff574431f01b8a0d79a7525a8e82fc80a67d48a2d22a44f20e346f1a145b5
+size 19535

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sparsevlm"
-version = "0.1.2"
+version = "0.1.3"
 description = "Training-free visual token sparsification for vision-language models (ICML 2025)"
 readme = "README.md"
 license = { text = "Apache-2.0" }

sparsevlm.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sparsevlm
-Version: 0.1.1
+Version: 0.1.2
 Summary: Training-free visual token sparsification for vision-language models (ICML 2025)
 Author-email: Aryan Chauhan <chauhanaryan31801@gmail.com>
 License: Apache-2.0
@@ -67,7 +67,7 @@ pip install sparsevlm
 ```python
 import torch
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-from sparsevlm import apply_sparsevlm, reset_n_vis, remove_hooks
+from sparsevlm import sparsevlm_generate
 
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     "Qwen/Qwen2.5-VL-7B-Instruct",
@@ -77,16 +77,22 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 )
 processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
 
-# Enable SparseVLM — no retraining needed
-state = apply_sparsevlm(model, n_vis=256)
-
-# Reset before each new image forward pass
-reset_n_vis(state, n_vis=256)
-inputs = processor(images=image, text=prompt, return_tensors="pt").to("cuda")
-output = model.generate(**inputs, max_new_tokens=256)
-
-# Remove hooks when done
-remove_hooks(state)
+# Prepare inputs normally
+messages = [{"role": "user", "content": [
+    {"type": "image", "image": image},
+    {"type": "text",  "text": "Describe this image."}
+]}]
+text   = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+inputs = processor(text=[text], images=[image], return_tensors="pt").to("cuda")
+
+# Run SparseVLM — keeps top-64 visual tokens out of 256 (25%)
+output = sparsevlm_generate(
+    model, processor, inputs,
+    n_vis=256,          # visual tokens in your sequence
+    keep_n_vis=64,      # keep 25% — tune this
+    max_new_tokens=256,
+)
+print(processor.decode(output[0][1:], skip_special_tokens=True))
 ```
 
 ---

sparsevlm.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ kernels/sparse_attn.py
 kernels/token_scorer.py
 kernels/varlen_packing.py
 sparsevlm/__init__.py
+sparsevlm/generate.py
 sparsevlm/patch.py
 sparsevlm/scheduler.py
 sparsevlm.egg-info/PKG-INFO

sparsevlm/__init__.py

@@ -46,4 +46,4 @@ def apply_sparsevlm(
 
 __all__ = ["apply_sparsevlm", "reset_n_vis", "unpatch_qwen2vl",
            "remove_hooks", "sparsevlm_generate"]
-__version__ = "0.1.2"
+__version__ = "0.1.3"