@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d247a931cc

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 08:00 UTC

[feat] Add FP8 per-tensor Q scale support to Triton attention backend (#42080)

Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>

Dom Brown committed May 19, 2026 at 17:02 UTC d247a931cc25e7253feccbd6260d48216ff5c081

3 files changed +166 -22

tests/kernels/attention/test_triton_unified_attention.py

+131 -8

@@ -9,6 +9,7 @@ from vllm.platforms import current_platform

9 from vllm.utils.math_utils import next_power_of_2

10 from vllm.utils.torch_utils import set_random_seed

11 from vllm.v1.attention.ops.triton_unified_attention import unified_attention

12 +from vllm.v1.kv_cache_interface import KVQuantMode

14 DEVICE_TYPE = current_platform.device_type

@@ -153,16 +154,20 @@ def test_triton_unified_attn(

154 q_descale = None

155 k_descale = None

156 v_descale = None

157 + kv_quant_mode = KVQuantMode.NONE

158 if q_dtype is not None:

157 - # QKV are drawn from N(0, 1): no need for a fp8 scaling factor

158 - maybe_quantized_query = query.to(q_dtype)

159 - maybe_quantized_key_cache = key_cache.to(q_dtype)

160 - maybe_quantized_value_cache = value_cache.to(q_dtype)

161 -

159 + # Use non-1 scales so FP8 Q/K/V descale handling is tested explicitly.

160 + q_scale = torch.tensor(0.75, dtype=torch.float32)

161 + k_scale = torch.tensor(0.5, dtype=torch.float32)

162 + v_scale = torch.tensor(0.25, dtype=torch.float32)

163 + q_descale = q_scale

164 scale_shape = (num_seqs, num_kv_heads)

163 - q_descale = None # Not yet supported

164 - k_descale = torch.rand(scale_shape, dtype=torch.float32)

165 - v_descale = torch.rand(scale_shape, dtype=torch.float32)

165 + k_descale = torch.full(scale_shape, k_scale.item(), dtype=torch.float32)

166 + v_descale = torch.full(scale_shape, v_scale.item(), dtype=torch.float32)

167 + maybe_quantized_query = (query / q_scale).to(q_dtype)

168 + maybe_quantized_key_cache = (key_cache / k_scale).to(q_dtype)

169 + maybe_quantized_value_cache = (value_cache / v_scale).to(q_dtype)

170 + kv_quant_mode = KVQuantMode.FP8_PER_TENSOR

171

172 num_par_softmax_segments = 16

173 head_size_padded = next_power_of_2(head_size)

@@ -201,6 +206,7 @@ def test_triton_unified_attn(

206 softmax_segm_output=softmax_segm_output,

207 softmax_segm_max=softmax_segm_max,

208 softmax_segm_expsum=softmax_segm_expsum,

209 + kv_quant_mode=kv_quant_mode,

210 )

211

212 ref_output = ref_paged_attn(

@@ -223,6 +229,123 @@ def test_triton_unified_attn(

229 )

230

231

232 +@pytest.mark.parametrize(

233 + "seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]

234 +)

235 +@pytest.mark.parametrize("num_heads", NUM_HEADS)

236 +@pytest.mark.parametrize("head_size", HEAD_SIZES)

237 +@pytest.mark.parametrize("block_size", BLOCK_SIZES)

238 +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)

239 +@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)

240 +@torch.inference_mode()

241 +def test_triton_unified_attn_bf16_query_fp8_kv(

242 + seq_lens: list[tuple[int, int]],

243 + num_heads: tuple[int, int],

244 + head_size: int,

245 + block_size: int,

246 + num_blocks: int,

247 + seq_threshold_3D: int,

248 +) -> None:

249 + """Test bf16 Q with FP8 per-tensor KV cache (dequant via _cast_kv_tile)."""

250 + torch.set_default_device(DEVICE_TYPE)

251 + set_random_seed(0)

252 +

253 + num_seqs = len(seq_lens)

254 + query_lens = [x[0] for x in seq_lens]

255 + kv_lens = [x[1] for x in seq_lens]

256 + num_query_heads = num_heads[0]

257 + num_kv_heads = num_heads[1]

258 + assert num_query_heads % num_kv_heads == 0

259 + max_query_len = max(query_lens)

260 + max_kv_len = max(kv_lens)

261 + window_size = (-1, -1)

262 + scale = head_size**-0.5

263 +

264 + dtype = torch.bfloat16

265 + query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)

266 + key_cache = torch.randn(

267 + num_blocks, block_size, num_kv_heads, head_size, dtype=dtype

268 + )

269 + value_cache = torch.randn_like(key_cache)

270 +

271 + k_scale = torch.tensor(0.5, dtype=torch.float32)

272 + v_scale = torch.tensor(0.25, dtype=torch.float32)

273 + fp8_key_cache = (key_cache / k_scale).to(FP8_DTYPE)

274 + fp8_value_cache = (value_cache / v_scale).to(FP8_DTYPE)

275 +

276 + scale_shape = (num_seqs, num_kv_heads)

