Repository: vllm-project/vllm
File: vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py (Python)

Code before:
            fused_activation=act,
        )
        intermediate_cache3 = matmul_ogs(
            intermediate_cache1,
            w2,
            quant_config.w2_bias,
            routing_data,
            scatter_indx=scatter_indx,
            precision_config=quant_config.w2_precision,
            gammas=None if apply_router_weight_on_input else gammas,
            y=output_tensor,
        )

    elif quant_config.use_mxfp4_w4a4:
        from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
        from aiter.ops.triton.quant_moe import downcast_to_static_fp8

        assert quant_config.w1_precision is not None, (
            "w1_precision in quant config can't be None"
        )
        assert quant_config.w2_precision is not None, (
            "w2_precision in quant config can't be None"
        )

        hidden_states = downcast_to_static_fp8(
            hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
        )

        intermediate_cache1 = moe_gemm_a8w4(
            hidden_states,
            w1.storage.data,
            None,
            quant_config.w1_precision.weight_scale.storage.data,
            quant_config.w1_precision.flex_ctx.lhs_data.scale,
            quant_config.w2_precision.flex_ctx.lhs_data.scale,
            quant_config.w1_bias,
            routing_data,
            gather_indx=gather_indx,
            gammas=gammas if apply_router_weight_on_input else None,
            swizzle_mx_scale="CDNA4_SCALE",
            out_dtype=torch.float8_e4m3fn,
            apply_swiglu=True,
            alpha=swiglu_alpha,
            limit=swiglu_limit,
            unpadded_N=unpadded_N_w1,
            unpadded_K=unpadded_K_w1,
        )

        intermediate_cache3 = moe_gemm_a8w4(
            intermediate_cache1,
            w2.storage.data,

Review comment (suggestion):
ah i see how input scale is passed, through within w1_precision.

I think by using `mxfp4_w4a8_moe_quant_config`, we can use `quant_config.a1_scale` here instead

Code after:
    assert quant_config.w2_precision is not None, (
        "w2_precision in quant config can't be None"
    )

    hidden_states = downcast_to_static_fp8(
        hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
    )

    intermediate_cache1 = moe_gemm_a8w4(
        hidden_states,
        w1.storage.data,
        None,
        quant_config.w1_precision.weight_scale.storage.data,
        quant_config.w1_precision.flex_ctx.lhs_data.scale,
        quant_config.w2_precision.flex_ctx.lhs_data.scale,
        quant_config.w1_bias,
        routing_data,
        gather_indx=gather_indx,
        gammas=gammas if apply_router_weight_on_input else None,
        swizzle_mx_scale="CDNA4_SCALE",
        out_dtype=torch.float8_e4m3fn,
        apply_swiglu=True,
        alpha=swiglu_alpha,
        limit=swiglu_limit,
        unpadded_N=unpadded_N_w1,
        unpadded_K=unpadded_K_w1,
    )

    intermediate_cache3 = moe_gemm_a8w4(
        intermediate_cache1,
        w2.storage.data,
        None,
        quant_config.w2_precision.weight_scale.storage.data,
        quant_config.w2_precision.flex_ctx.lhs_data.scale,
        None,
        quant_config.w2_bias,
        routing_data,
        scatter_indx=scatter_indx,
        gammas=None if apply_router_weight_on_input else gammas,
        swizzle_mx_scale="CDNA4_SCALE",
        unpadded_N=unpadded_N_w2,
        unpadded_K=unpadded_K_w2,
    )

    return intermediate_cache3


def make_routing_data(
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    num_local_experts: int,