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

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 14, 05:30 UTC

[oss] Init gpt-oss bf16 support (#22508)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Aug 10, 2025 at 11:19 UTC 0c5254b82acc625112ce7adc10811514f1a42d52

4 files changed +342 -125

vllm/model_executor/layers/fused_moe/config.py

+5 -1

@@ -324,6 +324,8 @@ class FusedMoEConfig:

324

325 max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE

326

327 + has_bias: bool = False

328 +

329 def __post_init__(self):

330 if self.dp_size > 1:

331 logger.debug_once("Using FusedMoEConfig::max_num_tokens=%d",

@@ -413,7 +415,8 @@ class FusedMoEConfig:

415 in_dtype: torch.dtype,

416 max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE,

417 quant_config: Optional[Union[FusedMoEQuantConfig,

416 - QuantizationConfig]] = None

418 + QuantizationConfig]] = None,

419 + has_bias: bool = False,

420 ) -> "FusedMoEConfig":

421

422 _quant_config: Optional[FusedMoEQuantConfig] = None

@@ -482,4 +485,5 @@ class FusedMoEConfig:

485 in_dtype=in_dtype,

486 quant_config=_quant_config,

487 max_num_tokens=max_num_tokens,

488 + has_bias=has_bias,

489 )

vllm/model_executor/layers/fused_moe/fused_moe.py

+160 -109

