[OP] add custom op aclnnMoeInitRoutingCustom

[jiazhengyi]

What this PR does / why we need it?

This pull request introduces a new custom operator aclnnMoeInitRoutingCustom for Mixture-of-Experts models.
It can be replaced by aclnnMoeInitRoutingV3 once CANN 8.5 becomes available.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

[gemini-code-assist]

Code Review

This pull request introduces a new custom operator moe_init_routing_custom for Mixture-of-Experts models. The changes are extensive, covering the operator's definition, implementation for both host and device, shape inference, tiling logic, and Torch bindings. I've identified a few critical issues related to attribute validation and shape inference that could lead to runtime errors or incorrect behavior. Additionally, there's a bug in the golden implementation within the test file that could affect test accuracy. My review provides suggestions to address these points to improve the robustness and correctness of the new operator.

[gemini-code-assist]

The validation for active_num is overly strict. The default value for this attribute is -1, which is a standard way to indicate a dynamic dimension. The current check activeNum < 0 incorrectly flags this default value as an error when dropPadMode is NO_DROP_PAD, preventing the operator from being used with dynamic shapes. Please adjust the check to permit the value -1, for instance, by changing it to activeNum < -1.

        if (activeNum < -1) {
            OPS_LOG_E(context->GetNodeName(), "The active_num should be greater than or equal to -1. But it is %ld.", activeNum);
            return ge::GRAPH_FAILED;
        }

[gemini-code-assist]

The shape for the expert_tokens_count_or_cumsum output is only set when expertTokenNumFlag is true. As this is a required output, its shape must be defined in all scenarios to prevent graph compilation failures. When the flag is false, you should set a valid shape, such as [0], to represent an empty tensor.

    if(expertTokenNumFlag){
        if (expertTokenNumType == ExpertTokenNumType::KEY_VALUE) {
            expertTokenCumsumOrCountShape->SetDimNum(DIM_TWO);
            expertTokenCumsumOrCountShape->SetDim(0U, experNum);
            expertTokenCumsumOrCountShape->SetDim(DIM_ONE, KEY_VALUE_MODE_DIM0_NUM);
        } else {
            expertTokenCumsumOrCountShape->SetDimNum(DIM_ONE);
            expertTokenCumsumOrCountShape->SetDim(0U, expertEnd - expertStart);
        }
    } else {
        expertTokenCumsumOrCountShape->SetDimNum(DIM_ONE);
        expertTokenCumsumOrCountShape->SetDim(0U, 0);
    }

[gemini-code-assist]

The expert_tokens_num_type parameter is not being validated. An invalid value could result in the expert_tokens_count_or_cumsum tensor not being initialized, leading to a runtime crash. It's important to add a TORCH_CHECK to confirm that expert_tokens_num_type is within the valid range of {0, 1, 2}, mirroring the validation in infershape.cpp.

    TORCH_CHECK(expert_tokens_num_type >= CUMSUM && expert_tokens_num_type <= KEY_VALUE, "expert_tokens_num_type must be one of {0, 1, 2}, but got ", expert_tokens_num_type);
    if (expert_tokens_num_type >= CUMSUM && expert_tokens_num_type <= COUNT) {
        // expert_tokens_count_or_cumsum in [end-start, ]
        expert_tokens_count_or_cumsum = at::empty({expert_length}, x.options().dtype(at::kLong));
    } else if (expert_tokens_num_type == KEY_VALUE) {
        // key_value in [2, end-start]
        expert_tokens_count_or_cumsum = at::empty({expert_num, 2}, x.options().dtype(at::kLong));
    }

[gemini-code-assist]

Similar to the main binding, this meta function is missing validation for expert_tokens_num_type. An invalid value will cause an uninitialized tensor to be returned, which can lead to issues during model tracing or exporting. Please add a TORCH_CHECK to ensure expert_tokens_num_type is valid.

    TORCH_CHECK(expert_tokens_num_type >= CUMSUM && expert_tokens_num_type <= KEY_VALUE, "expert_tokens_num_type must be one of {0, 1, 2}, but got ", expert_tokens_num_type);
    if (expert_tokens_num_type >= CUMSUM && expert_tokens_num_type <= COUNT) {
        // expert_tokens_count_or_cumsum in [end-start, ]
        expert_tokens_count_or_cumsum = at::empty({expert_length}, x.options().dtype(at::kLong));
    } else if (expert_tokens_num_type == KEY_VALUE) {
        // key_value in [2, end-start]
        expert_tokens_count_or_cumsum = at::empty({expert_num, 2}, x.options().dtype(at::kLong));
    }

[gemini-code-assist]

The expression expanded_scale == None is a comparison and does not perform an assignment. This appears to be a bug in the golden implementation, as the likely intent was to assign None to expanded_scale (expanded_scale = None). This could lead to incorrect test validations if expanded_scale retains a value from a previous execution path.

Suggested change

	expanded_scale == None
	expanded_scale = None

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.