diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 88eb3689a..95781121a 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -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}." diff --git a/test/registered/unit/test_dsa_tilelang_fp8_validation.py b/test/registered/unit/test_dsa_tilelang_fp8_validation.py new file mode 100644 index 000000000..212cd5a4f --- /dev/null +++ b/test/registered/unit/test_dsa_tilelang_fp8_validation.py @@ -0,0 +1,41 @@ +"""Rejecting an fp8_e4m3 KV cache with the tilelang DSA backend on CUDA. + +Regression: the combination used to boot the server and crash at decode +CUDA-graph capture with ``kernel main input KV dtype expected bfloat16, +but got float8_e4m3fn``. +""" + +import unittest + +from sglang.srt.arg_groups.overrides import _check_tilelang_dsa_fp8_kv +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestDsaTilelangFp8Validation(CustomTestCase): + + def test_cuda_fp8_tilelang_decode_rejected(self): + with self.assertRaises(ValueError): + _check_tilelang_dsa_fp8_kv("fp8_e4m3", "flashmla_kv", "tilelang", hip=False) + + def test_cuda_fp8_tilelang_prefill_rejected(self): + with self.assertRaises(ValueError): + _check_tilelang_dsa_fp8_kv("fp8_e4m3", "tilelang", "trtllm", hip=False) + + def test_hip_fp8_tilelang_allowed(self): + # ROCm has a real fp8 tilelang kernel + _check_tilelang_dsa_fp8_kv("fp8_e4m3", "tilelang", "tilelang", hip=True) + + def test_bf16_tilelang_allowed(self): + # what the CUDA kernel expects + _check_tilelang_dsa_fp8_kv("bfloat16", "tilelang", "tilelang", hip=False) + + def test_cuda_fp8_non_tilelang_allowed(self): + # fp8-capable backends must pass + _check_tilelang_dsa_fp8_kv("fp8_e4m3", "flashmla_kv", "trtllm", hip=False) + + +if __name__ == "__main__": + unittest.main()