@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f120bd42d3

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 08:45 UTC

[Kernel] Support Flashinfer trtllm fused MoE non gated FP8 & NVFP4 (#33506)

Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com>

amitz-nv committed Feb 12, 2026 at 23:06 UTC f120bd42d3daf733425d7feaaeffc2a23ba71c17

5 files changed +197 -45

tests/kernels/moe/test_flashinfer.py

+46 -12

@@ -71,7 +71,8 @@ def quant_fp8_per_tensor_batches(a):

72 for i in range(num_batches):

73 a_fp8, a_global_sf = input_to_float8(a[i])

74 - a_global_sf = 1.0 / a_global_sf

74 + if a_global_sf.numel() == 1:

75 + a_global_sf = a_global_sf.view(1, 1)

76 a_quant.append(a_fp8)

77 a_scales.append(a_global_sf)

@@ -81,6 +82,20 @@ def quant_fp8_per_tensor_batches(a):

82 return result_a_quant, result_a_scales

85 +def check_accuracy(ref_output, actual_output, atol=0.1, rtol=0.85, percent=0.925):

86 + close = torch.isclose(ref_output, actual_output, atol=atol, rtol=rtol)

87 + match_ratio = close.float().mean()

88 + assert match_ratio >= percent, (

89 + f"Match ratio {match_ratio:.4f} is below the threshold {percent:.4f}"

90 + )

91 +

92 + mismatch_percent = 1.0 - match_ratio.item()

93 + assert mismatch_percent <= 1 - percent, (

94 + f"Mismatch percentage {mismatch_percent:.4f} is above the threshold "

95 + f"{1 - percent:.4f}"

96 + )

97 +

98 +

99 @dataclass

100 class TestData:

101 hidden_states: torch.Tensor

@@ -104,14 +119,16 @@ class TestData:

119 is_gated = activation.is_gated

120

121 hidden_states = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10

107 - w13 = torch.randn(

108 - (e, (2 * n) if is_gated else n, k), device="cuda", dtype=torch.bfloat16

122 + w13 = (

123 + torch.randn(

124 + (e, (2 * n) if is_gated else n, k), device="cuda", dtype=torch.bfloat16

125 + )

126 + / 10

127 )

110 - w2 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16)

128 + w2 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16) / 10

129

130 # Scale to fp8

131 _, a1_scale = input_to_float8(hidden_states)

114 - a1_scale = 1.0 / a1_scale

132 a2_scale = torch.scalar_tensor(1.0).to(device="cuda").to(dtype=torch.float32)

133 w13_quantized, w13_weight_scale = quant_fp8_per_tensor_batches(w13)

134 w2_quantized, w2_weight_scale = quant_fp8_per_tensor_batches(w2)

@@ -124,14 +141,16 @@ class TestData:

141 layer.w2_input_scale = a2_scale

142 layer.w13_weight_scale = w13_weight_scale

143 layer.w2_weight_scale = w2_weight_scale

144 + layer.activation = activation

145 # Setup dummy config.

146 layer.moe_parallel_config = mk.FusedMoEParallelConfig.make_no_parallel()

147

148 # flashinfer expects swapped rows for w13

131 - layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)

149 + if is_gated:

150 + layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)

151 if is_trtllm:

152 rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(

134 - layer.w13_weight, layer.w2_weight

153 + layer.w13_weight, layer.w2_weight, is_gated

154 )

155 register_scales_for_trtllm_fp8_per_tensor_moe(

156 layer,

@@ -162,12 +181,14 @@ class TestData:

181 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)

182 @pytest.mark.parametrize("e", NUM_EXPERTS)

183 @pytest.mark.parametrize("topk", TOP_KS)

184 +@pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL])

185 def test_flashinfer_per_tensor_moe_fp8_no_graph(

186 m: int,

187 n: int,

188 k: int,

189 e: int,

190 topk: int,

191 + activation: MoEActivation,

192 monkeypatch,

193 ):

194 if not current_platform.has_device_capability(100):

@@ -175,7 +196,9 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

196 set_random_seed(7)

197 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

198 with set_current_vllm_config(vllm_config):

178 - td = TestData.make_moe_tensors_8bit(m, k, n, e, is_trtllm=True)

199 + td = TestData.make_moe_tensors_8bit(

200 + m, k, n, e, is_trtllm=True, activation=activation

201 + )

202

203 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)

204 topk_weights, topk_ids = Llama4MoE.custom_routing_function(

@@ -200,7 +223,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

223 topk_weights=topk_weights,

224 topk_ids=topk_ids,

225 inplace=False,

203 - activation=MoEActivation.SILU,

226 + activation=activation,

227 global_num_experts=e,

228 expert_map=None,

229 apply_router_weight_on_input=True,

@@ -219,7 +242,13 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

242 apply_router_weight_on_input=True,

243 )

