From 25e38216b6af2472949f17b1ec0f4e476768d334 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sat, 14 Mar 2026 16:45:54 +0800 Subject: [PATCH] [kernel slimming] Clean many useless sgl-kernel deprecated kernels (#20277) --- .../benchmark_fused_collective.py | 3 +- .../sglang/jit_kernel/benchmark/bench_rope.py | 62 +---- .../jit_kernel/benchmark/bench_store_cache.py | 21 -- sgl-kernel/CMakeLists.txt | 3 - sgl-kernel/benchmark/bench_fp4_gemm.py | 2 +- sgl-kernel/benchmark/bench_fp8_gemm.py | 5 +- .../benchmark/bench_nvfp4_scaled_gemm.py | 2 +- .../benchmark/bench_per_tensor_quant_fp8.py | 5 +- .../benchmark/bench_rotary_embedding.py | 2 +- sgl-kernel/csrc/common_extension.cc | 12 - sgl-kernel/csrc/elementwise/rope.cu | 168 ----------- sgl-kernel/csrc/gemm/per_tensor_quant_fp8.cu | 123 --------- sgl-kernel/csrc/memory/store.cu | 147 ---------- sgl-kernel/include/sgl_kernel_ops.h | 16 -- sgl-kernel/python/sgl_kernel/__init__.py | 9 +- sgl-kernel/python/sgl_kernel/elementwise.py | 116 -------- sgl-kernel/python/sgl_kernel/gemm.py | 128 +-------- sgl-kernel/python/sgl_kernel/memory.py | 17 -- .../sgl_kernel/testing/rotary_embedding.py | 43 ++- sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py | 6 +- sgl-kernel/tests/test_fp4_gemm.py | 154 ----------- sgl-kernel/tests/test_fp4_quantize.py | 260 ------------------ sgl-kernel/tests/test_per_tensor_quant_fp8.py | 67 ----- sgl-kernel/tests/test_rotary_embedding.py | 167 ----------- test/registered/kernels/test_fp4_moe.py | 3 +- test/registered/moe/test_cutedsl_moe.py | 2 +- 26 files changed, 60 insertions(+), 1483 deletions(-) delete mode 100644 sgl-kernel/csrc/elementwise/rope.cu delete mode 100644 sgl-kernel/csrc/gemm/per_tensor_quant_fp8.cu delete mode 100644 sgl-kernel/csrc/memory/store.cu delete mode 100644 sgl-kernel/tests/test_fp4_gemm.py delete mode 100644 sgl-kernel/tests/test_fp4_quantize.py delete mode 100644 sgl-kernel/tests/test_per_tensor_quant_fp8.py delete mode 100644 sgl-kernel/tests/test_rotary_embedding.py diff --git a/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py index 4aebf62b9..7050897c0 100644 --- a/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py +++ b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py @@ -42,7 +42,8 @@ from sglang.srt.layers.quantization.fp8_kernel import static_quant_fp8 try: from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM from sgl_kernel import rmsnorm as SGL_RMS_NORM - from sgl_kernel import scaled_fp4_quant as SGL_SCALED_FP4_QUANT + + from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as SGL_SCALED_FP4_QUANT except Exception: # pragma: no cover - fallback on non-supported platforms SGL_FUSED_ADD_RMS_NORM = None SGL_RMS_NORM = None diff --git a/python/sglang/jit_kernel/benchmark/bench_rope.py b/python/sglang/jit_kernel/benchmark/bench_rope.py index 74c4f1003..2ca71d4ca 100644 --- a/python/sglang/jit_kernel/benchmark/bench_rope.py +++ b/python/sglang/jit_kernel/benchmark/bench_rope.py @@ -83,25 +83,6 @@ def sglang_pos_enc_rope( ) -def sgl_kernel_rope( - q: torch.Tensor, - k: torch.Tensor, - positions: torch.Tensor, - is_neox: bool, -) -> None: - from sgl_kernel import apply_rope_with_cos_sin_cache_inplace - - head_size = q.shape[-1] - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=q.view(q.shape[0], -1), - key=k.view(k.shape[0], -1), - head_size=head_size, - cos_sin_cache=COS_SIN_CACHE, - is_neox=is_neox, - ) - - def sglang_fused_rope( q: torch.Tensor, k: torch.Tensor, @@ -150,37 +131,6 @@ def jit_rope_then_store( ) -def sgl_kernel_fused_rope_store( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - positions: torch.Tensor, - out_loc: torch.Tensor, - is_neox: bool, -) -> None: - from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace - - head_size = q.shape[-1] - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=q.view(q.shape[0], -1), - key=k.view(k.shape[0], -1), - head_size=head_size, - cos_sin_cache=COS_SIN_CACHE, - is_neox=is_neox, - fused_set_kv_buffer_arg=FusedSetKVBufferArg( - value=v.view(v.shape[0], -1), - k_buffer=k_cache, - v_buffer=v_cache, - k_scale=None, - v_scale=None, - cache_loc=out_loc, - ), - ) - - def jit_fused_rope_store( q: torch.Tensor, k: torch.Tensor, @@ -221,14 +171,13 @@ IS_NEOX_RANGE = get_benchmark_range( # Benchmark 1: RoPE only # --------------------------------------------------------------------------- -ROPE_LINE_VALS = ["flashinfer", "jit_pos_enc", "sgl_kernel", "jit_fused_rope"] +ROPE_LINE_VALS = ["flashinfer", "jit_pos_enc", "jit_fused_rope"] ROPE_LINE_NAMES = [ "FlashInfer", "SGL JIT PosEnc", - "sgl-kernel", "SGL JIT Fused RoPE", ] -ROPE_STYLES = [("green", "-."), ("red", "-"), ("orange", "-"), ("blue", "--")] +ROPE_STYLES = [("green", "-."), ("red", "-"), ("blue", "--")] rope_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE)) @@ -270,7 +219,6 @@ def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str) FN_MAP = { "flashinfer": flashinfer_rope, "jit_pos_enc": sglang_pos_enc_rope, - "sgl_kernel": sgl_kernel_rope, "jit_fused_rope": sglang_fused_rope, } fn = lambda: FN_MAP[provider](q, k, positions, is_neox) @@ -281,13 +229,12 @@ def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str) # Benchmark 2: RoPE + KV cache store # --------------------------------------------------------------------------- -STORE_LINE_VALS = ["jit_rope_then_store", "sgl_kernel_fused_store", "jit_fused_store"] +STORE_LINE_VALS = ["jit_rope_then_store", "jit_fused_store"] STORE_LINE_NAMES = [ "SGL JIT RoPE + Store", - "sgl-kernel Fused RoPE + Store", "SGL JIT Fused RoPE + Store", ] -STORE_STYLES = [("red", "-"), ("orange", "-"), ("blue", "--")] +STORE_STYLES = [("red", "-"), ("blue", "--")] store_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE)) @@ -343,7 +290,6 @@ def benchmark_store(batch_size: int, num_q_k_heads: str, is_neox: bool, provider FN_MAP = { "jit_rope_then_store": jit_rope_then_store, - "sgl_kernel_fused_store": sgl_kernel_fused_rope_store, "jit_fused_store": jit_fused_rope_store, } fn = lambda: FN_MAP[provider]( diff --git a/python/sglang/jit_kernel/benchmark/bench_store_cache.py b/python/sglang/jit_kernel/benchmark/bench_store_cache.py index 1f179242e..700274772 100644 --- a/python/sglang/jit_kernel/benchmark/bench_store_cache.py +++ b/python/sglang/jit_kernel/benchmark/bench_store_cache.py @@ -4,7 +4,6 @@ from typing import Tuple import torch import triton import triton.testing -from sgl_kernel import set_kv_buffer_kernel from sglang.jit_kernel.benchmark.utils import ( DEFAULT_DEVICE, @@ -14,19 +13,6 @@ from sglang.jit_kernel.benchmark.utils import ( ) from sglang.jit_kernel.kvcache import store_cache -_is_hip = bool(torch.version.hip) -HAS_AOT_STORE_CACHE = hasattr(torch.ops.sgl_kernel, "store_kv_cache") - - -def sglang_aot_store_cache( - k: torch.Tensor, - v: torch.Tensor, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - indices: torch.Tensor, -) -> None: - set_kv_buffer_kernel(k_cache, v_cache, indices, k, v) - def sglang_jit_store_cache( k: torch.Tensor, @@ -83,11 +69,6 @@ ITEM_SIZE = get_benchmark_range( LINE_VALS = ["jit", "torch_compile", "torch_streams"] LINE_NAMES = ["SGL JIT Kernel", "PyTorch Compile", "PyTorch 2 Stream"] STYLES = [("blue", "--"), ("red", ":"), ("green", "-.")] -# Keep non-HIP benchmark lines unchanged; only HIP tolerates missing AOT op. -if (not _is_hip) or HAS_AOT_STORE_CACHE: - LINE_VALS = ["aot"] + LINE_VALS - LINE_NAMES = ["SGL AOT Kernel"] + LINE_NAMES - STYLES = [("orange", "-")] + STYLES X_NAMES = ["item_size", "batch_size"] CONFIGS = list(itertools.product(ITEM_SIZE, BS_RANGE)) @@ -128,8 +109,6 @@ def benchmark( "torch_compile": torch_compile_store_cache, "torch_streams": torch_streams_store_cache, } - if (not _is_hip) or HAS_AOT_STORE_CACHE: - FN_MAP["aot"] = sglang_aot_store_cache def fn(): impl = FN_MAP[provider] diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index e735e9c2d..743c29104 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -270,7 +270,6 @@ set(SOURCES "csrc/elementwise/concat_mla.cu" "csrc/elementwise/copy.cu" "csrc/elementwise/fused_add_rms_norm_kernel.cu" - "csrc/elementwise/rope.cu" "csrc/elementwise/pos_enc.cu" "csrc/elementwise/topk.cu" "csrc/expert_specialization/es_fp8_blockwise.cu" @@ -286,7 +285,6 @@ set(SOURCES "csrc/gemm/fp8_blockwise_gemm_kernel.cu" "csrc/gemm/fp8_gemm_kernel.cu" "csrc/gemm/int8_gemm_kernel.cu" - "csrc/gemm/per_tensor_quant_fp8.cu" "csrc/gemm/per_token_group_quant_8bit.cu" "csrc/gemm/per_token_group_quant_8bit_v2.cu" "csrc/gemm/per_token_quant_fp8.cu" @@ -297,7 +295,6 @@ set(SOURCES "csrc/kvcacheio/transfer.cu" "csrc/mamba/causal_conv1d.cu" - "csrc/memory/store.cu" "csrc/memory/weak_ref_tensor.cpp" "csrc/moe/cutlass_moe/w4a8/scaled_mm_entry.cu" diff --git a/sgl-kernel/benchmark/bench_fp4_gemm.py b/sgl-kernel/benchmark/bench_fp4_gemm.py index 7c4e187a2..a7fc55e0b 100755 --- a/sgl-kernel/benchmark/bench_fp4_gemm.py +++ b/sgl-kernel/benchmark/bench_fp4_gemm.py @@ -5,8 +5,8 @@ import os import torch import triton from flashinfer import mm_fp4 -from sgl_kernel import cutlass_scaled_fp4_mm, scaled_fp4_quant +from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant from sglang.srt.utils import get_device_capability, is_sm100_supported # CI environment detection diff --git a/sgl-kernel/benchmark/bench_fp8_gemm.py b/sgl-kernel/benchmark/bench_fp8_gemm.py index a49f3b06f..1dbedb66a 100644 --- a/sgl-kernel/benchmark/bench_fp8_gemm.py +++ b/sgl-kernel/benchmark/bench_fp8_gemm.py @@ -7,7 +7,8 @@ from typing import Optional, Tuple import torch import triton from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm -from sgl_kernel import sgl_per_tensor_quant_fp8 + +from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 # Optional vLLM import try: @@ -97,7 +98,7 @@ def sglang_scaled_fp8_quant( if scale is None: scale = torch.zeros(1, device=input.device, dtype=torch.float32) is_static = False - sgl_per_tensor_quant_fp8(input, output, scale, is_static) + per_tensor_quant_fp8(input, output, scale, is_static) return output, scale diff --git a/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py b/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py index 3867f6093..eeb5842ed 100644 --- a/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py +++ b/sgl-kernel/benchmark/bench_nvfp4_scaled_gemm.py @@ -5,8 +5,8 @@ import os import torch import triton -from sgl_kernel import cutlass_scaled_fp4_mm, scaled_fp4_quant +from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant from sglang.srt.utils import get_device_capability # CI environment detection diff --git a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py index d44b017a0..43c2caddc 100644 --- a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py +++ b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py @@ -7,7 +7,8 @@ import numpy as np import torch import triton import triton.testing -from sgl_kernel import sgl_per_tensor_quant_fp8 + +from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 # Optional imports try: @@ -51,7 +52,7 @@ def sglang_scaled_fp8_quant( if scale is None: scale = torch.zeros(1, device=input.device, dtype=torch.float32) is_static = False - sgl_per_tensor_quant_fp8(input, output, scale, is_static) + per_tensor_quant_fp8(input, output, scale, is_static) return output, scale diff --git a/sgl-kernel/benchmark/bench_rotary_embedding.py b/sgl-kernel/benchmark/bench_rotary_embedding.py index 0cab8e653..30614355f 100644 --- a/sgl-kernel/benchmark/bench_rotary_embedding.py +++ b/sgl-kernel/benchmark/bench_rotary_embedding.py @@ -3,9 +3,9 @@ import os import torch import triton -from sgl_kernel import FusedSetKVBufferArg from sgl_kernel.testing.rotary_embedding import ( FlashInferRotaryEmbedding, + FusedSetKVBufferArg, MHATokenToKVPool, RotaryEmbedding, create_inputs, diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 5a572e092..cdce0064b 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -84,12 +84,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); m.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul); - m.def( - "apply_rope_pos_ids_cos_sin_cache(Tensor q, Tensor k, Tensor! q_rope, Tensor! k_rope, Tensor cos_sin_cache, " - "Tensor pos_ids, bool interleave, bool enable_pdl, " - "Tensor? v, Tensor!? k_buffer, Tensor!? v_buffer, Tensor? kv_cache_loc) -> ()"); - m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache); - m.def( "rotary_embedding(Tensor positions, Tensor! query," " Tensor!? key, int head_size," @@ -151,9 +145,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { " float eps, float fp8_min, float fp8_max, bool scale_ue8m0, bool fuse_silu_and_mul, Tensor? masked_m) -> ()"); m.impl("sgl_per_token_group_quant_8bit_v2", torch::kCUDA, &sgl_per_token_group_quant_8bit_v2); - m.def("sgl_per_tensor_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s, bool is_static) -> ()"); - m.impl("sgl_per_tensor_quant_fp8", torch::kCUDA, &sgl_per_tensor_quant_fp8); - m.def("sgl_per_token_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s) -> ()"); m.impl("sgl_per_token_quant_fp8", torch::kCUDA, &sgl_per_token_quant_fp8); @@ -355,9 +346,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { /* * From csrc/memory */ - m.def("store_kv_cache(Tensor k_cache, Tensor v_cache, Tensor out_loc, Tensor k, Tensor v) -> ()"); - m.impl("store_kv_cache", &store_kv_cache); - m.def("weak_ref_tensor(Tensor tensor) -> Tensor"); m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor); diff --git a/sgl-kernel/csrc/elementwise/rope.cu b/sgl-kernel/csrc/elementwise/rope.cu deleted file mode 100644 index b6abd1cce..000000000 --- a/sgl-kernel/csrc/elementwise/rope.cu +++ /dev/null @@ -1,168 +0,0 @@ -/* - * Copyright (c) 2024 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#include - -#include "pos_enc.cuh" -#include "utils.h" - -using namespace flashinfer; - -void apply_rope_pos_ids_cos_sin_cache( - at::Tensor q, - at::Tensor k, - at::Tensor q_rope, - at::Tensor k_rope, - at::Tensor cos_sin_cache, - at::Tensor pos_ids, - bool interleave, - bool enable_pdl, - const std::optional& v, - const std::optional& k_buffer, - const std::optional& v_buffer, - const std::optional& kv_cache_loc) { - CHECK_LAST_DIM_CONTIGUOUS(q); - CHECK_LAST_DIM_CONTIGUOUS(k); - - const bool save_kv_cache = v.has_value(); - if (save_kv_cache) { - TORCH_CHECK(v.has_value()); - TORCH_CHECK(k_buffer.has_value()); - TORCH_CHECK(v_buffer.has_value()); - TORCH_CHECK(kv_cache_loc.has_value()); - CHECK_LAST_DIM_CONTIGUOUS(v.value()); - CHECK_LAST_DIM_CONTIGUOUS(k_buffer.value()); - CHECK_LAST_DIM_CONTIGUOUS(v_buffer.value()); - CHECK_DIM(3, k_buffer.value()); // k_buffer: (nnz, H_K, D) - CHECK_DIM(3, v_buffer.value()); // v_buffer: (nnz, H_V, D) - CHECK_DIM(3, v.value()); // v: (nnz, H_V, D) - CHECK_DIM(1, kv_cache_loc.value()); // v: (n) - CHECK_INPUT(kv_cache_loc.value()); - } - size_t k_buffer_stride_n = save_kv_cache ? k_buffer->stride(0) : 0; - size_t k_buffer_stride_h = save_kv_cache ? k_buffer->stride(1) : 0; - size_t v_buffer_stride_n = save_kv_cache ? v_buffer->stride(0) : 0; - size_t v_buffer_stride_h = save_kv_cache ? v_buffer->stride(1) : 0; - size_t v_stride_n = save_kv_cache ? v->stride(0) : 0; - size_t v_stride_h = save_kv_cache ? v->stride(1) : 0; - auto kv_cache_loc_ptr = save_kv_cache ? static_cast(kv_cache_loc->data_ptr()) : nullptr; - - CHECK_INPUT(cos_sin_cache); - CHECK_INPUT(pos_ids); - auto device = q.device(); - CHECK_EQ(k.device(), device); - CHECK_EQ(cos_sin_cache.device(), device); - CHECK_EQ(pos_ids.device(), device); - CHECK_DIM(3, q); // q: (nnz, H_Q, D) - CHECK_DIM(3, k); // k: (nnz, H_K, D) - - // cos_sin_cache: (max_seq_len, R) - // First half of R is cos, second half is sin - CHECK_DIM(2, cos_sin_cache); - CHECK_EQ(q.size(0), k.size(0)); - CHECK_EQ(q.size(2), k.size(2)); - unsigned int rotary_dim = cos_sin_cache.size(1); - unsigned int num_qo_heads = q.size(1); - unsigned int num_kv_heads = k.size(1); - unsigned int head_dim = q.size(2); - unsigned int nnz = q.size(0); - size_t q_stride_n = q.stride(0); - size_t q_stride_h = q.stride(1); - size_t k_stride_n = k.stride(0); - size_t k_stride_h = k.stride(1); - - size_t q_rope_stride_n = q_rope.stride(0); - size_t q_rope_stride_h = q_rope.stride(1); - size_t k_rope_stride_n = k_rope.stride(0); - size_t k_rope_stride_h = k_rope.stride(1); - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(q.scalar_type(), c_type, [&] { - // TODO temporarily only use `BatchQKApplyRotaryPosIdsCosSinCacheEnhanced` when save_kv_cache - // to avoid changing original code path; but this branch is feature-complete and should switch to this later - if (save_kv_cache) { - cudaError_t status = BatchQKApplyRotaryPosIdsCosSinCacheEnhanced( - static_cast(q.data_ptr()), - static_cast(k.data_ptr()), - save_kv_cache ? static_cast(v->data_ptr()) : nullptr, - static_cast(q_rope.data_ptr()), - static_cast(k_rope.data_ptr()), - save_kv_cache ? static_cast(k_buffer->data_ptr()) : nullptr, - save_kv_cache ? static_cast(v_buffer->data_ptr()) : nullptr, - static_cast(cos_sin_cache.data_ptr()), - static_cast(pos_ids.data_ptr()), - nnz, - num_qo_heads, - num_kv_heads, - rotary_dim, - head_dim, - q_stride_n, - q_stride_h, - k_stride_n, - k_stride_h, - v_stride_n, - v_stride_h, - q_rope_stride_n, - q_rope_stride_h, - k_rope_stride_n, - k_rope_stride_h, - k_buffer_stride_n, - k_buffer_stride_h, - v_buffer_stride_n, - v_buffer_stride_h, - kv_cache_loc_ptr, - interleave, - save_kv_cache, - enable_pdl, - stream); - TORCH_CHECK( - status == cudaSuccess, - "BatchQKApplyRotaryPosIdsCosSinCacheEnhanced failed with error code " + - std::string(cudaGetErrorString(status))); - } else { - TORCH_CHECK(!enable_pdl); - cudaError_t status = BatchQKApplyRotaryPosIdsCosSinCache( - static_cast(q.data_ptr()), - static_cast(k.data_ptr()), - static_cast(q_rope.data_ptr()), - static_cast(k_rope.data_ptr()), - static_cast(cos_sin_cache.data_ptr()), - static_cast(pos_ids.data_ptr()), - nnz, - num_qo_heads, - num_kv_heads, - rotary_dim, - head_dim, - q_stride_n, - q_stride_h, - k_stride_n, - k_stride_h, - q_rope_stride_n, - q_rope_stride_h, - k_rope_stride_n, - k_rope_stride_h, - interleave, - stream); - TORCH_CHECK( - status == cudaSuccess, - "BatchQKApplyRotaryPosIdsCosSinCache failed with error code " + std::string(cudaGetErrorString(status))); - } - return true; - }); -} diff --git a/sgl-kernel/csrc/gemm/per_tensor_quant_fp8.cu b/sgl-kernel/csrc/gemm/per_tensor_quant_fp8.cu deleted file mode 100644 index 7afff7794..000000000 --- a/sgl-kernel/csrc/gemm/per_tensor_quant_fp8.cu +++ /dev/null @@ -1,123 +0,0 @@ -#include -#include - -#include -#include -#include - -#include "utils.h" - -template -__global__ void -per_tensor_absmax_kernel(const T* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) { - float max_value = 0.0f; - unsigned int tid = threadIdx.x; - unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x; - const int grid_size = blockDim.x * gridDim.x; - - constexpr uint32_t vec_size = 16 / sizeof(T); - using vec_t = flashinfer::vec_t; - - const int32_t num_vec_elems = num_elements / vec_size; - - for (int32_t i = gid; i < num_vec_elems; i += grid_size) { - vec_t input_vec; - input_vec.cast_load(input + i * vec_size); - -#pragma unroll - for (uint32_t j = 0; j < vec_size; ++j) { - float val = static_cast(input_vec[j]); - max_value = fmaxf(max_value, fabsf(val)); - } - } - - const int32_t remaining_start = num_vec_elems * vec_size; - for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) { - float val = static_cast(input[idx]); - max_value = fmaxf(max_value, fabsf(val)); - } - - max_value = blockReduceMax(max_value); - - if (tid == 0) { - atomicMaxFloat(output_s, max_value / FP8_E4M3_MAX); - } -} - -template -__global__ void per_tensor_quant_fp8_kernel( - const T* __restrict__ input, - DST_DTYPE* __restrict__ output, - const float* __restrict__ scale, - const int64_t num_elements) { - const int gid = blockIdx.x * blockDim.x + threadIdx.x; - const int grid_size = blockDim.x * gridDim.x; - const float scale_val = 1.0f / (*scale); - - // We want to store 128 bits of data at a time. 16 = 128 / 8 bits - // Load is already vectorized, so 16 elements work for T. - const uint32_t VEC_SIZE = 16; - using vec_t = flashinfer::vec_t; - - const int32_t num_vec_elems = num_elements / VEC_SIZE; - - for (int32_t i = gid; i < num_vec_elems; i += grid_size) { - vec_t input_vec; - input_vec.cast_load(input + i * VEC_SIZE); - - DST_DTYPE output_arr[VEC_SIZE]; -#pragma unroll - for (uint32_t j = 0; j < VEC_SIZE; ++j) { - float val = fmax(fmin(static_cast(input_vec[j]) * scale_val, FP8_E4M3_MAX), -FP8_E4M3_MAX); -#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3) - output_arr[j] = static_cast(val); -#else - output_arr[j] = c10::Float8_e4m3fnuz( - __hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret), - c10::Float8_e4m3fnuz::from_bits()); -#endif - } - *(uint4*)(output + i * VEC_SIZE) = *(uint4*)output_arr; - } - - const int32_t remaining_start = num_vec_elems * VEC_SIZE; - for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) { - float val = fmax(-FP8_E4M3_MAX, fmin(static_cast(input[idx]) * scale_val, FP8_E4M3_MAX)); -#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3) - output[idx] = static_cast(val); -#else - output[idx] = c10::Float8_e4m3fnuz( - __hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret), - c10::Float8_e4m3fnuz::from_bits()); -#endif - } -} - -void sgl_per_tensor_quant_fp8(torch::Tensor input, torch::Tensor output_q, torch::Tensor output_s, bool is_static) { - CHECK_INPUT(input); - CHECK_INPUT(output_q); - CHECK_INPUT(output_s); - - const int block_size = 256; - const int num_elements = input.numel(); - const int num_blocks = min((num_elements + block_size - 1) / block_size, 1024); - - dim3 grid(num_blocks); - dim3 block(block_size); - - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), scalar_t, [&] { - if (is_static == false) { - per_tensor_absmax_kernel<<>>( - static_cast(input.data_ptr()), static_cast(output_s.data_ptr()), num_elements); - } - - per_tensor_quant_fp8_kernel<<>>( - static_cast(input.data_ptr()), - static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()), - static_cast(output_s.data_ptr()), - num_elements); - return true; - }); -} diff --git a/sgl-kernel/csrc/memory/store.cu b/sgl-kernel/csrc/memory/store.cu deleted file mode 100644 index c6dd97ebd..000000000 --- a/sgl-kernel/csrc/memory/store.cu +++ /dev/null @@ -1,147 +0,0 @@ -#include -#include -#include -#include - -#include -#include -#include - -namespace { - -using std::size_t; -using std::uint64_t; - -// Each warp will process 256 bytes per loop iteration -template -__global__ void store_kv_cache_256x1( - uint64_t* __restrict__ k_cache, - uint64_t* __restrict__ v_cache, - const T* __restrict__ out_loc, - const size_t length, - const uint64_t* __restrict__ k, - const uint64_t* __restrict__ v, - const size_t kv_cache_stride, - const size_t kv_input_stride, - const size_t num_items) { - const auto idx = blockIdx.x * blockDim.x + threadIdx.x; - const auto warp_id = idx / 32; - const auto lane_id = idx % 32; - if (warp_id >= length) return; - const auto offset = out_loc[warp_id]; - const auto k_dst = k_cache + offset * kv_cache_stride; - const auto v_dst = v_cache + offset * kv_cache_stride; - const auto k_src = k + warp_id * kv_input_stride; - const auto v_src = v + warp_id * kv_input_stride; - for (size_t i = 0; i < num_items; ++i) { - k_dst[lane_id + i * 32] = k_src[lane_id + i * 32]; - v_dst[lane_id + i * 32] = v_src[lane_id + i * 32]; - } -} - -// Each warp will process 128 bytes per loop iteration -template -__global__ void store_kv_cache_128x2( - uint64_t* __restrict__ k_cache, - uint64_t* __restrict__ v_cache, - const T* __restrict__ out_loc, - const size_t length, - const uint64_t* __restrict__ k, - const uint64_t* __restrict__ v, - const size_t kv_cache_stride, - const size_t kv_input_stride, - const size_t num_items) { - const auto idx = blockIdx.x * blockDim.x + threadIdx.x; - const auto warp_id = idx / 32; - const auto lane_id = idx % 32; - if (warp_id >= length) return; - const auto offset = out_loc[warp_id]; - const auto copy_k = lane_id < 16; - const auto copy_id = lane_id % 16; - const auto cache = copy_k ? k_cache : v_cache; - const auto input = copy_k ? k : v; - const auto dst = cache + offset * kv_cache_stride; - const auto src = input + warp_id * kv_input_stride; - for (size_t i = 0; i < num_items; ++i) { - dst[copy_id + i * 16] = src[copy_id + i * 16]; - } -} - -} // namespace - -auto store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc, at::Tensor k, at::Tensor v) -> void { - const auto max_tokens = k_cache.size(0); - const auto num_tokens = out_loc.size(0); - k_cache = k_cache.view({max_tokens, -1}); - v_cache = v_cache.view({max_tokens, -1}); - k = k.view({num_tokens, -1}); - v = v.view({num_tokens, -1}); - - TORCH_CHECK( - k_cache.is_cuda() && v_cache.is_cuda() && out_loc.is_cuda() && k.is_cuda() && v.is_cuda(), - "All tensors must be CUDA tensors"); - TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have the same size"); - TORCH_CHECK(k_cache.strides() == v_cache.strides(), "k_cache and v_cache must have the same strides"); - TORCH_CHECK(k.sizes() == v.sizes(), "k and v must have the same size"); - TORCH_CHECK(k.strides() == v.strides(), "k and v must have the same strides"); - TORCH_CHECK(k.stride(-1) == 1 && k_cache.stride(-1) == 1, "k and k_cache must be contiguous in head."); - TORCH_CHECK(k.size(-1) == k_cache.size(-1), "k and k_cache must have the same head size"); - TORCH_CHECK(out_loc.dim() == 1 && out_loc.is_contiguous(), "out_loc must be a 1D contiguous tensor"); - static_assert(sizeof(uint64_t) == 8, "uint64_t must be 8 bytes, our code assumes that"); - - const auto length = out_loc.size(0); - const auto elem_size = k.element_size(); - const auto size_bytes = elem_size * k.size(-1); - const auto kv_cache_stride_bytes = elem_size * k_cache.stride(-2); - const auto kv_input_stride_bytes = elem_size * k.stride(-2); - const auto kv_cache_stride = kv_cache_stride_bytes / 8; - const auto kv_input_stride = kv_input_stride_bytes / 8; - - const auto k_cache_ptr = static_cast(k_cache.data_ptr()); - const auto v_cache_ptr = static_cast(v_cache.data_ptr()); - const auto k_ptr = static_cast(k.data_ptr()); - const auto v_ptr = static_cast(v.data_ptr()); - const auto num_threads = 256; - const auto num_warps = num_threads / 32; - const auto num_blocks = (length + num_warps - 1) / num_warps; - const auto stream = at::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_INTEGRAL_TYPES(out_loc.scalar_type(), "store_kv_cache", [&] { - if constexpr (!std::is_same_v && !std::is_same_v) { - // do not instantiate the kernel if out_loc is not int32 or int64 - TORCH_CHECK(false, "out_loc must be of type int32 or int64, got: ", out_loc.scalar_type()); - } else { - if (size_bytes % 256 == 0) { - const auto items_per_warp = size_bytes / 256; - store_kv_cache_256x1<<>>( - k_cache_ptr, - v_cache_ptr, - out_loc.data_ptr(), - length, - k_ptr, - v_ptr, - kv_cache_stride, - kv_input_stride, - items_per_warp); - } else if (size_bytes % 128 == 0) { - const auto items_per_warp = size_bytes / 128; - store_kv_cache_128x2<<>>( - k_cache_ptr, - v_cache_ptr, - out_loc.data_ptr(), - length, - k_ptr, - v_ptr, - kv_cache_stride, - kv_input_stride, - items_per_warp); - } else { - TORCH_CHECK( - false, - "The last dimension size bytes of k and v must be" - " divisible by 128 at least, got: ", - size_bytes); - } - } - }); -} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 92493599f..8bb8f4684 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -135,20 +135,6 @@ void silu_and_mul(at::Tensor& out, at::Tensor& input); void gelu_tanh_and_mul(at::Tensor& out, at::Tensor& input); void gelu_and_mul(at::Tensor& out, at::Tensor& input); -void apply_rope_pos_ids_cos_sin_cache( - at::Tensor q, - at::Tensor k, - at::Tensor q_rope, - at::Tensor k_rope, - at::Tensor cos_sin_cache, - at::Tensor pos_ids, - bool interleave, - bool enable_pdl, - const std::optional& v, - const std::optional& k_buffer, - const std::optional& v_buffer, - const std::optional& kv_cache_loc); - void rotary_embedding( torch::Tensor& positions, torch::Tensor& query, @@ -239,7 +225,6 @@ void sgl_per_token_group_quant_8bit_v2( bool scale_ue8m0, bool fuse_silu_and_mul, const std::optional& masked_m); -void sgl_per_tensor_quant_fp8(at::Tensor input, at::Tensor output_q, at::Tensor output_s, bool is_static); void sgl_per_token_quant_fp8(at::Tensor input, at::Tensor output_q, at::Tensor output_s); void bmm_fp8( at::Tensor A, @@ -609,7 +594,6 @@ void transfer_kv_all_layer_direct_lf_pf( * From csrc/memory */ at::Tensor weak_ref_tensor(const at::Tensor& tensor); -void store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc, at::Tensor k, at::Tensor v); /* * From FlashInfer diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index e84ff12f5..ca56261dd 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -18,8 +18,6 @@ from sgl_kernel.attention import ( ) from sgl_kernel.cutlass_moe import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data from sgl_kernel.elementwise import ( - FusedSetKVBufferArg, - apply_rope_with_cos_sin_cache_inplace, concat_mla_absorb_q, concat_mla_k, copy_to_gpu_no_ce, @@ -41,7 +39,6 @@ from sgl_kernel.expert_specialization import ( from sgl_kernel.gemm import ( awq_dequantize, bmm_fp8, - cutlass_scaled_fp4_mm, dsv3_fused_a_gemm, dsv3_router_gemm, fp8_blockwise_scaled_mm, @@ -51,15 +48,11 @@ from sgl_kernel.gemm import ( int8_scaled_mm, qserve_w4a8_per_chn_gemm, qserve_w4a8_per_group_gemm, - scaled_fp4_grouped_quant, - scaled_fp4_quant, - sgl_per_tensor_quant_fp8, sgl_per_token_group_quant_8bit, sgl_per_token_group_quant_fp8, sgl_per_token_group_quant_int8, sgl_per_token_quant_fp8, shuffle_rows, - silu_and_mul_scaled_fp4_grouped_quant, ) from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda from sgl_kernel.kvcacheio import ( @@ -75,7 +68,7 @@ from sgl_kernel.mamba import ( causal_conv1d_update_cpu, chunk_gated_delta_rule_cpu, ) -from sgl_kernel.memory import set_kv_buffer_kernel, weak_ref_tensor +from sgl_kernel.memory import weak_ref_tensor from sgl_kernel.moe import ( apply_shuffle_mul_sum, fp8_blockwise_scaled_grouped_mm, diff --git a/sgl-kernel/python/sgl_kernel/elementwise.py b/sgl-kernel/python/sgl_kernel/elementwise.py index 704cbf492..1ed1ae474 100644 --- a/sgl-kernel/python/sgl_kernel/elementwise.py +++ b/sgl-kernel/python/sgl_kernel/elementwise.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from typing import Optional import torch @@ -332,121 +331,6 @@ if torch.version.hip is not None: return out -@dataclass -class FusedSetKVBufferArg: - """ - value : Optional[torch.Tensor] - Value tensor, shape: ``(nnz, num_v_heads * head_size)``. - k_buffer : Optional[torch.Tensor] - Buffer for keys, shape: ``(nnz, num_k_heads * head_size)``. - v_buffer : Optional[torch.Tensor] - Buffer for values, shape: ``(nnz, num_v_heads * head_size)``. - k_scale : Optional[float] - Scale factor for keys. - v_scale : Optional[float] - Scale factor for values. - cache_loc : Optional[torch.Tensor] - Cache location tensor, used for indexing kv cache. - """ - - value: torch.Tensor - k_buffer: torch.Tensor - v_buffer: torch.Tensor - k_scale: Optional[float] - v_scale: Optional[float] - cache_loc: torch.Tensor - - -def _view_3d(x, head_size): - return x.view(x.shape[0], -1, head_size) - - -def apply_rope_with_cos_sin_cache_inplace( - positions: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor, - head_size: int, - cos_sin_cache: torch.Tensor, - is_neox: bool = True, - fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, - enable_pdl: Optional[bool] = None, -) -> None: - r""" - Apply rotary embedding to keys and queries with precomputed cos/sin values. - This is designed to be compatible with the SGL/vLLM implementation. - The result is inplace applied to the input tensors. - - Parameters - ---------- - positions : torch.Tensor - Position indices, shape: ``(nnz)``. - query : torch.Tensor - Query tensor, shape: ``(nnz, num_q_heads * head_size)``. - key : torch.Tensor - Key tensor, shape: ``(nnz, num_k_heads * head_size)``. - cos_sin_cache : torch.Tensor - Cosine and Sine cache tensor, shape: ``(max_seq_len, rotary_dim)``. - Cosine is the first half and Sine is the second half on rotary_dim. - is_neox : bool - Whether to use Neox style RoPE, default: ``True``. - - * If ``True``, the last dimension of the query/key tensor is not interleaved, i.e., - we rotate the first half dimensions ``([..., :head_dim//2])`` and the second half - dimensions ``([..., head_dim//2:])``. - - * If ``False``, the last dimension of the query/key tensor is interleaved, i.e., - we rotate the even dimensions ``([..., ::2])`` and odd dimensions ``([..., 1::2])``. - fused_set_kv_buffer_arg : FusedSetKVBufferArg - Fuse the set-kv-buffer operation into this kernel - - Note - ---- - The rotary dimension is determined by the cosine cache and sine cache. - """ - if cos_sin_cache.dtype != torch.float32: - raise ValueError("cos_sin_cache should be float32") - - if enable_pdl is None: - # the non-fused branch does not yet support PDL, but after we switch to our impl for that branch it will - enable_pdl = is_arch_support_pdl() and (fused_set_kv_buffer_arg is not None) - - if (a := fused_set_kv_buffer_arg) is not None: - assert a.k_scale is None, "k_scale is not yet supported" - assert a.v_scale is None, "v_scale is not yet supported" - assert a.cache_loc.dtype == torch.int64, f"{a.cache_loc.dtype=}" - - torch.ops.sgl_kernel.apply_rope_pos_ids_cos_sin_cache.default( - _view_3d(query, head_size), - _view_3d(key, head_size), - _view_3d(query, head_size), - _view_3d(key, head_size), - cos_sin_cache, - positions.long(), - (not is_neox), - enable_pdl, - ( - _view_3d(fused_set_kv_buffer_arg.value, head_size) - if fused_set_kv_buffer_arg is not None - else None - ), - ( - _view_3d(fused_set_kv_buffer_arg.k_buffer, head_size) - if fused_set_kv_buffer_arg is not None - else None - ), - ( - _view_3d(fused_set_kv_buffer_arg.v_buffer, head_size) - if fused_set_kv_buffer_arg is not None - else None - ), - ( - fused_set_kv_buffer_arg.cache_loc - if fused_set_kv_buffer_arg is not None - else None - ), - ) - - def rotary_embedding( positions: torch.Tensor, query: torch.Tensor, diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index 0d2389c96..6d320aa9d 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/sgl-kernel/python/sgl_kernel/gemm.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple +from typing import Optional import torch from sgl_kernel.utils import _get_cache_buf @@ -140,17 +140,6 @@ sgl_per_token_group_quant_fp8 = sgl_per_token_group_quant_8bit sgl_per_token_group_quant_int8 = sgl_per_token_group_quant_8bit -def sgl_per_tensor_quant_fp8( - input: torch.Tensor, - output_q: torch.Tensor, - output_s: torch.Tensor, - is_static: bool, -) -> None: - torch.ops.sgl_kernel.sgl_per_tensor_quant_fp8.default( - input, output_q, output_s, is_static - ) - - def sgl_per_token_quant_fp8( input: torch.Tensor, output_q: torch.Tensor, @@ -159,54 +148,6 @@ def sgl_per_token_quant_fp8( torch.ops.sgl_kernel.sgl_per_token_quant_fp8.default(input, output_q, output_s) -def cutlass_scaled_fp4_mm( - a: torch.Tensor, - b: torch.Tensor, - block_scale_a: torch.Tensor, - block_scale_b: torch.Tensor, - alpha: torch.Tensor, - out_dtype: torch.dtype, -) -> torch.Tensor: - from sglang.jit_kernel.nvfp4 import ( - cutlass_scaled_fp4_mm as jit_cutlass_scaled_fp4_mm, - ) - - return jit_cutlass_scaled_fp4_mm( - a, - b, - block_scale_a, - block_scale_b, - alpha, - out_dtype, - ) - - -def scaled_fp4_quant( - input: torch.Tensor, input_global_scale: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Quantize input tensor to FP4 and return quantized tensor and scale. - - This function quantizes the last dimension of the given tensor `input`. For - every 16 consecutive elements, a single dynamically computed scaling factor - is shared. This scaling factor is quantized using the `input_global_scale` - and is stored in a swizzled layout (see - https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x). - - Args: - input: The input tensor to be quantized to FP4 - input_global_scale: A scalar scaling factor for the entire tensor. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every - two values are packed into a uint8 and float8_e4m3 scaling factors - in a sizzled layout. - """ - from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as jit_scaled_fp4_quant - - return jit_scaled_fp4_quant(input, input_global_scale) - - def qserve_w4a8_per_chn_gemm( in_feats: torch.Tensor, kernel: torch.Tensor, @@ -280,73 +221,6 @@ def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape): return output_tensor -def scaled_fp4_grouped_quant( - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - mask: torch.Tensor, -): - """ - Quantize input tensor to FP4 and return quantized tensor and scale, for - grouped gemm inputs (e.g., grouped_gemm_nt_masked for flashinfer). - Args: - input: The input tensor to be quantized to FP4, with shape (l, m, k) - l is number of groups, m is number of tokens per group, k is number of features. - input_global_scale: A scalar scaling factor for the entire tensor, with - shape (l,). - Outputs: - output: The quantized tensor in FP4, with shape (m, k // 2, l) but the physical - layout is (l, m, k // 2). `// 2` is because two fp4 values are packed into - an uint8. - output_scales: The blockscale tensor in FP8-E4M3, with shape (32, 4, rm, 4, rk, l) - but the physical layout is (l, rm, rk, 32, 4, 4). - Note: - For the shape of output_scales, `32 * 4 * rm` is a padded m to nearest multiple of 128. - `4 * rk` is a padded `k // 16` to nearest multiple of 4. These layout constants are - required by the NVIDIA Blackwell MMA operations. - """ - from sglang.jit_kernel.nvfp4 import ( - scaled_fp4_grouped_quant as jit_scaled_fp4_grouped_quant, - ) - - return jit_scaled_fp4_grouped_quant(input_tensor, input_global_scale, mask) - - -def silu_and_mul_scaled_fp4_grouped_quant( - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - mask: torch.Tensor, -): - """ - Quantize input tensor to FP4 and return quantized tensor and scale, for - grouped gemm inputs (e.g., grouped_gemm_nt_masked for flashinfer). - Args: - input: The input tensor to be quantized to FP4, with shape (l, m, k * 2) - l is number of groups, m is number of tokens per group, k is number of features. - input_global_scale: A scalar scaling factor for the entire tensor, with - shape (l,). - mask: The mask tensor, with shape (l,) - Outputs: - output: The quantized tensor in FP4, with shape (m, k // 2, l) but the physical - layout is (l, m, k // 2). `// 2` is because two fp4 values are packed into - an uint8. - output_scales: The blockscale tensor in FP8-E4M3, with shape (32, 4, rm, 4, rk, l) - but the physical layout is (l, rm, rk, 32, 4, 4). - Note: - For the shape of output_scales, `32 * 4 * rm` is a padded m to nearest multiple of 128. - `4 * rk` is a padded `k // 16` to nearest multiple of 4. These layout constants are - required by the NVIDIA Blackwell MMA operations. - """ - from sglang.jit_kernel.nvfp4 import ( - silu_and_mul_scaled_fp4_grouped_quant as jit_silu_and_mul_scaled_fp4_grouped_quant, - ) - - return jit_silu_and_mul_scaled_fp4_grouped_quant( - input_tensor, - input_global_scale, - mask, - ) - - # GPTQ kernels def gptq_gemm( a: torch.Tensor, diff --git a/sgl-kernel/python/sgl_kernel/memory.py b/sgl-kernel/python/sgl_kernel/memory.py index 9ff4f957d..c195cb4d7 100644 --- a/sgl-kernel/python/sgl_kernel/memory.py +++ b/sgl-kernel/python/sgl_kernel/memory.py @@ -1,23 +1,6 @@ import torch -def set_kv_buffer_kernel( - k_cache: torch.Tensor, - v_cache: torch.Tensor, - loc: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - fallback: bool = False, -): - try: - if fallback: - raise RuntimeError("Fallback to torch implementation") - torch.ops.sgl_kernel.store_kv_cache(k_cache, v_cache, loc, k, v) - except RuntimeError: # ok, fallback to torch implementation - k_cache[loc] = k - v_cache[loc] = v - - def weak_ref_tensor(tensor): return ( torch.ops.sgl_kernel.weak_ref_tensor(tensor) diff --git a/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py b/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py index 109778f29..6d319f843 100644 --- a/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py +++ b/sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py @@ -1,7 +1,31 @@ +from dataclasses import dataclass from typing import Optional, Tuple, Union import torch -from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace + +from sglang.jit_kernel.rope import FusedSetKVBufferArg as _JitFusedSetKVBufferArg +from sglang.jit_kernel.rope import ( + apply_rope_with_cos_sin_cache_inplace as _jit_apply_rope_with_cos_sin_cache_inplace, +) + + +@dataclass +class FusedSetKVBufferArg: + value: torch.Tensor + k_buffer: torch.Tensor + v_buffer: torch.Tensor + cache_loc: torch.Tensor + # Kept for backward compatibility with old sgl_kernel test/bench callsites. + k_scale: Optional[float] = None + v_scale: Optional[float] = None + + def to_jit(self) -> _JitFusedSetKVBufferArg: + return _JitFusedSetKVBufferArg( + value=self.value, + k_buffer=self.k_buffer, + v_buffer=self.v_buffer, + cache_loc=self.cache_loc, + ) # vLLM torch native @@ -129,14 +153,19 @@ class FlashInferRotaryEmbedding(RotaryEmbedding): fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - fused_set_kv_buffer_arg=fused_set_kv_buffer_arg, - head_size=self.head_size, + query_view = query.view(query.shape[0], -1, self.head_size) + key_view = key.view(key.shape[0], -1, self.head_size) + _jit_apply_rope_with_cos_sin_cache_inplace( + q=query_view, + k=key_view, cos_sin_cache=self.cos_sin_cache, + positions=positions, is_neox=self.is_neox_style, + fused_args=( + fused_set_kv_buffer_arg.to_jit() + if fused_set_kv_buffer_arg is not None + else None + ), ) return query, key diff --git a/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py b/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py index 7acba566c..73274165f 100644 --- a/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py +++ b/sgl-kernel/tests/test_cutlass_w4a8_moe_mm.py @@ -1,8 +1,10 @@ import pytest import torch -from sgl_kernel import cutlass_w4a8_moe_mm, sgl_per_tensor_quant_fp8 +from sgl_kernel import cutlass_w4a8_moe_mm from utils import is_hopper +from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 + def pack_int4_values_to_int8(int4_values_interleaved: torch.Tensor) -> torch.Tensor: if int4_values_interleaved.shape[-1] % 2 != 0: @@ -148,7 +150,7 @@ def _per_tensor_quant_fp8( device=x.device, dtype=torch.float32, ) - sgl_per_tensor_quant_fp8(x, x_q, x_s, is_static=False) + per_tensor_quant_fp8(x, x_q, x_s, is_static=False) return x_q, x_s diff --git a/sgl-kernel/tests/test_fp4_gemm.py b/sgl-kernel/tests/test_fp4_gemm.py deleted file mode 100644 index 47401618b..000000000 --- a/sgl-kernel/tests/test_fp4_gemm.py +++ /dev/null @@ -1,154 +0,0 @@ -import pytest -import torch -from sgl_kernel import cutlass_scaled_fp4_mm, scaled_fp4_quant - -skip_condition = torch.cuda.get_device_capability() < (10, 0) - -DTYPES = [torch.float16, torch.bfloat16] -# m, n, k -SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)] -PAD_SHAPES = [(150, 128, 64), (128, 128, 96)] -SHAPES.extend(PAD_SHAPES) - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max - -kE2M1ToFloatArray = [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, -] - - -def e2m1_to_fp32(int4_value): - signBit = int4_value & 0x8 - int4_absValue = int4_value & 0x7 - float_result = kE2M1ToFloatArray[int4_absValue] - if signBit: - float_result = -float_result - return float_result - - -def break_fp4_bytes(a, dtype): - assert a.dtype == torch.uint8 - m, n = a.shape - a = a.flatten() - # Get upper 4 bits - highHalfByte = (a & 0xF0) >> 4 - # Get lower 4 bits - lowHalfByte = a & 0x0F - fH = torch.tensor([e2m1_to_fp32(x) for x in highHalfByte]).to(a.device) - fL = torch.tensor([e2m1_to_fp32(x) for x in lowHalfByte]).to(a.device) - # [0xAB, 0xCD] -> [0xB, 0xA, 0xD, 0xC] - out = torch.stack((fL, fH), dim=-1).reshape(m, n * 2) - return out - - -def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size): - sf_m, sf_k = a_sf_swizzled.shape - m_tiles = (m + 128 - 1) // 128 - f = block_size * 4 - k_tiles = (k + f - 1) // f - tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4)) - tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) - out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size) - return out[0:m, 0:k] - - -def dequantize_to_dtype( - tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16 -): - """Dequantize the fp4 tensor back to high precision.""" - # Two fp4 values are packed into one uint8. - assert tensor_fp4.dtype == torch.uint8 - m, packed_k = tensor_fp4.shape - k = packed_k * 2 - tensor_f32 = break_fp4_bytes(tensor_fp4, dtype) - tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size) - tensor_sf = tensor_sf.view(torch.float8_e4m3fn) - tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size) - tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale - - # scale the tensor - out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) - return out - - -def get_ref_results( - a_fp4, - b_fp4, - a_sf, - b_sf, - a_global_scale, - b_global_scale, - m, - n, - dtype, - block_size, - device, -): - _, m_k = a_fp4.shape - _, n_k = b_fp4.shape - assert m_k == n_k - a_in_dtype = dequantize_to_dtype( - a_fp4, a_sf, a_global_scale, dtype=dtype, device=device, block_size=block_size - ) - b_in_dtype = dequantize_to_dtype( - b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size - ) - return torch.matmul(a_in_dtype, b_in_dtype.t()) - - -@pytest.mark.skipif( - skip_condition, reason="Nvfp4 Requires compute capability of 10 or above." -) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("shape", SHAPES) -@torch.inference_mode() -def test_nvfp4_gemm( - dtype: torch.dtype, - shape: tuple[int, int], -) -> None: - m, n, packed_k = shape - k = packed_k * 2 - block_size = 16 - a_dtype = torch.randn((m, k), dtype=dtype, device="cuda") - b_dtype = torch.randn((n, k), dtype=dtype, device="cuda") - - a_global_scale = ( - (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1) - ).to(torch.float32) - b_global_scale = ( - (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1) - ).to(torch.float32) - alpha = 1.0 / (a_global_scale * b_global_scale) - a_fp4, a_scale_interleaved = scaled_fp4_quant(a_dtype, a_global_scale) - b_fp4, b_scale_interleaved = scaled_fp4_quant(b_dtype, b_global_scale) - - expected_out = get_ref_results( - a_fp4, - b_fp4, - a_scale_interleaved, - b_scale_interleaved, - a_global_scale, - b_global_scale, - m, - n, - dtype, - block_size, - "cuda", - ) - out = cutlass_scaled_fp4_mm( - a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype - ) - - torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/sgl-kernel/tests/test_fp4_quantize.py b/sgl-kernel/tests/test_fp4_quantize.py deleted file mode 100644 index e29bac211..000000000 --- a/sgl-kernel/tests/test_fp4_quantize.py +++ /dev/null @@ -1,260 +0,0 @@ -import pytest -import torch -from flashinfer import ( - scaled_fp4_grouped_quantize, - silu_and_mul_scaled_nvfp4_experts_quantize, -) -from sgl_kernel import scaled_fp4_quant, silu_and_mul - -skip_condition = torch.cuda.get_device_capability() < (10, 0) - -DTYPES = [torch.float16, torch.bfloat16] -SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)] -PAD_SHAPES = [ - (90, 64), - (150, 64), - (128, 48), - (128, 80), - (150, 80), - (90, 48), - (90, 128), - (150, 128), - (150, 48), - (90, 80), -] - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max - -# E2M1 to float -# 0111 -> 6 -# 0110 -> 4 -# 0101 -> 3 -# 0100 -> 2 -# 0011 -> 1.5 -# 0010 -> 1 -# 0001 -> 0.5 -# 0000 -> 0 -E2M1_TO_FLOAT32 = [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - 0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, -] -BLOCK_SIZE = 16 - - -def cast_from_fp4(x, m, n): - # The fp4 values are packed in uint8 as [v_1st | v_2nd] - v_2nd = x & 0xF - v_1st = (x >> 4) & 0xF - c = torch.stack((v_2nd, v_1st), dim=-1) - out = torch.tensor([E2M1_TO_FLOAT32[x] for x in c.flatten()]) - out = out.reshape(m, n).to(torch.float32) - return out - - -def cast_to_fp4(x): - sign = torch.sign(x) - x = torch.abs(x) - x[(x >= 0.0) & (x <= 0.25)] = 0.0 - x[(x > 0.25) & (x < 0.75)] = 0.5 - x[(x >= 0.75) & (x <= 1.25)] = 1.0 - x[(x > 1.25) & (x < 1.75)] = 1.5 - x[(x >= 1.75) & (x <= 2.5)] = 2.0 - x[(x > 2.5) & (x < 3.5)] = 3.0 - x[(x >= 3.5) & (x <= 5.0)] = 4.0 - x[x > 5.0] = 6.0 - return x * sign - - -def get_reciprocal(x): - if isinstance(x, torch.Tensor): - return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x) - elif isinstance(x, (float, int)): - return 0.0 if x == 0 else 1.0 / x - else: - raise TypeError("Input must be a float, int, or a torch.Tensor.") - - -def ref_nvfp4_quant(x, global_scale): - assert global_scale.dtype == torch.float32 - assert x.ndim == 2 - m, n = x.shape - x = torch.reshape(x, (m, n // BLOCK_SIZE, BLOCK_SIZE)) - vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) - scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX)) - scale = scale.to(torch.float8_e4m3fn).to(torch.float32) - output_scale = get_reciprocal(scale * get_reciprocal(global_scale)) - - scaled_x = x.to(torch.float32) * output_scale - clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n) - return cast_to_fp4(clipped_x), scale.squeeze(-1) - - -def recover_swizzled_scales(scale, m, n): - rounded_m = ((m + 128 - 1) // 128) * 128 - scale_n = n // BLOCK_SIZE - rounded_n = ((scale_n + 4 - 1) // 4) * 4 - # Recover the swizzled scaling factor to linear layout - tmp = torch.reshape(scale, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4)) - tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) - result = torch.reshape(tmp, (rounded_m, rounded_n)).to(torch.float32) - return result[:m, :scale_n] - - -@pytest.mark.skipif( - skip_condition, reason="Nvfp4 Requires compute capability of 10 or above." -) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("shape", SHAPES) -@torch.inference_mode() -def test_quantize_to_fp4( - dtype: torch.dtype, - shape: tuple[int, int], -) -> None: - torch.manual_seed(42) - torch.set_default_device("cuda:0") - - m, n = shape - - x = torch.randn((m, n), dtype=dtype) - tensor_amax = torch.abs(x).max().to(torch.float32) - global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax - out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) - - out, out_scale = scaled_fp4_quant(x, global_scale) - scale_ans = recover_swizzled_scales(out_scale, m, n) - out_ans = cast_from_fp4(out, m, n) - - torch.testing.assert_close(out_ans, out_ref) - torch.testing.assert_close(scale_ans, scale_ref) - - -@pytest.mark.skipif( - skip_condition, reason="Nvfp4 Requires compute capability of 10 or above." -) -@pytest.mark.parametrize("pad_shape", PAD_SHAPES) -@torch.inference_mode() -def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None: - torch.manual_seed(42) - dtype = torch.float16 - torch.set_default_device("cuda:0") - - m, n = pad_shape - - x = torch.randn((m, n), dtype=dtype) - - tensor_amax = torch.abs(x).max().to(torch.float32) - global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax - out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) - - out, out_scale = scaled_fp4_quant(x, global_scale) - - scale_ans = recover_swizzled_scales(out_scale, m, n) - out_ans = cast_from_fp4(out, m, n) - - torch.testing.assert_close(out_ans, out_ref) - torch.testing.assert_close(scale_ans, scale_ref) - - -@pytest.mark.skipif( - skip_condition, reason="Nvfp4 Requires compute capability of 10 or above." -) -@pytest.mark.parametrize("shape", [(2, 512, 2048), (2, 100, 128), (2, 128, 96)]) -def test_quantize_to_fp4_grouped(shape): - torch.manual_seed(42) - torch.set_default_device("cuda:0") - - l, m, k = shape - x = torch.randn((l, m, k), dtype=torch.bfloat16) - max_m = m // 2 - assert max_m <= m - mask = torch.randint(1, max_m, (l,), dtype=torch.int32) - tensor_amax = x.abs().amax(dim=(1, 2)).to(torch.float32) - x_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax - output, output_scales = scaled_fp4_grouped_quantize( - x, - mask, - x_sf_global, - ) - # output in logical (m, k, l), but its physical layout is (l, m, k). - # So permute first to (l, m, k). - output = output.permute(2, 0, 1) - # output_scale in logical (32, 4, rm, 4, rk, l), but its physical layout is (l, rm, rk, 32, 4, 4). - # So permute first to (l, rm, rk, 32, 4, 4). - padded_m = ((m + 128 - 1) // 128) * 128 - output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1) - for i in range(l): - a_fp4, a_scale_interleaved = scaled_fp4_quant(x[i], x_sf_global[i]) - torch.testing.assert_close(a_fp4[: mask[i]], output[i][: mask[i]]) - # Recover swizzled scales to linear layout and drop padded values, so - # no extra checks on padding are needed. - scale_ref = recover_swizzled_scales(a_scale_interleaved, m, k) - scale_ans = recover_swizzled_scales(output_scales[i], m, k) - torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]]) - - -@pytest.mark.skipif( - skip_condition, reason="Nvfp4 Requires compute capability of 10 or above." -) -@pytest.mark.parametrize("shape", [(32, 100, 2048), (32, 512, 2048), (6, 6144, 2048)]) -def test_silu_and_mul_quantize_to_fp4_grouped(shape): - torch.manual_seed(42) - torch.set_default_device("cuda:0") - - l, m, k = shape - x = torch.randn((l, m, k * 2), dtype=torch.bfloat16) - max_m = m // 2 - assert max_m <= m - mask = torch.randint(1, max_m, (l,), dtype=torch.int32) - - ref_y = silu_and_mul(x) - tensor_amax = ref_y.abs().amax(dim=(1, 2)).to(torch.float32) - y_sf_global = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax - ref_output, ref_output_scales = scaled_fp4_grouped_quantize( - ref_y, - mask, - y_sf_global, - ) - output, output_scales = silu_and_mul_scaled_nvfp4_experts_quantize( - x, - mask, - y_sf_global, - ) - - # output in logical (m, k, l), but its physical layout is (l, m, k). - # So permute first to (l, m, k). - output = output.permute(2, 0, 1) - ref_output = ref_output.permute(2, 0, 1) - - # output_scale in logical (32, 4, rm, 4, rk, l), but its physical layout is (l, rm, rk, 32, 4, 4). - # So permute first to (l, rm, rk, 32, 4, 4). - padded_m = ((m + 128 - 1) // 128) * 128 - output_scales = output_scales.permute(5, 2, 4, 0, 1, 3).view(l, padded_m, -1) - ref_output_scales = ref_output_scales.permute(5, 2, 4, 0, 1, 3).view( - l, padded_m, -1 - ) - - for i in range(l): - torch.testing.assert_close(ref_output[i, : mask[i]], output[i, : mask[i]]) - # We need to recover the swizzled scales to linear layout before applying mask slice. - scale_ref = recover_swizzled_scales(ref_output_scales[i], m, k) - scale_ans = recover_swizzled_scales(output_scales[i], m, k) - torch.testing.assert_close(scale_ref[: mask[i]], scale_ans[: mask[i]]) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/sgl-kernel/tests/test_per_tensor_quant_fp8.py b/sgl-kernel/tests/test_per_tensor_quant_fp8.py deleted file mode 100644 index 0840f298f..000000000 --- a/sgl-kernel/tests/test_per_tensor_quant_fp8.py +++ /dev/null @@ -1,67 +0,0 @@ -import itertools -from typing import Optional, Tuple - -import pytest -import torch -from sgl_kernel import sgl_per_tensor_quant_fp8 - -from sglang.srt.utils import is_hip - -_is_hip = is_hip() -fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn - - -def sglang_scaled_fp8_quant( - input: torch.Tensor, - scale: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - fp8_type_: torch.dtype = torch.float8_e4m3fn - output = torch.empty_like(input, device=input.device, dtype=fp8_type_) - is_static = True - if scale is None: - scale = torch.zeros(1, device=input.device, dtype=torch.float32) - is_static = False - sgl_per_tensor_quant_fp8(input, output, scale, is_static) - - return output, scale - - -def torch_scaled_fp8_quant(tensor, inv_scale): - # The reference implementation that fully aligns to - # the kernel being tested. - finfo = torch.finfo(torch.float8_e4m3fn) - scale = inv_scale.reciprocal() - qweight = (tensor.to(torch.float32) * scale).clamp(min=finfo.min, max=finfo.max) - qweight = qweight.to(torch.float8_e4m3fn) - return qweight - - -@pytest.mark.parametrize( - "num_tokens,hidden_dim", - list(itertools.product([128, 256, 512], [512, 2048, 4096])), -) -def test_per_tensor_quant_compare_implementations( - num_tokens: int, - hidden_dim: int, -): - device = torch.device("cuda") - x = torch.rand((num_tokens, hidden_dim), dtype=torch.float16, device=device) - - sglang_out, sglang_scale = sglang_scaled_fp8_quant(x) - torch_out = torch_scaled_fp8_quant(x, sglang_scale) - - torch.testing.assert_close( - sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3 - ) - - scale = torch.rand(1, dtype=torch.float32, device=device) - sglang_out, sglang_scale = sglang_scaled_fp8_quant(x, scale) - torch_out = torch_scaled_fp8_quant(x, scale) - - torch.testing.assert_close( - sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3 - ) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/sgl-kernel/tests/test_rotary_embedding.py b/sgl-kernel/tests/test_rotary_embedding.py deleted file mode 100644 index ed2cac32f..000000000 --- a/sgl-kernel/tests/test_rotary_embedding.py +++ /dev/null @@ -1,167 +0,0 @@ -from typing import Any, Dict, List, Optional, Tuple, Union - -import pytest -import torch -from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace -from sgl_kernel.testing.rotary_embedding import ( - FlashInferRotaryEmbedding, - MHATokenToKVPool, - RotaryEmbedding, - SglKernelRotaryEmbedding, - create_inputs, -) - - -@pytest.mark.parametrize( - "head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype, device, batch_size, seq_len, num_q_heads, num_kv_heads, save_kv_cache", - [ - # GPT-OSS cases - *[ - ( - 64, - 64, - 4096, - 8000, - True, - torch.bfloat16, - "cuda", - batch_size, - seq_len, - 64, - 8, - save_kv_cache, - ) - for batch_size, seq_len in ( - (1, 1), - (32, 1), - (128, 1), - (512, 1), - (2, 512), - (4, 4096), - ) - for save_kv_cache in (False, True) - ], - # Other cases - (64, 64, 32, 8000, True, torch.bfloat16, "cuda", 32, 32, 1, 1, False), - (256, 128, 4096, 10000, True, torch.bfloat16, "cuda", 2, 512, 4, 2, False), - (512, 128, 311, 10000, True, torch.bfloat16, "cuda", 3, 39, 4, 2, False), - (128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 32, 8, False), - (128, 128, 2048, 10000, False, torch.bfloat16, "cuda", 2, 512, 16, 4, False), - (512, 128, 311, 10000, False, torch.bfloat16, "cuda", 3, 39, 4, 2, False), - (64, 64, 32, 8000, True, torch.float32, "cuda", 32, 32, 1, 1, False), - (256, 128, 4096, 10000, True, torch.float32, "cuda", 2, 512, 4, 2, False), - (512, 128, 311, 10000, True, torch.float32, "cuda", 3, 39, 4, 2, False), - (128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 32, 8, False), - (128, 128, 2048, 10000, False, torch.float32, "cuda", 2, 512, 16, 4, False), - (512, 128, 311, 10000, False, torch.float32, "cuda", 3, 39, 4, 2, False), - ], -) -def test_correctness( - head_size: int, - rotary_dim: int, - max_position_embeddings: int, - base: int, - is_neox_style: bool, - dtype: torch.dtype, - device: str, - batch_size: int, - seq_len: int, - num_q_heads: int, - num_kv_heads: int, - save_kv_cache: bool, -): - config = dict( - head_size=head_size, - rotary_dim=rotary_dim, - max_position_embeddings=max_position_embeddings, - base=base, - is_neox_style=is_neox_style, - dtype=dtype, - ) - - rope_ref = RotaryEmbedding(**config).to(device) - rope_flashinfer = FlashInferRotaryEmbedding(**config).to(device) - rope_sglkernel = SglKernelRotaryEmbedding(**config).to(device) - inputs = create_inputs( - head_size=head_size, - batch_size=batch_size, - seq_len=seq_len, - device=device, - dtype=dtype, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - ) - - if save_kv_cache: - pool_ref_for_flashinfer = MHATokenToKVPool( - head_num=num_kv_heads, head_dim=head_size - ) - pool_flashinfer = MHATokenToKVPool(head_num=num_kv_heads, head_dim=head_size) - - query_ref, key_ref = inputs["query"].clone(), inputs["key"].clone() - query_flashinfer, key_flashinfer = inputs["query"].clone(), inputs["key"].clone() - query_sglkernel, key_sglkernel = inputs["query"].clone(), inputs["key"].clone() - - # This is to align with the flashinfer implementation, flashinfer uses float32 cos/sin cache - query_ref_for_flashinfer_out, key_ref_for_flashinfer_out = rope_ref.forward_native( - inputs["pos_ids"], query_ref.to(torch.float32), key_ref.to(torch.float32) - ) - - query_ref_for_sglkernel_out, key_ref_for_sglkernel_out = rope_ref.forward_native( - inputs["pos_ids"], query_ref, key_ref - ) - if save_kv_cache: - pool_ref_for_flashinfer.set_kv_buffer( - loc=inputs["out_cache_loc"], - cache_k=key_ref_for_flashinfer_out.view(-1, num_kv_heads, head_size), - cache_v=inputs["value"].view(-1, num_kv_heads, head_size), - ) - - query_flashinfer_out, key_flashinfer_out = rope_flashinfer.forward_cuda( - inputs["pos_ids"], - query_flashinfer, - key_flashinfer, - fused_set_kv_buffer_arg=( - FusedSetKVBufferArg( - value=inputs["value"], - k_buffer=pool_flashinfer.k_buffer[0].view(-1, num_kv_heads * head_size), - v_buffer=pool_flashinfer.v_buffer[0].view(-1, num_kv_heads * head_size), - k_scale=None, - v_scale=None, - cache_loc=inputs["out_cache_loc"], - ) - if save_kv_cache - else None - ), - ) - - query_sglkernel_out, key_sglkernel_out = rope_sglkernel.forward_cuda( - inputs["pos_ids"], - query_sglkernel, - key_sglkernel, - ) - - torch.testing.assert_close( - query_ref_for_flashinfer_out, query_flashinfer_out, atol=1e-2, rtol=1e-2 - ) - torch.testing.assert_close( - key_ref_for_flashinfer_out, key_flashinfer_out, atol=1e-2, rtol=1e-2 - ) - torch.testing.assert_close( - query_ref_for_sglkernel_out, query_sglkernel_out, atol=1e-2, rtol=1e-2 - ) - torch.testing.assert_close( - key_ref_for_sglkernel_out, key_sglkernel_out, atol=1e-2, rtol=1e-2 - ) - if save_kv_cache: - for field in ["k_buffer", "v_buffer"]: - x_ref = getattr(pool_ref_for_flashinfer, field)[0] - x_flashinfer = getattr(pool_flashinfer, field)[0] - torch.testing.assert_close(x_ref, x_flashinfer, atol=1e-2, rtol=1e-2) - nonzero_ref = x_ref != 0 - nonzero_flashinfer = x_ref != 0 - assert torch.all(nonzero_ref == nonzero_flashinfer) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/test/registered/kernels/test_fp4_moe.py b/test/registered/kernels/test_fp4_moe.py index a72e1e4b8..369741438 100644 --- a/test/registered/kernels/test_fp4_moe.py +++ b/test/registered/kernels/test_fp4_moe.py @@ -5,9 +5,10 @@ import pytest import torch from flashinfer import fp4_quantize, scaled_fp4_grouped_quantize from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe -from sgl_kernel import scaled_fp4_quant, silu_and_mul +from sgl_kernel import silu_and_mul from torch.nn import functional as F +from sglang.jit_kernel.nvfp4 import scaled_fp4_quant from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4 from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType from sglang.srt.layers.moe.topk import TopKConfig, select_experts diff --git a/test/registered/moe/test_cutedsl_moe.py b/test/registered/moe/test_cutedsl_moe.py index 90f1b7c7a..3846334d0 100644 --- a/test/registered/moe/test_cutedsl_moe.py +++ b/test/registered/moe/test_cutedsl_moe.py @@ -4,9 +4,9 @@ from typing import Callable import torch from flashinfer import fp4_quantize, scaled_fp4_grouped_quantize -from sgl_kernel import scaled_fp4_quant from torch.nn import functional as F +from sglang.jit_kernel.nvfp4 import scaled_fp4_quant from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.moe.flashinfer_cutedsl_moe import flashinfer_cutedsl_moe_masked from sglang.srt.layers.moe.topk import TopKConfig, select_experts