[Feature] Triton INT4 per-token-head KV cache quantization by JartX · Pull Request #40835 · vllm-project/vllm

JartX · 2026-04-24T21:28:37Z

Based on PR:

Changed to only INT4_PER_TOKEN_HEAD

INT4_PER_TOKEN_HEAD

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.912	±	0.0273
		strict-match	5	exact_match	↑	0.91	±	0.0273

Signed-off-by: JartX <sagformas@epdcenter.es>

JartX · 2026-06-22T09:49:18Z

Hi @tdoublep Could you please check it again? Only the int4 kernel remains, and the rest has been unified.

tdoublep

Thanks! I have some minor comments to try and minimize diff, but I think we are nearly there.

tdoublep · 2026-06-12T11:32:33Z

+        self._is_per_token_head = kv_cache_spec.kv_quant_mode.is_per_token_head
+        if self._is_per_token_head:
+            self._init_reorder_batch_threshold(1, supports_spec_as_decode=False)
+


Why is this change needed?

Yeah, this one's just leftover. It used to back a flash-attn varlen fastpath that needed decodes up front, but that path is gone now that everything goes through unified_attention. Removed it.

tdoublep · 2026-06-23T13:25:05Z

+        seq_lens_cpu = common_attn_metadata._seq_lens_cpu
+        qsl_cpu = (
+            common_attn_metadata.query_start_loc_cpu
+            if seq_lens_cpu is not None
+            else None
+        )
+


Why is this needed?

tdoublep · 2026-06-23T13:28:21Z

Do we really need this subdir - can't we just have the kernel under ops?

Yeah, fair — it didn't really earn its own package. Moved it up to ops/int4_per_token_head.py and dropped the subdir. Done in the latest push.

tdoublep · 2026-06-23T13:30:48Z

+            K = _cast_kv_tile(K_load, Q, k_scale, KV_QUANT_MODE)
+            V = _cast_kv_tile(V_load, Q, v_scale, KV_QUANT_MODE)


What is reason for moving this?

Good eye, nothing actually calls it, it's a leftover from an earlier pass. Removed, helpers file is back to matching main. I left it orphaned during the refactor

tdoublep · 2026-06-23T15:29:18Z



+@triton.jit
+def cast_kv_tile(data, Q, tensor_scale, KV_QUANT_MODE: tl.constexpr):


Maybe I'm missing something but where is this used?

tdoublep · 2026-06-23T15:30:12Z

    k_scale_cache: torch.Tensor,  # [num_blocks, block_size, num_kv_heads] float32
    v_scale_cache: torch.Tensor,  # [num_blocks, block_size, num_kv_heads] float32
    slot_mapping: torch.Tensor,  # [num_tokens]
+    *,


Do we need to add "*"?

tdoublep · 2026-06-23T15:32:17Z

+        full_block_f32 = 2 * kv_half_bytes // 4
+        slot_f32 = nkv * padded_hs * dtype_sz // 4
+        head_f32 = padded_hs * dtype_sz // 4
+        scale_off_f32 = hs * dtype_sz // 4


Why do we need this diff?

tdoublep · 2026-06-23T15:32:22Z

+        head_f32 = padded_hs * dtype_sz // 4
+        scale_off_f32 = hs * dtype_sz // 4

-        # K scales: kv_half=0


tdoublep · 2026-06-23T15:32:27Z

        )
        self._k_scale_cache.fill_(1.0)

-        # V scales: kv_half=1, offset by kv_half_bytes


tdoublep · 2026-06-23T15:33:56Z

 )
 @pytest.mark.parametrize(
-    "base_model,test_model",
+    "base_model,test_model,kv_cache_dtype",


How much does this increase the test time?

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

- Remove leftover reorder_batch_threshold init (dead since per-token-head moved to unified_attention) - Remove unused seq_lens_cpu / query_start_loc_cpu metadata fields - Flatten triton_quant_kv/ subdir into ops/int4_per_token_head.py - Revert _cast_kv_tile move in triton_unified_attention - Drop unused cast_kv_tile helper (triton_attention_helpers back to main) - Drop keyword-only * from reshape_and_cache_per_token_head_quant - Restore original _ensure_scale_caches comments - Use Llama-3.2-1B for the int4 e2e case (was Meta-Llama-3-8B) Signed-off-by: JartX <sagformas@epdcenter.es>

JartX · 2026-06-23T16:14:23Z

Thanks @tdoublep, really good catches .

Changes applied

`reorder_batch_threshold` — removed

Dead weight left over from an old flash-attn varlen mixed-batch path that assumed decodes-first ordering. That assumption was invalidated when we moved per-token-head to unified_attention, so the parameter had no effect. Dropped entirely.

`seq_lens_cpu` / `query_start_loc_cpu` — removed

Both were being written but never read downstream. Clean removal — no functional impact.

`triton_quant_kv/` subdirectory — flattened

Agreed it didn't justify its own package boundary. Moved up to ops/int4_per_token_head.py.

`_cast_kv_tile` placement — reverted

No good reason to move it. Restored to original location.

`cast_kv_tile` in helpers file — removed

Leftover from an earlier refactor, nothing was referencing it. Dropped — helpers file now matches main again.

Wildcard import (`*`) — removed

Comment churn in `_ensure_scale_caches` — reverted

Kept the original comments as-is.

tdoublep

Just did a last pass through the tests. Please revert the change of introducing "device" everywhere, it creates huge churn for no reason that I can see.

tdoublep · 2026-06-23T18:39:49Z

    )

    torch.set_default_device(DEVICE_TYPE)