@@ -275,6 +275,7 @@ def fused_moe_kernel(

275 a_ptr,

276 b_ptr,

277 c_ptr,

278 + b_bias_ptr,

279 a_scale_ptr,

280 b_scale_ptr,

281 topk_weights_ptr,

@@ -302,6 +303,8 @@ def fused_moe_kernel(

303 stride_bse,

304 stride_bsk,

305 stride_bsn,

306 + stride_bbe, # bias expert stride

307 + stride_bbn, # bias N stride

308 # Block size for block-wise quantization

309 group_n: tl.constexpr,

310 group_k: tl.constexpr,

@@ -317,6 +320,7 @@ def fused_moe_kernel(

320 use_int8_w8a8: tl.constexpr,

321 use_int8_w8a16: tl.constexpr,

322 per_channel_quant: tl.constexpr,

323 + HAS_BIAS: tl.constexpr,

324 ):

325 """

326 Implements the fused computation for a Mixture of Experts (MOE) using

@@ -414,7 +418,10 @@ def fused_moe_kernel(

418 else:

419 a_scale = tl.load(a_scale_ptr)

420 b_scale = tl.load(b_scale_ptr + off_experts)

417 -

421 + if HAS_BIAS:

422 + # bias shape: [num_experts, N]

423 + bias_ptrs = b_bias_ptr + off_experts * stride_bbe + offs_bn * stride_bbn

424 + bias = tl.load(bias_ptrs, mask=(offs_bn < N), other=0.0)

425 # -----------------------------------------------------------

426 # Iterate to compute a block of the C matrix.

427 # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block

@@ -456,7 +463,8 @@ def fused_moe_kernel(

463 # Advance the ptrs to the next K block.

464 a_ptrs += BLOCK_SIZE_K * stride_ak

465 b_ptrs += BLOCK_SIZE_K * stride_bk

459 -

466 + if HAS_BIAS:

467 + accumulator = accumulator + bias[None, :]

468 if MUL_ROUTED_WEIGHT:

469 moe_weight = tl.load(topk_weights_ptr + offs_token,

470 mask=token_mask,

@@ -471,6 +479,7 @@ def fused_moe_kernel(

479 accumulator = (accumulator * a_scale * b_scale).to(compute_type)

480 else:

481 accumulator = accumulator.to(compute_type)

482 +

483 # -----------------------------------------------------------

484 # Write back the block of the output

485 offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

@@ -499,7 +508,8 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

508 use_int8_w8a16: bool,

509 use_int4_w4a16: bool,

510 per_channel_quant: bool,

502 - block_shape: Optional[list[int]] = None) -> None:

511 + block_shape: Optional[list[int]] = None,

512 + B_bias: Optional[torch.Tensor] = None) -> None:

513 assert topk_weights is not None or not mul_routed_weight

514 assert topk_weights is None or topk_weights.stride(1) == 1

515 assert sorted_token_ids.stride(0) == 1

@@ -531,7 +541,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

541 A.size(0) * top_k * config['BLOCK_SIZE_M'])

542 grid = lambda META: (triton.cdiv(EM, META['BLOCK_SIZE_M']) * triton.cdiv(

543 B.size(1), META['BLOCK_SIZE_N']), )

534 -

544 + HAS_BIAS = B_bias is not None

545 if (use_int8_w8a16 or use_int4_w4a16) and \

546 block_shape is not None and block_shape[1] > 0:

547 assert B_scale is not None and B_scale.ndim == 3

@@ -611,6 +621,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

621 A,

622 B,

623 C,

624 + B_bias,

625 A_scale,

626 B_scale,

627 topk_weights,

@@ -638,6 +649,8 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

649 if B_scale is not None and B_scale.ndim == 3 else 0,

650 B_scale.stride(1)

651 if B_scale is not None and B_scale.ndim >= 2 else 0,

652 + B_bias.stride(0) if B_bias is not None else 0,

653 + B_bias.stride(1) if B_bias is not None else 0,

654 0 if block_shape is None else block_shape[0],

655 0 if block_shape is None else block_shape[1],

656 MUL_ROUTED_WEIGHT=mul_routed_weight,

@@ -647,6 +660,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,

660 use_int8_w8a8=use_int8_w8a8,

661 use_int8_w8a16=use_int8_w8a16,

662 per_channel_quant=per_channel_quant,

663 + HAS_BIAS=HAS_BIAS,

664 BLOCK_SIZE_K=BLOCK_SIZE_K,

665 **config,

666 )

@@ -1024,40 +1038,43 @@ def inplace_fused_experts(

1038 w2_zp: Optional[torch.Tensor] = None,

1039 a1_scale: Optional[torch.Tensor] = None,

1040 a2_scale: Optional[torch.Tensor] = None,

1027 - block_shape: Optional[List[int]] = None) -> None: #noqa: UP006

1041 + block_shape: Optional[List[int]] = None,

1042 + w1_bias: Optional[torch.Tensor] = None,

1043 + w2_bias: Optional[torch.Tensor] = None) -> None: #noqa: UP006

1044 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,

1045 activation, is_act_and_mul,

1046 apply_router_weight_on_input, use_fp8_w8a8,

1047 use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,

1048 use_mxfp4_w4a4, per_channel_quant, global_num_experts,

1049 expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,

1034 - a2_scale, block_shape)

1035 -

1036 -

1037 -def inplace_fused_experts_fake(

1038 - hidden_states: torch.Tensor,

1039 - w1: torch.Tensor,

1040 - w2: torch.Tensor,

1041 - topk_weights: torch.Tensor,

1042 - topk_ids: torch.Tensor,

1043 - activation: str = "silu",

1044 - is_act_and_mul: bool = True,

1045 - apply_router_weight_on_input: bool = False,

1046 - use_fp8_w8a8: bool = False,

1047 - use_int8_w8a8: bool = False,

1048 - use_int8_w8a16: bool = False,

1049 - use_int4_w4a16: bool = False,

1050 - use_mxfp4_w4a4: bool = False,

1051 - per_channel_quant: bool = False,

1052 - global_num_experts: int = -1,

1053 - expert_map: Optional[torch.Tensor] = None,

1054 - w1_scale: Optional[torch.Tensor] = None,

1055 - w2_scale: Optional[torch.Tensor] = None,

1056 - w1_zp: Optional[torch.Tensor] = None,

1057 - w2_zp: Optional[torch.Tensor] = None,

1058 - a1_scale: Optional[torch.Tensor] = None,

1059 - a2_scale: Optional[torch.Tensor] = None,

1060 - block_shape: Optional[list[int]] = None) -> None:

1050 + a2_scale, block_shape, w1_bias, w2_bias)

1051 +

1052 +

1053 +def inplace_fused_experts_fake(hidden_states: torch.Tensor,

1054 + w1: torch.Tensor,

1055 + w2: torch.Tensor,

1056 + topk_weights: torch.Tensor,

1057 + topk_ids: torch.Tensor,

1058 + activation: str = "silu",

1059 + is_act_and_mul: bool = True,

1060 + apply_router_weight_on_input: bool = False,

1061 + use_fp8_w8a8: bool = False,

1062 + use_int8_w8a8: bool = False,

1063 + use_int8_w8a16: bool = False,

1064 + use_int4_w4a16: bool = False,

1065 + use_mxfp4_w4a4: bool = False,

1066 + per_channel_quant: bool = False,

1067 + global_num_experts: int = -1,

1068 + expert_map: Optional[torch.Tensor] = None,

1069 + w1_scale: Optional[torch.Tensor] = None,

1070 + w2_scale: Optional[torch.Tensor] = None,

1071 + w1_zp: Optional[torch.Tensor] = None,

1072 + w2_zp: Optional[torch.Tensor] = None,

1073 + a1_scale: Optional[torch.Tensor] = None,

1074 + a2_scale: Optional[torch.Tensor] = None,

1075 + block_shape: Optional[list[int]] = None,

1076 + w1_bias: Optional[torch.Tensor] = None,

1077 + w2_bias: Optional[torch.Tensor] = None) -> None:

1078 pass

1079

1080

@@ -1246,36 +1263,38 @@ direct_register_custom_op(

1263

1264

1265 def outplace_fused_experts(

1249 - hidden_states: torch.Tensor,

1250 - w1: torch.Tensor,

1251 - w2: torch.Tensor,

1252 - topk_weights: torch.Tensor,

1253 - topk_ids: torch.Tensor,

1254 - activation: str = "silu",

1255 - is_act_and_mul: bool = True,

1256 - apply_router_weight_on_input: bool = False,

1257 - use_fp8_w8a8: bool = False,

1258 - use_int8_w8a8: bool = False,

1259 - use_int8_w8a16: bool = False,

1260 - use_int4_w4a16: bool = False,

1261 - use_mxfp4_w4a4: bool = False,

1262 - per_channel_quant: bool = False,

1263 - global_num_experts: int = -1,

1264 - expert_map: Optional[torch.Tensor] = None,

1265 - w1_scale: Optional[torch.Tensor] = None,

1266 - w2_scale: Optional[torch.Tensor] = None,

1267 - w1_zp: Optional[torch.Tensor] = None,

1268 - w2_zp: Optional[torch.Tensor] = None,

1269 - a1_scale: Optional[torch.Tensor] = None,

1270 - a2_scale: Optional[torch.Tensor] = None,

1271 - block_shape: Optional[List[int]] = None, #noqa: UP006

1266 + hidden_states: torch.Tensor,

1267 + w1: torch.Tensor,

1268 + w2: torch.Tensor,

1269 + topk_weights: torch.Tensor,

1270 + topk_ids: torch.Tensor,

1271 + activation: str = "silu",

1272 + is_act_and_mul: bool = True,

1273 + apply_router_weight_on_input: bool = False,

1274 + use_fp8_w8a8: bool = False,

1275 + use_int8_w8a8: bool = False,

1276 + use_int8_w8a16: bool = False,

1277 + use_int4_w4a16: bool = False,

1278 + use_mxfp4_w4a4: bool = False,

1279 + per_channel_quant: bool = False,

1280 + global_num_experts: int = -1,

1281 + expert_map: Optional[torch.Tensor] = None,

1282 + w1_scale: Optional[torch.Tensor] = None,

1283 + w2_scale: Optional[torch.Tensor] = None,

1284 + w1_zp: Optional[torch.Tensor] = None,

1285 + w2_zp: Optional[torch.Tensor] = None,

1286 + a1_scale: Optional[torch.Tensor] = None,

1287 + a2_scale: Optional[torch.Tensor] = None,

1288 + block_shape: Optional[List[int]] = None, #noqa: UP006

1289 + w1_bias: Optional[torch.Tensor] = None,

1290 + w2_bias: Optional[torch.Tensor] = None,

1291 ) -> torch.Tensor:

1292 return fused_experts_impl(

1293 hidden_states, w1, w2, topk_weights, topk_ids, False, activation,

1294 is_act_and_mul, apply_router_weight_on_input, use_fp8_w8a8,

1295 use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, use_mxfp4_w4a4,

1296 per_channel_quant, global_num_experts, expert_map, w1_scale, w2_scale,

1278 - w1_zp, w2_zp, a1_scale, a2_scale, block_shape)

1297 + w1_zp, w2_zp, a1_scale, a2_scale, block_shape, w1_bias, w2_bias)

1298

1299

1300 def outplace_fused_experts_fake(

@@ -1300,7 +1319,9 @@ def outplace_fused_experts_fake(

1319 w2_zp: Optional[torch.Tensor] = None,

1320 a1_scale: Optional[torch.Tensor] = None,

1321 a2_scale: Optional[torch.Tensor] = None,

1303 - block_shape: Optional[list[int]] = None) -> torch.Tensor:

1322 + block_shape: Optional[list[int]] = None,

1323 + w1_bias: Optional[torch.Tensor] = None,

1324 + w2_bias: Optional[torch.Tensor] = None) -> torch.Tensor:

1325 return torch.empty_like(hidden_states)

1326

1327

@@ -1332,33 +1353,34 @@ def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:

1353

1354 # TODO (bnell): replace this with modular op. Can get rid of inplace/outplace

1355 # torch ops.

1335 -def fused_experts(

1336 - hidden_states: torch.Tensor,

1337 - w1: torch.Tensor,

1338 - w2: torch.Tensor,

1339 - topk_weights: torch.Tensor,

1340 - topk_ids: torch.Tensor,

1341 - inplace: bool = False,

1342 - activation: str = "silu",

1343 - is_act_and_mul: bool = True,

1344 - apply_router_weight_on_input: bool = False,

1345 - use_fp8_w8a8: bool = False,

1346 - use_int8_w8a8: bool = False,

1347 - use_int8_w8a16: bool = False,

1348 - use_int4_w4a16: bool = False,

1349 - use_mxfp4_w4a4: bool = False,

1350 - per_channel_quant: bool = False,

1351 - global_num_experts: int = -1,

1352 - expert_map: Optional[torch.Tensor] = None,

1353 - w1_scale: Optional[torch.Tensor] = None,

1354 - w2_scale: Optional[torch.Tensor] = None,

1355 - w1_zp: Optional[torch.Tensor] = None,

1356 - w2_zp: Optional[torch.Tensor] = None,

1357 - a1_scale: Optional[torch.Tensor] = None,

1358 - a2_scale: Optional[torch.Tensor] = None,

1359 - block_shape: Optional[list[int]] = None,

1360 - allow_deep_gemm: bool = False,

1361 - allow_cutlass_block_scaled_grouped_gemm: bool = False) -> torch.Tensor:

1356 +def fused_experts(hidden_states: torch.Tensor,

1357 + w1: torch.Tensor,

1358 + w2: torch.Tensor,

1359 + topk_weights: torch.Tensor,

1360 + topk_ids: torch.Tensor,

1361 + inplace: bool = False,

1362 + activation: str = "silu",

1363 + is_act_and_mul: bool = True,

1364 + apply_router_weight_on_input: bool = False,

1365 + use_fp8_w8a8: bool = False,

1366 + use_int8_w8a8: bool = False,

1367 + use_int8_w8a16: bool = False,

1368 + use_int4_w4a16: bool = False,

1369 + use_mxfp4_w4a4: bool = False,

1370 + per_channel_quant: bool = False,

1371 + global_num_experts: int = -1,

1372 + expert_map: Optional[torch.Tensor] = None,

1373 + w1_scale: Optional[torch.Tensor] = None,

1374 + w2_scale: Optional[torch.Tensor] = None,

1375 + w1_zp: Optional[torch.Tensor] = None,

1376 + w2_zp: Optional[torch.Tensor] = None,

1377 + a1_scale: Optional[torch.Tensor] = None,

1378 + a2_scale: Optional[torch.Tensor] = None,

1379 + block_shape: Optional[list[int]] = None,

1380 + allow_deep_gemm: bool = False,

1381 + allow_cutlass_block_scaled_grouped_gemm: bool = False,

1382 + w1_bias: Optional[torch.Tensor] = None,

1383 + w2_bias: Optional[torch.Tensor] = None) -> torch.Tensor:

1384 # For now, disable DeepGemm for small N (<= 512) until better

1385 # permute/unpermute ops are available.

1386 # However, on B200, we use DeepGemm for all cases because they only support

@@ -1423,7 +1445,10 @@ def fused_experts(

1445 w2_zp=w2_zp,

1446 a1_scale=a1_scale,

1447 a2_scale=a2_scale,

1426 - block_shape=block_shape)

1448 + block_shape=block_shape,

1449 + w1_bias=w1_bias,

1450 + w2_bias=w2_bias,

1451 + )

1452

1453

1454 def fused_experts_impl(

@@ -1451,6 +1476,8 @@ def fused_experts_impl(

1476 a1_scale: Optional[torch.Tensor] = None,

1477 a2_scale: Optional[torch.Tensor] = None,

1478 block_shape: Optional[list[int]] = None,

1479 + w1_bias: Optional[torch.Tensor] = None,

1480 + w2_bias: Optional[torch.Tensor] = None,

1481 ) -> torch.Tensor:

1482 # Check constraints.

1483 if use_int4_w4a16:

@@ -1591,7 +1618,19 @@ def fused_experts_impl(

1618 use_int8_w8a16=use_int8_w8a16,

1619 use_int4_w4a16=use_int4_w4a16,

1620 per_channel_quant=per_channel_quant,

1594 - block_shape=block_shape)

1621 + block_shape=block_shape,

1622 + B_bias=w1_bias)

1623 +

1624 + # TODO fused kernel

1625 + def swiglu_oai(gate_up):

1626 + alpha = 1.702

1627 + limit = 7.0

1628 + gate, up = gate_up[..., ::2], gate_up[..., 1::2]

1629 + gate = gate.clamp(min=None, max=limit)

1630 + up = up.clamp(min=-limit, max=limit)

1631 + glu = gate * torch.sigmoid(gate * alpha)

1632 + gated_output = (up + 1) * glu

1633 + return gated_output

1634

1635 # Activation function with multiplication

1636 if activation == "silu" and is_act_and_mul:

@@ -1605,6 +1644,8 @@ def fused_experts_impl(

1644 intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))

1645 elif activation == "gelu":

1646 intermediate_cache2 = F.gelu(intermediate_cache1.view(-1, N))

1647 + elif activation == "swiglu_oai":

1648 + intermediate_cache2 = swiglu_oai(intermediate_cache1.view(-1, N))

1649 else:

1650 raise ValueError(f"Unsupported FusedMoe activation: {activation}, "

1651 f"with is_act_and_mul={is_act_and_mul}.")

@@ -1635,7 +1676,8 @@ def fused_experts_impl(

1676 use_int8_w8a16=use_int8_w8a16,

1677 use_int4_w4a16=use_int4_w4a16,

1678 per_channel_quant=per_channel_quant,

1638 - block_shape=block_shape)

1679 + block_shape=block_shape,

1680 + B_bias=w2_bias)

1681

1682 ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.size()),

1683 out_hidden_states[begin_chunk_idx:end_chunk_idx])

@@ -1672,6 +1714,8 @@ def fused_moe(

1714 a1_scale: Optional[torch.Tensor] = None,

1715 a2_scale: Optional[torch.Tensor] = None,

1716 block_shape: Optional[list[int]] = None,

1717 + w1_bias: Optional[torch.Tensor] = None,

1718 + w2_bias: Optional[torch.Tensor] = None,

1719 ) -> torch.Tensor:

1720 """

1721 This function computes a Mixture of Experts (MoE) layer using two sets of

@@ -1766,7 +1810,9 @@ def fused_moe(

1810 w2_zp=w2_zp,

1811 a1_scale=a1_scale,

1812 a2_scale=a2_scale,

1769 - block_shape=block_shape)

1813 + block_shape=block_shape,

1814 + w1_bias=w1_bias,

1815 + w2_bias=w2_bias)

1816

1817

1818 class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

@@ -1937,7 +1983,9 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1983 use_int8_w8a16=self.use_int8_w8a16,

1984 use_int4_w4a16=self.use_int4_w4a16,

1985 per_channel_quant=self.per_act_token_quant,

1940 - block_shape=self.block_shape)

