[sergqwer]

The int8-KV prompt attention (prompt_attn_i8_kernel, v2) drifts away from the FP32 kernel on short prompts. On IQ2_XS the first-token KL to the FP32 kernel is 0.025 after a 2K prompt, and the 32 greedy tokens change. This PR adds an opt-in kernel (STRATA_PROMPT_ATTN_IMMA=1) that puts both products on INT8 tensor cores with an operand split exact to 24 bits. It is closer to FP64 than v2 at every shape tested and 1.35-1.53x faster per launch. Unset (the default), nothing changes.

What it does

prompt_attn_i8v3_kernel keeps v2's block, its cp.async pipeline and its online softmax.

• Operands. QK^T and PV run as mma m16n8k32 on the int8 K and V codes as stored, with no int8 -> fp16 conversion. The other operand is split into three exact 8-bit parts:

  • q is scaled per (head, 64-dim group) and p·v_scale per (row, chunk), both to 127;
  • then X = rint(x · 65536) = h · 65536 + l · 256 + m, with h and l signed bytes and m an unsigned one (a u8 x s8 MMA).

  That is 24 bits, finer than v2's FP16 hi + lo split (~22 bits). Each part's int32 sum is exact, and the three parts are combined in FP32.

• Occupancy. One K stage and one V stage per warp (K(c+1) is gathered during softmax(c) and PV(c)), 12-row partial sums, and __launch_bounds__(128, 3): ~25.5 KB of shared memory and 168 registers, so 3 blocks per SM instead of v2's 2. The kernel waits on latency, not on its MMAs: with one block per SM it ran 1.54x slower than with two.

• Scope. sm_80+ only: the device code is behind the file's STRATA_PA_SM80, and Turing / Volta / HIP keep their kernels. The file compiles for 75/80/86/89/120. A test hook, qsa_prompt_attn_set_imma(), is added. qsa_prompt_attn_parity now checks IMMA against FP64 and v2 and times both.

Measured

RTX 5090, Ryzen 9 9950X3D, IQ2_XS (ISTA GSQ-RCO), --kv int8, --expert-cache 12000.

Kernel accuracy against FP64 (qsa_prompt_attn_parity, max relative error, 0 failures):

context / queries	IMMA	v2	FP32 kernel
32K / 2048	2.24e-6	3.17e-6	2.51e-6
1500 / 1500	2.15e-6	3.30e-6	1.96e-6
2100 / 256	2.64e-6	2.78e-6	1.90e-6

End to end, first-token KL to the FP32 kernel (STRATA_PROMPT_ATTN_OLD=1); the last column says whether the 32 greedy tokens match the FP32 kernel's:

prompt	v2	IMMA	tokens = FP32 (v2 / IMMA)
600	0.0010	0.0016	yes / yes
2K	0.0248	0.0026	no / yes
8K	0.0009	0.0028	no / yes
32K	0.0340	0.0196	yes / no

• Median KL: 0.0027 for IMMA against 0.013 for v2.
• The top token is the same everywhere.
• The greedy tokens match the FP32 kernel's on 3 of 4 prompts with IMMA, and on 2 of 4 with v2.

On the same model in NVFP4 (our fork) the gap at 2K was wider: KL 0.10 for v2 against 0.007 for IMMA.

Speed:

• Per launch (qsa_prompt_attn_parity): 32K / 2048 queries 2.09 -> 1.55 ms (1.35x); 1500 / 1500 0.61 -> 0.40 ms (1.53x); 2100 / 256 0.28 -> 0.26 ms (1.09x).

• In the engine, 32K prompt in 8K chunks (nsys, 48 launches): 358.1 -> 311.6 ms (1.15x).

• Prompt time, 3 interleaved pairs, which the machine's noise dominates:

  	v2	IMMA
  2K	701 / 764 / 853 ms	713 / 832 / 806 ms
  32K	5091 / 5063 / 5419 ms	5191 / 5208 / 4915 ms

  The kernel saves ~47 ms of a ~5 s 32K prompt, which is below that noise.

Identity. With the switch unset, the first-token logits and the 32 greedy tokens equal v0.1.40.2 on IQ2_XS after 2K and after 32K (identical bytes, repeated).

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[xyzzing]

