diff --git a/python/sglang/kernels/ops/mamba/lfm_short_conv.py b/python/sglang/kernels/ops/mamba/lfm_short_conv.py new file mode 100644 index 000000000..ffdbf807f --- /dev/null +++ b/python/sglang/kernels/ops/mamba/lfm_short_conv.py @@ -0,0 +1,369 @@ +"""Gate-fused BF16 short convolution used by LFM2 on CUDA. + +LFM2 materializes ``B * x``, transposes it for the generic causal-conv kernel, +then materializes ``C * conv(B * x)``. These kernels preserve both BF16 +materialization points while operating directly on the three projection views. +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +_DISABLE_LFM_FUSED_CONV = os.getenv("SGLANG_DISABLE_LFM_FUSED_CONV", "0") == "1" + + +@triton.jit +def _lfm_short_conv_prefill_kernel( + b_ptr, + c_ptr, + x_ptr, + weight_ptr, + state_ptr, + query_start_loc_ptr, + cache_indices_ptr, + has_initial_state_ptr, + out_ptr, + b_stride_t: tl.constexpr, + c_stride_t: tl.constexpr, + x_stride_t: tl.constexpr, + weight_stride_d: tl.constexpr, + state_stride_slot: tl.constexpr, + state_stride_d: tl.constexpr, + state_stride_w: tl.constexpr, + out_stride_t: tl.constexpr, + blocks_per_seq, + dim: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_D: tl.constexpr, +): + tile = tl.program_id(0) + seq = tile // blocks_per_seq + token_block = tile - seq * blocks_per_seq + start = tl.load(query_start_loc_ptr + seq) + end = tl.load(query_start_loc_ptr + seq + 1) + t = start + token_block * BLOCK_T + tl.arange(0, BLOCK_T) + d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) + tt = t[:, None] + dd = d[None, :] + valid_d = dd < dim + valid_t = tt < end + + slot = tl.load(cache_indices_ptr + seq) + active = slot != PAD_SLOT_ID + has_initial = tl.load(has_initial_state_ptr + seq).to(tl.int1) + state_base = slot * state_stride_slot + dd * state_stride_d + + token0 = tt - 2 + in_sequence0 = (token0 >= start) & (token0 < end) + in_mask0 = valid_d & in_sequence0 & active + b0 = tl.load(b_ptr + token0 * b_stride_t + dd, mask=in_mask0, other=0.0).to( + tl.float32 + ) + x0 = tl.load(x_ptr + token0 * x_stride_t + dd, mask=in_mask0, other=0.0).to( + tl.float32 + ) + new_bx0 = (b0 * x0).to(tl.bfloat16) + state_pos0 = token0 - start + 2 + old0 = tl.load( + state_ptr + state_base + state_pos0 * state_stride_w, + mask=( + valid_d + & active + & has_initial + & (state_pos0 >= 0) + & (state_pos0 < 2) + & ~in_sequence0 + ), + other=0.0, + ) + bx0 = tl.where(in_sequence0, new_bx0, old0).to(tl.bfloat16) + + token1 = tt - 1 + in_sequence1 = (token1 >= start) & (token1 < end) + in_mask1 = valid_d & in_sequence1 & active + b1 = tl.load(b_ptr + token1 * b_stride_t + dd, mask=in_mask1, other=0.0).to( + tl.float32 + ) + x1 = tl.load(x_ptr + token1 * x_stride_t + dd, mask=in_mask1, other=0.0).to( + tl.float32 + ) + new_bx1 = (b1 * x1).to(tl.bfloat16) + state_pos1 = token1 - start + 2 + old1 = tl.load( + state_ptr + state_base + state_pos1 * state_stride_w, + mask=( + valid_d + & active + & has_initial + & (state_pos1 >= 0) + & (state_pos1 < 2) + & ~in_sequence1 + ), + other=0.0, + ) + bx1 = tl.where(in_sequence1, new_bx1, old1).to(tl.bfloat16) + + in_mask2 = valid_d & valid_t + b2 = tl.load(b_ptr + tt * b_stride_t + dd, mask=in_mask2, other=0.0).to(tl.float32) + x2 = tl.load(x_ptr + tt * x_stride_t + dd, mask=in_mask2, other=0.0).to(tl.float32) + bx2 = (b2 * x2).to(tl.bfloat16) + + w0 = tl.load(weight_ptr + dd * weight_stride_d, mask=valid_d, other=0.0).to( + tl.float32 + ) + w1 = tl.load(weight_ptr + dd * weight_stride_d + 1, mask=valid_d, other=0.0).to( + tl.float32 + ) + w2 = tl.load(weight_ptr + dd * weight_stride_d + 2, mask=valid_d, other=0.0).to( + tl.float32 + ) + conv = tl.zeros((BLOCK_T, BLOCK_D), dtype=tl.float32) + conv += bx0.to(tl.float32) * w0 + conv += bx1.to(tl.float32) * w1 + conv += bx2.to(tl.float32) * w2 + conv = conv.to(tl.bfloat16) + c = tl.load( + c_ptr + tt * c_stride_t + dd, + mask=valid_d & valid_t, + other=0.0, + ).to(tl.float32) + active_y = (c * conv.to(tl.float32)).to(tl.bfloat16) + pad_y = (c * bx2.to(tl.float32)).to(tl.bfloat16) + y = tl.where(active, active_y, pad_y) + tl.store( + out_ptr + tt * out_stride_t + dd, + y, + mask=valid_d & valid_t, + ) + + # Exactly one token tile per sequence commits the final two gated inputs. + final_prev = tl.sum(tl.where(tt == end - 1, bx1.to(tl.float32), 0.0), axis=0).to( + tl.bfloat16 + ) + final_cur = tl.sum(tl.where(tt == end - 1, bx2.to(tl.float32), 0.0), axis=0).to( + tl.bfloat16 + ) + final_block = (token_block * BLOCK_T <= end - start - 1) & ( + end - start - 1 < (token_block + 1) * BLOCK_T + ) + final_mask = valid_d & active & final_block + tl.store( + state_ptr + state_base, + final_prev[None, :], + mask=final_mask, + ) + tl.store( + state_ptr + state_base + state_stride_w, + final_cur[None, :], + mask=final_mask, + ) + + +@triton.jit +def _lfm_short_conv_decode_kernel( + b_ptr, + c_ptr, + x_ptr, + weight_ptr, + state_ptr, + cache_indices_ptr, + out_ptr, + b_stride_t: tl.constexpr, + c_stride_t: tl.constexpr, + x_stride_t: tl.constexpr, + weight_stride_d: tl.constexpr, + state_stride_slot: tl.constexpr, + state_stride_d: tl.constexpr, + state_stride_w: tl.constexpr, + out_stride_t: tl.constexpr, + dim: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_D: tl.constexpr, +): + token = tl.program_id(0) + d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) + valid_d = d < dim + slot = tl.load(cache_indices_ptr + token) + active = slot != PAD_SLOT_ID + state_base = slot * state_stride_slot + d * state_stride_d + state_mask = valid_d & active + + prev0 = tl.load(state_ptr + state_base, mask=state_mask, other=0.0).to(tl.float32) + prev1 = tl.load( + state_ptr + state_base + state_stride_w, + mask=state_mask, + other=0.0, + ).to(tl.float32) + w0 = tl.load(weight_ptr + d * weight_stride_d, mask=valid_d, other=0.0).to( + tl.float32 + ) + w1 = tl.load(weight_ptr + d * weight_stride_d + 1, mask=valid_d, other=0.0).to( + tl.float32 + ) + w2 = tl.load(weight_ptr + d * weight_stride_d + 2, mask=valid_d, other=0.0).to( + tl.float32 + ) + b = tl.load(b_ptr + token * b_stride_t + d, mask=valid_d, other=0.0).to(tl.float32) + x = tl.load(x_ptr + token * x_stride_t + d, mask=valid_d, other=0.0).to(tl.float32) + bx_bf16 = (b * x).to(tl.bfloat16) + bx = bx_bf16.to(tl.float32) + conv = tl.zeros((BLOCK_D,), dtype=tl.float32) + conv += prev0 * w0 + conv += prev1 * w1 + conv += bx * w2 + conv_bf16 = conv.to(tl.bfloat16) + c = tl.load(c_ptr + token * c_stride_t + d, mask=valid_d, other=0.0).to(tl.float32) + active_y = (c * conv_bf16.to(tl.float32)).to(tl.bfloat16) + pad_y = (c * bx_bf16.to(tl.float32)).to(tl.bfloat16) + y = tl.where(active, active_y, pad_y) + + tl.store(out_ptr + token * out_stride_t + d, y, mask=valid_d) + tl.store(state_ptr + state_base, prev1.to(tl.bfloat16), mask=state_mask) + tl.store(state_ptr + state_base + state_stride_w, bx_bf16, mask=state_mask) + + +def can_use_fused_lfm_short_conv( + b: torch.Tensor, + c: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + state: torch.Tensor, +) -> bool: + """Return whether the exact LFM2 BF16 width-3 contract is satisfied.""" + return ( + not _DISABLE_LFM_FUSED_CONV + and b.is_cuda + and torch.version.hip is None + and torch.cuda.get_device_capability(b.device) == (9, 0) + and c.is_cuda + and x.is_cuda + and weight.is_cuda + and state.is_cuda + and b.device == c.device == x.device == weight.device == state.device + and b.dtype == c.dtype == x.dtype == torch.bfloat16 + and b.shape == c.shape == x.shape + and b.ndim == 2 + and b.stride(1) == c.stride(1) == x.stride(1) == 1 + and weight.dtype == b.dtype + and weight.ndim == 2 + and weight.shape == (b.shape[1], 3) + and weight.stride(1) == 1 + and bias is None + and state.dtype == b.dtype + and state.ndim == 3 + and state.shape[1:] == (b.shape[1], 2) + ) + + +def can_dispatch_fused_lfm_short_conv( + b: torch.Tensor, + c: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + state: torch.Tensor, +) -> bool: + """Return whether this call is in the end-to-end-qualified serving domain.""" + return ( + b.ndim == 2 + and b.shape[1] == 2048 + and can_use_fused_lfm_short_conv(b, c, x, weight, bias, state) + ) + + +def fused_lfm_short_conv_prefill( + b: torch.Tensor, + c: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + state: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + max_seq_len: int, +) -> torch.Tensor: + """Compute and cache ``C * conv(B * x)`` for packed prefill tokens.""" + tokens, dim = b.shape + out = torch.empty((tokens, dim), dtype=b.dtype, device=b.device) + block_t = 32 + block_d = 64 + blocks_per_seq = triton.cdiv(max_seq_len, block_t) + num_sequences = query_start_loc.numel() - 1 + _lfm_short_conv_prefill_kernel[ + (num_sequences * blocks_per_seq, triton.cdiv(dim, block_d)) + ]( + b, + c, + x, + weight, + state, + query_start_loc, + cache_indices, + has_initial_state, + out, + b.stride(0), + c.stride(0), + x.stride(0), + weight.stride(0), + state.stride(0), + state.stride(1), + state.stride(2), + out.stride(0), + blocks_per_seq, + dim, + -1, + BLOCK_T=block_t, + BLOCK_D=block_d, + num_warps=4, + ) + return out + + +def fused_lfm_short_conv_decode( + b: torch.Tensor, + c: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + state: torch.Tensor, + cache_indices: torch.Tensor, +) -> torch.Tensor: + """Compute, cache, and gate one decode token per active request.""" + tokens, dim = b.shape + out = torch.empty((tokens, dim), dtype=b.dtype, device=b.device) + block_d = 256 + _lfm_short_conv_decode_kernel[(tokens, triton.cdiv(dim, block_d))]( + b, + c, + x, + weight, + state, + cache_indices, + out, + b.stride(0), + c.stride(0), + x.stride(0), + weight.stride(0), + state.stride(0), + state.stride(1), + state.stride(2), + out.stride(0), + dim, + -1, + BLOCK_D=block_d, + num_warps=8, + ) + return out + + +__all__ = [ + "can_dispatch_fused_lfm_short_conv", + "can_use_fused_lfm_short_conv", + "fused_lfm_short_conv_decode", + "fused_lfm_short_conv_prefill", +] diff --git a/python/sglang/srt/models/lfm2_moe.py b/python/sglang/srt/models/lfm2_moe.py index b8ffaa2dd..cea383e87 100644 --- a/python/sglang/srt/models/lfm2_moe.py +++ b/python/sglang/srt/models/lfm2_moe.py @@ -17,6 +17,11 @@ from typing import Iterable, List, Optional, Set, Tuple import torch from torch import nn +from sglang.kernels.ops.mamba.lfm_short_conv import ( + can_dispatch_fused_lfm_short_conv, + fused_lfm_short_conv_decode, + fused_lfm_short_conv_prefill, +) from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig from sglang.srt.distributed import get_pp_group from sglang.srt.layers.activation import SiluAndMul @@ -342,9 +347,26 @@ class Lfm2MoeShortConv(nn.Module): proj, _ = self.in_proj(hidden_states) B_gate, C_gate, x = proj.chunk(3, dim=-1) - Bx = B_gate * x + use_fused = can_dispatch_fused_lfm_short_conv( + B_gate, + C_gate, + x, + self.conv_weight, + self.conv_bias, + conv_state, + ) - if forward_batch.forward_mode.is_decode(): + if forward_batch.forward_mode.is_decode() and use_fused: + gated_conv_out = fused_lfm_short_conv_decode( + B_gate, + C_gate, + x, + self.conv_weight, + conv_state, + meta.cache_indices, + ) + elif forward_batch.forward_mode.is_decode(): + Bx = B_gate * x conv_out = causal_conv1d_update( Bx, conv_state, @@ -353,11 +375,32 @@ class Lfm2MoeShortConv(nn.Module): activation=None, conv_state_indices=meta.cache_indices, ) + gated_conv_out = C_gate * conv_out elif forward_batch.forward_mode.is_target_verify(): + Bx = B_gate * x conv_out = shortconv_target_verify( self, Bx, meta, forward_batch.spec_info.draft_token_num, "LFM2-MoE" ) + gated_conv_out = C_gate * conv_out + elif ( + use_fused + and meta.query_start_loc is not None + and meta.has_initial_state is not None + and forward_batch.extend_seq_lens_cpu + ): + gated_conv_out = fused_lfm_short_conv_prefill( + B_gate, + C_gate, + x, + self.conv_weight, + conv_state, + meta.query_start_loc, + meta.cache_indices, + meta.has_initial_state, + max(forward_batch.extend_seq_lens_cpu), + ) else: + Bx = B_gate * x Bx_t = Bx.transpose(0, 1).contiguous() conv_out = causal_conv1d_fn( Bx_t, @@ -369,8 +412,9 @@ class Lfm2MoeShortConv(nn.Module): conv_states=conv_state, activation=None, ).transpose(0, 1) + gated_conv_out = C_gate * conv_out - output, _ = self.out_proj(C_gate * conv_out) + output, _ = self.out_proj(gated_conv_out) return output diff --git a/test/registered/kernels/ops/mamba/test_lfm_short_conv.py b/test/registered/kernels/ops/mamba/test_lfm_short_conv.py new file mode 100644 index 000000000..26f173d11 --- /dev/null +++ b/test/registered/kernels/ops/mamba/test_lfm_short_conv.py @@ -0,0 +1,242 @@ +"""Correctness coverage for the gate-fused LFM2 short convolution.""" + +import unittest + +import torch + +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from sglang.kernels.ops.mamba.lfm_short_conv import ( + can_dispatch_fused_lfm_short_conv, + can_use_fused_lfm_short_conv, + fused_lfm_short_conv_decode, + fused_lfm_short_conv_prefill, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large") + +PAD_SLOT_ID = -1 +requires_sm90 = unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0), + "requires SM90 (Hopper)", +) + + +def _inputs(tokens: int, dim: int, slots: int, seed: int): + generator = torch.Generator(device="cuda").manual_seed(seed) + + def make(shape): + return torch.randn( + shape, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + + return ( + make((tokens, dim)), + make((tokens, dim)), + make((tokens, dim)), + make((dim, 3)), + make((slots, dim, 2)), + ) + + +def _reference_prefill( + b, + c, + x, + weight, + state, + query_start_loc, + cache_indices, + has_initial_state, + seq_lens, +): + bx = b * x + conv = causal_conv1d_fn( + bx.transpose(0, 1).contiguous(), + weight, + None, + conv_states=state, + query_start_loc=query_start_loc, + seq_lens_cpu=seq_lens, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=None, + ).transpose(0, 1) + return c * conv + + +def _reference_decode(b, c, x, weight, state, cache_indices): + conv = causal_conv1d_update( + b * x, + state, + weight, + None, + activation=None, + conv_state_indices=cache_indices, + ) + return c * conv + + +@requires_sm90 +class TestLFMShortConv(CustomTestCase): + def test_prefill_matches_reference_for_variable_sequences(self): + seq_lens = [1, 2, 33, 65] + tokens = sum(seq_lens) + b, c, x, weight, initial_state = _inputs(tokens, 2051, 3, seed=1) + query_start_loc = torch.tensor( + [0, 1, 3, 36, tokens], device="cuda", dtype=torch.int32 + ) + cache_indices = torch.tensor( + [0, 1, PAD_SLOT_ID, 2], device="cuda", dtype=torch.int32 + ) + has_initial_state = torch.tensor( + [False, True, True, False], device="cuda", dtype=torch.bool + ) + + expected_state = initial_state.clone() + expected = _reference_prefill( + b, + c, + x, + weight, + expected_state, + query_start_loc, + cache_indices, + has_initial_state, + seq_lens, + ) + actual_state = initial_state.clone() + actual = fused_lfm_short_conv_prefill( + b, + c, + x, + weight, + actual_state, + query_start_loc, + cache_indices, + has_initial_state, + max(seq_lens), + ) + + active_tokens = torch.cat( + [ + torch.full( + (seq_len,), + cache_index != PAD_SLOT_ID, + device="cuda", + dtype=torch.bool, + ) + for seq_len, cache_index in zip(seq_lens, cache_indices.tolist()) + ] + ) + torch.testing.assert_close( + actual[active_tokens].float(), + expected[active_tokens].float(), + rtol=0.02, + atol=0.05, + ) + self.assertTrue(torch.equal(actual_state, expected_state)) + + def test_decode_matches_reference_and_preserves_padded_slots(self): + tokens, dim, slots = 5, 2051, 4 + b, c, x, weight, initial_state = _inputs(tokens, dim, slots, seed=2) + cache_indices = torch.tensor( + [2, PAD_SLOT_ID, 0, 3, PAD_SLOT_ID], + device="cuda", + dtype=torch.int32, + ) + + expected_state = initial_state.clone() + expected = _reference_decode(b, c, x, weight, expected_state, cache_indices) + actual_state = initial_state.clone() + actual = fused_lfm_short_conv_decode( + b, c, x, weight, actual_state, cache_indices + ) + + active_tokens = cache_indices != PAD_SLOT_ID + torch.testing.assert_close( + actual[active_tokens].float(), + expected[active_tokens].float(), + rtol=0.02, + atol=0.05, + ) + self.assertTrue(torch.equal(actual_state, expected_state)) + self.assertTrue(torch.equal(actual_state[1], initial_state[1])) + + def test_decode_cuda_graph_replay(self): + tokens, dim, slots = 4, 2048, 4 + b, c, x, weight, initial_state = _inputs(tokens, dim, slots, seed=3) + cache_indices = torch.arange(tokens, device="cuda", dtype=torch.int32) + + # Compile before capture, then prove the captured call reads live inputs and + # writes the caller-owned state/output on replay. + fused_lfm_short_conv_decode( + b, c, x, weight, initial_state.clone(), cache_indices + ) + graph_state = initial_state.clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = fused_lfm_short_conv_decode( + b, c, x, weight, graph_state, cache_indices + ) + + b.add_(torch.tensor(0.25, device="cuda", dtype=b.dtype)) + x.mul_(torch.tensor(0.75, device="cuda", dtype=x.dtype)) + expected_state = initial_state.clone() + expected = _reference_decode(b, c, x, weight, expected_state, cache_indices) + graph_state.copy_(initial_state) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close( + actual.float(), expected.float(), rtol=0.02, atol=0.05 + ) + self.assertTrue(torch.equal(graph_state, expected_state)) + + def test_dispatch_is_narrow(self): + b, c, x, weight, state = _inputs(4, 257, 4, seed=4) + self.assertTrue(can_use_fused_lfm_short_conv(b, c, x, weight, None, state)) + self.assertFalse( + can_dispatch_fused_lfm_short_conv(b, c, x, weight, None, state) + ) + + b_2048, c_2048, x_2048, weight_2048, state_2048 = _inputs(4, 2048, 4, seed=5) + self.assertTrue( + can_dispatch_fused_lfm_short_conv( + b_2048, c_2048, x_2048, weight_2048, None, state_2048 + ) + ) + self.assertFalse( + can_use_fused_lfm_short_conv( + b, c, x, weight, torch.zeros(257, device="cuda"), state + ) + ) + self.assertFalse( + can_use_fused_lfm_short_conv( + b, c, x, weight[:, :2].contiguous(), None, state + ) + ) + self.assertFalse( + can_use_fused_lfm_short_conv(b, c.float(), x, weight, None, state) + ) + self.assertFalse( + can_use_fused_lfm_short_conv( + b.transpose(0, 1), + c.transpose(0, 1), + x.transpose(0, 1), + weight, + None, + state, + ) + ) + + +if __name__ == "__main__": + unittest.main()