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

[Kernel] Register merged diffusion agent kernels with KDA backend (#37385)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-09-01 22:51:10 +08:00
committed by GitHub
parent 4c7ff0d906
commit 5993f91f84
6 changed files with 111 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/fused_op.py
+5
View File
@@ -108,6 +108,7 @@ BACKEND_METHODS: Dict[KernelBackend, str] = {
KernelBackend.AOT: "forward_aot",
KernelBackend.CUTE_DSL: "forward_cute_dsl",
KernelBackend.FLYDSL: "forward_flydsl",
KernelBackend.KDA: "forward_kda",
KernelBackend.FLASHINFER: "forward_flashinfer",
KernelBackend.DEEPGEMM: "forward_deepgemm",
KernelBackend.AITER: "forward_aiter",
@@ -122,6 +123,7 @@ _METHOD_BACKEND_LABELS: Dict[str, str] = {
# must never trigger a surprise compilation in a serving process; force it
# explicitly when wanted. Per-op priority overrides this (see BaseFusedOp).
DEFAULT_PRIORITY: Tuple[KernelBackend, ...] = (
KernelBackend.KDA,
KernelBackend.AOT,
KernelBackend.JIT,
KernelBackend.FLASHINFER,
@@ -434,6 +436,9 @@ class BaseFusedOp(nn.Module, ABC):
def forward_cute_dsl(self, *args, **kwargs):
raise NotImplementedError(f"{self._op_label()}: no cute_dsl backend")
def forward_kda(self, *args, **kwargs):
raise NotImplementedError(f"{self._op_label()}: no KDA backend")
def forward_flashinfer(self, *args, **kwargs):
raise NotImplementedError(f"{self._op_label()}: no flashinfer backend")
python/sglang/kernels/ops/diffusion/__init__.py
+32 -11
View File
@@ -82,6 +82,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
_CUDA,
"Bit-exact RMSNorm + adaLN scale/shift.",
),
(
"diffusion.scale_residual_norm_scale_shift",
KernelBackend.KDA,
"norm.norm_scale_shift_jit:kda_scale_residual_norm_scale_shift",
_CUDA_SM100_PLUS,
"KDA B200 native CUDA residual + LayerNorm + scale/shift (#27392).",
),
(
"diffusion.scale_residual_norm_scale_shift",
KernelBackend.TRITON,
@@ -110,6 +117,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
_CUDA,
"Qwen residual LayerNorm/modulation + NVFP4 quantization.",
),
(
"diffusion.norm_scale_shift",
KernelBackend.KDA,
"norm.norm_scale_shift_jit:kda_norm_scale_shift",
_CUDA_SM100_PLUS,
"KDA B200 native CUDA LayerNorm + scale/shift (#27392).",
),
(
"diffusion.norm_scale_shift",
KernelBackend.CUTE_DSL,
@@ -168,10 +182,10 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
),
(
"diffusion.residual_gate_add",
KernelBackend.JIT,
KernelBackend.KDA,
"modulate.residual_gate_add_jit:residual_gate_add",
_CUDA,
"Fused residual + gate * update.",
"KDA native CUDA residual + gate * update (#29361).",
),
(
"diffusion.timestep_embedding",
@@ -201,19 +215,26 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
_CUDA,
"Fused in-place QK RMS-norm + RoPE.",
),
(
"diffusion.flux2_layernorm_modulate_fp8_quant",
KernelBackend.KDA,
"norm.layernorm_modulate_triton:fused_layernorm_modulate_fp8_quant_raw",
_CUDA,
"KDA-generated FLUX.2 LayerNorm + adaLN modulation + static FP8 quantization.",
),
(
"diffusion.flux2_qkv_epilogue",
KernelBackend.JIT,
KernelBackend.KDA,
"rope.flux2_qkv_epilogue_jit:try_fused_flux2_qkv_epilogue",
_CUDA,
"FLUX.2 QK RMS-norm + RoPE + joint QKV packing.",
"KDA-generated FLUX.2 QK RMS-norm + RoPE + joint QKV packing.",
),
(
"diffusion.flux2_token_cat_fp8",
KernelBackend.TRITON,
KernelBackend.KDA,
"layout.flux2_token_cat_fp8_triton:try_flux2_token_cat_fp8",
_CUDA,
"FLUX.2 single-block token concatenation + static FP8 quantization.",
"KDA-generated FLUX.2 token concatenation + static FP8 quantization.",
),
(
"diffusion.qwen_qkv_epilogue",
@@ -224,10 +245,10 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
),
(
"diffusion.ltx2_qknorm_split_rope",
KernelBackend.JIT,
KernelBackend.KDA,
"rope.ltx2_qknorm_split_rope_jit:ltx2_qknorm_split_rope_cuda",
_CUDA,
"LTX-2 QK-norm + split RoPE.",
_CUDA_SM100_PLUS,
"KDA native CUDA LTX-2 QK-norm + split RoPE (#29708).",
),
(
"diffusion.ltx25_decoder_rope",
@@ -343,10 +364,10 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
),
(
"diffusion.causal_conv3d_cat_pad",
KernelBackend.JIT,
KernelBackend.KDA,
"layout.causal_conv3d_cat_pad_jit:fused_causal_conv3d_cat_pad_cuda",
_CUDA,
"Causal Conv3d cat + pad.",
"KDA native CUDA causal Conv3d cat + pad (#29281).",
),
(
"diffusion.causal_conv3d_cat_pad",
python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py
+27
View File
@@ -222,6 +222,33 @@ def try_fused_scale_residual_norm_scale_shift(
return y, residual_out
def kda_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps):
"""Run the KDA B200 native CUDA path introduced by PR #27392.
Unlike ``try_fused_norm_scale_shift``, this explicit backend entry point
fails on unsupported inputs instead of silently returning ``None`` for a
caller-owned fallback.
"""
out = try_fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps)
if out is None:
raise RuntimeError("unsupported input for KDA norm-scale-shift CUDA")
return out
def kda_scale_residual_norm_scale_shift(
residual, x, gate, weight, bias, scale, shift, norm_type, eps
):
"""Run the KDA B200 residual + norm + scale/shift path from PR #27392."""
out = try_fused_scale_residual_norm_scale_shift(
residual, x, gate, weight, bias, scale, shift, norm_type, eps
)
if out is None:
raise RuntimeError(
"unsupported input for KDA scale-residual-norm-scale-shift CUDA"
)
return out
def try_fused_norm_scale_shift_fp8(
x, weight, bias, scale, shift, input_scale, norm_type, eps
):
python/sglang/kernels/spec.py
+1
View File
@@ -44,6 +44,7 @@ class KernelBackend(str, Enum):
AOT = "aot" # sgl_kernel wheel (CUDA / ROCm builds)
CUTE_DSL = "cute_dsl"
FLYDSL = "flydsl" # FlyDSL MLIR compiler (device=HIP, gfx950)
KDA = "KDA" # Kernel Design Agents generated implementation
FLASHINFER = "flashinfer"
DEEPGEMM = "deepgemm"
AITER = "aiter" # AMD aiter library (device=HIP)
test/registered/kernels/ops/layernorm/test_fused_op.py
+9 -1
View File
@@ -10,7 +10,7 @@ import pytest
import torch
import sglang.kernels as K
from sglang.kernels.fused_op import BaseFusedOp
from sglang.kernels.fused_op import BACKEND_METHODS, DEFAULT_PRIORITY, BaseFusedOp
from sglang.kernels.registry import KernelRegistry
from sglang.kernels.spec import CapabilityRequirement as Cap
from sglang.kernels.spec import KernelBackend, KernelSpec
@@ -64,6 +64,14 @@ def test_available_backends():
}
def test_kda_backend_contract():
assert KernelBackend.KDA.value == "KDA"
assert BACKEND_METHODS[KernelBackend.KDA] == "forward_kda"
assert DEFAULT_PRIORITY[0] is KernelBackend.KDA
with pytest.raises(NotImplementedError, match="no KDA backend"):
_ToyAdd().forward(_t(1.0), _t(2.0), backend=KernelBackend.KDA)
def test_priority_dispatch():
# TRITON is first in priority and always eligible (no capability).
assert _ToyAdd()(_t(1.0), _t(2.0)).item() == 1003.0
test/registered/kernels/ops/layernorm/test_kernels_namespace.py
+37
View File
@@ -30,6 +30,19 @@ EXPECTED = {
"quantization.nvfp4_gemm_swiglu_nvfp4_quant": {"cute_dsl"},
"kvcache.reshape_and_cache_flash": {"triton"},
"diffusion.apply_group_norm_silu": {"triton"},
"diffusion.norm_scale_shift": {"KDA", "cute_dsl", "flydsl"},
"diffusion.scale_residual_norm_scale_shift": {
"KDA",
"triton",
"cute_dsl",
"flydsl",
},
"diffusion.residual_gate_add": {"KDA"},
"diffusion.ltx2_qknorm_split_rope": {"KDA"},
"diffusion.causal_conv3d_cat_pad": {"KDA", "triton"},
"diffusion.flux2_layernorm_modulate_fp8_quant": {"KDA"},
"diffusion.flux2_qkv_epilogue": {"KDA"},
"diffusion.flux2_token_cat_fp8": {"KDA"},
}
_CPU = PlatformInfo(device_type="cpu")
@@ -83,6 +96,30 @@ def test_sparse_linear_attention_registry_targets_forward_kernel():
assert spec.target.endswith(":_attn_fwd")
@pytest.mark.parametrize(
"op, target_suffix",
(
("diffusion.norm_scale_shift", ":kda_norm_scale_shift"),
(
"diffusion.scale_residual_norm_scale_shift",
":kda_scale_residual_norm_scale_shift",
),
("diffusion.residual_gate_add", ":residual_gate_add"),
(
"diffusion.ltx2_qknorm_split_rope",
":ltx2_qknorm_split_rope_cuda",
),
(
"diffusion.causal_conv3d_cat_pad",
":fused_causal_conv3d_cat_pad_cuda",
),
),
)
def test_merged_diffusion_kda_provenance_backend(op, target_suffix):
spec = K.registry.get_backend(op, KernelBackend.KDA)
assert spec.target.endswith(target_suffix)
def test_single_backend_resolves_without_backend():
assert (
K.select_kernel("kvcache.reshape_and_cache_flash").backend