Skip to content

[Linear Attn] Add CUSTOM enum and plugin extensibility for kernel backends #24937

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merrymercy merged 4 commits into from
May 12, 2026
Merged

[Linear Attn] Add CUSTOM enum and plugin extensibility for kernel backends #24937

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9b9a04d
[Linear Attn] Add CUSTOM enum and plugin extensibility for kernel bac…
minosfuture May 11, 2026
375b662
Merge remote-tracking branch 'origin/main' into linear_attn_backend
minosfuture May 12, 2026
31803b2
lint fix
minosfuture May 12, 2026
f465626
Merge branch 'main' into linear_attn_backend
minosfuture May 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 10 additions & 5 deletions python/sglang/srt/layers/attention/linear/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,11 @@ class LinearAttnKernelBackend(Enum):
TRITON = "triton"
CUTEDSL = "cutedsl"
FLASHINFER = "flashinfer"
CUSTOM = "custom"

@classmethod
def _missing_(cls, value):
return cls.CUSTOM

def is_triton(self):
return self == LinearAttnKernelBackend.TRITON
Expand All @@ -26,6 +31,9 @@ def is_cutedsl(self):
def is_flashinfer(self):
return self == LinearAttnKernelBackend.FLASHINFER

def is_custom(self):
return self == LinearAttnKernelBackend.CUSTOM


LINEAR_ATTN_DECODE_BACKEND: Optional[LinearAttnKernelBackend] = None
LINEAR_ATTN_PREFILL_BACKEND: Optional[LinearAttnKernelBackend] = None
Expand All @@ -41,11 +49,8 @@ def initialize_linear_attn_config(server_args: ServerArgs):

LINEAR_ATTN_DECODE_BACKEND = LinearAttnKernelBackend(decode)
LINEAR_ATTN_PREFILL_BACKEND = LinearAttnKernelBackend(prefill)
rank0_log(
f"Linear attention kernel backend: "
f"decode={LINEAR_ATTN_DECODE_BACKEND.value}, "
f"prefill={LINEAR_ATTN_PREFILL_BACKEND.value}"
)

rank0_log(f"Linear attention kernel backend: decode={decode}, prefill={prefill}")


def get_linear_attn_decode_backend() -> LinearAttnKernelBackend:
Expand Down
4 changes: 4 additions & 0 deletions python/sglang/srt/server_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,6 +305,10 @@ def add_rl_on_policy_target_choices(choices):
RL_ON_POLICY_TARGET_CHOICES.extend(choices)


def add_linear_attn_kernel_backend_choices(choices):
LINEAR_ATTN_KERNEL_BACKEND_CHOICES.extend(choices)


def _resolve_speculative_algorithm_alias(
speculative_algorithm: Optional[str],
speculative_draft_model_path: Optional[str],
Expand Down
Loading