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

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 6, 19:15 UTC

[Kernel] Add FP8 KV cache support to Triton MLA decode attention (#34597)

Signed-off-by: grimulkan <grimulkan@gmail.com>

grimulkan committed Mar 12, 2026 at 10:32 UTC a1257fd1ea93da6e27b31e4739ac2707781d8ba7

4 files changed +192 -8

docs/design/attention_backends.md

+1 -1

@@ -213,5 +213,5 @@ configuration.

213 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |

214 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |

215 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |

216 -| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |

216 +| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |

217 | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |

tests/kernels/attention/test_triton_decode_attention.py

+134

@@ -90,3 +90,137 @@ def test_decode_attention(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):

90 )

92 assert torch.allclose(o, o1)

93 +

94 +

95 +def _quantize_to_fp8(tensor: torch.Tensor):

96 + """Quantize a BF16 tensor to FP8 e4m3fn with per-tensor scale.

97 +

98 + Returns (fp8_tensor, scale) where:

99 + fp8_tensor ≈ tensor / scale (stored as float8_e4m3fn)

100 + tensor ≈ fp8_tensor.to(float32) * scale (dequantized)

101 + """

102 + amax = tensor.abs().amax()

103 + # float8_e4m3fn max representable value is 448.0

104 + scale = (amax / 448.0).clamp(min=1e-12).to(torch.float32)

105 + fp8_tensor = (

106 + (tensor.to(torch.float32) / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn)

107 + )

108 + return fp8_tensor, scale

109 +

110 +

111 +@pytest.mark.parametrize("B", [3])

112 +@pytest.mark.parametrize("L", [1025])

113 +@pytest.mark.parametrize("H_Q", [32])

114 +@pytest.mark.parametrize("H_KV", [32, 8])

115 +@pytest.mark.parametrize("D_QK", [128, 576])

116 +@pytest.mark.parametrize("D_V", [128, 512])

117 +@pytest.mark.parametrize("CACHE_SIZE", [16384])

118 +@pytest.mark.parametrize("PAGE_SIZE", [1, 16])

119 +def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE):

120 + """Test FP8 KV cache path: quantize K/V to FP8, run kernel with scales,

121 + and compare against BF16 reference output."""

122 + assert CACHE_SIZE % PAGE_SIZE == 0

123 + dtype = torch.bfloat16

124 + seq_len = L

125 + sm_scale = 1.0 / (D_QK**0.5)

126 + num_kv_splits = 8

127 +

128 + num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)

129 + req_to_page = torch.randint(

130 + 0, CACHE_SIZE // PAGE_SIZE, (B, num_pages_per_batch, 1), device="cuda"

131 + )

132 + req_to_token = req_to_page * PAGE_SIZE

133 + req_to_token = req_to_token.expand(B, num_pages_per_batch, PAGE_SIZE)

134 + req_to_token = req_to_token + torch.arange(PAGE_SIZE, device="cuda").view(1, 1, -1)

135 + req_to_token = req_to_token.view(B, -1)

136 + req_to_token = req_to_token[:, :seq_len].contiguous()

137 +

138 + q = torch.randn(B, H_Q, D_QK, dtype=dtype, device="cuda")

139 +

140 + # Create BF16 K/V as reference

141 + k_bf16 = torch.randn(CACHE_SIZE, H_KV, D_QK, dtype=dtype, device="cuda")

142 + v_bf16 = torch.randn(CACHE_SIZE, H_KV, D_V, dtype=dtype, device="cuda")

143 +

144 + # --- BF16 reference ---

145 + o_ref = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")

146 + lse_ref = torch.zeros(B, H_Q, dtype=dtype, device="cuda")

147 + attn_logits = torch.empty(

148 + (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"

149 + )

150 +

151 + if PAGE_SIZE == 1:

152 + decode_attention_fwd(

153 + q,

154 + k_bf16,

155 + v_bf16,

156 + o_ref,

157 + lse_ref,

158 + req_to_token,

159 + b_seq_len=torch.full((B,), seq_len, device="cuda"),

160 + attn_logits=attn_logits,

161 + num_kv_splits=num_kv_splits,

162 + sm_scale=sm_scale,

163 + )

164 + else:

165 + k_paged = k_bf16.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)

