Tune the INT4 prefill GEMM

Open
  • Is something wrong?
#15 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
55/100
Issue type
Feature
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch
Domain
machine-learning, performance

Research direction

Read src/bitsmith/kernels/int4_triton.py at _gemm and the benchmark setup in scripts/bench_kernels.py; compare the 512-row results in results/kernels_rtx3090.txt with cuBLAS fp16. Explore the listed tile, pipeline, and dequantization ideas, then benchmark the 7B shapes. Done means 512-row INT4 runs within 1.3x of cuBLAS fp16.

Written by the indexing model from the issue text.

Description

area: kernels enhancement performance

For 512 rows the fused INT4 GEMM is about 2x slower than cuBLAS fp16
(results/kernels_rtx3090.txt). Prefill is compute bound, so the goal is to get close to
cuBLAS, not to beat it.

Ideas: larger tiles and more pipeline stages in the autotune list, dequantizing a tile
once into shared memory and reusing it across row blocks, BLOCK_K larger than the group
size.

Files: src/bitsmith/kernels/int4_triton.py (_gemm), scripts/bench_kernels.py

Done when 512-row INT4 is within 1.3x of cuBLAS fp16 on the 7B shapes.

Dominant language
Python
Stars
0
Forks
0
No merged PRs in 30d

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from yashwanthdevavarapu/bitsmith

All issues in yashwanthdevavarapu/bitsmith

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.