diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index fdc6f9656..9e65c7f3c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -582,6 +582,8 @@ class Envs: # None = standard attention. See https://arxiv.org/abs/2512.12087 SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR = EnvFloat(None) SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR = EnvFloat(None) + # SM120 FlashMLA decode backend: "flashinfer" (default), "triton", or "torch". + SGLANG_SM120_FLASHMLA_BACKEND = EnvStr("flashinfer") # Triton SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/flash_mla_sm120.py b/python/sglang/srt/layers/attention/flash_mla_sm120.py index 8d49d2a6e..e6a6e5c96 100644 --- a/python/sglang/srt/layers/attention/flash_mla_sm120.py +++ b/python/sglang/srt/layers/attention/flash_mla_sm120.py @@ -12,9 +12,13 @@ separate region at the end of each page. """ import logging -import os +import math import torch +import triton +import triton.language as tl + +from sglang.srt.environ import envs logger = logging.getLogger(__name__) @@ -193,17 +197,15 @@ def _sm120_sparse_decode_fwd( return out.to(torch.bfloat16), lse.permute(0, 2, 1) -# Default SM120 FlashMLA backend: "triton" (optimized) or "torch" (pure-PyTorch fallback). -# Controlled by SGLANG_SM120_TRITON_FLASHMLA env var (1=triton, 0=torch). -_sm120_default_backend = ( - "triton" if os.environ.get("SGLANG_SM120_TRITON_FLASHMLA", "1") == "1" else "torch" -) +# SM120 FlashMLA: default FlashInfer (CUTLASS SM120 sparse MLA decode). +# Override with SGLANG_SM120_FLASHMLA_BACKEND=triton|torch to force fallback. +_sm120_default_backend = envs.SGLANG_SM120_FLASHMLA_BACKEND.get() def flash_mla_with_kvcache_sm120(**kwargs): """SM120 FlashMLA sparse decode entry point. - Dispatches to the Triton kernel (default) or PyTorch fallback. + Dispatches to FlashInfer (default if available), Triton, or PyTorch fallback. """ q = kwargs["q"] k_cache = kwargs["k_cache"] @@ -218,6 +220,20 @@ def flash_mla_with_kvcache_sm120(**kwargs): extra_indices = kwargs.get("extra_indices_in_kvcache") extra_topk_length = kwargs.get("extra_topk_length") + if _sm120_default_backend == "flashinfer": + return _flash_mla_flashinfer( + q, + k_cache, + indices, + topk_length, + attn_sink, + head_dim_v, + softmax_scale, + extra_k_cache, + extra_indices, + extra_topk_length, + ) + if _sm120_default_backend == "triton": from sglang.srt.layers.attention.flash_mla_sm120_triton import ( flash_mla_sparse_decode_triton, @@ -250,3 +266,199 @@ def flash_mla_with_kvcache_sm120(**kwargs): extra_topk_length, ) return (out, lse) + + +# --- Page-split utilities: pbs=256 → pbs=64 --- +# SGLang SWA KV cache footer layout per 256-token page: +# [data: 256 * 576 bytes] [scale: 256 * 8 bytes] [padding] +# FlashInfer decode_dsv4 expects per 64-token page: +# [data: 64 * 576 bytes] [scale: 64 * 8 bytes] [padding to 37440] +_PBS_SRC = 256 # SGLang physical page size +_PBS_DST = 64 # FlashInfer page_block_size +_NOPE_ROPE_STRIDE = 576 # bytes per token for nope+rope +_SCALE_STRIDE = 8 # bytes per token for scale (7 + 1 pad) +_BYTES_PER_DST_PAGE = ( + _PBS_DST * _NOPE_ROPE_STRIDE + _PBS_DST * _SCALE_STRIDE +) # 64*576 + 64*8 = 37376 + 512 = 37888 +# Padded to 576 alignment + +_BYTES_PER_DST_PAGE_PADDED = math.ceil(_BYTES_PER_DST_PAGE / 576) * 576 # 37440 + +# Pre-allocated buffer for page-split output per device (lazily sized). +_split_buf = {} # device -> tensor + + +@triton.jit +def _page_split_kernel( + src_ptr, + dst_ptr, + N_pages, + src_stride0: tl.constexpr, + dst_stride0: tl.constexpr, + DATA_PER_SUB: tl.constexpr, # 64 * 576 = 36864 + SCALE_PER_SUB: tl.constexpr, # 64 * 8 = 512 + SRC_SCALE_OFF: tl.constexpr, # 256 * 576 = 147456 + DST_SCALE_OFF: tl.constexpr, # 64 * 576 = 36864 + RATIO: tl.constexpr, # 4 + BLOCK_SIZE: tl.constexpr, +): + """Fused page-split: copy data+scale for all sub-pages in one kernel.""" + pid = tl.program_id(0) + page_idx = pid // RATIO + sub = pid % RATIO + + if page_idx >= N_pages: + return + + src_base = src_ptr + page_idx * src_stride0 + dst_base = dst_ptr + (page_idx * RATIO + sub) * dst_stride0 + + # Copy data region: DATA_PER_SUB bytes from src offset sub*DATA_PER_SUB + data_src_off = sub * DATA_PER_SUB + for start in tl.range(0, DATA_PER_SUB, BLOCK_SIZE): + offs = start + tl.arange(0, BLOCK_SIZE) + mask = offs < DATA_PER_SUB + vals = tl.load(src_base + data_src_off + offs, mask=mask) + tl.store(dst_base + offs, vals, mask=mask) + + # Copy scale region: SCALE_PER_SUB bytes + scale_src_off = SRC_SCALE_OFF + sub * SCALE_PER_SUB + for start in tl.range(0, SCALE_PER_SUB, BLOCK_SIZE): + offs = start + tl.arange(0, BLOCK_SIZE) + mask = offs < SCALE_PER_SUB + vals = tl.load(src_base + scale_src_off + offs, mask=mask) + tl.store(dst_base + DST_SCALE_OFF + offs, vals, mask=mask) + + +def _split_kv_pages_to_64(kv_u8: torch.Tensor, src_pbs: int) -> torch.Tensor: + """Split pbs=N footer-format pages into pbs=64 footer-format pages. + + Uses a fused Triton kernel to do all sub-page copies in a single launch + instead of 8 separate copy kernels (4 sub-pages × 2 regions). + """ + assert src_pbs % _PBS_DST == 0 and src_pbs >= _PBS_DST + if src_pbs == _PBS_DST: + return kv_u8 + + N = kv_u8.shape[0] + ratio = src_pbs // _PBS_DST + num_dst_pages = N * ratio + + dev = kv_u8.device + buf = _split_buf.get(dev) + if buf is None or buf.shape[0] < num_dst_pages: + buf = torch.empty( + num_dst_pages, + _BYTES_PER_DST_PAGE_PADDED, + dtype=torch.uint8, + device=dev, + ) + _split_buf[dev] = buf + out = buf[:num_dst_pages] + + # Get raw 2D view of source + src_2d = kv_u8 + if src_2d.ndim == 4: + src_stride0 = src_2d.stride(0) + src_2d = torch.as_strided(src_2d, (N, src_stride0), (src_stride0, 1)) + else: + src_stride0 = src_2d.stride(0) + + grid = (N * ratio,) + _page_split_kernel[grid]( + src_2d, + out, + N, + src_stride0, + _BYTES_PER_DST_PAGE_PADDED, + _PBS_DST * _NOPE_ROPE_STRIDE, # DATA_PER_SUB = 36864 + _PBS_DST * _SCALE_STRIDE, # SCALE_PER_SUB = 512 + src_pbs * _NOPE_ROPE_STRIDE, # SRC_SCALE_OFF = 147456 + _PBS_DST * _NOPE_ROPE_STRIDE, # DST_SCALE_OFF = 36864 + ratio, # RATIO = 4 + 1024, # BLOCK_SIZE + ) + + bpt = _NOPE_ROPE_STRIDE + _SCALE_STRIDE # 584 + return out.as_strided( + (num_dst_pages, _PBS_DST, 1, bpt), + (_BYTES_PER_DST_PAGE_PADDED, bpt, bpt, 1), + ) + + +def _flash_mla_flashinfer( + q, + k_cache, + indices, + topk_length, + attn_sink, + head_dim_v, + softmax_scale, + extra_k_cache, + extra_indices, + extra_topk_length, +): + """FlashInfer SM120 sparse MLA via sparse_mla_sm120_decode_dsv4. + + SGLang SWA pool uses page_size=256 (footer format: 256*576 bytes data + 256*8 bytes scale). + FlashInfer decode_dsv4 fast path requires page_block_size=64 (footer: 64*576 + 64*8). + We split 256-token pages into 4 virtual 64-token pages. + Token indices are invariant under page-split (identity mapping). + """ + from flashinfer.mla._sparse_mla_sm120 import sparse_mla_sm120_decode_dsv4 + + B, _, H, D = q.shape # (batch, 1, num_heads, head_dim) + dev = q.device + + # --- Page-split: convert pbs=N kv_cache to pbs=64 view --- + kv_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache + src_pbs = k_cache.shape[1] if k_cache.ndim >= 3 else _PBS_SRC + kv_64 = _split_kv_pages_to_64(kv_u8, src_pbs) if src_pbs != _PBS_DST else kv_u8 + + extra_kv_u8 = ( + extra_k_cache.view(torch.uint8) + if extra_k_cache is not None and extra_k_cache.dtype != torch.uint8 + else extra_k_cache + ) + extra_kv_64 = extra_kv_u8 + + # Indices: no remapping needed (page-split preserves token addressing). + idx = indices.squeeze(1) if indices.dim() == 3 else indices + extra_idx = ( + extra_indices.squeeze(1) + if extra_indices is not None and extra_indices.dim() == 3 + else extra_indices + ) + + output = torch.empty(B, H, head_dim_v, dtype=torch.bfloat16, device=dev) + out_lse = torch.empty(B, H, dtype=torch.float32, device=dev) + + # Pre-allocate split-K scratch for decode-dsv4 fast path. + topk = idx.shape[-1] + extra_topk = extra_idx.shape[-1] if extra_idx is not None else 0 + _BI = 64 + num_splits = (topk + _BI - 1) // _BI + ( + (extra_topk + _BI - 1) // _BI if extra_topk > 0 else 0 + ) + mid_out = torch.empty( + B, H, num_splits, head_dim_v, dtype=torch.bfloat16, device=dev + ) + mid_lse = torch.empty(B, H, num_splits, dtype=torch.float32, device=dev) + + sparse_mla_sm120_decode_dsv4( + q=q.squeeze(1) if q.ndim == 4 else q, + kv_cache=kv_64, + indices=idx, + mid_out=mid_out, + mid_lse=mid_lse, + output=output, + out_lse=out_lse, + sm_scale=softmax_scale, + topk_length=topk_length, + attn_sink=attn_sink, + extra_kv_cache=extra_kv_64, + extra_indices=extra_idx, + extra_topk_length=extra_topk_length, + ) + + return (output.unsqueeze(1), None) diff --git a/test/registered/kernels/test_sm120_flash_mla.py b/test/registered/kernels/test_flash_mla_backends.py similarity index 91% rename from test/registered/kernels/test_sm120_flash_mla.py rename to test/registered/kernels/test_flash_mla_backends.py index 82cd45c25..06ab30e85 100644 --- a/test/registered/kernels/test_sm120_flash_mla.py +++ b/test/registered/kernels/test_flash_mla_backends.py @@ -53,6 +53,8 @@ register_cuda_ci(est_time=45, stage="base-b", runner_config="1-gpu-large") # Per-token byte layout _BYTES_PER_TOKEN = _NOPE_ROPE_STRIDE + _SCALE_STRIDE # 576 + 8 = 584 +_IS_SM120 = torch.cuda.is_available() and torch.cuda.get_device_capability() == (12, 0) + def _build_kvcache( num_pages: int, @@ -179,6 +181,7 @@ def _build_q_indices( return q, indices +@unittest.skipUnless(_IS_SM120, "SM120 (compute capability 12.0) required") class TestGatherAndDequant(CustomTestCase): @classmethod def setUpClass(cls): @@ -228,6 +231,7 @@ class TestGatherAndDequant(CustomTestCase): ) +@unittest.skipUnless(_IS_SM120, "SM120 (compute capability 12.0) required") class TestSparseDecodeTritonVsTorch(CustomTestCase): @classmethod def setUpClass(cls): @@ -340,6 +344,7 @@ class TestSparseDecodeTritonVsTorch(CustomTestCase): ) +@unittest.skipUnless(_IS_SM120, "SM120 (compute capability 12.0) required") class TestApplyAttnSink(CustomTestCase): @classmethod def setUpClass(cls): @@ -382,6 +387,7 @@ class TestApplyAttnSink(CustomTestCase): ) +@unittest.skipUnless(_IS_SM120, "SM120 (compute capability 12.0) required") class TestMergePartialAttn(CustomTestCase): @classmethod def setUpClass(cls): @@ -423,6 +429,7 @@ class TestMergePartialAttn(CustomTestCase): torch.testing.assert_close(merged_lse, lse1, atol=1e-5, rtol=1e-5) +@unittest.skipUnless(_IS_SM120, "SM120 (compute capability 12.0) required") class TestEntryPointDispatch(CustomTestCase): @classmethod def setUpClass(cls): @@ -458,6 +465,39 @@ class TestEntryPointDispatch(CustomTestCase): rtol=5e-2, ) + def test_flashinfer_backend_matches_triton(self): + """FlashInfer SM120 sparse MLA decode matches Triton reference.""" + import importlib + + if importlib.util.find_spec("flashinfer.sparse_mla_sm120") is None: + self.skipTest("FlashInfer SM120 sparse MLA not available") + + k_cache, _ = _build_kvcache(4, 64, device=self.device, seed=5) + q, indices = _build_q_indices(1, 4, 32, 4, 64, device=self.device, seed=13) + topk_length = torch.tensor([32], dtype=torch.int32, device=self.device) + + kwargs = dict( + q=q, + k_cache=k_cache, + indices=indices, + topk_length=topk_length, + attn_sink=None, + head_dim_v=_D, + softmax_scale=_D**-0.5, + ) + + with mock.patch.object(fmod, "_sm120_default_backend", "triton"): + out_triton, _ = flash_mla_with_kvcache_sm120(**kwargs) + with mock.patch.object(fmod, "_sm120_default_backend", "flashinfer"): + out_fi, _ = flash_mla_with_kvcache_sm120(**kwargs) + + torch.testing.assert_close( + out_fi.to(torch.float32), + out_triton.to(torch.float32), + atol=5e-2, + rtol=5e-2, + ) + if __name__ == "__main__": import sys