166 + v_paged = v_bf16.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)

167 + decode_attention_fwd(

168 + q,

169 + k_paged,

170 + v_paged,

171 + o_ref,

172 + lse_ref,

173 + req_to_page,

174 + b_seq_len=torch.full((B,), seq_len, device="cuda"),

175 + attn_logits=attn_logits,

176 + num_kv_splits=num_kv_splits,

177 + sm_scale=sm_scale,

178 + page_size=PAGE_SIZE,

179 + )

180 +

181 + # --- FP8 path ---

182 + k_fp8, k_scale = _quantize_to_fp8(k_bf16)

183 + v_fp8, v_scale = _quantize_to_fp8(v_bf16)

184 +

185 + o_fp8 = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")

186 + lse_fp8 = torch.zeros(B, H_Q, dtype=dtype, device="cuda")

187 + attn_logits_fp8 = torch.empty(

188 + (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device="cuda"

189 + )

190 +

191 + if PAGE_SIZE == 1:

192 + decode_attention_fwd(

193 + q,

194 + k_fp8,

195 + v_fp8,

196 + o_fp8,

197 + lse_fp8,

198 + req_to_token,

199 + b_seq_len=torch.full((B,), seq_len, device="cuda"),

200 + attn_logits=attn_logits_fp8,

201 + num_kv_splits=num_kv_splits,

202 + sm_scale=sm_scale,

203 + k_scale=k_scale,

204 + v_scale=v_scale,

205 + )

206 + else:

207 + k_fp8_paged = k_fp8.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_QK)

208 + v_fp8_paged = v_fp8.view(CACHE_SIZE // PAGE_SIZE, PAGE_SIZE, H_KV, D_V)

209 + decode_attention_fwd(

210 + q,

211 + k_fp8_paged,

212 + v_fp8_paged,

213 + o_fp8,

214 + lse_fp8,

215 + req_to_page,

216 + b_seq_len=torch.full((B,), seq_len, device="cuda"),

217 + attn_logits=attn_logits_fp8,

218 + num_kv_splits=num_kv_splits,

219 + sm_scale=sm_scale,

220 + page_size=PAGE_SIZE,

221 + k_scale=k_scale,

222 + v_scale=v_scale,

223 + )

224 +

225 + # FP8 tolerances match test_mla_backends.py test_backend_correctness.

226 + torch.testing.assert_close(o_ref, o_fp8, atol=5e-1, rtol=1e-2)

vllm/v1/attention/backends/mla/triton_mla.py

+10 -7

@@ -32,6 +32,8 @@ class TritonMLABackend(MLACommonBackend):

32 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

33 "auto",

34 "bfloat16",

35 + "fp8",

36 + "fp8_e4m3",

37 ]

39 @classmethod

@@ -108,10 +110,11 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

110 "TritonMLAImpl"

111 )

112

113 + # For FP8 KV cache, we dequantize to BF16 on load inside the

114 + # Triton kernel. Tell the common layer not to quantize queries

115 + # to FP8 — we handle FP8 KV cache with BF16 queries (Mode 1).

116 if is_quantized_kv_cache(self.kv_cache_dtype):

112 - raise NotImplementedError(

113 - "TritonMLA V1 with FP8 KV cache not yet supported"

114 - )

117 + self.supports_quant_query_input = False

118

119 def _flash_attn_varlen_diff_headdims(

120 self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

@@ -135,9 +138,6 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

138 assert kv_c_and_k_pe_cache.numel() > 0

139 assert attn_metadata.decode is not None

140

138 - if self.kv_cache_dtype.startswith("fp8"):

139 - raise NotImplementedError("FP8 Triton MLA not yet supported")

140 -

141 if type(q) is tuple:

142 q = torch.cat(q, dim=-1)

143

@@ -171,7 +171,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

171 kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]