A datapoint from our int8 attention work on gfx1100 (fork 65a3c8f, fixture-tested — tests/hip/qsa_int8_q_fixture.cu): a design that quantizes the softmax-input path failed our pre-declared precision gates — the int8→i32 dots themselves were integer-exact vs double (4.9e-8), but softmax Δ 9.5e-2 / output Δ 1.55e-1 against 2e-2 gates, an int8-max-abs-vs-outlier structural class. Your design differs in exactly the way that matters: K/V enter as the already-stored int8 codes, accumulation is exact int32, and softmax stays fp32 — so I'd expect it to pass an independent parity gate where ours did not.

If AMD coverage is ever wanted: the RDNA3 primitive exists (v_wmma_i32_16x16x16_iu8), and __byte_perm already has a shim in upstream's hip_compat (v_perm_b32, #262). Expectation note: on gfx1100, prompt attention is ~10% of 128K prefill after 0.1.40.2's SIMT fp32 kernel landed. Happy to measure on a 7900 XTX if there's interest.

AI disclosure: this comment was drafted by an AI assistant (GLM-5.3, Z.ai) at the account owner's direction, from our recorded fixtures and docs (GFX1100.md rejected-list entry; verdict in fork commit 65a3c8f's message).

[edt-xx]

@xyzzing "Happy to measure on a 7900 XTX if there's interest.". I have a 7900xt (linux) here and run with 70-100K prompts (letta-code). I'd be interested to see if this helps.

[blange48]

Tested on 4× RTX 5080 (sm_120, 84 SMs, PCIe Gen3), IQ3_S, --kv int8 --kv-resident 32768, layer split 12,24,36, --batch 8 --batch-groups 4, on 0.1.40.2 plus #1368. Same image, STRATA_PROMPT_ATTN_IMMA unset vs =1:

	default	IMMA
45K-token prompt (2 runs)	11.15 / 11.07 s	11.32 / 11.30 s
fresh ~325 / ~585-token prompt, median first token	527 / 598 ms	577 / 592 ms
decode, 1 / 4 clients	141 / 270 tok/s	140 / 240 tok/s (one run)
greedy answers (5 fixed prompts, 220 tokens)		2/5 identical, the others diverge after 82-326 characters

No speed gain on this rig: on a 4-stage split the prompt time is dominated by expert streaming and the stage hand-offs, not by prompt attention. The changed greedy answers are expected, since it's a different (more accurate) kernel. I can't measure the KL to FP32 here; all answers read fine. It stays off in our production for now. If a long-context quality check would be useful (e.g. 64K-128K with int8 KV), I can run one.

[sergqwer]

@blange48 thank you for running it on a split rig. Two notes, and the long-context check you offered.

• Decode: the kernel runs only in the prompt path (Prefill, through qsa_prompt_attn_batch). Decode and the verify windows never call it. So 270 -> 240 tok/s with 4 clients (one run) can only come from other clients' prompts read in the meantime, or from run-to-run variance.

• Prompt: prompt attention is a small part of a prompt. On one RTX 5090 at 96K (IQ2_XS, int8 K/V, STRATA_PREFILL_TIMING=1):

  	v2 (default)	IMMA
  qsa attn, ms	1,060 / 1,152	893 / 967 (-16%)
  whole prompt, s	16.7 / 18.0	16.6 / 17.6

  So the kernel gain is about 1% of the prompt, and two runs of the same arm differ by 7% here. On a 4-stage split, where expert streaming and the hand-offs dominate, it is smaller still. Your 1.6% the other way is about the same size. If you have a minute, the qsa attn entry of the STRATA_PREFILL_TIMING=1 line for both arms would show whether the kernel itself is slower on the 5080's 84 SMs or whether the difference is elsewhere.

Long context, against the FP32 kernel (STRATA_PROMPT_ATTN_OLD=1). 96,399-token prompt, IQ2_XS, int8 K/V, first-token logits and 32 greedy tokens:

	first-token KL to FP32	greedy tokens
v2 (default)	0.00076	= FP32
IMMA	0.000075	= FP32
v1 (another FP32 summation order: the noise floor)	0.000015	= FP32

At 96K IMMA is 10x closer to the FP32 kernel than v2, as it was at 2K (0.0026 against 0.025). Greedy answers that part after a few hundred characters are what any FP32-level change does over 220 tokens, and they do not say which kernel is closer. Against the FP32 kernel, IMMA is the closer of the two.

[midhatn]

