[Feature] Adapt DispathGmmCombineDecode opertor to align with weight scale dtype of small operators. [RFC: issue 5476]

[wangqiankun13]

What this PR does / why we need it?

[Feature] Adapt DispathGmmCombineDecode opertor to align with weight scale dtype of small operators.

• Before: weight scale must be float32
• After: weight scale can be float32/float16 when x is float16, float32/bfloat16 when x is float32/bfloat16. And w1 scale can use different dtype with w2 scale.

More info about this operator, please refer to RFC: issue #5476

Does this PR introduce any user-facing change?

How was this patch tested?

Perf

When scale is of type fp16 or bf16, it will be cast to fp32 internally within the operator, while the subsequent computations remain unchanged. Therefore, this PR will introduce an additional cast operation but halve the memory copy operations for scale . Furthermore, since the scale data is only a few KB in size and participates in relatively few computations, its impact is almost negligible compared to major operations like matrix multiplication. Thus, the theoretical performance change should be minimal.

test single operator cases from qwen3-235b,

• single A3 node(ep16), 64 moe experts, 4 experts / die (like qwen3-235b ep32)
• batch=18/32, token_hidden_size 4096, moe_intermediate_size 1536

The test was conducted for 100 rounds, and the average of the last 95 rounds was taken.

	bs18(us)	bs32(us)
Without this PR	96.28	108.83
With this PR	96.06	107.90

Note: Single-operator benchmarks represent an ideal scenario. They are usually only useful for referencing relative changes and may not fully align with performance data observed within the full model.

Acc

test qwen3-235b eplb on a single A3 node(ep16),
with dispatch_gmm_combine_decode

dataset	version	metric	mode	vllm-api-stream-chat
aime2024	604a78	accuracy	gen	83.33

• vLLM version: v0.13.0
• vLLM main: vllm-project/vllm@2f4e654

[gemini-code-assist]

Code Review

This pull request adapts the DispathGmmCombineDecode operator to support different data types for weight scales, enhancing flexibility. The changes are consistently applied across C++ kernel definitions, templates, and the corresponding Python code. While the overall direction is good, I've identified a critical issue in two C++ headers related to uninitialized tensor access that could lead to undefined behavior, and a high-severity issue with incorrect static assertions for buffer sizing.

[gemini-code-assist]

There is a critical issue here. The ubRawScaleList member is not initialized in the constructor when ElementRawScale is the same as ElementFp32Scale (e.g., float). Accessing it here via ubRawScaleList[ubListId] leads to using an uninitialized LocalTensor, which is undefined behavior.

To fix this, you should move the declarations of ubRawScale and layoutRawUbScale inside the if constexpr blocks where they are actually used (at lines 307 and 334).

[gemini-code-assist]

There is a critical issue here. The ubRawScaleList member is not initialized in the constructor when ElementRawScale is the same as ElementFp32Scale (e.g., float). Accessing it here via ubRawScaleList[ubListId] leads to using an uninitialized LocalTensor, which is undefined behavior.

To fix this, you should move the declarations of ubRawScale and layoutRawUbScale inside the if constexpr blocks where ubRawScale is actually used (at lines 225 and 251).

[gemini-code-assist]

The static_assert for UB size is incorrect. When ElementRawScale is the same as ElementFp32Scale (i.e., float), only ubFp32ScaleList is allocated. However, this static_assert calculates the size as if both ubRawScaleList and ubFp32ScaleList are allocated, by using sizeof(ElementRawScale) + sizeof(ElementFp32Scale). This makes the check overly restrictive and may cause compilation to fail for valid configurations.

    static_assert((UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + (std::is_same_v<ElementRawScale, ElementFp32Scale> ? 0 : TileShape::COLUMN * sizeof(ElementRawScale)) + TileShape::COLUMN * sizeof(ElementFp32Scale) +
                                TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) +
                   (TileShape::COUNT + TileShape::COUNT) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <=
                      ArchTag::UB_SIZE,
                  "TileShape is too large to fit in UB");

[gemini-code-assist]

The static_assert for UB size is incorrect. When ElementRawScale is the same as ElementFp32Scale (i.e., float), only ubFp32ScaleList is allocated. However, this static_assert calculates the size as if both ubRawScaleList and ubFp32ScaleList are allocated, by using sizeof(ElementRawScale) + sizeof(ElementFp32Scale). This makes the check overly restrictive and may cause compilation to fail for valid configurations.

    static_assert((UB_STAGES * (TileShape::COUNT * sizeof(ElementC) + (std::is_same_v<ElementRawScale, ElementFp32Scale> ? 0 : TileShape::COLUMN * sizeof(ElementRawScale)) + TileShape::COLUMN * sizeof(ElementFp32Scale) +
                                TileShape::ROW * sizeof(ElementPerTokenScale) + TileShape::COUNT * sizeof(ElementD)) +
                   (TileShape::COUNT + TileShape::COUNT) * sizeof(float) + TileShape::ROW * BYTE_PER_BLK) <=
                      ArchTag::UB_SIZE,
                  "TileShape is too large to fit in UB");

[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.