From 4aadf94146b13770168eeb6f7d55206af8bf1314 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Wed, 15 Jul 2026 12:52:15 +0800 Subject: [PATCH] [Kernel] Relocate vendored fla and mamba kernel trees to sglang.kernels (RFC #29630, Phase 2.5, 7/7) (#30795) Co-authored-by: Claude Fable 5 --- .../bench_cutedsl_kda_decode.py | 4 +-- .../bench_fused_gate_cumsum.py | 6 ++-- .../bench_gdn_decode.py | 4 +-- .../bench_gdn_prefill.py | 4 +-- .../bench_gdn_prefill_cutedsl.py | 8 +++--- .../bench_kda_decode.py | 4 +-- .../bench_kda_prefill_cutedsl.py | 2 +- benchmark/fla/benchmark_layernorm_gated.py | 4 +-- .../triton_flashinfer_cudnn.py | 2 +- .../sglang/kernels/ops/attention/__init__.py | 16 +++++++++++ .../fla/bench_gdn_replayssm_decode.py | 6 ++-- .../ops}/attention/fla/chunk.py | 14 +++++----- .../ops}/attention/fla/chunk_delta_h.py | 6 ++-- .../ops}/attention/fla/chunk_fwd.py | 8 +++--- .../ops}/attention/fla/chunk_intra.py | 10 +++---- .../fla/chunk_intra_token_parallel.py | 4 +-- .../ops}/attention/fla/chunk_o.py | 6 ++-- .../attention/fla/chunk_scaled_dot_kkt.py | 4 +-- .../ops}/attention/fla/cumsum.py | 4 +-- .../ops}/attention/fla/fused_gdn_gating.py | 0 .../ops}/attention/fla/fused_norm_gate.py | 0 .../ops}/attention/fla/fused_recurrent.py | 4 +-- .../fla/fused_recurrent_linear_replayssm.py | 0 .../fla/fused_sigmoid_gating_recurrent.py | 0 .../ops}/attention/fla/index.py | 2 +- .../ops}/attention/fla/kda.py | 18 ++++++------ .../ops}/attention/fla/l2norm.py | 2 +- .../ops}/attention/fla/layernorm_gated.py | 0 .../ops}/attention/fla/op.py | 2 +- .../ops}/attention/fla/solve_tril.py | 4 +-- .../ops}/attention/fla/utils.py | 0 .../ops}/attention/fla/wy_fast.py | 2 +- .../linear/kda_blackwell/__init__.py | 2 +- python/sglang/kernels/ops/mamba/__init__.py | 17 +++++++++++ .../ops}/mamba/causal_conv1d_triton.py | 0 .../ops}/mamba/mamba_state_scatter_triton.py | 0 .../ops/mamba/triton_ops}/__init__.py | 0 .../ops/mamba/triton_ops}/layernorm_gated.py | 0 .../ops/mamba/triton_ops}/mamba_ssm.py | 0 .../ops/mamba/triton_ops}/ssd_bmm.py | 0 .../ops/mamba/triton_ops}/ssd_chunk_scan.py | 0 .../ops/mamba/triton_ops}/ssd_chunk_state.py | 0 .../ops/mamba/triton_ops}/ssd_combined.py | 0 .../mamba/triton_ops}/ssd_state_passing.py | 0 .../ops/mamba/triton_ops}/ssu_dispatch.py | 2 +- .../xpu/kernels/fla/chunk_delta_h.py | 6 ++-- .../xpu/kernels/fla/chunk_fwd.py | 8 +++--- .../fla/fused_sigmoid_gating_recurrent.py | 2 +- .../layers/attention/attention_registry.py | 10 +++++-- .../attention/hybrid_linear_attn_backend.py | 12 ++++---- .../layers/attention/linear/gdn_backend.py | 12 ++++---- .../layers/attention/linear/kda_backend.py | 8 +++--- .../attention/linear/kernels/gdn_cutedsl.py | 2 +- .../linear/kernels/gdn_flashinfer.py | 2 +- .../attention/linear/kernels/gdn_triton.py | 8 +++--- .../attention/linear/kernels/kda_cutedsl.py | 2 +- .../attention/linear/kernels/kda_flashkda.py | 2 +- .../attention/linear/kernels/kda_triton.py | 8 +++--- .../layers/attention/mamba/causal_conv1d.py | 12 ++++++-- .../srt/layers/attention/mamba/mamba.py | 28 +++++++++---------- .../attention/mamba/mixer2_rms_norm_gated.py | 2 +- python/sglang/srt/managers/scheduler.py | 6 ++-- .../sglang/srt/models/bailing_moe_linear.py | 4 +-- python/sglang/srt/models/jet_nemotron.py | 6 ++-- python/sglang/srt/models/kimi_linear.py | 2 +- python/sglang/srt/models/qwen3_5.py | 6 ++-- python/sglang/srt/models/qwen3_next.py | 4 +-- python/sglang/srt/server_args.py | 2 +- .../attention_methods/mamba2_attention.py | 2 +- .../attention/test_chunk_gated_delta_rule.py | 4 +-- .../test_gdn_noncontiguous_stride.py | 4 +-- .../attention/test_gdn_prefill_cutedsl.py | 14 +++++----- test/registered/attention/test_kda_kernels.py | 10 +++---- .../attention/test_kda_prefill_cutedsl.py | 10 +++---- .../attention/test_kda_prefill_flashkda.py | 2 +- .../gdn/test_linear_replayssm_decode.py | 4 +-- test/registered/jit/test_cutedsl_gdn.py | 2 +- .../jit/test_fused_verify_triton_gdn.py | 6 ++-- test/registered/layers/mamba/conftest.py | 4 +-- .../layers/mamba/test_causal_conv1d.py | 2 +- .../registered/layers/mamba/test_mamba_ssm.py | 4 +-- .../layers/mamba/test_mamba_ssm_ssd.py | 2 +- .../layers/test_fla_layernorm_guard.py | 4 +-- .../mem_cache/test_int8_checkpoint_store.py | 2 +- .../layers/test_mamba_state_scatter_triton.py | 2 +- .../unit/mem_cache/test_mamba_unittest.py | 2 +- .../test_unified_radix_cache_bench.py | 2 +- .../test_unified_radix_cache_unittest.py | 2 +- 88 files changed, 226 insertions(+), 183 deletions(-) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/bench_gdn_replayssm_decode.py (96%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/chunk.py (94%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/chunk_delta_h.py (98%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/chunk_fwd.py (98%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/chunk_intra.py (99%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/chunk_intra_token_parallel.py (97%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/chunk_o.py (95%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/chunk_scaled_dot_kkt.py (96%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/cumsum.py (98%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/fused_gdn_gating.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/fused_norm_gate.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/fused_recurrent.py (99%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/fused_recurrent_linear_replayssm.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/fused_sigmoid_gating_recurrent.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/index.py (93%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/kda.py (98%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/l2norm.py (98%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/layernorm_gated.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/op.py (95%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/solve_tril.py (99%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/utils.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/fla/wy_fast.py (98%) rename python/sglang/{srt/layers/attention => kernels/ops}/mamba/causal_conv1d_triton.py (100%) rename python/sglang/{srt/layers/attention => kernels/ops}/mamba/mamba_state_scatter_triton.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/__init__.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/layernorm_gated.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/mamba_ssm.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/ssd_bmm.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/ssd_chunk_scan.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/ssd_chunk_state.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/ssd_combined.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/ssd_state_passing.py (100%) rename python/sglang/{srt/layers/attention/mamba/ops => kernels/ops/mamba/triton_ops}/ssu_dispatch.py (99%) diff --git a/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py b/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py index ea124c487..af138d5d0 100644 --- a/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py +++ b/benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py @@ -18,10 +18,10 @@ import torch import triton from sglang.jit_kernel.cutedsl_kda import cutedsl_fused_sigmoid_gating_kda_update -from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) -from sglang.srt.layers.attention.fla.kda import chunk_kda +from sglang.kernels.ops.attention.fla.kda import chunk_kda def make_inputs( diff --git a/benchmark/bench_linear_attention/bench_fused_gate_cumsum.py b/benchmark/bench_linear_attention/bench_fused_gate_cumsum.py index 1b2c105c3..dc674a218 100644 --- a/benchmark/bench_linear_attention/bench_fused_gate_cumsum.py +++ b/benchmark/bench_linear_attention/bench_fused_gate_cumsum.py @@ -22,9 +22,9 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python") import torch import triton -from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.kda import kda_gate_chunk_cumsum +from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.kda import kda_gate_chunk_cumsum CHUNK_SIZE = 64 diff --git a/benchmark/bench_linear_attention/bench_gdn_decode.py b/benchmark/bench_linear_attention/bench_gdn_decode.py index f292168c7..58bd9d3c3 100644 --- a/benchmark/bench_linear_attention/bench_gdn_decode.py +++ b/benchmark/bench_linear_attention/bench_gdn_decode.py @@ -29,10 +29,10 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python")) import torch import triton -from sglang.srt.layers.attention.fla.fused_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_packed_decode, ) -from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/benchmark/bench_linear_attention/bench_gdn_prefill.py b/benchmark/bench_linear_attention/bench_gdn_prefill.py index 04fdb7c50..f73f846b4 100644 --- a/benchmark/bench_linear_attention/bench_gdn_prefill.py +++ b/benchmark/bench_linear_attention/bench_gdn_prefill.py @@ -29,10 +29,10 @@ from flashinfer.gdn_prefill import ( chunk_gated_delta_rule as flashinfer_chunk_gated_delta_rule, ) -from sglang.srt.layers.attention.fla.chunk import ( +from sglang.kernels.ops.attention.fla.chunk import ( chunk_gated_delta_rule as triton_chunk_gated_delta_rule, ) -from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd +from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd # --------------------------------------------------------------------------- # Helpers diff --git a/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py b/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py index bdcbfc6b5..7359ffffb 100644 --- a/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py +++ b/benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py @@ -26,14 +26,14 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python") import torch +from sglang.kernels.ops.attention.fla.chunk import ( + chunk_gated_delta_rule as triton_chunk_gated_delta_rule, +) +from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd from sglang.kernels.ops.attention.linear.gdn_blackwell import ( chunk_gated_delta_rule_cutedsl, prepare_metadata_cutedsl, ) -from sglang.srt.layers.attention.fla.chunk import ( - chunk_gated_delta_rule as triton_chunk_gated_delta_rule, -) -from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd # --------------------------------------------------------------------------- # Helpers (shared shape: pool layout [N, H, K, V] with K-last stride) diff --git a/benchmark/bench_linear_attention/bench_kda_decode.py b/benchmark/bench_linear_attention/bench_kda_decode.py index e1bd86f35..c06d254c8 100644 --- a/benchmark/bench_linear_attention/bench_kda_decode.py +++ b/benchmark/bench_linear_attention/bench_kda_decode.py @@ -22,10 +22,10 @@ import argparse import torch -from sglang.srt.layers.attention.fla.fused_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_kda_packed_decode, ) -from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py b/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py index e899c777d..9b438d015 100644 --- a/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py +++ b/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py @@ -29,6 +29,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python") import torch import torch.nn.functional as F +from sglang.kernels.ops.attention.fla.kda import chunk_kda, fused_recurrent_kda from sglang.kernels.ops.attention.linear.kda_blackwell import prepare_metadata from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_h import ( kda_h_cutedsl, @@ -42,7 +43,6 @@ from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_o import ( from sglang.kernels.ops.attention.linear.kda_blackwell.prologue import ( kda_prologue, ) -from sglang.srt.layers.attention.fla.kda import chunk_kda, fused_recurrent_kda BT = 64 # chunk size diff --git a/benchmark/fla/benchmark_layernorm_gated.py b/benchmark/fla/benchmark_layernorm_gated.py index e678d8c31..3ae0a291d 100644 --- a/benchmark/fla/benchmark_layernorm_gated.py +++ b/benchmark/fla/benchmark_layernorm_gated.py @@ -4,10 +4,10 @@ import numpy as np import torch # Import the function to benchmark -from sglang.srt.layers.attention.fla.layernorm_gated import ( +from sglang.kernels.ops.attention.fla.layernorm_gated import ( _layer_norm_fwd as layer_norm_fwd, ) -from sglang.srt.layers.attention.fla.layernorm_gated import ( +from sglang.kernels.ops.attention.fla.layernorm_gated import ( rms_norm_ref, ) diff --git a/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py b/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py index 96f839c4f..3a9f9be50 100644 --- a/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py +++ b/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py @@ -7,7 +7,7 @@ import torch.utils.benchmark as benchmark from flashinfer import BatchDecodeWithPagedKVCacheWrapper from sglang.kernels.ops.attention.decode_attention import decode_attention_fwd -from sglang.srt.layers.attention.flashinfer_backend import should_use_tensor_core +from sglang.kernels.ops.attention.flashinfer_backend import should_use_tensor_core def benchmark_forward( diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index 557e41f06..f72763785 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/python/sglang/kernels/ops/attention/__init__.py @@ -43,6 +43,22 @@ del _mod, _fn __all__ = [] +# Vendored linear-attention (flash-linear-attention port) kernels relocated +# in Phase 2.5 (RFC #29630); representative entry points for inventory. +for _mod, _fn in [ + ("fla.chunk", "chunk_gated_delta_rule"), + ("fla.fused_recurrent", "fused_recurrent_gated_delta_rule"), + ("fla.kda", "fused_recurrent_kda_fwd"), +]: + register_kernel( + KernelSpec( + op=f"attention.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.attention.{_mod}:{_fn}", + ) + ) +del _mod, _fn + # Linear-attention / MiniMax-sparse / diffusion kernels migrated in Phase 2.5 # (RFC #29630); registered for inventory. for _grp, _mod, _fn in [ diff --git a/python/sglang/srt/layers/attention/fla/bench_gdn_replayssm_decode.py b/python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py similarity index 96% rename from python/sglang/srt/layers/attention/fla/bench_gdn_replayssm_decode.py rename to python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py index 47e5f3227..58ecd41f3 100644 --- a/python/sglang/srt/layers/attention/fla/bench_gdn_replayssm_decode.py +++ b/python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py @@ -15,7 +15,7 @@ traffic ratio is reported per L. Run:: - python -m sglang.srt.layers.attention.fla.bench_gdn_replayssm_decode + python -m sglang.kernels.ops.attention.fla.bench_gdn_replayssm_decode Requires a GPU (Triton). """ @@ -27,10 +27,10 @@ import argparse import torch import triton -from sglang.srt.layers.attention.fla.fused_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_packed_decode, ) -from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import ( +from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import ( fused_recurrent_gdn_replayssm_decode, ) diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/kernels/ops/attention/fla/chunk.py similarity index 94% rename from python/sglang/srt/layers/attention/fla/chunk.py rename to python/sglang/kernels/ops/attention/fla/chunk.py index c243ba22f..8acb9d14e 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/kernels/ops/attention/fla/chunk.py @@ -7,15 +7,15 @@ from typing import Optional import torch from einops import rearrange -from sglang.srt.layers.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h -from sglang.srt.layers.attention.fla.chunk_fwd import chunk_gated_delta_rule_fwd_intra -from sglang.srt.layers.attention.fla.chunk_o import chunk_fwd_o -from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum -from sglang.srt.layers.attention.fla.index import ( +from sglang.kernels.ops.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from sglang.kernels.ops.attention.fla.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from sglang.kernels.ops.attention.fla.chunk_o import chunk_fwd_o +from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum +from sglang.kernels.ops.attention.fla.index import ( prepare_chunk_indices, ) -from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd -from sglang.srt.layers.attention.fla.utils import ( +from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd +from sglang.kernels.ops.attention.fla.utils import ( SUPPRESS_LEVEL, autocast_custom_fwd, input_guard, diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py similarity index 98% rename from python/sglang/srt/layers/attention/fla/chunk_delta_h.py rename to python/sglang/kernels/ops/attention/fla/chunk_delta_h.py index fec725d01..7626358e3 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_delta_h.py @@ -9,12 +9,12 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import ( +from sglang.kernels.ops.attention.fla.index import ( prepare_chunk_indices, prepare_chunk_offsets, ) -from sglang.srt.layers.attention.fla.op import exp, safe_exp -from sglang.srt.layers.attention.fla.utils import ( +from sglang.kernels.ops.attention.fla.op import exp, safe_exp +from sglang.kernels.ops.attention.fla.utils import ( autotune_cache_kwargs, is_nvidia_hopper, ) diff --git a/python/sglang/srt/layers/attention/fla/chunk_fwd.py b/python/sglang/kernels/ops/attention/fla/chunk_fwd.py similarity index 98% rename from python/sglang/srt/layers/attention/fla/chunk_fwd.py rename to python/sglang/kernels/ops/attention/fla/chunk_fwd.py index 432a274cd..828ddc4c9 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_fwd.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_fwd.py @@ -5,13 +5,13 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.op import safe_exp -from sglang.srt.layers.attention.fla.utils import ( +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.op import safe_exp +from sglang.kernels.ops.attention.fla.utils import ( autotune_cache_kwargs, is_tf32_supported, ) -from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd +from sglang.kernels.ops.attention.fla.wy_fast import recompute_w_u_fwd # TF32 for the block-merge dot products (16x16 matmuls) is safe and ~2x faster on SM90. # The numerically sensitive forward-substitution uses scalar ops, not tl.dot. diff --git a/python/sglang/srt/layers/attention/fla/chunk_intra.py b/python/sglang/kernels/ops/attention/fla/chunk_intra.py similarity index 99% rename from python/sglang/srt/layers/attention/fla/chunk_intra.py rename to python/sglang/kernels/ops/attention/fla/chunk_intra.py index a8d35629e..44a067d8c 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_intra.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_intra.py @@ -5,14 +5,14 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.chunk_intra_token_parallel import ( +from sglang.kernels.ops.attention.fla.chunk_intra_token_parallel import ( chunk_kda_fwd_intra_token_parallel, ) -from sglang.srt.layers.attention.fla.index import ( +from sglang.kernels.ops.attention.fla.index import ( prepare_chunk_indices, ) -from sglang.srt.layers.attention.fla.op import exp, exp2, gather -from sglang.srt.layers.attention.fla.utils import ( +from sglang.kernels.ops.attention.fla.op import exp, exp2, gather +from sglang.kernels.ops.attention.fla.utils import ( autotune_cache_kwargs, is_gather_supported, is_tf32_supported, @@ -1036,7 +1036,7 @@ def chunk_kda_fwd_intra( FUSE_DIAGONAL=fuse_diagonal, ) - from sglang.srt.layers.attention.fla.kda import ( + from sglang.kernels.ops.attention.fla.kda import ( recompute_w_u_fwd as kda_recompute_w_u_fwd, ) diff --git a/python/sglang/srt/layers/attention/fla/chunk_intra_token_parallel.py b/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py similarity index 97% rename from python/sglang/srt/layers/attention/fla/chunk_intra_token_parallel.py rename to python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py index ec8bc848c..7481b9923 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_intra_token_parallel.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py @@ -6,8 +6,8 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.op import exp2 -from sglang.srt.layers.attention.fla.utils import autotune_cache_kwargs +from sglang.kernels.ops.attention.fla.op import exp2 +from sglang.kernels.ops.attention.fla.utils import autotune_cache_kwargs @triton.heuristics( diff --git a/python/sglang/srt/layers/attention/fla/chunk_o.py b/python/sglang/kernels/ops/attention/fla/chunk_o.py similarity index 95% rename from python/sglang/srt/layers/attention/fla/chunk_o.py rename to python/sglang/kernels/ops/attention/fla/chunk_o.py index bac5e93a5..c2c04312a 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_o.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_o.py @@ -8,9 +8,9 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.op import exp, safe_exp -from sglang.srt.layers.attention.fla.utils import check_shared_mem, is_nvidia_hopper +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.op import exp, safe_exp +from sglang.kernels.ops.attention.fla.utils import check_shared_mem, is_nvidia_hopper BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8] diff --git a/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py b/python/sglang/kernels/ops/attention/fla/chunk_scaled_dot_kkt.py similarity index 96% rename from python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py rename to python/sglang/kernels/ops/attention/fla/chunk_scaled_dot_kkt.py index e1059ac29..f401ac6f5 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py +++ b/python/sglang/kernels/ops/attention/fla/chunk_scaled_dot_kkt.py @@ -8,8 +8,8 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.op import safe_exp +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.op import safe_exp # @triton.autotune( diff --git a/python/sglang/srt/layers/attention/fla/cumsum.py b/python/sglang/kernels/ops/attention/fla/cumsum.py similarity index 98% rename from python/sglang/srt/layers/attention/fla/cumsum.py rename to python/sglang/kernels/ops/attention/fla/cumsum.py index 911b78e00..775a8f800 100644 --- a/python/sglang/srt/layers/attention/fla/cumsum.py +++ b/python/sglang/kernels/ops/attention/fla/cumsum.py @@ -8,8 +8,8 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.utils import check_shared_mem, input_guard +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.utils import check_shared_mem, input_guard BS_LIST = [32, 64] if check_shared_mem() else [16, 32] diff --git a/python/sglang/srt/layers/attention/fla/fused_gdn_gating.py b/python/sglang/kernels/ops/attention/fla/fused_gdn_gating.py similarity index 100% rename from python/sglang/srt/layers/attention/fla/fused_gdn_gating.py rename to python/sglang/kernels/ops/attention/fla/fused_gdn_gating.py diff --git a/python/sglang/srt/layers/attention/fla/fused_norm_gate.py b/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py similarity index 100% rename from python/sglang/srt/layers/attention/fla/fused_norm_gate.py rename to python/sglang/kernels/ops/attention/fla/fused_norm_gate.py diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py similarity index 99% rename from python/sglang/srt/layers/attention/fla/fused_recurrent.py rename to python/sglang/kernels/ops/attention/fla/fused_recurrent.py index 67c58592f..e2880355e 100644 --- a/python/sglang/srt/layers/attention/fla/fused_recurrent.py +++ b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py @@ -8,8 +8,8 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.op import exp -from sglang.srt.layers.attention.fla.utils import input_guard +from sglang.kernels.ops.attention.fla.op import exp +from sglang.kernels.ops.attention.fla.utils import input_guard @triton.jit(do_not_specialize=["T"]) diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent_linear_replayssm.py b/python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py similarity index 100% rename from python/sglang/srt/layers/attention/fla/fused_recurrent_linear_replayssm.py rename to python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py similarity index 100% rename from python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py rename to python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py diff --git a/python/sglang/srt/layers/attention/fla/index.py b/python/sglang/kernels/ops/attention/fla/index.py similarity index 93% rename from python/sglang/srt/layers/attention/fla/index.py rename to python/sglang/kernels/ops/attention/fla/index.py index 31b2e524e..1ece69150 100644 --- a/python/sglang/srt/layers/attention/fla/index.py +++ b/python/sglang/kernels/ops/attention/fla/index.py @@ -5,7 +5,7 @@ import torch import triton -from sglang.srt.layers.attention.fla.utils import tensor_cache +from sglang.kernels.ops.attention.fla.utils import tensor_cache @tensor_cache diff --git a/python/sglang/srt/layers/attention/fla/kda.py b/python/sglang/kernels/ops/attention/fla/kda.py similarity index 98% rename from python/sglang/srt/layers/attention/fla/kda.py rename to python/sglang/kernels/ops/attention/fla/kda.py index 579ae1a03..32e1357f4 100644 --- a/python/sglang/srt/layers/attention/fla/kda.py +++ b/python/sglang/kernels/ops/attention/fla/kda.py @@ -12,19 +12,19 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h -from sglang.srt.layers.attention.fla.chunk_intra import chunk_kda_fwd_intra -from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum -from sglang.srt.layers.attention.fla.fused_norm_gate import layer_norm_gated_fwd -from sglang.srt.layers.attention.fla.fused_recurrent import ( +from sglang.kernels.ops.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from sglang.kernels.ops.attention.fla.chunk_intra import chunk_kda_fwd_intra +from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum +from sglang.kernels.ops.attention.fla.fused_norm_gate import layer_norm_gated_fwd +from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_fwd_kernel, ) -from sglang.srt.layers.attention.fla.index import ( +from sglang.kernels.ops.attention.fla.index import ( prepare_chunk_indices, ) -from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd -from sglang.srt.layers.attention.fla.op import exp, log -from sglang.srt.layers.attention.fla.utils import ( +from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd +from sglang.kernels.ops.attention.fla.op import exp, log +from sglang.kernels.ops.attention.fla.utils import ( check_shared_mem, is_intel, ) diff --git a/python/sglang/srt/layers/attention/fla/l2norm.py b/python/sglang/kernels/ops/attention/fla/l2norm.py similarity index 98% rename from python/sglang/srt/layers/attention/fla/l2norm.py rename to python/sglang/kernels/ops/attention/fla/l2norm.py index d6b6ae7f7..dd1605cd3 100644 --- a/python/sglang/srt/layers/attention/fla/l2norm.py +++ b/python/sglang/kernels/ops/attention/fla/l2norm.py @@ -9,7 +9,7 @@ import torch.nn as nn import triton import triton.language as tl -from sglang.srt.layers.attention.fla.utils import input_guard +from sglang.kernels.ops.attention.fla.utils import input_guard BT_LIST = [8, 16, 32, 64, 128] diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/kernels/ops/attention/fla/layernorm_gated.py similarity index 100% rename from python/sglang/srt/layers/attention/fla/layernorm_gated.py rename to python/sglang/kernels/ops/attention/fla/layernorm_gated.py diff --git a/python/sglang/srt/layers/attention/fla/op.py b/python/sglang/kernels/ops/attention/fla/op.py similarity index 95% rename from python/sglang/srt/layers/attention/fla/op.py rename to python/sglang/kernels/ops/attention/fla/op.py index 9b3191075..54cc9aaca 100644 --- a/python/sglang/srt/layers/attention/fla/op.py +++ b/python/sglang/kernels/ops/attention/fla/op.py @@ -8,7 +8,7 @@ import triton import triton.language as tl import triton.language.extra.libdevice as tldevice -from sglang.srt.layers.attention.fla.utils import is_gather_supported +from sglang.kernels.ops.attention.fla.utils import is_gather_supported if os.environ.get("FLA_USE_FAST_OPS", "0") == "1": exp = tldevice.fast_expf diff --git a/python/sglang/srt/layers/attention/fla/solve_tril.py b/python/sglang/kernels/ops/attention/fla/solve_tril.py similarity index 99% rename from python/sglang/srt/layers/attention/fla/solve_tril.py rename to python/sglang/kernels/ops/attention/fla/solve_tril.py index 4aedad342..0640a4019 100644 --- a/python/sglang/srt/layers/attention/fla/solve_tril.py +++ b/python/sglang/kernels/ops/attention/fla/solve_tril.py @@ -8,8 +8,8 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.utils import input_guard +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.utils import input_guard # @triton.autotune( diff --git a/python/sglang/srt/layers/attention/fla/utils.py b/python/sglang/kernels/ops/attention/fla/utils.py similarity index 100% rename from python/sglang/srt/layers/attention/fla/utils.py rename to python/sglang/kernels/ops/attention/fla/utils.py diff --git a/python/sglang/srt/layers/attention/fla/wy_fast.py b/python/sglang/kernels/ops/attention/fla/wy_fast.py similarity index 98% rename from python/sglang/srt/layers/attention/fla/wy_fast.py rename to python/sglang/kernels/ops/attention/fla/wy_fast.py index 980a475cc..9956b6a92 100644 --- a/python/sglang/srt/layers/attention/fla/wy_fast.py +++ b/python/sglang/kernels/ops/attention/fla/wy_fast.py @@ -8,7 +8,7 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices # @triton.autotune( diff --git a/python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py index 92399d7eb..02803ae40 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py +++ b/python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py @@ -163,7 +163,7 @@ def chunk_kda_cutedsl( # injected through the cutedsl KKT/Aqk MMAs as an identity-right-operand pass: # with kL'=M (M in the first 64 K-slots) and kR'=onehot(chunk-pos), the MMA # kL'@kR'.T == M, so kkt_inv_uw/kernel_o see the correct matrix without overflow. - from sglang.srt.layers.attention.fla.kda import chunk_kda_scaled_dot_kkt_fwd + from sglang.kernels.ops.attention.fla.kda import chunk_kda_scaled_dot_kkt_fwd ones_beta = q.new_ones(1, T, Hv, dtype=torch.float32) M_kk, M_qk = chunk_kda_scaled_dot_kkt_fwd( diff --git a/python/sglang/kernels/ops/mamba/__init__.py b/python/sglang/kernels/ops/mamba/__init__.py index eff57abe1..5b3252df2 100644 --- a/python/sglang/kernels/ops/mamba/__init__.py +++ b/python/sglang/kernels/ops/mamba/__init__.py @@ -84,3 +84,20 @@ def causal_conv1d_update( __all__ = ["causal_conv1d_fwd", "causal_conv1d_update"] + + +# Vendored mamba_ssm-derived kernels relocated in Phase 2.5 (RFC #29630). +for _mod, _fn in [ + ("triton_ops.ssd_combined", "mamba_chunk_scan_combined"), + ("triton_ops.mamba_ssm", "selective_state_update"), + ("causal_conv1d_triton", "causal_conv1d_fn"), + ("mamba_state_scatter_triton", "fused_mamba_state_scatter_with_mask"), +]: + register_kernel( + KernelSpec( + op=f"mamba.{_fn}", + backend=KernelBackend.TRITON, + target=f"sglang.kernels.ops.mamba.{_mod}:{_fn}", + ) + ) +del _mod, _fn diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py rename to python/sglang/kernels/ops/mamba/causal_conv1d_triton.py diff --git a/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py b/python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py rename to python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/__init__.py b/python/sglang/kernels/ops/mamba/triton_ops/__init__.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/__init__.py rename to python/sglang/kernels/ops/mamba/triton_ops/__init__.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py b/python/sglang/kernels/ops/mamba/triton_ops/layernorm_gated.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/layernorm_gated.py rename to python/sglang/kernels/ops/mamba/triton_ops/layernorm_gated.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py b/python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py rename to python/sglang/kernels/ops/mamba/triton_ops/mamba_ssm.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_bmm.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_bmm.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/ssd_bmm.py rename to python/sglang/kernels/ops/mamba/triton_ops/ssd_bmm.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_scan.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_scan.py rename to python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_scan.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py rename to python/sglang/kernels/ops/mamba/triton_ops/ssd_chunk_state.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_combined.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/ssd_combined.py rename to python/sglang/kernels/ops/mamba/triton_ops/ssd_combined.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py b/python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py similarity index 100% rename from python/sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py rename to python/sglang/kernels/ops/mamba/triton_ops/ssd_state_passing.py diff --git a/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py b/python/sglang/kernels/ops/mamba/triton_ops/ssu_dispatch.py similarity index 99% rename from python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py rename to python/sglang/kernels/ops/mamba/triton_ops/ssu_dispatch.py index df61896be..c8d28341f 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/ssu_dispatch.py +++ b/python/sglang/kernels/ops/mamba/triton_ops/ssu_dispatch.py @@ -51,7 +51,7 @@ class TritonSSUBackend(MambaSSUBackend): enable_stochastic_rounding: bool = False, cache_philox_rounds: int = 0, ) -> None: - from sglang.srt.layers.attention.mamba.ops.mamba_ssm import ( + from sglang.kernels.ops.mamba.triton_ops.mamba_ssm import ( selective_state_update, ) diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py index bdb52b1c3..009fff16c 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py @@ -4,12 +4,12 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import ( +from sglang.kernels.ops.attention.fla.index import ( prepare_chunk_indices, prepare_chunk_offsets, ) -from sglang.srt.layers.attention.fla.op import exp, make_tensor_descriptor, safe_exp -from sglang.srt.layers.attention.fla.utils import ( +from sglang.kernels.ops.attention.fla.op import exp, make_tensor_descriptor, safe_exp +from sglang.kernels.ops.attention.fla.utils import ( autotune_cache_kwargs, ) diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py index c2299a2f0..ead49963c 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py @@ -2,12 +2,12 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.op import safe_exp -from sglang.srt.layers.attention.fla.utils import ( +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.op import safe_exp +from sglang.kernels.ops.attention.fla.utils import ( autotune_cache_kwargs, ) -from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd +from sglang.kernels.ops.attention.fla.wy_fast import recompute_w_u_fwd _MERGE_DOT_PRECISION = tl.constexpr("ieee") diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py index 83a0fce3f..ab30b9b5e 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py @@ -3,7 +3,7 @@ from typing import Optional import torch import triton -from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update_kernel, ) diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 9060ecc68..49821659b 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -43,7 +43,9 @@ def create_flashinfer_backend(runner): import torch if not runner.use_mla_backend: - from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend + from sglang.srt.layers.attention.flashinfer_backend import ( + FlashInferAttnBackend, + ) # Init streams if runner.server_args.speculative_algorithm == "EAGLE": @@ -219,7 +221,9 @@ def create_flashattention_v3_backend(runner): @register_attention_backend("fa4") def create_flashattention_v4_backend(runner): - from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) return FlashAttentionBackend(runner, fa_impl_ver=4) @@ -279,7 +283,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac ) if cfg := mambaish_config(runner.model_config): - from sglang.srt.layers.attention.fla.utils import check_environments + from sglang.kernels.ops.attention.fla.utils import check_environments from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend from sglang.srt.layers.attention.linear.lightning_backend import ( LightningAttentionBackend, diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 58cc14837..562713643 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -3,19 +3,19 @@ from typing import Optional, Union import torch +from sglang.kernels.ops.mamba.causal_conv1d_triton import PAD_SLOT_ID +from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( + fused_conv_window_scatter_with_mask, + fused_mamba_state_scatter_with_mask, + track_mamba_states_if_needed, +) from sglang.srt.configs.hybrid_arch import mamba2_config from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID from sglang.srt.layers.attention.mamba.mamba import MambaMixer2 from sglang.srt.layers.attention.mamba.mamba2_metadata import ( ForwardMetadata, Mamba2Metadata, ) -from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( - fused_conv_window_scatter_with_mask, - fused_mamba_state_scatter_with_mask, - track_mamba_states_if_needed, -) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 1d73b9f65..ac5525ec1 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -2,8 +2,12 @@ from typing import Optional, Tuple, Union import torch +from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_fn, + causal_conv1d_update, +) from sglang.srt.configs.hybrid_arch import hybrid_gdn_config -from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel from sglang.srt.layers.attention.linear.utils import ( @@ -11,10 +15,6 @@ from sglang.srt.layers.attention.linear.utils import ( get_linear_attn_decode_backend, get_linear_attn_prefill_backend, ) -from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( - causal_conv1d_fn, - causal_conv1d_update, -) from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import MambaPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -23,7 +23,7 @@ from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_npu from sglang.srt.utils.common import rank0_log if not is_cpu(): - from sglang.srt.layers.attention.fla.chunk_delta_h import ( + from sglang.kernels.ops.attention.fla.chunk_delta_h import ( CHUNK_SIZE as FLA_CHUNK_SIZE, ) diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 8a4784d8b..7cc6b49b9 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -2,6 +2,10 @@ from typing import Optional, Tuple, Union import torch +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_fn, + causal_conv1d_update, +) from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel from sglang.srt.layers.attention.linear.utils import ( @@ -9,10 +13,6 @@ from sglang.srt.layers.attention.linear.utils import ( get_linear_attn_decode_backend, get_linear_attn_prefill_backend, ) -from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( - causal_conv1d_fn, - causal_conv1d_update, -) from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.utils import is_cpu, is_cuda, is_npu from sglang.srt.utils.common import rank0_log diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py index 80e271207..e29be0692 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py @@ -65,11 +65,11 @@ class CuteDSLGDNKernel(LinearAttnKernelBase): raise RuntimeError( f"CuTe DSL GDN prefill requires head_k_dim=128, got {head_k_dim}." ) + from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd from sglang.kernels.ops.attention.linear.gdn_blackwell import ( chunk_gated_delta_rule_cutedsl, prepare_metadata_cutedsl, ) - from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd self._extend_fn = chunk_gated_delta_rule_cutedsl self._prepare_meta_fn = prepare_metadata_cutedsl diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py index 9ce121f75..7b7a89801 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py @@ -233,7 +233,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase): query_start_loc: torch.Tensor, **kwargs, ) -> tuple: - from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd + from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd total_seq_len = q.shape[1] num_v_heads = v.shape[2] diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py index abdfa5e37..c6942226b 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py @@ -6,14 +6,14 @@ from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( from sglang.srt.utils import is_cpu, is_npu, is_xpu if not is_cpu(): - from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule - from sglang.srt.layers.attention.fla.fused_recurrent import ( + from sglang.kernels.ops.attention.fla.chunk import chunk_gated_delta_rule + from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_packed_decode, ) - from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import ( + from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import ( fused_recurrent_gdn_replayssm_decode, ) - from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( + from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py b/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py index ed3df69d1..ba0dd85b3 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py @@ -49,10 +49,10 @@ class CuteDSLKDAKernel(LinearAttnKernelBase): raise RuntimeError( f"CuTe DSL KDA prefill requires head_k_dim=128, got {head_k_dim}." ) + from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd from sglang.kernels.ops.attention.linear.kda_blackwell import ( chunk_kda_cutedsl, ) - from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd self._extend_fn = chunk_kda_cutedsl self._l2norm_fn = l2norm_fwd diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py index 2720dd228..cce92fe8e 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py @@ -47,7 +47,7 @@ def _triton_fallback( -- otherwise the fallback silently skips activation. chunk_kda updates the ssm state in-place via cache_indices and returns only the output tensor. """ - from sglang.srt.layers.attention.fla.kda import chunk_kda + from sglang.kernels.ops.attention.fla.kda import chunk_kda return chunk_kda( q=q, diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py index e8f0fc405..bf7e43714 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py @@ -8,16 +8,16 @@ from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( from sglang.srt.utils import is_cpu, is_npu if not is_cpu(): - from sglang.srt.layers.attention.fla.fused_recurrent import ( + from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_kda_packed_decode, ) - from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import ( + from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import ( fused_recurrent_linear_replayssm_decode, ) - from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( + from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) - from sglang.srt.layers.attention.fla.kda import chunk_kda + from sglang.kernels.ops.attention.fla.kda import chunk_kda class TritonKDAKernel(LinearAttnKernelBase): diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d.py b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py index ba94f3105..4601f1df5 100644 --- a/python/sglang/srt/layers/attention/mamba/causal_conv1d.py +++ b/python/sglang/srt/layers/attention/mamba/causal_conv1d.py @@ -9,9 +9,15 @@ from typing import Optional import torch -from .causal_conv1d_triton import PAD_SLOT_ID -from .causal_conv1d_triton import causal_conv1d_fn as _causal_conv1d_fn_triton -from .causal_conv1d_triton import causal_conv1d_update as _causal_conv1d_update_triton +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + PAD_SLOT_ID, +) +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_fn as _causal_conv1d_fn_triton, +) +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_update as _causal_conv1d_update_triton, +) try: from sgl_kernel import causal_conv1d_fwd diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index 80ec14481..632766518 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/python/sglang/srt/layers/attention/mamba/mamba.py @@ -4,6 +4,10 @@ from typing import Callable, List, Optional, Tuple import torch import torch.nn as nn +from sglang.kernels.ops.mamba.triton_ops import ( + mamba_chunk_scan_combined, + selective_state_update, +) from sglang.srt.configs.mamba_utils import ( Mamba2CacheParams, extra_groups_for_head_shards, @@ -13,10 +17,6 @@ from sglang.srt.distributed import ( ) from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata from sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated import Mixer2RMSNormGated -from sglang.srt.layers.attention.mamba.ops import ( - mamba_chunk_scan_combined, - selective_state_update, -) from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) @@ -42,16 +42,16 @@ from sglang.srt.utils import ( ) if is_cuda(): + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_fn as causal_conv1d_fn_triton, + ) + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_update as causal_conv1d_update_triton, + ) from sglang.srt.layers.attention.mamba.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, ) - from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( - causal_conv1d_fn as causal_conv1d_fn_triton, - ) - from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( - causal_conv1d_update as causal_conv1d_update_triton, - ) elif is_npu(): from sgl_kernel_npu.mamba.causal_conv1d import ( causal_conv1d_fn_npu as causal_conv1d_fn, @@ -63,16 +63,16 @@ elif is_xpu(): # XPU has no native causal_conv1d kernel yet; use the portable Triton # implementation for both the "native" and the "_triton" entry points so # `causal_conv1d_fn` / `causal_conv1d_fn_triton` are always bound on XPU. - from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( causal_conv1d_fn as causal_conv1d_fn, ) - from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( causal_conv1d_fn as causal_conv1d_fn_triton, ) - from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( causal_conv1d_update as causal_conv1d_update, ) - from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( + from sglang.kernels.ops.mamba.causal_conv1d_triton import ( causal_conv1d_update as causal_conv1d_update_triton, ) diff --git a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py index f42f9c801..7ceb3e161 100644 --- a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py +++ b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py @@ -2,11 +2,11 @@ from typing import Union import torch +from sglang.kernels.ops.attention.fla.layernorm_gated import rms_norm_gated from sglang.srt.distributed.communication_op import ( tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce, ) -from sglang.srt.layers.attention.fla.layernorm_gated import rms_norm_gated from sglang.srt.layers.dp_attention import ( attn_tp_all_reduce, is_dp_attention_enabled, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4d5078e61..7eeeeb5b9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -38,6 +38,9 @@ import torch.distributed from torch.cuda import Stream as CudaStream from torch.distributed import barrier +from sglang.kernels.ops.mamba.triton_ops import ( + initialize_mamba_selective_state_update_backend, +) from sglang.srt.configs.model_config import ModelConfig, ModelImpl, is_minimax_sparse from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix @@ -69,9 +72,6 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder -from sglang.srt.layers.attention.mamba.ops import ( - initialize_mamba_selective_state_update_backend, -) from sglang.srt.layers.dp_attention import compute_dp_attention_world_info from sglang.srt.layers.moe import initialize_moe_config from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 2877a155e..27b265814 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -9,6 +9,8 @@ import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig +from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated +from sglang.kernels.ops.attention.fla.layernorm_gated import layernorm_fn from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.distributed import ( get_pp_group, @@ -17,8 +19,6 @@ from sglang.srt.distributed import ( from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul -from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated -from sglang.srt.layers.attention.fla.layernorm_gated import layernorm_fn from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, diff --git a/python/sglang/srt/models/jet_nemotron.py b/python/sglang/srt/models/jet_nemotron.py index fec8c1fb6..5ebf473c2 100644 --- a/python/sglang/srt/models/jet_nemotron.py +++ b/python/sglang/srt/models/jet_nemotron.py @@ -5,11 +5,11 @@ import einops import torch import torch.nn as nn -from sglang.srt.configs.jet_nemotron import JetBlockConfig, JetNemotronConfig -from sglang.srt.layers.attention.fla.fused_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_update, ) -from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated +from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated +from sglang.srt.configs.jet_nemotron import JetBlockConfig, JetNemotronConfig from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( HybridLinearAttnBackend, MambaAttnBackendBase, diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 77c9e2384..9d51b3f77 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -8,6 +8,7 @@ from typing import Optional import torch from torch import nn +from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.distributed import ( divide, @@ -15,7 +16,6 @@ from sglang.srt.distributed import ( tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder -from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelBatchedLinear, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index e5d386f90..2541449fe 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -25,6 +25,9 @@ import triton from sglang.jit_kernel.triton.gdn_fused_proj import ( fused_qkvzba_split_reshape_cat_contiguous, ) + +# Layers - Attention +from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated from sglang.kernels.ops.layernorm.elementwise import fused_sigmoid_mul # Configs @@ -38,9 +41,6 @@ from sglang.srt.configs.qwen3_5 import ( from sglang.srt.distributed import get_pp_group from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation - -# Layers - Attention -from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.dp_attention import ( diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index e768df860..42e38db81 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -7,12 +7,12 @@ import triton from torch import nn from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkvzba_split_reshape_cat +from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated +from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated from sglang.srt.configs.qwen3_next import Qwen3NextConfig from sglang.srt.distributed import get_pp_group from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation -from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated -from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.dp_attention import ( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 17cb1b725..2c8525f71 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -32,6 +32,7 @@ from functools import cached_property from typing import Any, Callable, Dict, List, Literal, Optional, Union from sglang.jit_kernel.kv_canary.consts import RealKvHashMode +from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.arg_groups.arg_utils import A, Arg, add_cli_args_from_dataclass from sglang.srt.arg_groups.argparse_actions import ( DeprecatedAction, @@ -47,7 +48,6 @@ from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ) from sglang.srt.environ import envs from sglang.srt.function_call.function_call_parser import FunctionCallParser -from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.model_executor.cuda_graph_config import ( ALLOWED_BACKENDS_PER_PHASE, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index 3ae3e7a6d..6a98a1fcb 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py @@ -389,7 +389,7 @@ class MockMamba2ModelRunner(ModelRunner): # `MambaMixer2.forward_decode` requires. In production the # scheduler calls this during initialization; the fixture must # mirror that or DECODE crashes with a missing-backend error. - from sglang.srt.layers.attention.mamba.ops import ( + from sglang.kernels.ops.mamba.triton_ops import ( initialize_mamba_selective_state_update_backend, ) diff --git a/test/registered/attention/test_chunk_gated_delta_rule.py b/test/registered/attention/test_chunk_gated_delta_rule.py index 680c2b50e..a3d830688 100644 --- a/test/registered/attention/test_chunk_gated_delta_rule.py +++ b/test/registered/attention/test_chunk_gated_delta_rule.py @@ -2,8 +2,8 @@ import unittest import torch -from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule -from sglang.srt.layers.attention.fla.fused_recurrent import ( +from sglang.kernels.ops.attention.fla.chunk import chunk_gated_delta_rule +from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule, ) from sglang.srt.utils import get_device diff --git a/test/registered/attention/test_gdn_noncontiguous_stride.py b/test/registered/attention/test_gdn_noncontiguous_stride.py index f008804d6..f7debb162 100644 --- a/test/registered/attention/test_gdn_noncontiguous_stride.py +++ b/test/registered/attention/test_gdn_noncontiguous_stride.py @@ -8,8 +8,8 @@ import unittest import torch -from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating -from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci diff --git a/test/registered/attention/test_gdn_prefill_cutedsl.py b/test/registered/attention/test_gdn_prefill_cutedsl.py index 975b09e29..1d1eb1361 100644 --- a/test/registered/attention/test_gdn_prefill_cutedsl.py +++ b/test/registered/attention/test_gdn_prefill_cutedsl.py @@ -23,17 +23,17 @@ if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 1 allow_module_level=True, ) +from sglang.kernels.ops.attention.fla.fused_recurrent import ( # noqa: E402 + fused_recurrent_gated_delta_rule, +) +from sglang.kernels.ops.attention.fla.index import ( # noqa: E402 + prepare_chunk_indices, + prepare_chunk_offsets, +) from sglang.kernels.ops.attention.linear.gdn_blackwell import ( # noqa: E402 chunk_gated_delta_rule_cutedsl, prepare_metadata_cutedsl, ) -from sglang.srt.layers.attention.fla.fused_recurrent import ( # noqa: E402 - fused_recurrent_gated_delta_rule, -) -from sglang.srt.layers.attention.fla.index import ( # noqa: E402 - prepare_chunk_indices, - prepare_chunk_offsets, -) @pytest.mark.parametrize("num_seqs", [1, 5, 257]) diff --git a/test/registered/attention/test_kda_kernels.py b/test/registered/attention/test_kda_kernels.py index 4e273d6a4..41e1858d9 100644 --- a/test/registered/attention/test_kda_kernels.py +++ b/test/registered/attention/test_kda_kernels.py @@ -2,15 +2,15 @@ import unittest import torch -from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum -from sglang.srt.layers.attention.fla.fused_recurrent import ( +from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum +from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_kda_packed_decode, ) -from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices -from sglang.srt.layers.attention.fla.kda import ( +from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices +from sglang.kernels.ops.attention.fla.kda import ( fused_recurrent_kda, kda_gate_chunk_cumsum, ) diff --git a/test/registered/attention/test_kda_prefill_cutedsl.py b/test/registered/attention/test_kda_prefill_cutedsl.py index 83fa7eae6..2ca80ec19 100644 --- a/test/registered/attention/test_kda_prefill_cutedsl.py +++ b/test/registered/attention/test_kda_prefill_cutedsl.py @@ -25,15 +25,15 @@ if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 1 allow_module_level=True, ) +from sglang.kernels.ops.attention.fla.index import ( # noqa: E402 + prepare_chunk_indices, + prepare_chunk_offsets, +) +from sglang.kernels.ops.attention.fla.kda import fused_recurrent_kda # noqa: E402 from sglang.kernels.ops.attention.linear.kda_blackwell import ( # noqa: E402 chunk_kda_cutedsl, prepare_metadata, ) -from sglang.srt.layers.attention.fla.index import ( # noqa: E402 - prepare_chunk_indices, - prepare_chunk_offsets, -) -from sglang.srt.layers.attention.fla.kda import fused_recurrent_kda # noqa: E402 def _l2norm(x: torch.Tensor) -> torch.Tensor: diff --git a/test/registered/attention/test_kda_prefill_flashkda.py b/test/registered/attention/test_kda_prefill_flashkda.py index 954d1dadc..5bf57e85d 100644 --- a/test/registered/attention/test_kda_prefill_flashkda.py +++ b/test/registered/attention/test_kda_prefill_flashkda.py @@ -39,7 +39,7 @@ except ImportError: allow_module_level=True, ) -from sglang.srt.layers.attention.fla.kda import chunk_kda # noqa: E402 +from sglang.kernels.ops.attention.fla.kda import chunk_kda # noqa: E402 from sglang.srt.layers.attention.linear.kernels.kda_flashkda import ( # noqa: E402 FlashKDAKernel, ) diff --git a/test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py b/test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py index 754bc95c9..b5a278a75 100644 --- a/test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py +++ b/test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py @@ -98,11 +98,11 @@ class TestLinearReplaySSMDecode(CustomTestCase): L_SWEEP = (1, 4, 8, 16) def _run_one(self, cfg, L, dtype, force_flush_steps=(), is_kda=False): - from sglang.srt.layers.attention.fla.fused_recurrent import ( + from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_packed_decode, fused_recurrent_kda_packed_decode, ) - from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import ( + from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import ( fused_recurrent_linear_replayssm_decode, ) diff --git a/test/registered/jit/test_cutedsl_gdn.py b/test/registered/jit/test_cutedsl_gdn.py index d74b89ac8..d6f48a052 100644 --- a/test/registered/jit/test_cutedsl_gdn.py +++ b/test/registered/jit/test_cutedsl_gdn.py @@ -21,7 +21,7 @@ except ImportError: cutedsl_gdn = None try: - from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( + from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/test/registered/jit/test_fused_verify_triton_gdn.py b/test/registered/jit/test_fused_verify_triton_gdn.py index c417fa962..abca9a705 100644 --- a/test/registered/jit/test_fused_verify_triton_gdn.py +++ b/test/registered/jit/test_fused_verify_triton_gdn.py @@ -14,11 +14,11 @@ import torch from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci try: - from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating - from sglang.srt.layers.attention.fla.fused_recurrent import ( + from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating + from sglang.kernels.ops.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_update, ) - from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( + from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/test/registered/layers/mamba/conftest.py b/test/registered/layers/mamba/conftest.py index 606a5ee1a..d8d5df6d6 100644 --- a/test/registered/layers/mamba/conftest.py +++ b/test/registered/layers/mamba/conftest.py @@ -1,7 +1,7 @@ import pytest -from sglang.srt.layers.attention.mamba.ops import ssu_dispatch -from sglang.srt.layers.attention.mamba.ops.ssu_dispatch import ( +from sglang.kernels.ops.mamba.triton_ops import ssu_dispatch +from sglang.kernels.ops.mamba.triton_ops.ssu_dispatch import ( initialize_mamba_selective_state_update_backend, ) from sglang.srt.server_args import ServerArgs diff --git a/test/registered/layers/mamba/test_causal_conv1d.py b/test/registered/layers/mamba/test_causal_conv1d.py index e3e8d2a9d..0ca1ca421 100644 --- a/test/registered/layers/mamba/test_causal_conv1d.py +++ b/test/registered/layers/mamba/test_causal_conv1d.py @@ -15,7 +15,7 @@ import torch import torch.nn.functional as F from einops import rearrange -from sglang.srt.layers.attention.mamba.causal_conv1d_triton import ( +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( PAD_SLOT_ID, causal_conv1d_fn, causal_conv1d_update, diff --git a/test/registered/layers/mamba/test_mamba_ssm.py b/test/registered/layers/mamba/test_mamba_ssm.py index 165262178..250c9ae53 100644 --- a/test/registered/layers/mamba/test_mamba_ssm.py +++ b/test/registered/layers/mamba/test_mamba_ssm.py @@ -13,8 +13,8 @@ import torch import torch.nn.functional as F from einops import rearrange, repeat -from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID -from sglang.srt.layers.attention.mamba.ops.mamba_ssm import selective_state_update +from sglang.kernels.ops.mamba.causal_conv1d_triton import PAD_SLOT_ID +from sglang.kernels.ops.mamba.triton_ops.mamba_ssm import selective_state_update from sglang.srt.utils import get_device, is_sm100_supported diff --git a/test/registered/layers/mamba/test_mamba_ssm_ssd.py b/test/registered/layers/mamba/test_mamba_ssm_ssd.py index 2e9d93a4d..9df125a86 100644 --- a/test/registered/layers/mamba/test_mamba_ssm_ssd.py +++ b/test/registered/layers/mamba/test_mamba_ssm_ssd.py @@ -14,8 +14,8 @@ import torch import torch.nn.functional as F from einops import rearrange, repeat +from sglang.kernels.ops.mamba.triton_ops import mamba_chunk_scan_combined from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata -from sglang.srt.layers.attention.mamba.ops import mamba_chunk_scan_combined from sglang.srt.utils import get_device from sglang.srt.utils.common import is_hip from sglang.utils import is_in_ci diff --git a/test/registered/layers/test_fla_layernorm_guard.py b/test/registered/layers/test_fla_layernorm_guard.py index d860d3b71..e2398abc0 100644 --- a/test/registered/layers/test_fla_layernorm_guard.py +++ b/test/registered/layers/test_fla_layernorm_guard.py @@ -8,10 +8,10 @@ import pytest import torch import torch.nn.functional as F -from sglang.srt.layers.attention.fla.layernorm_gated import ( +from sglang.kernels.ops.attention.fla.layernorm_gated import ( _layer_norm_fwd as layer_norm_fwd, ) -from sglang.srt.layers.attention.fla.layernorm_gated import ( +from sglang.kernels.ops.attention.fla.layernorm_gated import ( layernorm_fn, rms_norm_ref, ) diff --git a/test/registered/mem_cache/test_int8_checkpoint_store.py b/test/registered/mem_cache/test_int8_checkpoint_store.py index e5c227e72..af6b03b6b 100644 --- a/test/registered/mem_cache/test_int8_checkpoint_store.py +++ b/test/registered/mem_cache/test_int8_checkpoint_store.py @@ -126,7 +126,7 @@ class TestInt8CheckpointCodec(unittest.TestCase): class TestInt8CheckpointDecodeError(unittest.TestCase): def test_decode_error_within_bound(self): try: - from sglang.srt.layers.attention.fla.kda import fused_recurrent_kda + from sglang.kernels.ops.attention.fla.kda import fused_recurrent_kda except (ImportError, ModuleNotFoundError) as e: self.skipTest(f"fla kernels unavailable: {e}") diff --git a/test/registered/unit/layers/test_mamba_state_scatter_triton.py b/test/registered/unit/layers/test_mamba_state_scatter_triton.py index 3fb04ec53..aaf5ce55a 100644 --- a/test/registered/unit/layers/test_mamba_state_scatter_triton.py +++ b/test/registered/unit/layers/test_mamba_state_scatter_triton.py @@ -8,7 +8,7 @@ import unittest import torch try: - from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( + from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( fused_mamba_state_scatter_with_mask, ) diff --git a/test/registered/unit/mem_cache/test_mamba_unittest.py b/test/registered/unit/mem_cache/test_mamba_unittest.py index f6690fe43..01a830342 100755 --- a/test/registered/unit/mem_cache/test_mamba_unittest.py +++ b/test/registered/unit/mem_cache/test_mamba_unittest.py @@ -3,10 +3,10 @@ from array import array import torch +from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.environ import envs -from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 03906f7b4..9202da96a 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -19,9 +19,9 @@ from typing import Callable import torch +from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape from sglang.srt.environ import envs -from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index d0d3dff92..7f0c972a8 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -12,6 +12,7 @@ from unittest import mock import torch +from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape from sglang.srt.disaggregation.kv_events import ( BlockRemoved, @@ -19,7 +20,6 @@ from sglang.srt.disaggregation.kv_events import ( StorageMedium, ) from sglang.srt.environ import envs -from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator