fix bug in dispatch_ffn_combine kernel

[LQLlulu]

fix bug in dispatch_ffn_combine kernel

What this PR does / why we need it?

Does this PR introduce any user-facing change?

How was this patch tested?

[gemini-code-assist]

Summary of Changes

Hello @LQLlulu, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a bug in the dispatch_ffn_combine kernel by overhauling its epilogue processing. The changes ensure that swglu operations are correctly applied across different stages, with precise token accumulation and proper cross-core synchronization, leading to more robust and accurate kernel execution.

Highlights

• Epilogue Logic Refactoring: The core epilogue processing logic within the DispatchFFNCombineKernel has been significantly refactored to correctly handle swglu (Swish-Gated Linear Unit) operations across different stages.
• Token Sum Tracking: New variables, dequantSum1 and dequantSum2, have been introduced to accurately track token counts for the distinct swglu processing phases.
• Synchronization Adjustment: Cross-core synchronization and flag setting mechanisms have been adjusted and streamlined, removing redundant calls within the main loop and aligning with the new two-phase epilogue structure.
• Corrected currentM Calculation: The calculation of currentM has been updated to ensure accurate token processing throughout the kernel execution.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the epilogue logic in the dispatch_ffn_combine kernel to fix a bug. While the new two-stage swiglu logic seems to address issues with token handling, I've found multiple critical issues introduced by the changes. A synchronization variable is not being incremented, which will likely cause a deadlock. Additionally, there are incorrect function calls that will lead to compilation errors. Please address these critical issues.

As per the repository's style guide, I'm suggesting an updated title and summary for this pull request:

Suggested PR Title:

[Ops][BugFix] Fix bug in dispatch_ffn_combine kernel

Suggested PR Summary:

### What this PR does / why we need it?
This PR fixes a bug in the `dispatch_ffn_combine` kernel's epilogue logic for AIV cores. The previous implementation had issues with token handling and synchronization, particularly when `epilogueGranularity` was used. This PR refactors the epilogue into a two-stage swiglu process, correcting the calculation of tokens for each stage and ensuring proper truncation at `maxOutputSize`.

### Does this PR introduce _any_ user-facing change?
No. This is a bug fix in a kernel implementation and does not alter any user-facing APIs or behaviors.

### How was this patch tested?
CI should be used to validate the fix. It is recommended to add a specific test case that reproduces the original bug to prevent regressions.

[gemini-code-assist]

This line introduces two issues:

1. The synchronization variable syncgmm1Idx is no longer incremented inside the loop. This will break the synchronization protocol with the GMM1 function, which expects a unique, incrementing flag for each iteration, likely leading to a deadlock or incorrect execution.
2. The statement prevGroupSum1 += currentM; is appended to the end of the CrossCoreSetFlag call. This is poor coding style and harms readability. It should be on its own line.

            AscendC::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncgmm1Idx / CROSS_CORE_FLAG_MAX_SET_COUNT);   // V notifies C that the current communication round is complete
            syncgmm1Idx++;
            prevGroupSum1 += currentM;

[gemini-code-assist]

The variable blockEpilogue is being called as blockEpilogue1, but it was declared as blockEpilogue on line 660. This will cause a compilation error.

            blockEpilogue(gmC[gmOffsetC], shapeC, gmPerTokenScale1[rowStartThisCore], gmPermutedToken[gmOffsetD], gmPerTokenScale2[rowStartThisCore], params.epilogueCoreNum);

[gemini-code-assist]

Similar to the previous comment, the variable blockEpilogue is being called as blockEpilogue1. It should be blockEpilogue as declared on line 660.

                blockEpilogue(gmC[gmOffsetC], shapeC, gmPerTokenScale1[rowStartThisCore], gmPermutedToken[gmOffsetD], gmPerTokenScale2[rowStartThisCore], coreNum);