1986 + block_shape=self.block_shape,

1987 + B_bias=None # TODO support B_bias

1988 + )

1989

1990 self.activation(activation, intermediate_cache2,

1991 intermediate_cache1.view(-1, N))

@@ -1948,26 +1996,29 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1996 intermediate_cache2, a2_scale, self.quant_dtype,

1997 self.per_act_token_quant, self.block_shape)

1998

1951 - invoke_fused_moe_kernel(qintermediate_cache2,

1952 - w2,

1953 - intermediate_cache3,

1954 - a2q_scale,

1955 - w2_scale,

1956 - w2_zp,

1957 - topk_weights,

1958 - sorted_token_ids,

1959 - expert_ids,

1960 - num_tokens_post_padded,

1961 - not apply_router_weight_on_input,

1962 - 1,

1963 - config,

1964 - compute_type=compute_type,

1965 - use_fp8_w8a8=self.use_fp8_w8a8,

1966 - use_int8_w8a8=self.use_int8_w8a8,

1967 - use_int8_w8a16=self.use_int8_w8a16,

1968 - use_int4_w4a16=self.use_int4_w4a16,

1969 - per_channel_quant=self.per_act_token_quant,

1970 - block_shape=self.block_shape)

1999 + invoke_fused_moe_kernel(

2000 + qintermediate_cache2,

2001 + w2,

2002 + intermediate_cache3,

2003 + a2q_scale,

2004 + w2_scale,

2005 + w2_zp,

2006 + topk_weights,

2007 + sorted_token_ids,

2008 + expert_ids,

2009 + num_tokens_post_padded,

2010 + not apply_router_weight_on_input,

2011 + 1,

2012 + config,

2013 + compute_type=compute_type,

2014 + use_fp8_w8a8=self.use_fp8_w8a8,

2015 + use_int8_w8a8=self.use_int8_w8a8,

2016 + use_int8_w8a16=self.use_int8_w8a16,

2017 + use_int4_w4a16=self.use_int4_w4a16,

2018 + per_channel_quant=self.per_act_token_quant,

2019 + block_shape=self.block_shape,

2020 + B_bias=None # TODO support B_bias

2021 + )

