diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index 81d9bfc17..2193411a7 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -19,8 +19,15 @@ import triton import triton.language as tl from sglang.srt.environ import envs +from sglang.srt.utils import is_hip logger = logging.getLogger(__name__) +_is_hip = is_hip() + +_GLM_DSA_MODEL_ARCHS = ( + "GlmMoeDsaForCausalLM", + "GlmMoeDsaForCausalLMNextN", +) # Page layout constants for DSv4-Flash (MODEL1): # nope_dim = 448, rope_dim = 64, quantize_block_size = 64 @@ -474,3 +481,78 @@ def _flash_mla_flashinfer( ) return (output.unsqueeze(1), None) + + +def _validate_flashinfer_sparse_mla_backend( + *, + model_arch: str, + device_sm_major: int, + kv_cache_dtype: torch.dtype, + prefill_impl: str, + decode_impl: str, +) -> bool: + selected = {prefill_impl, decode_impl} + uses_flashinfer_sparse_mla = "flashinfer_sparse_mla" in selected + is_glm_sm12_fp8 = ( + model_arch in _GLM_DSA_MODEL_ARCHS + and device_sm_major == 12 + and kv_cache_dtype == torch.float8_e4m3fn + and not _is_hip + ) + if uses_flashinfer_sparse_mla and not is_glm_sm12_fp8: + raise ValueError( + "flashinfer_sparse_mla supports only GLM DSA with FP8 KV cache " + "on NVIDIA SM120/SM121; " + f"got model_arch={model_arch!r}, sm_major={device_sm_major}, " + f"kv_cache_dtype={kv_cache_dtype}, prefill_impl={prefill_impl!r}, " + f"decode_impl={decode_impl!r}." + ) + if is_glm_sm12_fp8: + unsupported = selected - {"flashinfer_sparse_mla"} + if unsupported: + raise ValueError( + "GLM DSA with FP8 KV cache on NVIDIA SM120/SM121 supports " + "only flashinfer_sparse_mla, " + f"but got {sorted(unsupported)}." + ) + return uses_flashinfer_sparse_mla + + +def flashinfer_sparse_mla_forward( + q: torch.Tensor, + kv_cache: torch.Tensor, + indices: torch.Tensor, + seq_lens: torch.Tensor, + workspace_buffer: torch.Tensor, + *, + page_size: int, + kv_cache_dim: int, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + sm_scale: float, + skip_softmax_threshold_scale_factor: float | None, +) -> torch.Tensor: + """Run FlashInfer's SM120 sparse MLA kernel on SGLang's packed DSA cache.""" + from flashinfer.mla import trtllm_batch_decode_with_kv_cache_mla + + topk = indices.shape[1] + result = trtllm_batch_decode_with_kv_cache_mla( + query=q.unsqueeze(1), + kv_cache=kv_cache.view(torch.uint8) + .view(-1, page_size, kv_cache_dim) + .unsqueeze(1), + workspace_buffer=workspace_buffer, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + block_tables=indices.unsqueeze(1), + seq_lens=seq_lens, + max_seq_len=topk, + sparse_mla_top_k=topk, + bmm1_scale=float(sm_scale), + bmm2_scale=1.0, + kv_scale_format="arbitrary_fp32", + skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor, + ) + return result.squeeze(1) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 95781121a..842d7246d 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1293,6 +1293,25 @@ def _dsa_split_backend_resolution(view: Any) -> dict: user_set_prefill = view.dsa_prefill_backend is not None user_set_decode = view.dsa_decode_backend is not None declared: Dict[str, Any] = {} + model_arch = hf_config.architectures[0] + is_glm_sm12_fp8 = ( + model_arch == "GlmMoeDsaForCausalLM" + and major == 12 + and kv_cache_dtype == "fp8_e4m3" + and not is_hip() + ) + + if is_glm_sm12_fp8: + backend = "flashinfer_sparse_mla" + if not user_set_prefill: + declared["dsa_prefill_backend"] = backend + if not user_set_decode: + declared["dsa_decode_backend"] = backend + logger.warning( + "Set DSA backends for GLM FP8 KV Cache on SM120/SM121: " + f"prefill={backend}, decode={backend}." + ) + return declared if view.enable_hisparse: from sglang.srt.arg_groups.hisparse_hook import _hisparse_default_backend diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index f8871b74c..f326b3a45 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -331,7 +331,13 @@ class DSAIndexerMetadata(BaseIndexerMetadata): _DSA_IMPL_T: TypeAlias = Literal[ - "flashmla_sparse", "flashmla_sparse_q8", "flashmla_kv", "fa3", "tilelang", "trtllm" + "flashmla_sparse", + "flashmla_sparse_q8", + "flashmla_kv", + "flashinfer_sparse_mla", + "fa3", + "tilelang", + "trtllm", ] @@ -490,8 +496,29 @@ class DeepseekSparseAttnBackend( self._q8kv8_identity_scale: Optional[torch.Tensor] = None self._q8kv8_qpad_buf: Optional[torch.Tensor] = None + from sglang.kernels.ops.attention.flash_mla_sm120 import ( + _validate_flashinfer_sparse_mla_backend, + ) + + uses_flashinfer_sparse_mla = _validate_flashinfer_sparse_mla_backend( + model_arch=model_runner.model_config.hf_config.architectures[0], + device_sm_major=self.device_sm_major, + kv_cache_dtype=self.kv_cache_dtype, + prefill_impl=self.dsa_prefill_impl, + decode_impl=self.dsa_decode_impl, + ) + + if uses_flashinfer_sparse_mla: + self.workspace_buffer = get_buffer( + "dsa_flashinfer_sparse_mla_workspace", + lambda: torch.zeros( + envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), + dtype=torch.uint8, + device=model_runner.device, + ), + ) # Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode) - if self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm": + elif self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm": self.workspace_buffer = get_buffer( "dsa_trtllm_workspace", lambda: torch.empty( @@ -2090,6 +2117,19 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) + elif dsa_impl == "flashinfer_sparse_mla": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + if topk_transform_method == TopkTransformMethod.RAGGED: + page_table_1 = topk_indices + return self._forward_flashinfer_sparse_mla( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + seq_lens=metadata.dsa_cache_seqlens_int32, + sm_scale=layer.scaling, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(), + ) elif dsa_impl == "flashmla_kv": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) @@ -2233,6 +2273,17 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) + elif self.dsa_decode_impl == "flashinfer_sparse_mla": + if q_all is None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashinfer_sparse_mla( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + seq_lens=metadata.dsa_cache_seqlens_int32, + sm_scale=layer.scaling, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + ) elif self.dsa_decode_impl == "flashmla_kv": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) @@ -2502,6 +2553,35 @@ class DeepseekSparseAttnBackend( o = o[:, :num_heads, :] return o + def _forward_flashinfer_sparse_mla( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + page_table_1: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + skip_softmax_threshold_scale_factor: float | None, + ) -> torch.Tensor: + from sglang.kernels.ops.attention.flash_mla_sm120 import ( + flashinfer_sparse_mla_forward, + ) + + assert self.workspace_buffer is not None + return flashinfer_sparse_mla_forward( + q=q_all, + kv_cache=kv_cache, + indices=page_table_1, + seq_lens=seq_lens, + workspace_buffer=self.workspace_buffer, + page_size=self.real_page_size, + kv_cache_dim=self.kv_cache_dim, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + sm_scale=sm_scale, + skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor, + ) + def _forward_flashmla_kv( self, q_all: torch.Tensor, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f55e275ba..9bf05be7e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -324,6 +324,7 @@ DSA_CHOICES = [ "flashmla_sparse_q8", "flashmla_kv", "flashmla_auto", + "flashinfer_sparse_mla", "fa3", "tilelang", "aiter", diff --git a/scripts/release/bump_flashinfer_version.py b/scripts/release/bump_flashinfer_version.py index 6d873c16b..d4bff2d62 100755 --- a/scripts/release/bump_flashinfer_version.py +++ b/scripts/release/bump_flashinfer_version.py @@ -20,7 +20,8 @@ def read_current_flashinfer_version(repo_root: Path) -> str: pyproject = repo_root / "python" / "pyproject.toml" content = pyproject.read_text() match = re.search( - r"flashinfer_python==(\d+\.\d+\.\d+(?:rc\d+|\.post\d+)?)", content + r"flashinfer_python(?:\[[^\]]+\])?==" r"(\d+\.\d+\.\d+(?:rc\d+|\.post\d+)?)", + content, ) if not match: raise ValueError(f"Could not find flashinfer_python version in {pyproject}") @@ -39,8 +40,10 @@ def replace_flashinfer_version( name = file_path.name if name == "pyproject.toml": - new_content = new_content.replace( - f"flashinfer_python=={old_version}", f"flashinfer_python=={new_version}" + new_content = re.sub( + rf"(flashinfer_python(?:\[[^\]]+\])?==){re.escape(old_version)}", + rf"\g<1>{new_version}", + new_content, ) new_content = new_content.replace( f"flashinfer_cubin=={old_version}", f"flashinfer_cubin=={new_version}" diff --git a/test/manual/test_dsa_alias_cli_registry_env.py b/test/manual/test_dsa_alias_cli_registry_env.py index aa02eeb37..105b95193 100644 --- a/test/manual/test_dsa_alias_cli_registry_env.py +++ b/test/manual/test_dsa_alias_cli_registry_env.py @@ -37,6 +37,7 @@ class TestDSAChoicesAndFields(unittest.TestCase): def test_dsa_choices_is_canonical(self): self.assertIn("fa3", self.DSA_CHOICES) self.assertIn("tilelang", self.DSA_CHOICES) + self.assertIn("flashinfer_sparse_mla", self.DSA_CHOICES) def test_nsa_choices_is_alias(self): self.assertIs( diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 0e8dcda32..c49414abd 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -78,6 +78,8 @@ def _make_model_runner( disaggregation_mode="null", max_running_requests=None, disaggregation_decode_extra_slots=0, + kv_lora_rank=512, + qk_rope_head_dim=64, ): """Create a mock ModelRunner with the fields configurators need.""" mr = MagicMock() @@ -96,6 +98,8 @@ def _make_model_runner( mc = SimpleNamespace() mc.head_dim = head_dim mc.v_head_dim = v_head_dim + mc.kv_lora_rank = kv_lora_rank + mc.qk_rope_head_dim = qk_rope_head_dim mc.is_hybrid_swa = is_hybrid_swa mc.full_attention_layer_ids = ( full_attention_layer_ids @@ -113,7 +117,6 @@ def _make_model_runner( mc.hf_config.get_text_config = lambda: mc.hf_config mc.linear_attn_registry_result = None mr.model_config = mc - mr.kv_cache_dtype = "fake_bf16" sa = SimpleNamespace() @@ -132,6 +135,7 @@ def _make_model_runner( sa.disaggregation_mode = disaggregation_mode sa.max_running_requests = max_running_requests sa.disaggregation_decode_extra_slots = disaggregation_decode_extra_slots + sa.enable_hisparse = False sa.enable_dsa_cache_layer_split = False sa.kv_cache_dtype = "auto" mr.server_args = sa @@ -230,6 +234,44 @@ class TestDefaultConfigurator(unittest.TestCase): self.assertIsNone(config.full_max_total_num_tokens) self.assertIsNone(config.swa_max_total_num_tokens) + @patch( + "sglang.srt.model_executor.pool_configurator.get_dsa_index_head_dim", + return_value=128, + ) + @patch( + "sglang.srt.model_executor.pool_configurator.is_deepseek_dsa", + return_value=True, + ) + @patch( + "sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim", + side_effect=(576, 656), + ) + def test_dsa_mla_cell_size_uses_backend_kv_layout( + self, mock_calculate_mla_kv_cache_dim, _mock_is_dsa, _mock_index_head_dim + ): + num_layers = 2 + raw = _make_model_runner( + num_layers=num_layers, + use_mla_backend=True, + ) + packed = _make_model_runner( + num_layers=num_layers, + use_mla_backend=True, + ) + + with mock_cpu_env(kv_size=1): + from sglang.srt.model_executor.pool_configurator import ( + DefaultPoolConfigurator, + ) + + raw_configurator = DefaultPoolConfigurator(raw) + packed_configurator = DefaultPoolConfigurator(packed) + + # The DSA indexer adds 128 FP8 values and one FP32 scale (4 bytes). + self.assertEqual(raw_configurator._cell_size, (576 + 132) * num_layers) + self.assertEqual(packed_configurator._cell_size, (656 + 132) * num_layers) + self.assertEqual(mock_calculate_mla_kv_cache_dim.call_count, 2) + class TestHybridSWAConfigurator(unittest.TestCase): """Hybrid SWA: full/swa split, ratio, memory invariant.""" diff --git a/test/registered/unit/test_flashinfer_sparse_mla.py b/test/registered/unit/test_flashinfer_sparse_mla.py new file mode 100644 index 000000000..d9817d606 --- /dev/null +++ b/test/registered/unit/test_flashinfer_sparse_mla.py @@ -0,0 +1,123 @@ +import sys +import unittest +from types import ModuleType +from unittest.mock import patch + +import torch + +from sglang.kernels.ops.attention.flash_mla_sm120 import ( + _validate_flashinfer_sparse_mla_backend, + flashinfer_sparse_mla_forward, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestFlashInferSparseMLAAdapter(unittest.TestCase): + def _mock_flashinfer(self, op): + flashinfer = ModuleType("flashinfer") + flashinfer.__path__ = [] + mla = ModuleType("flashinfer.mla") + mla.trtllm_batch_decode_with_kv_cache_mla = op + flashinfer.mla = mla + return patch.dict( + sys.modules, + {"flashinfer": flashinfer, "flashinfer.mla": mla}, + ) + + def test_maps_sglang_layout_to_public_flashinfer_api(self): + captured = {} + + def fake_op(**kwargs): + captured.update(kwargs) + query = kwargs["query"] + return query.new_full((*query.shape[:-1], kwargs["kv_lora_rank"]), 2) + + with self._mock_flashinfer(fake_op): + output = flashinfer_sparse_mla_forward( + q=torch.zeros((2, 8, 576), dtype=torch.bfloat16), + kv_cache=torch.zeros((128, 1, 656), dtype=torch.uint8), + indices=torch.tensor( + [[7, 9, -1, -1], [4, 6, 8, -1]], dtype=torch.int32 + ), + seq_lens=torch.tensor([2, 3], dtype=torch.int32), + workspace_buffer=torch.zeros(1024, dtype=torch.uint8), + page_size=64, + kv_cache_dim=656, + qk_nope_head_dim=192, + kv_lora_rank=512, + qk_rope_head_dim=64, + sm_scale=0.125, + skip_softmax_threshold_scale_factor=0.25, + ) + + self.assertEqual(tuple(captured["query"].shape), (2, 1, 8, 576)) + self.assertEqual(tuple(captured["kv_cache"].shape), (2, 1, 64, 656)) + self.assertEqual(tuple(captured["block_tables"].shape), (2, 1, 4)) + self.assertEqual( + captured["block_tables"].tolist(), + [[[7, 9, -1, -1]], [[4, 6, 8, -1]]], + ) + self.assertEqual(captured["seq_lens"].tolist(), [2, 3]) + self.assertEqual(captured["max_seq_len"], 4) + self.assertEqual(captured["sparse_mla_top_k"], 4) + self.assertEqual(captured["qk_nope_head_dim"], 192) + self.assertEqual(captured["bmm1_scale"], 0.125) + self.assertEqual(captured["bmm2_scale"], 1.0) + self.assertEqual(captured["kv_scale_format"], "arbitrary_fp32") + self.assertEqual(captured["skip_softmax_threshold_scale_factor"], 0.25) + self.assertNotIn("backend", captured) + self.assertEqual(tuple(output.shape), (2, 8, 512)) + self.assertTrue(torch.all(output == 2)) + + +class TestFlashInferSparseMLABackendGate(unittest.TestCase): + def _validate(self, prefill, decode, model_arch="GlmMoeDsaForCausalLM"): + return _validate_flashinfer_sparse_mla_backend( + model_arch=model_arch, + device_sm_major=12, + kv_cache_dtype=torch.float8_e4m3fn, + prefill_impl=prefill, + decode_impl=decode, + ) + + def test_accepts_flashinfer_for_both_phases(self): + for model_arch in ( + "GlmMoeDsaForCausalLM", + "GlmMoeDsaForCausalLMNextN", + ): + with self.subTest(model_arch=model_arch): + self.assertTrue( + self._validate( + "flashinfer_sparse_mla", + "flashinfer_sparse_mla", + model_arch, + ) + ) + + def test_rejects_other_or_mixed_backends(self): + for prefill, decode in ( + ("trtllm", "trtllm"), + ("flashinfer_sparse_mla", "trtllm"), + ): + with self.subTest(prefill=prefill, decode=decode): + with self.assertRaisesRegex(ValueError, "only flashinfer_sparse_mla"): + self._validate(prefill, decode) + + def test_reports_unsupported_configuration(self): + with self.assertRaises(ValueError) as error: + self._validate( + "flashinfer_sparse_mla", + "flashinfer_sparse_mla", + "DeepseekV3ForCausalLM", + ) + + message = str(error.exception) + self.assertIn("model_arch='DeepseekV3ForCausalLM'", message) + self.assertIn("sm_major=12", message) + self.assertIn("kv_cache_dtype=torch.float8_e4m3fn", message) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 2701fe0b4..7fa9275af 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -990,6 +990,20 @@ class TestGoldenModelOverrides(_IsolatedPublish): self.assertEqual( _dsa_split_backend_resolution(_view(arch="LlamaForCausalLM")), {} ) + with ( + patch("sglang.srt.configs.model_config.is_deepseek_dsa", return_value=True), + patch.object(overrides_module, "is_npu", return_value=False), + patch.object(overrides_module, "is_xpu", return_value=False), + patch.object(overrides_module, "is_hip", return_value=False), + patch("torch.cuda.get_device_capability", return_value=(12, 0)), + ): + self.assertEqual( + _dsa_split_backend_resolution(_view(arch="GlmMoeDsaForCausalLM")), + { + "dsa_prefill_backend": "flashinfer_sparse_mla", + "dsa_decode_backend": "flashinfer_sparse_mla", + }, + ) with ( patch("sglang.srt.configs.model_config.is_deepseek_dsa", return_value=True), patch.object(overrides_module, "is_npu", return_value=False),