277 + k_descale = torch.full(scale_shape, k_scale.item(), dtype=torch.float32)

278 + v_descale = torch.full(scale_shape, v_scale.item(), dtype=torch.float32)

279 +

280 + cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(

281 + dim=0, dtype=torch.int32

282 + )

283 + kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)

284 +

285 + max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size

286 + block_tables = torch.randint(

287 + 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32

288 + )

289 +

290 + output = torch.empty_like(query)

291 +

292 + num_par_softmax_segments = 16

293 + head_size_padded = next_power_of_2(head_size)

294 + softmax_segm_output = torch.empty(

295 + (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),

296 + dtype=torch.float32,

297 + )

298 + softmax_segm_max = torch.empty(

299 + (seq_threshold_3D, num_query_heads, num_par_softmax_segments),

300 + dtype=torch.float32,

301 + )

302 + softmax_segm_expsum = torch.empty(

303 + (seq_threshold_3D, num_query_heads, num_par_softmax_segments),

304 + dtype=torch.float32,

305 + )

306 +

307 + unified_attention(

308 + q=query,

309 + k=fp8_key_cache,

310 + v=fp8_value_cache,

311 + out=output,

312 + cu_seqlens_q=cu_query_lens,

313 + seqused_k=kv_lens_t,

314 + max_seqlen_q=max_query_len,

315 + max_seqlen_k=max_kv_len,

316 + softmax_scale=scale,

317 + causal=True,

318 + window_size=window_size,

319 + block_table=block_tables,

320 + softcap=0,

321 + q_descale=None,

322 + k_descale=k_descale,

323 + v_descale=v_descale,

324 + seq_threshold_3D=seq_threshold_3D,

325 + num_par_softmax_segments=num_par_softmax_segments,

326 + softmax_segm_output=softmax_segm_output,

327 + softmax_segm_max=softmax_segm_max,

328 + softmax_segm_expsum=softmax_segm_expsum,

329 + kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,

330 + )

331 +

332 + ref_output = ref_paged_attn(

333 + query=query,

334 + key_cache=key_cache,

335 + value_cache=value_cache,

336 + query_lens=query_lens,

337 + kv_lens=kv_lens,

338 + block_tables=block_tables,

339 + scale=scale,

340 + )

341 +

342 + atol, rtol = 1.5e-1, 1.5e-1

343 + (

344 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),

345 + f"{torch.max(torch.abs(output - ref_output))}",

346 + )

347 +

348 +

