chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[sgl-kernel] rebase FlashMLA 0217 (#18902)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Fan Yin
2026-03-07 16:30:52 +08:00
committed by GitHub
parent c267bdb805
commit 43d6a32045
3 changed files with 164 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

173
python/sglang/srt/layers/attention/flashmla_backend.py
View File

@@ -1,9 +1,9 @@
from __future__ import annotations
"""
Support attention backend for FlashMLA.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union
@@ -23,11 +23,8 @@ if TYPE_CHECKING:
from sglang.srt.speculative.spec_info import SpecInput
# FlashMLA only supports pagesize=64
PAGE_SIZE = 64
# FlashMLA FP8 issue: https://github.com/deepseek-ai/FlashMLA/issues/56
@dataclass
class FlashMLADecodeMetadata:
@@ -47,8 +44,6 @@ class FlashMLADecodeMetadata:
class FlashMLABackend(FlashInferMLAAttnBackend):
"""Flashmla attention kernels."""
def __init__(
self,
model_runner: ModelRunner,
@@ -76,7 +71,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
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
# Check if KV cache is FP8 (supports both e4m3 and e5m2)
self.is_fp8_kvcache = self.data_type in {
torch.float8_e4m3fn,
torch.float8_e5m2,
@@ -84,8 +78,13 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
def init_forward_metadata(self, forward_batch: ForwardBatch):
self.cuda_graph_kv_indices = None
self.cuda_graph_mla_metadata = None
self.cuda_graph_num_splits = None
self.cuda_graph_mla_metadata_view = None
self.cuda_graph_num_splits_view = None
def init_forward_metadata(self, forward_batch: ForwardBatch):
bs = forward_batch.batch_size
if forward_batch.forward_mode.is_decode_or_idle():
max_seqlen_pad = triton.cdiv(
@@ -143,8 +142,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
# Use FlashMLADecodeMetadata which has the attributes forward_extend expects
self.forward_metadata = FlashMLADecodeMetadata(
mla_metadata,
num_splits,
@@ -160,34 +157,31 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
block_kv_indices: Optional[torch.Tensor] = None,
):
if block_kv_indices is None:
cuda_graph_kv_indices = torch.full(
self.cuda_graph_kv_indices = torch.full(
(max_bs, (self.max_context_len + PAGE_SIZE) // PAGE_SIZE),
1,
dtype=torch.int32,
device="cuda",
)
else:
cuda_graph_kv_indices = block_kv_indices
self.cuda_graph_kv_indices = block_kv_indices
if self.num_draft_tokens:
self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
torch.ones(
max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device
),
self.num_draft_tokens * self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
else:
self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
torch.ones(
max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device
),
self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_kv_indices = cuda_graph_kv_indices
device_props = torch.cuda.get_device_properties(self.req_to_token.device)
max_num_sm_parts = device_props.multi_processor_count
self.cuda_graph_mla_metadata = torch.empty(
(max_num_sm_parts, 8),
dtype=torch.int32,
device="cuda",
)
self.cuda_graph_num_splits = torch.empty(
max_bs + 1,
dtype=torch.int32,
device="cuda",
)
self.cuda_graph_mla_metadata_view = None
self.cuda_graph_num_splits_view = None
def init_forward_metadata_capture_cuda_graph(
self,
@@ -211,20 +205,35 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
num_q_heads = self.num_q_heads * (self.num_draft_tokens or 1)
num_q_heads = self.num_q_heads
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
assert actual_num_sm_parts <= self.cuda_graph_mla_metadata.shape[0], (
f"num_sm_parts {actual_num_sm_parts} exceeds preallocated max "
f"{self.cuda_graph_mla_metadata.shape[0]}"
)
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata = FlashMLADecodeMetadata(
self.cuda_graph_mla_metadata,
self.cuda_graph_num_splits[: bs + 1],
self.cuda_graph_mla_metadata_view,
self.cuda_graph_num_splits_view,
self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
)
elif forward_mode.is_target_verify():
seq_lens = seq_lens + self.num_draft_tokens
max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE)
@@ -238,17 +247,28 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
self.num_draft_tokens * self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
assert actual_num_sm_parts <= self.cuda_graph_mla_metadata.shape[0]
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata = FlashMLADecodeMetadata(
self.cuda_graph_mla_metadata,
self.cuda_graph_num_splits[: bs + 1],
self.cuda_graph_mla_metadata_view,
self.cuda_graph_num_splits_view,
self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
)
else:
@@ -273,12 +293,12 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
):
if forward_mode.is_decode_or_idle():
assert seq_lens_cpu is not None
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE)
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices[:bs],
@@ -288,24 +308,46 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
num_q_heads = self.num_q_heads * (self.num_draft_tokens or 1)
num_q_heads = self.num_q_heads
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
if actual_num_sm_parts != self.cuda_graph_mla_metadata_view.shape[0]:
import logging
logger = logging.getLogger(__name__)
logger.warning(
f"num_sm_parts mismatch in CUDA Graph replay: "
f"capture={self.cuda_graph_mla_metadata_view.shape[0]}, "
f"replay={actual_num_sm_parts}. "
f"This may indicate batch size changed between capture and replay."
)
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata
self.forward_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata_view
self.forward_metadata.num_splits = self.cuda_graph_num_splits_view
self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[
:bs, :max_seqlen_pad
]
elif forward_mode.is_target_verify():
seq_lens = seq_lens[:bs] + self.num_draft_tokens
seq_lens_cpu = seq_lens_cpu[:bs] + self.num_draft_tokens
max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE)
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices[:bs],
@@ -315,16 +357,27 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
)
mla_metadata, num_splits = get_mla_metadata(
seq_lens.to(torch.int32),
self.num_draft_tokens * self.num_q_heads,
1,
is_fp8_kvcache=self.is_fp8_kvcache,
)
self.cuda_graph_mla_metadata.copy_(mla_metadata)
actual_num_sm_parts = mla_metadata.shape[0]
if actual_num_sm_parts != self.cuda_graph_mla_metadata_view.shape[0]:
self.cuda_graph_mla_metadata_view = self.cuda_graph_mla_metadata[
:actual_num_sm_parts
]
self.cuda_graph_num_splits_view = self.cuda_graph_num_splits[: bs + 1]
self.cuda_graph_mla_metadata[:actual_num_sm_parts].copy_(mla_metadata)
self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata
self.forward_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1]
self.forward_metadata.mla_metadata = self.cuda_graph_mla_metadata_view
self.forward_metadata.num_splits = self.cuda_graph_num_splits_view
self.forward_metadata.block_kv_indices = self.cuda_graph_kv_indices[
:bs, :max_seqlen_pad
]
@@ -368,14 +421,11 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
if self.is_fp8_kvcache:
# For FP8 KV cache, Q needs to be converted to FP8 for FlashMLA kernel
# In SGLang, we use layer.k_scale for both q and k scales
if layer.k_scale is not None:
q_scale = layer.k_scale
descale_q = layer.k_scale.reshape(1)
descale_k = layer.k_scale.reshape(1)
else:
# Fallback to 1.0 if k_scale is not initialized
q_scale = torch.ones((1,), dtype=torch.float32, device=reshape_q.device)
descale_q = torch.ones(
(1,), dtype=torch.float32, device=reshape_q.device
@@ -384,7 +434,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
(1,), dtype=torch.float32, device=reshape_q.device
)
# Reshape to 2D for scaled_fp8_quant (which requires 2D input)
q_shape = reshape_q.shape
reshape_q_2d = reshape_q.reshape(-1, q_shape[-1])
reshape_q_fp8_2d, _ = scaled_fp8_quant(reshape_q_2d, q_scale)
@@ -394,7 +443,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim),
block_table=self.forward_metadata.block_kv_indices[:bs],
cache_seqlens=forward_batch.seq_lens.to(torch.int32),
head_dim_v=self.kv_lora_rank, # TODO Retrieve from config.
head_dim_v=self.kv_lora_rank,
tile_scheduler_metadata=self.forward_metadata.flashmla_metadata,
num_splits=self.forward_metadata.num_splits,
softmax_scale=layer.scaling,
@@ -405,13 +454,12 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
else:
# todo: need check all causal True or False?
o, _ = flash_mla_with_kvcache(
q=reshape_q,
k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim),
block_table=self.forward_metadata.block_kv_indices[:bs],
cache_seqlens=forward_batch.seq_lens.to(torch.int32),
head_dim_v=self.kv_lora_rank, # TODO Retrieve from config.
head_dim_v=self.kv_lora_rank,
tile_scheduler_metadata=self.forward_metadata.flashmla_metadata,
num_splits=self.forward_metadata.num_splits,
softmax_scale=layer.scaling,
@@ -447,14 +495,11 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim)
if self.is_fp8_kvcache:
# For FP8 KV cache, Q needs to be converted to FP8 for FlashMLA kernel
# In SGLang, we use layer.k_scale for both q and k scales
if layer.k_scale is not None:
q_scale = layer.k_scale
descale_q = layer.k_scale.reshape(1)
descale_k = layer.k_scale.reshape(1)
else:
# Fallback to 1.0 if k_scale is not initialized
q_scale = torch.ones(
(1,), dtype=torch.float32, device=reshape_q.device
)
@@ -465,8 +510,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
(1,), dtype=torch.float32, device=reshape_q.device
)
# Quantize Q using scaled_fp8_quant (matching vLLM's approach)
# Reshape to 2D for scaled_fp8_quant (which requires 2D input)
q_shape = reshape_q.shape
reshape_q_2d = reshape_q.reshape(-1, q_shape[-1])
reshape_q_fp8_2d, _ = scaled_fp8_quant(reshape_q_2d, q_scale)
@@ -501,13 +544,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
# TODO: multi step kv indices optimization
class FlashMLAMultiStepDraftBackend:
"""
Wrap multiple flashmla attention backends as one for multiple consecutive
draft decoding steps.
"""
def __init__(
self,
model_runner: ModelRunner,
@@ -566,6 +603,10 @@ class FlashMLAMultiStepDraftBackend:
def init_forward_metadata_capture_cuda_graph(self, forward_batch: ForwardBatch):
def call_fn(i, forward_batch):
# EAGLE draft worker uses DECODE mode for draft steps
from sglang.srt.model_executor.forward_batch_info import ForwardMode
# Create a dummy forward_mode for draft step
self.attn_backends[i].init_forward_metadata_capture_cuda_graph(
forward_batch.batch_size,
forward_batch.batch_size * self.topk,
@@ -582,6 +623,8 @@ class FlashMLAMultiStepDraftBackend:
self, forward_batch: ForwardBatch, bs: int
):
def call_fn(i, forward_batch):
from sglang.srt.model_executor.forward_batch_info import ForwardMode
self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
bs,
forward_batch.req_pool_indices,