flashinfer-ai/flashinfer

BF16 hidden_states for trtllm_fp4_block_scale_moe

Open

#2,657 建立於 2026年2月28日

在 GitHub 查看
 (5 留言) (0 反應) (1 負責人)Python (1,031 fork)github user discovery
good first issueop: moe

倉庫指標

Star
 (5,756 star)
PR 合併指標
 (PR 指標待抓取)

描述

The doc string of trtllm_fp4_block_scale_moe indicates that the hidden_states can be bf16 and the hidden_states_scale is optional. But running the reproducing script hits error if passing hidden_states as bf16 and hidden_states_scale as None:

  File "/sgl-workspace/sglang/repro)fp4.py", line 303, in <module>
    main()
  File "/sgl-workspace/sglang/reprofp4.py", line 227, in main
    fi_outputs = trtllm_fp4_block_scale_moe(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/flashinfer/fused_moe/core.py", line 2552, in trtllm_fp4_block_scale_moe
    return get_trtllm_moe_sm100_module().trtllm_fp4_block_scale_moe(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/flashinfer/fused_moe/core.py", line 1866, in trtllm_fp4_block_scale_moe_op
    intermediate_output = moe_op.trtllm_fp4_block_scale_moe(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python/tvm_ffi/cython/function.pxi", line 929, in tvm_ffi.core.Function.__call__
  File "<unknown>", line 0, in __tvm_ffi_trtllm_fp4_block_scale_moe
  File "<unknown>", line 0, in flashinfer::trtllm_fp4_block_scale_moe(tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::TensorView, tvm::ffi::TensorView, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::TensorView, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::TensorView, tvm::ffi::TensorView, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::TensorView, tvm::ffi::TensorView, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::Optional<tvm::ffi::TensorView, void>, tvm::ffi::Optional<tvm::ffi::TensorView, void>, long, long, tvm::ffi::Optional<long, void>, tvm::ffi::Optional<long, void>, long, long, long, tvm::ffi::Optional<double, void>, long, bool, bool, long, tvm::ffi::TensorView, tvm::ffi::Array<long, void>)
  File "<unknown>", line 0, in flashinfer::FP4BlockScaleLauncher::run(long, bool, bool, bool)
  File "/workspace/csrc/trtllm_fused_moe_kernel_launcher.cu", line 1302, in virtual void flashinfer::FP4BlockScaleLauncher::check_moe() const
RuntimeError: Check failed: (mDtypeWeights == btg::Dtype::MxE2m1) is false: Only MxE2m1 weights are supported by block scale MoE with Bfloat16, E4m3 or MxE4m3 activation

Without passing in hidden_states_scale, errors out as:

Running FlashInfer FP4 MoE kernel with bf16 hidden states...
2026-02-28 04:54:09,050 - INFO - autotuner.py:256 - flashinfer.jit: [Autotuner]: Autotuning process starts ...
2026-02-28 04:54:09,050 - INFO - autotuner.py:262 - flashinfer.jit: [Autotuner]: Autotuning process ends
Traceback (most recent call last):
  File "/sgl-workspace/sglang/reprofp4.py", line 300, in <module>
    main()
  File "/sgl-workspace/sglang/reprofp4.py", line 224, in main
    fi_outputs = trtllm_fp4_block_scale_moe(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: trtllm_fp4_block_scale_moe() missing 1 required positional argument: 'hidden_states_scale'


repro_fp4.py @aleozlx @nekorobov

貢獻者指南