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>

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}."

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()