244

222 - torch.testing.assert_close(output, flashinfer_output, atol=5.5e-2, rtol=1e-2)

245 + check_accuracy(

246 + ref_output=output,

247 + actual_output=flashinfer_output,

248 + atol=0.1,

249 + rtol=0.85,

250 + percent=0.925,

251 + )

252

253

254 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)

@@ -320,8 +349,13 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

349 expert_map=None,

350 apply_router_weight_on_input=True,

351 )

323 - torch.testing.assert_close(

324 - output, flashinfer_cutlass_output, atol=5.5e-2, rtol=1e-2

352 +

353 + check_accuracy(

354 + ref_output=output,

355 + actual_output=flashinfer_cutlass_output,

356 + atol=0.1,

357 + rtol=0.85,

358 + percent=0.925,

359 )

360

361

vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py

+8 -6

@@ -35,8 +35,8 @@ def _supports_current_device() -> bool:

37 def _supports_no_act_and_mul() -> bool:

38 - """Does not support non-gated MoE (i.e. Nanotron-Mini)."""

39 - return False

38 + """Supports non-gated MoE."""

39 + return True

42 def _supports_quant_scheme(

@@ -52,8 +52,7 @@ def _supports_quant_scheme(

54 def _supports_activation(activation: MoEActivation) -> bool:

55 - """Supports silu activation only."""

56 - return activation == MoEActivation.SILU

55 + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

58 def _supports_routing_method(

@@ -74,6 +73,7 @@ def _supports_routing_method(

73 elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):

74 # NOTE(dbari): as above, potentially allow others here.

75 return routing_method in [

76 + RoutingMethodType.DeepSeekV3,

77 RoutingMethodType.Llama4,

78 RoutingMethodType.Renormalize,

79 RoutingMethodType.RenormalizeNaive,

@@ -291,6 +291,7 @@ def fi_trtllm_fp8_per_tensor_moe(

291 local_num_experts: int,

292 use_routing_scales_on_input: bool,

293 routing_method_type: int,

294 + activation_type: int,

295 routed_scaling_factor: float = 1.0,

296 ) -> torch.Tensor:

297 num_expert_group = num_expert_group if num_expert_group is not None else 0

@@ -326,9 +327,9 @@ def fi_trtllm_fp8_per_tensor_moe(

327 routed_scaling_factor=routed_scaling_factor,

328 use_routing_scales_on_input=use_routing_scales_on_input,

329 routing_method_type=routing_method_type,

329 - # TODO: Required for flashinfer==0.6.3, remove with update

330 + # TODO: enum type Required for flashinfer==0.6.3, remove with update

331 # https://github.com/flashinfer-ai/flashinfer/pull/2508

331 - activation_type=ActivationType.Swiglu,

332 + activation_type=ActivationType(activation_type),

333 )

334

335

@@ -351,6 +352,7 @@ def fi_trtllm_fp8_per_tensor_moe_fake(

352 local_num_experts: int,

353 use_routing_scales_on_input: bool,

354 routing_method_type: int,

355 + activation_type: int,

356 routed_scaling_factor: float = 1.0,

357 ) -> torch.Tensor:

358 return torch.empty_like(hidden_states)

vllm/model_executor/layers/quantization/modelopt.py

+4 -3

@@ -937,10 +937,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

937 )

938 # TODO(rob): this validation should happen at kernel selection

939 # time in the oracle rather than here.

940 - assert layer.activation == MoEActivation.SILU, (

941 - f"Expected 'silu' activation but got {layer.activation}"

940 + SUPPORTED_ACTIVATIONS = [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

941 + assert layer.activation in SUPPORTED_ACTIVATIONS, (

942 + f"Only {SUPPORTED_ACTIVATIONS} activations are supported for FlashInfer "

943 + f"TRTLLM FP4 MoE, {layer.activation} found instead."

944 )

943 - assert not layer.renormalize

945 return apply_fi_trtllm_fp8_per_tensor_moe(

946 layer=layer,

947 hidden_states=x,

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

+51 -19

@@ -15,6 +15,10 @@ from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEParallelConfig,

16 RoutingMethodType,

17 )

18 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

19 + activation_to_flashinfer_int,

20 + align_fp4_moe_weights_for_fi,

21 +)

22 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

23 swizzle_blockscale,

24 )

@@ -50,8 +54,8 @@ def _supports_current_device() -> bool:

56 def _supports_no_act_and_mul() -> bool:

53 - """Does not support non-gated MoE (i.e. Nemotron-Nano)."""

54 - return False

57 + """Supports non-gated MoE."""

58 + return True

61 def _supports_quant_scheme(

@@ -66,8 +70,7 @@ def _supports_quant_scheme(

72 def _supports_activation(activation: MoEActivation) -> bool:

69 - """Supports silu activation only."""

70 - return activation in [MoEActivation.SILU]

73 + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

76 def _supports_routing_method(

@@ -150,6 +153,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(

153 hidden_size,

154 intermediate_size,

155 num_experts,

156 + is_gated_activation: bool,

157 ):

158 from flashinfer import nvfp4_block_scale_interleave

159 from flashinfer.fused_moe.core import (

@@ -160,15 +164,18 @@ def prepare_static_weights_for_trtllm_fp4_moe(

164 _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}

165 """Prepare quantized weights for kernel (done offline with weights)."""

166 epilogue_tile_m = 128 # FIXME: this depends on the kernel internals

167 + gemm1_intermediate_size = (

168 + 2 * intermediate_size if is_gated_activation else intermediate_size

169 + )

170

171 # Convert quantized weights to proper formats

172 gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape(

166 - num_experts, 2 * intermediate_size, hidden_size // 2

173 + num_experts, gemm1_intermediate_size, hidden_size // 2

174 ) # packed fp4

175 gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view(

176 torch.float8_e4m3fn

177 ).reshape(

171 - num_experts, 2 * intermediate_size, hidden_size // 16

178 + num_experts, gemm1_intermediate_size, hidden_size // 16

179 ) # fp8 scaling factors

180

181 gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape(

@@ -191,6 +198,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(

198 _cache_permute_indices,

199 gemm1_weights_fp4[i].view(torch.uint8),

200 epilogue_tile_m,

201 + is_gated_act_gemm=is_gated_activation,

202 )

203 gemm1_weights_fp4_shuffled.append(

204 gemm1_weights_fp4[i]

@@ -203,6 +211,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(

211 gemm1_scales_linear_fp4[i].view(torch.uint8),

212 epilogue_tile_m,

213 num_elts_per_sf=16,

214 + is_gated_act_gemm=is_gated_activation,

215 )

216 gemm1_scales_fp4_shuffled.append(

217 nvfp4_block_scale_interleave(

@@ -246,7 +255,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(

255 gemm1_scales_fp4_shuffled = (

256 torch.stack(gemm1_scales_fp4_shuffled)

257 .view(torch.float8_e4m3fn)

249 - .reshape(num_experts, 2 * intermediate_size, hidden_size // 16)

258 + .reshape(num_experts, gemm1_intermediate_size, hidden_size // 16)

259 )

260

261 gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled)

@@ -297,10 +306,10 @@ def flashinfer_trtllm_fp4_moe(

306

307 from vllm.model_executor.models.llama4 import Llama4MoE

308

300 - # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2404

301 - assert activation == MoEActivation.SILU, (

302 - "Only SiLU activation is supported for FlashInfer TRTLLM FP4 MoE. "

303 - f"{activation} found instead."

309 + SUPPORTED_ACTIVATIONS = [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

310 + assert activation in SUPPORTED_ACTIVATIONS, (

311 + f"Only {SUPPORTED_ACTIVATIONS} activations are supported for FlashInfer "

312 + f"TRTLLM FP4 MoE, {activation} found instead."

313 )

314

315 # Quantize input to FP4

@@ -325,6 +334,9 @@ def flashinfer_trtllm_fp4_moe(

334 else router_logits

335 )

336

337 + # Determine activation type

338 + activation_type = activation_to_flashinfer_int(layer.activation)

339 +

340 # Call TRT-LLM FP4 block-scale MoE kernel

341 out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(

342 routing_logits=router_logits,

@@ -355,6 +367,7 @@ def flashinfer_trtllm_fp4_moe(

367 routed_scaling_factor=None,

368 routing_method_type=routing_method_type,

369 do_finalize=True,

370 + activation_type=activation_type,

371 )[0]

372

373 return out

@@ -479,10 +492,16 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(

492 ]

493

494 # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels.

482 - if is_act_and_mul and backend in [

483 - NvFp4MoeBackend.FLASHINFER_CUTLASS,

484 - NvFp4MoeBackend.FLASHINFER_TRTLLM,

485 - ]:

495 + is_gated = layer.activation.is_gated

496 + if (

497 + is_gated

498 + and is_act_and_mul

499 + and backend

500 + in [

501 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

502 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

503 + ]

504 + ):

505 w13, w13_scale = reorder_w1w3_to_w3w1(w13, w13_scale)

506

507 # For some FI kernels, the input scales are shared by all experts.

@@ -495,19 +514,32 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(

514

515 # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels.

516 if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

517 + # Align weights for FI NVFP4 MoE kernels.

518 + min_alignment = 16 if is_gated else 128

519 + w13, w13_scale, w2, w2_scale, padded_intermediate = (

520 + align_fp4_moe_weights_for_fi(

521 + w13, w13_scale, w2, w2_scale, is_act_and_mul, min_alignment

522 + )

523 + )

524 + layer.intermediate_size_per_partition = padded_intermediate

525 +

526 w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe(

527 w13,

528 w2,

529 w13_scale,

530 w2_scale,

503 - w2.size(-2), # hidden_size

504 - w13.size(-2) // 2, # intermediate_size

505 - w13.size(0), # num_experts

531 + hidden_size=w2.size(-2),

532 + intermediate_size=w13.size(-2) // 2 if is_gated else w13.size(-2),

533 + num_experts=w13.size(0),

534 + is_gated_activation=is_gated,

535 )

536

537 # We do not need to make this a parameter, because

538 # it is not used during the weight (re)-loading process.

510 - layer.g1_scale_c = a13_scale * w13_scale_2 / a2_scale

539 + if is_gated:

540 + layer.g1_scale_c = a13_scale * w13_scale_2 / a2_scale

541 + else:

542 + layer.g1_scale_c = torch.ones_like(a13_scale) / a2_scale

543 layer.a1_gscale = 1.0 / a13_scale

544 layer.g1_alphas = a13_scale * w13_scale_2

545 layer.g2_alphas = a2_scale * w2_scale_2

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

+88 -5

@@ -6,6 +6,7 @@ import torch

7 from vllm import envs

8 from vllm.logger import init_logger

9 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

10 from vllm.platforms import current_platform

11 from vllm.utils.math_utils import round_up

@@ -18,6 +19,20 @@ class FlashinferMoeBackend(Enum):

19 CUTEDSL = "CUTEDSL"

22 +def activation_to_flashinfer_int(activation: MoEActivation) -> int:

23 + from flashinfer.fused_moe.core import ActivationType

24 +

25 + # silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively

26 + ACTIVATION_TO_FI_ACTIVATION = {

27 + MoEActivation.SILU_NO_MUL: ActivationType.Silu,

28 + MoEActivation.GELU_NO_MUL: ActivationType.Gelu,

29 + MoEActivation.SILU: ActivationType.Swiglu,

30 + MoEActivation.GELU: ActivationType.Geglu,

31 + MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,

32 + }

33 + return ACTIVATION_TO_FI_ACTIVATION[activation].value

34 +

35 +

36 def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:

37 return (

38 x.reshape(-1, 2, x.shape[-2] // 2, x.shape[-1]).flip(dims=[1]).reshape(x.shape)

@@ -25,7 +40,7 @@ def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:

42 def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(

28 - gemm1_weights: torch.Tensor, gemm2_weights: torch.Tensor

43 + gemm1_weights: torch.Tensor, gemm2_weights: torch.Tensor, is_gated_activation: bool

44 ):

45 """Shuffle weights for for FI TRT-LLM Format"""

46 from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a

@@ -40,6 +55,8 @@ def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(

55 for i in range(num_experts):

56 gemm1_weights_fp8_interleaved.append(

57 reorder_rows_for_gated_act_gemm(gemm1_weights[i])

58 + if is_gated_activation

59 + else gemm1_weights[i]

60 )

62 # Stack weights and scales for all experts

@@ -86,7 +103,13 @@ def register_scales_for_trtllm_fp8_per_tensor_moe(

103 )

104 layer.w2_input_scale_inv = 1.0 / w2_input_scale

105 layer.output1_scales_gate_scalar = g1_alphas

89 - layer.output1_scales_scalar = g1_alphas * layer.w2_input_scale_inv

106 +

107 + if layer.activation.is_gated:

108 + layer.output1_scales_scalar = g1_alphas * layer.w2_input_scale_inv

109 + else:

110 + layer.output1_scales_scalar = (

111 + torch.ones_like(g1_alphas) * layer.w2_input_scale_inv

112 + )

113 layer.output2_scales_scalar = g2_alphas

114

115

@@ -125,6 +148,7 @@ def apply_fi_trtllm_fp8_per_tensor_moe(

148 assert layer.custom_routing_function is None, (

149 "Custom routing function is only supported for Llama4"

150 )

151 + activation_type = activation_to_flashinfer_int(layer.activation)

152

153 return torch.ops.vllm.fi_trtllm_fp8_per_tensor_moe(

154 routing_logits=router_logits,

@@ -145,6 +169,7 @@ def apply_fi_trtllm_fp8_per_tensor_moe(

169 local_num_experts=layer.local_num_experts,

170 use_routing_scales_on_input=apply_router_weight_on_input,

171 routing_method_type=layer.routing_method_type,

172 + activation_type=activation_type,

173 )

174

175

@@ -274,8 +299,64 @@ def convert_moe_weights_to_flashinfer_trtllm_block_layout(

299 return w13_weights_shuffled_tensor, w2_weights_shuffled_tensor

300

301

302 +def align_fp4_moe_weights_for_fi(

303 + w13: torch.Tensor,

304 + w13_scale: torch.Tensor,

305 + w2: torch.Tensor,

306 + w2_scale: torch.Tensor,

307 + is_act_and_mul: bool,

308 + min_alignment: int = 16,

309 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:

310 + """Pad intermediate size so FlashInfer kernels' alignment constraints hold.

311 +

312 + Some FlashInfer FP4 MoE kernels require the intermediate size

313 + used for GEMM to be divisible by a small alignment value. When this is

314 + not satisfied (e.g. with certain tensor-parallel sizes), we pad the

315 + gate/up and down projection weights along the intermediate dim.

316 + """

317 +

318 + # Current local intermediate size (per partition) is the K dimension of

319 + # the down projection.

320 + num_experts, hidden_size, intermediate = w2.shape

321 + intermediate *= 2 # because of packed FP4

322 +

323 + padded_intermediate = round_up(intermediate, min_alignment)

324 +

325 + if padded_intermediate == intermediate:

326 + return w13, w13_scale, w2, w2_scale, intermediate

327 +

328 + logger.info_once(

329 + "Padding intermediate size from %d to %d for up/down projection weights.",

330 + intermediate,

331 + padded_intermediate,

332 + scope="local",

333 + )

334 +

335 + up_mult = 2 if is_act_and_mul else 1

336 + padded_gate_up_dim = up_mult * padded_intermediate

337 +

338 + # Pad w13 and w2 along its intermediate dimension.

339 + padded_w13 = w13.new_zeros((num_experts, padded_gate_up_dim, hidden_size // 2))

340 + padded_w13[:, : w13.shape[1], :] = w13

341 +

342 + padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate // 2))

343 + padded_w2[:, :, : w2.shape[2]] = w2

344 +

345 + padded_w13_scale = w13_scale.new_zeros(

346 + (num_experts, padded_gate_up_dim, hidden_size // 16)

347 + )

348 + padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale

349 +

350 + padded_w2_scale = w2_scale.new_zeros(

351 + (num_experts, hidden_size, padded_intermediate // 16)

352 + )

353 + padded_w2_scale[:, :, : w2_scale.shape[2]] = w2_scale

354 +

355 + return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate

356 +

357 +

358 def align_fp8_moe_weights_for_fi(

278 - w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool

359 + w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool, min_alignment: int = 16

360 ) -> tuple[torch.Tensor, torch.Tensor, int]:

361 """Pad intermediate size so FlashInfer kernels' alignment constraints hold.

362

@@ -289,7 +370,6 @@ def align_fp8_moe_weights_for_fi(

370 # the down projection.

371 num_experts, hidden_size, intermediate = w2.shape

372

292 - min_alignment = 16

373 padded_intermediate = round_up(intermediate, min_alignment)

374

375 if padded_intermediate == intermediate:

@@ -342,11 +422,14 @@ def prepare_fp8_moe_layer_for_fi(

422

423 # Some FI MoE kernels require internal alignment of 16

424 # for the gate-up proj. Pad the weights to respect this.

425 + is_gated = layer.activation.is_gated

426 if not block_quant:

427 + min_alignment = 16 if is_gated else 128

428 w13, w2, new_intermediate = align_fp8_moe_weights_for_fi(

429 w13,

430 w2,

431 layer.moe_config.is_act_and_mul,

432 + min_alignment,

433 )

434 layer.intermediate_size_per_partition = new_intermediate

435

@@ -363,7 +446,7 @@ def prepare_fp8_moe_layer_for_fi(

446 assert w13_input_scale is not None

447 assert w2_input_scale is not None

448

366 - rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(w13, w2)

449 + rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(w13, w2, is_gated)

450 register_scales_for_trtllm_fp8_per_tensor_moe(

451 layer,

452 w13_scale=w13_scale,