[Kernel] Add KDA NVFP4 GEMM for Qwen3.x on SM120 (#36865)

Co-authored-by: Song Bian <biansonghz@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

python/sglang/kernels/kda_kernels/.clang-format

@@ -0,0 +1,25 @@
+BasedOnStyle: Google
+IndentWidth: 2
+ColumnLimit: 120
+AllowShortFunctionsOnASingleLine: Empty
+DerivePointerAlignment: false
+PointerAlignment: Left
+NamespaceIndentation: None
+SortIncludes: true
+AllowShortLoopsOnASingleLine: false
+BinPackParameters: false
+BinPackArguments: false
+AlignAfterOpenBracket: AlwaysBreak
+AlignOperands: Align
+PenaltyBreakBeforeFirstCallParameter: 1
+PenaltyReturnTypeOnItsOwnLine: 100
+
+IncludeCategories:
+  - Regex: '^<sgl_kernel/.*\.h>$'
+    Priority: 0
+  - Regex: '^<sgl_kernel/.*/.*>$'
+    Priority: 2
+  - Regex: '^<sgl_kernel/.*\.cuh>$'
+    Priority: 1
+  - Regex: '^<.*/.*>$'
+    Priority: 3

python/sglang/kernels/kda_kernels/README.md

@@ -0,0 +1,25 @@
+# Kernel Design Agent kernels
+
+This directory is the implementation home for kernels produced or extended by
+the Humanize2 / Kernel Design Agents workflow. `KernelBackend.KDA` records that
+provenance; it does not identify the implementation language. A KDA kernel may
+use CUDA, Triton, or CuTe DSL.
+
+Runtime code must continue to import the stable operator facade under
+`sglang.kernels.ops`. The facade owns registration and fallback policy, while
+this directory owns generated implementation modules and their CUDA sources.
+Importing `sglang.kernels` therefore remains metadata-only and does not eagerly
+load Triton, CUTLASS, or compile a JIT extension.
+
+| Kernel family | Implementation | Provenance |
+|---|---|---|
+| Qwen3.x ModelOpt NVFP4 GEMM on SM120 | `qwen3x_nvfp4_gemm_sm120.py` | [BBuf/KDA-Pilot#195](https://github.com/BBuf/KDA-Pilot/pull/195) at `516c976cee824a236679adf6eb525275a0a9a120` |
+| Qwen-Image norm / residual-norm scale-shift | `norm_scale_shift_jit.py` | [sgl-project/sglang#27392](https://github.com/sgl-project/sglang/pull/27392), merge commit `26e1d4d847` |
+| Cosmos3 causal Conv3D cat-pad | `causal_conv3d_cat_pad_jit.py` | [sgl-project/sglang#29281](https://github.com/sgl-project/sglang/pull/29281), merge commit `5996b54bd3` |
+| Diffusion residual-gate add | `residual_gate_add_jit.py` | [sgl-project/sglang#29361](https://github.com/sgl-project/sglang/pull/29361), merge commit `495f13fa12` |
+| LTX2 QK-norm split-RoPE | `ltx2_qknorm_split_rope_jit.py` | [sgl-project/sglang#29708](https://github.com/sgl-project/sglang/pull/29708), merge commit `fcb9f229b3` |
+| FLUX.2 FP8 producer and QKV packing fusions | `layernorm_modulate_triton.py`, `flux2_qkv_epilogue_jit.py`, `flux2_token_cat_fp8_triton.py` | [sgl-project/sglang#37162](https://github.com/sgl-project/sglang/pull/37162), merge commit `1c3ad92438` |
+
+For JIT kernels, the Python entry module and the corresponding source under
+`csrc/` move together. The shared `sglang.kernels.jit` loader remains build
+infrastructure rather than an ownership directory.

python/sglang/kernels/kda_kernels/__init__.py

@@ -0,0 +1,20 @@
+"""Implementation home for kernels produced by Kernel Design Agents.
+
+Runtime code should keep importing kernels through :mod:`sglang.kernels.ops`.
+This package owns generated implementations and provenance metadata, while the
+operator facades own registration, dispatch, and stable public exports.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+_ROOT = Path(__file__).resolve().parent
+
+
+def _cuda_source(name: str) -> str:
+    """Return an absolute path to a KDA-owned JIT CUDA source file."""
+    return str(_ROOT / "csrc" / name)
+
+
+__all__: list[str] = []

python/sglang/kernels/kda_kernels/causal_conv3d_cat_pad_jit.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.kda_kernels import _cuda_source
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
@@ -20,7 +21,7 @@ def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module:
     return load_jit(
         "diffusion_causal_conv3d_cat_pad",
         *args,
-        cuda_files=["diffusion/causal_conv3d_cat_pad.cuh"],
+        cuda_files=[_cuda_source("diffusion/causal_conv3d_cat_pad.cuh")],
         cuda_wrappers=[
             (
                 "causal_conv3d_cat_pad",

python/sglang/kernels/kda_kernels/csrc/diffusion/causal_conv3d_cat_pad.cuh

@@ -1,3 +1,4 @@
+// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #29281.
 // Native CUDA fast path for Cosmos3 VAE causal-Conv3D cat/pad copy.
 //
 // The op writes the output of:

python/sglang/kernels/kda_kernels/csrc/diffusion/flux2_qkv_epilogue.cuh

@@ -1,3 +1,4 @@
+// KDA provenance: Humanize2 / Kernel Design Agents, SGLang PR #37162.
 #pragma once
 
 #include <sgl_kernel/tensor.h>

python/sglang/kernels/kda_kernels/csrc/diffusion/ltx2_qknorm_split_rope.cuh

@@ -1,3 +1,4 @@
+// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #29708.
 // CUDA fast path for LTX2 Q/K RMSNorm + split RoPE.
 //
 // Developed with MIT HAN Lab Kernel Design Agents:

python/sglang/kernels/kda_kernels/csrc/diffusion/norm_scale_shift.cuh

@@ -1,3 +1,4 @@
+// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #27392.
 // Minimal native-CUDA fast path for generic bf16 hidden=3072 norm-scale-shift.
 //
 // Supported shape family:

python/sglang/kernels/kda_kernels/csrc/diffusion/residual_gate_add.cuh

@@ -1,3 +1,4 @@
+// KDA provenance: BBuf/KDA-Pilot, merged in SGLang PR #29361.
 // CUDA fast path for bit-exact diffusion residual-gate updates:
 //   out = residual + update * gate
 

python/sglang/kernels/kda_kernels/flux2_qkv_epilogue_jit.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.kernels.jit.utils import cache_once, load_jit
+from sglang.kernels.kda_kernels import _cuda_source
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
@@ -18,7 +19,7 @@ _ALIGN = 32
 def flux2_qkv_epilogue_module() -> Module:
     return load_jit(
         "flux2_qkv_epilogue_bf16",
-        cuda_files=["diffusion/flux2_qkv_epilogue.cuh"],
+        cuda_files=[_cuda_source("diffusion/flux2_qkv_epilogue.cuh")],
         cuda_wrappers=[
             (
                 "flux2_qkv_epilogue",

python/sglang/kernels/kda_kernels/layernorm_modulate_triton.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# KDA extension provenance: FLUX.2 FP8 producer fusion from SGLang PR #37162.
 """Fused LayerNorm + adaLN modulate Triton kernels for bf16 activations.
 
 Two fusions, each replacing an eager multi-kernel chain with a single

python/sglang/kernels/kda_kernels/ltx2_qknorm_split_rope_jit.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.kernels.jit.utils import cache_once, load_jit
+from sglang.kernels.kda_kernels import _cuda_source
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
@@ -15,7 +16,7 @@ if TYPE_CHECKING:
 def _jit_ltx2_qknorm_split_rope_module() -> Module:
     return load_jit(
         "diffusion_ltx2_qknorm_split_rope",
-        cuda_files=["diffusion/ltx2_qknorm_split_rope.cuh"],
+        cuda_files=[_cuda_source("diffusion/ltx2_qknorm_split_rope.cuh")],
         cuda_wrappers=[
             (
                 "ltx2_qknorm_split_rope_pair",

python/sglang/kernels/kda_kernels/norm_scale_shift_jit.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.kernels.jit.utils import cache_once, load_jit
+from sglang.kernels.kda_kernels import _cuda_source
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
@@ -81,7 +82,7 @@ def norm_scale_shift_module() -> Module:
         )
     return load_jit(
         "norm_scale_shift_native",
-        cuda_files=["diffusion/norm_scale_shift.cuh"],
+        cuda_files=[_cuda_source("diffusion/norm_scale_shift.cuh")],
         cuda_wrappers=[
             (
                 "nss_bf16_row",
@@ -118,7 +119,7 @@ _module = norm_scale_shift_module
 def norm_scale_shift_nvfp4_module() -> Module:
     return load_jit(
         "norm_scale_shift_nvfp4_native",
-        cuda_files=["diffusion/norm_scale_shift.cuh"],
+        cuda_files=[_cuda_source("diffusion/norm_scale_shift.cuh")],
         cuda_wrappers=[
             (
                 "srnss_nvfp4_row",

python/sglang/kernels/kda_kernels/qwen3x_nvfp4_gemm.py

@@ -0,0 +1,117 @@
+"""Lazy dispatch gate for the KDA Qwen3.x SM120 NVFP4 GEMM."""
+
+from __future__ import annotations
+
+import functools
+import importlib.util
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    import torch
+
+
+_SUPPORTED_SHAPES = frozenset(
+    (m, k, n)
+    for m in (1, 2, 4, 8)
+    for k, n in (
+        (2560, 18432),
+        (9216, 2560),
+        (4096, 24576),
+        (12288, 4096),
+    )
+).union({(9, 17408, 5120)})
+
+
+@functools.cache
+def _has_runtime(device: torch.device) -> bool:
+    import torch
+
+    try:
+        return (
+            importlib.util.find_spec("cutlass") is not None
+            and importlib.util.find_spec("cuda.bindings.driver") is not None
+            and torch.cuda.get_device_capability(device) == (12, 0)
+        )
+    except ModuleNotFoundError:
+        return False
+
+
+def _supports(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_sf: Optional[torch.Tensor],
+    weight_sf: torch.Tensor,
+    alpha: torch.Tensor,
+    out_dtype: torch.dtype,
+    out_features: int,
+) -> bool:
+    import torch
+
+    if (
+        input.device.type != "cuda"
+        or input.ndim != 2
+        or input.dtype != torch.uint8
+        or input_sf is None
+        or out_dtype != torch.bfloat16
+    ):
+        return False
+
+    m, packed_k = input.shape
+    k = packed_k * 2
+    n = int(out_features)
+    if (m, k, n) not in _SUPPORTED_SHAPES:
+        return False
+    if input.stride() != (packed_k, 1):
+        return False
+    if (
+        weight.dtype != torch.uint8
+        or weight.shape != (packed_k, n)
+        or weight.stride() != (1, packed_k)
+    ):
+        return False
+
+    scale_k = k // 16
+    padded_m = ((m + 127) // 128) * 128
+    if (
+        input_sf.dtype not in (torch.uint8, torch.float8_e4m3fn)
+        or input_sf.shape != (padded_m, scale_k)
+        or input_sf.stride() != (scale_k, 1)
+    ):
+        return False
+    if (
+        weight_sf.dtype != torch.float8_e4m3fn
+        or weight_sf.shape != (scale_k, n)
+        or weight_sf.stride() != (1, scale_k)
+    ):
+        return False
+    if alpha.dtype != torch.float32 or alpha.numel() != 1 or not alpha.is_contiguous():
+        return False
+    if any(t.device != input.device for t in (weight, input_sf, weight_sf, alpha)):
+        return False
+    return _has_runtime(input.device)
+
+
+def try_qwen3x_nvfp4_gemm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_sf: Optional[torch.Tensor],
+    weight_sf: torch.Tensor,
+    alpha: torch.Tensor,
+    out_dtype: torch.dtype,
+    out_features: int,
+) -> torch.Tensor | None:
+    """Run the E2E-qualified KDA GEMM, or return ``None`` for fallback."""
+    if not _supports(
+        input, weight, input_sf, weight_sf, alpha, out_dtype, out_features
+    ):
+        return None
+
+    from sglang.kernels.kda_kernels.qwen3x_nvfp4_gemm_sm120 import (
+        _run_qwen3x_nvfp4_gemm,
+    )
+
+    assert input_sf is not None
+    return _run_qwen3x_nvfp4_gemm(input, weight, input_sf, weight_sf, alpha)
+
+
+__all__ = ["try_qwen3x_nvfp4_gemm"]

python/sglang/kernels/kda_kernels/residual_gate_add_jit.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.kda_kernels import _cuda_source
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
@@ -29,7 +30,7 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
     return load_jit(
         "diffusion_residual_gate_add",
         *args,
-        cuda_files=["diffusion/residual_gate_add.cuh"],
+        cuda_files=[_cuda_source("diffusion/residual_gate_add.cuh")],
         cuda_wrappers=[
             (
                 "residual_gate_add",

python/sglang/kernels/ops/diffusion/README.md

@@ -26,10 +26,11 @@ re-export would make all of them import-time requirements everywhere.
 
 ## Layout
 
-One subpackage per **operator domain**; the backend is a **filename suffix**
-(`_triton`, `_jit`, `_cutedsl`, `_flydsl`, or `_bitexact` where that says more).
-This matches `ops/attention` and `ops/gemm`, and it keeps every implementation
-of one logical op in one directory.
+Ordinary implementations use one subpackage per **operator domain**; the
+compiler is a **filename suffix** (`_triton`, `_jit`, `_cutedsl`, `_flydsl`, or
+`_bitexact` where that says more). Implementations with Kernel Design Agents
+provenance live under `sglang.kernels.kda_kernels`; this facade remains their
+only supported runtime import surface.
 
 ```
 norm/        RMSNorm / LayerNorm / GroupNorm and their fused epilogues
@@ -41,6 +42,7 @@ layout/      pure data movement: USP/Ulysses relayout, varlen pack, causal pad
 common/      numerics primitives, platform predicates, non-Triton fallbacks
 sites/       request-scoped mount policy — NOT kernels (see below)
 ext/         JIT C++/CUDA extensions (Hunyuan3D raster/inpaint) — NOT kernels
+../../kda_kernels/  agent-generated implementations and their JIT CUDA sources
 ```
 
 ## The two numerical contracts
@@ -115,7 +117,7 @@ Several norms look interchangeable and are not. Start here.
 
 | Entry point | Backend | Contract | Applies to |
 |---|---|---|---|
-| `residual_gate_add` | JIT CUDA | bit-exact `residual + update * gate` | contiguous tensors, or a transposed-dense `[B, tokens, hidden]` residual/output with contiguous update and row-broadcast gate (SANA-Video) |
+| `residual_gate_add` | KDA (JIT CUDA) | bit-exact `residual + update * gate` | contiguous tensors, or a transposed-dense `[B, tokens, hidden]` residual/output with contiguous update and row-broadcast gate (SANA-Video) |
 
 The transposed-dense path uses a shared-memory tile to read the update in
 logical row-major order while keeping residual reads and output writes
@@ -132,7 +134,7 @@ tensor copy per residual site.
 | `fused_rope_rotate_half_bitexact` | Triton | bit-exact (elementwise only) |
 | `fused_interleaved_rope_fp64` | JIT CUDA | bit-exact vs paired SANA-Video fp64 RoPE |
 | `fused_inplace_helios_qk_rope` | JIT CUDA | bit-exact paired in-place RoPE for Helios' transposed frequency layout |
-| `ltx2_qknorm_split_rope_cuda` | JIT CUDA | close; **validated on B200** |
+| `ltx2_qknorm_split_rope_cuda` | KDA (JIT CUDA) | close; **validated on B200** |
 | `fused_ltx25_decoder_rope` | JIT CUDA | bit-exact paired 3D RoPE from cached compact axis tables |
 | `apply_rotary_embedding` | Triton (+fallbacks) | close; the generic entry point |
 | `hunyuan_qkv_rope_pack` | Triton | bit-exact; packs QKV and applies RoPE in one pass |
@@ -162,7 +164,9 @@ inspecting model modules is its whole job.
 
 ## Adding a kernel
 
-1. Put it in the operator domain it belongs to, with a backend suffix.
+1. Put ordinary implementations in their operator domain. Put a kernel
+   generated by the KDA workflow in `sglang.kernels.kda_kernels`, together
+   with its source revision and any JIT CUDA source files.
 2. Export it from `__init__.py` (`_EXPORTS`) and register a `KernelSpec`
    (`_SPECS`) — `test_import_surface.py` checks both resolve.
 3. Give it a `can_use_*` predicate; raise, don't return `None`.

python/sglang/kernels/ops/diffusion/__init__.py

@@ -8,20 +8,20 @@ Importing a submodule directly (``...diffusion.norm.norm_triton``) couples the
 caller to the file layout; ``test_import_surface.py`` guards against it.  The
 one exception is a test that deliberately exercises a single backend.
 
-Layout -- one subpackage per **operator domain** (``norm``, ``modulate``,
-``rope``, ``activation``, ``attention``, ``layout``) with the backend carried
-as a filename suffix (``_triton`` / ``_jit`` / ``_cutedsl`` / ``_flydsl``, or
-``_bitexact`` where that is the more informative label), matching how
-``ops/attention`` and ``ops/gemm`` are organized.  ``common`` holds shared
-numerics and platform plumbing, ``sites`` the request-scoped mount policy, and
-``ext`` the JIT C++/CUDA extensions that are not kernels.  Start from
-``README.md``: several norms look interchangeable and are not.
+Layout -- ordinary implementations use one subpackage per **operator domain**
+(``norm``, ``modulate``, ``rope``, ``activation``, ``attention``, ``layout``).
+Implementations generated by kernel-design agents live in
+``sglang.kernels.kda_kernels`` and are still exported through this facade.
+``common`` holds shared numerics and platform plumbing, ``sites`` the
+request-scoped mount policy, and ``ext`` JIT C++/CUDA extensions that are not
+kernels. Start from ``README.md``: several norms look interchangeable and are
+not.
 
 Resolution is lazy (PEP 562).  The backends have disjoint, heavy dependencies
 -- Triton, CUTLASS/CuTe-DSL, and FlyDSL (ROCm) -- so an eager
 re-export would turn every one of them into a hard import-time requirement on
-every platform.  ``_EXPORTS`` maps a symbol to its module and the import
-happens on first attribute access.
+every platform. ``_EXPORTS`` maps a symbol to its relative or fully-qualified
+owner module and the import happens on first attribute access.
 """
 
 from __future__ import annotations
@@ -85,7 +85,7 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
     (
         "diffusion.scale_residual_norm_scale_shift",
         KernelBackend.KDA,
-        "norm.norm_scale_shift_jit:kda_scale_residual_norm_scale_shift",
+        "sglang.kernels.kda_kernels.norm_scale_shift_jit:kda_scale_residual_norm_scale_shift",
         _CUDA_SM100_PLUS,
         "KDA B200 native CUDA residual + LayerNorm + scale/shift (#27392).",
     ),
@@ -113,14 +113,14 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
     (
         "diffusion.scale_residual_norm_scale_shift_nvfp4",
         KernelBackend.JIT,
-        "norm.norm_scale_shift_jit:try_fused_scale_residual_norm_scale_shift_nvfp4",
+        "sglang.kernels.kda_kernels.norm_scale_shift_jit:try_fused_scale_residual_norm_scale_shift_nvfp4",
         _CUDA,
         "Qwen residual LayerNorm/modulation + NVFP4 quantization.",
     ),
     (
         "diffusion.norm_scale_shift",
         KernelBackend.KDA,
-        "norm.norm_scale_shift_jit:kda_norm_scale_shift",
+        "sglang.kernels.kda_kernels.norm_scale_shift_jit:kda_norm_scale_shift",
         _CUDA_SM100_PLUS,
         "KDA B200 native CUDA LayerNorm + scale/shift (#27392).",
     ),
@@ -141,14 +141,14 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
     (
         "diffusion.layernorm_modulate",
         KernelBackend.TRITON,
-        "norm.layernorm_modulate_triton:fused_layernorm_modulate",
+        "sglang.kernels.kda_kernels.layernorm_modulate_triton:fused_layernorm_modulate",
         _CUDA,
         "Bit-exact LayerNorm + adaLN modulate.",
     ),
     (
         "diffusion.qk_head_layernorm",
         KernelBackend.TRITON,
-        "norm.layernorm_modulate_triton:fused_qk_head_layernorm",
+        "sglang.kernels.kda_kernels.layernorm_modulate_triton:fused_qk_head_layernorm",
         _CUDA,
         "Bit-exact per-head LayerNorm for q/k.",
     ),
@@ -183,7 +183,7 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
     (
         "diffusion.residual_gate_add",
         KernelBackend.KDA,
-        "modulate.residual_gate_add_jit:residual_gate_add",
+        "sglang.kernels.kda_kernels.residual_gate_add_jit:residual_gate_add",
         _CUDA,
         "KDA native CUDA residual + gate * update (#29361).",
     ),
@@ -218,21 +218,21 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
     (
         "diffusion.flux2_layernorm_modulate_fp8_quant",
         KernelBackend.KDA,
-        "norm.layernorm_modulate_triton:fused_layernorm_modulate_fp8_quant_raw",
+        "sglang.kernels.kda_kernels.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.KDA,
-        "rope.flux2_qkv_epilogue_jit:try_fused_flux2_qkv_epilogue",
+        "sglang.kernels.kda_kernels.flux2_qkv_epilogue_jit:try_fused_flux2_qkv_epilogue",
         _CUDA,
         "KDA-generated FLUX.2 QK RMS-norm + RoPE + joint QKV packing.",
     ),
     (
         "diffusion.flux2_token_cat_fp8",
         KernelBackend.KDA,
-        "layout.flux2_token_cat_fp8_triton:try_flux2_token_cat_fp8",
+        "sglang.kernels.kda_kernels.flux2_token_cat_fp8_triton:try_flux2_token_cat_fp8",
         _CUDA,
         "KDA-generated FLUX.2 token concatenation + static FP8 quantization.",
     ),
@@ -246,7 +246,7 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
     (
         "diffusion.ltx2_qknorm_split_rope",
         KernelBackend.KDA,
-        "rope.ltx2_qknorm_split_rope_jit:ltx2_qknorm_split_rope_cuda",
+        "sglang.kernels.kda_kernels.ltx2_qknorm_split_rope_jit:ltx2_qknorm_split_rope_cuda",
         _CUDA_SM100_PLUS,
         "KDA native CUDA LTX-2 QK-norm + split RoPE (#29708).",
     ),
@@ -365,7 +365,7 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
     (
         "diffusion.causal_conv3d_cat_pad",
         KernelBackend.KDA,
-        "layout.causal_conv3d_cat_pad_jit:fused_causal_conv3d_cat_pad_cuda",
+        "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit:fused_causal_conv3d_cat_pad_cuda",
         _CUDA,
         "KDA native CUDA causal Conv3d cat + pad (#29281).",
     ),
@@ -404,7 +404,11 @@ for _op, _backend, _target, _caps, _description in _SPECS:
         KernelSpec(
             op=_op,
             backend=_backend,
-            target=f"sglang.kernels.ops.diffusion.{_target}",
+            target=(
+                _target
+                if _target.startswith("sglang.")
+                else f"sglang.kernels.ops.diffusion.{_target}"
+            ),
             capabilities=_caps,
             format_signature=FormatSignature(description=_description),
             description=_description,
@@ -430,19 +434,19 @@ _EXPORTS: dict[str, str] = {
     "can_use_group_norm_silu_rows": "norm.group_norm_silu_twopass_triton",
     "group_norm_silu_4d": "norm.group_norm_silu_twopass_triton",
     "group_norm_silu_rows": "norm.group_norm_silu_twopass_triton",
-    "can_use_fused_layernorm_modulate": "norm.layernorm_modulate_triton",
-    "can_use_fused_qk_head_layernorm": "norm.layernorm_modulate_triton",
-    "fused_layernorm_modulate": "norm.layernorm_modulate_triton",
-    "fused_layernorm_modulate_fp8_quant_raw": "norm.layernorm_modulate_triton",
-    "fused_layernorm_modulate_raw": "norm.layernorm_modulate_triton",
-    "fused_qk_head_layernorm": "norm.layernorm_modulate_triton",
-    "is_plain_layer_norm": "norm.layernorm_modulate_triton",
+    "can_use_fused_layernorm_modulate": "sglang.kernels.kda_kernels.layernorm_modulate_triton",
+    "can_use_fused_qk_head_layernorm": "sglang.kernels.kda_kernels.layernorm_modulate_triton",
+    "fused_layernorm_modulate": "sglang.kernels.kda_kernels.layernorm_modulate_triton",
+    "fused_layernorm_modulate_fp8_quant_raw": "sglang.kernels.kda_kernels.layernorm_modulate_triton",
+    "fused_layernorm_modulate_raw": "sglang.kernels.kda_kernels.layernorm_modulate_triton",
+    "fused_qk_head_layernorm": "sglang.kernels.kda_kernels.layernorm_modulate_triton",
+    "is_plain_layer_norm": "sglang.kernels.kda_kernels.layernorm_modulate_triton",
     "rmsnorm_scale": "norm.native_bf16_rmsnorm_triton",
     "rmsnorm_tanh_residual": "norm.native_bf16_rmsnorm_triton",
     "norm_infer": "norm.norm_triton",
     "rms_norm_fn": "norm.norm_triton",
-    "try_fused_bias_mul_add": "norm.norm_scale_shift_jit",
-    "try_fused_bias_scale_residual_norm_scale_shift": "norm.norm_scale_shift_jit",
+    "try_fused_bias_mul_add": "sglang.kernels.kda_kernels.norm_scale_shift_jit",
+    "try_fused_bias_scale_residual_norm_scale_shift": "sglang.kernels.kda_kernels.norm_scale_shift_jit",
     "triton_one_pass_rms_norm": "norm.rmsnorm_onepass_triton",
     "can_use_fused_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact",
     "can_use_fused_scale_residual_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact",
@@ -450,12 +454,12 @@ _EXPORTS: dict[str, str] = {
     "fused_scale_residual_rmsnorm_scale_shift_bitexact": "norm.rmsnorm_scale_shift_bitexact",
     "fused_norm_scale_shift": "norm.scale_residual_norm_cutedsl",
     "fused_scale_residual_norm_scale_shift": "norm.scale_residual_norm_cutedsl",
-    "fused_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
-    "fused_scale_residual_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
-    "try_fused_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
-    "try_fused_scale_residual_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
+    "fused_norm_scale_shift_fp8": "sglang.kernels.kda_kernels.norm_scale_shift_jit",
+    "fused_scale_residual_norm_scale_shift_fp8": "sglang.kernels.kda_kernels.norm_scale_shift_jit",
+    "try_fused_norm_scale_shift_fp8": "sglang.kernels.kda_kernels.norm_scale_shift_jit",
+    "try_fused_scale_residual_norm_scale_shift_fp8": "sglang.kernels.kda_kernels.norm_scale_shift_jit",
     "validate_scale_shift": "norm.scale_residual_norm_cutedsl",
-    "try_fused_scale_residual_norm_scale_shift_nvfp4": "norm.norm_scale_shift_jit",
+    "try_fused_scale_residual_norm_scale_shift_nvfp4": "sglang.kernels.kda_kernels.norm_scale_shift_jit",
     "can_use_wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton",
     "wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton",
     "can_use_qk_rmsnorm_native": "norm.zimage_qk_rmsnorm_triton",
@@ -468,9 +472,9 @@ _EXPORTS: dict[str, str] = {
     "can_use_modulate_scale_shift_cuda": "modulate.modulate_scale_shift_jit",
     "modulate_scale_shift": "modulate.modulate_scale_shift_jit",
     "modulate_scale_shift_cuda": "modulate.modulate_scale_shift_jit",
-    "can_use_residual_gate_add_cuda": "modulate.residual_gate_add_jit",
-    "residual_gate_add": "modulate.residual_gate_add_jit",
-    "residual_gate_add_cuda": "modulate.residual_gate_add_jit",
+    "can_use_residual_gate_add_cuda": "sglang.kernels.kda_kernels.residual_gate_add_jit",
+    "residual_gate_add": "sglang.kernels.kda_kernels.residual_gate_add_jit",
+    "residual_gate_add_cuda": "sglang.kernels.kda_kernels.residual_gate_add_jit",
     "fuse_layernorm_scale_shift_gate_select01_kernel": "modulate.scale_shift_triton",
     "fuse_residual_layernorm_scale_shift_gate_select01_kernel": "modulate.scale_shift_triton",
     "fuse_scale_shift_kernel": "modulate.scale_shift_triton",
@@ -479,10 +483,10 @@ _EXPORTS: dict[str, str] = {
     "can_use_fused_temb_table_slices": "modulate.wan_temb_table_slices_triton",
     "fused_temb_table_slices": "modulate.wan_temb_table_slices_triton",
     # Rotary embeddings and the QK-norm chains fused around them
-    "try_fused_flux2_qkv_epilogue": "rope.flux2_qkv_epilogue_jit",
+    "try_fused_flux2_qkv_epilogue": "sglang.kernels.kda_kernels.flux2_qkv_epilogue_jit",
     "hunyuan_qkv_rope_pack": "rope.hunyuan_qkv_pack_triton",
-    "can_use_ltx2_qknorm_split_rope_cuda": "rope.ltx2_qknorm_split_rope_jit",
-    "ltx2_qknorm_split_rope_cuda": "rope.ltx2_qknorm_split_rope_jit",
+    "can_use_ltx2_qknorm_split_rope_cuda": "sglang.kernels.kda_kernels.ltx2_qknorm_split_rope_jit",
+    "ltx2_qknorm_split_rope_cuda": "sglang.kernels.kda_kernels.ltx2_qknorm_split_rope_jit",
     "apply_ltx2_split_rotary_emb": "rope.ltx2_rotary_triton",
     "can_use_ltx25_decoder_rope": "rope.ltx25_decoder_rope_jit",
     "fused_ltx25_decoder_rope": "rope.ltx25_decoder_rope_jit",
@@ -498,7 +502,7 @@ _EXPORTS: dict[str, str] = {
     "fused_inplace_helios_qk_rope": "rope.helios_qk_rope_jit",
     "apply_rotary_embedding": "rope.rotary_triton",
     # Tensor layout transformations fused with downstream quantization
-    "try_flux2_token_cat_fp8": "layout.flux2_token_cat_fp8_triton",
+    "try_flux2_token_cat_fp8": "sglang.kernels.kda_kernels.flux2_token_cat_fp8_triton",
     # Activation-function fusions
     "can_use_fused_bias_glu": "activation.sana_conv_post_triton",
     "can_use_fused_bias_silu": "activation.sana_conv_post_triton",
@@ -515,8 +519,8 @@ _EXPORTS: dict[str, str] = {
     "_attn_fwd": "attention.sparse_linear_attn_triton",
     "get_block_map": "attention.sparse_linear_attn_triton",
     # Data movement: bitwise identical to the aten chains they replace
-    "can_use_fused_causal_conv3d_cat_pad_cuda": "layout.causal_conv3d_cat_pad_jit",
-    "fused_causal_conv3d_cat_pad_cuda": "layout.causal_conv3d_cat_pad_jit",
+    "can_use_fused_causal_conv3d_cat_pad_cuda": "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit",
+    "fused_causal_conv3d_cat_pad_cuda": "sglang.kernels.kda_kernels.causal_conv3d_cat_pad_jit",
     "fused_causal_conv3d_cat_pad": "layout.causal_conv3d_cat_pad_triton",
     "pack_qkv_destination_major": "layout.ulysses_qkv_triton",
     "can_use_usp_merge_heads": "layout.usp_relayout_jit",
@@ -593,7 +597,8 @@ def __getattr__(name: str) -> Any:
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
     from importlib import import_module
 
-    value = getattr(import_module(f"{__name__}.{module}"), name)
+    owner = module if module.startswith("sglang.") else f"{__name__}.{module}"
+    value = getattr(import_module(owner), name)
     globals()[name] = value  # cache; later lookups skip __getattr__ entirely
     return value
 

python/sglang/kernels/ops/diffusion/norm/scale_residual_norm_cutedsl.py

@@ -263,7 +263,7 @@ def fused_norm_scale_shift(
     D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
     thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
     """
-    from sglang.kernels.ops.diffusion.norm.norm_scale_shift_jit import (
+    from sglang.kernels.kda_kernels.norm_scale_shift_jit import (
         try_fused_norm_scale_shift as _try_qwen_native_norm_scale_shift,
     )
 
@@ -349,7 +349,7 @@ def fused_scale_residual_norm_scale_shift(
     D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
     thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
     """
-    from sglang.kernels.ops.diffusion.norm.norm_scale_shift_jit import (
+    from sglang.kernels.kda_kernels.norm_scale_shift_jit import (
         try_fused_scale_residual_norm_scale_shift as _try_qwen_native_residual_path,
     )
 

python/sglang/kernels/ops/gemm/__init__.py

@@ -19,6 +19,8 @@ if TYPE_CHECKING:
 
 _CUDA = frozenset({CapabilityRequirement.CUDA})
 _SM90 = frozenset({CapabilityRequirement.cuda(min_sm=(9, 0), max_sm=(9, 0))})
+_SM120 = frozenset({CapabilityRequirement.cuda(min_sm=(12, 0), max_sm=(12, 0))})
+_KDA_PACKAGE = "sglang.kernels.kda_kernels"
 
 
 def _prefer_torch_rowwise_fp8(
@@ -212,6 +214,22 @@ register_kernel(
         description="Tiny bf16 GEMM (sglang.kernels.jit, JIT-only).",
     )
 )
+register_kernel(
+    KernelSpec(
+        op="gemm.qwen3x_nvfp4",
+        backend=KernelBackend.KDA,
+        target=f"{_KDA_PACKAGE}.qwen3x_nvfp4_gemm:try_qwen3x_nvfp4_gemm",
+        capabilities=_SM120,
+        format_signature=FormatSignature(
+            supported_dtypes=("uint8", "float8_e4m3fn", "bfloat16"),
+            description="Qwen3.x ModelOpt NVFP4 decode GEMM on SM120",
+        ),
+        description=(
+            "Shape-specialized Qwen3.x NVFP4 GEMM generated by Kernel Design "
+            "Agents and implemented with CuTe DSL."
+        ),
+    )
+)
 
 
 def fp8_scaled_mm(
@@ -264,12 +282,28 @@ def tiny_gemm_bf16(
     return impl(x, w, out, out_dtype=out_dtype, max_m=max_m)
 
 
+def try_qwen3x_nvfp4_gemm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_sf: Optional[torch.Tensor],
+    weight_sf: torch.Tensor,
+    alpha: torch.Tensor,
+    out_dtype: torch.dtype,
+    out_features: int,
+) -> torch.Tensor | None:
+    """Run the validated SM120 fast path, or return ``None`` for fallback."""
+    return get_kernel("gemm.qwen3x_nvfp4", KernelBackend.KDA)(
+        input, weight, input_sf, weight_sf, alpha, out_dtype, out_features
+    )
+
+
 __all__ = [
     "Fp8ScaledMMOp",
-    "fp8_scaled_mm",
     "bmm_fp8",
     "dsv3_fused_a_gemm",
+    "fp8_scaled_mm",
     "tiny_gemm_bf16",
+    "try_qwen3x_nvfp4_gemm",
 ]
 
 

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -140,6 +140,20 @@ def fp4_gemm(
     out_features: int,
     quant_mode: str = "w4a4",
 ) -> torch.Tensor:
+    from sglang.kernels.ops.gemm import try_qwen3x_nvfp4_gemm
+
+    kda_output = try_qwen3x_nvfp4_gemm(
+        input,
+        weight,
+        input_sf,
+        weight_sf,
+        alpha,
+        out_dtype,
+        out_features,
+    )
+    if kda_output is not None:
+        return kda_output
+
     if not enable_flashinfer_fp4_gemm:
         raise RuntimeError(
             "NVFP4 GEMM requires flashinfer's mm_fp4; please install flashinfer."

test/registered/kernels/ops/diffusion/test_import_surface.py

@@ -49,10 +49,15 @@ def _module_defines(module_path: str) -> set[str]:
     Importing would pull in Triton / CuTe-DSL / FlyDSL, none of which are
     installed on the CPU CI lane -- so this reads the source instead.
     """
-    path = _PACKAGE_DIR / (module_path.replace(".", "/") + ".py")
-    if not path.exists():
-        path = _PACKAGE_DIR / module_path.replace(".", "/") / "__init__.py"
-    assert path.exists(), f"{PACKAGE}.{module_path} does not exist"
+    if module_path.startswith("sglang."):
+        spec = importlib.util.find_spec(module_path)
+        assert spec is not None and spec.origin is not None, module_path
+        path = pathlib.Path(spec.origin)
+    else:
+        path = _PACKAGE_DIR / (module_path.replace(".", "/") + ".py")
+        if not path.exists():
+            path = _PACKAGE_DIR / module_path.replace(".", "/") / "__init__.py"
+        assert path.exists(), f"{PACKAGE}.{module_path} does not exist"
 
     names: set[str] = set()
     for node in ast.parse(path.read_text(encoding="utf-8")).body:
@@ -85,7 +90,12 @@ def _scan_root(root: str) -> tuple[frozenset[str], tuple[str, ...]]:
 
     for path in root_dir.rglob("*.py"):
         rel = path.relative_to(_REPO_ROOT).as_posix()
-        if rel.startswith("python/sglang/kernels/ops/diffusion/"):
+        if rel.startswith(
+            (
+                "python/sglang/kernels/ops/diffusion/",
+                "python/sglang/kernels/kda_kernels/",
+            )
+        ):
             continue
         try:
             source = path.read_text(encoding="utf-8")

test/registered/kernels/ops/layernorm/test_kernels_namespace.py

@@ -43,6 +43,7 @@ EXPECTED = {
     "diffusion.flux2_layernorm_modulate_fp8_quant": {"KDA"},
     "diffusion.flux2_qkv_epilogue": {"KDA"},
     "diffusion.flux2_token_cat_fp8": {"KDA"},
+    "gemm.qwen3x_nvfp4": {"KDA"},
 }
 
 _CPU = PlatformInfo(device_type="cpu")
@@ -120,6 +121,14 @@ def test_merged_diffusion_kda_provenance_backend(op, target_suffix):
     assert spec.target.endswith(target_suffix)
 
 
+def test_kda_backend_implementations_live_in_kda_home():
+    specs = [
+        spec for spec in K.registry.all_specs() if spec.backend is KernelBackend.KDA
+    ]
+    assert specs
+    assert all(spec.target.startswith("sglang.kernels.kda_kernels.") for spec in specs)
+
+
 def test_single_backend_resolves_without_backend():
     assert (
         K.select_kernel("kvcache.reshape_and_cache_flash").backend