172 PAGE_SIZE = kv_c_and_k_pe_cache.size(1)

173

174 - # Run MQA

174 + # Run MQA — always pass layer scales. When KV cache is

175 + # BF16 the kernel's `if dtype.is_fp8()` check is a no-op.

176 decode_attention_fwd(

177 q,

178 kv_c_and_k_pe_cache,

@@ -184,6 +185,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

185 num_kv_splits,

186 self.scale,

187 PAGE_SIZE,

188 + k_scale=layer._k_scale,

189 + v_scale=layer._v_scale,

190 )

191

192 return o, lse

vllm/v1/attention/ops/triton_decode_attention.py

+47

@@ -31,6 +31,7 @@ It supports page size >= 1.

32 import logging

34 +import torch

35 from packaging import version

37 from vllm.platforms import current_platform

@@ -74,6 +75,8 @@ def _fwd_kernel_stage1(

75 stride_mid_ob,

76 stride_mid_oh,

77 stride_mid_os,

78 + k_scale,

79 + v_scale,

80 kv_group_num: tl.constexpr,

81 BLOCK_DMODEL: tl.constexpr,

82 BLOCK_DV: tl.constexpr,

@@ -109,6 +112,8 @@ def _fwd_kernel_stage1(

112 acc = tl.zeros([BLOCK_DV], dtype=tl.float32)

113

114 if split_kv_end > split_kv_start:

115 + ks = tl.load(k_scale)

116 + vs = tl.load(v_scale)

117 for start_n in range(split_kv_start, split_kv_end, BLOCK_N):

118 offs_n = start_n + tl.arange(0, BLOCK_N)

119 kv_page_number = tl.load(

@@ -129,6 +134,8 @@ def _fwd_kernel_stage1(

134 mask=(offs_n[:, None] < split_kv_end) & (mask_d[None, :]),

135 other=0.0,

136 )

137 + if k.dtype.is_fp8():

138 + k = (k.to(tl.float32) * ks).to(q.dtype)

139 qk = tl.sum(q[None, :] * k, 1)

140 qk *= sm_scale

141

@@ -147,6 +154,8 @@ def _fwd_kernel_stage1(

154 mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),

155 other=0.0,

156 )

157 + if v.dtype.is_fp8():

158 + v = (v.to(tl.float32) * vs).to(q.dtype)

159

160 n_e_max = tl.maximum(tl.max(qk, 0), e_max)

161 re_scale = tl.exp(e_max - n_e_max)

@@ -194,6 +203,8 @@ def _decode_att_m_fwd(

203 sm_scale,

204 page_size,

205 logit_cap,

206 + k_scale,

207 + v_scale,

208 ):

209 BLOCK = 64 if not is_hip_ else 8

210

@@ -231,6 +242,8 @@ def _decode_att_m_fwd(

242 att_out.stride(0),

243 att_out.stride(1),

244 att_out.stride(2),

245 + k_scale,

246 + v_scale,

247 kv_group_num=kv_group_num,

248 BLOCK_DMODEL=BLOCK_DMODEL,

249 BLOCK_DV=BLOCK_DV,

@@ -264,6 +277,8 @@ def _fwd_grouped_kernel_stage1(

277 stride_mid_ob,

278 stride_mid_oh,

279 stride_mid_os,

280 + k_scale,

281 + v_scale,

282 kv_group_num: tl.constexpr,

283 q_head_num: tl.constexpr,

284 BLOCK_DMODEL: tl.constexpr,

@@ -316,6 +331,8 @@ def _fwd_grouped_kernel_stage1(

331 acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)

332

333 if split_kv_end > split_kv_start:

334 + ks = tl.load(k_scale)

335 + vs = tl.load(v_scale)

336 for start_n in range(split_kv_start, split_kv_end, BLOCK_N):

337 offs_n = start_n + tl.arange(0, BLOCK_N)

338 kv_page_number = tl.load(

@@ -336,6 +353,8 @@ def _fwd_grouped_kernel_stage1(

353 mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]),

354 other=0.0,

355 )

356 + if k.dtype.is_fp8():

357 + k = (k.to(tl.float32) * ks).to(q.dtype)

358 qk = tl.dot(q, k.to(q.dtype))

359 if BLOCK_DPE > 0:

360 offs_buf_kpe = (

@@ -348,6 +367,8 @@ def _fwd_grouped_kernel_stage1(

367 mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]),

368 other=0.0,

369 )

370 + if kpe.dtype.is_fp8():

371 + kpe = (kpe.to(tl.float32) * ks).to(qpe.dtype)

372 qk += tl.dot(qpe, kpe.to(qpe.dtype))

373 qk *= sm_scale

374

@@ -368,6 +389,8 @@ def _fwd_grouped_kernel_stage1(

389 mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),

390 other=0.0,

391 )

392 + if v.dtype.is_fp8():

393 + v = (v.to(tl.float32) * vs).to(q.dtype)

394

395 n_e_max = tl.maximum(tl.max(qk, 1), e_max)

396 re_scale = tl.exp(e_max - n_e_max)

@@ -416,6 +439,8 @@ def _decode_grouped_att_m_fwd(

439 sm_scale,

440 page_size,

441 logit_cap,

442 + k_scale,

443 + v_scale,

444 ):

445 BLOCK = 32

446 Lk = k_buffer.shape[-1]

@@ -473,6 +498,8 @@ def _decode_grouped_att_m_fwd(

498 att_out.stride(0),

499 att_out.stride(1),

500 att_out.stride(2),

501 + k_scale,

502 + v_scale,

503 kv_group_num=kv_group_num,

504 q_head_num=head_num,

505 BLOCK_DMODEL=BLOCK_DMODEL,

@@ -609,6 +636,8 @@ def decode_attention_fwd_normal(

636 sm_scale,

637 page_size,

638 logit_cap=0.0,

639 + k_scale=None,

640 + v_scale=None,

641 ):

642 _decode_att_m_fwd(

643 q,

@@ -621,6 +650,8 @@ def decode_attention_fwd_normal(

650 sm_scale,

651 page_size,

652 logit_cap,

653 + k_scale,

654 + v_scale,

655 )

656 _decode_softmax_reducev_fwd(

657 attn_logits, q, o, lse, v_buffer, b_seq_len, num_kv_splits

@@ -640,6 +671,8 @@ def decode_attention_fwd_grouped(

671 sm_scale,

672 page_size,

673 logit_cap=0.0,

674 + k_scale=None,

675 + v_scale=None,

676 ):

677 _decode_grouped_att_m_fwd(

678 q,

@@ -652,6 +685,8 @@ def decode_attention_fwd_grouped(

685 sm_scale,

686 page_size,

687 logit_cap,

688 + k_scale,

689 + v_scale,

690 )

691 _decode_softmax_reducev_fwd(

692 attn_logits, q, o, lse, v_buffer, b_seq_len, num_kv_splits

@@ -671,8 +706,16 @@ def decode_attention_fwd(

706 sm_scale,

707 page_size=1,

708 logit_cap=0.0,

709 + k_scale=None,

710 + v_scale=None,

711 ):

712 assert num_kv_splits == attn_logits.shape[2]

713 +

714 + if k_scale is None:

715 + k_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device)

716 + if v_scale is None:

717 + v_scale = torch.tensor(1.0, dtype=torch.float32, device=q.device)

718 +

719 kv_group_num = q.shape[1] // v_buffer.shape[-2]

720

721 if kv_group_num == 1:

@@ -690,6 +733,8 @@ def decode_attention_fwd(

733 sm_scale,

734 page_size,

735 logit_cap,

736 + k_scale,

737 + v_scale,

738 )

739 else:

740 # GQA/MQA/MLA

@@ -706,4 +751,6 @@ def decode_attention_fwd(

751 sm_scale,

752 page_size,

753 logit_cap,

754 + k_scale,

755 + v_scale,

756 )