[Attention] Mirror Triton KV dtype checks in MLA (#47060)
Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com>
Mike G committed
Jul 15, 2026 at 18:54 UTC
915dffaa5f93f55b44c8ac630b472700d131e165
1 file changed
+26
vllm/v1/attention/backends/mla/triton_mla.py
+26
@@ -180,6 +180,32 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
180
"TritonMLAImpl"
181
)
182
183
+ if current_platform.is_cuda():
184
+ cap = current_platform.get_device_capability()
185
+ cap_str = cap.as_version_str() if cap is not None else "unknown"
186
+ dev = current_platform.get_device_name()
187
+ if self.kv_cache_dtype.startswith("fp8") and not (
188
+ current_platform.has_device_capability(89)
189
+ ):
190
+ suggested = (
191
+ "float16" if (cap is None or cap.to_int() < 80) else "bfloat16"
192
+ )
193
+ raise ValueError(
194
+ f"FP8 KV cache is not supported by the Triton MLA backend "
195
+ f"on {dev} (compute capability {cap_str}); native FP8 "
196
+ f"(fp8e4nv) requires SM89+. Re-run with "
197
+ f"--kv-cache-dtype {suggested}."
198
+ )
199
+ if self.kv_cache_dtype == "bfloat16" and not (
200
+ current_platform.has_device_capability(80)
201
+ ):
202
+ raise ValueError(
203
+ f"bfloat16 KV cache is not supported by the Triton MLA "
204
+ f"backend on {dev} (compute capability {cap_str}); "
205
+ f"bfloat16 requires SM80+. Re-run with "
206
+ f"--kv-cache-dtype float16."
207
+ )
208
+
209
# For FP8 KV cache, we dequantize to BF16 on load inside the
210
# Triton kernel. Tell the common layer not to quantize queries
211
# to FP8 — we handle FP8 KV cache with BF16 queries (Mode 1).