[cuda] Nestest Superscale Quantization for CUDA Weights (int4 / int6)

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.

[Gasoonjia]

Quant schema update

Weights are loaded from GGUF and kept at 4-bit (Q4_K) / 6-bit (Q6_K); we further-quantize the per-group metadata (the affine scale/zero) instead of storing it in bf16. Dequant stays affine: w = (q - zero) * scale. The grouping mirrors GGUF's Q4_K/Q6_K layout: group_size = 32 (int4) / 16 (int6), and a 256-weight super-block that shares one fp16 metadata step.

• int4 (Q4_K):group_size=32, super-block=256 (8 groups). scale = per-group uint8 code x per-256 fp16 step; zero = per-group uint8 code x per-256 fp16 step. Both fp16 steps are packed into one 32-bit word and broadcast to the 8 lanes of a super-block with a single warp shuffle in the decode kernel. 4.625 bpw.
• int6 (Q6_K): symmetric (no zero, exactly like Q6_K), group_size=16, super-block=256 (16 groups). scale = per-group int8 code x per-256 fp16 step; the per-group int8 codes are hoisted (loaded once per uint4) in the decode kernel. 6.5625 bpw.

Why this design

(a) Follow Q4_K / Q6_K. Weights come from GGUF, so we mirror llama.cpp's 256-weight super-block granularity (one fp16 step per 256 weights) and its symmetric-Q6_K / asymmetric-Q4_K split, keeping our representation close to the source format.

(b) int6 keeps its int8 scale code. int6 is symmetric, so scale is its only metadata; dropping it to a 6-bit code collapses accuracy (-9.5 dB) with no zero to compensate, so int6 keeps int8 scale + per-256 fp16 step. A dedicated int6 6-bit-planar-metadata path could revisit this in the future.

Ablation study

int4 / Q4_K whole-weight SNR vs the exact GGUF Q4_K dequant (same 8-tensor sample throughout; anchors reproduce exactly).

scale code	scale step	zero code	zero step	SNR (dB)	bpw
bf16 (no code quant)	-	bf16	-	48.04	-
uint8	per-256 fp16	uint8	per-256 fp16 (latest / z_pack)	45.89	4.625
uint8	per-256 fp16	uint8	per-256 bf16	45.67	4.625
uint8	per-256 fp16	uint8	per-row fp16	45.17	4.5625
uint8	per-256 fp16	uint8	per-row bf16	45.15	4.5625
uint8	per-256 bf16	uint8	per-row bf16	45.10	4.5625
uint8	per-row bf16	uint8	per-row bf16	44.37	4.500
uint8	per-row fp16	uint8	per-row bf16	44.36	4.500
6-bit	per-256 fp16	8-bit	per-256 fp16	46.19*	4.5625
6-bit	per-256 fp16	6-bit	per-256 fp16	38.5	4.500
6-bit	per-row	8-bit	per-row	37.85	4.441

*6-bit-scale variant: higher raw SNR but rejected on perf (6-bit planar scale reconstruct = +5.2% decode; see below).

Analysis. The accuracy lever is step granularity, not step dtype or code width: moving the scale step from per-row to per-256 buys +0.78 dB (44.37 -> 45.15), and moving the zero step to per-256 as well buys another +0.74 dB (-> 45.89); step dtype (fp16 vs bf16) is worth <=0.05 dB for a uint8 code. We therefore keep both codes at uint8 (a 6-bit code loses 6-10 dB, and its planar reconstruct also costs decode perf) and spend the accuracy budget on per-256 fp16 steps for both scale and zero — which is exactly GGUF Q4_K's per-super-block-fp16-d granularity. This keeps the weight codes byte-aligned and identical to the source 4-bit/6-bit values, holds bpw in a tight band (4.625, +0.125 over the 4.5 baseline), and lands at 45.89 dB (+1.52 dB over the prior 44.37 baseline). The one higher-raw-SNR alternative (6-bit scale + 8-bit zero, 46.19 dB) was rejected because the 6-bit planar scale reconstruct regresses decode +5.2% for +0.3 dB — not worth it.

