fix(dsa): fail fast on fp8_e4m3 KV with tilelang DSA backend on CUDA (#31346)
Co-authored-by: mosya415 <263250241+mosya415@users.noreply.github.com>
This commit is contained in:
@@ -1248,6 +1248,29 @@ def _dsa_kv_cache_dtype_default(view: Any) -> dict:
|
||||
return {}
|
||||
|
||||
|
||||
def _check_tilelang_dsa_fp8_kv(
|
||||
kv_cache_dtype: str,
|
||||
prefill_backend: Optional[str],
|
||||
decode_backend: Optional[str],
|
||||
*,
|
||||
hip: bool,
|
||||
) -> None:
|
||||
"""tilelang's fp8 KV path is ROCm-only; the CUDA kernel hardcodes bfloat16.
|
||||
Reject here instead of crashing at decode CUDA-graph capture."""
|
||||
if (
|
||||
not hip
|
||||
and kv_cache_dtype == "fp8_e4m3"
|
||||
and "tilelang" in {prefill_backend, decode_backend}
|
||||
):
|
||||
raise ValueError(
|
||||
"The tilelang DSA prefill/decode kernels only support an fp8_e4m3 KV "
|
||||
"cache on ROCm/HIP; on CUDA they require a bfloat16 KV cache. Use "
|
||||
"--kv-cache-dtype bfloat16 with the tilelang backend, or keep "
|
||||
"--kv-cache-dtype fp8_e4m3 and pick an fp8-capable DSA backend "
|
||||
"(flashmla_kv on Hopper, trtllm on Blackwell)."
|
||||
)
|
||||
|
||||
|
||||
@register_post_process
|
||||
def _dsa_split_backend_resolution(view: Any) -> dict:
|
||||
"""Slot pass in the DSA arm: default the DSA prefill/decode split
|
||||
@@ -1306,6 +1329,7 @@ def _dsa_split_backend_resolution(view: Any) -> dict:
|
||||
|
||||
prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend)
|
||||
decode = declared.get("dsa_decode_backend", view.dsa_decode_backend)
|
||||
_check_tilelang_dsa_fp8_kv(kv_cache_dtype, prefill, decode, hip=is_hip())
|
||||
logger.warning(
|
||||
f"Set DSA backends for {kv_cache_dtype} KV Cache: "
|
||||
f"prefill={prefill}, decode={decode}."
|
||||
|
||||
Reference in New Issue
Block a user