[Perf] Add Flashinfer DeepGEMM SM90 for SwapAB Optimization

[b8zhong]

Motivation

After flashinfer-ai/flashinfer#2131 in Flashinfer, we can benefit from SwapAB, where the input order is swapped to benefit when the M dimension is < 32 (e.g when BS < 32 in decoding). When it is larger, there is no benefit.

Modifications

(Requires Flashinfer nightly, and the backend currently only supports SM90)
Note that Flashinfer will compile it's own DeepGEMM. So it is separate from the DeepGEMM built in the Docker container.

Accuracy Tests

Benchmarking and Profiling

for ((N=1; N<=128; N*=2)); do
  python3 -m sglang.bench_serving \
    --backend sglang \
    --flush-cache \
    --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 1024 \
    --random-range-ratio 1.0 \
    --num-prompts $((6*N)) \
    --max-concurrency $N \
    --output-file res.jsonl
done

We can see that when the M dimension is small, there is around a 5-8% E2E benefit

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Fridge003]

@b8zhong Will the warmup process be handled by flashinfer for this case?
We know that original deepgemm kernels need a lot of time for warm up...

[b8zhong]

@Fridge003 I think, it uses the same DeepGEMM compiler under the hood. E.g during warmup you can see this process and a few similar ones. Although, I don't absolutely have the most context, so it may or may not be fully correct

root       41757  0.0  0.0   2808  1112 pts/1    S    00:55   0:00 \
/bin/sh -c \
/usr/local/cuda/bin/nvcc \
  --generate-dependencies-with-compile \
  --dependency-output fp8_blockscale_gemm_90/fp8_blockscale_gemm.cuda.o.d \
  -DPy_LIMITED_API=0x03090000 \
  -D_GLIBCXX_USE_CXX11_ABI=1 \
  -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/csrc/nv_internal \
  -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/csrc/nv_internal/include \
  -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/csrc/nv_internal/tensorrt_llm/cutlass_extensions/include \
  -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include \
  -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels \
  -isystem /usr/include/python3.12 \
  -isystem /usr/local/cuda/include \
  -isystem /usr/local/cuda/include/cccl \
  -isystem /usr/local/lib/python3.12/dist-packages/tvm_ffi/include \
  -isystem /usr/local/lib/python3.12/dist-packages/flashinfer/data/include \
  -isystem /usr/local/lib/python3.12/dist-packages/flashinfer/data/csrc \
  -isystem /usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/include \
  -isystem /usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/tools/util/include \
  -isystem /usr/local/lib/python3.12/dist-packages/flashinfer/data/spdlog/include \
  --compiler-options=-fPIC \
  --expt-relaxed-constexpr \
  -static-global-template-stub=false \
  -std=c++17 \
  --threads=1 \
  -use_fast_math \
  -DNDEBUG \
  -O3 \
  -gencode=arch=compute_90a,code=sm_90a \
  -DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED \
  -DCOMPILE_HOPPER_TMA_GEMMS \
  -DENABLE_BF16 \
  -DENABLE_FP8 \
  -DENABLE_FP8_BLOCK_SCALE \
  -DFLASHINFER_ENABLE_F16 \
  -DFLASHINFER_ENABLE_BF16 \
  -DFLASHINFER_ENABLE_FP8_E4M3 \
  -DFLASHINFER_ENABLE_FP8_E5M2 \
  -DFLASHINFER_ENABLE_FP8_E8M0 \
  -DFLASHINFER_ENABLE_FP4_E2M1 \
  -c /usr/local/lib/python3.12/dist-packages/flashinfer/data/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu \
  -o fp8_blockscale_gemm_90/fp8_blockscale_gemm.cuda.o

This is the regular --fp8-gemm-backend=deep_gemm

root       62105  0.0  0.0   2804  1072 pts/1    S+   01:08   0:00 sh -c -- /usr/local/cuda/bin/nvcc /root/.cache/deep_gemm/cache/kernel.sm90_fp8_gemm_1d2d.6f90a6ff4eb7dbd1eb1a6966f918905c/kernel.cu -o /root/.cache/deep_gemm/tmp/57456-d236c320-2b64367a-55587891 -std=c++20 --diag-suppress=39,161,174,177,186,940 --ptxas-options=--register-usage-level=10 -I/usr/local/lib/python3.12/dist-packages/deep_gemm/include --gpu-architecture=sm_90a --compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi -cubin -O3 --expt-relaxed-constexpr --expt-extended-lambda 2>&1

[b8zhong]

@Fridge003 Although, it does seem somewhat faster actually (maybe 2x faster?)

[b8zhong]

Added accuracy numbers too~

[b8zhong]

/tag-and-rerun-ci

[Fridge003]

Can we add a test for this new fp8 gemm kernel

[b8zhong]

Done @Fridge003

[Fridge003]

/rerun-failed-ci