2022

2023 ops.moe_sum(intermediate_cache3, output)

2024

vllm/model_executor/layers/fused_moe/layer.py

+28 -12

@@ -255,7 +255,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

255 self.fused_experts = fused_experts # type: ignore

256 self.topk_indices_dtype = None

257 self.moe = moe

258 -

258 + self.has_bias = self.moe.has_bias

259 self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()

260 if self.rocm_aiter_moe_enabled:

261 from .rocm_aiter_fused_moe import rocm_aiter_fused_experts

@@ -291,7 +291,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

291 requires_grad=False)

292 layer.register_parameter("w13_weight", w13_weight)

293 set_weight_attrs(w13_weight, extra_weight_attrs)

294 -

294 + if self.has_bias:

295 + w13_bias = torch.nn.Parameter(torch.zeros(

296 + num_experts,

297 + 2 * intermediate_size_per_partition,

298 + dtype=params_dtype),

299 + requires_grad=False)

300 + layer.register_parameter("w13_bias", w13_bias)

301 + set_weight_attrs(w13_bias, extra_weight_attrs)

302 # down_proj (row parallel)

303 w2_weight = torch.nn.Parameter(torch.empty(

304 num_experts,

@@ -301,6 +308,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

308 requires_grad=False)

309 layer.register_parameter("w2_weight", w2_weight)

310 set_weight_attrs(w2_weight, extra_weight_attrs)

311 + if self.has_bias:

312 + w2_bias = torch.nn.Parameter(torch.zeros(num_experts,

313 + hidden_size,

314 + dtype=params_dtype),

315 + requires_grad=False)

316 + layer.register_parameter("w2_bias", w2_bias)

317 + set_weight_attrs(w2_bias, extra_weight_attrs)

318

319 def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:

320 # Pad the weight tensor. This is an optimization on ROCm platform, which

@@ -465,6 +479,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

479 hidden_states=x,

480 w1=layer.w13_weight,

481 w2=layer.w2_weight,

482 + w1_bias=layer.w13_bias if self.has_bias else None,

483 + w2_bias=layer.w2_bias if self.has_bias else None,

484 topk_weights=topk_weights,

485 topk_ids=topk_ids,

486 inplace=True,

@@ -702,6 +718,7 @@ class FusedMoE(torch.nn.Module):

718 activation: str = "silu",

719 enable_eplb: bool = False,

720 num_redundant_experts: int = 0,

721 + has_bias: bool = False,

722 ):

723 super().__init__()

724 if params_dtype is None:

@@ -793,16 +810,15 @@ class FusedMoE(torch.nn.Module):

810 # since model_config is not set in the pytest test.

811 model_dtype = params_dtype

812

796 - moe = FusedMoEConfig.make(

797 - num_experts=self.global_num_experts,

798 - experts_per_token=top_k,

799 - hidden_dim=hidden_size,

800 - num_local_experts=self.local_num_experts,

801 - moe_parallel_config=self.moe_parallel_config,

802 - in_dtype=model_dtype,

803 - max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,

804 - quant_config=quant_config,

805 - )

813 + moe = FusedMoEConfig.make(num_experts=self.global_num_experts,

814 + experts_per_token=top_k,

815 + hidden_dim=hidden_size,

816 + num_local_experts=self.local_num_experts,

817 + moe_parallel_config=self.moe_parallel_config,

818 + in_dtype=model_dtype,

819 + max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,

820 + quant_config=quant_config,

821 + has_bias=has_bias)

822 self.moe_config = moe

823 self.quant_config = quant_config

824

vllm/model_executor/models/gpt_oss.py

+149 -3

@@ -160,7 +160,9 @@ class MLPBlock(torch.nn.Module):

160 renormalize=True,

161 quant_config=quant_config,

