Skip to content
cysong2025Public

About

Single-GPU LLM decode research prototype: paged KV cache, Triton attention, CUDA append, scheduling, shared prefixes, and multi-layer transactions.

Topics

Resources

Code of conduct

Contributing

Security policy

Stars

1 star

Watchers

0 watching

Forks

⚡ FlashDec

A transactional, paged-decode runtime for single-GPU LLM inference

从 Triton kernel、Paged KV Cache 到 vLLM/Qwen2.5-3B 端到端验证。

Repository checks Python PyTorch CUDA Triton License

Quick start · Architecture · Performance · API · Reproduce

FlashDec 优化大模型逐 token 生成阶段的数据路径:它用事务化的 Paged KV Cache 管理请求状态,以 Triton/CUDA kernel 完成 RoPE、KV append 和 paged attention,并通过 out-of-tree 插件接入 vLLM。vLLM 继续负责模型、prefill、调度、sampling 与 serving;FlashDec 只替换满足条件的 single-token decode attention,并在不支持的场景自动回退。

代表性结果 测量边界
🚀 Qwen2.5-3B 完整生成延迟 −4.58% B8 / input 8192 / output 4096,offline LLM.generate
⚡ 输出吞吐 +4.80% 与上面相同的 4 轮 balanced AB/BA 实验
🧩 vLLM decode-attention p50 −19.75% / −20.74% B8 / context 1024、2048,kernel-only
🧠 Shared-prefix KV 容量 −68.8% 固定 48-block pool,admission 9/16 → 16/16

Note

端到端结果来自 RTX 5070 上固定的 Qwen2.5-3B BF16 长上下文离线 workload,baseline 为显式 vLLM TRITON_ATTN,不代表所有模型、在线流量或 vLLM 默认最快配置。完整协议、范围和负结果见性能报告。

架构

FlashDec integration architecture from vLLM and the plugin router through the transactional runtime to Triton and CUDA kernels.

层次 FlashDec 提供的能力
vLLM 集成 general_plugins 注册、CUSTOM attention backend、严格 eligibility 检查和原生 Triton fallback
Decode runtime block-aware scheduler、DecodeEngine、跨层 token transaction、失败回滚与 lifecycle validation
KV memory physical block allocator、block table、shared-prefix refcount/LRU、request-private tail
GPU data path fused RoPE + KV append、grouped-GQA split-KV paged attention、query-head reducer
Evidence PyTorch reference、跨 backend parity、执行路径 attestation、commit-bound benchmark summaries

vLLM 是如何接入的?

FlashDec 没有维护修改版 vLLM,也没有接管整套推理引擎。插件复用 vLLM 的模型执行和 KV layout,在 uniform single-token decode、支持的 dtype/head shape 等条件满足时选择 FlashDec kernel;prefill、mixed batch 或其他不支持路径继续使用 vLLM Triton backend。正式实验还会校验 CUDA Graph capture marker,避免“配置了插件但实际没有执行”的假阳性。

实现合同与 fallback 边界见 vLLM backend design。

性能概览

FlashDec results across Qwen end-to-end generation, vLLM attention, KV capacity and runtime mechanisms.

结果层次 FlashDec 对比对象 结果 结论
完整 Qwen LLM.generate vLLM TRITON_ATTN 0.9542x latency 长上下文目标通过;短路径 guard 为 1.0029x
Qwen attention kernel vLLM Triton 0.8025x / 0.7926x p50 B8 两个目标 shape 通过
Cache-owned transaction checked dispatch 1.7307x p50 speedup 仅适用于 cache-owned metadata
Shared prefix private context blocks 68.8% KV capacity saved 容量收益,不宣称稳定 latency 收益
Boundary scheduler greedy step-only 100% vs 0% completion correctness/progress,不是 latency 对比
External kernel check FlashInfer 0.6.15 1.2003x / 1.2284x latency ratio FlashInfer 更快;只代表共同 kernel shape

图表由受版本控制的数据快照和 canonical summaries 确定性生成。不同层次的 ratio 方向和计时范围不能互相换算;原始结论见结果索引。

快速开始

推荐 Linux/WSL、Python 3.10+,以及与 PyTorch CUDA build 匹配的 CUDA Toolkit、NVCC 和 Ninja。

git clone https://github.com/cysong2025/flashdec.git
cd flashdec

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
python -m pip install -e ".[dev,cuda-extension]"

python scripts/check_env.py
python -m pytest -q -ra

先查看兼容性矩阵。FlashInfer 和 vLLM 使用隔离的固定环境,避免 pip 替换已经验证的 PyTorch/CUDA stack;对应安装和 Qwen 命令见复现指南。

API 示例

import flashdec

output = flashdec.decode(
    q,
    k_cache,
    v_cache,
    block_tables,
    seq_lens,
    sm_scale=q.shape[-1] ** -0.5,
    block_size=32,
)

事务化多层调用:

tx = engine.begin_step(request_ids)
for layer_idx, (q, k, v) in enumerate(zip(q_by_layer, k_by_layer, v_by_layer)):
    engine.step_layer(tx, layer_idx, q, k, v)
result = engine.commit_step(tx)

任一 layer 失败都会 abort 整个 token;调用方必须丢弃此前返回的 layer outputs。完整张量约定和状态契约见 API documentation。

文档

想了解什么 从这里开始
系统分层与数据流 Architecture
vLLM 插件和 fallback vLLM backend
API 与状态语义 API reference
性能数字和适用边界 Performance report
环境、测试与 benchmark Reproducibility
全部设计和实验文档 Documentation index

项目边界

FlashDec 核心面向单 GPU、single-token decode、FP16/BF16 和 caller-provided Q/K/V。它不自行实现 tokenizer、完整 Transformer/prefill、sampling、HTTP/RPC 服务、生产级 continuous batching、TP/PP、多机或 swap/offload;可选 vLLM 插件复用其中一部分外部能力,但不把它们变成 FlashDec 自有实现。

欢迎提交 issue 或 pull request。贡献前请阅读贡献指南、行为准则和安全政策。

License

FlashDec 采用 Apache License 2.0 开源。

About

Single-GPU LLM decode research prototype: paged KV cache, Triton attention, CUDA append, scheduling, shared prefixes, and multi-layer transactions.

Topics

Resources

Code of conduct

Contributing

Security policy

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Used by

Contributors

Languages