NVIDIA / NVIDIA/Megatron-LM

[BUG]MoE FA supporting for DS V2 Lite(16 attention header) is incorrect

Open
#1,749 1 comment 0 reactions 0 assignees View on GitHub
bug community-request module: moe
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**

for MoE, FA for deepseek V3(128 att header) is correct

name param: B param: N_Q param: H_Q param: N_KV param: H_KV param: d_h
FusedAttnFwd 1 4096 128 4096 128 192


While deepseek V2 lite is not, could you please help to fix? It mess seq_len with mbs
param: B param: N_Q param: H_Q param: N_KV param: H_KV param: d_h
4096 10 16 10 16 192

<!--table
{mso-displayed-decimal-separator:"\.";
mso-displayed-thousand-separator:"\,";}
@page
{mso-header-data:"&L&\0022Calibri\0022&10&K0000FF \[AMD Official Use Only - AMD Product\]&1\#\000D";
margin:.75in .7in .75in .7in;
mso-header-margin:.3in;
mso-footer-margin:.3in;}
tr
{mso-height-source:auto;}
col
{mso-width-source:auto;}
br
{mso-data-placement:same-cell;}
td
{padding-top:1px;
padding-right:1px;
padding-left:1px;
mso-ignore:padding;
color:black;
font-size:11.0pt;
font-weight:400;
font-style:normal;
text-decoration:none;
font-family:Calibri, sans-serif;
mso-font-charset:0;
mso-number-format:General;
text-align:general;
vertical-align:bottom;
border:none;
mso-background-source:auto;
mso-pattern:auto;
mso-protection:locked visible;
white-space:nowrap;
mso-rotate:0;}
.xl65
{font-weight:700;
text-align:center;
vertical-align:top;
border:.5pt solid windowtext;}
.xl66
{background:yellow;
mso-pattern:black none;}
-->

param: B | param: N_Q | param: H_Q | param: N_KV | param: H_KV | param: d_h
-- | -- | -- | -- | -- | --
4096 | 10 | 16 | 10 | 16 | 192

Could you please help to fix?

**To Reproduce**
Run deepseek V2 lite in megatron-lm
https://github.com/ROCm/MAD/blob/develop/benchmark/megatron_lm/README.md

**Expected behavior**
Will see mbs and seq-len parse is opposite.

**Stack trace/logs**
https://amdcloud.sharepoint.com/sites/DatacenterPerfTeam/_layouts/15/download.aspx?UniqueId=ffd5757e4d41435b9a39d5774129cddf&e=XvOE7V

https://amdcloud.sharepoint.com/:x:/r/sites/DatacenterPerfTeam/Shared%20Documents/General/AI/Training/AI_Training/MI350_PostLaunch/PerfAnalysis/Megatron-LM/DeepSeek_V2_Lite/MI355X-dsV2Lite-prbf16seq4096tp1pp1ep8mbs4gbs-betaImg/[3800835329.xlsx](https://amdcloud.sharepoint.com/:x:/r/sites/DatacenterPerfTeam/Shared%20Documents/General/AI/Training/AI_Training/MI350_PostLaunch/PerfAnalysis/Megatron-LM/DeepSeek_V2_Lite/MI355X-dsV2Lite-prbf16seq4096tp1pp1ep8mbs4gbs-betaImg/3800835329.xlsx?d=w0054799547a94b7f8628c94a4d284bbe&csf=1&web=1&e=0Nk5Yb)?d=w0054799547a94b7f8628c94a4d284bbe&csf=1&web=1&e=0Nk5Yb

**Environment (please complete the following information):**

**Proposed fix**
If you have a proposal for how to fix the issue state it here or link to a PR.

**Additional context**
Add any other context about the problem here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.