[DSA] Support TileLang FP8 KV cache on CUDA

python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py

@@ -1371,12 +1371,15 @@ def tilelang_sparse_fwd(
     topk = indices.shape[-1]
     assert topk % 64 == 0, "topk must be padded to a multiple of 64"
 
-    if _is_hip:
-        is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
+    is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
+    if _is_hip or is_fp8_kv:
         if is_fp8_kv:
             if q.dtype != kv.dtype:
                 q = q.to(kv.dtype)
-            if _is_gfx95_supported:
+            if not _is_hip:
+                block_I, threads, block_per_cu = 32, 128, 1
+                cu = torch.cuda.get_device_properties(q.device).multi_processor_count
+            elif _is_gfx95_supported:
                 block_I, threads, block_per_cu, cu = 64, 256, 2, 256
             else:
                 block_I, threads, block_per_cu, cu = 64, 256, 1, 304

python/sglang/srt/arg_groups/overrides.py

@@ -529,6 +529,7 @@ def _check_dsa_backend_constraints(
     decode_backend: Optional[str],
     *,
     hip: bool,
+    dcp_size: int = 1,
 ) -> None:
     """Validate DSA backend / platform / kv-cache-dtype constraints."""
     chosen = {prefill_backend, decode_backend}
@@ -541,15 +542,21 @@ def _check_dsa_backend_constraints(
             "(flashmla_kv on Hopper, trtllm on Blackwell)."
         )
 
-    cuda_fp8_unsupported = {"tilelang"} & chosen
-    if not hip and kv_cache_dtype == "fp8_e4m3" and cuda_fp8_unsupported:
+    if hip or kv_cache_dtype != "fp8_e4m3" or "tilelang" not in chosen:
+        return
+    if prefill_backend != "tilelang" or decode_backend != "tilelang":
         raise ValueError(
-            f"The {'/'.join(sorted(cuda_fp8_unsupported))} 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, or keep "
-            "--kv-cache-dtype fp8_e4m3 and pick an fp8-capable DSA backend "
-            "(flashmla_kv on Hopper, trtllm on Blackwell)."
+            "CUDA TileLang FP8 KV requires both DSA backends to be tilelang; "
+            "raw and scaled FP8 KV layouts cannot be mixed."
         )
+    if dcp_size > 1:
+        raise ValueError("CUDA TileLang FP8 KV does not support --dcp-size > 1.")
+
+    import torch
+
+    major, minor = torch.cuda.get_device_capability()
+    if (major, minor) < (8, 9):
+        raise ValueError("CUDA TileLang FP8 KV requires SM89 or newer.")
 
 
 def _check_tilelang_dsa_fp8_kv(
@@ -558,10 +565,11 @@ def _check_tilelang_dsa_fp8_kv(
     decode_backend: Optional[str],
     *,
     hip: bool,
+    dcp_size: int = 1,
 ) -> None:
     """Backward-compatible entry point for the TileLang DSA validation."""
     _check_dsa_backend_constraints(
-        kv_cache_dtype, prefill_backend, decode_backend, hip=hip
+        kv_cache_dtype, prefill_backend, decode_backend, hip=hip, dcp_size=dcp_size
     )
 
 
@@ -652,11 +660,21 @@ def _dsa_split_backend_resolution(view: Any) -> dict:
             declared["dsa_decode_backend"] = backend
         prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend)
         decode = declared.get("dsa_decode_backend", view.dsa_decode_backend)
+        if (
+            not get_platform().is_hip
+            and kv_cache_dtype == "fp8_e4m3"
+            and "tilelang" in (prefill, decode)
+        ):
+            raise ValueError("CUDA TileLang FP8 KV does not support HiSparse.")
         # The hisparse allow-list in hisparse_hook is platform- but not
         # dtype-aware, so an explicitly requested backend still has to clear the
         # shared backend/kv-cache-dtype rules before this arm returns early.
         _check_dsa_backend_constraints(
-            kv_cache_dtype, prefill, decode, hip=get_platform().is_hip
+            kv_cache_dtype,
+            prefill,
+            decode,
+            hip=get_platform().is_hip,
+            dcp_size=getattr(view, "dcp_size", 1) or 1,
         )
         logger.warning(
             f"HiSparse enabled ({kv_cache_dtype}): using DSA backends "
@@ -686,7 +704,11 @@ 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_dsa_backend_constraints(
-        kv_cache_dtype, prefill, decode, hip=get_platform().is_hip
+        kv_cache_dtype,
+        prefill,
+        decode,
+        hip=get_platform().is_hip,
+        dcp_size=getattr(view, "dcp_size", 1) or 1,
     )
     logger.warning(
         f"Set DSA backends for {kv_cache_dtype} KV Cache: "

python/sglang/srt/layers/attention/dsa_backend.py

@@ -473,6 +473,15 @@ def __init__(
         self.supports_mha_one_shot: bool = True
         self.dsa_prefill_impl: _DSA_IMPL_T = get_exec().kernel.dsa_prefill_backend
         self.dsa_decode_impl: _DSA_IMPL_T = get_exec().kernel.dsa_decode_backend
+        if (
+            not _is_hip
+            and self.token_to_kv_pool.dtype == torch.float8_e4m3fn
+            and not self.dsa_kv_cache_store_fp8
+            and "tilelang" in (self.dsa_prefill_impl, self.dsa_decode_impl)
+        ):
+            # The one-shot FP8 dequantizers require scaled KV rows; TileLang
+            # consumes raw FP8 KV without per-block scales.
+            self.supports_mha_one_shot = False
         # Opt-in exact fast path for the Triton sparse-MLA prefill kernel.
         self.dsa_triton_union: int = _resolve_dsa_triton_union(
             union=get_exec().kernel.dsa_triton_union,

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -3015,6 +3015,14 @@ def calculate_mla_kv_cache_dim(
     ):
         return kv_cache_dim
 
+    # CUDA TileLang uses raw FP8 rows without per-block scales. Argument
+    # validation requires both consumers to use the same TileLang layout.
+    if (
+        get_exec().kernel.dsa_prefill_backend == "tilelang"
+        and get_exec().kernel.dsa_decode_backend == "tilelang"
+    ):
+        return kv_cache_dim
+
     quant_block_size = DSATokenToKVPool.quant_block_size
     rope_storage_dtype = DSATokenToKVPool.rope_storage_dtype
     # Calculate override_kv_cache_dim for FP8 storage in backends that use scaled KV layout

test/registered/kernels/ops/attention/test_tilelang_dsa_zero_rope.py

@@ -1,4 +1,4 @@
-"""DSA packed-cache and HIP TileLang kernels must accept GLM's 256+0 geometry without changing 512+64."""
+"""DSA packed-cache and TileLang kernels must accept GLM's 256+0 geometry without changing 512+64."""
 
 import math
 import unittest
@@ -15,7 +15,6 @@
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-# backend-specific: the zero-tail specialization only exists in the HIP TileLang kernels
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_amd_ci(est_time=180, suite="stage-b-test-1-gpu-small-amd-mi35x")
 
@@ -68,8 +67,12 @@ def _torch_sparse_attention(q, kv, indices, scale, d_v):
 
 
 @unittest.skipUnless(
-    torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
-    "the zero-tail TileLang specialization is compiled for gfx950",
+    torch.cuda.is_available()
+    and (
+        (is_hip() and is_gfx95_supported())
+        or (not is_hip() and torch.cuda.get_device_capability() >= (8, 9))
+    ),
+    "TileLang DSA requires gfx950 or SM89+ for these tests",
 )
 class TestTileLangDSAZeroRope(CustomTestCase):
     def _assert_matches_torch(self, use_fp8, d_v, d_tail):
@@ -80,26 +83,36 @@ def _assert_matches_torch(self, use_fp8, d_v, d_tail):
 
         torch.manual_seed(7)
         tokens, heads, topk = 17, 64, 2112
+        if not is_hip() and not use_fp8 and d_tail:
+            # The existing CUDA BF16 tail kernel processes pairs of 128 rows.
+            topk = 2304
         dim = d_v + d_tail
         q = torch.randn(tokens, heads, dim, device="cuda", dtype=torch.bfloat16)
         kv = torch.randn(topk, 1, dim, device="cuda", dtype=torch.bfloat16)
         indices = torch.arange(topk, device="cuda", dtype=torch.int32)
         indices = indices.view(1, 1, topk).expand(tokens, -1, -1).clone()
         indices[..., 2051:] = -1  # padded rows must be masked, not gathered from slot 0
         if use_fp8:
-            q = q.to(FP8_DTYPE)
             kv = kv.to(FP8_DTYPE)
+        reference_q = q.to(FP8_DTYPE) if use_fp8 else q
 
         scale = 1.0 / math.sqrt(dim)
-        expected = _torch_sparse_attention(q, kv, indices, scale, d_v)
-        # the HIP combine kernel returns [batch=1, tokens, heads, d_v]
+        expected = _torch_sparse_attention(reference_q, kv, indices, scale, d_v)
+        # the partial/combine path returns [batch=1, tokens, heads, d_v]
         actual = tilelang_sparse_fwd(q, kv, indices, scale, d_v=d_v).squeeze(0)
         torch.testing.assert_close(
             actual,
             expected,
             atol=0.20 if use_fp8 else 0.04,
             rtol=0.12 if use_fp8 else 0.04,
         )
+        if use_fp8:
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                captured = tilelang_sparse_fwd(q, kv, indices, scale, d_v=d_v)
+            graph.replay()
+            torch.cuda.synchronize()
+            torch.testing.assert_close(captured.squeeze(0), actual, atol=0, rtol=0)
 
     def test_bf16_zero_rope_matches_torch(self):
         """Before the fix the BF16 partial kernel emitted zero-extent tail copies and failed to compile."""

test/registered/unit/test_dsa_tilelang_fp8_validation.py

@@ -1,11 +1,7 @@
-"""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``.
-"""
+"""Validate raw TileLang FP8 KV against incompatible backends and platforms."""
 
 import unittest
+from unittest.mock import patch
 
 from sglang.srt.arg_groups.overrides import _check_tilelang_dsa_fp8_kv
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -23,12 +19,26 @@ def test_cuda_fp8_tilelang_prefill_rejected(self):
         with self.assertRaises(ValueError):
             _check_tilelang_dsa_fp8_kv("fp8_e4m3", "tilelang", "trtllm", hip=False)
 
+    @patch("torch.cuda.get_device_capability", return_value=(9, 0))
+    def test_cuda_fp8_tilelang_allowed(self, _capability):
+        _check_tilelang_dsa_fp8_kv("fp8_e4m3", "tilelang", "tilelang", hip=False)
+
+    @patch("torch.cuda.get_device_capability", return_value=(8, 0))
+    def test_cuda_fp8_requires_fp8_tensor_cores(self, _capability):
+        with self.assertRaisesRegex(ValueError, "SM89"):
+            _check_tilelang_dsa_fp8_kv("fp8_e4m3", "tilelang", "tilelang", hip=False)
+
+    def test_cuda_fp8_tilelang_dcp_rejected(self):
+        with self.assertRaisesRegex(ValueError, "dcp-size"):
+            _check_tilelang_dsa_fp8_kv(
+                "fp8_e4m3", "tilelang", "tilelang", hip=False, dcp_size=2
+            )
+
     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):