+    device = "cuda"


why do we need to pass the device around in addition to setting the default device?

tdoublep · 2026-06-23T18:41:10Z


-    cache_dtype: torch.dtype  # torch.int8 or FP8_DTYPE
-    kv_cache_dtype_str: str  # "int8_per_token_head" or "fp8_per_token_head"
+    cache_dtype: torch.dtype  # torch.uint8, torch.int8, or FP8_DTYPE


Why does this specific comment change?

tdoublep · 2026-06-23T18:41:36Z


    set_random_seed(seed)
    torch.set_default_device(DEVICE_TYPE)
+    device = "cuda"


tdoublep · 2026-06-23T18:43:16Z

+    if is_int4:
+        deq_atol, deq_rtol = 0.5, 0.5
+    else:
+        deq_atol, deq_rtol = 0.1, 0.1


where is deq_atol and deq_rtol used in non-int4 path?

tdoublep · 2026-06-23T18:43:26Z


    torch.set_default_device(DEVICE_TYPE)
    set_random_seed(42)
+    device = "cuda"


tdoublep · 2026-06-23T18:44:00Z

    for i in range(num_tokens):
        blk = i // block_size
        off = i % block_size
-


nit: reduce unnecessary diff

tdoublep · 2026-06-23T18:44:47Z


    torch.set_default_device(DEVICE_TYPE)
    set_random_seed(0)
+    device = "cuda"


- Drop the device="cuda" / device=device churn (set_default_device already handles placement; was also hardcoding cuda) - Restore original QuantConfig dtype comments - Move int4 tolerance into the int4 branch (deq_atol/deq_rtol were unused in the int8/fp8 path) - Restore blank line in round-trip loop; recollapse calls to match main Signed-off-by: JartX <sagformas@epdcenter.es>

tdoublep

LGTM

JartX · 2026-06-23T19:38:53Z

Thanks again @tdoublep, all good catches

Thanks so much for the thorough reviews. Learned quite a bit going through them — good catches all around.

I have pushed the fixes:

`device=` keyword — reverted across all tests

Fair call — it was pure churn. set_default_device(DEVICE_TYPE) already handles placement, so the explicit device="cuda" was redundant at best and was hardcoding CUDA on top of the abstraction. Cleaned up everywhere.

`QuantConfig` dtype comments — reverted

No real justification for the change. Restored originals.

`deq_atol` / `deq_rtol` — scoped to INT4 branch

Only the INT4 path actually uses them, so moved the definitions into that branch. The INT8/FP8 path keeps its plain 0.1 — no need to pull in the wider tolerances there.

Blank line / "ditto" nits — fixed

Formatting — recollapsed to match `main`

A few torch.zeros/randn calls and a parametrize/fixture had gotten split across lines during iteration. Recollapsed them so the formatting is consistent with the rest of the codebase.

Two CI failures from stale call sites after merging main: - int4_per_token_head kernel called compute_kv_seq_mask without the seq_len arg added in main, shifting positionals and dropping MAX_MM_RANGES (Triton CompilationError at 238:19). - test_int8_per_token_head_raw_cache_matches_round_reference called triton_reshape_and_cache_flash_per_token_head_quant without the required kv_quant_mode arg. Signed-off-by: JartX <sagformas@epdcenter.es>

…ect#40835) Signed-off-by: JartX <sagformas@epdcenter.es> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

JartX added 30 commits April 15, 2026 20:38

precommit

b8d7e34

Signed-off-by: JartX <sagformas@epdcenter.es>

remove comment

dfe2fa4

Signed-off-by: JartX <sagformas@epdcenter.es>

precommit

bd231d0

Signed-off-by: JartX <sagformas@epdcenter.es>

fix fast path

dd11029

Signed-off-by: JartX <sagformas@epdcenter.es>

reduce comments

141be40

Signed-off-by: JartX <sagformas@epdcenter.es>

fix scheduler

1efb4ff

Signed-off-by: JartX <sagformas@epdcenter.es>

fix fast path mistral

dd0c0cc

Signed-off-by: JartX <sagformas@epdcenter.es>

dtype

ea38cc1

Signed-off-by: JartX <sagformas@epdcenter.es>

add guard

86b2d75

Signed-off-by: JartX <sagformas@epdcenter.es>

remove comments

0450bc0

Signed-off-by: JartX <sagformas@epdcenter.es>

Merge branch 'main' into refactor/prefill-fastpath-per-token-head

e5d73fc

use FLASH_ATTN if available

d32d7dc

Signed-off-by: JartX <sagformas@epdcenter.es>

added window

006eee3

