From dc0944d71167a50b10e723c1a705d8faed070c03 Mon Sep 17 00:00:00 2001 From: Zhongze Jiang Date: Wed, 7 Oct 2026 18:51:24 +0800 Subject: [PATCH] feat(dsa): support TileLang FP8 KV cache on CUDA --- .../ops/attention/dsa/tilelang_kernel.py | 9 ++-- python/sglang/srt/arg_groups/overrides.py | 42 ++++++++++++++----- .../srt/layers/attention/dsa_backend.py | 9 ++++ .../srt/mem_cache/kv_cache_configurator.py | 8 ++++ .../attention/test_tilelang_dsa_zero_rope.py | 27 ++++++++---- .../unit/test_dsa_tilelang_fp8_validation.py | 24 +++++++---- 6 files changed, 92 insertions(+), 27 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py b/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py index bd9a37b531c5..b6a1d6a7e465 100644 --- a/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py +++ b/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 diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index bb7641350252..3200f7d5d8aa 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/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 ) @@ -657,11 +665,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 " @@ -691,7 +709,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: " diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 4794c46022b0..e83ddfad6bba 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 1e1667f20bf5..e3d1a4123c07 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -2647,6 +2647,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 diff --git a/test/registered/kernels/ops/attention/test_tilelang_dsa_zero_rope.py b/test/registered/kernels/ops/attention/test_tilelang_dsa_zero_rope.py index 1458bd73bd86..5ccfe6010132 100644 --- a/test/registered/kernels/ops/attention/test_tilelang_dsa_zero_rope.py +++ b/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,6 +83,9 @@ 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) @@ -87,12 +93,12 @@ def _assert_matches_torch(self, use_fp8, d_v, d_tail): 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, @@ -100,6 +106,13 @@ def _assert_matches_torch(self, use_fp8, d_v, d_tail): 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.""" diff --git a/test/registered/unit/test_dsa_tilelang_fp8_validation.py b/test/registered/unit/test_dsa_tilelang_fp8_validation.py index 42ed533fc7c1..14888d42dd2d 100644 --- a/test/registered/unit/test_dsa_tilelang_fp8_validation.py +++ b/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):