Perf/kv fp8 read path

[naerymdan]

I used it for a 1 hour agentic coding session to make sure it... you knwo, still worked.

What this changes

Two independent costs on the fp8 KV read path — the one a chunked prefill
re-runs over the whole growing prefix, which is why the penalty scaled as N².

1. The dequant scale multiplied the largest tensor in the loop.
It is per (token, kv_head), so it is constant down each tl.dot's reduction
dim and never has to touch K or V:

scores[m,n] = (sum_d q[m,d] * k[d,n]) * s_k[n]
p @ (diag(s_v) @ v) = (p * s_v[None,:]) @ v

Scaling the BLOCK_M x BLOCK_N dot output instead of the BLOCK_D x BLOCK_N K
tile and BLOCK_N x BLOCK_DV V tile is head_dim/BLOCK_M fewer multiplies.
(p must stay unscaled — l_i accumulates it as the softmax denominator.)

That also removes the only reason the tile was widened to fp32:
kv_load_e4m3_tile_f32 widens purely so a * 256.0 can restore the true e4m3
scale, and 2^8 is a power of two, so it folds into the dequant scale exactly.
kv_load_e4m3_tile_scaled16 stops before the widen and keeps the tile 16-bit.

Accuracy improves. Today the general scale multiplies before the narrow to
the compute dtype, so the product rounds. Folding it out means the tile reaches
the dot exactly (3 mantissa bits, |x| <= 1.75) and the scale is applied in
fp32 after. Worst-case error in test_extend_paged_attention_decodes_fp8_scales
drops 0.281 -> 0.0996.

2. The tile heuristic charged K/V at 2 bytes/element regardless of cache dtype.
_select_extend_tile budgets (BLOCK_M + 2*BLOCK_N) * BLOCK_D * 2. The q tile
is always 2 bytes/element, but K and V follow the cache — so a 1-byte fp8 cache
is billed for twice the shared memory it uses and drops a rung. Billing K/V at
the real element size gives (BLOCK_M*2 + 2*BLOCK_N*kv_bytes) * BLOCK_D, which
at kv_bytes=2 is algebraically the old expression, so every 16-bit cache
keeps the tile it had. At head_dim 256 on a 99KB opt-in card the fp8 arm goes
BLOCK_N 32 -> 64. This is the larger half of the win, and it is plain
shared-memory accounting, so it should pay off anywhere the fp8 arm is a rung low.

Result

RTX 3070 (sm_86, 8GB), Qwen3.6-35B-A3B-NVFP4, --moe-backend hybrid, fp8 KV,
180k window. Base is 3e5bbdd; both builds resolved an identical memory plan
(moe_cache_size=1077), so none of this is expert-residency drift.

ctx	TTFT	Δ	decode	Δ
~32.7k	34.13 -> 28.18 s	−17.4%	47.04 -> 47.94 tok/s	+1.9%
~65.1k	88.79 -> 64.36 s	−27.5%	38.99 -> 40.92 tok/s	+4.9%
~99.3k	168.27 -> 112.78 s	−33.0%	33.59 -> 36.26 tok/s	+8.0%

Per commit at ~99k, so either can be taken alone: 168.27 -> 160.42 s (−4.7%,
commit 1, mostly a decode gain) -> 112.78 s (−29.7%, commit 2).

That takes fp8 prefill from 18–35% slower than bf16 to faster than it, so
there is no longer a break-even output length to respect.

Tests

pytest tests/kernels/test_e4m3_compat.py tests/kernels/test_triton_attention.py -m "not slow" -> 4 failed, 39 passed. The same 4 fail on 3e5bbdd unpatched
(two are a FakeKVCache double missing k_scale; two are the fp8 accuracy test
exceeding its 2e-2 tolerance on sm_86) — this branch introduces none. Two tests
added: all-256-code equivalence of the two loaders, and both halves of the tile
change (fp8 grows, 16-bit does not move).

Authored by Opus 5 seriously