minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support Gemma4 MoE NVFP4 (#25054)

Browse Source
This commit is contained in:
Charles Chen
2026-05-20 22:45:15 -07:00
committed by GitHub
parent 888a8794ef
commit 847cbada9c
9 changed files with 271 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/cutlass_moe.py
+1 -1
View File
@@ -468,7 +468,7 @@ def cutlass_moe_fp4(
)
del rep_a_fp4, rep_a_blockscale
# hidden size dimension is split to one halfpytho sized tensor.
# hidden size dimension is split to one half sized tensor.
intermediate = torch.empty(
(m_a * num_topk, w1_fp4.shape[1] // 2), device=device, dtype=out_dtype
)
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+21 -9
View File
@@ -509,18 +509,26 @@ def align_fp4_moe_weights_for_flashinfer_trtllm(layer: Module) -> None:
layer.intermediate_size_per_partition = intermediate_size
def get_activation_type(activation: str) -> int:
def get_activation_type(activation: str, is_gated: bool = True) -> int:
"""Map SGLang activation string to FlashInfer ActivationType int value."""
from flashinfer.fused_moe.core import ActivationType
_ACTIVATION_STR_TO_TYPE = {
"silu": ActivationType.Swiglu,
"relu2": ActivationType.Relu2,
}
if is_gated:
_ACTIVATION_STR_TO_TYPE = {
"silu": ActivationType.Swiglu,
"gelu": ActivationType.Geglu,
}
else:
_ACTIVATION_STR_TO_TYPE = {
"silu": ActivationType.Silu,
"gelu": ActivationType.Gelu,
"relu2": ActivationType.Relu2,
}
act = _ACTIVATION_STR_TO_TYPE.get(activation)
if act is None:
raise ValueError(
f"Unsupported activation '{activation}' for TRTLLM MoE. "
f"Unsupported activation '{activation}' for TRTLLM MoE "
f"(is_gated={is_gated}). "
f"Expected one of {list(_ACTIVATION_STR_TO_TYPE.keys())}."
)
return act.value
@@ -863,7 +871,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.moe.utils import RoutingMethodType
_SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2"}
_SUPPORTED_FP4_ACTIVATIONS = {"silu", "relu2", "gelu"}
assert runner_config.activation in _SUPPORTED_FP4_ACTIVATIONS, (
f"Only {_SUPPORTED_FP4_ACTIVATIONS} are supported for FP4 MoE, "
f"got '{runner_config.activation}'."
@@ -896,7 +904,9 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
hs_scale = hs_scale_linear.view(torch.float8_e4m3fn).reshape(
*hs_scale_linear.shape[:-1], -1
)
activation_type = get_activation_type(runner_config.activation)
activation_type = get_activation_type(
runner_config.activation, is_gated=runner_config.is_gated
)
num_tokens = hs_fp4.shape[0]
hidden_size = (
@@ -1070,7 +1080,9 @@ def fused_experts_none_to_flashinfer_trtllm_bf16(
assert (
runner_config.num_fused_shared_experts == 0
), "Fused shared experts are not supported for flashinfer trtllm moe"
activation_type = get_activation_type(runner_config.activation)
activation_type = get_activation_type(
runner_config.activation, is_gated=runner_config.is_gated
)
hidden_states = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py
+4 -1
View File
@@ -400,7 +400,10 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme):
get_activation_type,
)
activation_type = get_activation_type(moe_runner_config.activation)
activation_type = get_activation_type(
moe_runner_config.activation,
is_gated=moe_runner_config.is_gated,
)
quant_info = FlashInferTrtllmFp8MoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
python/sglang/srt/layers/quantization/fp8.py
+4 -1
View File
@@ -1930,7 +1930,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
get_activation_type,
)
activation_type = get_activation_type(self.moe_runner_config.activation)
activation_type = get_activation_type(
self.moe_runner_config.activation,
is_gated=self.moe_runner_config.is_gated,
)
quant_info = FlashInferTrtllmFp8MoeQuantInfo(
w13_weight=layer.w13_weight,
python/sglang/srt/layers/quantization/modelopt_quant.py
+57 -17
View File
@@ -102,7 +102,9 @@ except ImportError:
# Define a minimal ActivationType enum if flashinfer is not available
class ActivationType(IntEnum):
Swiglu = 3
Geglu = 4
Relu2 = 6
Identity = 7
# Initialize logger for the module
@@ -263,10 +265,8 @@ MOE_NVFP4_DISPATCH = envs.SGLANG_MOE_NVFP4_DISPATCH.get()
# Supported activation schemes for the current configuration
ACTIVATION_SCHEMES = ["static"]
ACT_STR_TO_TYPE_MAP = {
"silu": ActivationType.Swiglu, # This is the default
"relu2": ActivationType.Relu2,
}
_SUPPORTED_ACT_STRS = ("silu", "relu2", "gelu")
class ModelOptQuantConfig(QuantizationConfig):
@@ -1028,7 +1028,10 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
output1_scales_gate_scalar=layer.output1_scales_gate_scalar,
output2_scales_scalar=layer.output2_scales_scalar,
use_routing_scales_on_input=True,
activation_type=get_activation_type(self.moe_runner_config.activation),
activation_type=get_activation_type(
self.moe_runner_config.activation,
is_gated=self.moe_runner_config.is_gated,
),
)
return fused_experts_none_to_flashinfer_trtllm_fp8(
@@ -1036,15 +1039,33 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
)
if get_moe_runner_backend().is_flashinfer_cutlass():
activation = ACT_STR_TO_TYPE_MAP[self.moe_runner_config.activation]
assert (
(
activation is ActivationType.Relu2
and not self.moe_runner_config.is_gated
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
get_activation_type,
)
activation_str = self.moe_runner_config.activation
assert activation_str in _SUPPORTED_ACT_STRS, (
f"Activation {activation_str!r} is not supported for "
f"flashinfer cutlass fp8 moe (supported: {_SUPPORTED_ACT_STRS})."
)
activation = ActivationType(
get_activation_type(
activation_str, is_gated=self.moe_runner_config.is_gated
)
or activation is ActivationType.Swiglu
and self.moe_runner_config.is_gated
), "Only Relu2 non-gated or Swiglu gated are supported for flashinfer cutlass fp8 moe"
)
# FlashInfer CUTLASS MoE supports gated Swiglu/Geglu and non-gated
# Relu2/Identity. Non-gated Silu/Gelu are not implemented.
_CUTLASS_SUPPORTED = {
ActivationType.Swiglu,
ActivationType.Geglu,
ActivationType.Relu2,
ActivationType.Identity,
}
assert activation in _CUTLASS_SUPPORTED, (
f"Activation {activation_str!r} (is_gated="
f"{self.moe_runner_config.is_gated}) maps to {activation.name}, "
"which is not supported by flashinfer cutlass fp8 moe."
)
topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
x_fp8, _ = scaled_fp8_quant(x, layer.w13_input_scale)
output_dtype = x.dtype
@@ -1310,7 +1331,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
layer.output_size_per_partition = output_size_per_partition
if input_size_per_partition % 16 != 0:
raise ValueError(
"Unsupported model when in features size is " "not multiple of 16"
"Unsupported model when in features size is not multiple of 16"
)
weight_dtype = (
@@ -2021,8 +2042,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
activation = self.moe_runner_config.activation
assert (
activation in ACT_STR_TO_TYPE_MAP
), f"{activation=} missing from {ACT_STR_TO_TYPE_MAP.keys()=}"
activation in _SUPPORTED_ACT_STRS
), f"{activation=} not in supported {_SUPPORTED_ACT_STRS}"
moe_runner_config = self.moe_runner_config
# FlashInfer TRTLLM FP4 path
@@ -2102,11 +2123,30 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
return self.runner.run(dispatch_output, quant_info)
if self.enable_flashinfer_cutlass_moe:
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
get_activation_type,
)
from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker
assert (
not moe_runner_config.apply_router_weight_on_input
), "apply_router_weight_on_input is not supported for Flashinfer"
# Resolve the FlashInfer ActivationType honoring the gated flag,
# then verify the CUTLASS FP4 kernel supports it.
fi_activation = ActivationType(
get_activation_type(activation, is_gated=moe_runner_config.is_gated)
)
_CUTLASS_FP4_SUPPORTED = {
ActivationType.Swiglu,
ActivationType.Geglu,
ActivationType.Relu2,
ActivationType.Identity,
}
assert fi_activation in _CUTLASS_FP4_SUPPORTED, (
f"Activation {activation!r} (is_gated={moe_runner_config.is_gated}) "
f"maps to {fi_activation.name}, which is not supported by the "
"flashinfer cutlass fp4 moe kernel."
)
# TRTLLM Cutlass moe takes in activations in BF16/Half/nvfp4 precision
# and fp4 quantized weights loaded from the checkpoint
x = dispatch_output.hidden_states
@@ -2157,7 +2197,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
tp_size=layer.moe_tp_size,
tp_rank=layer.moe_tp_rank,
tune_max_num_tokens=next_power_of_2(x.shape[0]),
activation_type=ACT_STR_TO_TYPE_MAP[activation],
activation_type=fi_activation,
enable_alltoall=get_moe_a2a_backend().is_flashinfer(),
)[0]
python/sglang/srt/managers/scheduler.py
+8 -1
View File
@@ -826,7 +826,14 @@ class Scheduler(
self.model_config.hf_config, "text_config", self.model_config.hf_config
)
if hasattr(config_to_check, "num_experts_per_tok"):
# Different MoE architectures expose the per-token expert count under
# different attribute names (e.g. Gemma4 uses ``top_k_experts``).
moe_topk_attrs = (
"num_experts_per_tok",
"num_experts_per_token",
"top_k_experts",
)
if any(hasattr(config_to_check, attr) for attr in moe_topk_attrs):
initialize_moe_config(self.server_args)
# Initialize GEMM-related configuration for FP8 and FP4 backends.
python/sglang/srt/models/gemma4_causal.py
+81 -27
View File
@@ -42,6 +42,7 @@ from sglang.srt.layers.linear import (
)
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
@@ -1140,7 +1141,7 @@ class Gemma4ForCausalLM(PreTrainedModel):
("gate_up_proj", "up_proj", 1),
]
expert_params_mapping = [
fused_expert_params_mapping = [
# (param_name, ckpt_weight_name, shard_ids)
# gate_up_proj is fused [E, 2*I, H] — chunk into w1 (gate) + w3 (up)
("experts.w13_weight", "experts.gate_up_proj", ("w1", "w3")),
@@ -1148,6 +1149,23 @@ class Gemma4ForCausalLM(PreTrainedModel):
]
num_experts = self.config.num_experts
# Per-expert checkpoint format used by compressed-tensors / FP8
# (e.g. RedHatAI/*-FP8-Dynamic) and by ModelOpt NVFP4
# (e.g. nvidia/Gemma-4-*-NVFP4). Each expert is stored as a
# separate key with shape (out, in):
# experts.<id>.{gate,up,down}_proj.{weight,weight_scale,
# weight_scale_2,input_scale}
# `make_expert_params_mapping` emits tuples whose `weight_name` ends
# in a trailing dot, so the standard `name.replace(weight_name,
# param_name)` collapses every suffix uniformly to the fused
# FusedMoE params (experts.w13_*, experts.w2_*).
per_expert_params_mapping = FusedMoE.make_expert_params_mapping(
ckpt_gate_proj_name="gate_proj",
ckpt_down_proj_name="down_proj",
ckpt_up_proj_name="up_proj",
num_experts=num_experts,
)
k_eq_v_layers = self._get_k_eq_v_layers()
params_dict = dict(self.named_parameters())
@@ -1201,22 +1219,41 @@ class Gemma4ForCausalLM(PreTrainedModel):
# MoE expert weights checked first (gate_up_proj contains "up_proj"
# which would false-match the stacked dense MLP mapping).
orig_name = name
for param_name, weight_name, shard_ids in expert_params_mapping:
name = orig_name
if weight_name not in name:
# 1) Per-expert checkpoint layout (compressed-tensors FP8 like
# RedHatAI/*-FP8-Dynamic, ModelOpt NVFP4 like
# nvidia/Gemma-4-*-NVFP4): experts.<id>.{gate,up,down}_proj.*
# The trailing dot in `weight_name` lets a single mapping fold
# weight, weight_scale, weight_scale_2, and input_scale into
# their corresponding fused FusedMoE params (experts.w13_*,
# experts.w2_*).
for (
param_name,
weight_name,
expert_id,
shard_id,
) in per_expert_params_mapping:
if weight_name not in orig_name:
continue
name = name.replace(weight_name, param_name)
name = orig_name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
for i in range(num_experts):
chunks = loaded_weight[i].chunk(len(shard_ids), dim=0)
for chunk, sid in zip(chunks, shard_ids):
weight_loader(param, chunk, name, sid, i)
weight_loader(
param,
loaded_weight,
name,
shard_id=shard_id,
expert_id=expert_id,
)
loaded_params.add(name)
break
else:
for param_name, weight_name, shard_id in stacked_params_mapping:
# 2) BF16 fused checkpoint layout: experts.gate_up_proj is a
# [E, 2*I, H] tensor that needs per-expert chunking into
# w1 (gate) and w3 (up).
for param_name, weight_name, shard_ids in fused_expert_params_mapping:
name = orig_name
if weight_name not in name:
continue
@@ -1225,25 +1262,42 @@ class Gemma4ForCausalLM(PreTrainedModel):
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
if should_dup_k_to_v:
weight_loader(param, loaded_weight, "v")
for i in range(num_experts):
chunks = loaded_weight[i].chunk(len(shard_ids), dim=0)
for chunk, sid in zip(chunks, shard_ids):
weight_loader(param, chunk, name, sid, i)
loaded_params.add(name)
break
else:
name = orig_name
if name.endswith(".bias") and name not in params_dict:
continue
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
for param_name, weight_name, shard_id in stacked_params_mapping:
name = orig_name
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
if should_dup_k_to_v:
weight_loader(param, loaded_weight, "v")
loaded_params.add(name)
break
else:
name = orig_name
if name.endswith(".bias") and name not in params_dict:
continue
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
unloaded_params = params_dict.keys() - loaded_params
if unloaded_params:
param_names = set(dict(self.named_parameters()).keys())
python/sglang/srt/models/gemma4_mm.py
+88 -61
View File
@@ -33,6 +33,7 @@ from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
from sglang.srt.layers.layernorm import Gemma4RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
@@ -817,6 +818,27 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
("experts.w2_weight", "experts.down_proj", ("w2",)),
]
# Per-expert checkpoint format used by compressed-tensors / FP8
# (e.g. RedHatAI/*-FP8-Dynamic) and by ModelOpt NVFP4
# (e.g. nvidia/Gemma-4-*-NVFP4). Each expert is stored as a
# separate key with shape (out, in):
# experts.<id>.{gate,up,down}_proj.{weight,weight_scale,
# weight_scale_2,input_scale}
# `make_expert_params_mapping` emits tuples whose `weight_name` ends
# in a trailing dot, so the standard `name.replace(weight_name,
# param_name)` collapses every suffix uniformly to the fused
# FusedMoE params (experts.w13_*, experts.w2_*).
per_expert_params_mapping = (
FusedMoE.make_expert_params_mapping(
ckpt_gate_proj_name="gate_proj",
ckpt_down_proj_name="down_proj",
ckpt_up_proj_name="up_proj",
num_experts=num_experts,
)
if num_experts
else []
)
params_dict = dict(self.named_parameters())
params_dict.update(dict(self.named_buffers()))
non_persistent_buffers: Set[str] = set()
@@ -890,60 +912,44 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
and int(m.group(1)) in k_eq_v_layers
)
# Per-expert checkpoint format used by compressed-tensors / FP8
# (e.g. RedHatAI/*-FP8-Dynamic). Each expert is stored as a
# separate key with shape (out, in):
# experts.<id>.gate_proj.{weight,weight_scale}
# experts.<id>.up_proj.{weight,weight_scale}
# experts.<id>.down_proj.{weight,weight_scale}
# These need to be folded into sglang's fused FusedMoE params:
# experts.w13_weight[_scale] (gate->shard "w1", up->shard "w3")
# experts.w2_weight[_scale] (down->shard "w2")
per_expert_match = re.match(
r"^(.*?\.moe\.experts\.)(\d+)\.(gate_proj|up_proj|down_proj)"
r"\.(weight|weight_scale)$",
name,
)
if per_expert_match:
prefix = per_expert_match.group(1)
expert_id = int(per_expert_match.group(2))
proj = per_expert_match.group(3)
suffix = per_expert_match.group(4)
if proj == "gate_proj":
base, sid = "w13_weight", "w1"
elif proj == "up_proj":
base, sid = "w13_weight", "w3"
else: # down_proj
base, sid = "w2_weight", "w2"
if suffix == "weight_scale":
base += "_scale"
fused_name = prefix + base
if fused_name in params_dict:
param = params_dict[fused_name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, fused_name, sid, expert_id)
loaded_params.add(fused_name)
continue
# MoE expert weights checked first (gate_up_proj contains "up_proj"
# which would false-match the stacked dense MLP mapping).
orig_name = name
for param_name, weight_name, shard_ids in expert_params_mapping:
name = orig_name
if weight_name not in name:
# 1) Per-expert checkpoint layout (compressed-tensors FP8 like
# RedHatAI/*-FP8-Dynamic, ModelOpt NVFP4 like
# nvidia/Gemma-4-*-NVFP4): experts.<id>.{gate,up,down}_proj.*
# The trailing dot in `weight_name` lets a single mapping fold
# weight, weight_scale, weight_scale_2, and input_scale into
# their corresponding fused FusedMoE params (experts.w13_*,
# experts.w2_*).
for (
param_name,
weight_name,
expert_id,
shard_id,
) in per_expert_params_mapping:
if weight_name not in orig_name:
continue
name = name.replace(weight_name, param_name)
name = orig_name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
for i in range(num_experts):
chunks = loaded_weight[i].chunk(len(shard_ids), dim=0)
for chunk, sid in zip(chunks, shard_ids):
weight_loader(param, chunk, name, sid, i)
weight_loader(
param,
loaded_weight,
name,
shard_id=shard_id,
expert_id=expert_id,
)
loaded_params.add(name)
break
else:
for param_name, weight_name, shard_id in self.stacked_params_mapping:
# 2) BF16 fused checkpoint layout: experts.gate_up_proj is a
# [E, 2*I, H] tensor that needs per-expert chunking into
# w1 (gate) and w3 (up).
for param_name, weight_name, shard_ids in expert_params_mapping:
name = orig_name
if weight_name not in name:
continue
@@ -952,25 +958,46 @@ class Gemma4ForConditionalGeneration(PreTrainedModel):
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
if should_dup_k_to_v:
weight_loader(param, loaded_weight, "v")
for i in range(num_experts):
chunks = loaded_weight[i].chunk(len(shard_ids), dim=0)
for chunk, sid in zip(chunks, shard_ids):
weight_loader(param, chunk, name, sid, i)
loaded_params.add(name)
break
else:
name = orig_name
if name.endswith(".bias") and name not in params_dict:
continue
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
for (
param_name,
weight_name,
shard_id,
) in self.stacked_params_mapping:
name = orig_name
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
if should_dup_k_to_v:
weight_loader(param, loaded_weight, "v")
loaded_params.add(name)
break
else:
name = orig_name
if name.endswith(".bias") and name not in params_dict:
continue
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add(name)
unloaded_params = params_dict.keys() - loaded_params
if unloaded_params:
param_names = set(dict(self.named_parameters()).keys())
python/sglang/srt/server_args.py
+7
View File
@@ -2230,6 +2230,13 @@ class ServerArgs:
"Gemma4 only supports trtllm_mha or triton attention backend, "
f"got prefill={prefill_backend}, decode={decode_backend}"
)
if is_sm100_supported() and self.moe_runner_backend == "auto":
self.moe_runner_backend = "flashinfer_trtllm"
logger.info(
"Use flashinfer_trtllm as MoE runner backend on SM100 for Gemma-4 NVFP4"
)
elif model_arch == "MossVLForConditionalGeneration":
if self.is_attention_backend_not_set():
self.prefill_attention_backend = "flashinfer"