Add rms_norm related kernels. (sgl-project#4)

* Add rms_norm and fused_add_rms_norm kernel.

* Add gemma_rms_norm and gemma_fused_add_rms_norm.

* Move device check to utils.py.

* Add at::SGLXPUNorm for rmsnorm related ops.

* Use template parameter pack to replace get_update_vec_size.

* Change the max work group size on BMG from 1024 to 512. Replace SIMD with NUM_REDUCE_STAGES.

* Replace at::Tensor with torch::Tensor.

* Replace at::SGLXPUNorm with at::native::xpu. Remove redundant comments.

* Refactor preferred_vector_width. Remove 'default'.

* Refactor '_check_layer_norm_inputs'.

* Refactor template.

* Replace dpcpp with sycl.

* Refactor SYCL_DISPATCH_FLOATING_TYPES.

* weight only 1 dtype.

* Code Style.

* Repalce std::exp with sycl::exp and other ops when dtype is bf16.

Author: Hanyu-Jin

.github/workflows/lint.yml

@@ -19,4 +19,4 @@ jobs:
           pre-commit install
 
       - name: Linting
-        run: pre-commit run --all-files --show-diff-on-failure
+        run: pre-commit run --all-files --show-diff-on-failure

.isort.cfg

@@ -0,0 +1,3 @@
+[settings]
+profile=black
+known_first_party=sglang

.pre-commit-config.yaml

@@ -26,7 +26,7 @@ repos:
     rev: v0.11.7
     hooks:
       - id: ruff
-        args: [--select=F401, --fixable=F401]
+        args: [--select=F401, --fix] #able=F401]
         files: ^(benchmark/|docs/|examples/)
         exclude: \.ipynb$
   - repo: https://github.com/psf/black
@@ -38,7 +38,7 @@ repos:
     hooks:
       - id: codespell
         additional_dependencies: ['tomli']
-        args: ['--toml', 'python/pyproject.toml', '-L', 'cann']
+        args: ['--toml', 'pyproject.toml', '-L', 'cann']
         exclude: |
           (?x)^(
             test/srt/test_reasoning_parser\.py|

benchmark/bench_awq_dequant.py

@@ -1,5 +1,5 @@
 import itertools
-from typing import List, Tuple
+from typing import Tuple
 
 import torch
 import triton

benchmark/bench_cutlass_mla.py

@@ -1,5 +1,4 @@
 import argparse
-import copy
 import itertools
 
 import torch

benchmark/bench_fp8_blockwise_gemm.py

@@ -1,16 +1,16 @@
 import argparse
-import copy
 import itertools
 
 import deep_gemm
 import torch
 import triton
 from deep_gemm import get_col_major_tma_aligned_tensor
 from sgl_kernel import fp8_blockwise_scaled_mm
+from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
+
 from sglang.srt.layers.quantization.fp8_kernel import (
     w8a8_block_fp8_matmul_triton as w8a8_block_fp8_matmul,
 )
-from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
 
 
 def get_weight_shapes(args):

benchmark/bench_lightning_attention_decode.py

@@ -1,4 +1,3 @@
-import itertools
 import math
 
 import torch

benchmark/bench_moe_ep_post_reorder.py

@@ -1,6 +1,7 @@
 import torch
 import triton
 from sgl_kernel import ep_moe_post_reorder
+
 from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel
 
 batch_sizes = [64, 128, 256, 512, 640, 768, 1024, 2048, 4096]

benchmark/bench_moe_ep_pre_reorder.py

@@ -1,6 +1,7 @@
 import torch
 import triton
 from sgl_kernel import ep_moe_pre_reorder
+
 from sglang.srt.layers.moe.ep_moe.kernels import pre_reorder_triton_kernel
 
 batch_sizes = [64, 128, 256, 512, 640, 768, 1024, 2048, 4096]

benchmark/bench_moe_fused_gate.py

@@ -1,10 +1,7 @@
-import itertools
-import math
-
 import torch
 import triton
-import triton.language as tl
 from sgl_kernel import moe_fused_gate
+
 from sglang.srt.layers.moe.topk import biased_grouped_topk