potential bug in fp4 quantize kernel

[Qiaolin-Yu]

I found it extremely slow. In trtllm, it should be 1~2 us. Therefore, I think there's potential bug in this kernel (fp4_quantize)

[Qiaolin-Yu]

cc @Fridge003 @yzh119

[bkryu]

Hi @Qiaolin-Yu, thank you for bringing this up. The issue is closely related to #1734 where we see slowdown in quantize_with_block_size during fp4_quantize() in the cutlass MoE w.r.t. trtllm's FP4 MoE.

We are looking into the fixes to the kernel

[Qiaolin-Yu]

We are looking into the fixes to the kernel

cool

[bkryu]

Hi @Qiaolin-Yu , PR 2025 that speeds up fp4 quantization for cutlass MoE has been checked in. You should now be able to test it if you build from the main branch, or upon the next release if using releases.

Feel free to let us know whether you see a difference.

[vincentzed]

Note

Edit update: I just realized the original was 28 micros, so this result seems fine, and is comparable.

I investigated this on latest sgl, which has upgraded to flashinfer 0.5.2, which includes pr #2025

I use this patch: Check before (sgl_kernel) and after (flashinfer kernel)
https://github.com/sgl-project/sglang/compare/main...bzhng-development:sglang:vz/revert-two-prs-flashinfer-fp4?expand=1
Launch cmd to test:
python -m sglang.launch_server --model-path nvidia/Llama-3.1-8B-Instruct-FP4 --mem-fraction 0.8 --tp 2 --attention-backend flashinfer --quantization modelopt_fp4 --kv-cache-dtype auto --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 8}'

and python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 5 --random-input-len 4096 --random-output-len 512 --random-range-ratio 0.0 --max-concurrency 1 --flush-cache, the perf of sgl_kernels is only slightly better now, which is good.

BTW, we use no sf swizzle, which may explain it.
https://github.com/bzhng-development/sglang/blob/vz/revert-two-prs-flashinfer-fp4/python/sglang/srt/layers/quantization/modelopt_quant.py#L1623
On the other hand, https://github.com/bzhng-development/sglang/blob/39d18d2b3c647137252e4d79c802b201edd1a287/python/sglang/srt/layers/quantization/modelopt_quant.py#L1636 we dont use the swizzled result.

The profile of the last call of each, which is ~ comparable.
sgl-kernel

Flashinfer env:

flashinfer-cubin          0.5.2
flashinfer-jit-cache      0.5.2.dev20251112+cu129
flashinfer-python         0.5.2

Profile on 0.5.2:

So I checked the delta of average time of each:

flashinfer − sglang
= 0.0000029684889567669174 
  − 0.0000020088161125469924
= 0.0000009596728442199250

Flashinfer 50% slower (still small difference compared to large 28 micros

0.0000009596728442199250 / 0.0000020088161125469924
≈ 0.4777  (≈ 47.8% slower)

I haven't tested is_sf_swizzled_layout=True prior to the bug fix and after. #2025 does seem to fix the use case

[bkryu]

Thanks @vincentzed. To recapitulate what you are measuring to make sure I understood, it seems like you are finding that FlashInfer's fp4 quantization is slightly slower than SGLang's when is_sf_swizzled_layout=False, but not to an alarming degree.

At the moment, it seems like the end-to-end impact is minor, but please let us know if the difference is in fact significant. I can see if there are any further optimizations available