349 @pytest.mark.parametrize(

350 "seq_lens",

351 [

vllm/v1/attention/backends/triton_attn.py

+17 -8

@@ -39,6 +39,7 @@ from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (

39 from vllm.v1.attention.ops.triton_unified_attention import unified_attention

40 from vllm.v1.kv_cache_interface import (

41 AttentionSpec,

42 + KVQuantMode,

43 get_kv_quant_mode,

44 kv_cache_uses_per_token_head_scales,

45 )

@@ -585,6 +586,7 @@ class TritonAttentionImpl(AttentionImpl):

586 if key_cache.dtype == torch.uint8:

587 key_cache = key_cache.view(self.fp8_dtype)

588 value_cache = value_cache.view(self.fp8_dtype)

589 + q_descale = None

590 k_descale = None

591 v_descale = None

592 k_scale_cache = self._k_scale_cache

@@ -592,17 +594,24 @@ class TritonAttentionImpl(AttentionImpl):

594 # FP8 per-tensor / auto path (original flow).

595 else:

596 key_cache, value_cache = kv_cache.unbind(1)

595 - if is_quantized_kv_cache(self.kv_cache_dtype):

596 - if key_cache.dtype != self.fp8_dtype:

597 - key_cache = key_cache.view(self.fp8_dtype)

598 - value_cache = value_cache.view(self.fp8_dtype)

599 - assert layer._q_scale_float == 1.0, (

600 - "A non 1.0 q_scale is not currently supported."

601 - )

597 + if (

598 + is_quantized_kv_cache(self.kv_cache_dtype)

599 + and key_cache.dtype != self.fp8_dtype

600 + ):

601 + key_cache = key_cache.view(self.fp8_dtype)

602 + value_cache = value_cache.view(self.fp8_dtype)

603 descale_shape = (

604 attn_metadata.query_start_loc.shape[0] - 1,

605 key_cache.shape[2],

606 )

607 + q_descale = (

608 + layer._q_scale

609 + if (

610 + self._kv_quant_mode == KVQuantMode.FP8_PER_TENSOR

611 + and query.dtype == self.fp8_dtype

612 + )

613 + else None

614 + )

615 k_descale = layer._k_scale.expand(descale_shape)

616 v_descale = layer._v_scale.expand(descale_shape)

617 k_scale_cache = None

@@ -638,7 +647,7 @@ class TritonAttentionImpl(AttentionImpl):

647 window_size=self.sliding_window,

648 block_table=block_table,

649 softcap=self.logits_soft_cap,

641 - q_descale=None, # Not supported

650 + q_descale=q_descale,

651 k_descale=k_descale,

652 v_descale=v_descale,

653 seq_threshold_3D=seq_threshold_3D,

vllm/v1/attention/ops/triton_unified_attention.py

+18 -6

@@ -45,7 +45,8 @@ def _cast_kv_tile(data, Q, tensor_scale, KV_QUANT_MODE: tl.constexpr):

45 ``3`` (FP8 per-token-head): plain cast. Per-token-head modes apply

46 their scales separately on S/P inside the loop.

47 - ``KV_QUANT_MODE == 1`` (FP8 per-tensor): dequantize using the

48 - tensor-wide scale.

48 + tensor-wide scale, unless Q is also FP8 and the caller folds the scales

49 + into the attention score and output accumulator.

50 """

51 if KV_QUANT_MODE == 1:

52 if Q.dtype.is_fp8():

@@ -191,6 +192,7 @@ def kernel_unified_attention(

192 qq_bias_ptr,

193 # Scalars

194 scale,

195 + q_scale,

196 k_scale,

197 v_scale,

198 out_scale,

@@ -268,8 +270,10 @@ def kernel_unified_attention(

270 # (see ``unified_attention`` wrapper for the gating rules).

271 USE_TD: tl.constexpr = False,

272 USE_TD_QO: tl.constexpr = False,

273 + Q_IS_FP8: tl.constexpr = False,

274 ):

275 USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2

276 + USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8

277

278 if USE_TD:

279 tl.static_assert(

@@ -355,6 +359,11 @@ def kernel_unified_attention(

359 L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

360 # acc : (BLOCK_M, HEAD_SIZE_PADDED)

361 acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)

362 + score_scale = scale

363 + value_scale = 1.0

364 + if USE_FP8_Q_DESCALE:

365 + score_scale = scale * tl.load(q_scale) * tl.load(k_scale)

366 + value_scale = tl.load(v_scale)

367

368 context_len = seq_len - cur_batch_query_len

369

@@ -497,9 +506,9 @@ def kernel_unified_attention(

506 if USE_PER_TOKEN_HEAD_SCALES:

507 # Per-token-head quant: fuse softmax_scale with per-head k_scale

508 # to avoid a separate BLOCK_M × TILE_SIZE multiply on S.

500 - S += tl.dot(Q, K) * (scale * k_token_head_scales[None, :])

509 + S += tl.dot(Q, K) * (score_scale * k_token_head_scales[None, :])

510 else:

502 - S += scale * tl.dot(Q, K)

511 + S += score_scale * tl.dot(Q, K)

512

513 if USE_SOFTCAP:

514 S = apply_softcap(S, softcap)

@@ -537,6 +546,8 @@ def kernel_unified_attention(

546

547 # ---- Epilogue ---------------------------------------------------------

548 if IS_3D:

549 + if USE_FP8_Q_DESCALE:

550 + acc *= value_scale

551 # Store per-segment partials; finalized by ``reduce_segments``.

552 if USE_TD_QO:

553 # 3D target: segm_output[token, head, segm_idx, :]. Advance

@@ -592,6 +603,8 @@ def kernel_unified_attention(

603 )

604 else:

605 acc = acc / L[:, None]

606 + if USE_FP8_Q_DESCALE:

607 + acc *= value_scale

608 if USE_FP8:

609 acc = acc * tl.load(out_scale)

610 acc = tl.clamp(acc, FP8_MIN, FP8_MAX)

@@ -790,8 +803,6 @@ def unified_attention(

803 use_td: bool = False,

804 ):

805 assert causal, "Only causal attention is supported"

793 - assert q_descale is None, "Q scales not supported"

794 -

806 if sinks is not None:

807 assert sinks.shape[0] == q.shape[1], "Sinks must be num_query_heads size"

808

@@ -938,7 +949,6 @@ def unified_attention(

949 # Pass the K cache as a stand-in pointer; never dereferenced.

950 k_scale_ptr = k

951 v_scale_ptr = v

941 -

952 # 3D needs real segm tensors; 2D never touches them but Triton wants

953 # a non-null pointer. Reuse ``out`` as the placeholder.

954 segm_output_ptr = softmax_segm_output if use_3d else out

@@ -970,6 +980,7 @@ def unified_attention(

980 k_scale_cache_ptr=k_scale_ptr,

981 v_scale_cache_ptr=v_scale_ptr,

982 scale=softmax_scale,

983 + q_scale=q_descale,

984 k_scale=k_descale,

985 v_scale=v_descale,

986 out_scale=1 / output_scale if output_scale is not None else 1.0,

@@ -1017,6 +1028,7 @@ def unified_attention(

1028 USE_FP8=output_scale is not None,

1029 IS_3D=use_3d,

1030 KV_QUANT_MODE=kv_quant_mode,

1031 + Q_IS_FP8=(q.dtype == current_platform.fp8_dtype()),

1032 CHUNK_LOOKBACK=chunk_lookback,

1033 CHUNK_SIZE=chunk_size,

1034 USE_TD=use_td,