[DeepEP] increase MAX_K_VALUE_A2 from 10 to 16 to support topk > 10#593
[DeepEP] increase MAX_K_VALUE_A2 from 10 to 16 to support topk > 10#593Yuxin1999 wants to merge 1 commit into
Conversation
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Code Review
This pull request increases the maximum K value for A2 hardware (MAX_K_VALUE_A2) from 10 to 16. However, the K_MAX constant (used during A3 tiling validation) remains set to 10, which will cause tiling failures on A3 hardware when using configurations with a top-k value greater than 10. It is recommended to update K_MAX to 16 as well to ensure consistent support across both A2 and A3 hardware.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| constexpr uint32_t RANK_NUM_PER_NODE_A2 = 8; | ||
| constexpr uint32_t BLOCK_SIZE_A2 = 32; | ||
| constexpr uint32_t MAX_K_VALUE_A2 = 10; | ||
| constexpr uint32_t MAX_K_VALUE_A2 = 16; |
There was a problem hiding this comment.
While increasing MAX_K_VALUE_A2 to 16 successfully unblocks custom model configurations with topk > 10 on A2 hardware, please note that K_MAX (defined at line 127 as constexpr int64_t K_MAX = 10;) is still set to 10.
K_MAX is used in CheckTensorShape (line 599) to validate the k dimension of expertIdStorageShape:
OP_TILING_CHECK(
(expertIdsDim1 <= 0) || (expertIdsDim1 > K_MAX),
...
);Since CheckTensorShape is called during A3 tiling (MoeDistributeDispatchA3TilingFuncImpl), leaving K_MAX at 10 will cause tiling to fail on A3 hardware when using custom models with topk > 10 (e.g., topk=14).
To ensure consistent support for topk > 10 across both A2 and A3 hardware, please also update K_MAX to 16 at line 127.
When adapting a custom model with topk=14, dispatch tiling failed due to a hardcoded upper bound MAX_K_VALUE_A2 = 10 in the A2 dispatch path. This PR raises the limit to 16 to unblock custom model configurations with topk > 10.
Change[dispatch_normal_a2_tiling.cpp]increase MAX_K_VALUE_A2 from 10 to 16
Root Cause:MAX_K_VALUE_A2 was originally set to match the standard model's topk value. Custom models with higher topk (e.g., 14) would exceed this limit and hit an assertion/error during tiling.
Test:Verified dispatch tiling succeeds with a custom model using topk=14 on A2 hardware.