Tested a local v0.1.40.2 port on a single RTX 4070 Laptop (sm_89), Ryzen 9 7940HS, 64 GB RAM, Windows/MSVC and CUDA 13.4. IQ3_S, 64K context, int8 KV with 32K resident on the GPU.

The parity benchmark passed and IMMA was closer to FP64 at both main shapes:

Context / queries	v2 max relative error	IMMA	v2 time	IMMA time
32,768 / 256	3.89e-6	2.67e-6	1.248 ms	0.868 ms
65,536 / 64	3.51e-6	2.20e-6	0.458 ms	0.400 ms

Those are gains over v2, not the larger v1-to-v2 gains also printed by the benchmark. Both invocations ended with zero failures.

Whole prompt processing was essentially unchanged: 73.605 s with v2 versus 73.352 s with IMMA for 22,305 tokens (two samples per arm). The smaller 635/5,914-token prompts were also close. Same compiled binary, only IMMA toggled; scatter was held off in both arms. All bounded lookup, structured tool-call, cache-return and queue-trace checks passed. This was sampled high-reasoning generation, not a greedy-token identity or model-level KL comparison.

So this adds sm_89 parity/per-kernel evidence, but I would not claim an end-to-end speed improvement on this RAM/PCIe-heavy setup from these samples.

[blange48]

Thanks, that settles it. Here are the STRATA_PREFILL_TIMING=1 numbers you asked for, on 4× RTX 5080 (84 SMs, sm_120), IQ3_S, --kv int8 --kv-resident 32768, layer split 12,24,36, now on 0.1.40.3 plus this PR. A 45K-token prompt, chunks as --prefill auto picks them, the qsa attn entry of the timing line:

chunk	v2 (default)	IMMA
4,071 tokens	22 ms	15 ms (−32%)
8,192 tokens	44 ms	29 ms (−34%)
share of the chunk's GPU timeline	~3-4%	~2-3%
whole 45K prompt, 2 runs	11.43 / 11.45 s	11.40 / 11.44 s

So the kernel itself is about a third faster on the 5080 too. On this split the whole prompt doesn't move, as you expected: hand-offs and expert streaming dominate. You're also right about decode: with IMMA on, a later run gave 1 / 2 / 3 / 4 / 8 clients at 145 / 171 / 239 / 268 / 394 tok/s, the same as without it, and my earlier 270 → 240 was variance. batch_test stays 8/8 and 3/3 identical to solo.

Given your 96K KL numbers (10× closer to FP32 than v2), we turned it on in production (STRATA_PROMPT_ATTN_IMMA=1).

[sergqwer]

@blange48 @midhatn thank you both. For the maintainer, the evidence so far on three cards:

	kernel	whole prompt	quality
RTX 5090 (sm_120), IQ2_XS	qsa attn -16% at 96K	~1% of the prompt	KL to the FP32 kernel 0.000075 vs v2 0.00076 (96K); 0.0026 vs 0.025 (2K)
4x RTX 5080 (sm_120), IQ3_S, layer split	-32% / -34% (4K / 8K chunks)	unchanged (streaming and hand-offs dominate)	batch_test 8/8 and 3/3 identical to solo
RTX 4070 Laptop (sm_89), IQ3_S	1.15-1.44x, closer to FP64 than v2 (2.67e-6 vs 3.89e-6)	unchanged (RAM/PCIe-bound)	parity passes

Decode never calls the kernel. It stays opt-in (STRATA_PROMPT_ATTN_IMMA=1), so the default is byte-identical.

[Niko1221]

Thanks. Ported as an opt-in (STRATA_PROMPT_ATTN_IMMA=1). On an RTX 3060 the 20K-token prompt read is +2.5% (5 of 5 pairs), teacher-forced KL against the default is at the noise floor (0.0012) and 3 of 3 needles at 32K were found. We want one more card class and a long-prompt KL before making it the default.

[Niko1221]

Thanks for the pull request.

Your change is already in main for 0.2.0-beta, with your credit, so there is nothing left for you to do here. We're closing it because we're resetting the tracker for 0.2.0-beta: there are more than 1,300 open issues and pull requests, and we're closing everything that isn't being worked on right now.

If part of your work did not make it in and you want it considered, open a new pull request on top of main that follows CONTRIBUTING.md, or comment /reopen after fixing the title area prefix and the template.