CUDA graph capture for the decode step
Nobody has claimed this yet.
Assessment
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Newbie friendliness
- 52/100
- Issue type
- Feature
- Clarity
- Clearly specified
- Activity status
- Active
- Domain
- machine-learning, performance
Research direction
Start with decode_tokens_per_s in src/bitsmith/bench/hf.py and compare the two proposed approaches: torch.compile(mode="reduce-overhead") with a static KV cache, or torch.cuda.CUDAGraph around a forward pass. Check src/bitsmith/kernels/int4_triton.py and the GEMV capture example in scripts/bench_kernels.py for capture safety. Done means bitsmith hf reports graph-captured bf16 and fused INT4 speeds, with INT4 faster on Qwen2.5-1.5B.
Written by the indexing model from the issue text.
Description
Per layer, the fused INT4 kernel is up to 1.33x faster than fp16 cuBLAS, but end-to-end decoding
through Hugging Face generate is still slower than bf16 (25.4 vs 42.4 tokens/s on
Qwen2.5-1.5B), because launching about 200 small kernels per token from Python dominates.
Task: capture one decode step in a CUDA graph and replay it. Options:
- a static KV cache plus
torch.compile(mode="reduce-overhead") - a manual
torch.cuda.CUDAGrapharound a forward with a static cache
Files
src/bitsmith/bench/hf.py:decode_tokens_per_ssrc/bitsmith/kernels/int4_triton.py: kernels must be capture safe (they already are
for the GEMV path inscripts/bench_kernels.py)
Done when bitsmith hf reports a graph-captured decode speed for bf16 and fused INT4,
and INT4 is faster than bf16 on Qwen2.5-1.5B.
- Dominant language
- Python
- Stars
- 0
- Forks
- 0
- PR merge metrics
- No merged PRs in 30d
Getting set up
- No Dockerfile or Docker Compose file
- Has a pull request template
- Read the contributing guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from yashwanthdevavarapu/bitsmith
-
area: benchmarks good first issue help wanted
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
-
area: benchmarks good first issue
Difficulty 2/5 1-3 hours Newbie friendliness 86/100
-
area: cli area: formats good first issue
Difficulty 2/5 1-3 hours Newbie friendliness 92/100
-
area: cli good first issue
Difficulty 2/5 1-3 hours Newbie friendliness 88/100
-
area: benchmarks good first issue
Difficulty 2/5 1-3 hours Newbie friendliness 90/100
All issues in yashwanthdevavarapu/bitsmith
Similar issues
-
Difficulty 1/5 1-3 hours Newbie friendliness 88/100
aws-samples/sample-per-model-bedrock#857 · 1 comment ·
Maintainers usually reply within 1 day
-
Difficulty 1/5 Under an hour Newbie friendliness 70/100
521xueweihan/HelloGitHub#4328 ·
-
documentation maintainer-managed
Difficulty 2/5 1-3 hours Newbie friendliness 84/100
awslabs/visual-asset-management-system#1418 ·
Maintainers usually reply within 1 day
-
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
Maintainers usually reply within 1 day
-
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
Maintainers usually reply within 1 day