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)