[kernel slimming] Clean many useless sgl-kernel deprecated kernels (#20277)
This commit is contained in:
Xiaoyu Zhang
@@ -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
|
||||
|
||||
@@ -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](
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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 <ATen/cuda/Exceptions.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#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<at::Tensor>& v,
|
||||
const std::optional<at::Tensor>& k_buffer,
|
||||
const std::optional<at::Tensor>& v_buffer,
|
||||
const std::optional<at::Tensor>& 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<int64_t*>(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<c_type*>(q.data_ptr()),
|
||||
static_cast<c_type*>(k.data_ptr()),
|
||||
save_kv_cache ? static_cast<c_type*>(v->data_ptr()) : nullptr,
|
||||
static_cast<c_type*>(q_rope.data_ptr()),
|
||||
static_cast<c_type*>(k_rope.data_ptr()),
|
||||
save_kv_cache ? static_cast<c_type*>(k_buffer->data_ptr()) : nullptr,
|
||||
save_kv_cache ? static_cast<c_type*>(v_buffer->data_ptr()) : nullptr,
|
||||
static_cast<float*>(cos_sin_cache.data_ptr()),
|
||||
static_cast<int64_t*>(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<c_type*>(q.data_ptr()),
|
||||
static_cast<c_type*>(k.data_ptr()),
|
||||
static_cast<c_type*>(q_rope.data_ptr()),
|
||||
static_cast<c_type*>(k_rope.data_ptr()),
|
||||
static_cast<float*>(cos_sin_cache.data_ptr()),
|
||||
static_cast<int64_t*>(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;
|
||||
});
|
||||
}
|
||||
@@ -1,123 +0,0 @@
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/util/Float8_e4m3fn.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <cub/block/block_reduce.cuh>
|
||||
#include <flashinfer/vec_dtypes.cuh>
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
template <typename T>
|
||||
__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<T, vec_size>;
|
||||
|
||||
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<float>(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<float>(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 <typename T, typename DST_DTYPE>
|
||||
__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<T, VEC_SIZE>;
|
||||
|
||||
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<float>(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<DST_DTYPE>(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<float>(input[idx]) * scale_val, FP8_E4M3_MAX));
|
||||
#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
|
||||
output[idx] = static_cast<DST_DTYPE>(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<scalar_t><<<grid, block, 0, stream>>>(
|
||||
static_cast<scalar_t*>(input.data_ptr()), static_cast<float*>(output_s.data_ptr()), num_elements);
|
||||
}
|
||||
|
||||
per_tensor_quant_fp8_kernel<scalar_t, __nv_fp8_e4m3><<<grid, block, 0, stream>>>(
|
||||
static_cast<scalar_t*>(input.data_ptr()),
|
||||
static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
|
||||
static_cast<float*>(output_s.data_ptr()),
|
||||
num_elements);
|
||||
return true;
|
||||
});
|
||||
}
|
||||
@@ -1,147 +0,0 @@
|
||||
#include <ATen/Dispatch.h>
|
||||
#include <ATen/core/TensorBody.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <c10/util/Exception.h>
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
|
||||
namespace {
|
||||
|
||||
using std::size_t;
|
||||
using std::uint64_t;
|
||||
|
||||
// Each warp will process 256 bytes per loop iteration
|
||||
template <typename T>
|
||||
__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 <typename T>
|
||||
__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<uint64_t*>(k_cache.data_ptr());
|
||||
const auto v_cache_ptr = static_cast<uint64_t*>(v_cache.data_ptr());
|
||||
const auto k_ptr = static_cast<const uint64_t*>(k.data_ptr());
|
||||
const auto v_ptr = static_cast<const uint64_t*>(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<scalar_t, int32_t> && !std::is_same_v<scalar_t, int64_t>) {
|
||||
// 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<<<num_blocks, num_threads, 0, stream>>>(
|
||||
k_cache_ptr,
|
||||
v_cache_ptr,
|
||||
out_loc.data_ptr<scalar_t>(),
|
||||
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<<<num_blocks, num_threads, 0, stream>>>(
|
||||
k_cache_ptr,
|
||||
v_cache_ptr,
|
||||
out_loc.data_ptr<scalar_t>(),
|
||||
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);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -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<at::Tensor>& v,
|
||||
const std::optional<at::Tensor>& k_buffer,
|
||||
const std::optional<at::Tensor>& v_buffer,
|
||||
const std::optional<at::Tensor>& 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<torch::Tensor>& 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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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__])
|
||||
@@ -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__])
|
||||
@@ -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__])
|
||||
@@ -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__])
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user