Add model card, metadata, and usage examples

README.md

@@ -0,0 +1,91 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+---
+
+# JetSpec: Parallel Tree Drafting
+
+JetSpec is an implementation of **parallel tree drafting** designed to accelerate speculative decoding in Large Language Models (LLMs) with up to 10x acceptance length.
+
+This repository contains the speculative draft head presented in the paper [JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting](https://arxiv.org/abs/2606.18394).
+
+- **Project Page:** [hao-ai-lab.github.io/JetSpec](https://hao-ai-lab.github.io/JetSpec)
+- **Repository:** [github.com/hao-ai-lab/JetSpec](https://github.com/hao-ai-lab/JetSpec)
+
+## Method Overview
+
+Speculative decoding is fast when the target accepts many draft tokens and drafting remains cheap. JetSpec keeps the **one-pass drafting efficiency and restores causal branch conditioning**. The draft head reads fused hidden states from the frozen target and emits per-depth logits in a single parallel pass. Tree construction spends a draft budget over high-probability branches, and the target verifies every node in one batched/tree-masked forward.
+
+## Installation
+
+To run JetSpec, clone the repository and install the package:
+
+```bash
+git clone https://github.com/hao-ai-lab/JetSpec.git
+cd JetSpec
+pip install -e '.[bench,kernel]'
+```
+
+## Usage
+
+### Hugging Face Reference
+
+You can run speculative decoding references using the `LLM` entrypoint:
+
+```python
+from jetspec import LLM, SamplingParams
+
+llm = LLM("Qwen/Qwen3-8B", attn_implementation="flash_attention_2")
+out = llm.generate(
+    "The three primary colors are",
+    SamplingParams(temperature=0.0, max_new_tokens=64),
+)
+print(out["text"])
+```
+
+### JetSpec Inference Engine
+
+For optimal latency, the inference engine utilizes a Triton tree-attention backend, compiled verification, and CUDA graphs:
+
+```python
+from jetspec import load_draft_head, DraftHeadTreeDrafter
+from jetspec.inference_engine import JetSpecEngine, SamplingParams
+
+engine = JetSpecEngine(
+    "Qwen/Qwen3-8B",
+    attn_backend="triton_paged_tree_compiled",
+)
+head = load_draft_head("JetSpec/jetspec-qwen3-8b")
+drafter = DraftHeadTreeDrafter(
+    head,
+    target=engine.model,
+    block_size=head.block_size,
+    target_layer_ids=head.target_layer_ids,
+)
+out = engine.generate_tree(
+    "The three primary colors are",
+    drafter,
+    block_size=head.block_size,
+    tree_depth=20,
+    tree_width=7,
+    budget=128,
+    target_layer_ids=head.target_layer_ids,
+    sampling_params=SamplingParams(temperature=0.0, max_new_tokens=64),
+)
+print(out["text"])
+print("tokens per forward:", out["tpf"])
+```
+
+## Citation
+
+```bibtex
+@inproceedings{jetspec2026,
+  title = {JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting},
+  author = {Hu, Lanxiang and Feng, Zhaoxiang and Wu, Yulun and Yuan, Haoran and Zhao, Yujie and Qian, Yu-Yang and Wang, Bojun and Zhao, Peng and Jiang, Daxin ... [et al.]},
+  year = {2026},
+  url = {https://arxiv.org/abs/2606.18394},
+  eprint = {2606.18394},
+  note = {Preprint}
+}
+```