Signed-off-by: JartX <sagformas@epdcenter.es>

precommit

63373bc

Signed-off-by: JartX <sagformas@epdcenter.es>

is_flash_attn_varlen_func_available and mixed batch

80fe133

Signed-off-by: JartX <sagformas@epdcenter.es>

precommit

1afdbc5

Signed-off-by: JartX <sagformas@epdcenter.es>

remove FA

6672664

Signed-off-by: JartX <sagformas@epdcenter.es>

Merge branch 'main' into refactor/prefill-fastpath-per-token-head

98ce477

added mixed

e89cfda

Signed-off-by: JartX <sagformas@epdcenter.es>

add per_token_head_attention kernel

3d69c3e

Signed-off-by: JartX <sagformas@epdcenter.es>

pre-commit

affc8b3

Signed-off-by: JartX <sagformas@epdcenter.es>

missing seq_lens_cpu and query

9425464

Signed-off-by: JartX <sagformas@epdcenter.es>

working on kernel

e7ae0a1

Signed-off-by: JartX <sagformas@epdcenter.es>

working on kernel

090a7a8

Signed-off-by: JartX <sagformas@epdcenter.es>

precommit

1338413

Signed-off-by: JartX <sagformas@epdcenter.es>

precommit

7f4e7ec

Signed-off-by: JartX <sagformas@epdcenter.es>

prefill kernel

724f341

Signed-off-by: JartX <sagformas@epdcenter.es>

int8 wmma

0e283d4

Signed-off-by: JartX <sagformas@epdcenter.es>

int8 wmma

55cac18

Signed-off-by: JartX <sagformas@epdcenter.es>

remove dead code

24650f6

Signed-off-by: JartX <sagformas@epdcenter.es>

mergify Bot removed the needs-rebase label Jun 22, 2026

tdoublep reviewed Jun 23, 2026

View reviewed changes

tdoublep approved these changes Jun 23, 2026

View reviewed changes

tdoublep added the ready ONLY add when PR is ready to merge/full CI is needed label Jun 23, 2026

Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2

c5cbc66

tdoublep enabled auto-merge (squash) June 23, 2026 19:32

auto-merge was automatically disabled June 23, 2026 20:44
Head branch was pushed to by a user without write access

JartX added 2 commits June 24, 2026 00:34

Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2

e1b12e3

Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2

df5f0f6

tdoublep enabled auto-merge (squash) June 24, 2026 06:15

Merge branch 'main' into refactor/prefill-fastpath-per-token-head-v2

46a7bb5

tdoublep merged commit 70749fd into vllm-project:main Jun 24, 2026
114 checks passed

mgoin mentioned this pull request Jun 30, 2026

[Roadmap] vLLM Roadmap Q2 2026 #39749

Closed

29 of 75 tasks

LiRunGuo mentioned this pull request Sep 15, 2026

[Bugfix][Attention] Support int4_per_token_head on non-power-of-two head sizes #56198

Merged

JartX mentioned this pull request Sep 21, 2026

[ROCm][DO NOT MERGE] RDNA3 (gfx1100) full inference stack — tracking branch for splitting into landable PRs #57925

Draft

dudilester mentioned this pull request Oct 21, 2026

[Bugfix][Test] Rotate KV into RHT domain in the INT4 unified attn test #62881

Merged

6 tasks done

		K = _cast_kv_tile(K_load, Q, k_scale, KV_QUANT_MODE)
		V = _cast_kv_tile(V_load, Q, v_scale, KV_QUANT_MODE)



		@triton.jit
		def cast_kv_tile(data, Q, tensor_scale, KV_QUANT_MODE: tl.constexpr):

Conversation

JartX commented Apr 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JartX commented Jun 22, 2026

Uh oh!

tdoublep left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JartX commented Jun 23, 2026

Changes applied

reorder_batch_threshold — removed

seq_lens_cpu / query_start_loc_cpu — removed

triton_quant_kv/ subdirectory — flattened

_cast_kv_tile placement — reverted

cast_kv_tile in helpers file — removed

Wildcard import (*) — removed

Comment churn in _ensure_scale_caches — reverted

Uh oh!

tdoublep left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tdoublep left a comment

Choose a reason for hiding this comment

Uh oh!

JartX commented Jun 23, 2026

device= keyword — reverted across all tests

QuantConfig dtype comments — reverted

deq_atol / deq_rtol — scoped to INT4 branch

Blank line / "ditto" nits — fixed

Formatting — recollapsed to match main

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

JartX commented Apr 24, 2026 •

edited

Loading

`reorder_batch_threshold` — removed

`seq_lens_cpu` / `query_start_loc_cpu` — removed

`triton_quant_kv/` subdirectory — flattened

`_cast_kv_tile` placement — reverted

`cast_kv_tile` in helpers file — removed

Wildcard import (`*`) — removed

Comment churn in `_ensure_scale_caches` — reverted

`device=` keyword — reverted across all tests

`QuantConfig` dtype comments — reverted

`deq_atol` / `deq_rtol` — scoped to INT4 branch

Formatting — recollapsed to match `main`