{"id":42485993,"url":"https://github.com/alepot55/flash-reasoning","last_synced_at":"2026-01-28T11:28:19.782Z","repository":{"id":334798912,"uuid":"1142826951","full_name":"alepot55/flash-reasoning","owner":"alepot55","description":"Tree-Aware Attention for System 2 Reasoning. Reduces KV-Cache VRAM by 96% and exceeds physical HBM bandwidth (1.33x) via Fused GQA Triton kernels.","archived":false,"fork":false,"pushed_at":"2026-01-26T23:37:06.000Z","size":492,"stargazers_count":0,"open_issues_count":1,"forks_count":0,"subscribers_count":0,"default_branch":"master","last_synced_at":"2026-01-27T09:42:38.508Z","etag":null,"topics":["attention-mechanism","fused-gqa","kv-cache","reasoning","triton"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/alepot55.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-01-26T22:40:38.000Z","updated_at":"2026-01-26T23:00:31.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/alepot55/flash-reasoning","commit_stats":null,"previous_names":["alepot55/flash-reasoning"],"tags_count":null,"template":false,"template_full_name":null,"purl":"pkg:github/alepot55/flash-reasoning","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alepot55%2Fflash-reasoning","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alepot55%2Fflash-reasoning/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alepot55%2Fflash-reasoning/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alepot55%2Fflash-reasoning/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/alepot55","download_url":"https://codeload.github.com/alepot55/flash-reasoning/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/alepot55%2Fflash-reasoning/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":28845088,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-01-28T10:53:21.605Z","status":"ssl_error","status_checked_at":"2026-01-28T10:53:20.789Z","response_time":57,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["attention-mechanism","fused-gqa","kv-cache","reasoning","triton"],"created_at":"2026-01-28T11:28:19.071Z","updated_at":"2026-01-28T11:28:19.778Z","avatar_url":"https://github.com/alepot55.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003ch1 align=\"center\"\u003e⚡ Flash-Reasoning\u003c/h1\u003e\n\n\u003cp align=\"center\"\u003e\n  \u003cstrong\u003eTree-Aware KV-Cache Attention for Reasoning LLMs\u003c/strong\u003e\n\u003c/p\u003e\n\n\u003cp align=\"center\"\u003e\n  \u003ca href=\"https://www.python.org/downloads/\"\u003e\u003cimg src=\"https://img.shields.io/badge/python-3.10%2B-blue?logo=python\u0026logoColor=white\" alt=\"Python 3.10+\"\u003e\u003c/a\u003e\n  \u003ca href=\"https://pytorch.org/\"\u003e\u003cimg src=\"https://img.shields.io/badge/PyTorch-2.4%2B-ee4c2c?logo=pytorch\u0026logoColor=white\" alt=\"PyTorch\"\u003e\u003c/a\u003e\n  \u003ca href=\"https://triton-lang.org/\"\u003e\u003cimg src=\"https://img.shields.io/badge/Triton-3.0%2B-76B900?logo=nvidia\u0026logoColor=white\" alt=\"Triton\"\u003e\u003c/a\u003e\n  \u003ca href=\"https://github.com/alepot55/flash-reasoning/blob/main/LICENSE\"\u003e\u003cimg src=\"https://img.shields.io/badge/license-MIT-blue.svg\" alt=\"License\"\u003e\u003c/a\u003e\n\u003c/p\u003e\n\n---\n\n## The Problem\n\nReasoning LLMs (DeepSeek-R1, o1, Tree-of-Thought) generate **decision trees** during inference. Standard engines treat sequences linearly → **O(n × b)** memory waste when `b` branches share a prefix.\n\n```mermaid\nflowchart LR\n    subgraph Standard[\"❌ Standard (vLLM)\"]\n        direction TB\n        B1[\"Branch 1: [Prefix | Suffix]\"] --\u003e M1[\"2112 tokens\"]\n        B2[\"Branch 2: [Prefix | Suffix]\"] --\u003e M2[\"2112 tokens\"]\n        B3[\"Branch 256: [Prefix | Suffix]\"] --\u003e M3[\"2112 tokens\"]\n    end\n    \n    subgraph Tree[\"✅ Flash-Reasoning\"]\n        direction TB\n        P[\"Shared Prefix\u003cbr/\u003e2048 tokens (1×)\"] --\u003e S1[\"Suffix 1\u003cbr/\u003e64 tok\"]\n        P --\u003e S2[\"Suffix 2\u003cbr/\u003e64 tok\"]\n        P --\u003e S3[\"Suffix 256\u003cbr/\u003e64 tok\"]\n    end\n    \n    Standard -.-\u003e|\"541K tokens\"| X[\"❌\"]\n    Tree -.-\u003e|\"18K tokens\"| Y[\"✅ 29× less\"]\n```\n\n---\n\n## Results\n\n**2.54× faster** | **96.6% less VRAM** | **L2 cache exploitation**\n\n### Speedup\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"assets/speedup.png\" alt=\"Speedup Chart\" width=\"700\"\u003e\n\u003c/p\u003e\n\n### Memory Bandwidth\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"assets/bandwidth.png\" alt=\"Bandwidth Chart\" width=\"700\"\u003e\n\u003c/p\u003e\n\n\u003e Effective bandwidth exceeds HBM limit (900 GB/s) because shared prefix blocks hit L2 cache (~5 TB/s).\n\n### VRAM Reduction\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"assets/vram.png\" alt=\"VRAM Chart\" width=\"700\"\u003e\n\u003c/p\u003e\n\n### Benchmark Table\n\n| Batch | Tree | Linear | Speedup | VRAM Reduction |\n|------:|-----:|-------:|--------:|---------------:|\n| 1 | 0.089 ms | 0.091 ms | 1.02× | 3% |\n| 16 | 0.184 ms | 0.301 ms | 1.64× | 74% |\n| 64 | 0.483 ms | 1.102 ms | 2.28× | 92% |\n| 128 | 0.912 ms | 2.281 ms | 2.50× | 95% |\n| **256** | **1.859 ms** | **4.729 ms** | **2.54×** | **96.6%** |\n\n---\n\n## Features\n\n| Feature | Description |\n|---------|-------------|\n| **Physical Prefix Sharing** | Branches share KV blocks via reference counting |\n| **Fused GQA Kernel** | K/V loaded once per KV-head group (4-8× traffic reduction) |\n| **Online Softmax** | FlashAttention-style O(1) memory per query |\n| **Triton Autotuning** | Automatic optimization for A100/H100/RTX |\n\n---\n\n## Installation\n\n```bash\ngit clone https://github.com/alepot55/flash-reasoning.git\ncd flash-reasoning\nuv sync --all-extras\n```\n\n---\n\n## Quick Start\n\n```python\nimport torch\nfrom flash_reasoning import PhysicalKVAllocator, tree_attention\n\n# Initialize allocator\nallocator = PhysicalKVAllocator(\n    num_blocks=1024, block_size=16, num_kv_heads=8, head_dim=128, device=\"cuda\"\n)\n\n# Allocate shared prefix\nroot = allocator.alloc_branch(num_tokens=2048)\n\n# Fork into 256 branches (all share prefix)\nbranches = [\n    allocator.alloc_branch(2112, parent_branch_id=root.branch_id, fork_position=2048)\n    for _ in range(256)\n]\n\n# Compute attention\nq = torch.randn(256, 32, 128, device=\"cuda\", dtype=torch.float16)\noutput = tree_attention(q, allocator, [b.branch_id for b in branches])\n```\n\n---\n\n## Architecture\n\n```\nsrc/flash_reasoning/\n├── core/memory.py          # PhysicalKVAllocator (block alloc + refcount)\n├── kernels/tree_attention.py   # Triton kernels (fused GQA + online softmax)\n└── ops/attention.py        # tree_attention() wrapper\n```\n\n---\n\n## Benchmarks\n\n```bash\nuv run python benchmarks/benchmark_throughput.py --batch-sizes 64 128 256\n```\n\n---\n\n## Tests\n\n```bash\nuv run pytest tests/ -v\n# 9 passed ✓\n```\n\n---\n\n## Citation\n\n```bibtex\n@software{flash_reasoning_2025,\n  title = {Flash-Reasoning: Tree-Aware KV-Cache Attention},\n  author = {Potenza, Alessandro},\n  year = {2025},\n  url = {https://github.com/alepot55/flash-reasoning}\n}\n```\n\n---\n\n## Related Work\n\n- [FlashAttention](https://github.com/Dao-AILab/flash-attention) — IO-aware attention\n- [vLLM](https://github.com/vllm-project/vllm) — PagedAttention\n- [SGLang](https://github.com/sgl-project/sglang) — RadixAttention\n\n---\n\n\u003cp align=\"center\"\u003e\n  \u003csub\u003eMIT License • Built with Triton\u003c/sub\u003e\n\u003c/p\u003e\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Falepot55%2Fflash-reasoning","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Falepot55%2Fflash-reasoning","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Falepot55%2Fflash-reasoning/lists"}