Skip to content

[DeepEP] increase MAX_K_VALUE_A2 from 10 to 16 to support topk > 10#593

Open
Yuxin1999 wants to merge 1 commit into
sgl-project:mainfrom
Yuxin1999:main
Open

[DeepEP] increase MAX_K_VALUE_A2 from 10 to 16 to support topk > 10#593
Yuxin1999 wants to merge 1 commit into
sgl-project:mainfrom
Yuxin1999:main

Conversation

@Yuxin1999

Copy link
Copy Markdown

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.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant