Repository: sgl-project/sglang
Pull request #4472: Support FlashMLA backend
Author: login: sleepcoo
Created: 2025-03-16T10:07:56Z | Merged: 2025-03-16T16:07:07Z

<!-- Thank you for your contribution! We appreciate it. The following guidelines will help improve your pull request and facilitate feedback. If anything is unclear, don't hesitate to submit your pull request and ask the maintainers for assistance. -->
## Motivation
Integrate flashmla for decoding, and the accuracy test is currently okay. The current implementation is quite simple, directly integrating flashmla as the backend. Later, we need to abstract a fastmla_backend, using fa3 for prefill and flashmla for decode
## Modifications
* FlashMLABackend inherits from FlashInferMLAAttnBackend, using FlashInferMLAAttnBackend for prefill and FlashMLABackend for decoding.
* Add the create_flashmla_kv_indices_triton function to be compatible with the block table format of flashmla.
## command
```
 python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --trust-remote --tp 8 --page-size 64 --disable-cuda-graph --enable-flashmla
```
## Todo
* Support FlashMLA decode with cudagraph
* Enable speculative sampling in FlashMLA
* Add unit test
* Performance Analysis and Optimization
* Integrate FA3 prefill


Diff:
diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py
new file mode 100644
index 000000000000..61a67a8d8f22
--- /dev/null
+++ b/python/sglang/srt/layers/attention/flashmla_backend.py
@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+"""
+Support attention backend for flashMLA.
+
+Current initial integration of FlashMLA shows normal accuracy, but performance is slightly lacking.
+#TODO
+Support FlashMLA decode with cudagraph
+Enable speculative sampling in FlashMLA
+Integrate FA3 prefill
+"""
+
+
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+import triton
+from flash_mla import flash_mla_with_kvcache, get_mla_metadata
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
+from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+
+
+# FlashMLA only supports pagesize=64
+PAGE_SIZE = 64
+
+
+class FlashMLABackend(FlashInferMLAAttnBackend):
+    """Flashinfer attention kernels."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+        kv_last_page_len_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(
+            model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
+        )
+
+        self.num_q_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.num_local_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.kv_lora_rank = model_runner.model_config.kv_lora_rank
+        self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
+        self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
+        self.v_head_dim = model_runner.model_config.v_head_dim
+        self.scaling = model_runner.model_config.scaling
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        cache_loc = forward_batch.out_cache_loc
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer,
+                    cache_loc,
+                    k,
+                    v,
+                )
+        bs = forward_batch.batch_size
+
+        max_seqlen_pad = triton.cdiv(forward_batch.seq_lens.max().item(), PAGE_SIZE)
+        flashmla_index = torch.full(
+            (bs, max_seqlen_pad), -1, dtype=torch.int32, device=q.device
+        )
+        create_flashmla_kv_indices_triton[(bs,)](
+            self.indices_updater_decode.req_to_token,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            None,
+            flashmla_index,
+            self.indices_updater_decode.req_to_token.size(1),
+            flashmla_index.size(1),
+            max_seqlen_pad,
+        )
+
+        mla_metadata, mla_splits = get_mla_metadata(
+            forward_batch.seq_lens.to(torch.int32),
+            1 * self.num_q_heads // self.num_kv_heads,
+            self.num_kv_heads,
+        )
+
+        k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+        reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
+
+        o, _ = flash_mla_with_kvcache(
+            q=reshape_q,
+            k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim),
+            block_table=flashmla_index,
+            cache_seqlens=forward_batch.seq_lens.to(torch.int32),
+            head_dim_v=self.kv_lora_rank,  # TODO Retrieve from config.
+            tile_scheduler_metadata=mla_metadata,
+            num_splits=mla_splits,
+            softmax_scale=layer.scaling,
+            causal=False,
+        )
+
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py
index c64a6c78223d..7fdee7bfea65 100644
--- a/python/sglang/srt/layers/attention/utils.py
+++ b/python/sglang/srt/layers/attention/utils.py
@@ -15,6 +15,7 @@ def create_flashinfer_kv_indices_triton(
     BLOCK_SIZE: tl.constexpr = 512
     pid = tl.program_id(axis=0)
 
+    # find the req pool idx, this is for batch to token
     req_pool_index = tl.load(req_pool_indices_ptr + pid)
     kv_indices_offset = tl.load(kv_indptr + pid)
 
@@ -37,3 +38,56 @@ def create_flashinfer_kv_indices_triton(
             mask=mask,
         )
         tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask)
+
+
+@triton.jit
+def create_flashmla_kv_indices_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices_ptr,
+    page_kernel_lens_ptr,
+    kv_start_idx,
+    kv_indices_ptr,
+    req_to_token_ptr_stride: tl.constexpr,
+    kv_indices_ptr_stride: tl.constexpr,
+    max_pagesize: tl.constexpr,
+):
+    PAGED_SIZE: tl.constexpr = 64
+    BLOCK_SIZE: tl.constexpr = 4096
+    NUM_PAGE_PER_BLOCK: tl.constexpr = 64
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+
+    kv_start = 0
+    kv_end = 0
+    if kv_start_idx:
+        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
+        kv_end = kv_start
+
+    kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
+
+    num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE)
+    num_pages_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+
+    for i in range(num_pages_loop):
+        paged_offset = (
+            tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK
+        ) * PAGED_SIZE
+        paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK
+
+        mask = paged_offset <= num_paged * PAGED_SIZE
+        mask_out = paged_offset_out <= num_paged
+
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + kv_start
+            + paged_offset,
+            mask=mask,
+        )
+        tl.store(
+            kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out,
+            data // PAGED_SIZE,
+            mask=mask_out,
+        )
diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py
index 5cfd6d244ad5..011ff3ed0590 100644
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -71,6 +71,7 @@
     "speculative_accept_threshold_single": ServerArgs.speculative_accept_threshold_single,
     "speculative_accept_threshold_acc": ServerArgs.speculative_accept_threshold_acc,
     "enable_flashinfer_mla": ServerArgs.enable_flashinfer_mla,
+    "enable_flashmla": ServerArgs.enable_flashmla,
     "disable_radix_cache": ServerArgs.disable_radix_cache,
     "flashinfer_mla_disable_ragged": ServerArgs.flashinfer_mla_disable_ragged,
 }
@@ -1273,7 +1274,10 @@ def merge_batch(self, other: "ScheduleBatch"):
 
     def get_model_worker_batch(self) -> ModelWorkerBatch:
         if self.forward_mode.is_decode_or_idle():
-            if global_server_args_dict["enable_flashinfer_mla"]:
+            if (
+                global_server_args_dict["enable_flashinfer_mla"]
+                or global_server_args_dict["enable_flashmla"]
+            ):
                 decode_seq_lens = self.seq_lens.cpu()
             else:
                 decode_seq_lens = None
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index c4150bda139b..71d720afa6ac 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -149,6 +149,7 @@ def __init__(
                 "speculative_accept_threshold_single": server_args.speculative_accept_threshold_single,
                 "speculative_accept_threshold_acc": server_args.speculative_accept_threshold_acc,
                 "enable_flashinfer_mla": server_args.enable_flashinfer_mla,
+                "enable_flashmla": server_args.enable_flashmla,
                 "disable_radix_cache": server_args.disable_radix_cache,
                 "flashinfer_mla_disable_ragged": server_args.flashinfer_mla_disable_ragged,
                 "debug_tensor_dump_output_folder": server_args.debug_tensor_dump_output_folder,
@@ -223,6 +224,9 @@ def model_specific_adjustment(self):
                         "MLA optimization is turned on. Use flashinfer mla backend."
                     )
                     server_args.attention_backend = "flashinfer_mla"
+                elif server_args.enable_flashmla:
+                    logger.info("MLA optimization is turned on. Use flashmla decode.")
+                    server_args.attention_backend = "flashmla"
                 else:
                     logger.info("MLA optimization is turned on. Use triton backend.")
                     server_args.attention_backend = "triton"
@@ -840,6 +844,10 @@ def init_attention_backend(self):
             )
 
             self.attn_backend = FlashInferMLAAttnBackend(self)
+        elif self.server_args.attention_backend == "flashmla":
+            from sglang.srt.layers.attention.flashmla_backend import FlashMLABackend
+
+            self.attn_backend = FlashMLABackend(self)
         else:
             raise ValueError(
                 f"Invalid attention backend: {self.server_args.attention_backend}"
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 9bdcc2d3b96b..545cb62e58e6 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -173,6 +173,7 @@ class ServerArgs:
     tool_call_parser: str = None
     enable_hierarchical_cache: bool = False
     enable_flashinfer_mla: bool = False
+    enable_flashmla: bool = False
     flashinfer_mla_disable_ragged: bool = False
     warmups: Optional[str] = None
 
@@ -227,6 +228,8 @@ def __post_init__(self):
 
         assert self.chunked_prefill_size % self.page_size == 0
 
+        if self.enable_flashmla is True:
+            assert self.page_size == 64, "FlashMLA only support page_size=64"
         # Set cuda graph max batch size
         if self.cuda_graph_max_bs is None:
             # Based on detailed statistics, when serving TP1/TP2 models on lower-end GPUs with HBM<25G, you can either disable cuda graph or set `cuda_graph_max_bs` to a very small value to reduce the memory overhead of creating cuda graphs, with almost no impact on performance. However, when serving models with TP4 or TP8, we need to enable cuda graph to maintain high performance. In this case, we can set `cuda_graph_max_bs` to 80 (half of the default value 160) to reduce the memory overhead of creating cuda graphs. Looking at the logs from TP4 serving of qwen2-72b, a value of 80 is sufficient and can reduce the memory overhead of creating cuda graphs on lower-end GPUs compared to the original 160, avoiding OOM issues.
@@ -753,6 +756,11 @@ def add_cli_args(parser: argparse.ArgumentParser):
             action="store_true",
             help="Enable FlashInfer MLA optimization",
         )
+        parser.add_argument(
+            "--enable-flashmla",
+            action="store_true",
+            help="Enable FlashMLA decode optimization",
+        )
         parser.add_argument(
             "--flashinfer-mla-disable-ragged",
             action="store_true",
diff --git a/scripts/playground/bench_speculative.py b/scripts/playground/bench_speculative.py
index 74370b78d064..72faf706a3ba 100644
--- a/scripts/playground/bench_speculative.py
+++ b/scripts/playground/bench_speculative.py
@@ -182,6 +182,12 @@ def main(args, server_args):
                     "--enable-flashinfer-mla",
                 ]
             )
+        if server_args.enable_flashmla:
+            other_args.extend(
+                [
+                    "--enable-flashmla",
+                ]
+            )
 
         if server_args.quantization:
             other_args.extend(


Reviews:
user: login: zhyncs
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-16T15:57:39Z
user: login: merrymercy
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-16T23:38:30Z
user: login: sleepcoo
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-17T03:14:28Z

Review comments:
diff_hunk: @@ -227,6 +228,8 @@ def __post_init__(self):
 
         assert self.chunked_prefill_size % self.page_size == 0
 
+        if self.enable_flashmla is True:
+            assert self.page_size == 64, "FlashMLA only support page_size=64"
path: python/sglang/srt/server_args.py
user: login: merrymercy
body: automatically set this
created_at: 2025-03-16T23:37:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,128 @@
+from __future__ import annotations
+
+"""
+Support attention backend for flashMLA.
+
+Current initial integration of FlashMLA shows normal accuracy, but performance is slightly lacking.
+#TODO
+Support FlashMLA decode with cudagraph
+Enable speculative sampling in FlashMLA
+Integrate FA3 prefill
+"""
+
+
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+import triton
+from flash_mla import flash_mla_with_kvcache, get_mla_metadata
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
+from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+
+
+# FlashMLA only supports pagesize=64
+PAGE_SIZE = 64
+
+
+class FlashMLABackend(FlashInferMLAAttnBackend):
+    """Flashinfer attention kernels."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+        kv_last_page_len_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(
+            model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
+        )
+
+        self.num_q_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.num_local_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.kv_lora_rank = model_runner.model_config.kv_lora_rank
+        self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
+        self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
+        self.v_head_dim = model_runner.model_config.v_head_dim
+        self.scaling = model_runner.model_config.scaling
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        cache_loc = forward_batch.out_cache_loc
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer,
+                    cache_loc,
+                    k,
+                    v,
+                )
+        bs = forward_batch.batch_size
+
+        max_seqlen_pad = triton.cdiv(forward_batch.seq_lens.max().item(), PAGE_SIZE)
+        flashmla_index = torch.full(
+            (bs, max_seqlen_pad), -1, dtype=torch.int32, device=q.device
+        )
+        create_flashmla_kv_indices_triton[(bs,)](
path: python/sglang/srt/layers/attention/flashmla_backend.py
user: login: merrymercy
body: this metadata is the same for all layers, so we can process them only once in `init_forward_metadata`
created_at: 2025-03-16T23:38:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,128 @@
+from __future__ import annotations
+
+"""
+Support attention backend for flashMLA.
+
+Current initial integration of FlashMLA shows normal accuracy, but performance is slightly lacking.
+#TODO
+Support FlashMLA decode with cudagraph
+Enable speculative sampling in FlashMLA
+Integrate FA3 prefill
+"""
+
+
+from typing import TYPE_CHECKING, Optional, Union
+
+import torch
+import triton
+from flash_mla import flash_mla_with_kvcache, get_mla_metadata
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
+from sglang.srt.layers.attention.utils import create_flashmla_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+
+
+# FlashMLA only supports pagesize=64
+PAGE_SIZE = 64
+
+
+class FlashMLABackend(FlashInferMLAAttnBackend):
+    """Flashinfer attention kernels."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+        kv_last_page_len_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__(
+            model_runner, skip_prefill, kv_indptr_buf, kv_last_page_len_buf
+        )
+
+        self.num_q_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.num_local_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.kv_lora_rank = model_runner.model_config.kv_lora_rank
+        self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
+        self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
+        self.v_head_dim = model_runner.model_config.v_head_dim
+        self.scaling = model_runner.model_config.scaling
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.kv_cache_dim = self.kv_lora_rank + self.qk_rope_head_dim
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
+    ):
+        cache_loc = forward_batch.out_cache_loc
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer,
+                    cache_loc,
+                    k,
+                    v,
+                )
+        bs = forward_batch.batch_size
+
+        max_seqlen_pad = triton.cdiv(forward_batch.seq_lens.max().item(), PAGE_SIZE)
+        flashmla_index = torch.full(
+            (bs, max_seqlen_pad), -1, dtype=torch.int32, device=q.device
+        )
+        create_flashmla_kv_indices_triton[(bs,)](
path: python/sglang/srt/layers/attention/flashmla_backend.py
user: login: sleepcoo
body: I will  fix the issue in the PR of CUDA graph.
created_at: 2025-03-17T03:14:28Z
author_association: COLLABORATOR

PR comments:
user: login: lishicheng1996
created_at: 2025-03-17T13:26:40Z
body: Hi, I test --enable-flashmla on H20*16, and it is slower than normal version. Do you have data about  E2E inference speed gain with   flashmla? Thanks very much!
author_association: NONE
user: login: sleepcoo
created_at: 2025-03-17T13:29:00Z
body: > Hi, I test --enable-flashmla on H20*16, and it become slower than normal version. Do you have data about E2E inference speed gain with flashmla? Thanks very much!

Wait for my cudagraph implementation
author_association: COLLABORATOR