minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Refactor FP4 quantization and remove deprecated JIT kernels (#30448)

Browse Source 
Co-authored-by: root <root@sgl-b300-inference.datacrunch.io>
This commit is contained in:
Brayden Zhong
2026-07-14 09:22:07 +08:00
committed by GitHub
co-authored by root
parent 0c01971eeb
commit 9756f768a6
38 changed files with 120 additions and 5705 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/flashinfer_allreduce_fusion/README.md
+1 -1
View File
@@ -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.
benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py
+8 -8
View File
@@ -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
benchmark/kernels/quantization/bench_fp4_quant.py
-137
View File
@@ -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()
docs_new/docs/advanced_features/quantization.mdx
+1 -6
View File
@@ -284,11 +284,6 @@ Backend selection applies to **blockwise FP8**, **MXFP8** (dense linear), and **
<td>SM80+</td>
<td>Auto-selects: <code>flashinfer_cutedsl</code> on SM100; <code>marlin</code> on SM80-SM90; <code>flashinfer_cutlass</code> otherwise (including SM120)</td>
</tr>
<tr>
<td><code>cutlass</code></td>
<td>SM100/120</td>
<td>SGLang CUTLASS kernel</td>
</tr>
<tr>
<td><code>flashinfer_cutlass</code></td>
<td>SM100/120</td>
@@ -317,7 +312,7 @@ Backend selection applies to **blockwise FP8**, **MXFP8** (dense linear), and **
</tbody>
</table>
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
docs_new/docs/advanced_features/server_arguments.mdx
+2 -2
View File
@@ -1436,9 +1436,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--fp4-gemm-backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default; selects <code>flashinfer_cutedsl</code> on SM100, <code>marlin</code> on SM80-SM90, <code>flashinfer_cutlass</code> 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.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Choose the runner backend for NVFP4 GEMM operations. Options: 'auto' (default; selects <code>flashinfer_cutedsl</code> on SM100, <code>marlin</code> on SM80-SM90, <code>flashinfer_cutlass</code> 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.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`auto`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>cutlass</code>, <code>flashinfer_cudnn</code>, <code>flashinfer_cutedsl</code>, <code>flashinfer_cutlass</code>, <code>flashinfer_trtllm</code>, <code>marlin</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>flashinfer_cudnn</code>, <code>flashinfer_cutedsl</code>, <code>flashinfer_cutlass</code>, <code>flashinfer_trtllm</code>, <code>marlin</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--disable-flashinfer-autotune`</td>
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_expert_quant.cuh
-806
View File
@@ -1,806 +0,0 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include "nvfp4_quant.cuh"
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
using namespace host;
// Quantizes the provided PackedVec into the uint32_t output
template <class Type, bool UE8M0_SF = false>
SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& 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<uint32_t&>(SFValue) >> 23;
fp8SFVal = tmp & 0xff;
// Convert back to fp32.
reinterpret_cast<uint32_t&>(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<float2>(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 <class Type>
SGL_DEVICE void silu_and_mul(PackedVec<Type>& x_vec, const PackedVec<Type>& 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<float2>(x_vec.elts[i]);
y[i] = device::cast<float2>(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<packed_t<Type>>(x[i]);
}
}
// Use UE4M3 by default.
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
__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<Type>;
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<int4*>(local_offsets) =
__ldca(reinterpret_cast<const int4*>(&input_offset_by_experts[chunk_start]));
*reinterpret_cast<int4*>(local_offsets + 4) =
__ldca(reinterpret_cast<const int4*>(&input_offset_by_experts[chunk_start + 4]));
*reinterpret_cast<int4*>(local_offsets + 8) =
__ldca(reinterpret_cast<const int4*>(&input_offset_by_experts[chunk_start + 8]));
*reinterpret_cast<int4*>(local_offsets + 12) =
__ldca(reinterpret_cast<const int4*>(&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<PackedVec const*>(in)[inOffset];
if (use_mask || use_silu_and_mul) {
PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(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<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(
rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
}
#endif
}
// Use UE4M3 by default.
template <class Type, bool UE8M0_SF = false>
__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<Type>;
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<PackedVec const*>(in)[inOffset];
if (use_silu_and_mul) {
PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(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<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(
rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
}
#endif
}
// Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
__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<Type>;
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<int4*>(&shared_input_offsets[i]) = *reinterpret_cast<const int4*>(&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<PackedVec const*>(in)[inOffset];
if (use_mask || use_silu_and_mul) {
PackedVec in_vec_mul = reinterpret_cast<PackedVec const*>(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<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(
rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
}
#endif
}
template <typename T>
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<int>((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<T, false><<<grid, block, 0, stream>>>(
m_topk,
k,
reinterpret_cast<T*>(input),
reinterpret_cast<float*>(input_global_scale),
reinterpret_cast<uint32_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<int32_t*>(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<T, false, false><<<grid, block, shared_mem_size, stream>>>(
m_topk,
k,
reinterpret_cast<T*>(input),
reinterpret_cast<float*>(input_global_scale),
reinterpret_cast<uint32_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(mask),
n_experts,
use_silu_and_mul);
} else {
cvt_fp16_to_fp4<T, false, true><<<grid, block, shared_mem_size, stream>>>(
m_topk,
k,
reinterpret_cast<T*>(input),
reinterpret_cast<float*>(input_global_scale),
reinterpret_cast<uint32_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(mask),
n_experts,
use_silu_and_mul);
}
} else {
if (n_experts >= 16) {
cvt_fp16_to_fp4<T, false, false><<<grid, block, 0, stream>>>(
m_topk,
k,
reinterpret_cast<T*>(input),
reinterpret_cast<float*>(input_global_scale),
reinterpret_cast<uint32_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(mask),
n_experts,
/* bool low_latency */ true,
use_silu_and_mul);
} else {
cvt_fp16_to_fp4<T, false, true><<<grid, block, 0, stream>>>(
m_topk,
k,
reinterpret_cast<T*>(input),
reinterpret_cast<float*>(input_global_scale),
reinterpret_cast<uint32_t*>(output),
reinterpret_cast<uint32_t*>(output_scale),
reinterpret_cast<uint32_t*>(input_offset_by_experts),
reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
reinterpret_cast<int32_t*>(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<fp16_t, bf16_t>()
.template with_device<kDLCUDA>(device)
.verify(input);
TensorMatcher({MTopK, OutputCols}) //
.with_dtype<uint8_t>()
.with_device(device)
.verify(output);
TensorMatcher({OutputScaleRows, OutputScaleCols}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(output_scale);
TensorMatcher({NExperts}) //
.with_dtype<float>()
.with_device(device)
.verify(input_global_scale);
TensorMatcher({OffsetSize}) //
.with_dtype<int32_t>()
.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<int>(MTopK.unwrap());
const auto k = static_cast<int>(K.unwrap());
RuntimeCheck(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
const auto n_experts = static_cast<int>(NExperts.unwrap());
const auto offset_size = static_cast<int>(OffsetSize.unwrap());
RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch");
RuntimeCheck(static_cast<int>(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<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
if (host::is_type<fp16_t>(input.dtype())) {
quant_impl<half>(
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<fp16_t, bf16_t>()
.template with_device<kDLCUDA>(device)
.verify(input);
TensorMatcher({MTopK, OutputCols}) //
.with_dtype<uint8_t>()
.with_device(device)
.verify(output);
TensorMatcher({OutputScaleRows, OutputScaleCols}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(output_scale);
TensorMatcher({NExperts}) //
.with_dtype<float>()
.with_device(device)
.verify(input_global_scale);
TensorMatcher({NExperts}) //
.with_dtype<int32_t>()
.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<int>(MTopK.unwrap());
const auto k_by_2 = static_cast<int>(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<int>(NExperts.unwrap());
RuntimeCheck(static_cast<int>(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<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
if (host::is_type<fp16_t>(input.dtype())) {
quant_impl<half>(
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<fp16_t, bf16_t>()
.template with_device<kDLCUDA>(device)
.verify(input);
TensorMatcher({MTopK, OutputCols}) //
.with_dtype<uint8_t>()
.with_device(device)
.verify(output);
TensorMatcher({OutputScaleRows, OutputScaleCols}) //
.with_dtype<int32_t>()
.with_device(device)
.verify(output_scale);
TensorMatcher({NExperts}) //
.with_dtype<float>()
.with_device(device)
.verify(input_global_scale);
TensorMatcher({OffsetSize}) //
.with_dtype<int32_t>()
.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<int>(MTopK.unwrap());
const auto k_by_2 = static_cast<int>(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<int>(NExperts.unwrap());
const auto offset_size = static_cast<int>(OffsetSize.unwrap());
RuntimeCheck(offset_size == n_experts + 1, "input/output offset size mismatch");
RuntimeCheck(static_cast<int>(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<int>(OutputScaleCols.unwrap()) * 4 == padded_k, "output_scale second dim mismatch");
const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
if (host::is_type<fp16_t>(input.dtype())) {
quant_impl<half>(
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);
}
}
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant.cuh
-160
View File
@@ -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 <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <cutlass/arch/config.h>
#include <cuda.h>
#include <cuda_fp8.h>
#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 <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
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<uint8_t*>(SFout) + SFOffset;
}
#endif
return nullptr;
}
// Define a 16 bytes packed data type.
template <class Type>
struct PackedVec {
packed_t<Type> elts[4];
};
template <>
struct PackedVec<__nv_fp8_e4m3> {
__nv_fp8x2_e4m3 elts[8];
};
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_entry.cuh
-87
View File
@@ -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 <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
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);
}
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_quant_kernels.cuh
-241
View File
@@ -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 <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include "nvfp4_quant.cuh"
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
using namespace host;
// Quantizes the provided PackedVec into the uint32_t output
template <class Type, bool UE8M0_SF = false>
SGL_DEVICE uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& 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<float>(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<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++) {
if constexpr (std::is_same_v<Type, half>) {
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 <class Type, bool UE8M0_SF = false>
__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<Type>;
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<PackedVec const*>(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<uint32_t, CVT_FP4_NUM_THREADS_PER_SF>(rowIdx, colIdx, numCols, SFout);
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
}
}
#endif
}
template <typename T>
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<T, true><<<grid, block, 0, stream>>>(
m, n, input, SFScale, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(SFOuput));
} else {
cvt_fp16_to_fp4<T, false><<<grid, block, 0, stream>>>(
m, n, input, SFScale, reinterpret_cast<uint32_t*>(output), reinterpret_cast<uint32_t*>(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<uint8_t>(output.dtype()), "output must be uint8");
RuntimeCheck(host::is_type<int32_t>(output_sf.dtype()), "output_sf must be int32");
RuntimeCheck(host::is_type<float>(input_sf.dtype()), "input_sf must be float32");
RuntimeCheck(
host::is_type<fp16_t>(input.dtype()) || host::is_type<bf16_t>(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<int32_t>(input.size(0));
const int32_t n = static_cast<int32_t>(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<int>(runtime::get_sm_count(device_id));
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
const cudaStream_t stream = LaunchKernel::resolve_device(input.device());
constexpr bool useUE8M0 = false;
if (host::is_type<fp16_t>(input.dtype())) {
auto input_ptr = reinterpret_cast<half const*>(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);
}
}
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_common.cuh
-66
View File
@@ -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 <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include <cstddef>
#include <cstdint>
#include <cuda_runtime.h>
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<int64_t>(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;
}
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_entry.cuh
-34
View File
@@ -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 <sgl_kernel/tensor.h>
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);
}
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_kernels.cuh
-146
View File
@@ -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<uint8_t>(A.dtype()), "a must be uint8");
RuntimeCheck(host::is_type<uint8_t>(B.dtype()), "b must be uint8");
RuntimeCheck(host::is_type<fp8_e4m3_t>(A_sf.dtype()), "scale_a must be float8_e4m3fn");
RuntimeCheck(host::is_type<fp8_e4m3_t>(B_sf.dtype()), "scale_b must be float8_e4m3fn");
RuntimeCheck(host::is_type<float>(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<int64_t>(A.size(0));
const auto n = static_cast<int64_t>(B.size(0));
const auto k = static_cast<int64_t>(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<fp16_t>(D.dtype())) {
cutlass_fp4_f16_gemm_dispatch_sm120(
D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
} else if (host::is_type<bf16_t>(D.dtype())) {
cutlass_fp4_bf16_gemm_dispatch_sm120(
D, A, B, A_sf, B_sf, alpha, static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), stream);
} else {
Panic("Unsupported output data type of nvfp4 mm sm120");
}
} else {
if (host::is_type<fp16_t>(D.dtype())) {
cutlassFp4GemmDispatchSm100<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (host::is_type<bf16_t>(D.dtype())) {
cutlassFp4GemmDispatchSm100<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (host::is_type<float>(D.dtype())) {
cutlassFp4GemmDispatchSm100<float>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else {
Panic("Unsupported output data type of nvfp4 mm");
}
}
}
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm100.cuh
-305
View File
@@ -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 <typename T>
struct KernelConfigM128 {
using OutputType = T;
using MmaTileShape = Shape<_128, _256, _256>;
using ClusterShape = Shape<int, int, _1>;
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 <typename T>
const dim3 KernelConfigM128<T>::preferred_cluster(1, 4, 1);
template <typename T>
const dim3 KernelConfigM128<T>::fallback_cluster(1, 2, 1);
// Config(half_t/bfloat16_t) for M <= 256
template <typename T>
struct KernelConfigM256 {
using OutputType = T;
using MmaTileShape = Shape<_256, _256, _256>;
using ClusterShape = Shape<int, int, _1>;
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 <typename T>
const dim3 KernelConfigM256<T>::preferred_cluster(2, 4, 1);
template <typename T>
const dim3 KernelConfigM256<T>::fallback_cluster(2, 1, 1);
// Config(half_t/bfloat16_t) for 256 < M <= 1024
template <typename T>
struct KernelConfigDefault {
using OutputType = T;
using MmaTileShape = Shape<_256, _256, _256>;
using ClusterShape = Shape<int, int, _1>;
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 <typename T>
const dim3 KernelConfigDefault<T>::preferred_cluster(2, 4, 1);
template <typename T>
const dim3 KernelConfigDefault<T>::fallback_cluster(2, 1, 1);
// Config(half_t/bfloat16_t) for M > 1024: 1x4 cluster reduces M-tail waste.
template <typename T>
struct KernelConfigLargeM {
using OutputType = T;
using MmaTileShape = Shape<_256, _256, _256>;
using ClusterShape = Shape<int, int, _1>;
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 <typename T>
const dim3 KernelConfigLargeM<T>::preferred_cluster(1, 4, 1);
template <typename T>
const dim3 KernelConfigLargeM<T>::fallback_cluster(1, 2, 1);
struct KernelConfigFp32 {
using OutputType = float;
using MmaTileShape = Shape<_128, _128, _256>;
using ClusterShape = Shape<int, int, _1>;
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 <typename KernelConfig>
struct Fp4GemmSm100 {
using Config = KernelConfig;
using OutputType = typename KernelConfig::OutputType;
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using LayoutATag = cutlass::layout::RowMajor;
static constexpr int AlignmentA = 32;
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_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<ElementD>::value;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::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<ElementD, float, void, float>>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag,
OperatorClass,
ElementA,
LayoutATag,
AlignmentA,
ElementB,
LayoutBTag,
AlignmentB,
ElementAccumulator,
MmaTileShape,
ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
sizeof(typename CollectiveEpilogue::SharedStorage))>,
MainloopSchedule>::CollectiveOp;
using GemmKernel =
cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
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>
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<int>(M);
int n = static_cast<int>(N);
int k = static_cast<int>(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<ElementA const*>(A.data_ptr()),
stride_A,
static_cast<ElementB const*>(B.data_ptr()),
stride_B,
static_cast<ElementSFA const*>(A_sf.data_ptr()),
layout_SFA,
static_cast<ElementSFB const*>(B_sf.data_ptr()),
layout_SFB},
{ // Epilogue arguments
{}, // epilogue.thread
nullptr,
stride_D,
static_cast<ElementD*>(D.data_ptr()),
stride_D}};
auto& fusion_args = arguments.epilogue.thread;
fusion_args.alpha_ptr = static_cast<ElementCompute const*>(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 <typename T>
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<T>(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 <typename OutType>
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<Fp4GemmSm100<KernelConfigM128<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (m <= 256) {
runGemm<Fp4GemmSm100<KernelConfigM256<OutType>>>(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<Fp4GemmSm100<KernelConfigDefault<OutType>>>(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<Fp4GemmSm100<KernelConfigLargeM<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
}
}
template <>
void cutlassFp4GemmDispatchSm100<float>(
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<Fp4GemmSm100<KernelConfigFp32>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
}
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm120.cuh
-228
View File
@@ -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 <typename Config, typename OutType>
struct Fp4GemmSm120 {
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using LayoutATag = cutlass::layout::RowMajor;
static constexpr int AlignmentA = 32;
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_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<ElementD>::value;
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::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<static_cast<int>(
sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
using GemmKernel =
cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
};
template <typename Gemm>
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<ElementA const*>(A.data_ptr()),
stride_A,
static_cast<ElementB const*>(B.data_ptr()),
stride_B,
static_cast<ElementSFA const*>(A_sf.data_ptr()),
layout_SFA,
static_cast<ElementSFB const*>(B_sf.data_ptr()),
layout_SFB},
{{}, nullptr, stride_D, static_cast<ElementD*>(D.data_ptr()), stride_D}};
auto& fusion_args = arguments.epilogue.thread;
fusion_args.alpha_ptr = static_cast<ElementCompute const*>(alpha.data_ptr());
return arguments;
}
template <typename Gemm>
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<Gemm>(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<uint32_t>(16), next_pow_2(m));
if (mp2 <= 32) {
runGemmSm120<Fp4GemmSm120<sm120_fp4_config_small_m, cutlass::bfloat16_t>::Gemm>(
D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (mp2 <= 256) {
runGemmSm120<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::bfloat16_t>::Gemm>(
D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else {
runGemmSm120<Fp4GemmSm120<sm120_fp4_config_default, cutlass::bfloat16_t>::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<uint32_t>(16), next_pow_2(m));
if (mp2 <= 32) {
runGemmSm120<Fp4GemmSm120<sm120_fp4_config_small_m, cutlass::half_t>::Gemm>(
D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else if (mp2 <= 256) {
runGemmSm120<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::half_t>::Gemm>(
D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
} else {
runGemmSm120<Fp4GemmSm120<sm120_fp4_config_default, cutlass::half_t>::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)
python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh
-882
View File
@@ -1,882 +0,0 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include <cutlass/arch/arch.h>
#include <cutlass/cutlass.h>
#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 <algorithm>
#include <cassert>
#include <cstdint>
#include <limits>
#include <unordered_map>
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<int>{}(key.device_id);
size_t h2 = std::hash<uintptr_t>{}(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<WorkspaceKey, WorkspaceState, WorkspaceKeyHash> cache;
WorkspaceKey key{device_id, reinterpret_cast<uintptr_t>(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<int64_t>(expert_offsets[expert_id]);
int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
// size for block in block scale.
int64_t group_size = 16;
int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
int64_t k = static_cast<int64_t>(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<int64_t>(k / 2);
int64_t group_k = static_cast<int64_t>(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<uintptr_t>(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<uintptr_t>(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<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
*layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(
cute::make_shape(static_cast<int>(m), static_cast<int>(n), static_cast<int>(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<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, LayoutSFA, LayoutSFB, ScaleConfig> \
<<<1, num_experts, 0, stream>>>( \
static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \
static_cast<C_TYPE**>(out_starts.data_ptr()), \
static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \
static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \
static_cast<float**>(alpha_starts.data_ptr()), \
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \
static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \
static_cast<C_TYPE*>(out_tensors.data_ptr()), \
static_cast<SF_TYPE*>(a_scales.data_ptr()), \
static_cast<SF_TYPE*>(b_scales.data_ptr()), \
static_cast<float*>(alphas.data_ptr()), \
static_cast<int32_t*>(expert_offsets.data_ptr()), \
static_cast<int32_t*>(sf_offsets.data_ptr()), \
static_cast<int32_t*>(problem_sizes.data_ptr()), \
K, \
N); \
}
template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
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<int>(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<bf16_t>(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<fp16_t>(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<Shape<int32_t, int32_t, int32_t>>;
using ElementType = cutlass::float_e2m1_t;
using ElementSFType = cutlass::float_ue4m3_t;
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_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<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::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<ElementD, ElementAccumulator, ElementC, ElementAccumulator>;
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<static_cast<int>(
sizeof(typename CollectiveEpilogue::SharedStorage))>,
cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop, CollectiveEpilogue>;
using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
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<int>(expert_offsets.size(0));
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
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<UnderlyingProblemShape*>(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<int, int> 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<int>::max());
// Mainloop Arguments
typename GemmKernel::MainloopArguments mainloop_args{
static_cast<const ElementType**>(a_ptrs.data_ptr()),
static_cast<StrideA*>(ab_strides.data_ptr()),
static_cast<const ElementType**>(b_ptrs.data_ptr()),
static_cast<StrideB*>(ab_strides.data_ptr()),
static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
// Epilogue Arguments
typename GemmKernel::EpilogueArguments epilogue_args{
{}, // epilogue.thread
nullptr,
static_cast<StrideC*>(c_strides.data_ptr()),
static_cast<ElementD**>(out_ptrs.data_ptr()),
static_cast<StrideC*>(c_strides.data_ptr())};
auto& fusion_args = epilogue_args.thread;
fusion_args.alpha_ptr_array = reinterpret_cast<float**>(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 <typename OutType>
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<Shape<int32_t, int32_t, int32_t>>;
using ElementType = cutlass::float_e2m1_t;
using ElementSFType = cutlass::float_ue4m3_t;
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_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<ElementC>::value;
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::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<static_cast<int>(
sizeof(typename CollectiveEpilogue::SharedStorage))>,
typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop, CollectiveEpilogue>;
using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
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<int>(expert_offsets.size(0));
run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
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<UnderlyingProblemShape*>(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<int, int> 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<int>::max());
// Mainloop Arguments
typename GemmKernel::MainloopArguments mainloop_args{
static_cast<const ElementType**>(a_ptrs.data_ptr()),
static_cast<StrideA*>(ab_strides.data_ptr()),
static_cast<const ElementType**>(b_ptrs.data_ptr()),
static_cast<StrideB*>(ab_strides.data_ptr()),
static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
// Epilogue Arguments
typename GemmKernel::EpilogueArguments epilogue_args{
{}, // epilogue.thread
nullptr,
static_cast<StrideC*>(c_strides.data_ptr()),
static_cast<ElementD**>(out_ptrs.data_ptr()),
static_cast<StrideC*>(c_strides.data_ptr())};
auto& fusion_args = epilogue_args.thread;
fusion_args.alpha_ptr_array = reinterpret_cast<float**>(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<uint8_t>(a.dtype()), "a must be uint8");
RuntimeCheck(host::is_type<uint8_t>(b.dtype()), "b must be uint8");
RuntimeCheck(host::is_type<fp8_e4m3_t>(a_blockscale.dtype()), "a_blockscale must be float8_e4m3fn");
RuntimeCheck(host::is_type<fp8_e4m3_t>(b_blockscales.dtype()), "b_blockscales must be float8_e4m3fn");
RuntimeCheck(host::is_type<float>(alphas.dtype()), "alphas must be float32");
RuntimeCheck(host::is_type<int64_t>(ab_strides.dtype()), "ab_strides must be int64");
RuntimeCheck(host::is_type<int64_t>(c_strides.dtype()), "c_strides must be int64");
RuntimeCheck(host::is_type<int32_t>(problem_sizes.dtype()), "problem_sizes must be int32");
RuntimeCheck(host::is_type<int32_t>(expert_offsets.dtype()), "expert_offsets must be int32");
RuntimeCheck(host::is_type<int32_t>(sf_offsets.dtype()), "sf_offsets must be int32");
RuntimeCheck(host::is_type<int64_t>(a_ptrs.dtype()), "a_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(b_ptrs.dtype()), "b_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(out_ptrs.dtype()), "out_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(a_scales_ptrs.dtype()), "a_scales_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(b_scales_ptrs.dtype()), "b_scales_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(alpha_ptrs.dtype()), "alpha_ptrs must be int64");
RuntimeCheck(host::is_type<int64_t>(layout_sfa.dtype()), "layout_sfa must be int64");
RuntimeCheck(host::is_type<int64_t>(layout_sfb.dtype()), "layout_sfb must be int64");
RuntimeCheck(
host::is_type<bf16_t>(output.dtype()) || host::is_type<fp16_t>(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<int>(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<int>(a.size(0));
int N = static_cast<int>(b.size(1));
int K = static_cast<int>(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<bf16_t>(output.dtype())) {
run_fp4_blockwise_scaled_group_mm_sm100<cutlass::bfloat16_t>(
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<cutlass::half_t>(
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<bf16_t>(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);
}
python/sglang/jit_kernel/nvfp4.py
-636
View File
@@ -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
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py
+3 -15
View File
@@ -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,7 +652,8 @@ 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:
if fp4_gemm is None:
raise RuntimeError("No FP4 GEMM kernel available. Install flashinfer.")
out = fp4_gemm(
x_fp4,
w.T,
@@ -662,19 +663,6 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
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."
)
out = slice_nvfp4_output(out, output_size)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
-18
View File
@@ -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"
python/sglang/multimodal_gen/runtime/platforms/cuda.py
+1 -15
View File
@@ -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
python/sglang/srt/layers/moe/cutlass_moe.py
-163
View File
@@ -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
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
-3
View File
@@ -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
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+1 -4
View File
@@ -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
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4_moe.py
+33 -26
View File
@@ -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
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)
python/sglang/srt/layers/quantization/fp4_utils.py
-4
View File
@@ -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
python/sglang/srt/layers/quantization/modelopt_quant.py
+15 -75
View File
@@ -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:
if not enable_flashinfer_fp4_gemm:
raise RuntimeError(
"NVFP4 GEMM requires flashinfer's mm_fp4; please install flashinfer."
)
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
)
else:
return cutlass_fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype)
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,10 +2417,12 @@ 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():
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(
@@ -2470,7 +2430,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
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."
)
python/sglang/srt/server_args.py
+1 -2
View File
@@ -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,
),
sgl-kernel/benchmark/bench_fp4_gemm.py
+22 -26
View File
@@ -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,
test/registered/jit/benchmark/bench_nvfp4_blockwise_moe.py
-263
View File
@@ -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)
test/registered/jit/benchmark/bench_nvfp4_quant.py
-197
View File
@@ -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)
test/registered/jit/benchmark/bench_nvfp4_scaled_mm.py
-189
View File
@@ -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)
test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py
+1 -10
View File
@@ -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,
test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py
+1 -14
View File
@@ -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,16 +303,6 @@ 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(),
test/registered/jit/test_nvfp4_blockwise_moe.py
-137
View File
@@ -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"]))
test/registered/jit/test_nvfp4_gemm.py
-152
View File
@@ -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"]))
test/registered/jit/test_nvfp4_quant.py
-225
View File
@@ -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"]))
test/registered/jit/test_silu_and_mul_scaled_fp4_experts_quant_packed.py
-337
View File
@@ -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"]))
test/registered/kernels/test_fp4_moe.py
+3 -58
View File
@@ -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)
test/registered/unit/test_server_args_migration.py
+2 -2
View File
@@ -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."""