diff --git a/benchmark/kernels/flashinfer_allreduce_fusion/README.md b/benchmark/kernels/flashinfer_allreduce_fusion/README.md index e651604c7..f1379cb19 100644 --- a/benchmark/kernels/flashinfer_allreduce_fusion/README.md +++ b/benchmark/kernels/flashinfer_allreduce_fusion/README.md @@ -97,6 +97,6 @@ If `--output-file` is specified, all configurations will be summarized in Markdo - The fused operator internally uses "oneshot"/"twoshot" two trigger methods; oneshot is enabled by default and twoshot is tested simultaneously. - FP8/FP4: - FP8 uses sglang's FP8 tools and dtype, with underlying platform selection of `e4m3`/`e4m3fnuz` etc. - - FP4 uses sgl-kernel's `scaled_fp4_quant`, requiring corresponding platform support. + - FP4 uses sglang's `fp4_quantize` (FlashInfer-backed), requiring corresponding platform support. - CUDA Graph: - Uses sglang's `graph_capture()` to prepare capture-ready state for communication, then uses `torch.cuda.graph` to capture kernels, reducing measurement jitter. diff --git a/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py index d1c3f545c..2905a62a1 100644 --- a/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py +++ b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py @@ -43,11 +43,11 @@ try: from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM from sgl_kernel import rmsnorm as SGL_RMS_NORM - from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as SGL_SCALED_FP4_QUANT + from sglang.srt.layers.quantization.fp4_utils import fp4_quantize as SGL_FP4_QUANT except Exception: # pragma: no cover - fallback on non-supported platforms SGL_FUSED_ADD_RMS_NORM = None SGL_RMS_NORM = None - SGL_SCALED_FP4_QUANT = None + SGL_FP4_QUANT = None FP8_DTYPE = SGLANG_FP8_DTYPE @@ -386,9 +386,9 @@ def standard_allreduce_rmsnorm_fp4_quant( residual_out = allreduce_out # Finally FP4 quantization - if SGL_SCALED_FP4_QUANT is None: - raise RuntimeError("scaled_fp4_quant is not available on this platform") - quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale) + if SGL_FP4_QUANT is None: + raise RuntimeError("fp4_quantize is not available on this platform") + quant_res, output_scale_res = SGL_FP4_QUANT(quant_input, input_global_scale) if residual is not None: return quant_res, residual_out, output_scale_res else: @@ -464,9 +464,9 @@ def standard_allreduce_rmsnorm_fp4_quant_native( residual_out = allreduce_out # Apply FP4 quantization (still using fused CUDA op as there's no native FP4) - if SGL_SCALED_FP4_QUANT is None: - raise RuntimeError("scaled_fp4_quant is not available on this platform") - quant_res, output_scale_res = SGL_SCALED_FP4_QUANT(quant_input, input_global_scale) + if SGL_FP4_QUANT is None: + raise RuntimeError("fp4_quantize is not available on this platform") + quant_res, output_scale_res = SGL_FP4_QUANT(quant_input, input_global_scale) if residual is not None: return quant_res, residual_out, output_scale_res diff --git a/benchmark/kernels/quantization/bench_fp4_quant.py b/benchmark/kernels/quantization/bench_fp4_quant.py deleted file mode 100644 index 0d5b54aeb..000000000 --- a/benchmark/kernels/quantization/bench_fp4_quant.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Benchmark FP4 quantize: sglang jit_kernel vs flashinfer. - -Compares ``sglang.jit_kernel.nvfp4.scaled_fp4_quant`` against -``flashinfer.fp4_quantize`` over a sweep of (M, K) shapes. - -Timing uses ``flashinfer.testing.bench_gpu_time`` (CUDA-graph based with -rotating-buffer cold-L2). -""" - -import argparse -import itertools - -import numpy as np -import torch -from flashinfer import fp4_quantize as flashinfer_fp4_quantize -from flashinfer.testing import bench_gpu_time - -from sglang.jit_kernel.nvfp4 import scaled_fp4_quant - -Ms = [1, 8, 32, 128, 512, 1024, 2048, 4096, 8192, 16384, 32768] -Ks = [128, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 8192, 16384] - - -def _bench(fn, input_args) -> float: - times = bench_gpu_time( - fn=fn, - input_args=input_args, - use_cuda_graph=True, - dry_run_time_ms=25, - repeat_time_ms=100, - ) - return float(np.median(times)) - - -def benchmark(M: int, K: int, dtype: torch.dtype, device: str): - x = torch.randn(M, K, device=device, dtype=dtype) - global_scale = torch.ones(1, device=device, dtype=torch.float32) - - sglang_ms = _bench( - lambda x, gs: scaled_fp4_quant(x, gs), - input_args=(x, global_scale), - ) - flashinfer_ms = _bench( - lambda x, gs: flashinfer_fp4_quantize(x, gs, backend="cute-dsl"), - input_args=(x, global_scale), - ) - - return sglang_ms, flashinfer_ms - - -def plot_speedup(rows, path): - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - - Ms_unique = sorted({int(r[0]) for r in rows}) - Ks_unique = sorted({int(r[1]) for r in rows}) - grid = np.full((len(Ms_unique), len(Ks_unique)), np.nan) - m_idx = {m: i for i, m in enumerate(Ms_unique)} - k_idx = {k: i for i, k in enumerate(Ks_unique)} - for M, K, _, _, sp in rows: - grid[m_idx[int(M)], k_idx[int(K)]] = float(sp) - - fig, ax = plt.subplots(figsize=(12, 8)) - vmax = max(2.0, np.nanmax(grid)) - vmin = min(0.5, np.nanmin(grid)) - im = ax.imshow( - grid, - aspect="auto", - cmap="RdYlGn", - vmin=vmin, - vmax=vmax, - origin="lower", - ) - ax.set_xticks(range(len(Ks_unique))) - ax.set_xticklabels(Ks_unique, rotation=45) - ax.set_yticks(range(len(Ms_unique))) - ax.set_yticklabels(Ms_unique) - ax.set_xlabel("K") - ax.set_ylabel("M") - ax.set_title("Speedup: flashinfer / sglang (>1 means sglang faster)") - for i in range(len(Ms_unique)): - for j in range(len(Ks_unique)): - v = grid[i, j] - if np.isfinite(v): - ax.text(j, i, f"{v:.2f}", ha="center", va="center", fontsize=7) - fig.colorbar(im, ax=ax, label="speedup") - fig.tight_layout() - fig.savefig(path, dpi=130) - print(f"Saved plot to {path}") - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") - parser.add_argument("--device", default="cuda") - parser.add_argument("--csv", type=str, default=None) - parser.add_argument("--plot", type=str, default=None) - args = parser.parse_args() - - dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 - - rows = [] - header = ( - f"{'M':>8} {'K':>8} {'sglang(us)':>12} {'flashinfer(us)':>16} {'speedup':>10}" - ) - print(header) - print("-" * len(header)) - - for M, K in itertools.product(Ms, Ks): - try: - sglang_ms, flashinfer_ms = benchmark(M, K, dtype, args.device) - except Exception as e: - print(f"{M:>8} {K:>8} skipped: {e}") - continue - sglang_us = sglang_ms * 1e3 - flashinfer_us = flashinfer_ms * 1e3 - speedup = flashinfer_us / sglang_us - print( - f"{M:>8} {K:>8} {sglang_us:>12.3f} {flashinfer_us:>16.3f} {speedup:>10.3f}" - ) - rows.append((M, K, sglang_us, flashinfer_us, speedup)) - - if args.csv: - with open(args.csv, "w") as f: - f.write("M,K,sglang_us,flashinfer_us,speedup_flashinfer_over_sglang\n") - for M, K, s, fi, sp in rows: - f.write(f"{M},{K},{s:.6f},{fi:.6f},{sp:.6f}\n") - print(f"Saved CSV to {args.csv}") - - if args.plot: - plot_speedup(rows, args.plot) - - -if __name__ == "__main__": - main() diff --git a/docs_new/docs/advanced_features/quantization.mdx b/docs_new/docs/advanced_features/quantization.mdx index ee0d4b7e9..f913a0d3e 100644 --- a/docs_new/docs/advanced_features/quantization.mdx +++ b/docs_new/docs/advanced_features/quantization.mdx @@ -284,11 +284,6 @@ Backend selection applies to **blockwise FP8**, **MXFP8** (dense linear), and ** SM80+ Auto-selects: flashinfer_cutedsl on SM100; marlin on SM80-SM90; flashinfer_cutlass otherwise (including SM120) - - cutlass - SM100/120 - SGLang CUTLASS kernel - flashinfer_cutlass SM100/120 @@ -317,7 +312,7 @@ Backend selection applies to **blockwise FP8**, **MXFP8** (dense linear), and ** -On Blackwell, when FlashInfer is unavailable for NVFP4, the SGLang CUTLASS kernel is used as an automatic fallback. On SM80-SM90, `auto` selects Marlin for NVFP4. +On SM80-SM90, `auto` selects Marlin for NVFP4. NVFP4 GEMM requires FlashInfer to be installed. ## Offline Quantization diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index 7a71a1a02..7694eac40 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -1436,9 +1436,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--fp4-gemm-backend` - Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default; selects flashinfer_cutedsl on SM100, marlin on SM80-SM90, flashinfer_cutlass otherwise (including SM120)), 'cutlass' (SGLang CUTLASS kernel), 'flashinfer_cutlass' (FlashInfer CUTLASS backend), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_cutedsl' (FlashInfer CuTe DSL backend), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling), 'marlin' (weight-only W4A16 fallback for SM80-SM90). All FlashInfer backends fall back to sgl-kernel CUTLASS when FlashInfer is unavailable. + Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default; selects flashinfer_cutedsl on SM100, marlin on SM80-SM90, flashinfer_cutlass otherwise (including SM120)), 'flashinfer_cutlass' (FlashInfer CUTLASS backend), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_cutedsl' (FlashInfer CuTe DSL backend), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling), 'marlin' (weight-only W4A16 fallback for SM80-SM90). Requires FlashInfer to be installed. `auto` - auto, cutlass, flashinfer_cudnn, flashinfer_cutedsl, flashinfer_cutlass, flashinfer_trtllm, marlin + auto, flashinfer_cudnn, flashinfer_cutedsl, flashinfer_cutlass, flashinfer_trtllm, marlin `--disable-flashinfer-autotune` diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh deleted file mode 100644 index 6378825da..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh +++ /dev/null @@ -1,806 +0,0 @@ -#include -#include - -#include -#include -#include - -#include "nvfp4_quant.cuh" -#include -#include - -using namespace host; - -// Quantizes the provided PackedVec into the uint32_t output -template -SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - // Get absolute maximum values among the local 8 values. - auto localMax = __habs2(vec.elts[0]); - -// Local maximum value. -#pragma unroll - for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { - localMax = __hmax2(localMax, __habs2(vec.elts[i])); - } - - // Get the absolute maximum among all 16 values (two threads). - localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax); - // Get the final absolute maximum values. - float vecMax = float(__hmax(localMax.x, localMax.y)); - - // Get the SF (max value of the vector / max value of e2m1). - // maximum value of e2m1 = 6.0. - // TODO: use half as compute data type. - float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); - // 8 bits representation of the SF. - uint8_t fp8SFVal; - // Write the SF to global memory (STG.8). - if constexpr (UE8M0_SF) { - // Extract the 8 exponent bits from float32. - // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits. - uint32_t tmp = reinterpret_cast(SFValue) >> 23; - fp8SFVal = tmp & 0xff; - // Convert back to fp32. - reinterpret_cast(SFValue) = tmp << 23; - } else { - // Here SFValue is always positive, so E4M3 is the same as UE4M3. - __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); - reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp; - // Convert back to fp32. - SFValue = float(tmp); - } - // Get the output scale. - // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * - // reciprocal(SFScaleVal)) - float outputScale = - SFValue != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f; - - if (SFout) { - // Write the SF to global memory (STG.8). - *SFout = fp8SFVal; - } - - // Convert the input to float. - float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; - -#pragma unroll - for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { - fp2Vals[i] = device::cast(vec.elts[i]); - fp2Vals[i].x *= outputScale; - fp2Vals[i].y *= outputScale; - } - - // Convert to e2m1 values. - uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals); - - // Write the e2m1 values to global memory. - return e2m1Vec; -#else - return 0; -#endif -} - -SGL_DEVICE float silu(const float& val) { - return val / (1.0f + __expf(-val)); -} - -template -SGL_DEVICE void silu_and_mul(PackedVec& x_vec, const PackedVec& y_vec) { - float2 x[CVT_FP4_ELTS_PER_THREAD / 2]; - float2 y[CVT_FP4_ELTS_PER_THREAD / 2]; - -#pragma unroll - for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { - x[i] = device::cast(x_vec.elts[i]); - y[i] = device::cast(y_vec.elts[i]); - x[i].x = silu(x[i].x) * y[i].x; - x[i].y = silu(x[i].y) * y[i].y; - x_vec.elts[i] = device::cast>(x[i]); - } -} - -// Use UE4M3 by default. -template -__global__ void -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -__launch_bounds__(512, 4) cvt_fp16_to_fp4( -#else -cvt_fp16_to_fp4( -#endif - int32_t numRows, - int32_t numCols, - Type const* in, - float const* SFScale, - uint32_t* out, - uint32_t* SFout, - uint32_t* input_offset_by_experts, - uint32_t* output_scale_offset_by_experts, - int32_t* mask, - int n_experts, - bool low_latency, - bool use_silu_and_mul) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using PackedVec = PackedVec; - static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); - static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); - - // Input tensor row/col loops. - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; - bool use_mask = mask != nullptr; - // When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated). - int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow; - - // Each global thread processes one element - for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { - // Calculate which row and column this global thread should process - int rowIdx = globalIdx / colsPerRow; - int colIdx = globalIdx % colsPerRow; - - // Find index within the experts using different strategies based on expert - // count - int rowIdx_in_expert = 0; - int expert_idx = 0; - - if constexpr (SMALL_NUM_EXPERTS) { - for (int i = 0; i < n_experts; i++) { - uint32_t current_offset = __ldca(&input_offset_by_experts[i]); - uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]); - if (rowIdx >= current_offset && rowIdx < next_offset) { - rowIdx_in_expert = rowIdx - current_offset; - expert_idx = i; - break; - } - } - } else { - // Load input offsets into registers first, then do the computation. - // Local array size set to 17 because of register limit. - uint32_t local_offsets[17]; - for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) { - *reinterpret_cast(local_offsets) = - __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start])); - *reinterpret_cast(local_offsets + 4) = - __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start + 4])); - *reinterpret_cast(local_offsets + 8) = - __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start + 8])); - *reinterpret_cast(local_offsets + 12) = - __ldca(reinterpret_cast(&input_offset_by_experts[chunk_start + 12])); - local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]); - -// Check against the 16 loaded offsets -#pragma unroll - for (int i = 0; i < 16; i++) { - if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) { - rowIdx_in_expert = rowIdx - local_offsets[i]; - expert_idx = chunk_start + i; - break; - } - } - } - } - - // Early exit when using masks. - if (use_mask && rowIdx_in_expert >= mask[expert_idx]) { - continue; - } - - int64_t inOffset = rowIdx * actualColsPerRow + colIdx; - PackedVec in_vec = reinterpret_cast(in)[inOffset]; - if (use_mask || use_silu_and_mul) { - PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; - silu_and_mul(in_vec, in_vec_mul); - } - - // Get the output tensor offset. - // Same as inOffset because 8 elements are packed into one uint32_t. - int64_t outOffset = rowIdx * colsPerRow + colIdx; - auto& out_pos = out[outOffset]; - - // Get the global scaling factor, which will be applied to the SF. - // Note SFScale is the same as next GEMM's alpha, which is - // (448.f / (Alpha_A / 6.f)). - float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx]; - - int factor = CVT_FP4_SF_VEC_SIZE * 4; - // The actual output_scales dim is computed from the padded numCols. - int32_t numCols_padded = (numCols + factor - 1) / factor * factor; - int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4; - uint32_t* SFout_in_expert = SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout; - - auto sf_out = cvt_quant_to_fp4_get_sf_out_offset( - rowIdx_in_expert, colIdx, numCols, SFout_in_expert); - - out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); - } -#endif -} - -// Use UE4M3 by default. -template -__global__ void -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -__launch_bounds__(512, 4) cvt_fp16_to_fp4_expert( -#else -cvt_fp16_to_fp4_expert( -#endif - int32_t numRows, - int32_t numCols, - Type const* in, - float const* SFScale, - uint32_t* out, - uint32_t* SFout, - int32_t* mask, - bool use_silu_and_mul, - int n_experts) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using PackedVec = PackedVec; - static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); - static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); - - // Input tensor row/col loops. - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int stride = (gridDim.x * blockDim.x) / n_experts; - int remainder = (gridDim.x * blockDim.x) % n_experts; - int expert_idx; - int tid_in_expert; - int actual_stride; - if (remainder > 0) { - int bound = remainder * (stride + 1); - if (tid < bound) { - expert_idx = tid / (stride + 1); - tid_in_expert = tid % (stride + 1); - actual_stride = stride + 1; - } else { - expert_idx = remainder + (tid - bound) / stride; - tid_in_expert = (tid - bound) % stride; - actual_stride = stride; - } - } else { - expert_idx = tid / stride; - tid_in_expert = tid % stride; - actual_stride = stride; - } - int m = numRows / n_experts; - int padded_m = (m + (128 - 1)) / 128 * 128; - - int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; - // TODO(kaixih@nvidia): For now, we assume mask is used together with - // silu_and_mal. Maybe we want a more general behavior of mask later. In the - // silu case, the input last dim doubles. - bool use_mask = mask != nullptr; - int actualColsPerRow = use_silu_and_mul ? colsPerRow * 2 : colsPerRow; - - // Each global thread processes one element - for (int globalIdx = tid_in_expert + expert_idx * m * colsPerRow; globalIdx < (expert_idx + 1) * m * colsPerRow; - globalIdx += actual_stride) { - // Calculate which row and column this global thread should process - int rowIdx = globalIdx / colsPerRow; - int colIdx = globalIdx % colsPerRow; - - // Find index within the experts - int rowIdx_in_expert = rowIdx - expert_idx * m; - - // Early exit when using masks. - if (use_mask && rowIdx_in_expert >= mask[expert_idx]) { - break; - } - - int64_t inOffset = rowIdx * actualColsPerRow + colIdx; - PackedVec in_vec = reinterpret_cast(in)[inOffset]; - if (use_silu_and_mul) { - PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; - silu_and_mul(in_vec, in_vec_mul); - } - - // Get the output tensor offset. - // Same as inOffset because 8 elements are packed into one uint32_t. - int64_t outOffset = rowIdx * colsPerRow + colIdx; - auto& out_pos = out[outOffset]; - - // Get the global scaling factor, which will be applied to the SF. - // Note SFScale is the same as next GEMM's alpha, which is - // (448.f / (Alpha_A / 6.f)). - float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx]; - - int factor = CVT_FP4_SF_VEC_SIZE * 4; - // The actual output_scales dim is computed from the padded numCols. - int32_t numCols_padded = (numCols + factor - 1) / factor * factor; - int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4; - uint32_t* SFout_in_expert = SFout + expert_idx * padded_m * numCols_SFout; - - auto sf_out = cvt_quant_to_fp4_get_sf_out_offset( - rowIdx_in_expert, colIdx, numCols, SFout_in_expert); - - out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); - } -#endif -} - -// Kernel for LARGE_M_TOPK = true (large m_topk optimized version) -template -__global__ void -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -__launch_bounds__(1024, 4) cvt_fp16_to_fp4( -#else -cvt_fp16_to_fp4( -#endif - int32_t numRows, - int32_t numCols, - Type const* in, - float const* SFScale, - uint32_t* out, - uint32_t* SFout, - uint32_t* input_offset_by_experts, - uint32_t* output_scale_offset_by_experts, - int32_t* mask, - int n_experts, - bool use_silu_and_mul) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using PackedVec = PackedVec; - static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); - static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); - extern __shared__ uint32_t shared_input_offsets[]; - - // Load input offsets into shared memory. - // If n_experts is larger than 4, use vectorized int4 to save instructions. - // If n_experts is smaller than 4, read directly. - if constexpr (SMALL_NUM_EXPERTS) { - for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) { - shared_input_offsets[i] = input_offset_by_experts[i]; - } - } else { - for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) { - *reinterpret_cast(&shared_input_offsets[i]) = *reinterpret_cast(&input_offset_by_experts[i]); - } - if (threadIdx.x == 0) { - shared_input_offsets[n_experts] = input_offset_by_experts[n_experts]; - } - } - - __syncthreads(); - - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; - bool use_mask = mask != nullptr; - // When use_silu_and_mul is true, input last dim is 2*k (gate+up concatenated). - int actualColsPerRow = (use_mask || use_silu_and_mul) ? colsPerRow * 2 : colsPerRow; - - // Each global thread processes one element - for (int globalIdx = tid; globalIdx < numRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { - // Calculate which row and column this global thread should process - int rowIdx = globalIdx / colsPerRow; - int colIdx = globalIdx % colsPerRow; - - // Find expert using binary search for better performance with large m_topk - int rowIdx_in_expert = 0; - int expert_idx = 0; - - // Binary search through experts using shared memory - int left = 0, right = n_experts - 1; - while (left <= right) { - int mid = (left + right) / 2; - // Get offsets: shared_input_offsets[i] corresponds to - // input_offset_by_experts[i] - uint32_t mid_offset = shared_input_offsets[mid]; - uint32_t next_offset = shared_input_offsets[mid + 1]; - - if (rowIdx >= mid_offset && rowIdx < next_offset) { - rowIdx_in_expert = rowIdx - mid_offset; - expert_idx = mid; - break; - } else if (rowIdx < mid_offset) { - right = mid - 1; - } else { - left = mid + 1; - } - } - - if (use_mask && rowIdx_in_expert >= mask[expert_idx]) { - continue; - } - - int64_t inOffset = rowIdx * actualColsPerRow + colIdx; - - PackedVec in_vec = reinterpret_cast(in)[inOffset]; - if (use_mask || use_silu_and_mul) { - PackedVec in_vec_mul = reinterpret_cast(in)[inOffset + colsPerRow]; - silu_and_mul(in_vec, in_vec_mul); - } - - int64_t outOffset = rowIdx * colsPerRow + colIdx; - auto& out_pos = out[outOffset]; - - float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx]; - - int factor = CVT_FP4_SF_VEC_SIZE * 4; - int32_t numCols_padded = (numCols + factor - 1) / factor * factor; - int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4; - uint32_t* SFout_in_expert = SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout; - - auto sf_out = cvt_quant_to_fp4_get_sf_out_offset( - rowIdx_in_expert, colIdx, numCols, SFout_in_expert); - - out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); - } -#endif -} - -template -void quant_impl( - void* output, - void* output_scale, - void* input, - void* input_global_scale, - void* input_offset_by_experts, - void* output_scale_offset_by_experts, - void* mask, - bool use_silu_and_mul, - int m_topk, - int k, - int n_experts, - cudaStream_t stream) { - // TODO: this multiProcessorCount should be cached. - int device; - cudaGetDevice(&device); - int multiProcessorCount; - cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, device); - - // Grid, Block size. - // Each thread converts 8 values. - int const workSizePerRow = k / ELTS_PER_THREAD; - int const totalWorkSize = m_topk * workSizePerRow; - dim3 block(std::min(workSizePerRow, 512)); - // Get number of blocks per SM (assume we can fully utilize the SM). - int const numBlocksPerSM = 2048 / block.x; - dim3 grid(std::min(static_cast((totalWorkSize + block.x - 1) / block.x), multiProcessorCount * numBlocksPerSM)); - while (grid.x <= multiProcessorCount && block.x > 64) { - grid.x *= 2; - block.x = (block.x + 1) / 2; - } - - // TODO(kaixih@nvidia): Should relax this to allow any grid size. - if (mask != nullptr) { - grid.x = (grid.x + n_experts - 1) / n_experts * n_experts; - cvt_fp16_to_fp4_expert<<>>( - m_topk, - k, - reinterpret_cast(input), - reinterpret_cast(input_global_scale), - reinterpret_cast(output), - reinterpret_cast(output_scale), - reinterpret_cast(mask), - use_silu_and_mul, - n_experts); - return; - } - - int const blockRepeat = (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x); - if (blockRepeat > 1) { - size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t); - if (n_experts >= 4) { - cvt_fp16_to_fp4<<>>( - m_topk, - k, - reinterpret_cast(input), - reinterpret_cast(input_global_scale), - reinterpret_cast(output), - reinterpret_cast(output_scale), - reinterpret_cast(input_offset_by_experts), - reinterpret_cast(output_scale_offset_by_experts), - reinterpret_cast(mask), - n_experts, - use_silu_and_mul); - } else { - cvt_fp16_to_fp4<<>>( - m_topk, - k, - reinterpret_cast(input), - reinterpret_cast(input_global_scale), - reinterpret_cast(output), - reinterpret_cast(output_scale), - reinterpret_cast(input_offset_by_experts), - reinterpret_cast(output_scale_offset_by_experts), - reinterpret_cast(mask), - n_experts, - use_silu_and_mul); - } - } else { - if (n_experts >= 16) { - cvt_fp16_to_fp4<<>>( - m_topk, - k, - reinterpret_cast(input), - reinterpret_cast(input_global_scale), - reinterpret_cast(output), - reinterpret_cast(output_scale), - reinterpret_cast(input_offset_by_experts), - reinterpret_cast(output_scale_offset_by_experts), - reinterpret_cast(mask), - n_experts, - /* bool low_latency */ true, - use_silu_and_mul); - } else { - cvt_fp16_to_fp4<<>>( - m_topk, - k, - reinterpret_cast(input), - reinterpret_cast(input_global_scale), - reinterpret_cast(output), - reinterpret_cast(output_scale), - reinterpret_cast(input_offset_by_experts), - reinterpret_cast(output_scale_offset_by_experts), - reinterpret_cast(mask), - n_experts, - /* bool low_latency */ true, - use_silu_and_mul); - } - } -} - -inline int getSMVersion(int device_id) { - int sm_major = 0; - int sm_minor = 0; - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); - return sm_major * 10 + sm_minor; -} - -void scaled_fp4_experts_quant_sm100a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView input_offset_by_experts, - tvm::ffi::TensorView output_scale_offset_by_experts) { - auto MTopK = SymbolicSize{"m_topk"}; - auto K = SymbolicSize{"k"}; - auto OutputCols = SymbolicSize{"output_cols"}; - auto OutputScaleRows = SymbolicSize{"output_scale_rows"}; - auto OutputScaleCols = SymbolicSize{"output_scale_cols"}; - auto NExperts = SymbolicSize{"n_experts"}; - auto OffsetSize = SymbolicSize{"offset_size"}; - auto device = SymbolicDevice{}; - - TensorMatcher({MTopK, K}) // - .with_dtype() - .template with_device(device) - .verify(input); - TensorMatcher({MTopK, OutputCols}) // - .with_dtype() - .with_device(device) - .verify(output); - TensorMatcher({OutputScaleRows, OutputScaleCols}) // - .with_dtype() - .with_device(device) - .verify(output_scale); - TensorMatcher({NExperts}) // - .with_dtype() - .with_device(device) - .verify(input_global_scale); - TensorMatcher({OffsetSize}) // - .with_dtype() - .with_device(device) - .verify(input_offset_by_experts) - .verify(output_scale_offset_by_experts); - - const int device_id = input.device().device_id; - RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); - - const int BLOCK_SIZE = 16; - const auto m_topk = static_cast(MTopK.unwrap()); - const auto k = static_cast(K.unwrap()); - RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16"); - const auto n_experts = static_cast(NExperts.unwrap()); - const auto offset_size = static_cast(OffsetSize.unwrap()); - RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch"); - RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); - const int scales_k = k / BLOCK_SIZE; - const int padded_k = (scales_k + 3) / 4 * 4; - RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); - - const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); - if (host::is_type(input.dtype())) { - quant_impl( - output.data_ptr(), - output_scale.data_ptr(), - input.data_ptr(), - input_global_scale.data_ptr(), - input_offset_by_experts.data_ptr(), - output_scale_offset_by_experts.data_ptr(), - nullptr, // mask - false, // use_silu_and_mul - m_topk, - k, - n_experts, - stream); - } else { - quant_impl<__nv_bfloat16>( - output.data_ptr(), - output_scale.data_ptr(), - input.data_ptr(), - input_global_scale.data_ptr(), - input_offset_by_experts.data_ptr(), - output_scale_offset_by_experts.data_ptr(), - nullptr, // mask - false, // use_silu_and_mul - m_topk, - k, - n_experts, - stream); - } -} - -void silu_and_mul_scaled_fp4_experts_quant_sm100a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView mask, - bool use_silu_and_mul) { - auto MTopK = SymbolicSize{"m_topk"}; - auto KBy2 = SymbolicSize{"k_by_2"}; - auto OutputCols = SymbolicSize{"output_cols"}; - auto OutputScaleRows = SymbolicSize{"output_scale_rows"}; - auto OutputScaleCols = SymbolicSize{"output_scale_cols"}; - auto NExperts = SymbolicSize{"n_experts"}; - auto device = SymbolicDevice{}; - - TensorMatcher({MTopK, KBy2}) // - .with_dtype() - .template with_device(device) - .verify(input); - TensorMatcher({MTopK, OutputCols}) // - .with_dtype() - .with_device(device) - .verify(output); - TensorMatcher({OutputScaleRows, OutputScaleCols}) // - .with_dtype() - .with_device(device) - .verify(output_scale); - TensorMatcher({NExperts}) // - .with_dtype() - .with_device(device) - .verify(input_global_scale); - TensorMatcher({NExperts}) // - .with_dtype() - .with_device(device) - .verify(mask); - - const int device_id = input.device().device_id; - RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); - - const int BLOCK_SIZE = 16; - const auto m_topk = static_cast(MTopK.unwrap()); - const auto k_by_2 = static_cast(KBy2.unwrap()); - int k = k_by_2; - if (use_silu_and_mul) { - RuntimeCheck(k_by_2 % 2 == 0, "k must be a multiple of 2"); - k = k_by_2 / 2; - } - const auto n_experts = static_cast(NExperts.unwrap()); - RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); - const int scales_k = k / BLOCK_SIZE; - const int padded_k = (scales_k + 3) / 4 * 4; - RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); - - const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); - if (host::is_type(input.dtype())) { - quant_impl( - output.data_ptr(), - output_scale.data_ptr(), - input.data_ptr(), - input_global_scale.data_ptr(), - nullptr, // input_offset_by_experts - nullptr, // output_scale_offset_by_experts - mask.data_ptr(), - use_silu_and_mul, - m_topk, - k, - n_experts, - stream); - } else { - quant_impl<__nv_bfloat16>( - output.data_ptr(), - output_scale.data_ptr(), - input.data_ptr(), - input_global_scale.data_ptr(), - nullptr, // input_offset_by_experts - nullptr, // output_scale_offset_by_experts - mask.data_ptr(), - use_silu_and_mul, - m_topk, - k, - n_experts, - stream); - } -} - -void silu_and_mul_scaled_fp4_experts_quant_packed_sm100a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView input_offset_by_experts, - tvm::ffi::TensorView output_scale_offset_by_experts) { - auto MTopK = SymbolicSize{"m_topk"}; - auto KBy2 = SymbolicSize{"k_by_2"}; - auto OutputCols = SymbolicSize{"output_cols"}; - auto OutputScaleRows = SymbolicSize{"output_scale_rows"}; - auto OutputScaleCols = SymbolicSize{"output_scale_cols"}; - auto NExperts = SymbolicSize{"n_experts"}; - auto OffsetSize = SymbolicSize{"offset_size"}; - auto device = SymbolicDevice{}; - - TensorMatcher({MTopK, KBy2}) // - .with_dtype() - .template with_device(device) - .verify(input); - TensorMatcher({MTopK, OutputCols}) // - .with_dtype() - .with_device(device) - .verify(output); - TensorMatcher({OutputScaleRows, OutputScaleCols}) // - .with_dtype() - .with_device(device) - .verify(output_scale); - TensorMatcher({NExperts}) // - .with_dtype() - .with_device(device) - .verify(input_global_scale); - TensorMatcher({OffsetSize}) // - .with_dtype() - .with_device(device) - .verify(input_offset_by_experts) - .verify(output_scale_offset_by_experts); - - const int device_id = input.device().device_id; - RuntimeCheck(getSMVersion(device_id) >= 100, "fp4_quant is only supported on sm100+"); - - const int BLOCK_SIZE = 16; - const auto m_topk = static_cast(MTopK.unwrap()); - const auto k_by_2 = static_cast(KBy2.unwrap()); - // Input last dim is 2*k (gate+up concatenated). The kernel does SiLU(gate)*up - // then FP4-quantizes the k-dim result. - RuntimeCheck(k_by_2 % 2 == 0, "input last dim must be even (2*k)"); - const int k = k_by_2 / 2; - RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16"); - const auto n_experts = static_cast(NExperts.unwrap()); - const auto offset_size = static_cast(OffsetSize.unwrap()); - RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch"); - RuntimeCheck(static_cast(OutputCols.unwrap()) == k / 2, "output second dim mismatch"); - const int scales_k = k / BLOCK_SIZE; - const int padded_k = (scales_k + 3) / 4 * 4; - RuntimeCheck(static_cast(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch"); - - const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); - if (host::is_type(input.dtype())) { - quant_impl( - output.data_ptr(), - output_scale.data_ptr(), - input.data_ptr(), - input_global_scale.data_ptr(), - input_offset_by_experts.data_ptr(), - output_scale_offset_by_experts.data_ptr(), - nullptr, // mask - true, // use_silu_and_mul - m_topk, - k, - n_experts, - stream); - } else { - quant_impl<__nv_bfloat16>( - output.data_ptr(), - output_scale.data_ptr(), - input.data_ptr(), - input_global_scale.data_ptr(), - input_offset_by_experts.data_ptr(), - output_scale_offset_by_experts.data_ptr(), - nullptr, // mask - true, // use_silu_and_mul - m_topk, - k, - n_experts, - stream); - } -} diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh deleted file mode 100644 index e2d696673..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh +++ /dev/null @@ -1,160 +0,0 @@ -/* Copyright 2025 SGLang Team. All Rights Reserved. - -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 - -#define ELTS_PER_THREAD 8 - -constexpr int CVT_FP4_ELTS_PER_THREAD = 8; -constexpr int CVT_FP4_SF_VEC_SIZE = 16; - -// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t). -SGL_DEVICE uint32_t fp32_vec_to_e2m1(float (&array)[8]) { - // PTX instructions used here requires >= sm100f. -#if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \ - (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)) - uint32_t val; - asm volatile( - "{\n" - ".reg .b8 byte0;\n" - ".reg .b8 byte1;\n" - ".reg .b8 byte2;\n" - ".reg .b8 byte3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" - "}" - : "=r"(val) - : "f"(array[0]), - "f"(array[1]), - "f"(array[2]), - "f"(array[3]), - "f"(array[4]), - "f"(array[5]), - "f"(array[6]), - "f"(array[7])); - return val; -#else - printf("fp32_vec_to_e2m1 is not supported on this architecture\n"); - __trap(); - return 0; -#endif -} - -// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t). -SGL_DEVICE uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) { - // PTX instructions used here requires >= sm100f. -#if CUTLASS_ARCH_MMA_SM100A_ENABLED || CUTLASS_ARCH_MMA_SM103A_ENABLED || CUTLASS_ARCH_MMA_SM120A_ENABLED || \ - (defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)) - uint32_t val; - asm volatile( - "{\n" - ".reg .b8 byte0;\n" - ".reg .b8 byte1;\n" - ".reg .b8 byte2;\n" - ".reg .b8 byte3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" - "}" - : "=r"(val) - : "f"(array[0].x), - "f"(array[0].y), - "f"(array[1].x), - "f"(array[1].y), - "f"(array[2].x), - "f"(array[2].y), - "f"(array[3].x), - "f"(array[3].y)); - return val; -#else - printf("fp32_vec_to_e2m1 is not supported on this architecture\n"); - __trap(); - return 0; -#endif -} - -// Fast reciprocal. -SGL_DEVICE float reciprocal_approximate_ftz(float a) { - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -template -SGL_DEVICE uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx, int numCols, SFType* SFout) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || CVT_FP4_NUM_THREADS_PER_SF == 2); - - // One pair of threads write one SF to global memory. - // TODO: stage through smem for packed STG.32 - // is it better than STG.8 from 4 threads ? - if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) { - // SF vector index (16 elements share one SF in the K dimension). - int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; - int32_t mIdx = rowIdx; - - // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] - // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] - - int32_t mTileIdx = mIdx / (32 * 4); - // SF vector size 16. - int factor = CVT_FP4_SF_VEC_SIZE * 4; - int32_t numKTiles = (numCols + factor - 1) / factor; - int64_t mTileStride = numKTiles * 32 * 4 * 4; - - int32_t kTileIdx = (kIdx / 4); - int64_t kTileStride = 32 * 4 * 4; - - // M tile layout [32, 4] is column-major. - int32_t outerMIdx = (mIdx % 32); - int64_t outerMStride = 4 * 4; - - int32_t innerMIdx = (mIdx % (32 * 4)) / 32; - int64_t innerMStride = 4; - - int32_t innerKIdx = (kIdx % 4); - int64_t innerKStride = 1; - - // Compute the global offset. - int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride + outerMIdx * outerMStride + - innerMIdx * innerMStride + innerKIdx * innerKStride; - - return reinterpret_cast(SFout) + SFOffset; - } -#endif - return nullptr; -} - -// Define a 16 bytes packed data type. -template -struct PackedVec { - packed_t elts[4]; -}; - -template <> -struct PackedVec<__nv_fp8_e4m3> { - __nv_fp8x2_e4m3 elts[8]; -}; diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh deleted file mode 100644 index f2bd76787..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh +++ /dev/null @@ -1,87 +0,0 @@ -/* Copyright 2025 SGLang Team. All Rights Reserved. - -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 - -void scaled_fp4_quant_sm100a_sm120a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView input, - tvm::ffi::TensorView output_sf, - tvm::ffi::TensorView input_sf); - -void scaled_fp4_experts_quant_sm100a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView input_offset_by_experts, - tvm::ffi::TensorView output_scale_offset_by_experts); - -void silu_and_mul_scaled_fp4_experts_quant_sm100a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView mask, - bool use_silu_and_mul); - -void silu_and_mul_scaled_fp4_experts_quant_packed_sm100a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView input_offset_by_experts, - tvm::ffi::TensorView output_scale_offset_by_experts); - -void scaled_fp4_quant( - tvm::ffi::TensorView output, - tvm::ffi::TensorView input, - tvm::ffi::TensorView output_sf, - tvm::ffi::TensorView input_sf) { - scaled_fp4_quant_sm100a_sm120a(output, input, output_sf, input_sf); -} - -void scaled_fp4_experts_quant( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView input_offset_by_experts, - tvm::ffi::TensorView output_scale_offset_by_experts) { - scaled_fp4_experts_quant_sm100a( - output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts); -} - -void silu_and_mul_scaled_fp4_experts_quant( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView mask, - bool use_silu_and_mul) { - silu_and_mul_scaled_fp4_experts_quant_sm100a(output, output_scale, input, input_global_scale, mask, use_silu_and_mul); -} - -void silu_and_mul_scaled_fp4_experts_quant_packed( - tvm::ffi::TensorView output, - tvm::ffi::TensorView output_scale, - tvm::ffi::TensorView input, - tvm::ffi::TensorView input_global_scale, - tvm::ffi::TensorView input_offset_by_experts, - tvm::ffi::TensorView output_scale_offset_by_experts) { - silu_and_mul_scaled_fp4_experts_quant_packed_sm100a( - output, output_scale, input, input_global_scale, input_offset_by_experts, output_scale_offset_by_experts); -} diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh deleted file mode 100644 index bac38d83e..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh +++ /dev/null @@ -1,241 +0,0 @@ -/* Copyright 2025 SGLang Team. All Rights Reserved. - -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 "nvfp4_quant.cuh" -#include -#include - -using namespace host; - -// Quantizes the provided PackedVec into the uint32_t output -template -SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec& vec, float SFScaleVal, uint8_t* SFout) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - // Get absolute maximum values among the local 8 values. - auto localMax = __habs2(vec.elts[0]); - -// Local maximum value. -#pragma unroll - for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { - localMax = __hmax2(localMax, __habs2(vec.elts[i])); - } - - // Get the absolute maximum among all 16 values (two threads). - localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax); - // Get the final absolute maximum values. - float vecMax = float(__hmax(localMax.x, localMax.y)); - - // Get the SF (max value of the vector / max value of e2m1). - // maximum value of e2m1 = 6.0. - // TODO: use half as compute data type. - float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); - // 8 bits representation of the SF. - uint8_t fp8SFVal; - // Write the SF to global memory (STG.8). - if constexpr (UE8M0_SF) { - __nv_fp8_e8m0 tmp; - tmp.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf); - SFValue = static_cast(tmp); - fp8SFVal = tmp.__x; - } else { - // Here SFValue is always positive, so E4M3 is the same as UE4M3. - __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); - fp8SFVal = tmp.__x; - SFValue = static_cast(tmp); - } - // Get the output scale. - // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * - // reciprocal(SFScaleVal)) - float outputScale = - SFValue != 0 ? reciprocal_approximate_ftz(SFValue * reciprocal_approximate_ftz(SFScaleVal)) : 0.0f; - - if (SFout) { - // Write the SF to global memory (STG.8). - *SFout = fp8SFVal; - } - - // Convert the input to float. - float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; - -#pragma unroll - for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { - if constexpr (std::is_same_v) { - fp2Vals[i] = __half22float2(vec.elts[i]); - } else { - fp2Vals[i] = __bfloat1622float2(vec.elts[i]); - } - fp2Vals[i].x *= outputScale; - fp2Vals[i].y *= outputScale; - } - - // Convert to e2m1 values. - uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals); - - // Write the e2m1 values to global memory. - return e2m1Vec; -#else - return 0; -#endif -} - -// Use UE4M3 by default. -template -__global__ void -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -__launch_bounds__(512, 4) cvt_fp16_to_fp4( -#else -cvt_fp16_to_fp4( -#endif - int32_t numRows, int32_t numCols, Type const* in, float const* SFScale, uint32_t* out, uint32_t* SFout) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using PackedVec = PackedVec; - static constexpr int CVT_FP4_NUM_THREADS_PER_SF = (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); - static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); - - // Get the global scaling factor, which will be applied to the SF. - // Note SFScale is the same as next GEMM's alpha, which is - // (448.f / (Alpha_A / 6.f)). - float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0]; - - // Input tensor row/col loops. - for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { - for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD; colIdx += blockDim.x) { - int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; - PackedVec in_vec = reinterpret_cast(in)[inOffset]; - // Get the output tensor offset. - // Same as inOffset because 8 elements are packed into one uint32_t. - int64_t outOffset = inOffset; - auto& out_pos = out[outOffset]; - - auto sf_out = - cvt_quant_to_fp4_get_sf_out_offset(rowIdx, colIdx, numCols, SFout); - - out_pos = cvt_warp_fp16_to_fp4(in_vec, SFScaleVal, sf_out); - } - } -#endif -} - -template -void invokeFP4Quantization( - int m, - int n, - T const* input, - float const* SFScale, - int64_t* output, - int32_t* SFOuput, - bool useUE8M0, - int multiProcessorCount, - cudaStream_t stream) { - // Grid, Block size. - // Each thread converts 8 values. - dim3 block(std::min(int(n / ELTS_PER_THREAD), 512)); - // Get number of blocks per SM (assume we can fully utilize the SM). - int const numBlocksPerSM = 2048 / block.x; - dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM)); - - // Launch the cvt kernel. - if (useUE8M0) { - cvt_fp16_to_fp4<<>>( - m, n, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); - } else { - cvt_fp16_to_fp4<<>>( - m, n, input, SFScale, reinterpret_cast(output), reinterpret_cast(SFOuput)); - } -} - -// Instantiate the function. -template void invokeFP4Quantization( - int m, - int n, - half const* input, - float const* SFScale, - int64_t* output, - int32_t* SFOuput, - bool useUE8M0, - int multiProcessorCount, - cudaStream_t stream); - -template void invokeFP4Quantization( - int m, - int n, - __nv_bfloat16 const* input, - float const* SFScale, - int64_t* output, - int32_t* SFOuput, - bool useUE8M0, - int multiProcessorCount, - cudaStream_t stream); - -inline int getSMVersion(int device_id) { - int sm_major = 0; - int sm_minor = 0; - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); - return sm_major * 10 + sm_minor; -} - -void scaled_fp4_quant_sm100a_sm120a( - tvm::ffi::TensorView output, - tvm::ffi::TensorView input, - tvm::ffi::TensorView output_sf, - tvm::ffi::TensorView input_sf) { - RuntimeCheck(input.device().device_type == kDLCUDA, "input must be a CUDA tensor"); - RuntimeCheck(output.device() == input.device(), "output and input must be on same device"); - RuntimeCheck(output_sf.device() == input.device(), "output_sf and input must be on same device"); - RuntimeCheck(input_sf.device() == input.device(), "input_sf and input must be on same device"); - RuntimeCheck(input.dim() == 2, "input must be a 2D tensor"); - RuntimeCheck(output.dim() == 2, "output must be a 2D tensor"); - RuntimeCheck(output_sf.dim() == 2, "output_sf must be a 2D tensor"); - RuntimeCheck(input_sf.numel() == 1, "input_sf must have exactly one element"); - RuntimeCheck(host::is_type(output.dtype()), "output must be uint8"); - RuntimeCheck(host::is_type(output_sf.dtype()), "output_sf must be int32"); - RuntimeCheck(host::is_type(input_sf.dtype()), "input_sf must be float32"); - RuntimeCheck( - host::is_type(input.dtype()) || host::is_type(input.dtype()), "input dtype must be fp16 or bf16"); - - const int device_id = input.device().device_id; - const auto sm_version = getSMVersion(device_id); - RuntimeCheck(sm_version >= 100, "fp4_quant is only supported on sm100+"); - - const int32_t m = static_cast(input.size(0)); - const int32_t n = static_cast(input.size(1)); - - RuntimeCheck(output.size(0) == m, "output row size mismatch"); - RuntimeCheck(output.size(1) == n / 2, "output column size mismatch"); - RuntimeCheck(n % 16 == 0, "The N dimension must be multiple of 16."); - - const int multiProcessorCount = static_cast(runtime::get_sm_count(device_id)); - - auto input_sf_ptr = static_cast(input_sf.data_ptr()); - auto sf_out = static_cast(output_sf.data_ptr()); - auto output_ptr = static_cast(output.data_ptr()); - const cudaStream_t stream = LaunchKernel::resolve_device(input.device()); - - constexpr bool useUE8M0 = false; - if (host::is_type(input.dtype())) { - auto input_ptr = reinterpret_cast(input.data_ptr()); - invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream); - } else { - auto input_ptr = reinterpret_cast<__nv_bfloat16 const*>(input.data_ptr()); - invokeFP4Quantization(m, n, input_ptr, input_sf_ptr, output_ptr, sf_out, useUE8M0, multiProcessorCount, stream); - } -} diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_common.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_common.cuh deleted file mode 100644 index f5ebca05b..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_common.cuh +++ /dev/null @@ -1,66 +0,0 @@ -/* Copyright 2026 SGLang Team. All Rights Reserved. - -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. -==============================================================================*/ - -#pragma once - -#include -#include -#include - -#include -#include - -#include -#include -#include - -using namespace host; - -// clang-format off -#include "cutlass/cutlass.h" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/util/packed_stride.hpp" -// clang-format on - -#define CUTLASS_CHECK(status) \ - { \ - cutlass::Status error = status; \ - RuntimeCheck(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \ - } - -using namespace cute; - -inline uint32_t next_pow_2(uint32_t x) noexcept { - if (x <= 1) return 1; - return 1u << (32 - __builtin_clz(x - 1)); -} - -inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor { - if (required_bytes == 0) return {}; - DLDataType u8 = {kDLUInt, 8, 1}; - int64_t shape[] = {static_cast(required_bytes)}; - return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device); -} - -inline int getSMVersion(int device_id) { - int sm_major = 0; - int sm_minor = 0; - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); - return sm_major * 10 + sm_minor; -} diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh deleted file mode 100644 index 72d68f7d5..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh +++ /dev/null @@ -1,34 +0,0 @@ -/* Copyright 2025 SGLang Team. All Rights Reserved. - -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 - -void cutlass_scaled_fp4_mm_sm100a_sm120a( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha); - -void cutlass_scaled_fp4_mm( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha) { - cutlass_scaled_fp4_mm_sm100a_sm120a(D, A, B, A_sf, B_sf, alpha); -} diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh deleted file mode 100644 index 8c5cfefd7..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh +++ /dev/null @@ -1,146 +0,0 @@ -/* Copyright 2026 SGLang Team. All Rights Reserved. - -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 "nvfp4_scaled_mm_common.cuh" -#include "nvfp4_scaled_mm_sm100.cuh" -#include "nvfp4_scaled_mm_sm120.cuh" - -void cutlass_scaled_fp4_mm_sm100a_sm120a( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha) { - RuntimeCheck(A.device().device_type == kDLCUDA, "a must be a CUDA tensor"); - RuntimeCheck(B.device().device_type == kDLCUDA, "b must be a CUDA tensor"); - RuntimeCheck(A_sf.device().device_type == kDLCUDA, "scale_a must be a CUDA tensor"); - RuntimeCheck(B_sf.device().device_type == kDLCUDA, "scale_b must be a CUDA tensor"); - RuntimeCheck(alpha.device().device_type == kDLCUDA, "alpha must be a CUDA tensor"); - RuntimeCheck(D.device().device_type == kDLCUDA, "out must be a CUDA tensor"); - - RuntimeCheck(A.device() == B.device(), "a and b must be on same device"); - RuntimeCheck(A.device() == A_sf.device(), "a and scale_a must be on same device"); - RuntimeCheck(A.device() == B_sf.device(), "a and scale_b must be on same device"); - RuntimeCheck(A.device() == alpha.device(), "a and alpha must be on same device"); - RuntimeCheck(A.device() == D.device(), "a and out must be on same device"); - - RuntimeCheck(A.is_contiguous(), "a must be contiguous"); - RuntimeCheck(B.is_contiguous(), "b must be contiguous"); - RuntimeCheck(A_sf.is_contiguous(), "scale_a must be contiguous"); - RuntimeCheck(B_sf.is_contiguous(), "scale_b must be contiguous"); - RuntimeCheck(alpha.is_contiguous(), "alpha must be contiguous"); - RuntimeCheck(D.is_contiguous(), "out must be contiguous"); - - RuntimeCheck(host::is_type(A.dtype()), "a must be uint8"); - RuntimeCheck(host::is_type(B.dtype()), "b must be uint8"); - RuntimeCheck(host::is_type(A_sf.dtype()), "scale_a must be float8_e4m3fn"); - RuntimeCheck(host::is_type(B_sf.dtype()), "scale_b must be float8_e4m3fn"); - RuntimeCheck(host::is_type(alpha.dtype()), "alpha must be float32"); - - RuntimeCheck(A.dim() == 2, "a must be a matrix"); - RuntimeCheck(B.dim() == 2, "b must be a matrix"); - RuntimeCheck(A_sf.dim() == 2, "scale_a must be a matrix"); - RuntimeCheck(B_sf.dim() == 2, "scale_b must be a matrix"); - RuntimeCheck(alpha.numel() == 1, "alpha must have exactly one element"); - - RuntimeCheck( - A.size(1) == B.size(1), - "a and b shapes cannot be multiplied (", - A.size(0), - "x", - A.size(1), - " and ", - B.size(0), - "x", - B.size(1), - ")"); - - const auto m = static_cast(A.size(0)); - const auto n = static_cast(B.size(0)); - const auto k = static_cast(A.size(1) * 2); - - RuntimeCheck(D.dim() == 2, "out must be 2D"); - RuntimeCheck(D.size(0) == m, "out first dim must equal m"); - RuntimeCheck(D.size(1) == n, "out second dim must equal n"); - - constexpr int alignment = 32; - RuntimeCheck(k % alignment == 0, "Expected k to be divisible by ", alignment, ", but got k: ", k); - RuntimeCheck(n % alignment == 0, "Expected n to be divisible by ", alignment, ", but got n: ", n); - - auto round_up = [](int64_t x, int64_t y) { return (x + y - 1) / y * y; }; - const int64_t rounded_m = round_up(m, 128); - const int64_t rounded_n = round_up(n, 128); - const int64_t rounded_k = round_up(k / 16, 4); - - RuntimeCheck( - A_sf.size(1) == B_sf.size(1), - "scale_a and scale_b shapes cannot be multiplied (", - A_sf.size(0), - "x", - A_sf.size(1), - " and ", - B_sf.size(0), - "x", - B_sf.size(1), - ")"); - RuntimeCheck( - A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k, - "scale_a must be padded/swizzled to shape (", - rounded_m, - "x", - rounded_k, - "), got (", - A_sf.size(0), - "x", - A_sf.size(1), - ")"); - RuntimeCheck( - B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k, - "scale_b must be padded/swizzled to shape (", - rounded_n, - "x", - rounded_k, - "), got (", - B_sf.size(0), - "x", - B_sf.size(1), - ")"); - - const cudaStream_t stream = LaunchKernel::resolve_device(A.device()); - const int sm_version = getSMVersion(A.device().device_id); - - if (sm_version >= 120) { - if (host::is_type(D.dtype())) { - cutlass_fp4_f16_gemm_dispatch_sm120( - D, A, B, A_sf, B_sf, alpha, static_cast(m), static_cast(n), static_cast(k), stream); - } else if (host::is_type(D.dtype())) { - cutlass_fp4_bf16_gemm_dispatch_sm120( - D, A, B, A_sf, B_sf, alpha, static_cast(m), static_cast(n), static_cast(k), stream); - } else { - Panic("Unsupported output data type of nvfp4 mm sm120"); - } - } else { - if (host::is_type(D.dtype())) { - cutlassFp4GemmDispatchSm100(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (host::is_type(D.dtype())) { - cutlassFp4GemmDispatchSm100(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (host::is_type(D.dtype())) { - cutlassFp4GemmDispatchSm100(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else { - Panic("Unsupported output data type of nvfp4 mm"); - } - } -} diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm100.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm100.cuh deleted file mode 100644 index 699bb6236..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm100.cuh +++ /dev/null @@ -1,305 +0,0 @@ -/* Copyright 2026 SGLang Team. All Rights Reserved. - -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. -==============================================================================*/ - -#pragma once - -#include "nvfp4_scaled_mm_common.cuh" - -#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) - -// Config(half_t/bfloat16_t) for M <= 128 -template -struct KernelConfigM128 { - using OutputType = T; - using MmaTileShape = Shape<_128, _256, _256>; - using ClusterShape = Shape; - using EpilogueTile = Shape<_128, _64>; // Avoid register spilling - using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; - using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100; - const static dim3 preferred_cluster; - const static dim3 fallback_cluster; -}; -template -const dim3 KernelConfigM128::preferred_cluster(1, 4, 1); -template -const dim3 KernelConfigM128::fallback_cluster(1, 2, 1); - -// Config(half_t/bfloat16_t) for M <= 256 -template -struct KernelConfigM256 { - using OutputType = T; - using MmaTileShape = Shape<_256, _256, _256>; - using ClusterShape = Shape; - using EpilogueTile = Shape<_128, _64>; // Avoid register spilling - using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; - using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100; - const static dim3 preferred_cluster; - const static dim3 fallback_cluster; -}; -template -const dim3 KernelConfigM256::preferred_cluster(2, 4, 1); -template -const dim3 KernelConfigM256::fallback_cluster(2, 1, 1); - -// Config(half_t/bfloat16_t) for 256 < M <= 1024 -template -struct KernelConfigDefault { - using OutputType = T; - using MmaTileShape = Shape<_256, _256, _256>; - using ClusterShape = Shape; - using EpilogueTile = Shape<_128, _64>; // Avoid register spilling - using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; - using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100; - const static dim3 preferred_cluster; - const static dim3 fallback_cluster; -}; -template -const dim3 KernelConfigDefault::preferred_cluster(2, 4, 1); -template -const dim3 KernelConfigDefault::fallback_cluster(2, 1, 1); - -// Config(half_t/bfloat16_t) for M > 1024: 1x4 cluster reduces M-tail waste. -template -struct KernelConfigLargeM { - using OutputType = T; - using MmaTileShape = Shape<_256, _256, _256>; - using ClusterShape = Shape; - using EpilogueTile = Shape<_128, _64>; - using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; - using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100; - const static dim3 preferred_cluster; - const static dim3 fallback_cluster; -}; -template -const dim3 KernelConfigLargeM::preferred_cluster(1, 4, 1); -template -const dim3 KernelConfigLargeM::fallback_cluster(1, 2, 1); - -struct KernelConfigFp32 { - using OutputType = float; - using MmaTileShape = Shape<_128, _128, _256>; - using ClusterShape = Shape; - using EpilogueTile = cutlass::epilogue::collective::EpilogueTileAuto; - using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; - using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100; - const static dim3 preferred_cluster; - const static dim3 fallback_cluster; -}; -const dim3 KernelConfigFp32::preferred_cluster = dim3(1, 4, 1); -const dim3 KernelConfigFp32::fallback_cluster = dim3(1, 2, 1); - -template -struct Fp4GemmSm100 { - using Config = KernelConfig; - using OutputType = typename KernelConfig::OutputType; - - using ElementA = cutlass::nv_float4_t; - using LayoutATag = cutlass::layout::RowMajor; - static constexpr int AlignmentA = 32; - - using ElementB = cutlass::nv_float4_t; - using LayoutBTag = cutlass::layout::ColumnMajor; - static constexpr int AlignmentB = 32; - - using ElementD = OutputType; - using ElementC = OutputType; - using LayoutCTag = cutlass::layout::RowMajor; - using LayoutDTag = cutlass::layout::RowMajor; - static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; - static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; - - using ElementAccumulator = float; - using ArchTag = cutlass::arch::Sm100; - using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; - - using MmaTileShape = typename KernelConfig::MmaTileShape; - using ClusterShape = typename KernelConfig::ClusterShape; - using EpilogueTile = typename KernelConfig::EpilogueTile; - using EpilogueSchedule = typename KernelConfig::EpilogueSchedule; - using MainloopSchedule = typename KernelConfig::MainloopSchedule; - - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - MmaTileShape, - ClusterShape, - EpilogueTile, - ElementAccumulator, - ElementAccumulator, - void, - LayoutCTag, - AlignmentC, - ElementD, - LayoutDTag, - AlignmentD, - EpilogueSchedule, - cutlass::epilogue::fusion::LinearCombination>::CollectiveOp; - - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - LayoutATag, - AlignmentA, - ElementB, - LayoutBTag, - AlignmentB, - ElementAccumulator, - MmaTileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - MainloopSchedule>::CollectiveOp; - - using GemmKernel = - cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue, void>; - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - using StrideA = typename Gemm::GemmKernel::StrideA; - using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{})); - using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; - using StrideB = typename Gemm::GemmKernel::StrideB; - using LayoutB = decltype(cute::make_layout(make_shape(0, 0, 0), StrideB{})); - using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; - using StrideC = typename Gemm::GemmKernel::StrideC; - using LayoutC = decltype(cute::make_layout(make_shape(0, 0, 0), StrideC{})); - using StrideD = typename Gemm::GemmKernel::StrideD; - using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{})); -}; - -template -typename T::Gemm::Arguments args_from_options( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int64_t M, - int64_t N, - int64_t K) { - using ElementA = typename T::Gemm::ElementA; - using ElementB = typename T::Gemm::ElementB; - using ElementSFA = cutlass::float_ue4m3_t; - using ElementSFB = cutlass::float_ue4m3_t; - using ElementD = typename T::Gemm::ElementD; - using ElementCompute = float; - using StrideA = typename T::StrideA; - using StrideB = typename T::StrideB; - using StrideD = typename T::StrideD; - using Sm1xxBlkScaledConfig = typename T::Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - - int m = static_cast(M); - int n = static_cast(N); - int k = static_cast(K); - auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1}); - auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1}); - auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1}); - - auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, 1)); - auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, 1)); - - typename T::Gemm::Arguments arguments{ - cutlass::gemm::GemmUniversalMode::kGemm, - {m, n, k, 1}, - {// Mainloop arguments - static_cast(A.data_ptr()), - stride_A, - static_cast(B.data_ptr()), - stride_B, - static_cast(A_sf.data_ptr()), - layout_SFA, - static_cast(B_sf.data_ptr()), - layout_SFB}, - { // Epilogue arguments - {}, // epilogue.thread - nullptr, - stride_D, - static_cast(D.data_ptr()), - stride_D}}; - auto& fusion_args = arguments.epilogue.thread; - fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); - using KernelConfig = typename T::Config; - arguments.hw_info.cluster_shape = KernelConfig::preferred_cluster; - arguments.hw_info.cluster_shape_fallback = KernelConfig::fallback_cluster; - return arguments; -} - -template -void runGemm( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int64_t m, - int64_t n, - int64_t k, - cudaStream_t stream) { - typename T::Gemm gemm; - auto arguments = args_from_options(D, A, B, A_sf, B_sf, alpha, m, n, k); - - size_t workspace_size = T::Gemm::get_workspace_size(arguments); - auto workspace_tensor = alloc_workspace_tensor(workspace_size, A.device()); - void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr(); - - CUTLASS_CHECK(gemm.can_implement(arguments)); - - CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream)); - - CUTLASS_CHECK(gemm.run(arguments, workspace, stream)); -} - -template -void cutlassFp4GemmDispatchSm100( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int64_t m, - int64_t n, - int64_t k, - cudaStream_t stream) { - if (m <= 128) { - runGemm>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (m <= 256) { - runGemm>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (m <= 1024) { - // m in (256, 1024]: 2x4 cluster balances SM occupancy and data reuse - runGemm>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else { - // m in (1024, inf): 1x4 cluster eliminates M-tail waste for FLUX-class shapes - runGemm>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } -} - -template <> -void cutlassFp4GemmDispatchSm100( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int64_t m, - int64_t n, - int64_t k, - cudaStream_t stream) { - runGemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); -} - -#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) diff --git a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm120.cuh b/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm120.cuh deleted file mode 100644 index cdb159061..000000000 --- a/python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm120.cuh +++ /dev/null @@ -1,228 +0,0 @@ -/* Copyright 2026 SGLang Team. All Rights Reserved. - -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. -==============================================================================*/ - -#pragma once - -#include "nvfp4_scaled_mm_common.cuh" - -#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) - -struct sm120_fp4_config_small_m { - using ClusterShape = Shape<_1, _1, _1>; - using MmaTileShape = Shape<_128, _128, _256>; - using PerSmTileShape_MNK = Shape<_128, _128, _256>; -}; - -struct sm120_fp4_config_M256 { - using ClusterShape = Shape<_1, _1, _1>; - using MmaTileShape = Shape<_128, _128, _128>; - using PerSmTileShape_MNK = Shape<_128, _128, _128>; -}; - -struct sm120_fp4_config_default { - using ClusterShape = Shape<_1, _1, _1>; - using MmaTileShape = Shape<_256, _128, _128>; - using PerSmTileShape_MNK = Shape<_256, _128, _128>; -}; - -template -struct Fp4GemmSm120 { - using ElementA = cutlass::nv_float4_t; - using LayoutATag = cutlass::layout::RowMajor; - static constexpr int AlignmentA = 32; - - using ElementB = cutlass::nv_float4_t; - using LayoutBTag = cutlass::layout::ColumnMajor; - static constexpr int AlignmentB = 32; - - using ElementD = OutType; - using ElementC = OutType; - using LayoutCTag = cutlass::layout::RowMajor; - using LayoutDTag = cutlass::layout::RowMajor; - static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; - static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; - - using ElementAccumulator = float; - using ArchTag = cutlass::arch::Sm120; - using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; - - using MmaTileShape = typename Config::MmaTileShape; - using ClusterShape = typename Config::ClusterShape; - using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK; - - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - PerSmTileShape_MNK, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulator, - ElementAccumulator, - void, - LayoutCTag, - AlignmentC, - ElementD, - LayoutDTag, - AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; - - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - LayoutATag, - AlignmentA, - ElementB, - LayoutBTag, - AlignmentB, - ElementAccumulator, - MmaTileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; - - using GemmKernel = - cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue, void>; - - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; -}; - -template -typename Gemm::Arguments args_from_options_sm120( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int M, - int N, - int K) { - using ElementA = typename Gemm::ElementA; - using ElementB = typename Gemm::ElementB; - using ElementD = typename Gemm::ElementD; - using ElementSFA = cutlass::float_ue4m3_t; - using ElementSFB = cutlass::float_ue4m3_t; - using ElementCompute = float; - - using StrideA = typename Gemm::GemmKernel::StrideA; - using StrideB = typename Gemm::GemmKernel::StrideB; - using StrideC = typename Gemm::GemmKernel::StrideC; - using StrideD = typename Gemm::GemmKernel::StrideD; - - using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - - auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1}); - auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1}); - auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1}); - - auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1)); - auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); - - typename Gemm::Arguments arguments{ - cutlass::gemm::GemmUniversalMode::kGemm, - {M, N, K, 1}, - {static_cast(A.data_ptr()), - stride_A, - static_cast(B.data_ptr()), - stride_B, - static_cast(A_sf.data_ptr()), - layout_SFA, - static_cast(B_sf.data_ptr()), - layout_SFB}, - {{}, nullptr, stride_D, static_cast(D.data_ptr()), stride_D}}; - auto& fusion_args = arguments.epilogue.thread; - fusion_args.alpha_ptr = static_cast(alpha.data_ptr()); - - return arguments; -} - -template -void runGemmSm120( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int M, - int N, - int K, - cudaStream_t stream) { - Gemm gemm; - - auto arguments = args_from_options_sm120(D, A, B, A_sf, B_sf, alpha, M, N, K); - - size_t workspace_size = Gemm::get_workspace_size(arguments); - auto workspace_tensor = alloc_workspace_tensor(workspace_size, A.device()); - void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr(); - - CUTLASS_CHECK(gemm.can_implement(arguments)); - - CUTLASS_CHECK(gemm.initialize(arguments, workspace, stream)); - - CUTLASS_CHECK(gemm.run(arguments, workspace, stream)); -} - -void cutlass_fp4_bf16_gemm_dispatch_sm120( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int m, - int n, - int k, - cudaStream_t stream) { - uint32_t const mp2 = std::max(static_cast(16), next_pow_2(m)); - if (mp2 <= 32) { - runGemmSm120::Gemm>( - D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (mp2 <= 256) { - runGemmSm120::Gemm>( - D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else { - runGemmSm120::Gemm>( - D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } -} - -void cutlass_fp4_f16_gemm_dispatch_sm120( - tvm::ffi::TensorView D, - tvm::ffi::TensorView A, - tvm::ffi::TensorView B, - tvm::ffi::TensorView A_sf, - tvm::ffi::TensorView B_sf, - tvm::ffi::TensorView alpha, - int m, - int n, - int k, - cudaStream_t stream) { - uint32_t const mp2 = std::max(static_cast(16), next_pow_2(m)); - if (mp2 <= 32) { - runGemmSm120::Gemm>( - D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else if (mp2 <= 256) { - runGemmSm120::Gemm>( - D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } else { - runGemmSm120::Gemm>( - D, A, B, A_sf, B_sf, alpha, m, n, k, stream); - } -} - -#endif // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED) diff --git a/python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh b/python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh deleted file mode 100644 index c3293fbfd..000000000 --- a/python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh +++ /dev/null @@ -1,882 +0,0 @@ -#include -#include - -#include -#include - -#include -#include - -#include "cute/tensor.hpp" -#include "cutlass/epilogue/collective/collective_builder.hpp" -#include "cutlass/epilogue/collective/default_epilogue.hpp" -#include "cutlass/epilogue/thread/linear_combination.h" -#include "cutlass/gemm/collective/collective_builder.hpp" -#include "cutlass/gemm/device/gemm_universal_adapter.h" -#include "cutlass/gemm/dispatch_policy.hpp" -#include "cutlass/gemm/group_array_problem_shape.hpp" -#include "cutlass/gemm/kernel/gemm_universal.hpp" -#include "cutlass/tensor_ref.h" -#include "cutlass/util/command_line.h" -#include "cutlass/util/distribution.h" -#include "cutlass/util/host_tensor.h" -#include "cutlass/util/packed_stride.hpp" -#include "cutlass/util/reference/device/gemm.h" -#include "cutlass/util/reference/device/tensor_compare.h" -#include "cutlass/util/reference/host/gett.hpp" -#include "cutlass/util/reference/host/tensor_compare.h" -#include "cutlass/util/reference/host/tensor_fill.h" -#include "cutlass/util/reference/host/tensor_norm.h" -#include "cutlass/util/tensor_view_io.h" -#include -#include -#include -#include -#include - -using namespace host; -using namespace cute; - -struct WorkspaceKey { - int device_id; - uintptr_t stream; - auto operator==(const WorkspaceKey&) const -> bool = default; -}; - -struct WorkspaceKeyHash { - auto operator()(const WorkspaceKey& key) const -> size_t { - size_t h1 = std::hash{}(key.device_id); - size_t h2 = std::hash{}(key.stream); - return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); - } -}; - -struct WorkspaceState { - void* ptr = nullptr; - size_t bytes = 0; -}; - -inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* { - if (required_bytes == 0) { - return nullptr; - } - - thread_local std::unordered_map cache; - WorkspaceKey key{device_id, reinterpret_cast(stream)}; - auto& ws = cache[key]; - - if (ws.ptr != nullptr && ws.bytes >= required_bytes) { - return ws.ptr; - } - - RuntimeDeviceCheck(cudaSetDevice(device_id)); - if (ws.ptr != nullptr) { - RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream)); - ws.ptr = nullptr; - ws.bytes = 0; - } - RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream)); - ws.bytes = required_bytes; - return ws.ptr; -} - -inline int getSMVersion(int device_id) { - int sm_major = 0; - int sm_minor = 0; - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id)); - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id)); - return sm_major * 10 + sm_minor; -} - -template < - typename ElementAB, - typename ElementC, - typename ElementSF, - typename ElementAccumulator, - typename LayoutSFA, - typename LayoutSFB, - typename ScaleConfig> -__global__ void __get_group_gemm_starts( - ElementAB** a_offsets, - ElementAB** b_offsets, - ElementC** out_offsets, - ElementSF** a_scales_offsets, - ElementSF** b_scales_offsets, - ElementAccumulator** alpha_offsets, - LayoutSFA* layout_sfa_base_as_int, - LayoutSFB* layout_sfb_base_as_int, - ElementAB* a_base_as_int, - ElementAB* b_base_as_int, - ElementC* out_base_as_int, - ElementSF* a_scales_base_as_int, - ElementSF* b_scales_base_as_int, - ElementAccumulator* alphas_base_as_int, - const int32_t* expert_offsets, - const int32_t* sf_offsets, - const int32_t* problem_sizes_as_shapes, - const int K, - const int N) { - int64_t expert_id = threadIdx.x; - if (expert_id >= gridDim.x * blockDim.x) { - return; - } - // Originally int32_t but upcasting to int64_t to avoid overflow - // during offset calculations - int64_t expert_offset = static_cast(expert_offsets[expert_id]); - int64_t sf_offset = static_cast(sf_offsets[expert_id]); - // size for block in block scale. - int64_t group_size = 16; - int64_t m = static_cast(problem_sizes_as_shapes[expert_id * 3]); - int64_t n = static_cast(problem_sizes_as_shapes[expert_id * 3 + 1]); - int64_t k = static_cast(problem_sizes_as_shapes[expert_id * 3 + 2]); - assert((m >= 0 && n == N && k == K && k % 2 == 0) && "unexpected problem sizes"); - - int64_t half_k = static_cast(k / 2); - int64_t group_k = static_cast(k / group_size); - // Shape of A as uint8/byte = [M, K // 2] - // Shape of B as uint8/byte = [E, N, K // 2] - a_offsets[expert_id] = a_base_as_int + expert_offset * half_k; - - b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k; - // Shape of C = [M, N] - out_offsets[expert_id] = out_base_as_int + expert_offset * n; - // Shape of a_scale = [sum(sf_sizes), K // group_size] - a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k; - - assert((reinterpret_cast(a_scales_offsets[expert_id]) % 128) == 0 && "TMA requires 128-byte alignment"); - - // Shape of B scale = [E, N, K // group_size] - b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k; - assert((reinterpret_cast(b_scales_offsets[expert_id]) % 128) == 0 && "TMA requires 128-byte alignment"); - // Shape of alpha = [E] - alpha_offsets[expert_id] = alphas_base_as_int + expert_id; - - LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id; - LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id; - - *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA( - cute::make_shape(static_cast(m), static_cast(n), static_cast(k), 1)); - *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB( - cute::make_shape(static_cast(m), static_cast(n), static_cast(k), 1)); -} - -#define __CALL_GET_STARTS_KERNEL_BLOCKSCALE( \ - ELEMENT_AB_TYPE, SF_TYPE, TYPE_CHECK, C_TYPE, LayoutSFA, LayoutSFB, ScaleConfig) \ - else if (TYPE_CHECK) { \ - __get_group_gemm_starts \ - <<<1, num_experts, 0, stream>>>( \ - static_cast(a_starts.data_ptr()), \ - static_cast(b_starts.data_ptr()), \ - static_cast(out_starts.data_ptr()), \ - static_cast(a_scales_starts.data_ptr()), \ - static_cast(b_scales_starts.data_ptr()), \ - static_cast(alpha_starts.data_ptr()), \ - reinterpret_cast(layout_sfa.data_ptr()), \ - reinterpret_cast(layout_sfb.data_ptr()), \ - static_cast(a_tensors.data_ptr()), \ - static_cast(b_tensors.data_ptr()), \ - static_cast(out_tensors.data_ptr()), \ - static_cast(a_scales.data_ptr()), \ - static_cast(b_scales.data_ptr()), \ - static_cast(alphas.data_ptr()), \ - static_cast(expert_offsets.data_ptr()), \ - static_cast(sf_offsets.data_ptr()), \ - static_cast(problem_sizes.data_ptr()), \ - K, \ - N); \ - } - -template -void run_get_group_gemm_starts( - const tvm::ffi::TensorView a_starts, - const tvm::ffi::TensorView b_starts, - const tvm::ffi::TensorView out_starts, - const tvm::ffi::TensorView a_scales_starts, - const tvm::ffi::TensorView b_scales_starts, - const tvm::ffi::TensorView alpha_starts, - const tvm::ffi::TensorView layout_sfa, - const tvm::ffi::TensorView layout_sfb, - /*these are used for their base addresses*/ - tvm::ffi::TensorView const& a_tensors, - tvm::ffi::TensorView const& b_tensors, - tvm::ffi::TensorView const& out_tensors, - tvm::ffi::TensorView const& a_scales, - tvm::ffi::TensorView const& b_scales, - tvm::ffi::TensorView const& alphas, - tvm::ffi::TensorView const& expert_offsets, - tvm::ffi::TensorView const& sf_offsets, - tvm::ffi::TensorView const& problem_sizes, - int M, - int N, - int K) { - int num_experts = static_cast(expert_offsets.size(0)); - auto stream = LaunchKernel::resolve_device(a_tensors.device()); - - RuntimeCheck(out_tensors.size(1) == N, "Output tensor shape doesn't match expected shape"); - RuntimeCheck( - K / 2 == b_tensors.size(2), - "b_tensors(dim = 2) and a_tensors(dim = 1) trailing" - " dimension must match"); - if (false) { - } - //(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB, - // ScaleConfig) - __CALL_GET_STARTS_KERNEL_BLOCKSCALE( - cutlass::float_e2m1_t, - cutlass::float_ue4m3_t, - host::is_type(out_tensors.dtype()), - cutlass::bfloat16_t, - LayoutSFA, - LayoutSFB, - ScaleConfig) - __CALL_GET_STARTS_KERNEL_BLOCKSCALE( - cutlass::float_e2m1_t, - cutlass::float_ue4m3_t, - host::is_type(out_tensors.dtype()), - cutlass::half_t, - LayoutSFA, - LayoutSFB, - ScaleConfig) - else { - Panic("Invalid output type (must be float16 or bfloat16)"); - } -} - -void run_fp4_blockwise_scaled_group_mm_sm120( - tvm::ffi::TensorView output, - const tvm::ffi::TensorView a, - const tvm::ffi::TensorView b, - const tvm::ffi::TensorView a_blockscale, - const tvm::ffi::TensorView b_blockscales, - const tvm::ffi::TensorView alphas, - const tvm::ffi::TensorView ab_strides, - const tvm::ffi::TensorView c_strides, - const tvm::ffi::TensorView problem_sizes, - const tvm::ffi::TensorView expert_offsets, - const tvm::ffi::TensorView sf_offsets, - const tvm::ffi::TensorView a_ptrs, - const tvm::ffi::TensorView b_ptrs, - const tvm::ffi::TensorView out_ptrs, - const tvm::ffi::TensorView a_scales_ptrs, - const tvm::ffi::TensorView b_scales_ptrs, - const tvm::ffi::TensorView alpha_ptrs, - const tvm::ffi::TensorView layout_sfa, - const tvm::ffi::TensorView layout_sfb, - int M, - int N, - int K) { - using ProblemShape = cutlass::gemm::GroupProblemShape>; - using ElementType = cutlass::float_e2m1_t; - using ElementSFType = cutlass::float_ue4m3_t; - using ElementA = cutlass::nv_float4_t; - using ElementB = cutlass::nv_float4_t; - - using ElementC = cutlass::bfloat16_t; - using ElementD = cutlass::bfloat16_t; - using ElementAccumulator = float; - // Layout definitions - using LayoutA = cutlass::layout::RowMajor; - using LayoutB = cutlass::layout::ColumnMajor; - using LayoutC = cutlass::layout::RowMajor; - using LayoutD = cutlass::layout::RowMajor; - - // Alignment constraints - static constexpr int AlignmentA = 32; - static constexpr int AlignmentB = 32; - static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; - static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; - - // Architecture definitions - using ArchTag = cutlass::arch::Sm120; - using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; - using StageCountType = cutlass::gemm::collective::StageCountAuto; - using ThreadBlockShape = Shape<_128, _128, _128>; - // on the tile size - - using ClusterShape = Shape<_1, _1, _1>; - - using FusionOperation = - cutlass::epilogue::fusion::LinearCombination; - - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ThreadBlockShape, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - ElementAccumulator, - ElementAccumulator, - ElementC, - LayoutC*, - AlignmentC, - ElementD, - LayoutC*, - AlignmentD, - cutlass::epilogue::collective::EpilogueScheduleAuto, - FusionOperation>::CollectiveOp; - - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - OperatorClass, - ElementA, - LayoutA*, - AlignmentA, - ElementB, - LayoutB*, - AlignmentB, - ElementAccumulator, - ThreadBlockShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong>::CollectiveOp; - - using GemmKernel = cutlass::gemm::kernel::GemmUniversal; - - using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter; - using Gemm = Gemm1SM; - using StrideA = typename Gemm::GemmKernel::InternalStrideA; - using StrideB = typename Gemm::GemmKernel::InternalStrideB; - using StrideC = typename Gemm::GemmKernel::InternalStrideC; - using StrideD = typename Gemm::GemmKernel::InternalStrideD; - - using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; - using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; - using ScaleConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - - using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; - int num_experts = static_cast(expert_offsets.size(0)); - - run_get_group_gemm_starts( - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb, - a, - b, - output, - a_blockscale, - b_blockscales, - alphas, - expert_offsets, - sf_offsets, - problem_sizes, - M, - N, - K); - - // Create an instance of the GEMM - Gemm gemm_op; - - // Initialize problem_sizes_as_shapes correctly - UnderlyingProblemShape* problem_sizes_as_shapes = static_cast(problem_sizes.data_ptr()); - - // Set the Scheduler info - cutlass::KernelHardwareInfo hw_info; - - using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions; - typename Gemm::GemmKernel::TileSchedulerArguments scheduler; - scheduler.raster_order = RasterOrderOptions::AlongM; - hw_info.device_id = a.device().device_id; - static std::unordered_map cached_sm_counts; - if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { - cached_sm_counts[hw_info.device_id] = - cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); - } - hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits::max()); - - // Mainloop Arguments - typename GemmKernel::MainloopArguments mainloop_args{ - static_cast(a_ptrs.data_ptr()), - static_cast(ab_strides.data_ptr()), - static_cast(b_ptrs.data_ptr()), - static_cast(ab_strides.data_ptr()), - static_cast(a_scales_ptrs.data_ptr()), - reinterpret_cast(layout_sfa.data_ptr()), - static_cast(b_scales_ptrs.data_ptr()), - reinterpret_cast(layout_sfb.data_ptr())}; - - // Epilogue Arguments - typename GemmKernel::EpilogueArguments epilogue_args{ - {}, // epilogue.thread - nullptr, - static_cast(c_strides.data_ptr()), - static_cast(out_ptrs.data_ptr()), - static_cast(c_strides.data_ptr())}; - auto& fusion_args = epilogue_args.thread; - fusion_args.alpha_ptr_array = reinterpret_cast(alpha_ptrs.data_ptr()); - fusion_args.dAlpha = {_0{}, _0{}, 1}; - fusion_args.beta = 0.0f; - - // Gemm Arguments - typename GemmKernel::Arguments args{ - cutlass::gemm::GemmUniversalMode::kGrouped, - {num_experts, problem_sizes_as_shapes, nullptr}, - mainloop_args, - epilogue_args, - hw_info, - scheduler}; - - size_t workspace_size = Gemm::get_workspace_size(args); - const cudaStream_t stream = LaunchKernel::resolve_device(a.device()); - void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream); - - auto can_implement_status = gemm_op.can_implement(args); - RuntimeCheck( - can_implement_status == cutlass::Status::kSuccess, - "Failed to implement GEMM: ", - cutlassGetStatusString(can_implement_status)); - - // Run the GEMM - auto status = gemm_op.initialize(args, workspace); - RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status)); - - status = gemm_op.run(args, workspace, stream); - RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status)); -} - -template -void run_fp4_blockwise_scaled_group_mm_sm100( - tvm::ffi::TensorView output, - const tvm::ffi::TensorView a, - const tvm::ffi::TensorView b, - const tvm::ffi::TensorView a_blockscale, - const tvm::ffi::TensorView b_blockscales, - const tvm::ffi::TensorView alphas, - const tvm::ffi::TensorView ab_strides, - const tvm::ffi::TensorView c_strides, - const tvm::ffi::TensorView problem_sizes, - const tvm::ffi::TensorView expert_offsets, - const tvm::ffi::TensorView sf_offsets, - const tvm::ffi::TensorView a_ptrs, - const tvm::ffi::TensorView b_ptrs, - const tvm::ffi::TensorView out_ptrs, - const tvm::ffi::TensorView a_scales_ptrs, - const tvm::ffi::TensorView b_scales_ptrs, - const tvm::ffi::TensorView alpha_ptrs, - const tvm::ffi::TensorView layout_sfa, - const tvm::ffi::TensorView layout_sfb, - int M, - int N, - int K) { - using ProblemShape = cutlass::gemm::GroupProblemShape>; - using ElementType = cutlass::float_e2m1_t; - using ElementSFType = cutlass::float_ue4m3_t; - using ElementA = cutlass::nv_float4_t; - using ElementB = cutlass::nv_float4_t; - - using ElementC = OutType; - using ElementD = ElementC; - using ElementAccumulator = float; - // Layout definitions - using LayoutA = cutlass::layout::RowMajor; - using LayoutB = cutlass::layout::ColumnMajor; - using LayoutC = cutlass::layout::RowMajor; - using LayoutD = LayoutC; - - // Alignment constraints - static constexpr int AlignmentA = 32; - static constexpr int AlignmentB = 32; - static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; - static constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; - - // Architecture definitions - using ArchTag = cutlass::arch::Sm100; - using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp; // Epilogue Operator class tag - using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Mainloop Operator class tag - using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based - // on the tile size - - using ClusterShape = Shape<_1, _1, _1>; - struct MMA1SMConfig { - using MmaTileShape = Shape<_128, _128, _128>; - using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100; // Kernel to launch - using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; // Epilogue to launch - }; - - using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - ArchTag, - EpilogueOperatorClass, - typename MMA1SMConfig::MmaTileShape, - ClusterShape, - Shape<_128, _64>, - ElementAccumulator, - ElementAccumulator, - ElementC, - LayoutC*, - AlignmentC, - ElementD, - LayoutC*, - AlignmentD, - typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp; - - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< - ArchTag, - MainloopOperatorClass, - ElementA, - LayoutA*, - AlignmentA, - ElementB, - LayoutB*, - AlignmentB, - ElementAccumulator, - typename MMA1SMConfig::MmaTileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - typename MMA1SMConfig::KernelSchedule>::CollectiveOp; - - using GemmKernel = cutlass::gemm::kernel::GemmUniversal; - - using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter; - using Gemm = Gemm1SM; - using StrideA = typename Gemm::GemmKernel::InternalStrideA; - using StrideB = typename Gemm::GemmKernel::InternalStrideB; - using StrideC = typename Gemm::GemmKernel::InternalStrideC; - using StrideD = typename Gemm::GemmKernel::InternalStrideD; - - using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; - using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; - using ScaleConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; - - using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; - int num_experts = static_cast(expert_offsets.size(0)); - - run_get_group_gemm_starts( - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb, - a, - b, - output, - a_blockscale, - b_blockscales, - alphas, - expert_offsets, - sf_offsets, - problem_sizes, - M, - N, - K); - - // Create an instance of the GEMM - Gemm gemm_op; - - // Initialize problem_sizes_as_shapes correctly - UnderlyingProblemShape* problem_sizes_as_shapes = static_cast(problem_sizes.data_ptr()); - - // Set the Scheduler info - cutlass::KernelHardwareInfo hw_info; - using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm100GroupParams< - typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions; - typename Gemm::GemmKernel::TileSchedulerArguments scheduler; - scheduler.raster_order = RasterOrderOptions::AlongM; - hw_info.device_id = a.device().device_id; - static std::unordered_map cached_sm_counts; - if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { - cached_sm_counts[hw_info.device_id] = - cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); - } - hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits::max()); - - // Mainloop Arguments - typename GemmKernel::MainloopArguments mainloop_args{ - static_cast(a_ptrs.data_ptr()), - static_cast(ab_strides.data_ptr()), - static_cast(b_ptrs.data_ptr()), - static_cast(ab_strides.data_ptr()), - static_cast(a_scales_ptrs.data_ptr()), - reinterpret_cast(layout_sfa.data_ptr()), - static_cast(b_scales_ptrs.data_ptr()), - reinterpret_cast(layout_sfb.data_ptr())}; - - // Epilogue Arguments - typename GemmKernel::EpilogueArguments epilogue_args{ - {}, // epilogue.thread - nullptr, - static_cast(c_strides.data_ptr()), - static_cast(out_ptrs.data_ptr()), - static_cast(c_strides.data_ptr())}; - auto& fusion_args = epilogue_args.thread; - fusion_args.alpha_ptr_array = reinterpret_cast(alpha_ptrs.data_ptr()); - fusion_args.dAlpha = {_0{}, _0{}, 1}; - - // Gemm Arguments - typename GemmKernel::Arguments args{ - cutlass::gemm::GemmUniversalMode::kGrouped, - {num_experts, problem_sizes_as_shapes, nullptr}, - mainloop_args, - epilogue_args, - hw_info, - scheduler}; - - size_t workspace_size = Gemm::get_workspace_size(args); - const cudaStream_t stream = LaunchKernel::resolve_device(a.device()); - void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream); - - auto can_implement_status = gemm_op.can_implement(args); - RuntimeCheck( - can_implement_status == cutlass::Status::kSuccess, - "Failed to implement GEMM: ", - cutlassGetStatusString(can_implement_status)); - - // Run the GEMM - auto status = gemm_op.initialize(args, workspace); - RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status)); - - status = gemm_op.run(args, workspace, stream); - RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status)); -} - -void cutlass_fp4_group_mm_sm100a_sm120a( - tvm::ffi::TensorView output, - const tvm::ffi::TensorView a, - const tvm::ffi::TensorView b, - const tvm::ffi::TensorView a_blockscale, - const tvm::ffi::TensorView b_blockscales, - const tvm::ffi::TensorView alphas, - const tvm::ffi::TensorView ab_strides, - const tvm::ffi::TensorView c_strides, - const tvm::ffi::TensorView problem_sizes, - const tvm::ffi::TensorView expert_offsets, - const tvm::ffi::TensorView sf_offsets, - const tvm::ffi::TensorView a_ptrs, - const tvm::ffi::TensorView b_ptrs, - const tvm::ffi::TensorView out_ptrs, - const tvm::ffi::TensorView a_scales_ptrs, - const tvm::ffi::TensorView b_scales_ptrs, - const tvm::ffi::TensorView alpha_ptrs, - const tvm::ffi::TensorView layout_sfa, - const tvm::ffi::TensorView layout_sfb) { - auto check_cuda_contig = [](const tvm::ffi::TensorView t, const char* name) { - RuntimeCheck(t.device().device_type == kDLCUDA, name, " must be a CUDA tensor"); - RuntimeCheck(t.is_contiguous(), name, " must be contiguous"); - }; - - check_cuda_contig(output, "output"); - check_cuda_contig(a, "a"); - check_cuda_contig(b, "b"); - check_cuda_contig(a_blockscale, "a_blockscale"); - check_cuda_contig(b_blockscales, "b_blockscales"); - check_cuda_contig(alphas, "alphas"); - check_cuda_contig(ab_strides, "ab_strides"); - check_cuda_contig(c_strides, "c_strides"); - check_cuda_contig(problem_sizes, "problem_sizes"); - check_cuda_contig(expert_offsets, "expert_offsets"); - check_cuda_contig(sf_offsets, "sf_offsets"); - check_cuda_contig(a_ptrs, "a_ptrs"); - check_cuda_contig(b_ptrs, "b_ptrs"); - check_cuda_contig(out_ptrs, "out_ptrs"); - check_cuda_contig(a_scales_ptrs, "a_scales_ptrs"); - check_cuda_contig(b_scales_ptrs, "b_scales_ptrs"); - check_cuda_contig(alpha_ptrs, "alpha_ptrs"); - check_cuda_contig(layout_sfa, "layout_sfa"); - check_cuda_contig(layout_sfb, "layout_sfb"); - - RuntimeCheck( - output.device() == a.device() && a.device() == b.device() && a.device() == a_blockscale.device() && - a.device() == b_blockscales.device() && a.device() == alphas.device() && a.device() == ab_strides.device() && - a.device() == c_strides.device() && a.device() == problem_sizes.device() && - a.device() == expert_offsets.device() && a.device() == sf_offsets.device() && a.device() == a_ptrs.device() && - a.device() == b_ptrs.device() && a.device() == out_ptrs.device() && a.device() == a_scales_ptrs.device() && - a.device() == b_scales_ptrs.device() && a.device() == alpha_ptrs.device() && - a.device() == layout_sfa.device() && a.device() == layout_sfb.device(), - "all tensors must be on the same device"); - - RuntimeCheck(host::is_type(a.dtype()), "a must be uint8"); - RuntimeCheck(host::is_type(b.dtype()), "b must be uint8"); - RuntimeCheck(host::is_type(a_blockscale.dtype()), "a_blockscale must be float8_e4m3fn"); - RuntimeCheck(host::is_type(b_blockscales.dtype()), "b_blockscales must be float8_e4m3fn"); - RuntimeCheck(host::is_type(alphas.dtype()), "alphas must be float32"); - RuntimeCheck(host::is_type(ab_strides.dtype()), "ab_strides must be int64"); - RuntimeCheck(host::is_type(c_strides.dtype()), "c_strides must be int64"); - RuntimeCheck(host::is_type(problem_sizes.dtype()), "problem_sizes must be int32"); - RuntimeCheck(host::is_type(expert_offsets.dtype()), "expert_offsets must be int32"); - RuntimeCheck(host::is_type(sf_offsets.dtype()), "sf_offsets must be int32"); - RuntimeCheck(host::is_type(a_ptrs.dtype()), "a_ptrs must be int64"); - RuntimeCheck(host::is_type(b_ptrs.dtype()), "b_ptrs must be int64"); - RuntimeCheck(host::is_type(out_ptrs.dtype()), "out_ptrs must be int64"); - RuntimeCheck(host::is_type(a_scales_ptrs.dtype()), "a_scales_ptrs must be int64"); - RuntimeCheck(host::is_type(b_scales_ptrs.dtype()), "b_scales_ptrs must be int64"); - RuntimeCheck(host::is_type(alpha_ptrs.dtype()), "alpha_ptrs must be int64"); - RuntimeCheck(host::is_type(layout_sfa.dtype()), "layout_sfa must be int64"); - RuntimeCheck(host::is_type(layout_sfb.dtype()), "layout_sfb must be int64"); - RuntimeCheck( - host::is_type(output.dtype()) || host::is_type(output.dtype()), - "output must be bfloat16 or float16"); - - RuntimeCheck(a.dim() == 2, "a must be 2D"); - RuntimeCheck(b.dim() == 3, "b must be 3D"); - RuntimeCheck(a_blockscale.dim() == 2, "a_blockscale must be 2D"); - RuntimeCheck(b_blockscales.dim() == 3, "b_blockscales must be 3D"); - RuntimeCheck(alphas.dim() == 1, "alphas must be 1D"); - RuntimeCheck(ab_strides.dim() == 1, "ab_strides must be 1D"); - RuntimeCheck(c_strides.dim() == 1, "c_strides must be 1D"); - RuntimeCheck(problem_sizes.dim() == 2, "problem_sizes must be 2D"); - RuntimeCheck(expert_offsets.dim() == 1, "expert_offsets must be 1D"); - RuntimeCheck(sf_offsets.dim() == 1, "sf_offsets must be 1D"); - RuntimeCheck(a_ptrs.dim() == 1, "a_ptrs must be 1D"); - RuntimeCheck(b_ptrs.dim() == 1, "b_ptrs must be 1D"); - RuntimeCheck(out_ptrs.dim() == 1, "out_ptrs must be 1D"); - RuntimeCheck(a_scales_ptrs.dim() == 1, "a_scales_ptrs must be 1D"); - RuntimeCheck(b_scales_ptrs.dim() == 1, "b_scales_ptrs must be 1D"); - RuntimeCheck(alpha_ptrs.dim() == 1, "alpha_ptrs must be 1D"); - RuntimeCheck(layout_sfa.dim() == 2, "layout_sfa must be 2D"); - RuntimeCheck(layout_sfb.dim() == 2, "layout_sfb must be 2D"); - RuntimeCheck(problem_sizes.size(1) == 3, "problem_sizes must have shape (num_experts, 3)"); - - const int num_experts = static_cast(expert_offsets.size(0)); - RuntimeCheck(problem_sizes.size(0) == num_experts, "problem_sizes size mismatch with expert_offsets"); - RuntimeCheck(sf_offsets.size(0) == num_experts, "sf_offsets size mismatch with expert_offsets"); - RuntimeCheck(alphas.size(0) == num_experts, "alphas size mismatch with expert_offsets"); - RuntimeCheck(ab_strides.size(0) == num_experts, "ab_strides size mismatch with expert_offsets"); - RuntimeCheck(c_strides.size(0) == num_experts, "c_strides size mismatch with expert_offsets"); - RuntimeCheck(a_ptrs.size(0) == num_experts, "a_ptrs size mismatch with expert_offsets"); - RuntimeCheck(b_ptrs.size(0) == num_experts, "b_ptrs size mismatch with expert_offsets"); - RuntimeCheck(out_ptrs.size(0) == num_experts, "out_ptrs size mismatch with expert_offsets"); - RuntimeCheck(a_scales_ptrs.size(0) == num_experts, "a_scales_ptrs size mismatch with expert_offsets"); - RuntimeCheck(b_scales_ptrs.size(0) == num_experts, "b_scales_ptrs size mismatch with expert_offsets"); - RuntimeCheck(alpha_ptrs.size(0) == num_experts, "alpha_ptrs size mismatch with expert_offsets"); - RuntimeCheck(layout_sfa.size(0) == num_experts && layout_sfa.size(1) == 5, "layout_sfa must be [num_experts, 5]"); - RuntimeCheck(layout_sfb.size(0) == num_experts && layout_sfb.size(1) == 5, "layout_sfb must be [num_experts, 5]"); - - int M = static_cast(a.size(0)); - int N = static_cast(b.size(1)); - int K = static_cast(2 * b.size(2)); - RuntimeCheck(output.dim() == 2, "output must be 2D"); - RuntimeCheck(output.size(0) == M && output.size(1) == N, "output shape mismatch"); - - auto sm_version = getSMVersion(a.device().device_id); - if (sm_version == 100 || sm_version == 103) { - if (host::is_type(output.dtype())) { - run_fp4_blockwise_scaled_group_mm_sm100( - output, - a, - b, - a_blockscale, - b_blockscales, - alphas, - ab_strides, - c_strides, - problem_sizes, - expert_offsets, - sf_offsets, - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb, - M, - N, - K); - } else { - run_fp4_blockwise_scaled_group_mm_sm100( - output, - a, - b, - a_blockscale, - b_blockscales, - alphas, - ab_strides, - c_strides, - problem_sizes, - expert_offsets, - sf_offsets, - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb, - M, - N, - K); - } - } else if (sm_version >= 120) { - if (host::is_type(output.dtype())) { - run_fp4_blockwise_scaled_group_mm_sm120( - output, - a, - b, - a_blockscale, - b_blockscales, - alphas, - ab_strides, - c_strides, - problem_sizes, - expert_offsets, - sf_offsets, - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb, - M, - N, - K); - } else { - Panic("SM120 path currently supports only bfloat16 output"); - } - } else { - RuntimeCheck(false, "Unsupported SM version: ", sm_version); - } -} - -void cutlass_fp4_group_mm( - tvm::ffi::TensorView output, - const tvm::ffi::TensorView a, - const tvm::ffi::TensorView b, - const tvm::ffi::TensorView a_blockscale, - const tvm::ffi::TensorView b_blockscales, - const tvm::ffi::TensorView alphas, - const tvm::ffi::TensorView ab_strides, - const tvm::ffi::TensorView c_strides, - const tvm::ffi::TensorView problem_sizes, - const tvm::ffi::TensorView expert_offsets, - const tvm::ffi::TensorView sf_offsets, - const tvm::ffi::TensorView a_ptrs, - const tvm::ffi::TensorView b_ptrs, - const tvm::ffi::TensorView out_ptrs, - const tvm::ffi::TensorView a_scales_ptrs, - const tvm::ffi::TensorView b_scales_ptrs, - const tvm::ffi::TensorView alpha_ptrs, - const tvm::ffi::TensorView layout_sfa, - const tvm::ffi::TensorView layout_sfb) { - cutlass_fp4_group_mm_sm100a_sm120a( - output, - a, - b, - a_blockscale, - b_blockscales, - alphas, - ab_strides, - c_strides, - problem_sizes, - expert_offsets, - sf_offsets, - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb); -} diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py deleted file mode 100644 index a5d4e9d20..000000000 --- a/python/sglang/jit_kernel/nvfp4.py +++ /dev/null @@ -1,636 +0,0 @@ -from __future__ import annotations - -import os -from typing import TYPE_CHECKING, Optional, Tuple - -import torch - -from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch -from sglang.kernel_api_logging import debug_kernel_api -from sglang.srt.utils.custom_op import register_custom_op - -if TYPE_CHECKING: - from tvm_ffi.module import Module - - -_FLOAT4_E2M1_MAX = 6.0 -_FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max - - -def _nvfp4_cuda_flags() -> list[str]: - return [ - "-DNDEBUG", - "-DFLASHINFER_ENABLE_F16", - "-DCUTE_USE_PACKED_TUPLE=1", - "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", - "-DCUTLASS_VERSIONS_GENERATED", - "-DCUTLASS_TEST_LEVEL=0", - "-DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1", - "-DCUTLASS_DEBUG_TRACE_LEVEL=0", - "--expt-extended-lambda", - ] - - -def _nvfp4_arch_env(): - if not torch.cuda.is_available(): - raise RuntimeError("NVFP4 JIT kernels require CUDA.") - major, minor = torch.cuda.get_device_capability() - if major < 10: - raise RuntimeError( - f"NVFP4 JIT kernels require compute capability >= 10.0, got {major}.{minor}." - ) - # NVFP4 kernels use architecture-family-specific instructions and must be - # compiled for `sm_*a` targets (e.g. sm_100a), not plain sm_100. - # JIT compilation targets only the current device, unlike AOT fat-binaries; - # adding extra architectures here would clash with the single SGL_CUDA_ARCH - # value injected by load_jit(). - return override_jit_cuda_arch(major, minor, suffix="a") - - -@torch.compiler.disable -def prewarm_nvfp4_jit_modules( - *, include_expert_quant: bool = False, include_blockwise_moe: bool = False -) -> None: - """Materialize NVFP4 JIT modules before torch.compile traces the model.""" - _jit_nvfp4_quant_module() - _jit_nvfp4_scaled_mm_module() - if include_expert_quant: - _jit_nvfp4_expert_quant_module() - if include_blockwise_moe: - _jit_nvfp4_blockwise_moe_module() - - -@cache_once -def _jit_nvfp4_quant_module() -> Module: - with _nvfp4_arch_env(): - return load_jit( - "nvfp4_quant", - cuda_files=[ - "gemm/nvfp4/nvfp4_quant_kernels.cuh", - ], - cuda_wrappers=[ - ("scaled_fp4_quant", "scaled_fp4_quant_sm100a_sm120a"), - ], - extra_cuda_cflags=_nvfp4_cuda_flags(), - extra_dependencies=["cutlass"], - ) - - -@cache_once -def _jit_nvfp4_expert_quant_module() -> Module: - with _nvfp4_arch_env(): - return load_jit( - "nvfp4_expert_quant", - cuda_files=[ - "gemm/nvfp4/nvfp4_expert_quant.cuh", - ], - cuda_wrappers=[ - ("scaled_fp4_experts_quant", "scaled_fp4_experts_quant_sm100a"), - ( - "silu_and_mul_scaled_fp4_experts_quant", - "silu_and_mul_scaled_fp4_experts_quant_sm100a", - ), - ( - "silu_and_mul_scaled_fp4_experts_quant_packed", - "silu_and_mul_scaled_fp4_experts_quant_packed_sm100a", - ), - ], - extra_dependencies=["cutlass"], - extra_cuda_cflags=_nvfp4_cuda_flags(), - ) - - -@cache_once -def _jit_nvfp4_scaled_mm_module() -> Module: - with _nvfp4_arch_env(): - return load_jit( - "nvfp4_scaled_mm", - cuda_files=[ - "gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh", - "gemm/nvfp4/nvfp4_scaled_mm_entry.cuh", - ], - cuda_wrappers=[("cutlass_scaled_fp4_mm", "cutlass_scaled_fp4_mm")], - extra_dependencies=["cutlass"], - extra_cuda_cflags=_nvfp4_cuda_flags(), - ) - - -@cache_once -def _jit_nvfp4_blockwise_moe_module() -> Module: - with _nvfp4_arch_env(): - return load_jit( - "nvfp4_blockwise_moe", - cuda_files=[ - "moe/nvfp4_blockwise_moe.cuh", - ], - cuda_wrappers=[ - ("cutlass_fp4_group_mm", "cutlass_fp4_group_mm_sm100a_sm120a") - ], - extra_dependencies=["cutlass"], - extra_cuda_cflags=_nvfp4_cuda_flags(), - ) - - -@debug_kernel_api -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: - assert a.ndim == 2 and b.ndim == 2 - m, n = a.shape[0], b.shape[0] - out = torch.empty((m, n), dtype=out_dtype, device=a.device) - module = _jit_nvfp4_scaled_mm_module() - module.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b, alpha) - return out - - -@debug_kernel_api -def cutlass_fp4_group_mm( - a_fp4: torch.Tensor, - b_fp4: torch.Tensor, - a_blockscale: torch.Tensor, - b_blockscale: torch.Tensor, - alphas: torch.Tensor, - out_dtype: torch.dtype, - params: dict[str, torch.Tensor], -) -> torch.Tensor: - m_topk = a_fp4.shape[0] - n = b_fp4.shape[1] - output = torch.empty((m_topk, n), device=a_fp4.device, dtype=out_dtype) - num_experts = int(params["expert_offsets"].numel()) - device = a_fp4.device - - # Backward compatibility: older callers may not pass scratch tensors. - a_ptrs = params.get( - "a_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) - ) - b_ptrs = params.get( - "b_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) - ) - out_ptrs = params.get( - "out_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) - ) - a_scales_ptrs = params.get( - "a_scales_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) - ) - b_scales_ptrs = params.get( - "b_scales_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) - ) - alpha_ptrs = params.get( - "alpha_ptrs", torch.empty((num_experts,), dtype=torch.int64, device=device) - ) - layout_sfa = params.get( - "layout_sfa", torch.empty((num_experts, 5), dtype=torch.int64, device=device) - ) - layout_sfb = params.get( - "layout_sfb", torch.empty((num_experts, 5), dtype=torch.int64, device=device) - ) - - _cutlass_fp4_group_mm_custom_op( - output, - a_fp4, - b_fp4, - a_blockscale, - b_blockscale, - alphas, - params["ab_strides"], - params["c_strides"], - params["problem_sizes"], - params["expert_offsets"], - params["blockscale_offsets"], - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb, - ) - return output - - -@register_custom_op( - op_name="scaled_fp4_quant", - mutates_args=["output", "output_scale"], -) -def _scaled_fp4_quant_custom_op( - input: torch.Tensor, - output: torch.Tensor, - output_scale: torch.Tensor, - input_global_scale: torch.Tensor, -) -> None: - module = _jit_nvfp4_quant_module() - module.scaled_fp4_quant(output, input, output_scale, input_global_scale) - - -@debug_kernel_api -def scaled_fp4_quant( - input: torch.Tensor, input_global_scale: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - """Quantize input tensor to FP4 and return packed FP4 tensor + swizzled scales.""" - assert input.ndim >= 1, f"input.ndim needs to be >= 1, but got {input.ndim}." - other_dims = 1 if input.ndim == 1 else -1 - input = input.reshape(other_dims, input.shape[-1]) - m, n = input.shape - block_size = 16 - device = input.device - - assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}." - assert input.dtype in ( - torch.float16, - torch.bfloat16, - ), f"input.dtype needs to be fp16 or bf16 but got {input.dtype}." - - output = torch.empty((m, n // 2), device=device, dtype=torch.uint8) - - rounded_m = ((m + 128 - 1) // 128) * 128 - scale_n = n // block_size - rounded_n = ((scale_n + 4 - 1) // 4) * 4 - if rounded_n > scale_n: - output_scale = torch.zeros( - (rounded_m, rounded_n // 4), device=device, dtype=torch.int32 - ) - else: - output_scale = torch.empty( - (rounded_m, rounded_n // 4), device=device, dtype=torch.int32 - ) - - _scaled_fp4_quant_custom_op(input, output, output_scale, input_global_scale) - output_scale = output_scale.view(torch.float8_e4m3fn) - return output, output_scale - - -def _shuffle_rows_torch( - input_tensor: torch.Tensor, - dst2src_map: torch.Tensor, - output_tensor_shape: tuple[int, int], -) -> torch.Tensor: - # Keep compatibility when sgl-kernel is slimmed and shuffle_rows may not be present. - output = input_tensor.index_select(0, dst2src_map.to(dtype=torch.int64)) - return output.view(output_tensor_shape) - - -@register_custom_op( - op_name="scaled_fp4_experts_quant", - mutates_args=["output", "output_scales"], -) -def _scaled_fp4_experts_quant_custom_op( - output: torch.Tensor, - output_scales: torch.Tensor, - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - expert_offsets: torch.Tensor, - blockscale_offsets: torch.Tensor, -) -> None: - module = _jit_nvfp4_expert_quant_module() - module.scaled_fp4_experts_quant( - output, - output_scales, - input_tensor, - input_global_scale, - expert_offsets, - blockscale_offsets, - ) - - -@debug_kernel_api -def scaled_fp4_experts_quant( - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - expert_offsets: torch.Tensor, - blockscale_offsets: torch.Tensor, - topk: int, - expert_map: Optional[torch.Tensor] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize packed MoE activations to NVFP4.""" - assert ( - input_tensor.ndim == 2 - ), f"input.ndim needs to be == 2, but got {input_tensor.ndim}." - if expert_map is not None: - m, k = input_tensor.shape - output_tensor_shape = (m * topk, k) - input_tensor = _shuffle_rows_torch( - input_tensor, expert_map, output_tensor_shape - ) - - m_numtopk, k = input_tensor.shape - max_tokens_per_expert = int(os.environ.get("MODELOPT_MAX_TOKENS_PER_EXPERT", 65536)) - assert m_numtopk <= max_tokens_per_expert * topk, ( - f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT({max_tokens_per_expert})" - f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" - " MODELOPT_MAX_TOKENS_PER_EXPERT to set this value." - ) - scales_k = k // 16 - # output_scales is int32-packed FP8 scales, so second dim is in int32 units. - padded_k_in_int32 = (scales_k + 3) // 4 - - output = torch.empty( - m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 - ) - if padded_k_in_int32 * 4 > scales_k: - output_scales = torch.zeros( - max_tokens_per_expert * topk, - padded_k_in_int32, - dtype=torch.int32, - device=input_tensor.device, - ) - else: - output_scales = torch.empty( - max_tokens_per_expert * topk, - padded_k_in_int32, - dtype=torch.int32, - device=input_tensor.device, - ) - - _scaled_fp4_experts_quant_custom_op( - output, - output_scales, - input_tensor, - input_global_scale, - expert_offsets, - blockscale_offsets, - ) - output_scales = output_scales.view(torch.float8_e4m3fn) - return output, output_scales - - -@register_custom_op( - op_name="silu_and_mul_scaled_fp4_experts_quant_packed", - mutates_args=["output", "output_scales"], -) -def _silu_and_mul_scaled_fp4_experts_quant_packed_custom_op( - output: torch.Tensor, - output_scales: torch.Tensor, - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - expert_offsets: torch.Tensor, - blockscale_offsets: torch.Tensor, -) -> None: - module = _jit_nvfp4_expert_quant_module() - module.silu_and_mul_scaled_fp4_experts_quant_packed( - output, - output_scales, - input_tensor, - input_global_scale, - expert_offsets, - blockscale_offsets, - ) - - -@debug_kernel_api -def silu_and_mul_scaled_fp4_experts_quant_packed( - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - expert_offsets: torch.Tensor, - blockscale_offsets: torch.Tensor, - topk: int, - expert_map: Optional[torch.Tensor] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Fused SiLU+mul then FP4 quant for packed MoE inputs (expert_offsets aware). - - Input shape is (m, 2*k) — gate+up concatenated. The kernel does SiLU(gate)*up - then FP4-quantizes the k-dim result. - """ - assert ( - input_tensor.ndim == 2 - ), f"input.ndim needs to be == 2, but got {input_tensor.ndim}." - if expert_map is not None: - m, k = input_tensor.shape - output_tensor_shape = (m * topk, k) - input_tensor = _shuffle_rows_torch( - input_tensor, expert_map, output_tensor_shape - ) - - m_numtopk, k_input_doubled = input_tensor.shape - k = k_input_doubled // 2 - - max_tokens_per_expert = int(os.environ.get("MODELOPT_MAX_TOKENS_PER_EXPERT", 65536)) - assert m_numtopk <= max_tokens_per_expert * topk, ( - f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT({max_tokens_per_expert})" - f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use" - " MODELOPT_MAX_TOKENS_PER_EXPERT to set this value." - ) - scales_k = k // 16 - padded_k_in_int32 = (scales_k + 3) // 4 - - output = torch.empty( - m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8 - ) - if padded_k_in_int32 * 4 > scales_k: - output_scales = torch.zeros( - max_tokens_per_expert * topk, - padded_k_in_int32, - dtype=torch.int32, - device=input_tensor.device, - ) - else: - output_scales = torch.empty( - max_tokens_per_expert * topk, - padded_k_in_int32, - dtype=torch.int32, - device=input_tensor.device, - ) - - _silu_and_mul_scaled_fp4_experts_quant_packed_custom_op( - output, - output_scales, - input_tensor, - input_global_scale, - expert_offsets, - blockscale_offsets, - ) - output_scales = output_scales.view(torch.float8_e4m3fn) - return output, output_scales - - -@register_custom_op( - op_name="scaled_fp4_grouped_quant", - mutates_args=["output", "output_scales"], -) -def _scaled_fp4_grouped_quant_custom_op( - input_tensor: torch.Tensor, - output: torch.Tensor, - output_scales: torch.Tensor, - input_global_scale: torch.Tensor, - mask: torch.Tensor, -) -> None: - l, m, k = input_tensor.shape - del l, m - module = _jit_nvfp4_expert_quant_module() - module.silu_and_mul_scaled_fp4_experts_quant( - output.view(-1, k // 2), - output_scales.view(-1, output_scales.shape[-1]), - input_tensor.view(-1, k), - input_global_scale, - mask, - False, - ) - - -@debug_kernel_api -def scaled_fp4_grouped_quant( - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - mask: torch.Tensor, -): - """Quantize grouped GEMM inputs to FP4 and return logical (m, k//2, l).""" - device = input_tensor.device - l, m, k = input_tensor.shape - sf_vec_size = 16 - assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}." - - scale_k = k // sf_vec_size - padded_k = (scale_k + (4 - 1)) // 4 * 4 - padded_k_int32 = padded_k // 4 - padded_m = (m + (128 - 1)) // 128 * 128 - output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8) - output_scales = torch.empty( - l, padded_m, padded_k_int32, device=device, dtype=torch.int32 - ) - - _scaled_fp4_grouped_quant_custom_op( - input_tensor, - output, - output_scales, - input_global_scale, - mask, - ) - - output = output.permute(1, 2, 0) - output_scales = output_scales.view(torch.float8_e4m3fn).view( - l, padded_m // 128, padded_k // 4, 32, 4, 4 - ) - output_scales = output_scales.permute(3, 4, 1, 5, 2, 0) - return output, output_scales - - -@register_custom_op( - op_name="silu_and_mul_scaled_fp4_grouped_quant", - mutates_args=["output", "output_scales"], -) -def _silu_and_mul_scaled_fp4_grouped_quant_custom_op( - input_tensor: torch.Tensor, - output: torch.Tensor, - output_scales: torch.Tensor, - input_global_scale: torch.Tensor, - mask: torch.Tensor, -) -> None: - l, m, k_by_2 = input_tensor.shape - del l, m - module = _jit_nvfp4_expert_quant_module() - module.silu_and_mul_scaled_fp4_experts_quant( - output.view(-1, output.shape[-1]), - output_scales.view(-1, output_scales.shape[-1]), - input_tensor.view(-1, k_by_2), - input_global_scale, - mask, - True, - ) - - -@debug_kernel_api -def silu_and_mul_scaled_fp4_grouped_quant( - input_tensor: torch.Tensor, - input_global_scale: torch.Tensor, - mask: torch.Tensor, -): - """Apply SiLU-and-mul then quantize grouped GEMM inputs to FP4.""" - device = input_tensor.device - l, m, k_by_2 = input_tensor.shape - k = k_by_2 // 2 - sf_vec_size = 16 - assert k % sf_vec_size == 0, f"k must be multiple of 16, but got {k}." - - scale_k = k // sf_vec_size - padded_k = (scale_k + (4 - 1)) // 4 * 4 - padded_k_int32 = padded_k // 4 - padded_m = (m + (128 - 1)) // 128 * 128 - output = torch.empty(l, m, k // 2, device=device, dtype=torch.uint8) - output_scales = torch.empty( - l, padded_m, padded_k_int32, device=device, dtype=torch.int32 - ) - - _silu_and_mul_scaled_fp4_grouped_quant_custom_op( - input_tensor, - output, - output_scales, - input_global_scale, - mask, - ) - - output = output.permute(1, 2, 0) - output_scales = output_scales.view(torch.float8_e4m3fn).view( - l, padded_m // 128, padded_k // 4, 32, 4, 4 - ) - output_scales = output_scales.permute(3, 4, 1, 5, 2, 0) - return output, output_scales - - -@register_custom_op( - op_name="cutlass_fp4_group_mm", - mutates_args=[ - "output", - "a_ptrs", - "b_ptrs", - "out_ptrs", - "a_scales_ptrs", - "b_scales_ptrs", - "alpha_ptrs", - "layout_sfa", - "layout_sfb", - ], -) -def _cutlass_fp4_group_mm_custom_op( - output: torch.Tensor, - a_fp4: torch.Tensor, - b_fp4: torch.Tensor, - a_blockscale: torch.Tensor, - b_blockscale: torch.Tensor, - alphas: torch.Tensor, - ab_strides: torch.Tensor, - c_strides: torch.Tensor, - problem_sizes: torch.Tensor, - expert_offsets: torch.Tensor, - blockscale_offsets: torch.Tensor, - a_ptrs: torch.Tensor, - b_ptrs: torch.Tensor, - out_ptrs: torch.Tensor, - a_scales_ptrs: torch.Tensor, - b_scales_ptrs: torch.Tensor, - alpha_ptrs: torch.Tensor, - layout_sfa: torch.Tensor, - layout_sfb: torch.Tensor, -) -> None: - module = _jit_nvfp4_blockwise_moe_module() - module.cutlass_fp4_group_mm( - output, - a_fp4, - b_fp4, - a_blockscale, - b_blockscale, - alphas, - ab_strides, - c_strides, - problem_sizes, - expert_offsets, - blockscale_offsets, - a_ptrs, - b_ptrs, - out_ptrs, - a_scales_ptrs, - b_scales_ptrs, - alpha_ptrs, - layout_sfa, - layout_sfb, - ) - - -def suggest_nvfp4_global_scale(x: torch.Tensor) -> torch.Tensor: - """Utility for tests/benchmarks: return global scale used by NVFP4 quantization.""" - tensor_amax = torch.abs(x).max().to(torch.float32) - return _FLOAT8_E4M3_MAX * _FLOAT4_E2M1_MAX / tensor_amax diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index 9f91f06a8..4e2b118a9 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -637,7 +637,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase): fp4_quantize = _get_fp4_quantize_op() if fp4_quantize is None: raise RuntimeError( - "No FP4 quantization kernel available. Install flashinfer or sgl_kernel." + "No FP4 quantization kernel available. Install flashinfer." ) x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv) @@ -652,29 +652,17 @@ class ModelOptFp4LinearMethod(LinearMethodBase): if w_scale_interleaved.dtype == torch.uint8: w_scale_interleaved = w_scale_interleaved.view(torch.float8_e4m3fn) fp4_gemm, flashinfer_backend = _get_fp4_gemm_op() - if flashinfer_backend is not None: - out = fp4_gemm( - x_fp4, - w.T, - x_scale_interleaved, - w_scale_interleaved.T, - layer.alpha, - output_dtype, - backend=flashinfer_backend, - ) - elif fp4_gemm is not None: - out = fp4_gemm( - x_fp4, - w, - x_scale_interleaved, - w_scale_interleaved, - layer.alpha, - output_dtype, - ) - else: - raise RuntimeError( - "No FP4 GEMM kernel available. Install flashinfer or sgl_kernel." - ) + if fp4_gemm is None: + raise RuntimeError("No FP4 GEMM kernel available. Install flashinfer.") + out = fp4_gemm( + x_fp4, + w.T, + x_scale_interleaved, + w_scale_interleaved.T, + layer.alpha, + output_dtype, + backend=flashinfer_backend, + ) out = slice_nvfp4_output(out, output_size) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 0e11d466b..1df3b172e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -19,7 +19,6 @@ from typing import Any import torch import torch.nn as nn -from sglang.jit_kernel.nvfp4 import prewarm_nvfp4_jit_modules from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode from sglang.multimodal_gen.configs.pipeline_configs.flux import ( @@ -403,13 +402,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): compile_kwargs = build_torch_compile_kwargs(mode=mode) logger.info(f"Compiling transformer with mode: {mode}") - if self._needs_nvfp4_jit_prewarm(module): - logger.info( - "Prewarming NVFP4 JIT modules before torch.compile to avoid " - "Dynamo tracing JIT initialization." - ) - prewarm_nvfp4_jit_modules() - # TODO(triple-mu): support customized fullgraph and dynamic in the future self._torch_compile_registry.compile_once( module, @@ -424,16 +416,6 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): for transformer in filter(None, [self.transformer, self.transformer_2]): self._maybe_torch_compile(transformer) - @staticmethod - def _needs_nvfp4_jit_prewarm(module: nn.Module) -> bool: - for submodule in module.modules(): - quant_method = getattr(submodule, "quant_method", None) - if quant_method is None: - continue - if type(quant_method).__name__ == "ModelOptFp4LinearMethod": - return True - return False - def _cache_dit_dual_model_name(self) -> str: return "wan2.2" diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 4139d20c8..0cb5d4a7a 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -325,13 +325,6 @@ class CudaPlatformBase(Platform): try: from flashinfer import fp4_quantize - return fp4_quantize - except ImportError: - pass - - try: - from sgl_kernel import scaled_fp4_quant as fp4_quantize - return fp4_quantize except ImportError: return None @@ -375,16 +368,9 @@ class CudaPlatformBase(Platform): except ImportError: logger.warning( "Requested SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND=%r " - "but flashinfer.mm_fp4 is unavailable. Falling back to " - "cutlass.", + "but flashinfer.mm_fp4 is unavailable.", requested_backend or "flashinfer_trtllm (default)", ) - - try: - from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm - - return cutlass_fp4_gemm, None - except ImportError: return None, None @classmethod diff --git a/python/sglang/srt/layers/moe/cutlass_moe.py b/python/sglang/srt/layers/moe/cutlass_moe.py index b48653b11..ad16e3c7e 100755 --- a/python/sglang/srt/layers/moe/cutlass_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_moe.py @@ -4,7 +4,6 @@ from typing import Optional, Tuple import torch -from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams from sglang.srt.utils import is_cuda, is_sm90_supported, is_sm100_supported _is_cuda = is_cuda() @@ -20,11 +19,6 @@ if _is_cuda: ) from sglang.jit_kernel.activation import silu_and_mul - from sglang.jit_kernel.nvfp4 import ( - cutlass_fp4_group_mm, - scaled_fp4_experts_quant, - silu_and_mul_scaled_fp4_experts_quant_packed, - ) def cutlass_fused_experts_fp8( @@ -340,160 +334,3 @@ def cutlass_fused_experts_fp8( apply_shuffle_mul_sum(c2, output, c_map, topk_weights.to(out_dtype)) return output - - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = 448.0 - - -def cutlass_moe_fp4( - a: torch.Tensor, - a1_gscale: torch.Tensor, - w1_fp4: torch.Tensor, - w1_blockscale: torch.Tensor, - w1_alphas: torch.Tensor, - a2_gscale: torch.Tensor, - w2_fp4: torch.Tensor, - w2_blockscale: torch.Tensor, - w2_alphas: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - params: CutlassMoEParams, - apply_router_weight_on_input: bool = False, - no_combine: bool = False, -): - """ - MoE implementation for FP4 Inputs - - # Gemm 1 - a: Input tensor: [m, k] (half/bfloat16) - a1_gscale: Activation scale per expert: [e] (float32) - w1(gate up) (not an argument to cutlass_moe_fp4): [e, 2 * n, k] - w1_fp4: [e, 2 * n, k // 2], dtype: torch.uint8 (stacked fp4: E2M1) - (Note: `n` is the up projection output dim, `k` is the input dim in - full precision) - w1_blockscale: [e, 2 * n, k // block_size] (float8_e4m3) - (Block size = 16 for NVFP4) - - # Gemm 2 - a2_gscale: Activation scale per expert: [e] - w2(down projection) (not an argument to cutlass_moe_fp4): [e, k, n] - w2_fp4: [e, k, n // 2], dtype: torch.uint8 (stacked E2M1) - w2_blockscale: [e, k, n // block_size], dtype: float8_e4m3 - - Strides for activations, weights and output in logical number of elements. - The activations & output stride is the number of elements to the next row. - The weights stride is the number of elements to the next row per expert. - For example, if the weight is [e, n, k], then the b_stride is a tensor of - shape [e] with each element being k. Similarly for activations, if the - shape is [m, k], then the a_stride has shape [e] with each value k. - Similarly for output, if the output is [m, n], then the c_stride is a - tensor of shape [e] with each element being k. - - Note: cutlass_fp4_group_mm is designed to accept the strides of - activations and weights to be the same, so it is passed in as a single - tensor. - ab_strides_13: [e] dtype: int64 [Gemm 1: Activation / Weight strides] - ab_strides_2: [e] dtype: int64 [Gemm 2: Activation / Weight strides] - c_strides_13: [e] dtype: int64 [Gemm 1: Output Strides] - c_strides_2: [e] dtype: int64 [Gemm 1: Output Strides] - - topk_weights: [m, topk] dtype: float8 - topk_ids: [m, topk] dtype: float8 - - m, n, k: Unquantized weight shapes, dtype: int - e: number of experts for the current rank, dtype: int - assumes that topk < k < n to satisfy - up/down projection expectations. - """ - assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" - assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8" - assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8" - assert ( - w1_fp4.ndim == 3 - and w2_fp4.ndim == 3 - and w1_blockscale.ndim == 3 - and w2_blockscale.ndim == 3 - ), "All Weights must be of rank 3 for cutlass_moe_fp4" - m_a, k_a = a.shape - e_w1, nx2_w1, half_k_w1 = w1_fp4.shape - e_w2, k_w2, half_n_w2 = w2_fp4.shape - - assert e_w1 == e_w2 and e_w1 == params.num_experts, ( - "Number of experts must match", - " between weights.", - ) - assert ( - k_a // 2 == half_k_w1 and params.hidden_size == k_w2 - ), "Hidden size mismatch between a, w1 and w2" - assert ( - nx2_w1 == params.intermediate_size_per_partition * 2 - and half_n_w2 == params.intermediate_size_per_partition // 2 - ), ("mismatch in " "expected `n`") - assert 2 * half_k_w1 == k_w2, "Hidden size mismatch w2 and w1" - assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype" - - out_dtype = a.dtype - num_topk = topk_ids.shape[1] - device = a.device - a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) - c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device) - prepare_moe_input( - topk_ids, - params.expert_offsets, - params.problem_sizes1, - params.problem_sizes2, - a_map, - c_map, - params.num_experts, - params.intermediate_size_per_partition, - params.hidden_size, - params.blockscale_offsets, - ) - - rep_a_fp4, rep_a_blockscale = scaled_fp4_experts_quant( - a, - a1_gscale, - params.expert_offsets, - params.blockscale_offsets, - num_topk, - expert_map=a_map, - ) - c1 = cutlass_fp4_group_mm( - rep_a_fp4, - w1_fp4, - rep_a_blockscale, - w1_blockscale, - w1_alphas, - out_dtype, - params.to_gemm1_args(), - ) - del rep_a_fp4, rep_a_blockscale - - # fused: SiLU + mul then FP4 quant (expert-packed) - int_fp4, int_blockscale = silu_and_mul_scaled_fp4_experts_quant_packed( - c1, - a2_gscale, - params.expert_offsets, - params.blockscale_offsets, - num_topk, - ) - - c2 = cutlass_fp4_group_mm( - int_fp4, - w2_fp4, - int_blockscale, - w2_blockscale, - w2_alphas, - out_dtype, - params.to_gemm2_args(), - ) - del int_fp4, int_blockscale - - if no_combine: - c2 = shuffle_rows(c2, c_map, (m_a * num_topk, params.hidden_size)) - c2 = c2.view(m_a, num_topk, params.hidden_size) - return c2.to(out_dtype) - output = torch.empty((m_a, k_a), device=device, dtype=out_dtype) - weights = topk_weights.to(out_dtype) if not apply_router_weight_on_input else None - apply_shuffle_mul_sum(c2, output, c_map, weights) - return output diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index c814fe270..e84b61916 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -36,7 +36,6 @@ from sglang.srt.layers.moe.moe_runner.base import ( ) from sglang.srt.layers.utils import copy_or_rebind_param from sglang.srt.utils.common import ( - is_cuda_alike, is_flashinfer_available, next_power_of_2, ) @@ -103,8 +102,6 @@ if TYPE_CHECKING: if is_flashinfer_available(): from sglang.srt.layers.quantization.fp4_utils import fp4_quantize -elif is_cuda_alike(): - from sglang.jit_kernel.nvfp4 import scaled_fp4_quant as fp4_quantize else: fp4_quantize = None diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py index 0dd4dcea2..bb5e7722e 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py @@ -151,10 +151,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsLinearScheme): w = layer.weight_packed w_blockscale = layer.weight_scale - if ( - enable_flashinfer_fp4_gemm - and not get_fp4_gemm_runner_backend().is_cutlass() - ): + if enable_flashinfer_fp4_gemm: w = layer.weight_packed.T w_blockscale = layer.weight_scale.T diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py index 0a9921874..84bd4ae8f 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py @@ -11,7 +11,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( ) from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig -from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType from sglang.srt.layers.moe.utils import RoutingMethodType, get_moe_runner_backend from sglang.srt.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsMoEScheme, @@ -278,19 +277,18 @@ class CompressedTensorsW4A4Nvfp4MoE(CompressedTensorsMoEScheme): swizzle_blockscale(layer.w2_weight_scale), requires_grad=False ) - layer.cutlass_moe_params = CutlassMoEParams( - CutlassMoEType.BlockscaledFP4, - layer.w13_weight.device, - num_experts=layer.num_experts, - intermediate_size_per_partition=layer.w2_weight.shape[2] * 2, - hidden_size=layer.w13_weight.shape[2] * 2, - ) - def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config - self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + if self.use_flashinfer_trtllm: + self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) + else: + import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401 – triggers @register_fused_func + + self.runner = MoeRunner( + MoeRunnerBackend.FLASHINFER_CUTLASS, moe_runner_config + ) def apply_weights( self, @@ -385,24 +383,33 @@ class CompressedTensorsW4A4Nvfp4MoE(CompressedTensorsMoEScheme): output=symm_output, )[0] else: - from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4 + from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import ( + FlashInferCutlassMoeQuantInfo, + ) - topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids + assert ( + not self.moe_runner_config.apply_router_weight_on_input + ), "apply_router_weight_on_input is not supported for Flashinfer" - output = cutlass_moe_fp4( - a=x, - a1_gscale=layer.w13_input_scale_quant, - w1_fp4=layer.w13_weight, - w1_blockscale=layer.w13_weight_scale, - w1_alphas=layer.g1_alphas, - a2_gscale=layer.w2_input_scale_quant, - w2_fp4=layer.w2_weight, - w2_blockscale=layer.w2_weight_scale, - w2_alphas=layer.g2_alphas, - topk_weights=topk_weights, - topk_ids=topk_ids, - params=layer.cutlass_moe_params, - apply_router_weight_on_input=self.moe_runner_config.apply_router_weight_on_input, - ).to(x.dtype) + quant_info = FlashInferCutlassMoeQuantInfo( + quant_type="fp4", + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + output_dtype=x.dtype, + quant_scales=[ + layer.w13_input_scale_quant, + layer.w13_weight_scale, + layer.g1_alphas, + layer.w2_input_scale_quant, + layer.w2_weight_scale, + layer.g2_alphas, + ], + moe_ep_size=layer.moe_ep_size, + moe_ep_rank=layer.moe_ep_rank, + moe_tp_size=layer.moe_tp_size, + moe_tp_rank=layer.moe_tp_rank, + apply_routed_scaling_factor=False, + ) + return self.runner.run(dispatch_output, quant_info) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/fp4_utils.py b/python/sglang/srt/layers/quantization/fp4_utils.py index 39ffbb935..1b5858d71 100644 --- a/python/sglang/srt/layers/quantization/fp4_utils.py +++ b/python/sglang/srt/layers/quantization/fp4_utils.py @@ -94,7 +94,6 @@ class Fp4GemmRunnerBackend(Enum): """Enum for FP4 GEMM runner backend selection.""" AUTO = "auto" - CUTLASS = "cutlass" FLASHINFER_CUDNN = "flashinfer_cudnn" FLASHINFER_CUTEDSL = "flashinfer_cutedsl" FLASHINFER_CUTLASS = "flashinfer_cutlass" @@ -104,9 +103,6 @@ class Fp4GemmRunnerBackend(Enum): def is_auto(self) -> bool: return self == Fp4GemmRunnerBackend.AUTO - def is_cutlass(self) -> bool: - return self == Fp4GemmRunnerBackend.CUTLASS - def is_flashinfer_cudnn(self) -> bool: return self == Fp4GemmRunnerBackend.FLASHINFER_CUDNN diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 3e2a0bedc..26f3385a6 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -18,7 +18,6 @@ from sglang.srt.layers.moe import ( MoeRunnerConfig, get_moe_runner_backend, ) -from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.moe.utils import ( is_flashinfer_cutedsl_v1_path, @@ -108,14 +107,6 @@ except ImportError: shuffle_matrix_a = None shuffle_matrix_sf_a = None -if is_cuda(): - try: - from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm as cutlass_fp4_gemm - except ImportError: - cutlass_fp4_gemm = None -else: - cutlass_fp4_gemm = None - # Initialize logger for the module logger = logging.getLogger(__name__) @@ -144,23 +135,16 @@ def fp4_gemm( out_dtype: torch.dtype, out_features: int, ) -> torch.Tensor: - fp4_backend = get_fp4_gemm_runner_backend() - if fp4_backend.is_cutlass() and cutlass_fp4_gemm is not None: - # flashinfer.fp4_quantize returns scale factors as uint8 (e4m3fn bits - # stored in uint8 memory). The JIT kernel requires float8_e4m3fn dtype. - if input_sf.dtype != torch.float8_e4m3fn: - input_sf = input_sf.view(torch.float8_e4m3fn) - if weight_sf.dtype != torch.float8_e4m3fn: - weight_sf = weight_sf.view(torch.float8_e4m3fn) - return cutlass_fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype) - elif enable_flashinfer_fp4_gemm: - # Use the remapping logic to convert SGLang backend names to FlashInfer API names - backend = fp4_backend.get_flashinfer_backend() - return flashinfer_fp4_gemm( - input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend + if not enable_flashinfer_fp4_gemm: + raise RuntimeError( + "NVFP4 GEMM requires flashinfer's mm_fp4; please install flashinfer." ) - else: - return cutlass_fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype) + fp4_backend = get_fp4_gemm_runner_backend() + # Use the remapping logic to convert SGLang backend names to FlashInfer API names + backend = fp4_backend.get_flashinfer_backend() + return flashinfer_fp4_gemm( + input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend + ) if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None): @@ -1709,10 +1693,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase): w = layer.weight w_scale_interleaved = layer.weight_scale_interleaved - if ( - enable_flashinfer_fp4_gemm - and not get_fp4_gemm_runner_backend().is_cutlass() - ): + if enable_flashinfer_fp4_gemm: w = layer.weight.T w_scale_interleaved = layer.weight_scale_interleaved.T @@ -2406,29 +2387,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): requires_grad=False, ) - # Both flashinfer cutlass and regular cutlass use same processing for w2 - - # Set up CUTLASS MoE parameters (reuse to keep CUDA graph stable) - device = layer.w13_weight.device - inter_size = layer.w2_weight.shape[2] * 2 - hidden_size = layer.w13_weight.shape[2] * 2 - existing_params = getattr(layer, "cutlass_moe_params", None) - if ( - existing_params is None - or existing_params.cutlass_moe_type != CutlassMoEType.BlockscaledFP4 - or existing_params.num_experts != layer.num_experts - or existing_params.intermediate_size_per_partition != inter_size - or existing_params.hidden_size != hidden_size - or existing_params.device != device - ): - layer.cutlass_moe_params = CutlassMoEParams( - CutlassMoEType.BlockscaledFP4, - device, - num_experts=layer.num_experts, # global num experts - intermediate_size_per_partition=inter_size, # n - hidden_size=hidden_size, - ) # k - @property def load_up_proj_weight_first(self) -> bool: # Load W13 as [Up, Gate] for FlashInfer CUTLASS and CuteDSL v2 kernels. @@ -2459,18 +2417,19 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): if moe_runner_backend.is_flashinfer_cutlass(): import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401 - # The plain CUTLASS backend uses the direct cutlass_moe_fp4 fused path - # (see apply()), not a registered MoeRunner fused func, so skip creating - # a MoeRunner for it -- constructing one would fail the fused-func check. - if not moe_runner_backend.is_cutlass(): - self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + if moe_runner_backend.is_cutlass(): + raise NotImplementedError( + "moe_runner_backend=cutlass is not supported for NVFP4 MoE. " + "Use --moe-runner-backend flashinfer_cutlass instead." + ) + + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) def apply( self, layer: FusedMoE, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput # Note: dispatch_output may be a DeepEPLLDispatchOutput (no topk_output # attribute -- topk_ids/topk_weights live directly on the dispatch @@ -2621,26 +2580,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) return self.runner.run(dispatch_output, quant_info) - from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4 - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids - output = cutlass_moe_fp4( - a=x, - a1_gscale=layer.w13_input_scale_quant, - w1_fp4=layer.w13_weight, - w1_blockscale=layer.w13_blockscale_swizzled, - w1_alphas=layer.g1_alphas, - a2_gscale=layer.w2_input_scale_quant, - w2_fp4=layer.w2_weight, - w2_blockscale=layer.w2_blockscale_swizzled, - w2_alphas=layer.g2_alphas, - topk_weights=topk_weights, - topk_ids=topk_ids, - params=layer.cutlass_moe_params, - apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input, - no_combine=moe_runner_config.no_combine, - ).to(x.dtype) - # Scale by routed_scaling_factor is fused into select_experts. - return StandardCombineInput(hidden_states=output) + raise NotImplementedError( + f"Unsupported moe_runner_backend for NVFP4 MoE: {moe_runner_backend}. " + "Use --moe-runner-backend flashinfer_cutlass instead." + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index dbfb3884a..a9b6340c7 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -296,7 +296,6 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [ FP4_GEMM_RUNNER_BACKEND_CHOICES = [ "auto", - "cutlass", "flashinfer_cudnn", "flashinfer_cutedsl", "flashinfer_cutlass", @@ -1436,7 +1435,7 @@ class ServerArgs: fp4_gemm_runner_backend: A[ str, Arg( - help="Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default; selects flashinfer_cutedsl on SM100, marlin on SM80-SM90, flashinfer_cutlass otherwise (including SM120)), 'cutlass' (SGLang CUTLASS kernel), 'flashinfer_cutlass' (FlashInfer CUTLASS backend), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_cutedsl' (FlashInfer CuTe DSL backend), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling), 'marlin' (weight-only W4A16 fallback for SM80+). ", + help="Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default; selects flashinfer_cutedsl on SM100, marlin on SM80-SM90, flashinfer_cutlass otherwise (including SM120)), 'flashinfer_cutlass' (FlashInfer CUTLASS backend), 'flashinfer_cudnn' (FlashInfer cuDNN backend, optimal on CUDA 13+ with cuDNN 9.15+), 'flashinfer_cutedsl' (FlashInfer CuTe DSL backend), 'flashinfer_trtllm' (FlashInfer TensorRT-LLM backend, requires different weight preparation with shuffling), 'marlin' (weight-only W4A16 fallback for SM80+). ", cli_name="--fp4-gemm-backend", choices=FP4_GEMM_RUNNER_BACKEND_CHOICES, ), diff --git a/sgl-kernel/benchmark/bench_fp4_gemm.py b/sgl-kernel/benchmark/bench_fp4_gemm.py index c742c8e45..df536f508 100755 --- a/sgl-kernel/benchmark/bench_fp4_gemm.py +++ b/sgl-kernel/benchmark/bench_fp4_gemm.py @@ -6,14 +6,13 @@ from typing import List, Tuple import torch import triton -from flashinfer import mm_fp4 +from flashinfer import fp4_quantize, mm_fp4 from flashinfer.autotuner import autotune from flashinfer.jit.core import logger as flashinfer_logger from flashinfer.testing import bench_gpu_time flashinfer_logger.setLevel(logging.ERROR) -from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant from sglang.srt.utils import ( get_device_capability, is_sm100_supported, @@ -154,13 +153,12 @@ def _run_mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend): x_log=False, line_arg="provider", line_vals=( - ["sglang_cutlass", "cutlass", "cudnn", "trtllm", "cute-dsl", "auto"] + ["cutlass", "cudnn", "trtllm", "cute-dsl", "auto"] if is_sm100_supported() - else ["sglang_cutlass", "cutlass", "cudnn", "cute-dsl", "auto"] + else ["cutlass", "cudnn", "cute-dsl", "auto"] ), line_names=( [ - "sglang cutlass fp4", "flashinfer cutlass fp4", "cudnn fp4", "trtllm fp4", @@ -169,7 +167,6 @@ def _run_mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend): ] if is_sm100_supported() else [ - "sglang cutlass fp4", "flashinfer cutlass fp4", "cudnn fp4", "cute-dsl fp4", @@ -178,7 +175,6 @@ def _run_mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend): ), styles=( [ - ("red", "solid"), ("orange", "solid"), ("blue", "solid"), ("green", "solid"), @@ -187,7 +183,6 @@ def _run_mm_fp4(a_fp4, b_fp4_T, a_sf, b_sf_T, alpha, dtype, res_fi, backend): ] if is_sm100_supported() else [ - ("red", "solid"), ("orange", "solid"), ("blue", "solid"), ("brown", "solid"), @@ -212,26 +207,19 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): (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) + a_fp4, a_scale_interleaved = fp4_quantize(a_dtype, a_global_scale) + b_fp4, b_scale_interleaved = fp4_quantize(b_dtype, b_global_scale) + # flashinfer.fp4_quantize returns scale factors as uint8 (e4m3fn bits stored + # in uint8 memory); the JIT cutlass kernel requires float8_e4m3fn dtype. + if a_scale_interleaved.dtype != torch.float8_e4m3fn: + a_scale_interleaved = a_scale_interleaved.view(torch.float8_e4m3fn) + if b_scale_interleaved.dtype != torch.float8_e4m3fn: + b_scale_interleaved = b_scale_interleaved.view(torch.float8_e4m3fn) b_fp4_T = b_fp4.T b_sf_T = b_scale_interleaved.T res_fi = torch.empty((M, N), dtype=dtype, device="cuda") - if provider == "sglang_cutlass": - times_ms = bench_gpu_time( - fn=cutlass_scaled_fp4_mm, - input_args=( - a_fp4, - b_fp4, - a_scale_interleaved, - b_scale_interleaved, - alpha, - dtype, - ), - use_cuda_graph=True, - ) - elif provider == "cutlass": + if provider == "cutlass": with autotune(): _run_mm_fp4( a_fp4, @@ -359,8 +347,16 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file): bandwidth_gbs = total_bytes / (ms * 1e-3) / 1e9 if correctness: - res_cutlass = cutlass_scaled_fp4_mm( - a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype + res_cutlass = torch.empty((M, N), dtype=dtype, device="cuda") + mm_fp4( + a_fp4, + b_fp4_T, + a_scale_interleaved, + b_sf_T, + alpha, + dtype, + res_cutlass, + backend="cutlass", ) mm_fp4( a_fp4, diff --git a/test/registered/jit/benchmark/bench_nvfp4_blockwise_moe.py b/test/registered/jit/benchmark/bench_nvfp4_blockwise_moe.py deleted file mode 100644 index d0381f35f..000000000 --- a/test/registered/jit/benchmark/bench_nvfp4_blockwise_moe.py +++ /dev/null @@ -1,263 +0,0 @@ -from __future__ import annotations - -import sys -from typing import Any - -import torch -import triton - -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark -from sglang.jit_kernel.nvfp4 import ( - cutlass_fp4_group_mm, - scaled_fp4_experts_quant, - scaled_fp4_quant, -) -from sglang.srt.utils import is_sm100_supported -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" -) - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max -_NVFP4_SUPPORTED = is_sm100_supported() - - -def _round_up(x: int, y: int) -> int: - return ((x + y - 1) // y) * y - - -def _expert_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor: - offsets = [0] - for m in m_per_expert: - offsets.append(offsets[-1] + m) - return torch.tensor(offsets, dtype=torch.int32, device=device) - - -def _blockscale_offsets(m_per_expert: list[int], device: torch.device) -> torch.Tensor: - offsets = [0] - for m in m_per_expert: - offsets.append(offsets[-1] + _round_up(m, 128)) - return torch.tensor(offsets, dtype=torch.int32, device=device) - - -def _prepare_case( - total_tokens: int, n: int, k: int, num_experts: int, dtype: torch.dtype -) -> dict[str, Any]: - device = torch.device("cuda") - base = total_tokens // num_experts - rem = total_tokens % num_experts - m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)] - - expert_offsets_full = _expert_offsets(m_per_expert, device) - blockscale_offsets_full = _blockscale_offsets(m_per_expert, device) - - a = torch.randn((total_tokens, k), device=device, dtype=dtype) * 0.1 - b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1 - - a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) - for i in range(num_experts): - start = int(expert_offsets_full[i].item()) - end = int(expert_offsets_full[i + 1].item()) - a_global_scale[i] = ( - FLOAT8_E4M3_MAX - * FLOAT4_E2M1_MAX - / a[start:end].abs().max().to(torch.float32) - ) - - b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) - for i in range(num_experts): - b_global_scale[i] = ( - FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b[i].abs().max().to(torch.float32) - ) - - a_fp4, a_blockscale = scaled_fp4_experts_quant( - a, - a_global_scale, - expert_offsets_full, - blockscale_offsets_full, - topk=1, - ) - - b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8) - b_blockscale = torch.empty( - (num_experts, _round_up(n, 128), _round_up(k // 16, 4)), - device=device, - dtype=torch.float8_e4m3fn, - ) - for i in range(num_experts): - b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i]) - b_fp4[i].copy_(b_fp4_i) - b_blockscale[i].copy_(b_scale_i) - - alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32) - params = { - "ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device), - "c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device), - "problem_sizes": torch.tensor( - [[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device - ), - "expert_offsets": expert_offsets_full[:-1].contiguous(), - "blockscale_offsets": blockscale_offsets_full[:-1].contiguous(), - "a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device), - "layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device), - } - - expert_ranges: list[tuple[int, int]] = [] - start = 0 - for m in m_per_expert: - end = start + m - expert_ranges.append((start, end)) - start = end - - return { - "a": a, - "b": b, - "a_fp4": a_fp4, - "b_fp4": b_fp4, - "a_blockscale": a_blockscale, - "b_blockscale": b_blockscale, - "alphas": alphas, - "params": params, - "expert_offsets_full": expert_offsets_full, - "expert_ranges": expert_ranges, - "dtype": dtype, - } - - -def _torch_ref_group_mm(case: dict[str, Any]) -> torch.Tensor: - a = case["a"] - b = case["b"] - dtype = case["dtype"] - expert_ranges = case["expert_ranges"] - total_tokens = a.shape[0] - n = b.shape[1] - out = torch.empty((total_tokens, n), device=a.device, dtype=dtype) - for i, (start, end) in enumerate(expert_ranges): - out[start:end] = torch.matmul(a[start:end], b[i].t()) - return out - - -def _aot_cutlass_fp4_group_mm(case: dict[str, Any]) -> torch.Tensor: - a_fp4 = case["a_fp4"] - b_fp4 = case["b_fp4"] - a_blockscale = case["a_blockscale"] - b_blockscale = case["b_blockscale"] - alphas = case["alphas"] - params = case["params"] - out_dtype = case["dtype"] - - out = torch.empty( - (a_fp4.shape[0], b_fp4.shape[1]), device=a_fp4.device, dtype=out_dtype - ) - torch.ops.sgl_kernel.cutlass_fp4_group_mm.default( - out, - a_fp4, - b_fp4, - a_blockscale, - b_blockscale, - alphas, - params["ab_strides"], - params["c_strides"], - params["problem_sizes"], - params["expert_offsets"], - params["blockscale_offsets"], - ) - return out - - -def _probe_legacy_aot_group_mm() -> tuple[bool, str]: - if not torch.cuda.is_available(): - return False, "CUDA is not available." - if not _NVFP4_SUPPORTED: - return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." - try: - import sgl_kernel # noqa: F401 - except Exception as e: - return False, f"import sgl_kernel failed: {e}" - if not hasattr(torch.ops, "sgl_kernel"): - return False, "torch.ops.sgl_kernel is not registered." - op = getattr(torch.ops.sgl_kernel, "cutlass_fp4_group_mm", None) - if op is None or not hasattr(op, "default"): - return False, "torch.ops.sgl_kernel.cutlass_fp4_group_mm.default is missing." - try: - case = _prepare_case(64, 256, 128, 4, torch.bfloat16) - _aot_cutlass_fp4_group_mm(case) - torch.cuda.synchronize() - except Exception as e: - return False, f"calling AOT grouped_mm op failed: {e}" - return True, "" - - -_AOT_GROUP_MM_AVAILABLE, _AOT_GROUP_MM_REASON = _probe_legacy_aot_group_mm() - -shape_range = get_benchmark_range( - full_range=[(128, 256, 128, 4), (256, 512, 128, 8), (512, 512, 256, 8)], - ci_range=[(128, 256, 128, 4)], -) - -line_vals = ["jit"] -line_names = ["JIT NVFP4 MoE GroupMM"] -styles = [("green", "-")] -if _AOT_GROUP_MM_AVAILABLE: - line_vals.append("aot_sgl_kernel") - line_names.append("AOT NVFP4 MoE GroupMM") - styles.append(("orange", "-")) -line_vals.append("torch_ref") -line_names.append("Torch Ref") -styles.append(("blue", "-")) - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["total_tokens", "n", "k", "num_experts"], - x_vals=shape_range, - x_log=False, - line_arg="provider", - line_vals=line_vals, - line_names=line_names, - styles=styles, - ylabel="us", - plot_name="nvfp4-blockwise-moe-groupmm-performance", - args={}, - ) -) -def benchmark(total_tokens, n, k, num_experts, provider): - case = _prepare_case(total_tokens, n, k, num_experts, torch.bfloat16) - - if provider == "jit": - fn = lambda: cutlass_fp4_group_mm( - case["a_fp4"], - case["b_fp4"], - case["a_blockscale"], - case["b_blockscale"], - case["alphas"], - case["dtype"], - case["params"], - ) - elif provider == "aot_sgl_kernel": - fn = lambda: _aot_cutlass_fp4_group_mm(case) - elif provider == "torch_ref": - fn = lambda: _torch_ref_group_mm(case) - else: - raise ValueError(f"Unknown provider: {provider}") - - return run_benchmark(fn) - - -if __name__ == "__main__": - if not _NVFP4_SUPPORTED: - print("[skip] NVFP4 blockwise MoE benchmark requires sm100+ with CUDA 12.8+.") - sys.exit(0) - if not _AOT_GROUP_MM_AVAILABLE: - print( - f"[info] legacy AOT grouped_mm baseline unavailable: {_AOT_GROUP_MM_REASON}" - ) - benchmark.run(print_data=True) diff --git a/test/registered/jit/benchmark/bench_nvfp4_quant.py b/test/registered/jit/benchmark/bench_nvfp4_quant.py deleted file mode 100644 index cd25e3d90..000000000 --- a/test/registered/jit/benchmark/bench_nvfp4_quant.py +++ /dev/null @@ -1,197 +0,0 @@ -from __future__ import annotations - -import sys - -import torch -import triton - -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark -from sglang.jit_kernel.nvfp4 import scaled_fp4_quant -from sglang.srt.utils import is_sm100_supported -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" -) - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max -BLOCK_SIZE = 16 -_NVFP4_SUPPORTED = is_sm100_supported() - -try: - from flashinfer import fp4_quantize as flashinfer_fp4_quantize -except Exception: - flashinfer_fp4_quantize = None - - -def _torch_ref_quant(input: torch.Tensor, input_global_scale: torch.Tensor): - m, n = input.shape - x = input.view(m, n // BLOCK_SIZE, BLOCK_SIZE) - vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) - scale = input_global_scale * (vec_max / FLOAT4_E2M1_MAX) - scale = scale.to(torch.float8_e4m3fn).to(torch.float32) - output_scale = torch.where(scale == 0, torch.zeros_like(scale), 1.0 / scale) - - scaled_x = x.to(torch.float32) * output_scale - clipped = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n) - - rounded = clipped.clone() - rounded[(rounded >= 0.0) & (rounded <= 0.25)] = 0.0 - rounded[(rounded > 0.25) & (rounded < 0.75)] = 0.5 - rounded[(rounded >= 0.75) & (rounded <= 1.25)] = 1.0 - rounded[(rounded > 1.25) & (rounded < 1.75)] = 1.5 - rounded[(rounded >= 1.75) & (rounded <= 2.5)] = 2.0 - rounded[(rounded > 2.5) & (rounded < 3.5)] = 3.0 - rounded[(rounded >= 3.5) & (rounded <= 5.0)] = 4.0 - rounded[rounded > 5.0] = 6.0 - - # This baseline intentionally keeps work on GPU but does not pack to uint8. - return rounded, scale - - -def _aot_scaled_fp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor): - m, n = input.shape - output = torch.empty((m, n // 2), device=input.device, dtype=torch.uint8) - rounded_m = ((m + 128 - 1) // 128) * 128 - scale_n = n // BLOCK_SIZE - rounded_n = ((scale_n + 4 - 1) // 4) * 4 - output_scale = torch.empty( - (rounded_m, rounded_n // 4), device=input.device, dtype=torch.int32 - ) - torch.ops.sgl_kernel.scaled_fp4_quant.default( - output, input, output_scale, input_global_scale - ) - return output, output_scale.view(torch.float8_e4m3fn) - - -def _probe_legacy_aot_quant() -> tuple[bool, str]: - if not torch.cuda.is_available(): - return False, "CUDA is not available." - if not _NVFP4_SUPPORTED: - return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." - try: - import sgl_kernel # noqa: F401 - except Exception as e: - return False, f"import sgl_kernel failed: {e}" - if not hasattr(torch.ops, "sgl_kernel"): - return False, "torch.ops.sgl_kernel is not registered." - op = getattr(torch.ops.sgl_kernel, "scaled_fp4_quant", None) - if op is None or not hasattr(op, "default"): - return False, "torch.ops.sgl_kernel.scaled_fp4_quant.default is missing." - try: - x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda") - global_scale = ( - FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32) - ) - _aot_scaled_fp4_quant(x, global_scale) - torch.cuda.synchronize() - except Exception as e: - return False, f"calling AOT quant op failed: {e}" - return True, "" - - -_AOT_QUANT_AVAILABLE, _AOT_QUANT_REASON = _probe_legacy_aot_quant() - - -def _probe_flashinfer_quant() -> tuple[bool, str]: - if flashinfer_fp4_quantize is None: - return False, "import flashinfer.fp4_quantize failed." - if not torch.cuda.is_available(): - return False, "CUDA is not available." - if not _NVFP4_SUPPORTED: - return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." - try: - x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda") - global_scale = ( - FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.abs(x).max().to(torch.float32) - ) - flashinfer_fp4_quantize( - x, - global_scale, - BLOCK_SIZE, # sf_vec_size - False, # use_ue8m0 - True, # is_sf_swizzled_layout - ) - torch.cuda.synchronize() - except Exception as e: - return False, f"calling flashinfer.fp4_quantize failed: {e}" - return True, "" - - -_FLASHINFER_QUANT_AVAILABLE, _FLASHINFER_QUANT_REASON = _probe_flashinfer_quant() - -shape_range = get_benchmark_range( - full_range=[(128, 2048), (512, 4096), (1024, 4096), (2048, 8192)], - ci_range=[(128, 2048)], -) - -line_vals = [] -line_names = [] -styles = [] -if _FLASHINFER_QUANT_AVAILABLE: - line_vals.append("flashinfer") - line_names.append("FlashInfer FP4 Quant") - styles.append(("purple", "-")) -line_vals.append("jit") -line_names.append("JIT NVFP4 Quant") -styles.append(("green", "-")) -if _AOT_QUANT_AVAILABLE: - line_vals.append("aot_sgl_kernel") - line_names.append("AOT NVFP4 Quant") - styles.append(("orange", "-")) -line_vals.append("torch_ref") -line_names.append("Torch Ref") -styles.append(("blue", "-")) - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["m", "n"], - x_vals=shape_range, - x_log=False, - line_arg="provider", - line_vals=line_vals, - line_names=line_names, - styles=styles, - ylabel="us", - plot_name="nvfp4-quant-performance", - args={}, - ) -) -def benchmark(m, n, provider): - x = torch.randn((m, n), dtype=torch.bfloat16, device="cuda") - tensor_amax = torch.abs(x).max().to(torch.float32) - global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax - - if provider == "jit": - fn = lambda: scaled_fp4_quant(x, global_scale) - elif provider == "flashinfer": - fn = lambda: flashinfer_fp4_quantize( - x, - global_scale, - BLOCK_SIZE, # sf_vec_size - False, # use_ue8m0 - True, # is_sf_swizzled_layout - ) - elif provider == "aot_sgl_kernel": - fn = lambda: _aot_scaled_fp4_quant(x, global_scale) - elif provider == "torch_ref": - fn = lambda: _torch_ref_quant(x, global_scale) - else: - raise ValueError(f"Unknown provider: {provider}") - - return run_benchmark(fn) - - -if __name__ == "__main__": - if not _NVFP4_SUPPORTED: - print("[skip] NVFP4 quant benchmark requires sm100+ with CUDA 12.8+.") - sys.exit(0) - if not _FLASHINFER_QUANT_AVAILABLE: - print( - f"[info] flashinfer quant baseline unavailable: {_FLASHINFER_QUANT_REASON}" - ) - if not _AOT_QUANT_AVAILABLE: - print(f"[info] legacy AOT quant baseline unavailable: {_AOT_QUANT_REASON}") - benchmark.run(print_data=True) diff --git a/test/registered/jit/benchmark/bench_nvfp4_scaled_mm.py b/test/registered/jit/benchmark/bench_nvfp4_scaled_mm.py deleted file mode 100644 index feaacc5c0..000000000 --- a/test/registered/jit/benchmark/bench_nvfp4_scaled_mm.py +++ /dev/null @@ -1,189 +0,0 @@ -from __future__ import annotations - -import sys - -import torch -import triton - -from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark -from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant -from sglang.srt.utils import is_sm100_supported, is_sm120_supported -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci( - est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" -) - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max -BLOCK_SIZE = 16 -_NVFP4_SUPPORTED = is_sm100_supported() or is_sm120_supported() - -K_E2M1_TO_FLOAT = [ - 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, -] - - -def _dequantize_to_fp16( - tensor_fp4: torch.Tensor, tensor_sf: torch.Tensor, global_scale: torch.Tensor -): - m, packed_k = tensor_fp4.shape - k = packed_k * 2 - flat = tensor_fp4.flatten() - high = (flat & 0xF0) >> 4 - low = flat & 0x0F - f_h = torch.tensor([K_E2M1_TO_FLOAT[x] for x in high], device=tensor_fp4.device) - f_l = torch.tensor([K_E2M1_TO_FLOAT[x] for x in low], device=tensor_fp4.device) - val = torch.stack((f_l, f_h), dim=-1).reshape(m, k) - - rounded_m = ((m + 128 - 1) // 128) * 128 - scale_n = k // BLOCK_SIZE - rounded_n = ((scale_n + 4 - 1) // 4) * 4 - sf = tensor_sf.view(torch.float8_e4m3fn) - tmp = torch.reshape(sf, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4)) - tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) - scale = torch.reshape(tmp, (rounded_m, rounded_n))[:m, :scale_n].to(torch.float32) - scale = scale / global_scale - - return (val.view(m, scale_n, BLOCK_SIZE) * scale.unsqueeze(-1)).reshape(m, k) - - -def _aot_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: - out = torch.empty((a.shape[0], b.shape[0]), dtype=out_dtype, device=a.device) - torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default( - out, a, b, block_scale_a, block_scale_b, alpha - ) - return out - - -def _probe_legacy_aot_scaled_mm() -> tuple[bool, str]: - if not torch.cuda.is_available(): - return False, "CUDA is not available." - if not _NVFP4_SUPPORTED: - return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+." - try: - import sgl_kernel # noqa: F401 - except Exception as e: - return False, f"import sgl_kernel failed: {e}" - if not hasattr(torch.ops, "sgl_kernel"): - return False, "torch.ops.sgl_kernel is not registered." - op = getattr(torch.ops.sgl_kernel, "cutlass_scaled_fp4_mm", None) - if op is None or not hasattr(op, "default"): - return False, "torch.ops.sgl_kernel.cutlass_scaled_fp4_mm.default is missing." - try: - m, n, k = 16, 32, 64 - a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") - b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") - a_global_scale = ( - FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1) - ).to(torch.float32) - b_global_scale = ( - FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1) - ).to(torch.float32) - alpha = 1.0 / (a_global_scale * b_global_scale) - a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale) - b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale) - _aot_cutlass_scaled_fp4_mm(a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16) - torch.cuda.synchronize() - except Exception as e: - return False, f"calling AOT scaled_mm op failed: {e}" - return True, "" - - -_AOT_SCALED_MM_AVAILABLE, _AOT_SCALED_MM_REASON = _probe_legacy_aot_scaled_mm() - -shape_range = get_benchmark_range( - full_range=[(128, 4096, 4096), (512, 4096, 4096), (1024, 8192, 4096)], - ci_range=[(128, 4096, 4096)], -) - -line_vals = ["jit"] -line_names = ["JIT NVFP4 GEMM"] -styles = [("green", "-")] -if _AOT_SCALED_MM_AVAILABLE: - line_vals.append("aot_sgl_kernel") - line_names.append("AOT NVFP4 GEMM") - styles.append(("orange", "-")) -line_vals.append("torch_ref") -line_names.append("Torch Ref") -styles.append(("blue", "-")) - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["m", "n", "k"], - x_vals=shape_range, - x_log=False, - line_arg="provider", - line_vals=line_vals, - line_names=line_names, - styles=styles, - ylabel="us", - plot_name="nvfp4-scaled-mm-performance", - args={}, - ) -) -def benchmark(m, n, k, provider): - a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") - b = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") - - a_global_scale = ( - FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(a.flatten(), dim=-1) - ).to(torch.float32) - b_global_scale = ( - FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / torch.amax(b.flatten(), dim=-1) - ).to(torch.float32) - alpha = 1.0 / (a_global_scale * b_global_scale) - - a_fp4, a_sf = scaled_fp4_quant(a, a_global_scale) - b_fp4, b_sf = scaled_fp4_quant(b, b_global_scale) - - if provider == "jit": - fn = lambda: cutlass_scaled_fp4_mm( - a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16 - ) - elif provider == "aot_sgl_kernel": - fn = lambda: _aot_cutlass_scaled_fp4_mm( - a_fp4, b_fp4, a_sf, b_sf, alpha, torch.bfloat16 - ) - elif provider == "torch_ref": - a_ref = _dequantize_to_fp16(a_fp4, a_sf, a_global_scale) - b_ref = _dequantize_to_fp16(b_fp4, b_sf, b_global_scale) - fn = lambda: torch.matmul(a_ref, b_ref.t()) - else: - raise ValueError(f"Unknown provider: {provider}") - - return run_benchmark(fn) - - -if __name__ == "__main__": - if not _NVFP4_SUPPORTED: - print("[skip] NVFP4 scaled_mm benchmark requires sm100/sm120 with CUDA 12.8+.") - sys.exit(0) - if not _AOT_SCALED_MM_AVAILABLE: - print( - f"[info] legacy AOT scaled_mm baseline unavailable: {_AOT_SCALED_MM_REASON}" - ) - benchmark.run(print_data=True) diff --git a/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py index cc7a865d3..cfabb8ff5 100644 --- a/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py +++ b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py @@ -8,7 +8,6 @@ from pathlib import Path from typing import Any, Callable import flashinfer -import sgl_kernel import torch from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE @@ -38,7 +37,7 @@ WARMUP = 8 ITERS = 20 FLOAT4_E2M1_MAX = 6.0 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max -METHODS = ("cutlass", "flashinfer_auto", "flashinfer_cudnn") +METHODS = ("flashinfer_auto", "flashinfer_cudnn") def benchmark_provider( @@ -256,14 +255,6 @@ def run_shape_suite(shape_cases: list[dict[str, Any]]) -> list[dict[str, Any]]: } providers: dict[str, Callable[[], torch.Tensor]] = { - "cutlass": lambda: sgl_kernel.cutlass_scaled_fp4_mm( - quantized["x_fp4"], - quantized["w_fp4"], - quantized["x_sf"], - quantized["w_sf"], - quantized["alpha"], - DTYPE, - ), "flashinfer_auto": lambda: flashinfer.mm_fp4( quantized["x_fp4"], quantized["w_fp4"].T, diff --git a/test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py b/test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py index 9dbb19ae7..f7e15ef90 100644 --- a/test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py +++ b/test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py @@ -4,7 +4,6 @@ import flashinfer import pytest import torch -from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant from sglang.multimodal_gen.runtime.layers.quantization import ( modelopt_quant as diffusion_modelopt_quant, ) @@ -241,8 +240,6 @@ def _build_layer( def _resolve_mode(mode: str): - if mode == "jit_cutlass": - return scaled_fp4_quant, cutlass_scaled_fp4_mm, None if mode == "flashinfer2": return flashinfer.fp4_quantize, flashinfer.mm_fp4, "cudnn" if mode == "flashinfer_trtllm": @@ -281,7 +278,7 @@ def test_checkpoint_processing( not _nvfp4_supported(), reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs", ) -@pytest.mark.parametrize("mode", ["jit_cutlass", "flashinfer2"]) +@pytest.mark.parametrize("mode", ["flashinfer2"]) def test_flux2_shape_correctness(mode: str) -> None: m, n, k = FLUX2_PROJECTION_SHAPE quantize_op, gemm_op, gemm_backend = _resolve_mode(mode) @@ -306,25 +303,15 @@ def test_flux2_shape_correctness(mode: str) -> None: _dequantize_nvfp4(weight_fp4, weight_scale_swizzled, weight_global_scale).t(), ) - if gemm_backend is None: - actual = gemm_op( - x_fp4, - weight_fp4, - x_scale_swizzled, - weight_scale_swizzled, - alpha, - DTYPE, - ) - else: - actual = gemm_op( - x_fp4, - weight_fp4.t(), - x_scale_swizzled, - weight_scale_swizzled.t(), - alpha, - DTYPE, - backend=gemm_backend, - ) + actual = gemm_op( + x_fp4, + weight_fp4.t(), + x_scale_swizzled, + weight_scale_swizzled.t(), + alpha, + DTYPE, + backend=gemm_backend, + ) diff = _calc_diff(actual, expected.to(dtype=DTYPE)) assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{mode=}, {m=}, {n=}, {k=}, {diff=:.6f}" diff --git a/test/registered/jit/test_nvfp4_blockwise_moe.py b/test/registered/jit/test_nvfp4_blockwise_moe.py deleted file mode 100644 index c022628e3..000000000 --- a/test/registered/jit/test_nvfp4_blockwise_moe.py +++ /dev/null @@ -1,137 +0,0 @@ -import sys - -import pytest -import torch - -from sglang.jit_kernel.nvfp4 import ( - cutlass_fp4_group_mm, - scaled_fp4_experts_quant, - scaled_fp4_quant, -) -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large") -register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max - - -def _nvfp4_supported() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) - - -def _round_up(x: int, y: int) -> int: - return ((x + y - 1) // y) * y - - -def _build_expert_offsets( - m_per_expert: list[int], device: torch.device -) -> torch.Tensor: - offsets = [0] - for m in m_per_expert: - offsets.append(offsets[-1] + m) - return torch.tensor(offsets, dtype=torch.int32, device=device) - - -def _build_blockscale_offsets( - m_per_expert: list[int], device: torch.device -) -> torch.Tensor: - offsets = [0] - for m in m_per_expert: - offsets.append(offsets[-1] + _round_up(m, 128)) - return torch.tensor(offsets, dtype=torch.int32, device=device) - - -@pytest.mark.skipif( - not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" -) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_nvfp4_blockwise_moe_grouped_mm(dtype: torch.dtype) -> None: - torch.manual_seed(0) - device = torch.device("cuda") - - num_experts = 4 - m_per_expert = [33, 17, 48, 29] - n = 256 - k = 128 - - expert_offsets_full = _build_expert_offsets(m_per_expert, device) - blockscale_offsets_full = _build_blockscale_offsets(m_per_expert, device) - - total_m = int(expert_offsets_full[-1].item()) - a = torch.randn((total_m, k), device=device, dtype=dtype) * 0.1 - b = torch.randn((num_experts, n, k), device=device, dtype=dtype) * 0.1 - - a_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) - for i in range(num_experts): - start = int(expert_offsets_full[i].item()) - end = int(expert_offsets_full[i + 1].item()) - amax = a[start:end].abs().max().to(torch.float32) - a_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax - - b_global_scale = torch.empty((num_experts,), device=device, dtype=torch.float32) - for i in range(num_experts): - bmax = b[i].abs().max().to(torch.float32) - b_global_scale[i] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / bmax - - a_fp4, a_blockscale = scaled_fp4_experts_quant( - a, - a_global_scale, - expert_offsets_full, - blockscale_offsets_full, - topk=1, - ) - - b_fp4 = torch.empty((num_experts, n, k // 2), device=device, dtype=torch.uint8) - b_blockscale = torch.empty( - (num_experts, _round_up(n, 128), _round_up(k // 16, 4)), - device=device, - dtype=torch.float8_e4m3fn, - ) - for i in range(num_experts): - b_fp4_i, b_scale_i = scaled_fp4_quant(b[i], b_global_scale[i]) - b_fp4[i].copy_(b_fp4_i) - b_blockscale[i].copy_(b_scale_i) - - alphas = (1.0 / (a_global_scale * b_global_scale)).to(torch.float32) - - params = { - "ab_strides": torch.full((num_experts,), k, dtype=torch.int64, device=device), - "c_strides": torch.full((num_experts,), n, dtype=torch.int64, device=device), - "problem_sizes": torch.tensor( - [[m, n, k] for m in m_per_expert], dtype=torch.int32, device=device - ), - "expert_offsets": expert_offsets_full[:-1].contiguous(), - "blockscale_offsets": blockscale_offsets_full[:-1].contiguous(), - "a_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "b_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "out_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "a_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "b_scales_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "alpha_ptrs": torch.empty((num_experts,), dtype=torch.int64, device=device), - "layout_sfa": torch.empty((num_experts, 5), dtype=torch.int64, device=device), - "layout_sfb": torch.empty((num_experts, 5), dtype=torch.int64, device=device), - } - - out = cutlass_fp4_group_mm( - a_fp4, - b_fp4, - a_blockscale, - b_blockscale, - alphas, - dtype, - params, - ) - - ref = torch.empty((total_m, n), device=device, dtype=dtype) - for i in range(num_experts): - start = int(expert_offsets_full[i].item()) - end = int(expert_offsets_full[i + 1].item()) - ref[start:end] = torch.matmul(a[start:end], b[i].t()) - - torch.testing.assert_close(out, ref, atol=1e-1, rtol=1e-1) - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/jit/test_nvfp4_gemm.py b/test/registered/jit/test_nvfp4_gemm.py deleted file mode 100644 index f0fad1b86..000000000 --- a/test/registered/jit/test_nvfp4_gemm.py +++ /dev/null @@ -1,152 +0,0 @@ -import sys - -import pytest -import torch - -from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large") -register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) - - -def _nvfp4_supported() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) - - -DTYPES = [torch.float16, torch.bfloat16] -SHAPES = [ - (128, 128, 64), - (128, 128, 128), - (256, 128, 64), - (128, 256, 128), - (150, 128, 64), -] - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max - -K_E2M1_TO_FLOAT = [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, -] - - -def e2m1_to_fp32(int4_value: int) -> float: - sign_bit = int4_value & 0x8 - int4_abs_value = int4_value & 0x7 - float_result = K_E2M1_TO_FLOAT[int4_abs_value] - return -float_result if sign_bit else float_result - - -def break_fp4_bytes(a: torch.Tensor) -> torch.Tensor: - assert a.dtype == torch.uint8 - m, n = a.shape - a = a.flatten() - high_half_byte = (a & 0xF0) >> 4 - low_half_byte = a & 0x0F - f_h = torch.tensor([e2m1_to_fp32(x) for x in high_half_byte], device=a.device) - f_l = torch.tensor([e2m1_to_fp32(x) for x in low_half_byte], device=a.device) - return torch.stack((f_l, f_h), dim=-1).reshape(m, n * 2) - - -def convert_swizzled_to_linear( - a_sf_swizzled: torch.Tensor, m: int, k: int, block_size: int -) -> torch.Tensor: - sf_m, sf_k = a_sf_swizzled.shape - del sf_m, sf_k - 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 // block_size] - - -def dequantize_to_dtype( - tensor_fp4: torch.Tensor, - tensor_sf: torch.Tensor, - global_scale: torch.Tensor, - block_size: int = 16, -) -> torch.Tensor: - assert tensor_fp4.dtype == torch.uint8 - m, packed_k = tensor_fp4.shape - k = packed_k * 2 - tensor_f32 = break_fp4_bytes(tensor_fp4) - 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 - return (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) - - -def get_ref_results( - a_fp4: torch.Tensor, - b_fp4: torch.Tensor, - a_sf: torch.Tensor, - b_sf: torch.Tensor, - a_global_scale: torch.Tensor, - b_global_scale: torch.Tensor, - block_size: int, -) -> torch.Tensor: - a_in_dtype = dequantize_to_dtype(a_fp4, a_sf, a_global_scale, block_size=block_size) - b_in_dtype = dequantize_to_dtype(b_fp4, b_sf, b_global_scale, block_size=block_size) - return torch.matmul(a_in_dtype, b_in_dtype.t()) - - -@pytest.mark.skipif( - not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" -) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("shape", SHAPES) -def test_nvfp4_gemm(dtype: torch.dtype, shape: tuple[int, 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, - block_size, - ) - - 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__": - sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/jit/test_nvfp4_quant.py b/test/registered/jit/test_nvfp4_quant.py deleted file mode 100644 index 1ca9c68aa..000000000 --- a/test/registered/jit/test_nvfp4_quant.py +++ /dev/null @@ -1,225 +0,0 @@ -import sys - -import pytest -import torch - -from sglang.jit_kernel.nvfp4 import ( - scaled_fp4_grouped_quant, - scaled_fp4_quant, - silu_and_mul_scaled_fp4_grouped_quant, -) - -try: - from sgl_kernel import silu_and_mul as _sgl_silu_and_mul -except Exception: - _sgl_silu_and_mul = None - -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large") -register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) - - -def _nvfp4_supported() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) - - -def _silu_and_mul_reference(x: torch.Tensor) -> torch.Tensor: - if _sgl_silu_and_mul is not None: - return _sgl_silu_and_mul(x) - k = x.shape[-1] // 2 - return torch.nn.functional.silu(x[:, :, :k]) * x[:, :, k:] - - -DTYPES = [torch.float16, torch.bfloat16] -SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)] -PAD_SHAPES = [ - (90, 64), - (150, 64), - (128, 48), - (128, 80), -] - -FLOAT4_E2M1_MAX = 6.0 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max -BLOCK_SIZE = 16 - -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, -] - - -def cast_from_fp4(x: torch.Tensor, m: int, n: int) -> torch.Tensor: - v_2nd = (x & 0xF).to(torch.long) - v_1st = ((x >> 4) & 0xF).to(torch.long) - c = torch.stack((v_2nd, v_1st), dim=-1).flatten() - lut = torch.tensor(E2M1_TO_FLOAT32, device=x.device, dtype=torch.float32) - return lut[c].reshape(m, n) - - -def cast_to_fp4(x: torch.Tensor) -> torch.Tensor: - 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) - return 0.0 if x == 0 else 1.0 / x - - -def ref_nvfp4_quant(x: torch.Tensor, global_scale: torch.Tensor): - 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: torch.Tensor, m: int, n: int) -> torch.Tensor: - rounded_m = ((m + 128 - 1) // 128) * 128 - scale_n = n // BLOCK_SIZE - rounded_n = ((scale_n + 4 - 1) // 4) * 4 - 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( - not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" -) -@pytest.mark.parametrize("dtype", DTYPES) -@pytest.mark.parametrize("shape", SHAPES) -def test_quantize_to_fp4(dtype: torch.dtype, shape: tuple[int, int]) -> None: - torch.manual_seed(42) - m, n = shape - - x = torch.randn((m, n), dtype=dtype, device="cuda") - 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( - not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" -) -@pytest.mark.parametrize("shape", PAD_SHAPES) -def test_quantize_to_fp4_padded(shape: tuple[int, int]) -> None: - torch.manual_seed(42) - m, n = shape - x = torch.randn((m, n), dtype=torch.float16, device="cuda") - - 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( - not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" -) -@pytest.mark.parametrize("shape", [(2, 128, 512), (2, 100, 128)]) -def test_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None: - torch.manual_seed(42) - l, m, k = shape - - x = torch.randn((l, m, k), dtype=torch.bfloat16, device="cuda") - mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda") - 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_quant(x, x_sf_global, mask) - output = output.permute(2, 0, 1) - 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]]) - 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( - not _nvfp4_supported(), reason="NVFP4 requires compute capability >= 10.0" -) -@pytest.mark.parametrize("shape", [(4, 96, 256), (8, 128, 512)]) -def test_silu_and_mul_quantize_to_fp4_grouped(shape: tuple[int, int, int]) -> None: - torch.manual_seed(42) - l, m, k = shape - - x = torch.randn((l, m, k * 2), dtype=torch.bfloat16, device="cuda") - mask = torch.randint(1, max(2, m // 2), (l,), dtype=torch.int32, device="cuda") - - ref_y = _silu_and_mul_reference(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_quant(ref_y, y_sf_global, mask) - output, output_scales = silu_and_mul_scaled_fp4_grouped_quant(x, y_sf_global, mask) - - output = output.permute(2, 0, 1) - ref_output = ref_output.permute(2, 0, 1) - - 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]]) - 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__": - sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py b/test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py deleted file mode 100644 index f1c71bc65..000000000 --- a/test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py +++ /dev/null @@ -1,337 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Unit test for the fused JIT op ``silu_and_mul_scaled_fp4_experts_quant_packed`` -(introduced in PR #18612). - -On the CUTLASS NVFP4 MoE intermediate, the op fuses the previous two-step path - - intermediate = silu_and_mul(c1) # SiLU(gate) * up - fp4, sf = scaled_fp4_experts_quant(intermediate) # NVFP4 expert quant - -into a single kernel - - fp4, sf = silu_and_mul_scaled_fp4_experts_quant_packed(c1, ...) - -This test compares the fused op against that exact unfused -``silu_and_mul`` + ``scaled_fp4_experts_quant`` path **with uneven expert offsets**: -experts deliberately receive very different token counts, including tiny experts and -experts whose row count is not a multiple of the 128-row block-scale padding. That is -precisely the regime that stresses the per-expert ``expert_offsets`` / -``blockscale_offsets`` indexing the fusion has to get right. - -It follows the two existing siblings: - * ``test_silu_and_mul_quantize_to_fp4_grouped`` (the grouped/masked variant) -- the - unfused path is the reference, and the fused output must match it bit-exactly - (packed FP4 nibbles + recovered block scales), and - * ``test_nvfp4_blockwise_moe`` (the expert-offset variant) -- offsets are built from - an explicit, non-uniform per-expert token list. - -A high-precision ``F.silu(gate) * up`` check additionally grounds the unfused path so a -bug shared by both kernels cannot produce a false (vacuous) pass. - - pytest python/sglang/jit_kernel/tests/test_silu_and_mul_scaled_fp4_experts_quant_packed.py -v -""" - -import sys - -import pytest -import torch -import triton -from torch.nn import functional as F - -from sglang.jit_kernel.activation import silu_and_mul -from sglang.jit_kernel.nvfp4 import ( - scaled_fp4_experts_quant, - silu_and_mul_scaled_fp4_experts_quant_packed, -) -from sglang.test.ci.ci_register import register_cuda_ci - -# The NVFP4 expert-quant kernels are Blackwell-only (sm100a), so this runs on -# the B200 unit suite. -register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200") - -FLOAT8_E4M3_MAX = 448.0 -FLOAT4_E2M1_MAX = 6.0 -BLOCK_SIZE = 16 -kE2M1ToFloat = torch.tensor( - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 -) - - -def _nvfp4_supported() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0) - - -def _round_up(x: int, y: int) -> int: - return ((x + y - 1) // y) * y - - -# --------------------------------------------------------------------------- # -# Offset builders (mirror test_nvfp4_blockwise_moe.py). -# expert_offsets: cumulative *actual* per-expert rows ([E+1] int32) -# blockscale_offsets: cumulative rows padded up to 128 per expert ([E+1] int32) -# A non-uniform ``m_per_expert`` makes both offset tensors uneven. -# --------------------------------------------------------------------------- # -def _build_expert_offsets(m_per_expert, device) -> torch.Tensor: - offsets = [0] - for m in m_per_expert: - offsets.append(offsets[-1] + m) - return torch.tensor(offsets, dtype=torch.int32, device=device) - - -def _build_blockscale_offsets(m_per_expert, device) -> torch.Tensor: - offsets = [0] - for m in m_per_expert: - offsets.append(offsets[-1] + _round_up(m, 128)) - return torch.tensor(offsets, dtype=torch.int32, device=device) - - -# --------------------------------------------------------------------------- # -# FP4 dequant / scale-recovery helpers (mirror test/registered/kernels/test_fp4_moe.py) -# --------------------------------------------------------------------------- # -def break_fp4_bytes(a: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - assert a.dtype == torch.uint8 - m, n = a.shape - a_flat = a.flatten() - high = (a_flat & 0xF0) >> 4 - low = a_flat & 0x0F - combined = torch.stack((low, high), dim=1).flatten() - signs = (combined & 0x08).to(torch.bool) - abs_vals = (combined & 0x07).to(torch.long) - kE2M1 = kE2M1ToFloat.to(device=a.device) - values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0) - return values.reshape(m, n * 2).to(dtype=dtype) - - -def convert_swizzled_to_linear( - a_sf_swizzled: torch.Tensor, m: int, k: int, block_size: int -) -> torch.Tensor: - """De-swizzle one expert's block-scale region and drop the 128-row padding tail.""" - 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_nvfp4_to_dtype( - tensor_fp4: torch.Tensor, - tensor_sf: torch.Tensor, - global_scale: torch.Tensor, - dtype: torch.dtype, - device: torch.device, - block_size: int = 16, -) -> torch.Tensor: - """Dequantize one expert's packed FP4 (m, k//2) + swizzled block scales.""" - 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 - out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k) - return out.to(dtype=dtype) - - -def _recover_block_scales( - sf: torch.Tensor, s0: int, s1: int, m_e: int, n: int -) -> torch.Tensor: - """De-swizzled, un-padded block scales (float32) for one expert's region.""" - block = sf[s0:s1].contiguous().view(torch.float8_e4m3fn) - return convert_swizzled_to_linear(block, m_e, n, BLOCK_SIZE).to(torch.float32) - - -def _rel_l2(a: torch.Tensor, b: torch.Tensor) -> float: - return (a.float() - b.float()).norm().item() / b.float().norm().clamp_min( - 1e-9 - ).item() - - -# --------------------------------------------------------------------------- # -# Uneven per-expert token counts. Each list is deliberately NON-uniform so that -# expert_offsets / blockscale_offsets are uneven, exercising: -# * tiny experts (1, 5, 7 tokens), -# * an exactly-128 expert (no padding), -# * experts straddling the 128-row block-scale padding (130, 200, 384). -# --------------------------------------------------------------------------- # -UNEVEN_M_PER_EXPERT = [ - [33, 17, 48, 29], # all < 128 (matches test_nvfp4_blockwise_moe) - [1, 128, 200, 5, 64], # tiny + exactly-128 + cross-128 - [130, 1, 384, 17, 96, 7], # heavy skew, large dynamic range -] -NS = [256, 768] # 768 == Qwen3-30B-A3B moe_intermediate_size -DTYPES = [torch.bfloat16, torch.float16] - - -@pytest.mark.skipif( - not _nvfp4_supported(), - reason="NVFP4 fused expert-quant kernel requires compute capability >= 10.0 (B200/SM100).", -) -@pytest.mark.parametrize("m_per_expert", UNEVEN_M_PER_EXPERT) -@pytest.mark.parametrize("n", NS) -@pytest.mark.parametrize("dtype", DTYPES) -@torch.inference_mode() -def test_fused_matches_unfused_uneven_offsets(m_per_expert, n, dtype): - torch.manual_seed(0) - device = torch.device("cuda") - num_experts = len(m_per_expert) - - # --- uneven expert offsets --- - expert_offsets = _build_expert_offsets(m_per_expert, device) - blockscale_offsets = _build_blockscale_offsets(m_per_expert, device) - total_m = int(expert_offsets[-1].item()) - counts = torch.tensor(m_per_expert) - assert ( - counts.max() >= 2 * counts.min() - ), "expert offsets must be uneven for this test" - - # gate+up concatenated input (m, 2n); /5 keeps values in a sane FP4 range. - c1 = torch.randn((total_m, 2 * n), dtype=dtype, device=device) / 5.0 - gate, up = c1[:, :n].float(), c1[:, n:].float() - ref = F.silu(gate) * up # high-precision SiLU(gate) * up, (total_m, n) fp32 - - # Per-expert global scale, exactly like cutlass_moe builds a2_gscale. - gscale = torch.empty(num_experts, dtype=torch.float32, device=device) - for e in range(num_experts): - r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1]) - amax = ref[r0:r1].abs().max().clamp_min(1e-6) - gscale[e] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax - - # topk only gates the buffer-size assertion in the wrapper; the per-expert - # layout is driven entirely by the offsets. Both paths use the same value. - topk = 1 - - # ---- fused (new op) ---- - fused_fp4, fused_sf = silu_and_mul_scaled_fp4_experts_quant_packed( - c1, gscale, expert_offsets, blockscale_offsets, topk - ) - - # ---- unfused (the exact path the op replaced) ---- - intermediate = torch.empty((total_m, n), dtype=dtype, device=device) - silu_and_mul(c1, intermediate) - unf_fp4, unf_sf = scaled_fp4_experts_quant( - intermediate, gscale, expert_offsets, blockscale_offsets, topk - ) - - assert fused_fp4.shape == unf_fp4.shape == (total_m, n // 2) - - # Per-expert, bit-exact comparison honoring the uneven offsets. - for e in range(num_experts): - r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1]) - s0, s1 = int(blockscale_offsets[e]), int(blockscale_offsets[e + 1]) - m_e = r1 - r0 - - # (1) Packed FP4 nibbles are identical: same NVFP4 quantizer, same fused - # SiLU(gate)*up rounded to the storage dtype before quantization. - torch.testing.assert_close( - fused_fp4[r0:r1], unf_fp4[r0:r1], msg=f"FP4 bytes differ for expert {e}" - ) - - # (2) Recovered (de-swizzled, un-padded) block scales are identical. - torch.testing.assert_close( - _recover_block_scales(fused_sf, s0, s1, m_e, n), - _recover_block_scales(unf_sf, s0, s1, m_e, n), - msg=f"block scales differ for expert {e}", - ) - - # (3) Grounding: the unfused path really reproduces SiLU(gate)*up within FP4 - # error, so (1)/(2) cannot pass vacuously on a bug shared by both kernels. - deq = dequantize_nvfp4_to_dtype( - unf_fp4[r0:r1].contiguous(), - unf_sf[s0:s1].contiguous(), - gscale[e], - dtype, - device, - BLOCK_SIZE, - ) - assert ( - _rel_l2(deq, ref[r0:r1]) < 0.2 - ), f"expert {e}: unfused dequant does not match SiLU(gate)*up reference" - - -# --------------------------------------------------------------------------- # -# Performance. The fusion removes, on the MoE down-projection input, one -# intermediate buffer allocation, one extra kernel launch, and a full HBM -# round-trip of the SiLU(gate)*up result. The speedup is measured under CUDA -# graphs -- the steady-state GPU memory-traffic saving, matching how SGLang -# executes graphed decode (the credible, low-noise number; eager wall-clock is -# dominated by launch/dispatch overhead and is too noisy to assert on). The -# assert is only a conservative regression floor; the printed speedup is the real -# result. Mirrors test_cutedsl_gdn_performance, in the same kernel unit suite. -# -# Tokens are spread evenly across experts here (the representative throughput -# case) at realistic Qwen3-30B-A3B MoE dims (n=768, 128 experts), swept from a -# decode batch up to a prefill chunk; the uneven-offset corner cases are covered -# by the correctness test above. -# --------------------------------------------------------------------------- # -PERF_SHAPES = [ - (1024, 768, 128), - (4096, 768, 128), - (16384, 768, 128), -] - - -def _even_offsets(total_tokens, num_experts, device): - base, rem = divmod(total_tokens, num_experts) - m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)] - return ( - _build_expert_offsets(m_per_expert, device), - _build_blockscale_offsets(m_per_expert, device), - ) - - -@pytest.mark.skipif( - not _nvfp4_supported(), - reason="NVFP4 fused expert-quant kernel requires compute capability >= 10.0 (B200/SM100).", -) -@pytest.mark.parametrize("total_tokens,n,num_experts", PERF_SHAPES) -@torch.inference_mode() -def test_fused_perf_not_regressed(total_tokens, n, num_experts): - device = torch.device("cuda") - dtype = torch.bfloat16 - expert_offsets, blockscale_offsets = _even_offsets( - total_tokens, num_experts, device - ) - c1 = torch.randn((total_tokens, 2 * n), dtype=dtype, device=device) / 5.0 - gscale = torch.empty(num_experts, dtype=torch.float32, device=device) - for e in range(num_experts): - r0, r1 = int(expert_offsets[e]), int(expert_offsets[e + 1]) - amax = c1[r0:r1].abs().max().to(torch.float32).clamp_min(1e-6) - gscale[e] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / amax - topk = 1 - - def fused(): - silu_and_mul_scaled_fp4_experts_quant_packed( - c1, gscale, expert_offsets, blockscale_offsets, topk - ) - - def unfused(): - # The exact path the op replaced: alloc the intermediate, SiLU*mul into - # it, then quantize it -- one extra buffer + kernel + HBM round-trip. - intermediate = torch.empty((total_tokens, n), dtype=dtype, device=device) - silu_and_mul(c1, intermediate) - scaled_fp4_experts_quant( - intermediate, gscale, expert_offsets, blockscale_offsets, topk - ) - - g_f = triton.testing.do_bench_cudagraph(fused) # ms, median - g_u = triton.testing.do_bench_cudagraph(unfused) - cuda_graph_speedup = g_u / g_f - print( - f"\n [PERF] tokens={total_tokens:>6} n={n} E={num_experts}: " - f"unfused {g_u * 1e3:6.1f}us fused {g_f * 1e3:6.1f}us " - f"cuda-graph speedup = {cuda_graph_speedup:.2f}x" - ) - # Regression guard only: the fusion must not make this op slower. The actual - # win carries a wide margin over this floor, so shared-runner noise cannot - # flake it. - assert ( - cuda_graph_speedup >= 1.05 - ), f"fused regressed under cuda-graph: {cuda_graph_speedup:.2f}x" - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/kernels/test_fp4_moe.py b/test/registered/kernels/test_fp4_moe.py index 369741438..9fb940eba 100644 --- a/test/registered/kernels/test_fp4_moe.py +++ b/test/registered/kernels/test_fp4_moe.py @@ -8,9 +8,6 @@ from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_m 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 from sglang.test.ci.ci_register import register_cuda_ci @@ -282,13 +279,9 @@ def check_moe( w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax - w1_q[expert], w1_blockscale[expert] = scaled_fp4_quant( - w1[expert], w1_gs[expert] - ) + w1_q[expert], w1_blockscale[expert] = fp4_quantize(w1[expert], w1_gs[expert]) - w2_q[expert], w2_blockscale[expert] = scaled_fp4_quant( - w2[expert], w2_gs[expert] - ) + w2_q[expert], w2_blockscale[expert] = fp4_quantize(w2[expert], w2_gs[expert]) score = torch.randn((m, e), device="cuda", dtype=dtype) @@ -319,7 +312,7 @@ def check_moe( a_global_scale = ( (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1) ).to(torch.float32) - a_fp4, a_scale_interleaved = scaled_fp4_quant(a, a_global_scale) + a_fp4, a_scale_interleaved = fp4_quantize(a, a_global_scale) _, m_k = a_fp4.shape a_in_dtype = dequantize_nvfp4_to_dtype( a_fp4, @@ -365,53 +358,6 @@ def check_moe( torch.testing.assert_close(torch_output, test_output, atol=1e-1, rtol=1e-1) -@pytest.mark.parametrize("m,n,k", MNK_FACTORS) -@pytest.mark.parametrize("e", [40, 64, 256]) -@pytest.mark.parametrize("topk", [1, 6, 8]) -@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) -@torch.inference_mode() -def test_cutlass_fp4_moe_no_graph( - m: int, n: int, k: int, e: int, topk: int, dtype: torch.dtype -): - def cutlass_moe_impl( - a, - topk_weights, - topk_ids, - w1_q, - w2_q, - a1_gs, - w1_blockscale, - w1_alphas, - a2_gs, - w2_blockscale, - w2_alphas, - ): - params = CutlassMoEParams( - CutlassMoEType.BlockscaledFP4, - device=a.device, - num_experts=e, - intermediate_size_per_partition=n, # n - hidden_size=k, - ) # k - return cutlass_moe_fp4( - a=a, - a1_gscale=a1_gs, - w1_fp4=w1_q, - w1_blockscale=w1_blockscale, - w1_alphas=w1_alphas, - a2_gscale=a2_gs, - w2_fp4=w2_q, - w2_blockscale=w2_blockscale, - w2_alphas=w2_alphas, - topk_weights=topk_weights, - topk_ids=topk_ids, - params=params, - apply_router_weight_on_input=False, - ) - - check_moe(m, n, k, e, topk, dtype, cutlass_moe_impl, flip_w13=False) - - @pytest.mark.parametrize("m,n,k", MNK_FACTORS) @pytest.mark.parametrize("e", [40, 64, 256]) @pytest.mark.parametrize("topk", [1, 6, 8]) @@ -454,5 +400,4 @@ def test_flashinfer_fp4_moe_no_graph( if __name__ == "__main__": - test_cutlass_fp4_moe_no_graph(224, 1024, 1024, 256, 8, torch.half) test_flashinfer_fp4_moe_no_graph(224, 1024, 1024, 256, 8, torch.half) diff --git a/test/registered/unit/test_server_args_migration.py b/test/registered/unit/test_server_args_migration.py index a03b3a224..59f6a0a96 100644 --- a/test/registered/unit/test_server_args_migration.py +++ b/test/registered/unit/test_server_args_migration.py @@ -45,10 +45,10 @@ class TestServerArgsAnnotatedCli(CustomTestCase): def test_cli_name_differs_from_field_name(self): """cli_name maps a different CLI flag to the dataclass field via dest.""" sa = self._parse( - ["--fp8-gemm-backend", "triton", "--fp4-gemm-backend", "cutlass"] + ["--fp8-gemm-backend", "triton", "--fp4-gemm-backend", "marlin"] ) self.assertEqual(sa.fp8_gemm_runner_backend, "triton") - self.assertEqual(sa.fp4_gemm_runner_backend, "cutlass") + self.assertEqual(sa.fp4_gemm_runner_backend, "marlin") def test_nargs_question_with_const(self): """nargs='?' + const='' for --model-checksum."""