flashinfer-ai/flashinfer

BF16 hidden_states for trtllm_fp4_block_scale_moe

Open

#2 657 ouverte le 28 févr. 2026

Voir sur GitHub
 (5 commentaires) (0 réactions) (1 assigné)Python (1 031 forks)github user discovery
good first issueop: moe

Métriques du dépôt

Stars
 (5 756 stars)
Métriques de merge PR
 (Métriques PR en attente)

Description

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

Guide contributeur