flashinfer-ai/flashinfer
Voir sur GitHubBF16 hidden_states for trtllm_fp4_block_scale_moe
Open
#2 657 ouverte le 28 févr. 2026
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