162 prefix=f"{prefix}.experts",

163 - apply_router_weight_on_input=False)

163 + apply_router_weight_on_input=False,

164 + has_bias=True,

165 + activation="swiglu_oai")

166

167 def forward(self, x: torch.Tensor) -> torch.Tensor:

168 t = self.norm(x)

@@ -262,8 +264,8 @@ class GptOssForCausalLM(nn.Module):

264 sampling_metadata)

265 return logits

266

265 - def load_weights(self, weights: Iterable[tuple[str,

266 - torch.Tensor]]) -> set[str]:

267 + def _load_weights_mxfp4(

268 + self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

269 rename_mapping = {

270 "self_attn": "attn",

271 "input_layernorm.weight": "attn.norm.weight",

@@ -469,3 +471,147 @@ class GptOssForCausalLM(nn.Module):

471 loaded_params.add(renamed_name)

472

473 return loaded_params

474 +

475 + def _load_weights_other(

476 + self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

477 + rename_mapping = {

478 + "self_attn": "attn",

479 + "input_layernorm.weight": "attn.norm.weight",

480 + "post_attention_layernorm.weight": "mlp.norm.weight",

481 + "embed_tokens": "embedding",

482 + }

483 +

484 + def maybe_rename(name: str) -> str:

485 + for remap_name, new_name in rename_mapping.items():

486 + if remap_name in name:

487 + return name.replace(remap_name, new_name)

488 + return name

489 +

490 + params_dict = dict(self.named_parameters())

491 + loaded_params: set[str] = set()

492 +

493 + tp_rank = get_tensor_model_parallel_rank()

494 + tp_size = get_tensor_model_parallel_world_size()

495 + intermediate_size = self.model_config.intermediate_size

496 +

497 + per_rank_intermediate_size = cdiv(intermediate_size, tp_size)

498 + # Calculate common slicing bounds for current rank

499 + tp_rank_start = tp_rank * per_rank_intermediate_size

500 + tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size,

501 + intermediate_size)

502 +

503 + # Attention heads per rank

504 + heads_per_rank = self.model_config.num_attention_heads // tp_size

505 + head_start = tp_rank * heads_per_rank

506 +

507 + use_ep = self.vllm_config.parallel_config.enable_expert_parallel

508 + ep_size = get_ep_group().world_size

509 + ep_rank = get_ep_group().rank

510 + num_experts = self.model_config.num_local_experts

511 + experts_per_rank = num_experts // ep_size

512 + ep_rank_start = ep_rank * experts_per_rank

513 + ep_rank_end = (ep_rank + 1) * experts_per_rank

514 +

515 + for name, weight in weights:

516 + if ".experts.gate_up_proj" in name and "bias" not in name:

517 + # Handle MLP gate and up projection weights

518 + new_name = name.replace(".experts.gate_up_proj",

519 + ".experts.w13_weight")

520 +

521 + # Extract gate and up projection parts

522 + # since the weight is shuffled, we can slice directly

523 + if use_ep:

524 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

525 + else:

526 + narrow_weight = weight[:, :,

527 + 2 * tp_rank_start:2 * tp_rank_end]

528 +

529 + narrow_weight = narrow_weight.permute(0, 2, 1).contiguous()

530 + param = params_dict[new_name]

531 +

532 + param.copy_(narrow_weight)

533 + loaded_params.add(new_name)

534 +

535 + elif ".experts.down_proj" in name and "bias" not in name:

536 + # Handle MLP down projection weights

537 + new_name = name.replace(".experts.down_proj",

538 + ".experts.w2_weight")

539 +

540 + if use_ep:

541 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

542 + else:

543 + narrow_weight = weight[:, tp_rank_start:tp_rank_end, :]

544 + narrow_weight = narrow_weight.permute(0, 2, 1).contiguous()

545 + param = params_dict[new_name]

546 +

547 + param.copy_(narrow_weight)

548 + loaded_params.add(new_name)

549 +

550 + elif "gate_up_proj_bias" in name:

551 + # Handle MLP gate and up projection biases

552 + new_name = name.replace("gate_up_proj_bias", "w13_bias")

553 +

554 + # Extract gate and up projection bias parts

555 + if use_ep:

556 + narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

557 + else:

558 + narrow_weight = weight[:,

559 + 2 * tp_rank_start:2 * tp_rank_end]

560 +

561 + param = params_dict[new_name]

562 +

563 + param.copy_(narrow_weight)

564 + loaded_params.add(new_name)

565 +

566 + elif "down_proj_bias" in name:

567 + # Handle MLP down projection bias

568 + new_name = name.replace("down_proj_bias", "w2_bias")

569 +

570 + if use_ep:

571 + weight = weight[ep_rank_start:ep_rank_end, ...]

572 + else:

573 + # (only load on rank 0 to avoid duplication)

574 + if tp_rank != 0:

575 + weight.zero_()

576 + param = params_dict[new_name]

577 + param.copy_(weight)

578 + loaded_params.add(new_name)

579 + elif "sinks" in name:

580 + # Handle attention sinks (distributed across ranks)

581 + name = name.replace("self_attn", "attn")

582 + param = params_dict[name]

583 + narrow_weight = weight.narrow(0, head_start, heads_per_rank)

584 + param.data.copy_(narrow_weight)

585 + loaded_params.add(name)

586 + elif "q_proj" in name or "k_proj" in name or "v_proj" in name:

587 + shard_id = ("q" if "q_proj" in name else

588 + "k" if "k_proj" in name else "v")

589 + name = name.replace("self_attn", "attn")

590 + param_name = name.replace(f"{shard_id}_proj", "qkv")

591 + param = params_dict[param_name]

592 + weight_loader = param.weight_loader

593 + weight_loader(param, weight, loaded_shard_id=shard_id)

594 + loaded_params.add(param_name)

595 + else:

596 + # Handle all other weights with potential renaming

597 +

598 + renamed_name = maybe_rename(name)

599 + if renamed_name not in params_dict:

600 + continue

601 + param = params_dict[renamed_name]

602 + weight_loader = getattr(param, "weight_loader",

603 + default_weight_loader)

604 + weight_loader(param, weight)

605 + loaded_params.add(renamed_name)

606 +

607 + return loaded_params

608 +

609 + def load_weights(self, weights: Iterable[tuple[str,

610 + torch.Tensor]]) -> set[str]:

611 + quant_method = (self.model_config.quantization_config['quant_method']

612 + if hasattr(self.model_config, "quantization_config")

613 + else None)

614 + if quant_method == "mxfp4":

615 + return self._load_weights_mxfp4(weights)

616 + else:

617 + return self._load_weights_other(weights)