Skip to content

[Feature] Triton INT4 per-token-head KV cache quantization #40835

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tdoublep merged 72 commits into from
Jun 24, 2026
Merged

[Feature] Triton INT4 per-token-head KV cache quantization #40835

Show file tree
Hide file tree
Changes from 62 commits
Commits
Show all changes
72 commits
Select commit Hold shift + click to select a range
b8d7e34
precommit
JartX Apr 15, 2026
dfe2fa4
remove comment
JartX Apr 15, 2026
bd231d0
precommit
JartX Apr 15, 2026
dd11029
fix fast path
JartX Apr 15, 2026
141be40
reduce comments
JartX Apr 15, 2026
1efb4ff
fix scheduler
JartX Apr 15, 2026
dd0c0cc
fix fast path mistral
JartX Apr 15, 2026
ea38cc1
dtype
JartX Apr 15, 2026
86b2d75
add guard
JartX Apr 15, 2026
0450bc0
remove comments
JartX Apr 15, 2026
e5d73fc
Merge branch 'main' into refactor/prefill-fastpath-per-token-head
JartX Apr 15, 2026
d32d7dc
use FLASH_ATTN if available
JartX Apr 15, 2026
006eee3
added window
JartX Apr 15, 2026
63373bc
precommit
JartX Apr 15, 2026
80fe133
is_flash_attn_varlen_func_available and mixed batch
JartX Apr 15, 2026
1afdbc5
precommit
JartX Apr 15, 2026
6672664
remove FA
JartX Apr 15, 2026
98ce477
Merge branch 'main' into refactor/prefill-fastpath-per-token-head
JartX Apr 15, 2026
e89cfda
added mixed
JartX Apr 16, 2026
3d69c3e
add per_token_head_attention kernel
JartX Apr 16, 2026
affc8b3
pre-commit
JartX Apr 16, 2026
9425464
missing seq_lens_cpu and query
JartX Apr 16, 2026
e7ae0a1
working on kernel
JartX Apr 16, 2026
090a7a8
working on kernel
JartX Apr 16, 2026
1338413
precommit
JartX Apr 16, 2026
7f4e7ec
precommit
JartX Apr 16, 2026
724f341
prefill kernel
JartX Apr 16, 2026
0e283d4
int8 wmma
JartX Apr 16, 2026
55cac18
int8 wmma
JartX Apr 16, 2026
24650f6
remove dead code
JartX Apr 16, 2026
3e2116e
remove comment
JartX Apr 16, 2026
11eed26
remove comment
JartX Apr 16, 2026
f8bb990
Merge branch 'main' into refactor/prefill-fastpath-per-token-head
JartX Apr 19, 2026
d485ea1
feat: INT4 / INT2 per-token-head KV cache quantization
JartX Apr 23, 2026
b3a14c6
refactor(triton_quant_kv): split packed_per_token_head.py
JartX Apr 24, 2026
0c07c67
Merge branch 'main' into refactor/prefill-fastpath-per-token-head
JartX Apr 24, 2026
320602f
Merge refactor/prefill-fastpath-per-token-head into int4-int2 v2 mirror
JartX Apr 24, 2026
81ba439
fix(pth_attn_stage1): use other=0.0 on K/V loads
JartX Apr 24, 2026
8325949
feat(pth_attn): extend decode/prefill kernels to INT4/INT2 packed modes
JartX Apr 24, 2026
39f5ac0
perf(pth_attn): GQA-grouped packed decode kernel for INT4 / INT2
JartX Apr 24, 2026
a289666
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX May 4, 2026
03f0529
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX May 25, 2026
0039fcd
Merge remote-tracking branch 'origin/main' into refactor/prefill-fast…
JartX May 28, 2026
27715b1
Merge remote-tracking branch 'origin/main' into refactor/prefill-fast…
JartX Jun 2, 2026
3629cf0
Remove int2_per_token_head KV cache quantization mode
JartX Jun 2, 2026
27a5a1b
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 3, 2026
c838db1
Remove int8 WMMA/MFMA QK fast-path from per-token-head attention
JartX Jun 3, 2026
d2a4c37
Address mgoin's review on per-token-head KV cache
JartX Jun 3, 2026
2eff7d8
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 4, 2026
3e3dcaf
Merge remote-tracking branch 'origin/main' into refactor/prefill-fast…
JartX Jun 4, 2026
75c90b9
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 8, 2026
d934e23
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 9, 2026
a763450
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 10, 2026
4063939
[Attention] Replace KV-quant factory with explicit dispatch
JartX Jun 10, 2026
0ab5f18
[Attention] Collapse per-token-head forward dispatch
JartX Jun 10, 2026
986cbe7
[Attention] Consolidate INT4 KV-quant package into one module
JartX Jun 10, 2026
158fba6
[Attention] Trim verbose docstrings and stale factory references
JartX Jun 10, 2026
8bace36
[Attention] Fold _hadamard into int4_per_token_head
JartX Jun 10, 2026
4abc52a
[Attention] Route INT8/FP8 per-token-head KV through unified_attention
JartX Jun 11, 2026
184f28c
[Attention] Reuse the core per-token-head reshape kernel for INT8/FP8
JartX Jun 11, 2026
b6b0e89
[Attention] Restore the inline per-head scale padding comment
JartX Jun 11, 2026
dbb3fc8
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 11, 2026
6cfb19c
fix: INT8 per-token-head KV cache quantizer rounds to nearest
JartX Jun 12, 2026
1c4c469
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 18, 2026
6858251
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 22, 2026
769c279
Address review: minimize per-token-head INT4 diff
JartX Jun 23, 2026
6a5b2ef
Address review: revert device churn in per-token-head kernel tests
JartX Jun 23, 2026
c5cbc66
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
mergify[bot] Jun 23, 2026
0e5749e
Fix per-token-head call-site signatures after main merge
JartX Jun 23, 2026
e1b12e3
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 23, 2026
df5f0f6
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 24, 2026
46a7bb5
Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2
JartX Jun 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/design/attention_backends.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ Priority is **1 = highest** (tried first).
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

Expand Down
16 changes: 12 additions & 4 deletions tests/models/quantization/test_per_token_kv_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,17 +22,25 @@
reason="Per-token-head KV cache requires CUDA or ROCm GPU.",
)
@pytest.mark.parametrize(
"base_model,test_model",
"base_model,test_model,kv_cache_dtype",

@tdoublep tdoublep Jun 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How much does this increase the test time?

@JartX JartX Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Switched int4 to Llama-3.2-1B-Instruct so it doesn't drag in the 8B anymore, now it's just one extra dtype on the existing param.I originally moved it up to keep it in line with int2 (which genuinely fell apart at 1B), but int4 should be fine on the small model. If logprobs end up too tight we'll bump just that case back to 8B

[
(
"meta-llama/Meta-Llama-3-8B",
"meta-llama/Meta-Llama-3-8B",
"int4_per_token_head",
),
(
"meta-llama/Llama-3.2-1B-Instruct",
"meta-llama/Llama-3.2-1B-Instruct",
"int8_per_token_head",
),
(
"meta-llama/Llama-3.2-1B-Instruct",
"meta-llama/Llama-3.2-1B-Instruct",
"fp8_per_token_head",
),
],
)
@pytest.mark.parametrize(
"kv_cache_dtype", ["int8_per_token_head", "fp8_per_token_head"]
)
@pytest.mark.parametrize("max_tokens", [4])
@pytest.mark.parametrize("enforce_eager", [True])
@pytest.mark.parametrize("backend", ["TRITON_ATTN"])
Expand Down
Loading
Loading