chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] refactor: refactor diffusion triton kernels (#18966)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-02-19 17:03:44 +08:00
committed by GitHub
parent 48642d5384
commit 19aa19b111
9 changed files with 612 additions and 598 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/elementwise.py
+1 -1
View File
@@ -1,7 +1,7 @@
import torch
from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
from sglang.multimodal_gen.runtime.layers.triton_ops import fuse_scale_shift_kernel
class MulAdd(CustomOp):
python/sglang/multimodal_gen/runtime/layers/layernorm.py
+3 -6
View File
@@ -20,6 +20,9 @@ if _is_cuda:
if _is_npu:
import torch_npu
from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_tensor_model_parallel_rank,
@@ -27,12 +30,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_tp_group,
)
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
from sglang.multimodal_gen.runtime.layers.triton_ops import (
fuse_scale_shift_kernel,
norm_infer,
rms_norm_fn,
triton_one_pass_rms_norm,
)
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py
+1 -1
View File
@@ -32,9 +32,9 @@ from typing import Any, Optional, Tuple
import torch
from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
from sglang.multimodal_gen.runtime.layers.triton_ops import apply_rotary_embedding
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
python/sglang/multimodal_gen/runtime/layers/triton_ops.py
-1207
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+3 -3
View File
@@ -15,6 +15,9 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.normalization import AdaLayerNormContinuous
from sglang.jit_kernel.diffusion.triton.scale_shift import (
fuse_scale_shift_gate_select01_kernel,
)
from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
@@ -41,9 +44,6 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
apply_flashinfer_rope_qk_inplace,
)
from sglang.multimodal_gen.runtime.layers.triton_ops import (
fuse_scale_shift_gate_select01_kernel,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,