e2e results

gemma4-31b Q4_K_M @131072, greedy, cuda_graph, 3-rep median.

A100, turboquant OFF / bf16 KV

Comparing ET latest (this PR) vs ET prior (pre-metadata-opt baseline) vs llama.cpp.

Prefill (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	1729.7	1701.0	1266.2
2048	2438.1	2494.5	1513.7
8192	2273.7	2275.6	1514.9
32768	1674.6	1710.0	1269.1
130048	851.6	855.9	768.8

Decode (tok/s, higher better)

prompt len	ET latest	ET prior	llama.cpp
512	49.93	50.5	43.63
2048	47.93	48.4	42.73
8192	47.42	48.1	42.12
32768	46.79	47.0	39.74
130048	36.98	37.2	33.04

Peak VRAM (GB, lower better)

prompt len	ET latest	ET prior	llama.cpp
512	31.56	31.38	19.38
2048	31.85	31.66	19.70
8192	31.88	31.69	20.17
32768	31.88	31.69	22.04
130048	31.88	31.69	29.46

A100 takeaways. vs ET prior: decode neutral (-0.45% to -1.41%, mean ~-0.9%), prefill neutral (+-2.3%), VRAM +~0.19 GB (the expected +0.06 bpw of the per-256 zero step) — i.e. +1.52 dB SNR at negligible runtime cost. vs llama.cpp: ET decode is +12-18% faster and ET prefill is faster at every length; llama.cpp uses ~12 GB less VRAM at short/mid lengths (packed Q4_K_M weights vs ET's bf16 runtime weights), a gap that narrows to ~2.4 GB at 130048 as the KV cache dominates.

RTX 5090, turboquant ON / tq4 KV cache

Comparing ET latest (New, this PR) vs ET prior (Orig, pre-metadata-opt baseline), decode d=512.

Decode (tok/s, higher better) & Peak VRAM (GiB, lower better)

prompt len	Orig decode	New decode	Decode improvement	Orig VRAM	New VRAM	VRAM improvement
512	57.16	60.78	+6.33%	25.06	23.29	-7.06%
2048	56.39	59.79	+6.03%	25.06	23.29	-7.06%
8192	55.60	58.88	+5.90%	25.06	23.29	-7.06%
32768	55.29	58.60	+5.99%	25.06	23.29	-7.06%

5090 takeaways. With turboquant tq4 KV, the metadata-opt is a clear win on both axes: decode +5.9-6.3% faster and peak VRAM -7.06% (-1.77 GiB) at every measured length. (The VRAM drop here is a turboquant-path effect, not the +0.06 bpw weight-metadata cost seen on the A100 bf16-KV config.)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20571

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

@claude review this PR

[claude]

Claude finished @digantdesai's task in 7m 15s —— View job

---

PR Review: Superscale Quantization for CUDA Weights (int4 / int6)

Reviewed the full diff against origin/main. This is a clean, well-documented change: the quant metadata (scale/zero) moves from per-row bf16 to per-group uint8/int8 codes + a per-256-super-block fp16 step, mirroring GGUF Q4_K/Q6_K. The op signatures, meta kernels, C shims, .cu/.cuh, dispatch, and tests are all updated consistently, and the ablation study justifies the design well. Overall it looks correct and mergeable. A few notes below.

Correctness

1. _encode_uint8_per_super assumes zero_point >= 0 (safe today, but undocumented invariant).backends/cuda/coalesced_int4_tensor.py:147-150 uses block_max = xb.amax(...) (max, not absmax) and clamp_(0, 255), so any negative value silently clamps to 0. I confirmed torchao's Int4Tensor always produces non-negative zero_point (unsigned [0,15], and the symmetric path pins it to 8.0), so this is correct for all real inputs. But the old code stored t.zero_point.t() as bf16 and would have preserved a negative sign if one ever appeared. Consider a one-line assert or docstring note in _encode_uint8_per_super that the input must be non-negative, so a future non-torchao packer fails loudly instead of silently degrading. The scale encoder has the same max-based assumption, which is fine since scales are positive.

2. int6 kernel: unconditional scale_row[g_base + 1] read is OOB for gs=32.backends/cuda/runtime/shims/int6_plain_mm.cuh:251 always loads scale_row[g_base + 1]. For the shipped gs=16 this is safe (g_base is always even and g_base+1 < n_groups). But the comment at lines 246-248 explicitly claims gs in {16, 32} is supported — and for gs=32 a uint4 spans exactly one group, so g_base runs up to n_groups-1 and scale_row[g_base+1] reads one past the row (the ws1 value is then unused, since wpg_shift = gs_shift-3 = 2 makes every w>>2 == 0). This is a latent OOB read, not currently exercised because only gs=16 is used/tested. Either guard the ws1 load (g_base+1 < n_groups) or tighten the comment to state only gs=16 is supported. Fix this →

3. Stale docstring in _concat_coalesced_int4_along_n.examples/models/gemma4_31b/cuda_source_transformations.py:700 describes zero_step as [N, 1] bf16, but the actual encoding is [N, K/256] fp16 (as the rest of the PR documents). Minor, but worth fixing since the surrounding docstring is otherwise precise about shapes/dtypes. Fix this →

Things done well

• int4 warp-shuffle deadlock avoidance (int4_plain_mm.cuh:189-197): the warp-aligned trip count so all 32 lanes reach __shfl_sync(0xffffffff, ...) even when K_half_16 < 32 is the correct fix for full-mask shuffle on tiny shapes, and the "safe dummy load (index 0), contribute 0" for inactive lanes is right. Good that the test suite includes small shapes to exercise this.
• z_pack (packing both fp16 steps into one 32-bit word, one leader shuffle per 8-lane subgroup) keeps the dp4a dot products bit-identical to the scale-only baseline while adding zero-step decode at no extra shuffle — nicely reasoned and matches the "+0.74 dB at op-level-neutral" claim.
• Kernel-side validation (n_super = K/256, zero_step.size(1) == scale_step.size(1), K % 256 == 0) is thorough, and the error messages are detailed.
• Encode/decode agreement: rounding uses the fp16-rounded step (step_f) so pack-time and kernel-time reconstruction match — this is the subtle correctness detail that's easy to get wrong, and it's handled.

Test coverage

Solid — test_int4_dispatch.py / test_int6_dispatch.py and the two .cpp kernel tests were substantially expanded (single/multi super-block, wide-N, small shapes). One gap: no test exercises gs=32 for int6, which is why note #2 slipped through. If gs=32 is genuinely intended to be supported, add a case; otherwise the comment should say gs=16 only.

CI

The listed failures (test_int4_weights_preserved, the export-model-cuda-artifact jobs) are worth confirming are unrelated infra/merge-base noise vs. this metadata change — the TestInt4Inference::test_int4_weights_preserved failure in particular is on a path this PR touches. Recommend confirming it's green after the merge with main.

None of these are blocking; #1 and #2 are the ones I'd address (or explicitly wontfix) before merge.
· branch cuda-int4-int6-metadata-opt

[digantdesai]

what about non k-quant users?

[digantdesai]

check in the script to recreate these

[digantdesai]

s/zero_step/zero_point_step - Nit

[digantdesai]

Is this Tensor only for k-quants?

[digantdesai]

Looks good except the fact that I am not sure if we want to use the plain kernel for other 4/6b weights or just k-quant with quantized/nested scales.

[Gasoonjia]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

[digantdesai]

Thanks @digantdesai for review.

Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

[digantdesai]

stamping it to unblock you

[Gasoonjia]

Thanks @digantdesai for review.
Tensor now is k-quant only but the pipeline doesn't. We convert the original hdd int4 and int6 tensor into q4_k and q6_k and reuse k-quant pipeline. The precision are the same and only have ~0.7db weight during the conversion.

What about performance? Running Q4_0 as nested k-quant kernel must cost some perf, no?

for kernel-base test perf is even better cuz now bpw is lower.